diff --git a/src/codeflare_sdk/cluster/cluster.py b/src/codeflare_sdk/cluster/cluster.py index 0e0e73c06..c09e981b3 100644 --- a/src/codeflare_sdk/cluster/cluster.py +++ b/src/codeflare_sdk/cluster/cluster.py @@ -85,6 +85,7 @@ def create_app_wrapper(self): instance_types = self.config.machine_types env = self.config.envs local_interactive = self.config.local_interactive + image_pull_secrets = self.config.image_pull_secrets return generate_appwrapper( name=name, namespace=namespace, @@ -100,6 +101,7 @@ def create_app_wrapper(self): instance_types=instance_types, env=env, local_interactive=local_interactive, + image_pull_secrets=image_pull_secrets, ) # creates a new cluster with the provided or default spec diff --git a/src/codeflare_sdk/cluster/config.py b/src/codeflare_sdk/cluster/config.py index 25392db75..a00d128a2 100644 --- a/src/codeflare_sdk/cluster/config.py +++ b/src/codeflare_sdk/cluster/config.py @@ -49,3 +49,4 @@ class ClusterConfiguration: envs: dict = field(default_factory=dict) image: str = "ghcr.io/foundation-model-stack/base:ray2.1.0-py38-gpu-pytorch1.12.0cu116-20221213-193103" local_interactive: bool = False + image_pull_secrets: list = field(default_factory=list) diff --git a/src/codeflare_sdk/utils/generate_yaml.py b/src/codeflare_sdk/utils/generate_yaml.py index fffc4fa9a..9538a1e8e 100755 --- a/src/codeflare_sdk/utils/generate_yaml.py +++ b/src/codeflare_sdk/utils/generate_yaml.py @@ -141,6 +141,13 @@ def update_image(spec, image): container["image"] = image +def update_image_pull_secrets(spec, image_pull_secrets): + template_secrets = spec.get("imagePullSecrets", []) + spec["imagePullSecrets"] = template_secrets + [ + {"name": x} for x in image_pull_secrets + ] + + def update_env(spec, env): containers = spec.get("containers") for container in containers: @@ -178,6 +185,7 @@ def update_nodes( image, instascale, env, + image_pull_secrets, ): if "generictemplate" in item.keys(): head = item.get("generictemplate").get("spec").get("headGroupSpec") @@ -193,6 +201,7 @@ def update_nodes( for comp in [head, worker]: spec = comp.get("template").get("spec") update_affinity(spec, appwrapper_name, instascale) + update_image_pull_secrets(spec, image_pull_secrets) update_image(spec, image) update_env(spec, env) if comp == head: @@ -295,6 +304,7 @@ def generate_appwrapper( instance_types: list, env, local_interactive: bool, + image_pull_secrets: list, ): user_yaml = read_template(template) appwrapper_name, cluster_name = gen_names(name) @@ -318,6 +328,7 @@ def generate_appwrapper( image, instascale, env, + image_pull_secrets, ) update_dashboard_route(route_item, cluster_name, namespace) if local_interactive: @@ -409,6 +420,12 @@ def main(): # pragma: no cover default=False, help="Enable local interactive mode", ) + parser.add_argument( + "--image-pull-secrets", + required=False, + default=[], + help="Set image pull secrets for private registries", + ) args = parser.parse_args() name = args.name @@ -425,6 +442,7 @@ def main(): # pragma: no cover namespace = args.namespace local_interactive = args.local_interactive env = {} + image_pull_secrets = args.image_pull_secrets outfile = generate_appwrapper( name, @@ -441,6 +459,7 @@ def main(): # pragma: no cover instance_types, local_interactive, env, + image_pull_secrets, ) return outfile diff --git a/tests/test-case-cmd.yaml b/tests/test-case-cmd.yaml index d82096e8f..ea235ec9a 100644 --- a/tests/test-case-cmd.yaml +++ b/tests/test-case-cmd.yaml @@ -96,6 +96,7 @@ spec: cpu: 2 memory: 8G nvidia.com/gpu: 0 + imagePullSecrets: [] rayVersion: 2.1.0 workerGroupSpecs: - groupName: small-group-unit-cmd-cluster @@ -144,6 +145,7 @@ spec: cpu: 1 memory: 2G nvidia.com/gpu: 1 + imagePullSecrets: [] initContainers: - command: - sh diff --git a/tests/test-case.yaml b/tests/test-case.yaml index 0d85428df..da6058698 100644 --- a/tests/test-case.yaml +++ b/tests/test-case.yaml @@ -107,6 +107,8 @@ spec: cpu: 2 memory: 8G nvidia.com/gpu: 0 + imagePullSecrets: + - name: unit-test-pull-secret rayVersion: 2.1.0 workerGroupSpecs: - groupName: small-group-unit-test-cluster @@ -164,6 +166,8 @@ spec: cpu: 3 memory: 5G nvidia.com/gpu: 7 + imagePullSecrets: + - name: unit-test-pull-secret initContainers: - command: - sh diff --git a/tests/unit_test.py b/tests/unit_test.py index ead3521c8..d6e6c6ef6 100644 --- a/tests/unit_test.py +++ b/tests/unit_test.py @@ -220,6 +220,7 @@ def test_config_creation(): gpu=7, instascale=True, machine_types=["cpu.small", "gpu.large"], + image_pull_secrets=["unit-test-pull-secret"], ) assert config.name == "unit-test-cluster" and config.namespace == "ns" @@ -234,6 +235,7 @@ def test_config_creation(): assert config.template == f"{parent}/src/codeflare_sdk/templates/base-template.yaml" assert config.instascale assert config.machine_types == ["cpu.small", "gpu.large"] + assert config.image_pull_secrets == ["unit-test-pull-secret"] return config