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
+        pull_secret = self.config.pull_secret
         return generate_appwrapper(
             name=name,
             namespace=namespace,
@@ -98,6 +99,7 @@ def create_app_wrapper(self):
             instascale=instascale,
             instance_types=instance_types,
             env=env,
+            pull_secret=pull_secret,
         )
 
     # 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"
+    pull_secret: dict = field(default_factory=dict)

src/codeflare_sdk/utils/generate_yaml.py

@@ -23,6 +23,11 @@
 import uuid
 
 
+class NoAliasDumper(yaml.Dumper):
+    def ignore_aliases(self, data):
+        return True
+
+
 def read_template(template):
     with open(template, "r") as stream:
         try:
@@ -130,6 +135,13 @@ def update_image(spec, image):
         container["image"] = image
 
 
+def update_pull_secret(spec, pull_secret):
+    if pull_secret:
+        if "imagePullSecrets" not in spec:
+            spec["imagePullSecrets"] = []
+        spec["imagePullSecrets"].append(pull_secret)
+
+
 def update_env(spec, env):
     containers = spec.get("containers")
     for container in containers:
@@ -167,6 +179,7 @@ def update_nodes(
     image,
     instascale,
     env,
+    pull_secret,
 ):
     if "generictemplate" in item.keys():
         head = item.get("generictemplate").get("spec").get("headGroupSpec")
@@ -182,6 +195,7 @@ def update_nodes(
         for comp in [head, worker]:
             spec = comp.get("template").get("spec")
             update_affinity(spec, appwrapper_name, instascale)
+            update_pull_secret(spec, pull_secret)
             update_image(spec, image)
             update_env(spec, env)
             if comp == head:
@@ -192,7 +206,7 @@ def update_nodes(
 
 def write_user_appwrapper(user_yaml, output_file_name):
     with open(output_file_name, "w") as outfile:
-        yaml.dump(user_yaml, outfile, default_flow_style=False)
+        yaml.dump(user_yaml, outfile, default_flow_style=False, Dumper=NoAliasDumper)
     print(f"Written to: {output_file_name}")
 
 
@@ -210,6 +224,7 @@ def generate_appwrapper(
     instascale: bool,
     instance_types: list,
     env,
+    pull_secret: str,
 ):
     user_yaml = read_template(template)
     appwrapper_name, cluster_name = gen_names(name)
@@ -233,6 +248,7 @@ def generate_appwrapper(
         image,
         instascale,
         env,
+        pull_secret,
     )
     update_dashboard_route(route_item, cluster_name, namespace)
     outfile = appwrapper_name + ".yaml"
@@ -314,6 +330,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(
+        "--pull-secret",
+        required=False,
+        default="",
+        help="Set pull secret for a private registry",
+    )
 
     args = parser.parse_args()
     name = args.name
@@ -329,6 +351,7 @@ def main():  # pragma: no cover
     instance_types = args.instance_types
     namespace = args.namespace
     env = {}
+    pull_secret = args.pull_secret
 
     outfile = generate_appwrapper(
         name,
@@ -344,6 +367,7 @@ def main():  # pragma: no cover
         instascale,
         instance_types,
         env,
+        pull_secret,
     )
     return outfile