From 6c6915dc67d92dc0a6d7b7aad4cc1b21d4fe8915 Mon Sep 17 00:00:00 2001 From: Luis Perez Date: Mon, 23 Oct 2023 11:59:52 -0700 Subject: [PATCH] Respect a service's docker_registry for adhoc spark-runs There are some services where we want to ensure that a specialized docker registry is always used - even for adhoc development runs. This change is loosely based on what we do in push-to-registry, where we read service.yaml to see if a specialized registry needs to be used. --- paasta_tools/cli/cmds/spark_run.py | 18 ++++++++++++++++-- 1 file changed, 16 insertions(+), 2 deletions(-) diff --git a/paasta_tools/cli/cmds/spark_run.py b/paasta_tools/cli/cmds/spark_run.py index 26c07b170b..042475039c 100644 --- a/paasta_tools/cli/cmds/spark_run.py +++ b/paasta_tools/cli/cmds/spark_run.py @@ -17,6 +17,7 @@ import yaml from boto3.exceptions import Boto3Error +from service_configuration_lib import read_service_configuration from service_configuration_lib import spark_config from service_configuration_lib.spark_config import get_aws_credentials from service_configuration_lib.spark_config import get_grafana_url @@ -218,7 +219,7 @@ def add_subparser(subparsers): list_parser.add_argument( "--docker-registry", help="Docker registry to push the Spark image built.", - default=DEFAULT_SPARK_DOCKER_REGISTRY, + default=None, ) list_parser.add_argument( @@ -1036,6 +1037,14 @@ def get_docker_cmd( return inject_spark_conf_str(original_docker_cmd, spark_conf_str) +def _get_adhoc_docker_registry(service: str, soa_dir: str = DEFAULT_SOA_DIR) -> str: + if service is None: + raise NotImplementedError('"None" is not a valid service') + + service_configuration = read_service_configuration(service, soa_dir) + return service_configuration.get("docker_registry", DEFAULT_SPARK_DOCKER_REGISTRY) + + def build_and_push_docker_image(args: argparse.Namespace) -> Optional[str]: """ Build an image if the default Spark service image is not preferred. @@ -1059,7 +1068,12 @@ def build_and_push_docker_image(args: argparse.Namespace) -> Optional[str]: if cook_return != 0: return None - docker_url = f"{args.docker_registry}/{docker_tag}" + registry_uri = args.docker_registry or _get_adhoc_docker_registry( + service=args.service, + soa_dir=args.yelpsoa_config_root, + ) + + docker_url = f"{registry_uri}/{docker_tag}" command = f"docker tag {docker_tag} {docker_url}" print(PaastaColors.grey(command)) retcode, _ = _run(command, stream=True)