Docker builds for other CUDA versions

CHANGELOG.md

@@ -10,6 +10,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ### Added
 
+- Added Docker builds for other torch-supported versions of CUDA.
 - Adds [`allennlp-semparse`](https://github.com/allenai/allennlp-semparse) as an official, default plugin.
 
 ### Fixed

Dockerfile

@@ -18,9 +18,14 @@ LABEL com.nvidia.volumes.needed="nvidia_driver"
 
 WORKDIR /stage/allennlp
 
+# Install torch first. This build arg should be in the from of a version requirement,
+# like '==1.7' or '==1.7+cu102'.
+ARG TORCH
+RUN pip install --no-cache-dir torch${TORCH} -f https://download.pytorch.org/whl/torch_stable.html
+
 # Install the wheel of AllenNLP.
 COPY dist dist/
-RUN pip install $(ls dist/*.whl)
+RUN pip install --no-cache-dir $(ls dist/*.whl)
 # TODO(epwalsh): In PyTorch 1.7, dataclasses is an unconditional dependency, when it should
 # only be a conditional dependency for Python < 3.7.
 # This has been fixed on PyTorch master branch, so we should be able to

Dockerfile.test

@@ -17,6 +17,11 @@ LABEL com.nvidia.volumes.needed="nvidia_driver"
 
 WORKDIR /stage/allennlp
 
+# Install torch first. This build arg should be in the from of a version requirement,
+# like '==1.7' or '==1.7+cu102'.
+ARG TORCH
+RUN pip install --no-cache-dir torch${TORCH} -f https://download.pytorch.org/whl/torch_stable.html
+
 # Installing AllenNLP's dependencies is the most time-consuming part of building
 # this Docker image, so we make use of layer caching here by adding the minimal files
 # necessary to install the dependencies. Since most of the dependencies are defined

Makefile

@@ -13,6 +13,8 @@ MD_DOCS_EXTRAS = $(addprefix $(MD_DOCS_ROOT),README.md CHANGELOG.md CONTRIBUTING
 DOCKER_TAG = latest
 DOCKER_IMAGE_NAME = allennlp/allennlp:$(DOCKER_TAG)
 DOCKER_TEST_IMAGE_NAME = allennlp/test:$(DOCKER_TAG)
+# Our self-hosted runner currently has CUDA 11.0.
+DOCKER_TEST_TORCH_VERSION = '==1.7.0+cu110'
 DOCKER_RUN_CMD = docker run --rm \
 		-v $$HOME/.allennlp:/root/.allennlp \
 		-v $$HOME/.cache/torch:/root/.cache/torch \
@@ -139,17 +141,22 @@ clean :
 .PHONY : docker-image
 docker-image :
 	docker build \
-			--pull \
-			-f Dockerfile \
-			-t $(DOCKER_IMAGE_NAME) .
+		--pull \
+		-f Dockerfile \
+		--build-arg TORCH=$(DOCKER_TORCH_VERSION) \
+		-t $(DOCKER_IMAGE_NAME) .
 
 .PHONY : docker-run
 docker-run :
 	$(DOCKER_RUN_CMD) $(DOCKER_IMAGE_NAME) $(ARGS)
 
 .PHONY : docker-test-image
 docker-test-image :
-	docker build --pull -f Dockerfile.test -t $(DOCKER_TEST_IMAGE_NAME) .
+	docker build \
+		--pull \
+		-f Dockerfile.test \
+		--build-arg TORCH=$(DOCKER_TEST_TORCH_VERSION) \
+		-t $(DOCKER_TEST_IMAGE_NAME) .
 
 .PHONY : docker-test-run
 docker-test-run :