Convert Go raycluster_sdk_upgrade_test to Python test

Author: abhijeet-dhumal

tests/e2e/raycluster_sdk_upgrade_test.py

@@ -0,0 +1,172 @@
+import requests
+from time import sleep
+
+from torchx.specs.api import AppState, is_terminal
+
+from codeflare_sdk import Cluster, ClusterConfiguration, TokenAuthentication
+from codeflare_sdk.job.jobs import DDPJobDefinition
+
+import pytest
+
+from support import *
+from codeflare_sdk.cluster.cluster import get_cluster
+
+namespace = "test-ns-rayupgrade"
+
+
+# Creates a Ray cluster
+@pytest.mark.openshift
+class TestMNISTRayClusterUp:
+    def setup_method(self):
+        initialize_kubernetes_client(self)
+
+    def test_mnist_ray_cluster_sdk_auth(self):
+        create_namespace_with_name(self, namespace)
+        self.run_mnist_raycluster_sdk_oauth()
+
+    def run_mnist_raycluster_sdk_oauth(self):
+        ray_image = get_ray_image()
+
+        auth = TokenAuthentication(
+            token=run_oc_command(["whoami", "--show-token=true"]),
+            server=run_oc_command(["whoami", "--show-server=true"]),
+            skip_tls=True,
+        )
+        auth.login()
+
+        cluster = Cluster(
+            ClusterConfiguration(
+                name="mnist",
+                namespace=self.namespace,
+                num_workers=1,
+                head_cpus="1",
+                head_memory=2,
+                min_cpus="1",
+                max_cpus=1,
+                min_memory=1,
+                max_memory=2,
+                num_gpus=0,
+                instascale=False,
+                image=ray_image,
+                openshift_oauth=True,
+                write_to_file=True,
+            )
+        )
+
+        try:
+            cluster.up()
+            # check if Appwrapper has been created in namespace
+            self.assert_appwrapper_exists()
+            cluster.status()
+            # wait for raycluster to be Ready
+            cluster.wait_ready()
+            cluster.status()
+            # Check cluster details
+            cluster.details()
+            # Assert the cluster status is READY
+            _, ready = cluster.status()
+            assert ready, "Cluster is not ready!"
+
+        except Exception as e:
+            print(f"An unexpected error occurred. Error: ", e)
+            delete_namespace(self)
+
+    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}'"
+            )
+            assert True
+        except Exception as e:
+            print(f"AppWrapper 'mnist' has not been created. Error: {e}")
+            assert False
+
+
+class TestMnistJobSubmit:
+    def setup_method(self):
+        initialize_kubernetes_client(self)
+        self.namespace = namespace
+        self.cluster = get_cluster("mnist", self.namespace)
+        if not self.cluster:
+            raise RuntimeError("TestRayClusterUp needs to be run before this test")
+
+    def teardown_method(self):
+        delete_namespace(self)
+
+    def test_mnist_job_submission(self):
+        auth = TokenAuthentication(
+            token=run_oc_command(["whoami", "--show-token=true"]),
+            server=run_oc_command(["whoami", "--show-server=true"]),
+            skip_tls=True,
+        )
+        auth.login()
+
+        self.assert_jobsubmit_withoutLogin(self.cluster)
+        self.assert_jobsubmit_withlogin(self.cluster)
+        self.cluster.down()
+
+    # Assertions
+    def assert_jobsubmit_withoutLogin(self, cluster):
+        dashboard_url = cluster.cluster_dashboard_uri()
+        jobdata = {
+            "entrypoint": "python mnist.py",
+            "runtime_env": {
+                "working_dir": "./tests/e2e/",
+                "pip": "mnist_pip_requirements.txt",
+            },
+        }
+        try:
+            response = requests.post(
+                dashboard_url + "/api/jobs/", verify=False, json=jobdata
+            )
+            if response.status_code == 403:
+                assert True
+            else:
+                response.raise_for_status()
+                assert False
+
+        except Exception as e:
+            print(f"An unexpected error occurred. Error: {e}")
+            assert False
+
+    def assert_jobsubmit_withlogin(self, cluster):
+        jobdef = DDPJobDefinition(
+            name="mnist",
+            script="./tests/e2e/mnist.py",
+            scheduler_args={"requirements": "./tests/e2e/mnist_pip_requirements.txt"},
+        )
+        job = jobdef.submit(cluster)
+
+        done = False
+        time = 0
+        timeout = 900
+        while not done:
+            status = job.status()
+            if is_terminal(status.state):
+                break
+            if not done:
+                print(status)
+                if timeout and time >= timeout:
+                    raise TimeoutError(f"job has timed out after waiting {timeout}s")
+                sleep(5)
+                time += 5
+
+        print(job.status())
+        self.assert_job_completion(status)
+
+        print(job.logs())
+
+    def assert_job_completion(self, status):
+        if status.state == AppState.SUCCEEDED:
+            print(f"Job has completed: '{status.state}'")
+            assert True
+        else:
+            print(f"Job has completed: '{status.state}'")
+            assert False

tests/e2e/support.py

@@ -24,6 +24,14 @@ def create_namespace(self):
     self.api_instance.create_namespace(namespace_body)
 
 
+def create_namespace_with_name(self, namespace_name):
+    self.namespace = namespace_name
+    namespace_body = client.V1Namespace(
+        metadata=client.V1ObjectMeta(name=self.namespace)
+    )
+    self.api_instance.create_namespace(namespace_body)
+
+
 def delete_namespace(self):
     if hasattr(self, "namespace"):
         self.api_instance.delete_namespace(self.namespace)