From e63178f228305a865871d7c0133b0018edb9199f Mon Sep 17 00:00:00 2001
From: Sanjay Agravat <agravat@google.com>
Date: Sun, 14 Jun 2020 22:32:00 -0400
Subject: [PATCH 01/13] initial commit for DLVM feature for notebook, local,
 and cloud execution.

---
 caliban/cli.py           |  10 ++++
 caliban/cloud/core.py    |  22 ++++++++-
 caliban/docker.py        | 101 +++++++++++++++++++++++++++++++++++----
 caliban/history/utils.py |   5 +-
 caliban/main.py          |   6 +++
 5 files changed, 131 insertions(+), 13 deletions(-)

diff --git a/caliban/cli.py b/caliban/cli.py
index ea886b9..3ef2aa5 100644
--- a/caliban/cli.py
+++ b/caliban/cli.py
@@ -210,6 +210,13 @@ def image_tag_arg(parser):
       "Caliban will skip the build and push steps and use this image tag.")
 
 
+def dlvm_arg(parser):
+  parser.add_argument(
+      "--dlvm",
+      help="DLVM base image tag accessible via Container Registry. If supplied, "
+      "Caliban will skip the build and push steps and use this image tag.")
+
+
 def machine_type_arg(parser):
   machine_types = u.enum_vals(ct.MachineType)
   cpu_default = conf.DEFAULT_MACHINE_TYPE[conf.JobMode.CPU].value
@@ -282,6 +289,7 @@ def notebook_parser(base):
                            help="Run a local Jupyter notebook instance.")
   base_parser(parser)
   docker_run_arg(parser)
