Env var for custom Ray Image

Author: ChristianZaccaria

tests/e2e/mnist_raycluster_sdk_test.py

@@ -14,39 +14,46 @@
 
 import pytest
 
-from support import random_choice
+from support import random_choice, get_ray_image
 
 # Creates a Ray cluster, and trains the MNIST dataset using the CodeFlare SDK.
 # Asserts creation of AppWrapper, RayCluster, and successful completion of the training job.
 # Covers successfull installation of CodeFlare-SDK
 
+
 class TestMNISTRayClusterSDK:
     def setup_method(self):
         # Load the kube config from the environment or Kube config file.
         config.load_kube_config()
 
         # Initialize Kubernetes client
         self.api_instance = client.CoreV1Api()
-        self.custom_api = kubernetes.client.CustomObjectsApi(self.api_instance.api_client)
+        self.custom_api = kubernetes.client.CustomObjectsApi(
+            self.api_instance.api_client
+        )
 
     def teardown_method(self):
-        if hasattr(self, 'namespace'):
+        if hasattr(self, "namespace"):
             self.api_instance.delete_namespace(self.namespace)
-        if hasattr(self, 'configmap'):
-            self.api_instance.delete_namespaced_config_map(self.configmap.metadata.name, self.namespace)
+        if hasattr(self, "configmap"):
+            self.api_instance.delete_namespaced_config_map(
+                self.configmap.metadata.name, self.namespace
+            )
 
     def test_mnist_ray_cluster_sdk(self):
         self.create_test_namespace()
         self.run_mnist_raycluster_sdk()
 
     def create_test_namespace(self):
         self.namespace = f"test-ns-{random_choice()}"
-        namespace_body = client.V1Namespace(metadata=client.V1ObjectMeta(name=self.namespace))
+        namespace_body = client.V1Namespace(
+            metadata=client.V1ObjectMeta(name=self.namespace)
+        )
         self.api_instance.create_namespace(namespace_body)
         return self.namespace
 
     def run_mnist_raycluster_sdk(self):
-        ray_image = "quay.io/project-codeflare/ray:latest-py39-cu118"
+        ray_image = get_ray_image()
         host = os.getenv("CLUSTER_HOSTNAME")
 
         ingress_options = {}
@@ -61,7 +68,7 @@ def run_mnist_raycluster_sdk(self):
                         "host": host,
                         "annotations": {
                             "nginx.ingress.kubernetes.io/proxy-body-size": "100M",
-                        }
+                        },
                     },
                 ]
             }
@@ -127,17 +134,29 @@ def run_mnist_raycluster_sdk(self):
     # Assertions
     def assert_appwrapper_exists(self):
         try:
-            self.custom_api.get_namespaced_custom_object("workload.codeflare.dev", "v1beta1", self.namespace, "appwrappers", "mnist")
-            print(f"AppWrapper 'mnist' has been created in the namespace: '{self.namespace}'")
+            self.custom_api.get_namespaced_custom_object(
+                "workload.codeflare.dev",
+                "v1beta1",
+                self.namespace,
+                "appwrappers",
+                "mnist",
+            )
+            print(
+                f"AppWrapper 'mnist' has been created in the namespace: '{self.namespace}'"
+            )
             assert True
         except Exception as e:
             print(f"AppWrapper 'mnist' has not been created. Error: {e}")
             assert False
 
     def assert_raycluster_exists(self):
         try:
-            self.custom_api.get_namespaced_custom_object("ray.io", "v1", self.namespace, "rayclusters", "mnist")
-            print(f"RayCluster 'mnist' created successfully in the namespace: '{self.namespace}'")
+            self.custom_api.get_namespaced_custom_object(
+                "ray.io", "v1", self.namespace, "rayclusters", "mnist"
+            )
+            print(
+                f"RayCluster 'mnist' created successfully in the namespace: '{self.namespace}'"
+            )
             assert True
         except Exception as e:
             print(f"RayCluster 'mnist' has not been created. Error: {e}")

tests/e2e/support.py

@@ -1,14 +1,13 @@
+import os
 import random
 import string
 
-def read_file(file_name):
-    try:
-        with open(file_name, 'rb') as file:
-            return file.read()
-    except IOError as e:
-        raise e
+
+def get_ray_image():
+    default_ray_image = "quay.io/project-codeflare/ray:latest-py39-cu118"
+    return os.getenv("RAY_IMAGE", default_ray_image)
 
 
-alphabet = string.ascii_lowercase + string.digits
 def random_choice():
-    return ''.join(random.choices(alphabet, k=5))
+    alphabet = string.ascii_lowercase + string.digits
+    return "".join(random.choices(alphabet, k=5))