py build notebook images (kubeflow#791)

StatCan · May 11, 2018 · 5bffc37 · 5bffc37
1 parent 4ea0ce7
commit 5bffc37
Show file tree

Hide file tree

Showing 11 changed files with 88 additions and 13 deletions.
diff --git a/components/README.md b/components/README.md
@@ -5,8 +5,8 @@ This allows us to use `build_image.py` to build all images.
 
 For example
 ```
-python build_image.py tf_serving --tf_version=1.6 --platform=gpu
-python build_image.py tf_serving --tf_version=1.7  // default is cpu image
+python build_image.py --tf_version=1.6 --platform=gpu tf_serving
+python build_image.py --tf_version=1.4.1 tf_notebook
 ```
 
 See `build_image.py` for details.
@@ -16,3 +16,7 @@ See `build_image.py` for details.
 ### TF Serving
 - CPU: 1.4, 1.5, 1.6, 1.7
 - GPU: 1.6, 1.7
+
+### TF notebook
+- CPU: 1.4.1, 1.5.1, 1.6.0, 1.7.0
+- GPU: 1.4.1, 1.5.1, 1.6.0, 1.7.0
diff --git a/components/build_image.py b/components/build_image.py
@@ -1,7 +1,7 @@
 """Script to build images.
 
 For example,
-python build_image.py tf_serving --tf_version=1.6 --platform=gpu
+python build_image.py --tf_version=1.6 --platform=gpu tf_serving
 """
 import argparse
 import datetime
@@ -70,14 +70,18 @@ def get_build_args(config):
   config_list = [key + "=" + val for key, val in config.items()]
   return list(chain.from_iterable([["--build-arg", x] for x in config_list]))
 
-def build_tf_serving(args):
-  context_dir = "k8s-model-server/images"
-  version = args.tf_version if args.platform == "cpu" else args.tf_version + "gpu"
-
+def get_config(context_dir, version):
+  """Returns a dict of configuration from the version-config file."""
   config_file = os.path.join(context_dir, "versions", version, "version-config.json")
   with open(config_file) as f:
     config = yaml.load(f)
+  return config
+
+def build_tf_serving(args):
+  context_dir = "k8s-model-server/images"
+  version = args.tf_version if args.platform == "cpu" else args.tf_version + "gpu"
 
+  config = get_config(context_dir, version)
   build_args = get_build_args(config)
 
   command = list(chain(
@@ -88,6 +92,22 @@ def build_tf_serving(args):
   ))
   run(command, cwd=context_dir)
 
+def build_tf_notebook(args):
+  context_dir = "tensorflow-notebook-image"
+  version = args.tf_version if args.platform == "cpu" else args.tf_version + "gpu"
+
+  config = get_config(context_dir, version)
+  build_args = get_build_args(config)
+
+  command = list(chain(
+      ["docker", "build", "--pull"],
+      build_args,
+      ["-t", "{}/tensorflow-{}-notebook-{}:{}".format(
+          args.registry, args.tf_version, args.platform, args.tag),
+       "-f", "Dockerfile", "."]
+  ))
+  run(command, cwd=context_dir)
+
 def main():
   parser = argparse.ArgumentParser()
   subparsers = parser.add_subparsers()
@@ -102,21 +122,23 @@ def main():
     default="latest",
     help="The image tag"
   )
