Skip docker push if image exists remotely (#36)

CHANGELOG.md

@@ -3,6 +3,11 @@
 - Caliban now authenticates AI Platform job submissions using the authentication
   provided by `gcloud auth login`, rather than requiring a service account key.
   This significantly simplifies the setup required for a first time user.
+- `caliban cloud` now checks if the image exists remotely before issuing a
+  `docker push` command on the newly built image
+  (https://github.com/google/caliban/pull/36)
+- Big internal refactor to make it easier to work on code, increase test
+  coverage, add new backends (https://github.com/google/caliban/pull/32)
 
 # 0.2.6
 

caliban/docker/push.py

@@ -18,10 +18,15 @@
 
 """
 
+import json
 import subprocess
 
+from absl import logging
 
-def _image_tag_for_project(project_id: str, image_id: str) -> str:
+
+def _image_tag_for_project(project_id: str,
+                           image_id: str,
+                           include_tag: bool = True) -> str:
   """Generate the GCR Docker image tag for the supplied pair of project_id and
   image_id.
 
@@ -31,21 +36,48 @@ def _image_tag_for_project(project_id: str, image_id: str) -> str:
 
   """
   project_s = project_id.replace(":", "/")
-  return "gcr.io/{}/{}:latest".format(project_s, image_id)
+  base = f"gcr.io/{project_s}/{image_id}"
+  return f"{base}:latest" if include_tag else base
+
+
+def _gcr_list_tags(project_id: str, image_id: str):
+  """Returns a sequence of metadata for all tags of the supplied image_id in the
+  supplied project.
+
+  """
+  image_tag = _image_tag_for_project(project_id, image_id, include_tag=False)
+  cmd = [
+      "gcloud", "container", "images", "list-tags", f"--project={project_id}",
+      "--format=json", image_tag
+  ]
+  return json.loads(subprocess.check_output(cmd))
+
+
+def gcr_image_pushed(project_id: str, image_id: str) -> bool:
+  """Returns true if the supplied image has been pushed to the container registry
+  for the supplied project, false otherwise.
+
+  """
+  return len(_gcr_list_tags(project_id, image_id)) > 0
 
 
-def push_uuid_tag(project_id: str, image_id: str) -> str:
+def push_uuid_tag(project_id: str, image_id: str, force: bool = False) -> str:
   """Takes a base image and tags it for upload, then pushes it to a remote Google
   Container Registry.
 
   Returns the tag on a successful push.
-
-  TODO should this just check first before attempting to push if the image
-  exists? Immutable names means that if the tag is up there, we're done.
-  Potentially use docker-py for this.
-
   """
   image_tag = _image_tag_for_project(project_id, image_id)
-  subprocess.run(["docker", "tag", image_id, image_tag], check=True)
-  subprocess.run(["docker", "push", image_tag], check=True)
+
+  def missing_remotely():
+    missing = not gcr_image_pushed(project_id, image_id)
+    if not missing:
+      logging.info(
+          f"Skipping docker push, as {image_tag} already exists remotely.")
+    return missing
+
+  if force or missing_remotely():
+    subprocess.run(["docker", "tag", image_id, image_tag], check=True)
+    subprocess.run(["docker", "push", image_tag], check=True)
+
   return image_tag

caliban/platform/cloud/core.py

@@ -20,14 +20,11 @@
 import datetime
 from copy import deepcopy
 from pprint import pformat
-from subprocess import CalledProcessError, check_output
 from typing import Any, Dict, Iterable, List, Optional, Tuple
 
 import tqdm
 from absl import logging
 from blessings import Terminal
-from google.oauth2 import service_account
-from google.oauth2.credentials import Credentials
 from googleapiclient import discovery
 from googleapiclient.errors import HttpError
 

caliban/util/auth.py

@@ -30,8 +30,9 @@ def auth_access_token() -> Optional[str]:
 
   """
   try:
-    return check_output(['gcloud', 'auth', 'print-access-token'],
-                        encoding='utf8').rstrip()
+    ret = check_output(['gcloud', 'auth', 'print-access-token'],
+                       encoding='utf8').rstrip()
+    return ret if len(ret) > 0 else None
   except CalledProcessError:
     return None
 

requirements-dev.txt

@@ -4,4 +4,5 @@ ipython
 pre-commit
 pytest==5.4.3
 pytest-cov==2.10.0
+pytest-subprocess==0.1.5
 twine

tests/caliban/docker/test_push.py

@@ -17,6 +17,13 @@
 import caliban.docker.push as p
 
 
+def register_list_tags(process, project_id, tag, **kwargs):
+  process.register_subprocess([
+      "gcloud", "container", "images", "list-tags", f"--project={project_id}",
+      "--format=json", tag
+  ], **kwargs)
+
+
 def test_image_tag_for_project():
   """Tests that we generate a valid image tag for domain-scoped and modern
     project IDs.
@@ -27,3 +34,52 @@ def test_image_tag_for_project():
 
   assert p._image_tag_for_project(
       "google.com:face", "imageid") == "gcr.io/google.com/face/imageid:latest"
+
+
+def test_force_push_uuid_tag(fake_process):
+  """Check that the push command actually attempts to tag and push."""
+  project_id = "project"
+  image_id = "imageid"
+
+  tag = p._image_tag_for_project(project_id, image_id)
+
+  fake_process.register_subprocess(["docker", "tag", image_id, tag])
+  fake_process.register_subprocess(["docker", "push", tag])
+
+  assert p.push_uuid_tag(project_id, image_id, force=True) == tag
+
+
+def test_already_pushed_uuid_tag(fake_process):
+  """Check that push_uuid_tag does NOT attempt to push if the process already
+  exists.."""
+  project_id = "project"
+  image_id = "imageid"
+
+  base_tag = p._image_tag_for_project(project_id, image_id, include_tag=False)
+  tag = p._image_tag_for_project(project_id, image_id)
+
+  register_list_tags(fake_process,
+                     project_id,
+                     base_tag,
+                     stdout="[{\"metadata\": []}]")
+
+  assert p.push_uuid_tag(project_id, image_id) == tag
+
+
+def test_push_uuid_tag_if_no_remote_image(fake_process):
+  """Check that push_uuid_tag DOES attempt to push if the image doesn't exist in
+  the remote container registry already.
+
+  """
+  project_id = "project"
+  image_id = "imageid"
+
+  base_tag = p._image_tag_for_project(project_id, image_id, include_tag=False)
+  tag = p._image_tag_for_project(project_id, image_id)
+
+  register_list_tags(fake_process, project_id, base_tag, stdout="[]")
+
+  fake_process.register_subprocess(["docker", "tag", image_id, tag])
+  fake_process.register_subprocess(["docker", "push", tag])
+
+  assert p.push_uuid_tag(project_id, image_id) == tag

tests/caliban/util/test_auth.py

@@ -0,0 +1,69 @@
+#!/usr/bin/python
+#
+# Copyright 2020 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from subprocess import CalledProcessError
+
+from google.oauth2.credentials import Credentials
+
+import caliban.util.auth as a
+
+
+def register_auth(process, **kwargs):
+  process.register_subprocess(["gcloud", "auth", "print-access-token"],
+                              **kwargs)
+
+
+def fail_process(process):
+  process.returncode = 1
+  raise CalledProcessError("cmd", "exception! Not logged in!")
+
+
+def test_auth_access_token(fake_process):
+  """Check that if the user has logged in with `gcloud auth login`,
+  `auth_access_token` returns the correct token.
+
+  """
+  token = "token"
+  register_auth(fake_process, stdout=token)
+  assert a.auth_access_token() == token
+
+
+def test_missing_auth_access_token(fake_process):
+  """Check that if the user has NOT logged in with `gcloud auth login`,
+  `auth_access_token` returns None.
+
+  """
+  register_auth(fake_process, callback=fail_process)
+  assert a.auth_access_token() is None
+
+
+def test_gcloud_auth_credentials(fake_process):
+  """Check that if the user has logged in with `gcloud auth login`,
+  a proper instance of Credentials is returned.
+
+  """
+  token = "token"
+  register_auth(fake_process, stdout=token)
+  assert isinstance(a.gcloud_auth_credentials(), Credentials)
+
+
+def test_missing_gcloud_auth_credentials(fake_process):
+  """Check that if the user has logged in with `gcloud auth login`,
+  `auth_access_token` returns the correct token.
+
+  """
+  register_auth(fake_process, callback=fail_process)
+  assert a.gcloud_auth_credentials() is None