define env var in compose file to restrict gpus

ArneBinder · Apr 26, 2018 · a929661 · a929661
1 parent 2d75f5a
commit a929661
Show file tree

Hide file tree

Showing 5 changed files with 4 additions and 14 deletions.
diff --git a/docker/tools/visualize/Dockerfile.tf1_3_gpu b/docker/tools/visualize/Dockerfile.tf1_3_gpu
@@ -170,8 +170,3 @@ VOLUME /root/train
 VOLUME /root/corpora_out
 
 #ENTRYPOINT ["/root/recursive-embedding/set-user-with-folder.sh", "/root/train", "python", "train_fold.py"]
-
-# pass as env variable
-ARG _NVIDIA_VISIBLE_DEVICES=0
-RUN echo NVIDIA_VISIBLE_DEVICES="$_NVIDIA_VISIBLE_DEVICES"
-ENV NVIDIA_VISIBLE_DEVICES="$_NVIDIA_VISIBLE_DEVICES"
diff --git a/docker/tools/visualize/docker-compose.yml b/docker/tools/visualize/docker-compose.yml
@@ -8,7 +8,8 @@ services:
       args:
           OWN_LOCATION: docker/tools/visualize
           PROJECT_ROOT: /root/recursive-embedding
-          _NVIDIA_VISIBLE_DEVICES: ${NVIDIA_VISIBLE_DEVICES}
+    environment:
+      - NVIDIA_VISIBLE_DEVICES=${NVIDIA_VISIBLE_DEVICES}
     ports:
       - ${HOST_PORT}:5000
     volumes:

diff --git a/docker/train/tensorflow-fold/Dockerfile.tf1_3_cpu_mkl b/docker/train/tensorflow-fold/Dockerfile.tf1_3_cpu_mkl
@@ -120,7 +120,5 @@ EXPOSE 8888
 
 VOLUME /root/train
 
-ENV NVIDIA_VISIBLE_DEVICES=""
-
 ENTRYPOINT ["/root/set-user-with-folder.sh", "/root/train", "python", "train_fold.py"]
 
diff --git a/docker/train/tensorflow-fold/Dockerfile.tf1_3_gpu b/docker/train/tensorflow-fold/Dockerfile.tf1_3_gpu
@@ -160,10 +160,5 @@ WORKDIR "$PROJECT_ROOT"/src
 
 VOLUME /root/train
 
-# pass as env variable
-ARG _NVIDIA_VISIBLE_DEVICES=0
-RUN echo NVIDIA_VISIBLE_DEVICES="$_NVIDIA_VISIBLE_DEVICES"
-ENV NVIDIA_VISIBLE_DEVICES="$_NVIDIA_VISIBLE_DEVICES"
-
 ENTRYPOINT ["/root/set-user-with-folder.sh", "/root/train", "python", "train_fold.py"]
 
diff --git a/docker/train/tensorflow-fold/docker-compose.yml b/docker/train/tensorflow-fold/docker-compose.yml
@@ -8,7 +8,8 @@ services:
       args:
         OWN_LOCATION: docker/train/tensorflow-fold
         PROJECT_ROOT: /root/recursive-embedding
-        _NVIDIA_VISIBLE_DEVICES: ${NVIDIA_VISIBLE_DEVICES}
+    environment:
+      - NVIDIA_VISIBLE_DEVICES=${NVIDIA_VISIBLE_DEVICES}
     image: tensorflowfold:${DOCKERFILE_EXT}
     volumes:
       - ../../../src:/root/recursive-embedding/src
Original file line number	Diff line number	Diff line change
Expand Up		@@ -120,7 +120,5 @@ EXPOSE 8888

		VOLUME /root/train

		ENV NVIDIA_VISIBLE_DEVICES=""

		ENTRYPOINT ["/root/set-user-with-folder.sh", "/root/train", "python", "train_fold.py"]