Default to py3 for TF docker images.

tensorflow · Jul 6, 2020 · 9fca004 · 9fca004
1 parent 59eca3f
commit 9fca004
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 6 deletions.
diff --git a/tensorflow_cloud/containerize.py b/tensorflow_cloud/containerize.py
@@ -139,6 +139,7 @@ def _create_docker_file(self):
                 != machine_config.AcceleratorType.NO_ACCELERATOR
             ):
                 self.docker_base_image += "-gpu"
+            self.docker_base_image += "-py3"
 
         if not self._base_image_exist():
             if "dev" in self.docker_base_image:
@@ -344,7 +345,8 @@ def _get_logs(self, logs_generator, name, image_uri):
             if "error" in chunk:
                 raise RuntimeError(
                     "Docker image {} failed: {}\nImage URI: {}".format(
-                        name, str(chunk["error"]), image_uri)
+                        name, str(chunk["error"]), image_uri
+                    )
                 )
 
 

diff --git a/tests/tensorflow_cloud/containerize_test.py b/tests/tensorflow_cloud/containerize_test.py
@@ -59,7 +59,7 @@ def test_create_docker_file_defaults(self):
         lcb._create_docker_file()
 
         expected_docker_file_lines = [
-            "FROM tensorflow/tensorflow:{}-gpu\n".format(VERSION),
+            "FROM tensorflow/tensorflow:{}-gpu-py3\n".format(VERSION),
             "WORKDIR /app/\n",
             "COPY /app/ /app/\n",
             'ENTRYPOINT ["python", "mnist_example_using_fit.py"]',
@@ -85,7 +85,7 @@ def test_create_docker_with_requirements(self):
         lcb._create_docker_file()
 
         expected_docker_file_lines = [
-            "FROM tensorflow/tensorflow:{}-gpu\n".format(VERSION),
+            "FROM tensorflow/tensorflow:{}-gpu-py3\n".format(VERSION),
             "WORKDIR /app/\n",
             "COPY /app/requirements.txt /app/requirements.txt\n",
             "RUN if [ -e requirements.txt ]; "
@@ -112,7 +112,7 @@ def test_create_docker_file_with_destination_dir(self):
         lcb._create_docker_file()
 
         expected_docker_file_lines = [
-            "FROM tensorflow/tensorflow:{}-gpu\n".format(VERSION),
+            "FROM tensorflow/tensorflow:{}-gpu-py3\n".format(VERSION),
             "WORKDIR /my_app/temp/\n",
             "COPY /my_app/temp/ /my_app/temp/\n",
             'ENTRYPOINT ["python", "mnist_example_using_fit.py"]',
@@ -199,7 +199,7 @@ def test_create_docker_file_with_cpu_config(self):
         lcb._create_docker_file()
 
         expected_docker_file_lines = [
-            "FROM tensorflow/tensorflow:{}\n".format(VERSION),
+            "FROM tensorflow/tensorflow:{}-py3\n".format(VERSION),
             "WORKDIR /app/\n",
             "COPY /app/ /app/\n",
             'ENTRYPOINT ["python", "mnist_example_using_fit.py"]',
@@ -220,7 +220,7 @@ def test_create_docker_file_with_tpu_config(self):
         lcb._create_docker_file()
 
         expected_docker_file_lines = [
-            "FROM tensorflow/tensorflow:{}\n".format(VERSION),
+            "FROM tensorflow/tensorflow:{}-py3\n".format(VERSION),
             "WORKDIR /app/\n",
             "RUN pip install cloud-tpu-client\n",
             "COPY /app/ /app/\n",