remove: DDPJobDefinition from SDK (#498)

* remove: DDPJobDefinition and update tests

* add: address comments

Author: VanillaSpoon

Diff for: src/codeflare_sdk/__init__.py

@@ -16,6 +16,6 @@
     list_all_clusters,
 )
 
-from .job import JobDefinition, Job, DDPJobDefinition, DDPJob, RayJobClient
+from .job import RayJobClient
 
 from .utils import generate_cert

Diff for: src/codeflare_sdk/cluster/cluster.py

@@ -21,10 +21,8 @@
 from time import sleep
 from typing import List, Optional, Tuple, Dict
 
-import openshift as oc
 from kubernetes import config
 from ray.job_submission import JobSubmissionClient
-import urllib3
 
 from .auth import config_check, api_config_handler
 from ..utils import pretty_print
@@ -58,8 +56,6 @@ class Cluster:
     Note that currently, the underlying implementation is a Ray cluster.
     """
 
-    torchx_scheduler = "ray"
-
     def __init__(self, config: ClusterConfiguration):
         """
         Create the resource cluster object by passing in a ClusterConfiguration
@@ -477,20 +473,6 @@ def job_logs(self, job_id: str) -> str:
         """
         return self.job_client.get_job_logs(job_id)
 
-    def torchx_config(
-        self, working_dir: str = None, requirements: str = None
-    ) -> Dict[str, str]:
-        dashboard_address = urllib3.util.parse_url(self.cluster_dashboard_uri()).host
-        to_return = {
-            "cluster_name": self.config.name,
-            "dashboard_address": dashboard_address,
-        }
-        if working_dir:
-            to_return["working_dir"] = working_dir
-        if requirements:
-            to_return["requirements"] = requirements
-        return to_return
-
     def from_k8_cluster_object(
         rc,
         mcad=True,

Diff for: src/codeflare_sdk/job/__init__.py

@@ -1,3 +1 @@
-from .jobs import JobDefinition, Job, DDPJobDefinition, DDPJob
-
 from .ray_jobs import RayJobClient

Diff for: src/codeflare_sdk/job/jobs.py

@@ -2,10 +2,8 @@
 
 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
+from codeflare_sdk.job import RayJobClient
 
 import pytest
 
@@ -79,7 +77,7 @@ def assert_jobsubmit_withoutLogin(self, cluster):
             "entrypoint": "python mnist.py",
             "runtime_env": {
                 "working_dir": "./tests/e2e/",
-                "pip": "mnist_pip_requirements.txt",
+                "pip": "./tests/e2e/mnist_pip_requirements.txt",
             },
         }
         try:
@@ -98,19 +96,26 @@ def assert_jobsubmit_withoutLogin(self, cluster):
 
     def assert_jobsubmit_withlogin(self, cluster):
         self.assert_appwrapper_exists()
-        jobdef = DDPJobDefinition(
-            name="mnist",
-            script="./tests/e2e/mnist.py",
-            scheduler_args={"requirements": "./tests/e2e/mnist_pip_requirements.txt"},
+        auth_token = run_oc_command(["whoami", "--show-token=true"])
+        ray_dashboard = cluster.cluster_dashboard_uri()
+        header = {"Authorization": f"Bearer {auth_token}"}
+        client = RayJobClient(address=ray_dashboard, headers=header, verify=True)
+
+        # Submit the job
+        submission_id = client.submit_job(
+            entrypoint="python mnist.py",
+            runtime_env={
+                "working_dir": "./tests/e2e/",
+                "pip": "mnist_pip_requirements.txt",
+            },
         )
-        job = jobdef.submit(cluster)
-
+        print(f"Submitted job with ID: {submission_id}")
         done = False
         time = 0
         timeout = 900
         while not done:
-            status = job.status()
-            if is_terminal(status.state):
+            status = client.get_job_status(submission_id)
+            if status.is_terminal():
                 break
             if not done:
                 print(status)
@@ -119,11 +124,12 @@ def assert_jobsubmit_withlogin(self, cluster):
                 sleep(5)
                 time += 5
 
-        print(job.status())
-        self.assert_job_completion(status)
+        logs = client.get_job_logs(submission_id)
+        print(logs)
 
-        print(job.logs())
+        self.assert_job_completion(status)
 
+        client.delete_job(submission_id)
         cluster.down()
 
     def assert_appwrapper_exists(self):
@@ -144,9 +150,9 @@ def assert_appwrapper_exists(self):
             assert False
 
     def assert_job_completion(self, status):
-        if status.state == AppState.SUCCEEDED:
-            print(f"Job has completed: '{status.state}'")
+        if status == "SUCCEEDED":
+            print(f"Job has completed: '{status}'")
             assert True
         else:
-            print(f"Job has completed: '{status.state}'")
+            print(f"Job has completed: '{status}'")
             assert False