Switch to single-pip-package and install nvprune

angerson · Jan 28, 2021 · 4853be0 · 4853be0
1 parent 8bc0c84
commit 4853be0
Show file tree

Hide file tree

Showing 6 changed files with 11 additions and 4 deletions.
diff --git a/tensorflow/tools/dockerfiles/dockerfiles/devel-gpu-jupyter.Dockerfile b/tensorflow/tools/dockerfiles/dockerfiles/devel-gpu-jupyter.Dockerfile
@@ -41,6 +41,7 @@ RUN apt-get update && apt-get install -y --no-install-recommends \
         cuda-command-line-tools-${CUDA/./-} \
         libcublas-${CUDA/./-} \
         libcublas-dev-${CUDA/./-} \
+        cuda-nvprune-${CUDA/./-} \
         cuda-nvrtc-${CUDA/./-} \
         cuda-nvrtc-dev-${CUDA/./-} \
         cuda-cudart-dev-${CUDA/./-} \

diff --git a/tensorflow/tools/dockerfiles/dockerfiles/devel-gpu.Dockerfile b/tensorflow/tools/dockerfiles/dockerfiles/devel-gpu.Dockerfile
@@ -41,6 +41,7 @@ RUN apt-get update && apt-get install -y --no-install-recommends \
         cuda-command-line-tools-${CUDA/./-} \
         libcublas-${CUDA/./-} \
         libcublas-dev-${CUDA/./-} \
+        cuda-nvprune-${CUDA/./-} \
         cuda-nvrtc-${CUDA/./-} \
         cuda-nvrtc-dev-${CUDA/./-} \
         cuda-cudart-dev-${CUDA/./-} \

diff --git a/tensorflow/tools/dockerfiles/dockerfiles/ppc64le/devel-gpu-ppc64le-jupyter.Dockerfile b/tensorflow/tools/dockerfiles/dockerfiles/ppc64le/devel-gpu-ppc64le-jupyter.Dockerfile
@@ -41,6 +41,7 @@ RUN apt-get update && apt-get install -y --no-install-recommends \
         cuda-command-line-tools-${CUDA/./-} \
         libcublas-${CUDA/./-} \
         libcublas-dev-${CUDA/./-} \
+        cuda-nvprune-${CUDA/./-} \
         cuda-nvrtc-${CUDA/./-} \
         cuda-nvrtc-dev-${CUDA/./-} \
         cuda-cudart-dev-${CUDA/./-} \

diff --git a/tensorflow/tools/dockerfiles/dockerfiles/ppc64le/devel-gpu-ppc64le.Dockerfile b/tensorflow/tools/dockerfiles/dockerfiles/ppc64le/devel-gpu-ppc64le.Dockerfile
@@ -41,6 +41,7 @@ RUN apt-get update && apt-get install -y --no-install-recommends \
         cuda-command-line-tools-${CUDA/./-} \
         libcublas-${CUDA/./-} \
         libcublas-dev-${CUDA/./-} \
+        cuda-nvprune-${CUDA/./-} \
         cuda-nvrtc-${CUDA/./-} \
         cuda-nvrtc-dev-${CUDA/./-} \
         cuda-cudart-dev-${CUDA/./-} \

diff --git a/tensorflow/tools/dockerfiles/partials/ubuntu/devel-nvidia.partial.Dockerfile b/tensorflow/tools/dockerfiles/partials/ubuntu/devel-nvidia.partial.Dockerfile
@@ -18,6 +18,7 @@ RUN apt-get update && apt-get install -y --no-install-recommends \
         cuda-command-line-tools-${CUDA/./-} \
         libcublas-${CUDA/./-} \
         libcublas-dev-${CUDA/./-} \
+        cuda-nvprune-${CUDA/./-} \
         cuda-nvrtc-${CUDA/./-} \
         cuda-nvrtc-dev-${CUDA/./-} \
         cuda-cudart-dev-${CUDA/./-} \

diff --git a/tensorflow/tools/dockerfiles/spec.yml b/tensorflow/tools/dockerfiles/spec.yml
@@ -94,6 +94,8 @@ slice_sets:
     ubuntu:
         - add_to_name: ""
           dockerfile_exclusive_name: "cpu"
+          args:
+              - TF_PACKAGE=tensorflow-cpu
           partials:
               - ubuntu/version
               - ubuntu/cpu
@@ -103,7 +105,7 @@ slice_sets:
         - add_to_name: "-gpu"
           dockerfile_exclusive_name: "gpu"
           args:
-              - TF_PACKAGE=tensorflow-gpu
+              - TF_PACKAGE=tensorflow
           partials:
               - ubuntu/version
               - ubuntu/nvidia
@@ -595,7 +597,7 @@ slice_sets:
               - UBUNTU_VERSION=18.04
               - ARCH=ppc64le
               - CUDA=10.0
-              - TF_PACKAGE=tensorflow-gpu
+              - TF_PACKAGE=tensorflow
           partials:
               - ubuntu/version
               - ubuntu/nvidia
@@ -649,7 +651,7 @@ slice_sets:
               - tensorflow
               - shell
           args:
-              - TF_PACKAGE=tf-nightly
+              - TF_PACKAGE=tf-nightly-cpu
           tests:
               - import.sh
         - add_to_name: "nightly-gpu"
@@ -663,4 +665,4 @@ slice_sets:
           tests:
               - import-gpu.sh
           args:
-              - TF_PACKAGE=tf-nightly-gpu
+              - TF_PACKAGE=tf-nightly