+  dlvm_arg(parser)
 
   # Custom notebook arguments.
   parser.add_argument(
@@ -314,6 +322,7 @@ def local_run_parser(base):
   """Configure the subparser for `caliban run`."""
   parser = base.add_parser("run", help="Run a job inside a Docker container.")
   executing_parser(parser)
+  dlvm_arg(parser)
   image_id_arg(parser)
   docker_run_arg(parser)
   xgroup_submit_arg(parser)
@@ -380,6 +389,7 @@ def container_parser(parser):
   image_tag_arg(parser)
   project_id_arg(parser)
   region_arg(parser)
+  dlvm_arg(parser)
   machine_type_arg(parser)
   gpu_spec_arg(parser)
   tpu_spec_arg(parser)
diff --git a/caliban/cloud/core.py b/caliban/cloud/core.py
index b14743c..a1c4f91 100644
--- a/caliban/cloud/core.py
+++ b/caliban/cloud/core.py
@@ -486,7 +486,7 @@ def build_job_specs(
                     experiments=experiments)
 
 
-def generate_image_tag(project_id, docker_args, dry_run: bool = False):
+def generate_image_tag(project_id, docker_args, dlvm: str = None, dry_run: bool = False):
   """Generates a new Docker image and pushes an image to the user's GCloud
   Container Repository, tagged using the UUID of the generated image.
 
@@ -500,6 +500,11 @@ def generate_image_tag(project_id, docker_args, dry_run: bool = False):
   if dry_run:
     logging.info("Dry run - skipping actual 'docker build' and 'docker push'.")
     image_tag = "dry_run_tag"
+  elif dlvm is not None:
+    logging.info("DLVM - {}".format(dlvm))
+    #image_tag = dlvm
+    image_id = d.build_dlvm_image(dlvm=dlvm, **docker_args)
+    image_tag = d.push_uuid_tag(project_id, image_id)
   else:
     image_id = d.build_image(**docker_args)
     image_tag = d.push_uuid_tag(project_id, image_id)
@@ -543,6 +548,14 @@ def submit_job_specs(
 
   execute_requests(requests, num_specs, num_retries=request_retries)
 
+def _dlvm_id(image: str) -> str:
+  if image == "pytorch":
+    return "gcr.io/deeplearning-platform-release/pytorch-cpu:latest"
+  elif image == "tf-21":
+    return "gcr.io/deeplearning-platform-release/tf2-cpu.2-1"
+  else:
+    return None
+
 
 def submit_ml_job(job_mode: conf.JobMode,
                   docker_args: Dict[str, Any],
@@ -551,6 +564,7 @@ def submit_ml_job(job_mode: conf.JobMode,
                   credentials_path: Optional[str] = None,
                   dry_run: bool = False,
                   job_name: Optional[str] = None,
+                  dlvm: Optional[str] = None,
                   machine_type: Optional[ct.MachineType] = None,
                   gpu_spec: Optional[ct.GPUSpec] = None,
                   tpu_spec: Optional[ct.TPUSpec] = None,
@@ -633,7 +647,11 @@ def submit_ml_job(job_mode: conf.JobMode,
   with session_scope(engine) as session:
     container_spec = generate_container_spec(session, docker_args, image_tag)
 
-    if image_tag is None:
+    dlvm_tag = None
+    if dlvm is not None:
+      dlvm_tag = dlvm # _dlvm_id(dlvm)
+      image_tag = generate_image_tag(project_id, docker_args, dlvm_tag, dry_run=dry_run)
+    elif image_tag is None:
       image_tag = generate_image_tag(project_id, docker_args, dry_run=dry_run)
 
     experiments = create_experiments(
diff --git a/caliban/docker.py b/caliban/docker.py
index be5299e..38d06bd 100644
--- a/caliban/docker.py
+++ b/caliban/docker.py
@@ -308,8 +308,11 @@ def _service_account_entry(user_id: int, user_group: int, credentials_path: str,
 COPY --chown={user_id}:{user_group} {credentials_path} {container_creds}
 
 # Use the credentials file to activate gcloud, gsutil inside the container.
-RUN gcloud auth activate-service-account --key-file={container_creds} && \
-  git config --global credential.'https://source.developers.google.com'.helper gcloud.sh
+#RUN mkdir -p /home/agravat/.config
+#RUN chown agravat -R /home/agravat/.config
+#RUN gcloud init
+#RUN gcloud auth activate-service-account --key-file={container_creds} && \
+#  git config --global credential.'https://source.developers.google.com'.helper gcloud.sh
 
 ENV GOOGLE_APPLICATION_CREDENTIALS={container_creds}
 """.format_map({
@@ -389,9 +392,14 @@ def _notebook_entries(lab: bool = False, version: Optional[str] = None) -> str:
 
   library = "jupyterlab" if lab else "jupyter"
 
-  return """
-RUN pip install {}{}
-""".format(library, version_suffix)
+  #return """
+#RUN pip install {}{}
+#""".format(library, version_suffix)
+  cmds = """
+  RUN /opt/conda/bin/pip install https://storage.googleapis.com/deeplearning-platform-ui-public/jupyterlab_gcpscheduler-1.0.0.tar.gz
+  RUN /opt/conda/bin/jupyter lab build
+  """
+  return cmds
 
 
 def _custom_packages(
@@ -458,10 +466,18 @@ def _extra_dir_entries(workdir: str, user_id: int, user_group: int,
     ret += "\n{}".format(_copy_dir_entry(workdir, user_id, user_group, d))
   return ret
 
+def _dlvm_id(image: str) -> str:
+  if image == "pytorch":
+    return "gcr.io/deeplearning-platform-release/pytorch-cpu:latest"
+  elif image == "tf-21":
+    return "gcr.io/deeplearning-platform-release/tf2-cpu.2-1"
+  else:
+    return None
 
 def _dockerfile_template(
     job_mode: c.JobMode,
     workdir: Optional[str] = None,
+    dlvm: Optional[str] = None,
     base_image_fn: Optional[Callable[[c.JobMode], str]] = None,
     package: Optional[Union[List, u.Package]] = None,
     requirements_path: Optional[str] = None,
@@ -507,7 +523,12 @@ def _dockerfile_template(
   if base_image_fn is None:
     base_image_fn = base_image_id
 
-  base_image = base_image_fn(job_mode)
+  if dlvm is None:
+    base_image = base_image_fn(job_mode)
+  else:
+    base_image = _dlvm_id(dlvm)
+
+  logging.info("base image ------ {}".format(base_image))
 
   dockerfile = """
 FROM {base_image}
@@ -578,6 +599,52 @@ def docker_image_id(output: str) -> ImageId:
   return ImageId(output.splitlines()[-1].split()[-1])
 
 
+def build_dlvm_image(job_mode: c.JobMode,
+                build_path: str,
+                credentials_path: Optional[str] = None,
+                adc_path: Optional[str] = None,
+                no_cache: bool = False,
+                dlvm: str = None,
+                **kwargs) -> str:
+  """Builds a Docker image by generating a Dockerfile and passing it to `docker
+  build` via stdin. All output from the `docker build` process prints to
+  stdout.
+
+  Returns the image ID of the new docker container; if the command fails,
+  throws on error with information about the command and any issues that caused
+  the problem.
+
+  """
+  with u.TempCopy(credentials_path,
+                  tmp_name=".caliban_default_creds.json") as creds:
+    with u.TempCopy(adc_path, tmp_name=".caliban_adc_creds.json") as adc:
+      cache_args = ["--no-cache"] if no_cache else []
+      cmd = ["docker", "build"] + cache_args + ["--rm", "-f-", build_path]
+      logging.info("extra: {}".format(kwargs))
+      #spec = {k: v for k, v in build_image_kwargs.items()}
+
+      dockerfile = _dockerfile_template(job_mode,
+                                        credentials_path=creds,
+                                        adc_path=adc,
+                                        dlvm=dlvm,
+                                        **kwargs)
+
+      joined_cmd = " ".join(cmd)
+      logging.info("Running command: {}".format(joined_cmd))
+
+      try:
+        output, ret_code = u.capture_stdout(cmd, input_str=dockerfile)
+        if ret_code == 0:
+          return docker_image_id(output)
+        else:
+          error_msg = "Docker failed with error code {}.".format(ret_code)
+          raise DockerError(error_msg, cmd, ret_code)
+
+      except subprocess.CalledProcessError as e:
+        logging.error(e.output)
+        logging.error(e.stderr)
+
+
 def build_image(job_mode: c.JobMode,
                 build_path: str,
                 credentials_path: Optional[str] = None,
@@ -883,6 +950,7 @@ def run(job_mode: c.JobMode,
         run_args: Optional[List[str]] = None,
         script_args: Optional[List[str]] = None,
         image_id: Optional[str] = None,
+        dlvm: Optional[str] = None,
         **build_image_kwargs) -> None:
   """Builds an image using the supplied **build_image_kwargs and calls `docker
   run` on the resulting image using sensible defaults.
@@ -901,7 +969,15 @@ def run(job_mode: c.JobMode,
   if script_args is None:
     script_args = []
 
-  if image_id is None:
+  logging.info("run_args {}".format(run_args))
+  logging.info("script_args {}".format(script_args))
+
+  if dlvm is not None:
+    logging.info("dlvm id: {}".format(dlvm))
+    image_id = build_dlvm_image(job_mode, dlvm=dlvm, **build_image_kwargs)
+    logging.info("image_id: {}".format(image_id))
+  elif image_id is None:
+    logging.info("def run - image__id is none")
     image_id = build_image(job_mode, **build_image_kwargs)
 
   base_cmd = _run_cmd(job_mode, run_args)
@@ -916,6 +992,7 @@ def run(job_mode: c.JobMode,
 def run_interactive(job_mode: c.JobMode,
                     workdir: Optional[str] = None,
                     image_id: Optional[str] = None,
+                    dlvm: Optional[str] = None,
                     run_args: Optional[List[str]] = None,
                     mount_home: Optional[bool] = None,
                     shell: Optional[Shell] = None,
@@ -963,13 +1040,14 @@ def run_interactive(job_mode: c.JobMode,
 
   interactive_run_args = _interactive_opts(workdir) + [
       "-it", \
-      "--entrypoint", entrypoint
+      #"--entrypoint", entrypoint
   ] + _home_mount_cmds(mount_home) + run_args
 
   run(job_mode=job_mode,
       run_args=interactive_run_args,
       script_args=entrypoint_args,
       image_id=image_id,
+      dlvm=dlvm,
       shell=shell,
       workdir=workdir,
       **build_image_kwargs)
@@ -980,6 +1058,7 @@ def run_notebook(job_mode: c.JobMode,
                  lab: Optional[bool] = None,
                  version: Optional[bool] = None,
                  run_args: Optional[List[str]] = None,
+                 dlvm: Optional[str] = None,
                  **run_interactive_kwargs) -> None:
   """Start a notebook in the current working directory; the process will run
   inside of a Docker container that's identical to the environment available to
@@ -1019,8 +1098,10 @@ def run_notebook(job_mode: c.JobMode,
   docker_args = ["-p", "{}:{}".format(port, port)] + run_args
 
   run_interactive(job_mode,
-                  entrypoint="/opt/venv/bin/python",
-                  entrypoint_args=jupyter_args,
+                  dlvm=dlvm,
+                  #entrypoint="/opt/venv/bin/python",
+                  entrypoint="/bin/bash",
+                  #entrypoint_args=jupyter_args,
                   run_args=docker_args,
                   inject_notebook=inject_arg,
                   jupyter_version=version,
diff --git a/caliban/history/utils.py b/caliban/history/utils.py
index 5b7737c..543e83f 100644
--- a/caliban/history/utils.py
+++ b/caliban/history/utils.py
@@ -180,6 +180,7 @@ def generate_container_spec(
     session: Session,
     docker_args: Dict[str, Any],
     image_tag: Optional[str] = None,
+    dlvm_tag: Optional[str] = None,
 ) -> ContainerSpec:
   '''generates a container spec
 
@@ -192,7 +193,9 @@ def generate_container_spec(
   ContainerSpec instance
   '''
 
-  if image_tag is None:
+  if dlvm_tag is not None:
+    spec = {'image_id': dlvm_tag}
+  elif image_tag is None:
     spec = docker_args
   else:
     spec = {'image_id': image_tag}
diff --git a/caliban/main.py b/caliban/main.py
index d1120ff..4762514 100644
--- a/caliban/main.py
+++ b/caliban/main.py
@@ -68,9 +68,11 @@ def run_app(arg_input):
   elif command == "notebook":
     port = args.get("port")
     lab = args.get("lab")
+    dlvm = args.get("dlvm")
     version = args.get("jupyter_version")
     mount_home = not args['bare']
     docker.run_notebook(job_mode,
+                        dlvm=dlvm,
                         port=port,
                         lab=lab,
                         version=version,
@@ -95,6 +97,7 @@ def run_app(arg_input):
     dry_run = args["dry_run"]
     package = args["module"]
     image_id = args.get("image_id")
+    dlvm = args.get("dlvm")
     exp_config = args.get("experiment_config")
     xgroup = args.get('xgroup')
 
@@ -102,6 +105,7 @@ def run_app(arg_input):
                            run_args=docker_run_args,
                            script_args=script_args,
                            image_id=image_id,
+                           dlvm=dlvm,
                            experiment_config=exp_config,
                            dry_run=dry_run,
                            package=package,
@@ -120,6 +124,7 @@ def run_app(arg_input):
     tpu_spec = args.get("tpu_spec")
     image_tag = args.get("image_tag")
     machine_type = args.get("machine_type")
+    dlvm = args.get("dlvm")
     exp_config = args.get("experiment_config")
     labels = u.sanitize_labels(args.get("label") or [])
     xgroup = args.get('xgroup')
@@ -135,6 +140,7 @@ def run_app(arg_input):
         credentials_path=cloud_key,
         dry_run=dry_run,
         job_name=job_name,
+        dlvm=dlvm,
         machine_type=machine_type,
         gpu_spec=gpu_spec,
         tpu_spec=tpu_spec,

From ac8eb97e96d7f6d41377c21956173e4838ba6f7f Mon Sep 17 00:00:00 2001
From: Sanjay Agravat <agravat@google.com>
Date: Wed, 17 Jun 2020 01:20:36 +0000
Subject: [PATCH 02/13] fixed issue for DLVM notebook scheduler extension
 installation.

---
 caliban/docker.py | 17 +++++++++++------
 1 file changed, 11 insertions(+), 6 deletions(-)

diff --git a/caliban/docker.py b/caliban/docker.py
index 38d06bd..27f4917 100644
--- a/caliban/docker.py
+++ b/caliban/docker.py
@@ -396,7 +396,7 @@ def _notebook_entries(lab: bool = False, version: Optional[str] = None) -> str:
 #RUN pip install {}{}
 #""".format(library, version_suffix)
   cmds = """
-  RUN /opt/conda/bin/pip install https://storage.googleapis.com/deeplearning-platform-ui-public/jupyterlab_gcpscheduler-1.0.0.tar.gz
+  RUN /opt/conda/bin/pip install --user --no-cache-dir https://storage.googleapis.com/deeplearning-platform-ui-public/jupyterlab_gcpscheduler-1.0.0.tar.gz
   RUN /opt/conda/bin/jupyter lab build
   """
   return cmds
@@ -546,8 +546,6 @@ def _dockerfile_template(
 ENV HOME={c_home}
 
 WORKDIR {workdir}
-
-USER {uid}:{gid}
 """.format_map({
       "base_image": base_image,
       "username": username,
@@ -561,6 +559,16 @@ def _dockerfile_template(
                                      gid,
                                      adc_path=adc_path,
                                      credentials_path=credentials_path)
+  if inject_notebook.value != 'none':
+    install_lab = inject_notebook == NotebookInstall.lab
+    dockerfile += _notebook_entries(lab=install_lab, version=jupyter_version)
+
+  dockerfile += """
+
+USER {uid}:{gid}
+""".format_map({"uid": uid,
+                "gid": gid
+   })
 
   dockerfile += _dependency_entries(workdir,
                                     uid,
@@ -568,9 +576,6 @@ def _dockerfile_template(
                                     requirements_path=requirements_path,
                                     setup_extras=setup_extras)
 
-  if inject_notebook.value != 'none':
-    install_lab = inject_notebook == NotebookInstall.lab
-    dockerfile += _notebook_entries(lab=install_lab, version=jupyter_version)
 
   if extra_dirs is not None:
     dockerfile += _extra_dir_entries(workdir, uid, gid, extra_dirs)

From 80ca15bcd48ca63360e1f2fa4d2112fea1453521 Mon Sep 17 00:00:00 2001
From: Sanjay Agravat <agravat@google.com>
Date: Wed, 17 Jun 2020 02:40:39 +0000
Subject: [PATCH 03/13] Updated entrypoint for interactive mode for regular and
 DLVM mode.

---
 caliban/docker.py | 40 +++++++++++++++++++++++++++-------------
 1 file changed, 27 insertions(+), 13 deletions(-)

diff --git a/caliban/docker.py b/caliban/docker.py
index 27f4917..b55362c 100644
--- a/caliban/docker.py
+++ b/caliban/docker.py
@@ -379,7 +379,7 @@ def _credentials_entries(user_id: int,
   return ret
 
 
-def _notebook_entries(lab: bool = False, version: Optional[str] = None) -> str:
+def _notebook_entries(lab: bool = False, version: Optional[str] = None, dlvm: bool = False) -> str:
   """Returns the Dockerfile entries necessary to install Jupyter{lab}.
 
   Optionally takes a version string.
@@ -392,14 +392,15 @@ def _notebook_entries(lab: bool = False, version: Optional[str] = None) -> str:
 
   library = "jupyterlab" if lab else "jupyter"
 
-  #return """
-#RUN pip install {}{}
-#""".format(library, version_suffix)
-  cmds = """
+  if not dlvm:
+    return """
+RUN pip install {}{}
+""".format(library, version_suffix)
+  else:
+    return """
   RUN /opt/conda/bin/pip install --user --no-cache-dir https://storage.googleapis.com/deeplearning-platform-ui-public/jupyterlab_gcpscheduler-1.0.0.tar.gz
   RUN /opt/conda/bin/jupyter lab build
   """
-  return cmds
 
 
 def _custom_packages(
@@ -561,7 +562,10 @@ def _dockerfile_template(
                                      credentials_path=credentials_path)
   if inject_notebook.value != 'none':
     install_lab = inject_notebook == NotebookInstall.lab
-    dockerfile += _notebook_entries(lab=install_lab, version=jupyter_version)
+    if dlvm is None:
+        dockerfile += _notebook_entries(lab=install_lab, version=jupyter_version, dlvm=False)
+    else:
+        dockerfile += _notebook_entries(lab=install_lab, version=jupyter_version, dlvm=True)
 
   dockerfile += """
 
@@ -1043,10 +1047,20 @@ def run_interactive(job_mode: c.JobMode,
   if entrypoint is None:
     entrypoint = SHELL_DICT[shell].executable
 
-  interactive_run_args = _interactive_opts(workdir) + [
+  if dlvm is None:
+    # Pass the default entrypoint if not using DLVM
+    # Otherwise the DLVM automatically runs jupyterlab so don't set an
+    # entrypoint.
+    interactive_run_args = _interactive_opts(workdir) + [
+      "-it", \
+      "--entrypoint", entrypoint
+    ] + _home_mount_cmds(mount_home) + run_args
+  else:
+    interactive_run_args = _interactive_opts(workdir) + [
       "-it", \
-      #"--entrypoint", entrypoint
-  ] + _home_mount_cmds(mount_home) + run_args
+    ] + _home_mount_cmds(mount_home) + run_args
+
+    entrypoint_args = []
 
   run(job_mode=job_mode,
       run_args=interactive_run_args,
@@ -1104,9 +1118,9 @@ def run_notebook(job_mode: c.JobMode,
 
   run_interactive(job_mode,
                   dlvm=dlvm,
-                  #entrypoint="/opt/venv/bin/python",
-                  entrypoint="/bin/bash",
-                  #entrypoint_args=jupyter_args,
+                  entrypoint="python",
+                  #entrypoint="/bin/bash",
+                  entrypoint_args=jupyter_args,
                   run_args=docker_args,
                   inject_notebook=inject_arg,
                   jupyter_version=version,

From c8e158213f1eea05a7296702b1455eee942dd2d3 Mon Sep 17 00:00:00 2001
From: Sanjay Agravat <agravat@google.com>
Date: Wed, 17 Jun 2020 03:04:12 +0000
Subject: [PATCH 04/13] Implemented shell mode for DLVM.

---
 caliban/cli.py    |  1 +
 caliban/docker.py | 21 ++++++++++++++++-----
 caliban/main.py   |  2 ++
 3 files changed, 19 insertions(+), 5 deletions(-)

diff --git a/caliban/cli.py b/caliban/cli.py
index 3ef2aa5..8bff0ce 100644
--- a/caliban/cli.py
+++ b/caliban/cli.py
@@ -267,6 +267,7 @@ def shell_parser(base):
   parser = base.add_parser(
       "shell", help="Start an interactive shell with this dir mounted.")
   base_parser(parser)
+  dlvm_arg(parser)
   image_id_arg(parser)
   docker_run_arg(parser)
   parser.add_argument(
diff --git a/caliban/docker.py b/caliban/docker.py
index b55362c..f02c239 100644
--- a/caliban/docker.py
+++ b/caliban/docker.py
@@ -1049,16 +1049,27 @@ def run_interactive(job_mode: c.JobMode,
 
   if dlvm is None:
     # Pass the default entrypoint if not using DLVM
-    # Otherwise the DLVM automatically runs jupyterlab so don't set an
-    # entrypoint.
+    # Otherwise set the DLVM entrypoint depending on if we are in notebook 
+    # mode or shell mode.
     interactive_run_args = _interactive_opts(workdir) + [
       "-it", \
       "--entrypoint", entrypoint
     ] + _home_mount_cmds(mount_home) + run_args
   else:
-    interactive_run_args = _interactive_opts(workdir) + [
-      "-it", \
-    ] + _home_mount_cmds(mount_home) + run_args
+    # Don't set an entrypoint if we are running the DLVM notebook
+    # Otherwise we are running in shell mode so we pass in the shell
+    # entrypoint args.
+    if shell is None:
+        interactive_run_args = _interactive_opts(workdir) + [
+            "-it", \
+        ] + _home_mount_cmds(mount_home) + run_args
+    else:
+        entrypoint = SHELL_DICT[shell].executable
+        interactive_run_args = _interactive_opts(workdir) + [
+            "-it", \
+            "--entrypoint", entrypoint
+        ] + _home_mount_cmds(mount_home) + run_args
+
 
     entrypoint_args = []
 
diff --git a/caliban/main.py b/caliban/main.py
index 4762514..3ee0cc7 100644
--- a/caliban/main.py
+++ b/caliban/main.py
@@ -57,8 +57,10 @@ def run_app(arg_input):
   if command == "shell":
     mount_home = not args['bare']
     image_id = args.get("image_id")
+    dlvm = args.get("dlvm")
     shell = args['shell']
     docker.run_interactive(job_mode,
+                           dlvm=dlvm,
                            image_id=image_id,
                            run_args=docker_run_args,
                            mount_home=mount_home,

From 20c34ef2b6bd91686af76e29871c0524da34051d Mon Sep 17 00:00:00 2001
From: Sanjay Agravat <agravat@google.com>
Date: Wed, 17 Jun 2020 10:43:40 -0400
Subject: [PATCH 05/13] Code cleanup to move functionality from
 build_dlvm_image to build_image

---
 caliban/cloud/core.py | 11 ++-------
 caliban/docker.py     | 54 +++++--------------------------------------
 2 files changed, 8 insertions(+), 57 deletions(-)

diff --git a/caliban/cloud/core.py b/caliban/cloud/core.py
index a1c4f91..09620c9 100644
--- a/caliban/cloud/core.py
+++ b/caliban/cloud/core.py
@@ -500,13 +500,8 @@ def generate_image_tag(project_id, docker_args, dlvm: str = None, dry_run: bool
   if dry_run:
     logging.info("Dry run - skipping actual 'docker build' and 'docker push'.")
     image_tag = "dry_run_tag"
-  elif dlvm is not None:
-    logging.info("DLVM - {}".format(dlvm))
-    #image_tag = dlvm
-    image_id = d.build_dlvm_image(dlvm=dlvm, **docker_args)
-    image_tag = d.push_uuid_tag(project_id, image_id)
   else:
-    image_id = d.build_image(**docker_args)
+    image_id = d.build_image(dlvm=dlvm, **docker_args)
     image_tag = d.push_uuid_tag(project_id, image_id)
 
   return image_tag
@@ -647,10 +642,8 @@ def submit_ml_job(job_mode: conf.JobMode,
   with session_scope(engine) as session:
     container_spec = generate_container_spec(session, docker_args, image_tag)
 
-    dlvm_tag = None
     if dlvm is not None:
-      dlvm_tag = dlvm # _dlvm_id(dlvm)
-      image_tag = generate_image_tag(project_id, docker_args, dlvm_tag, dry_run=dry_run)
+      image_tag = generate_image_tag(project_id, docker_args, dlvm=dlvm, dry_run=dry_run)
     elif image_tag is None:
       image_tag = generate_image_tag(project_id, docker_args, dry_run=dry_run)
 
diff --git a/caliban/docker.py b/caliban/docker.py
index f02c239..cee47e6 100644
--- a/caliban/docker.py
+++ b/caliban/docker.py
@@ -398,8 +398,10 @@ def _notebook_entries(lab: bool = False, version: Optional[str] = None, dlvm: bo
 """.format(library, version_suffix)
   else:
     return """
-  RUN /opt/conda/bin/pip install --user --no-cache-dir https://storage.googleapis.com/deeplearning-platform-ui-public/jupyterlab_gcpscheduler-1.0.0.tar.gz
-  RUN /opt/conda/bin/jupyter lab build
+RUN /opt/conda/bin/pip install --user --no-cache-dir \
+  https://storage.googleapis.com/deeplearning-platform-ui-public/jupyterlab_gcpscheduler-1.0.0.tar.gz
+
+RUN /opt/conda/bin/jupyter lab build
   """
 
 
@@ -608,7 +610,7 @@ def docker_image_id(output: str) -> ImageId:
   return ImageId(output.splitlines()[-1].split()[-1])
 
 
-def build_dlvm_image(job_mode: c.JobMode,
+def build_image(job_mode: c.JobMode,
                 build_path: str,
                 credentials_path: Optional[str] = None,
                 adc_path: Optional[str] = None,
@@ -629,8 +631,6 @@ def build_dlvm_image(job_mode: c.JobMode,
     with u.TempCopy(adc_path, tmp_name=".caliban_adc_creds.json") as adc:
       cache_args = ["--no-cache"] if no_cache else []
       cmd = ["docker", "build"] + cache_args + ["--rm", "-f-", build_path]
-      logging.info("extra: {}".format(kwargs))
-      #spec = {k: v for k, v in build_image_kwargs.items()}
 
       dockerfile = _dockerfile_template(job_mode,
                                         credentials_path=creds,
@@ -654,48 +654,6 @@ def build_dlvm_image(job_mode: c.JobMode,
         logging.error(e.stderr)
 
 
-def build_image(job_mode: c.JobMode,
-                build_path: str,
-                credentials_path: Optional[str] = None,
-                adc_path: Optional[str] = None,
-                no_cache: bool = False,
-                **kwargs) -> str:
-  """Builds a Docker image by generating a Dockerfile and passing it to `docker
-  build` via stdin. All output from the `docker build` process prints to
-  stdout.
-
-  Returns the image ID of the new docker container; if the command fails,
-  throws on error with information about the command and any issues that caused
-  the problem.
-
-  """
-  with u.TempCopy(credentials_path,
-                  tmp_name=".caliban_default_creds.json") as creds:
-    with u.TempCopy(adc_path, tmp_name=".caliban_adc_creds.json") as adc:
-      cache_args = ["--no-cache"] if no_cache else []
-      cmd = ["docker", "build"] + cache_args + ["--rm", "-f-", build_path]
-
-      dockerfile = _dockerfile_template(job_mode,
-                                        credentials_path=creds,
-                                        adc_path=adc,
-                                        **kwargs)
-
-      joined_cmd = " ".join(cmd)
-      logging.info("Running command: {}".format(joined_cmd))
-
-      try:
-        output, ret_code = u.capture_stdout(cmd, input_str=dockerfile)
-        if ret_code == 0:
-          return docker_image_id(output)
-        else:
-          error_msg = "Docker failed with error code {}.".format(ret_code)
-          raise DockerError(error_msg, cmd, ret_code)
-
-      except subprocess.CalledProcessError as e:
-        logging.error(e.output)
-        logging.error(e.stderr)
-
-
 def _image_tag_for_project(project_id: str, image_id: str) -> str:
   """Generate the GCR Docker image tag for the supplied pair of project_id and
   image_id.
@@ -983,7 +941,7 @@ def run(job_mode: c.JobMode,
 
   if dlvm is not None:
     logging.info("dlvm id: {}".format(dlvm))
-    image_id = build_dlvm_image(job_mode, dlvm=dlvm, **build_image_kwargs)
+    image_id = build_image(job_mode, dlvm=dlvm, **build_image_kwargs)
     logging.info("image_id: {}".format(image_id))
   elif image_id is None:
     logging.info("def run - image__id is none")

From 0c96fc8c001ba857e1bc14054af4426f525f4e38 Mon Sep 17 00:00:00 2001
From: Sanjay Agravat <agravat@google.com>
Date: Wed, 17 Jun 2020 23:15:46 -0400
Subject: [PATCH 06/13] Fixed issue with run_interactive toggling between shell
 and DLVM modes.

---
 caliban/docker.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/caliban/docker.py b/caliban/docker.py
index cee47e6..496a6ab 100644
--- a/caliban/docker.py
+++ b/caliban/docker.py
@@ -1017,7 +1017,7 @@ def run_interactive(job_mode: c.JobMode,
     # Don't set an entrypoint if we are running the DLVM notebook
     # Otherwise we are running in shell mode so we pass in the shell
     # entrypoint args.
-    if shell is None:
+    if entrypoint_args == []:
         interactive_run_args = _interactive_opts(workdir) + [
             "-it", \
         ] + _home_mount_cmds(mount_home) + run_args

From aecb17e012f02066fb08571e47f8d589dc72bfe8 Mon Sep 17 00:00:00 2001
From: Sanjay Agravat <agravat@google.com>
Date: Thu, 18 Jun 2020 00:53:31 -0400
Subject: [PATCH 07/13] Cleanup run_interactive and scheduler extension build.

---
 caliban/docker.py | 98 ++++++++++++++++++++++++++++++++++++++++-------
 1 file changed, 85 insertions(+), 13 deletions(-)

diff --git a/caliban/docker.py b/caliban/docker.py
index 496a6ab..d03ef7c 100644
--- a/caliban/docker.py
+++ b/caliban/docker.py
@@ -398,7 +398,7 @@ def _notebook_entries(lab: bool = False, version: Optional[str] = None, dlvm: bo
 """.format(library, version_suffix)
   else:
     return """
-RUN /opt/conda/bin/pip install --user --no-cache-dir \
+RUN /opt/conda/bin/pip install \
   https://storage.googleapis.com/deeplearning-platform-ui-public/jupyterlab_gcpscheduler-1.0.0.tar.gz
 
 RUN /opt/conda/bin/jupyter lab build
@@ -1007,7 +1007,7 @@ def run_interactive(job_mode: c.JobMode,
 
   if dlvm is None:
     # Pass the default entrypoint if not using DLVM
-    # Otherwise set the DLVM entrypoint depending on if we are in notebook 
+    # Otherwise set the DLVM entrypoint depending on if we are in notebook
     # mode or shell mode.
     interactive_run_args = _interactive_opts(workdir) + [
       "-it", \
@@ -1017,16 +1017,11 @@ def run_interactive(job_mode: c.JobMode,
     # Don't set an entrypoint if we are running the DLVM notebook
     # Otherwise we are running in shell mode so we pass in the shell
     # entrypoint args.
-    if entrypoint_args == []:
-        interactive_run_args = _interactive_opts(workdir) + [
-            "-it", \
-        ] + _home_mount_cmds(mount_home) + run_args
-    else:
-        entrypoint = SHELL_DICT[shell].executable
-        interactive_run_args = _interactive_opts(workdir) + [
-            "-it", \
-            "--entrypoint", entrypoint
-        ] + _home_mount_cmds(mount_home) + run_args
+    entrypoint = SHELL_DICT[shell].executable
+    interactive_run_args = _interactive_opts(workdir) + [
+        "-it", \
+        "--entrypoint", entrypoint
+    ] + _home_mount_cmds(mount_home) + run_args
 
 
     entrypoint_args = []
@@ -1041,6 +1036,83 @@ def run_interactive(job_mode: c.JobMode,
       **build_image_kwargs)
 
 
+def run_notebook_interactive(job_mode: c.JobMode,
+    workdir: Optional[str] = None,
+    image_id: Optional[str] = None,
+    dlvm: Optional[str] = None,
+    run_args: Optional[List[str]] = None,
+    mount_home: Optional[bool] = None,
+    shell: Optional[Shell] = None,
+    entrypoint: Optional[str] = None,
+    entrypoint_args: Optional[List[str]] = None,
+    **build_image_kwargs) -> None:
+  """Start a live shell in the terminal, with all dependencies installed and the
+  current working directory (and optionally the user's home directory) mounted.
+
+  Keyword args:
+
+  - job_mode: c.JobMode.
+  - image_id: ID of the image to run. Supplying this will skip an image build.
+  - run_args: extra arguments to supply to `docker run`.
+  - mount_home: if true, mounts the user's $HOME directory into the container
+    to `/home/$USERNAME`. If False, nothing.
+  - shell: name of the shell to install into the container. Also configures the
+    entrypoint if that's not supplied.
+  - entrypoint: command to run. Defaults to the executable command for the
+    supplied shell.
+  - entrypoint_args: extra arguments to supply to the entrypoint.
+
+  any extra kwargs supplied are passed through to build_image.
+
+  """
+  if workdir is None:
+    workdir = DEFAULT_WORKDIR
+
+  if run_args is None:
+    run_args = []
+
+  if entrypoint_args is None:
+    entrypoint_args = []
+
+  if mount_home is None:
+    mount_home = True
+
+  if shell is None:
+    # Only set a default shell if we're also mounting the home volume.
+    # Otherwise a custom shell won't have access to the user's profile.
+    shell = default_shell() if mount_home else Shell.bash
+
+  if entrypoint is None:
+    entrypoint = SHELL_DICT[shell].executable
+
+  if dlvm is None:
+    # Pass the default entrypoint if not using DLVM
+    # Otherwise set the DLVM entrypoint depending on if we are in notebook
+    # mode or shell mode.
+    interactive_run_args = _interactive_opts(workdir) + [
+        "-it", \
+        "--entrypoint", entrypoint
+    ] + _home_mount_cmds(mount_home) + run_args
+  else:
+    # Don't set an entrypoint if we are running the DLVM notebook
+    # Otherwise we are running in shell mode so we pass in the shell
+    # entrypoint args.
+    interactive_run_args = _interactive_opts(workdir) + [
+        "-it", \
+      ] + _home_mount_cmds(mount_home) + run_args
+
+
+    entrypoint_args = []
+
+  run(job_mode=job_mode,
+      run_args=interactive_run_args,
+      script_args=entrypoint_args,
+      image_id=image_id,
+      dlvm=dlvm,
+      shell=shell,
+      workdir=workdir,
+      **build_image_kwargs)
+
 def run_notebook(job_mode: c.JobMode,
                  port: Optional[int] = None,
                  lab: Optional[bool] = None,
@@ -1085,7 +1157,7 @@ def run_notebook(job_mode: c.JobMode,
   ]
   docker_args = ["-p", "{}:{}".format(port, port)] + run_args
 
-  run_interactive(job_mode,
+  run_notebook_interactive(job_mode,
                   dlvm=dlvm,
                   entrypoint="python",
                   #entrypoint="/bin/bash",

From 77d0f5a6bb976abfdbdf78b59ab23aa3535665cb Mon Sep 17 00:00:00 2001
From: Sanjay Agravat <agravat@google.com>
Date: Thu, 18 Jun 2020 09:49:57 -0400
Subject: [PATCH 08/13] Set minimize to False for jupyter lab build and code
 cleanup.

---
 caliban/docker.py | 30 ++++++++++--------------------
 1 file changed, 10 insertions(+), 20 deletions(-)

diff --git a/caliban/docker.py b/caliban/docker.py
index d03ef7c..bb149de 100644
--- a/caliban/docker.py
+++ b/caliban/docker.py
@@ -401,7 +401,7 @@ def _notebook_entries(lab: bool = False, version: Optional[str] = None, dlvm: bo
 RUN /opt/conda/bin/pip install \
   https://storage.googleapis.com/deeplearning-platform-ui-public/jupyterlab_gcpscheduler-1.0.0.tar.gz
 
-RUN /opt/conda/bin/jupyter lab build
+RUN /opt/conda/bin/jupyter lab build --minimize==False
   """
 
 
@@ -940,11 +940,8 @@ def run(job_mode: c.JobMode,
   logging.info("script_args {}".format(script_args))
 
   if dlvm is not None:
-    logging.info("dlvm id: {}".format(dlvm))
     image_id = build_image(job_mode, dlvm=dlvm, **build_image_kwargs)
-    logging.info("image_id: {}".format(image_id))
   elif image_id is None:
-    logging.info("def run - image__id is none")
     image_id = build_image(job_mode, **build_image_kwargs)
 
   base_cmd = _run_cmd(job_mode, run_args)
@@ -974,6 +971,8 @@ def run_interactive(job_mode: c.JobMode,
   - job_mode: c.JobMode.
   - image_id: ID of the image to run. Supplying this will skip an image build.
   - run_args: extra arguments to supply to `docker run`.
+  - dlvm: key of the base DLVM image to run. Supplying this perform an image
+    build using DLVM as the base image instead of blueshift
   - mount_home: if true, mounts the user's $HOME directory into the container
     to `/home/$USERNAME`. If False, nothing.
   - shell: name of the shell to install into the container. Also configures the
@@ -1006,24 +1005,18 @@ def run_interactive(job_mode: c.JobMode,
     entrypoint = SHELL_DICT[shell].executable
 
   if dlvm is None:
-    # Pass the default entrypoint if not using DLVM
-    # Otherwise set the DLVM entrypoint depending on if we are in notebook
-    # mode or shell mode.
+    # Pass in the entrypoint if not using DLVM
+    # Otherwise set the DLVM entrypoint from the shell argument
     interactive_run_args = _interactive_opts(workdir) + [
       "-it", \
       "--entrypoint", entrypoint
     ] + _home_mount_cmds(mount_home) + run_args
   else:
-    # Don't set an entrypoint if we are running the DLVM notebook
-    # Otherwise we are running in shell mode so we pass in the shell
-    # entrypoint args.
     entrypoint = SHELL_DICT[shell].executable
     interactive_run_args = _interactive_opts(workdir) + [
         "-it", \
         "--entrypoint", entrypoint
     ] + _home_mount_cmds(mount_home) + run_args
-
-
     entrypoint_args = []
 
   run(job_mode=job_mode,
@@ -1053,6 +1046,8 @@ def run_notebook_interactive(job_mode: c.JobMode,
 
   - job_mode: c.JobMode.
   - image_id: ID of the image to run. Supplying this will skip an image build.
+  - dlvm: key of the base DLVM image to run. Supplying this perform an image
+    build using DLVM as the base image instead of blueshift
   - run_args: extra arguments to supply to `docker run`.
   - mount_home: if true, mounts the user's $HOME directory into the container
     to `/home/$USERNAME`. If False, nothing.
@@ -1086,22 +1081,17 @@ def run_notebook_interactive(job_mode: c.JobMode,
     entrypoint = SHELL_DICT[shell].executable
 
   if dlvm is None:
-    # Pass the default entrypoint if not using DLVM
-    # Otherwise set the DLVM entrypoint depending on if we are in notebook
-    # mode or shell mode.
+    # Run vanilla jupyter/jupyterlab notebook if not using DLVM
+    # Otherwise use the jupyterlab notebook from DLVM with the scheduler
+    # extension.
     interactive_run_args = _interactive_opts(workdir) + [
         "-it", \
         "--entrypoint", entrypoint
     ] + _home_mount_cmds(mount_home) + run_args
   else:
-    # Don't set an entrypoint if we are running the DLVM notebook
-    # Otherwise we are running in shell mode so we pass in the shell
-    # entrypoint args.
     interactive_run_args = _interactive_opts(workdir) + [
         "-it", \
       ] + _home_mount_cmds(mount_home) + run_args
-
-
     entrypoint_args = []
 
   run(job_mode=job_mode,

From e4cba55aef51be96d12c2c56e6544116a577a715 Mon Sep 17 00:00:00 2001
From: Sanjay Agravat <agravat@google.com>
Date: Thu, 18 Jun 2020 22:17:22 -0400
Subject: [PATCH 09/13] Added DLVM config map from Google Container Registry
 and set up the help function to list the DLVM types.

---
 caliban/cli.py    |  4 +++-
 caliban/config.py | 27 +++++++++++++++++++++++++++
 caliban/docker.py | 14 +++++++-------
 3 files changed, 37 insertions(+), 8 deletions(-)

diff --git a/caliban/cli.py b/caliban/cli.py
index 8bff0ce..797fd8f 100644
--- a/caliban/cli.py
+++ b/caliban/cli.py
@@ -211,9 +211,11 @@ def image_tag_arg(parser):
 
 
 def dlvm_arg(parser):
+  dlvm_types = {key for key in conf._dlvm_config('CPU').keys()}
   parser.add_argument(
       "--dlvm",
-      help="DLVM base image tag accessible via Container Registry. If supplied, "
+      help="DLVM base image type. Must be one of  "
+           "{}".format(dlvm_types) + ". If supplied, "
       "Caliban will skip the build and push steps and use this image tag.")
 
 
diff --git a/caliban/config.py b/caliban/config.py
index 8bad789..8efd070 100644
--- a/caliban/config.py
+++ b/caliban/config.py
@@ -78,6 +78,33 @@ class JobMode(str, Enum):
 }
 
 
+def extract_dlvm_image(job_mode: JobMode, dlvm_arg: str) -> str:
+  """Returns the DLVM image url for the job model and the comand line arg
+  or returns None if the key doesn't exist in the config.
+
+  """
+  return _dlvm_config(job_mode).get(dlvm_arg)
+
+
+def _dlvm_config(job_mode: JobMode) -> Dict[str, str]:
+  job_mode_str = job_mode.lower()
+  return {
+    "pytorch": f"gcr.io/deeplearning-platform-release/pytorch-{job_mode_str}",
+    "pytorch-1.0": f"gcr.io/deeplearning-platform-release/pytorch-{job_mode_str}.1-0",
+    "pytorch-1.1": f"gcr.io/deeplearning-platform-release/pytorch-{job_mode_str}.1-1",
+    "pytorch-1.2": f"gcr.io/deeplearning-platform-release/pytorch-{job_mode_str}.1-2",
+    "pytorch-1.3": f"gcr.io/deeplearning-platform-release/pytorch-{job_mode_str}.1-3",
+    "pytorch-1.4": f"gcr.io/deeplearning-platform-release/pytorch-{job_mode_str}.1-4",
+    "tf": f"gcr.io/deeplearning-platform-release/tf-{job_mode_str}",
+    "tf-1.13": f"gcr.io/deeplearning-platform-release/tf-{job_mode_str}.1-13",
+    "tf-1.14": f"gcr.io/deeplearning-platform-release/tf-{job_mode_str}.1-14",
+    "tf-1.15": f"gcr.io/deeplearning-platform-release/tf-{job_mode_str}.1-15",
+    "tf2": f"gcr.io/deeplearning-platform-release/tf2-{job_mode_str}",
+    "tf-2.0": f"gcr.io/deeplearning-platform-release/tf2-{job_mode_str}.2-0",
+    "tf-2.1": f"gcr.io/deeplearning-platform-release/tf2-{job_mode_str}.2-1",
+    "tf-2.2": f"gcr.io/deeplearning-platform-release/tf2-{job_mode_str}.2-2"
+  }
+
 def gpu(job_mode: JobMode) -> bool:
   """Returns True if the supplied JobMode is JobMode.GPU, False otherwise.
 
diff --git a/caliban/docker.py b/caliban/docker.py
index bb149de..d82e4be 100644
--- a/caliban/docker.py
+++ b/caliban/docker.py
@@ -469,13 +469,13 @@ def _extra_dir_entries(workdir: str, user_id: int, user_group: int,
     ret += "\n{}".format(_copy_dir_entry(workdir, user_id, user_group, d))
   return ret
 
-def _dlvm_id(image: str) -> str:
-  if image == "pytorch":
-    return "gcr.io/deeplearning-platform-release/pytorch-cpu:latest"
-  elif image == "tf-21":
-    return "gcr.io/deeplearning-platform-release/tf2-cpu.2-1"
-  else:
-    return None
+
+def _dlvm_id(job_mode: c.JobMode, dlvm_arg: str) -> str:
+  """Returns the DLVM image url for job mode and the command line parameter
+
+  """
+  return c.extract_dlvm_image(job_mode, dlvm_arg)
+
 
 def _dockerfile_template(
     job_mode: c.JobMode,

From 6ba10e02e105e367f42fda42900f15e1013ca784 Mon Sep 17 00:00:00 2001
From: Sanjay Agravat <agravat@google.com>
Date: Fri, 19 Jun 2020 00:17:40 -0400
Subject: [PATCH 10/13] Update required google api packages to specific
 versions due to google-auth release issue.

---
 setup.py | 7 ++++---
 1 file changed, 4 insertions(+), 3 deletions(-)

diff --git a/setup.py b/setup.py
index aae8196..656dea6 100644
--- a/setup.py
+++ b/setup.py
@@ -45,9 +45,10 @@ def readme():
     'pyyaml',
     'tqdm',
     'kubernetes>=10.0.1',
-    'google-auth>=1.7.0',
-    'google-cloud-core>=1.0.3',
-    'google-cloud-container>=0.3.0',
+    'google-auth==1.17.2',
+    'google-api-core==1.20.1',
+    'google-cloud-core==1.3.0',
+    'google-cloud-container==1.0.1',
     'psycopg2-binary==2.8.5',
     'urllib3>=1.25.7',
     'yaspin>=0.16.0',

From baf0aff1b43adb47252de56281b7b6545083fd16 Mon Sep 17 00:00:00 2001
From: Sanjay Agravat <agravat@google.com>
Date: Fri, 19 Jun 2020 08:59:42 -0400
Subject: [PATCH 11/13] Sorted DLVM arg types returned to help command.

---
 caliban/cli.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/caliban/cli.py b/caliban/cli.py
index 797fd8f..db5b460 100644
--- a/caliban/cli.py
+++ b/caliban/cli.py
@@ -211,7 +211,7 @@ def image_tag_arg(parser):
 
 
 def dlvm_arg(parser):
-  dlvm_types = {key for key in conf._dlvm_config('CPU').keys()}
+  dlvm_types = [key for key in sorted(conf._dlvm_config('CPU').keys())]
   parser.add_argument(
       "--dlvm",
       help="DLVM base image type. Must be one of  "

From b1cb30e89b152b6dc71886bec15c695bd6a6b69b Mon Sep 17 00:00:00 2001
From: Sanjay Agravat <agravat@google.com>
Date: Fri, 19 Jun 2020 12:54:17 -0400
Subject: [PATCH 12/13] Cleanup and update setup.py with upstream versions.

---
 caliban/cloud/core.py |  8 --------
 caliban/docker.py     | 11 +++--------
 setup.py              |  9 ++++-----
 3 files changed, 7 insertions(+), 21 deletions(-)

diff --git a/caliban/cloud/core.py b/caliban/cloud/core.py
index 09620c9..d895fee 100644
--- a/caliban/cloud/core.py
+++ b/caliban/cloud/core.py
@@ -543,14 +543,6 @@ def submit_job_specs(
 
   execute_requests(requests, num_specs, num_retries=request_retries)
 
-def _dlvm_id(image: str) -> str:
-  if image == "pytorch":
-    return "gcr.io/deeplearning-platform-release/pytorch-cpu:latest"
-  elif image == "tf-21":
-    return "gcr.io/deeplearning-platform-release/tf2-cpu.2-1"
-  else:
-    return None
-
 
 def submit_ml_job(job_mode: conf.JobMode,
                   docker_args: Dict[str, Any],
diff --git a/caliban/docker.py b/caliban/docker.py
index d82e4be..c876c5f 100644
--- a/caliban/docker.py
+++ b/caliban/docker.py
@@ -308,11 +308,8 @@ def _service_account_entry(user_id: int, user_group: int, credentials_path: str,
 COPY --chown={user_id}:{user_group} {credentials_path} {container_creds}
 
 # Use the credentials file to activate gcloud, gsutil inside the container.
-#RUN mkdir -p /home/agravat/.config
-#RUN chown agravat -R /home/agravat/.config
-#RUN gcloud init
-#RUN gcloud auth activate-service-account --key-file={container_creds} && \
-#  git config --global credential.'https://source.developers.google.com'.helper gcloud.sh
+RUN gcloud auth activate-service-account --key-file={container_creds} && \
+ git config --global credential.'https://source.developers.google.com'.helper gcloud.sh
 
 ENV GOOGLE_APPLICATION_CREDENTIALS={container_creds}
 """.format_map({
@@ -529,9 +526,7 @@ def _dockerfile_template(
   if dlvm is None:
     base_image = base_image_fn(job_mode)
   else:
-    base_image = _dlvm_id(dlvm)
-
-  logging.info("base image ------ {}".format(base_image))
+    base_image = _dlvm_id(job_mode, dlvm)
 
   dockerfile = """
 FROM {base_image}
diff --git a/setup.py b/setup.py
index 656dea6..9a24802 100644
--- a/setup.py
+++ b/setup.py
@@ -45,10 +45,9 @@ def readme():
     'pyyaml',
     'tqdm',
     'kubernetes>=10.0.1',
-    'google-auth==1.17.2',
-    'google-api-core==1.20.1',
-    'google-cloud-core==1.3.0',
-    'google-cloud-container==1.0.1',
+    'google-auth>=1.18.0',
+    'google-cloud-core>=1.0.3',
+    'google-cloud-container>=0.3.0',
     'psycopg2-binary==2.8.5',
     'urllib3>=1.25.7',
     'yaspin>=0.16.0',
@@ -65,7 +64,7 @@ def readme():
       description='Docker-based job runner for AI research.',
       long_description=readme(),
       long_description_content_type="text/markdown",
-      python_requires='>=3.7.0',
+      python_requires='>=3.6.0',
       author='Caliban Team',
       author_email='samritchie@google.com',
       url='https://github.com/google/caliban',

From 3a417e996e3c8f8a28ec0bfc26a755cd9ec1387a Mon Sep 17 00:00:00 2001
From: Sanjay Agravat <agravat@google.com>
Date: Fri, 19 Jun 2020 13:55:32 -0400
Subject: [PATCH 13/13] Set python to 3.6

---
 Makefile | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/Makefile b/Makefile
index 98f343e..a31756d 100644
--- a/Makefile
+++ b/Makefile
@@ -37,7 +37,7 @@ clean-files:
 .PHONY: install
 install:
 	rm -rf $(ENV_NAME)
-	virtualenv -p python3.7 $(ENV_NAME)
+	virtualenv -p python3.6 $(ENV_NAME)
 	$(PIP) install -r requirements-dev.txt
 	$(PIP) install -r docs/requirements.txt
 	$(PIP) install -e .