add: Support for image pull secrets for Ray Cluster images

Author: VanillaSpoon

src/codeflare_sdk/cluster/cluster.py

@@ -84,6 +84,7 @@ def create_app_wrapper(self):
         instascale = self.config.instascale
         instance_types = self.config.machine_types
         env = self.config.envs
+        image_pull_secrets = self.config.image_pull_secrets
         return generate_appwrapper(
             name=name,
             namespace=namespace,
@@ -98,6 +99,7 @@ def create_app_wrapper(self):
             instascale=instascale,
             instance_types=instance_types,
             env=env,
+            image_pull_secrets=image_pull_secrets,
         )
 
     # creates a new cluster with the provided or default spec

src/codeflare_sdk/cluster/config.py

@@ -48,3 +48,4 @@ class ClusterConfiguration:
     instascale: bool = False
     envs: dict = field(default_factory=dict)
     image: str = "ghcr.io/foundation-model-stack/base:ray2.1.0-py38-gpu-pytorch1.12.0cu116-20221213-193103"
+    image_pull_secrets: list = field(default_factory=list)

src/codeflare_sdk/utils/generate_yaml.py

@@ -130,6 +130,14 @@ def update_image(spec, image):
         container["image"] = image
 
 
+def update_image_pull_secrets(spec, image_pull_secrets):
+    if image_pull_secrets:
+        if "imagePullSecrets" not in spec:
+            spec["imagePullSecrets"] = []
+        for image_pull_secret in image_pull_secrets:
+            spec["imagePullSecrets"].append({"name": image_pull_secret})
+
+
 def update_env(spec, env):
     containers = spec.get("containers")
     for container in containers:
@@ -167,6 +175,7 @@ def update_nodes(
     image,
     instascale,
     env,
+    image_pull_secrets,
 ):
     if "generictemplate" in item.keys():
         head = item.get("generictemplate").get("spec").get("headGroupSpec")
@@ -182,6 +191,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:
@@ -211,6 +221,7 @@ def generate_appwrapper(
     instascale: bool,
     instance_types: list,
     env,
+    image_pull_secrets: list,
 ):
     user_yaml = read_template(template)
     appwrapper_name, cluster_name = gen_names(name)
@@ -234,6 +245,7 @@ def generate_appwrapper(
         image,
         instascale,
         env,
+        image_pull_secrets,
     )
     update_dashboard_route(route_item, cluster_name, namespace)
     outfile = appwrapper_name + ".yaml"
@@ -315,6 +327,12 @@ def main():  # pragma: no cover
         default="default",
         help="Set the kubernetes namespace you want to deploy your cluster to. Default. If left blank, uses the 'default' namespace",
     )
+    parser.add_argument(
+        "--image-pull-secrets",
+        required=False,
+        default=[],
+        help="Set image pull secrets for private registries",
+    )
 
     args = parser.parse_args()
     name = args.name
@@ -330,6 +348,7 @@ def main():  # pragma: no cover
     instance_types = args.instance_types
     namespace = args.namespace
     env = {}
+    image_pull_secrets = args.image_pull_secrets
 
     outfile = generate_appwrapper(
         name,
@@ -345,6 +364,7 @@ def main():  # pragma: no cover
         instascale,
         instance_types,
         env,
+        image_pull_secrets,
     )
     return outfile