Add platform arg and other improvements to dagster-docker CLI (#7698)

python_modules/automation/automation/docker/cli.py

@@ -1,4 +1,4 @@
-from typing import List
+from typing import List, Optional
 
 import click
 
@@ -26,68 +26,81 @@ def list():  # pylint: disable=redefined-builtin
         print(image.image)  # pylint: disable=print-call
 
 
-@cli.command()
-@click.option("--name", required=True, help="Name of image to build")
-@click.option(
+# Shared options between `build` and `build_all`
+opt_build_name = click.option("--name", required=True, help="Name of image to build")
+opt_build_dagster_version = click.option(
     "--dagster-version",
     required=True,
     help="Version of image to build",
 )
-@click.option(
+opt_build_platform = click.option(
+    "--platform",
+    required=False,
+    help="Target platform name to pass to `docker build`",
+)
+opt_build_timestamp = click.option(
     "-t",
     "--timestamp",
     type=click.STRING,
     required=False,
     default=current_time_str(),
     help="Timestamp to build in format 2020-07-11T040642 (defaults to now UTC)",
 )
+
+
+@cli.command()
+@opt_build_name
+@opt_build_dagster_version
+@opt_build_timestamp
 @click.option("-v", "--python-version", type=click.STRING, required=True)
-def build(name, dagster_version, timestamp, python_version):
-    get_image(name).build(timestamp, dagster_version, python_version)
+@opt_build_platform
+def build(
+    name: str, dagster_version: str, timestamp: str, python_version: str, platform: Optional[str]
+):
+    get_image(name).build(timestamp, dagster_version, python_version, platform)
 
 
 @cli.command()
-@click.option("--name", required=True, help="Name of image to build")
-@click.option(
-    "--dagster-version",
-    required=True,
-    help="Version of image to build, must match current dagster version",
-)
-@click.option(
-    "-t",
-    "--timestamp",
-    type=click.STRING,
-    required=False,
-    default=current_time_str(),
-    help="Timestamp to build in format 2020-07-11T040642 (defaults to now UTC)",
-)
-def build_all(name, dagster_version, timestamp):
+@opt_build_name
+@opt_build_dagster_version
+@opt_build_timestamp
+@opt_build_platform
+def build_all(name: str, dagster_version: str, timestamp: str, platform: Optional[str]):
     """Build all supported python versions for image"""
     image = get_image(name)
 
     for python_version in image.python_versions:
-        image.build(timestamp, dagster_version, python_version)
+        image.build(timestamp, dagster_version, python_version, platform)
+
+
+# Shared push options
+opt_push_name = click.option("--name", required=True, help="Name of image to push")
+opt_push_dagster_version = click.option(
+    "--dagster-version",
+    required=True,
+    help="Version of image to push",
+)
 
 
 @cli.command()
-@click.option("--name", required=True, help="Name of image to push")
+@opt_push_name
 @click.option("-v", "--python-version", type=click.STRING, required=True)
-@click.option("-v", "--custom-tag", type=click.STRING, required=False)
-def push(name, python_version, custom_tag):
+@click.option("--custom-tag", type=click.STRING, required=False)
+def push(name: str, python_version: str, custom_tag: Optional[str]):
     ensure_ecr_login()
     get_image(name).push(python_version, custom_tag=custom_tag)
 
 
 @cli.command()
-@click.option("--name", required=True, help="Name of image to push")
-def push_all(name):
+@opt_push_name
+def push_all(name: str):
     ensure_ecr_login()
     image = get_image(name)
     for python_version in image.python_versions:
         image.push(python_version)
 
 
-def push_to_registry(name: str, tags: List[str]):
+def push_to_registry(name: str, tags: List[str]) -> None:
     check.str_param(name, "name")
     check.list_param(tags, "tags", of_type=str)
 
@@ -103,13 +116,9 @@ def push_to_registry(name: str, tags: List[str]):
 
 
 @cli.command()
-@click.option("--name", required=True, help="Name of image to push")
-@click.option(
-    "--dagster-version",
-    required=True,
-    help="Version of image to push",
-)
-def push_dockerhub(name, dagster_version):
+@opt_push_name
+@opt_push_dagster_version
+def push_dockerhub(name: str, dagster_version: str):
     """Used for pushing k8s images to Docker Hub. Must be logged in to Docker Hub for this to
     succeed.
     """
@@ -120,13 +129,9 @@ def push_dockerhub(name, dagster_version):
 
 
 @cli.command()
-@click.option("--name", required=True, help="Name of image to push")
-@click.option(
-    "--dagster-version",
-    required=True,
-    help="Version of image to push",
-)
-def push_ecr(name, dagster_version):
+@opt_push_name
+@opt_push_dagster_version
+def push_ecr(name: str, dagster_version: str):
     """Used for pushing k8s images to our public ECR.
 
     You must be authed for ECR. Run:

python_modules/automation/automation/docker/dagster_docker.py

@@ -1,6 +1,6 @@
 import contextlib
 import os
-from typing import Callable, NamedTuple, Optional
+from typing import Callable, Dict, List, NamedTuple, Optional
 
 import yaml
 
@@ -49,20 +49,20 @@ def __new__(cls, image: str, build_cm: Callable = do_nothing, path: Optional[str
         )
 
     @property
-    def python_versions(self):
+    def python_versions(self) -> List[str]:
         """List of Python versions supported for this image."""
         with open(os.path.join(self.path, "versions.yaml"), "r", encoding="utf8") as f:
             versions = yaml.safe_load(f.read())
         return list(versions.keys())
 
-    def _get_last_updated_for_python_version(self, python_version):
+    def _get_last_updated_for_python_version(self, python_version: str) -> str:
         """Retrieve the last_updated timestamp for a particular python_version of this image."""
         check.str_param(python_version, "python_version")
         with open(os.path.join(self.path, "last_updated.yaml"), "r", encoding="utf8") as f:
             last_updated = yaml.safe_load(f.read())
             return last_updated[python_version]
 
-    def _set_last_updated_for_python_version(self, timestamp, python_version):
+    def _set_last_updated_for_python_version(self, timestamp: str, python_version: str) -> None:
         """Update the last_updated timestamp for a particular python_version of this image."""
         check.str_param(timestamp, "timestamp")
         check.str_param(python_version, "python_version")
@@ -79,15 +79,17 @@ def _set_last_updated_for_python_version(self, timestamp, python_version):
         with open(os.path.join(self.path, "last_updated.yaml"), "w", encoding="utf8") as f:
             yaml.dump(last_updated, f, default_flow_style=False)
 
-    def local_image(self, python_version):
+    def local_image(self, python_version: str) -> str:
         """Generates the local image name, like: "dagster/foo:some-tag" """
         check.str_param(python_version, "python_version")
 
         last_updated = self._get_last_updated_for_python_version(python_version)
         tag = python_version_image_tag(python_version, last_updated)
         return "{}/{}:{}".format(DEFAULT_LOCAL_PREFIX, self.image, tag)
 
-    def aws_image(self, python_version=None, custom_tag=None):
+    def aws_image(
+        self, python_version: Optional[str] = None, custom_tag: Optional[str] = None
+    ) -> str:
         """Generates the AWS ECR image name, like:
         "1234567890.dkr.ecr.us-west-1.amazonaws.com/foo:some-tag"
         """
@@ -108,7 +110,7 @@ def aws_image(self, python_version=None, custom_tag=None):
             aws_region=get_aws_region(),
         )
 
-    def _get_docker_args(self, dagster_version, python_version):
+    def _get_docker_args(self, dagster_version: str, python_version: str) -> Dict[str, str]:
         """Retrieve Docker arguments from this image's versions.yaml, and update with latest Dagster
         version.
 
@@ -141,19 +143,23 @@ def _get_docker_args(self, dagster_version, python_version):
         docker_args["DAGSTER_VERSION"] = dagster_version
         return docker_args
 
-    def build(self, timestamp, dagster_version, python_version):
+    def build(
+        self, timestamp, dagster_version: str, python_version: str, platform: Optional[str] = None
+    ) -> None:
         check.str_param(timestamp, "timestamp")
         check.str_param(python_version, "python_version")
+        check.opt_str_param(platform, "platform")
         with self.build_cm(self.path):
             self._set_last_updated_for_python_version(timestamp, python_version)
 
             execute_docker_build(
                 self.local_image(python_version),
                 docker_args=self._get_docker_args(dagster_version, python_version),
                 cwd=self.path,
+                platform=platform,
             )
 
-    def push(self, python_version, custom_tag=None):
+    def push(self, python_version: str, custom_tag: Optional[str] = None) -> None:
         """Push this image to ECR."""
 
         if custom_tag:

python_modules/automation/automation/docker/image_defs.py

@@ -2,6 +2,7 @@
 import contextlib
 import os
 import shutil
+from typing import Callable, Dict, Iterator, List, Optional
 
 from automation.git import git_repo_root
 
@@ -10,12 +11,14 @@
 from .dagster_docker import DagsterDockerImage
 
 
-def get_dagster_repo():
+def get_dagster_repo() -> str:
     return git_repo_root()
 
 
 @contextlib.contextmanager
-def copy_directories(paths, cwd, destination="build_cache"):
+def copy_directories(
+    paths: List[str], cwd: str, destination: str = "build_cache"
+) -> Iterator[None]:
     check.invariant(os.path.exists(cwd), "Image directory does not exist")
     build_cache_dir = os.path.join(cwd, destination)
 
@@ -47,7 +50,7 @@ def copy_directories(paths, cwd, destination="build_cache"):
 
 
 @contextlib.contextmanager
-def k8s_example_cm(cwd):
+def k8s_example_cm(cwd: str) -> Iterator[None]:
     with copy_directories(
         [
             "examples/deploy_k8s/example_project",
@@ -57,7 +60,7 @@ def k8s_example_cm(cwd):
         yield
 
 
-def get_core_celery_k8s_dirs():
+def get_core_celery_k8s_dirs() -> List[str]:
     return [
         "python_modules/dagster",
         "python_modules/libraries/dagster-postgres",
@@ -67,7 +70,7 @@ def get_core_celery_k8s_dirs():
     ]
 
 
-def get_core_k8s_dirs():
+def get_core_k8s_dirs() -> List[str]:
     return [
         "python_modules/dagster",
         "python_modules/libraries/dagster-postgres",
@@ -76,7 +79,7 @@ def get_core_k8s_dirs():
 
 
 @contextlib.contextmanager
-def k8s_example_editable_cm(cwd):
+def k8s_example_editable_cm(cwd: str) -> Iterator[None]:
     with copy_directories(
         get_core_celery_k8s_dirs()
         + [
@@ -91,7 +94,7 @@ def k8s_example_editable_cm(cwd):
 
 
 @contextlib.contextmanager
-def k8s_dagit_editable_cm(cwd):
+def k8s_dagit_editable_cm(cwd: str) -> Iterator[None]:
     print("!!!!! WARNING: You must call `make rebuild_dagit` after making changes to Dagit !!!!\n")
     with copy_directories(
         get_core_celery_k8s_dirs()
@@ -105,7 +108,7 @@ def k8s_dagit_editable_cm(cwd):
 
 
 @contextlib.contextmanager
-def k8s_dagit_example_cm(cwd):
+def k8s_dagit_example_cm(cwd: str) -> Iterator[None]:
     with copy_directories(
         get_core_celery_k8s_dirs()
         + [
@@ -122,7 +125,7 @@ def k8s_dagit_example_cm(cwd):
 
 
 @contextlib.contextmanager
-def k8s_celery_worker_editable_cm(cwd):
+def k8s_celery_worker_editable_cm(cwd: str) -> Iterator[None]:
     with copy_directories(
         get_core_celery_k8s_dirs(),
         cwd,
@@ -131,7 +134,7 @@ def k8s_celery_worker_editable_cm(cwd):
 
 
 @contextlib.contextmanager
-def user_code_example_cm(cwd):
+def user_code_example_cm(cwd: str) -> Iterator[None]:
     with copy_directories(
         [
             "examples/deploy_k8s/example_project",
@@ -142,7 +145,7 @@ def user_code_example_cm(cwd):
 
 
 @contextlib.contextmanager
-def user_code_example_editable_cm(cwd):
+def user_code_example_editable_cm(cwd: str) -> Iterator[None]:
     with copy_directories(
         get_core_celery_k8s_dirs() + ["python_modules/libraries/dagster-aws"],
         cwd,
@@ -154,7 +157,7 @@ def user_code_example_editable_cm(cwd):
 
 
 @contextlib.contextmanager
-def dagster_k8s_editable_cm(cwd):
+def dagster_k8s_editable_cm(cwd: str) -> Iterator[None]:
     print("!!!!! WARNING: You must call `make rebuild_dagit` after making changes to Dagit !!!!\n")
     with copy_directories(
         get_core_k8s_dirs()
@@ -169,7 +172,7 @@ def dagster_k8s_editable_cm(cwd):
 
 
 @contextlib.contextmanager
-def dagster_celery_k8s_editable_cm(cwd):
+def dagster_celery_k8s_editable_cm(cwd: str) -> Iterator[None]:
     print("!!!!! WARNING: You must call `make rebuild_dagit` after making changes to Dagit !!!!\n")
     with copy_directories(
         get_core_celery_k8s_dirs()
@@ -184,7 +187,7 @@ def dagster_celery_k8s_editable_cm(cwd):
 
 
 # Some images have custom build context manager functions, listed here
-CUSTOM_BUILD_CONTEXTMANAGERS = {
+CUSTOM_BUILD_CONTEXTMANAGERS: Dict[str, Callable] = {
     "k8s-example": k8s_example_cm,
     "k8s-example-editable": k8s_example_editable_cm,
     "k8s-dagit-editable": k8s_dagit_editable_cm,
@@ -197,7 +200,7 @@ def dagster_celery_k8s_editable_cm(cwd):
 }
 
 
-def list_images(images_path=None):
+def list_images(images_path: Optional[str] = None) -> List[DagsterDockerImage]:
     """List all images that we manage.
 
     Returns:
@@ -216,8 +219,7 @@ def list_images(images_path=None):
     return images
 
 
-def get_image(name, images_path=None):
+def get_image(name: str, images_path: Optional[str] = None) -> DagsterDockerImage:
     """Retrieve the image information from the list defined above."""
     image = next((img for img in list_images(images_path=images_path) if img.image == name), None)
-    check.invariant(image is not None, "could not find image {}".format(name))
-    return image
+    return check.not_none(image, "could not find image {}".format(name))