-
-  parser_tf_serving = subparsers.add_parser("tf_serving")
-  parser_tf_serving.set_defaults(func=build_tf_serving)
-
-  parser_tf_serving.add_argument(
+  parser.add_argument(
     "--tf_version",
     default="1.6",
     help="Tensorflow version"
   )
-  parser_tf_serving.add_argument(
+  parser.add_argument(
     "--platform",
     default="cpu",
     help="cpu or gpu"
   )
 
+  parser_tf_serving = subparsers.add_parser("tf_serving")
+  parser_tf_serving.set_defaults(func=build_tf_serving)
+
+  parser_tf_notebook = subparsers.add_parser("tf_notebook")
+  parser_tf_notebook.set_defaults(func=build_tf_notebook)
+
   args = parser.parse_args()
   args.func(args)
 

diff --git a/components/tensorflow-notebook-image/Dockerfile b/components/tensorflow-notebook-image/Dockerfile
@@ -30,6 +30,7 @@ RUN apt-get update && apt-get install -yq --no-install-recommends \
   curl \
   g++ \
   git \
+  gnupg \
   graphviz \
   locales \
   lsb-release \

diff --git a/components/tensorflow-notebook-image/versions/1.4.1/version-config.json b/components/tensorflow-notebook-image/versions/1.4.1/version-config.json
@@ -0,0 +1,6 @@
+{
+  "BASE_IMAGE": "ubuntu:latest",
+  "TF_PACKAGE": "https://storage.googleapis.com/tensorflow/linux/cpu/tensorflow-1.4.1-cp36-cp36m-linux_x86_64.whl",
+  "TF_PACKAGE_PY_27": "https://storage.googleapis.com/tensorflow/linux/cpu/tensorflow-1.4.1-cp27-none-linux_x86_64.whl",
+  "INSTALL_TFMA": "no"
+}
diff --git a/components/tensorflow-notebook-image/versions/1.4.1gpu/version-config.json b/components/tensorflow-notebook-image/versions/1.4.1gpu/version-config.json
@@ -0,0 +1,6 @@
+{
+  "BASE_IMAGE": "nvidia/cuda:8.0-cudnn6-devel-ubuntu16.04",
+  "TF_PACKAGE": "https://storage.googleapis.com/tensorflow/linux/gpu/tensorflow_gpu-1.4.1-cp36-cp36m-linux_x86_64.whl",
+  "TF_PACKAGE_PY_27": "https://storage.googleapis.com/tensorflow/linux/gpu/tensorflow_gpu-1.4.1-cp27-none-linux_x86_64.whl",
+  "INSTALL_TFMA": "no"
+}
diff --git a/components/tensorflow-notebook-image/versions/1.5.1/version-config.json b/components/tensorflow-notebook-image/versions/1.5.1/version-config.json
@@ -0,0 +1,6 @@
+{
+  "BASE_IMAGE": "ubuntu:latest",
+  "TF_PACKAGE": "https://storage.googleapis.com/tensorflow/linux/cpu/tensorflow-1.5.1-cp36-cp36m-linux_x86_64.whl",
+  "TF_PACKAGE_PY_27": "https://storage.googleapis.com/tensorflow/linux/cpu/tensorflow-1.5.1-cp27-none-linux_x86_64.whl",
+  "INSTALL_TFMA": "no"
+}
diff --git a/components/tensorflow-notebook-image/versions/1.5.1gpu/version-config.json b/components/tensorflow-notebook-image/versions/1.5.1gpu/version-config.json
@@ -0,0 +1,6 @@
+{
+  "BASE_IMAGE": "nvidia/cuda:9.0-cudnn7-devel-ubuntu16.04",
+  "TF_PACKAGE": "https://storage.googleapis.com/tensorflow/linux/gpu/tensorflow_gpu-1.5.1-cp36-cp36m-linux_x86_64.whl",
+  "TF_PACKAGE_PY_27": "https://storage.googleapis.com/tensorflow/linux/gpu/tensorflow_gpu-1.5.1-cp27-none-linux_x86_64.whl",
+  "INSTALL_TFMA": "no"
+}
diff --git a/components/tensorflow-notebook-image/versions/1.6.0/version-config.json b/components/tensorflow-notebook-image/versions/1.6.0/version-config.json
@@ -0,0 +1,6 @@
+{
+  "BASE_IMAGE": "ubuntu:latest",
+  "TF_PACKAGE": "https://storage.googleapis.com/tensorflow/linux/cpu/tensorflow-1.6.0-cp36-cp36m-linux_x86_64.whl",
+  "TF_PACKAGE_PY_27": "https://storage.googleapis.com/tensorflow/linux/cpu/tensorflow-1.6.0-cp27-none-linux_x86_64.whl",
+  "INSTALL_TFMA": "yes"
+}
diff --git a/components/tensorflow-notebook-image/versions/1.6.0gpu/version-config.json b/components/tensorflow-notebook-image/versions/1.6.0gpu/version-config.json
@@ -0,0 +1,6 @@
+{
+  "BASE_IMAGE": "nvidia/cuda:9.0-cudnn7-devel-ubuntu16.04",
+  "TF_PACKAGE": "https://storage.googleapis.com/tensorflow/linux/gpu/tensorflow_gpu-1.6.0-cp36-cp36m-linux_x86_64.whl",
+  "TF_PACKAGE_PY_27": "https://storage.googleapis.com/tensorflow/linux/gpu/tensorflow_gpu-1.6.0-cp27-none-linux_x86_64.whl",
+  "INSTALL_TFMA": "yes"
+}
diff --git a/components/tensorflow-notebook-image/versions/1.7.0/version-config.json b/components/tensorflow-notebook-image/versions/1.7.0/version-config.json
@@ -0,0 +1,6 @@
+{
+  "BASE_IMAGE": "ubuntu:latest",
+  "TF_PACKAGE": "https://storage.googleapis.com/tensorflow/linux/cpu/tensorflow-1.7.0-cp36-cp36m-linux_x86_64.whl",
+  "TF_PACKAGE_PY_27": "https://storage.googleapis.com/tensorflow/linux/cpu/tensorflow-1.7.0-cp27-none-linux_x86_64.whl",
+  "INSTALL_TFMA": "yes"
+}
diff --git a/components/tensorflow-notebook-image/versions/1.7.0gpu/version-config.json b/components/tensorflow-notebook-image/versions/1.7.0gpu/version-config.json
@@ -0,0 +1,6 @@
+{
+  "BASE_IMAGE": "nvidia/cuda:9.0-cudnn7-devel-ubuntu16.04",
+  "TF_PACKAGE": "https://storage.googleapis.com/tensorflow/linux/gpu/tensorflow_gpu-1.7.0-cp36-cp36m-linux_x86_64.whl",
+  "TF_PACKAGE_PY_27": "https://storage.googleapis.com/tensorflow/linux/gpu/tensorflow_gpu-1.7.0-cp27-none-linux_x86_64.whl",
+  "INSTALL_TFMA": "yes"
+}