expose more options to the user optionally

Signed-off-by: Kevin <[email protected]>

Author: KPostOffice

src/codeflare_sdk/jobs/config.py

@@ -19,8 +19,7 @@
 """
 
 from dataclasses import dataclass, field
-from ..cluster.cluster import Cluster
-
+from typing import Optional, Dict
 
 @dataclass
 class JobConfiguration:
@@ -29,7 +28,10 @@ class JobConfiguration:
     is passed in as an argument when creating a Job object.
     """
 
-    name: str = None
-    script: str = None
-    requirements: str = None
-    scheduler: str = "ray"
+    name: Optional[str] = None
+    script: Optional[str] = None
+    m: Optional[str] = None
+    h: Optional[str] = None  # custom resource types
+    env: Optional[Dict[str, str]] = None
+    working_dir: Optional[str] = None
+    requirements: Optional[str] = None

src/codeflare_sdk/jobs/jobs.py

@@ -20,11 +20,12 @@
 
 import abc
 from typing import List
+from pathlib import Path
 
 from ray.job_submission import JobSubmissionClient
 from torchx.components.dist import ddp
 from torchx.runner import get_runner, Runner
-from torchx.specs import AppHandle, parse_app_handle
+from torchx.specs import AppHandle, parse_app_handle, AppDryRunInfo
 
 from .config import JobConfiguration
 
@@ -35,16 +36,19 @@
 all_jobs: List["Job"] = []
 torchx_runner: Runner = get_runner()
 
-torchx_runner.run(app, scheduler, cfg, workspace)
-
-torchx_runner.run_component(component, component_args, scheduler, cfg, workspace)
-
 class JobDefinition(metaclass=abc.ABCMeta):
     """
     A job definition to be submitted to a generic backend cluster.
     """
 
-    def submit(self, cluster: Cluster):
+    def _dry_run(self, cluster) -> str:
+        """
+        Create job definition, but do not submit.
+
+        The primary purpose of this function is to facilitate unit testing.
+        """
+
+    def submit(self, cluster: "Cluster"):
         """
         Method for creating a job on a specific cluster
         """
@@ -76,59 +80,70 @@ def __init__(self, config: JobConfiguration):
         """
         self.config = config
 
-    def submit(self, cluster: Cluster) -> "TorchXRayJob":
-        """
-        Submit the job definition to a specific cluster, resulting in a Job object.
+    def _dry_run(self, cluster: "Cluster", *script_args) -> AppDryRunInfo:
         """
-        return TorchXRayJob(self, cluster)
-
+        Create job definition, but do not submit.
 
-class TorchXRayJob(Job):
-    """
-    Active submission of a dist.ddp job to a Ray cluster which can be used to get logs and status.
-    """
-    def __init__(self, job_definition: TorchXJobDefinition, cluster: Cluster, *script_args):
+        The primary purpose of this function is to facilitate unit testing.
         """
-        TODO
-        """
-        self.job_definition: TorchXJobDefinition = job_definition
-        self.cluster: Cluster = cluster
         j = f"{cluster.config.max_worker}x{max(cluster.config.gpu, 1)}"  # # of proc. = # of gpus
-        # TODO: allow user to override resource allocation for job
-        _app_handle: AppHandle = torchx_runner.run(
+        dashboard_address = f"{cluster.cluster_dashboard_uri(cluster.config.namespace).lstrip('http://')}:8265"
+        return torchx_runner.dryrun(
             app=ddp(
                 *script_args,
-                script = job_definition.config.script,
-                m=None,  # python module to run (might be worth exposing)
-                name = job_definition.config.name,
+                script = self.config.script,
+                m=self.config.m,
+                name=self.config.name,
                 h=None,  # for custom resource types
-                cpu=cluster.config.max_cpus,  # TODO: get from cluster config
+                cpu=cluster.config.max_cpus,
                 gpu = cluster.config.gpu,
                 memMB = 1024 * cluster.config.max_memory,  # cluster memory is in GB
                 j=j,
-                env=None,  # TODO: should definitely expose Dict[str, str]
-                max_retries = 0,  # TODO: maybe expose
-                mounts=None,  # TODO: maybe expose
-                debug=False  # TODO: expose
+                env=self.config.env,
+                # max_retries=0,  # default
+                # mounts=None,  # default
             ),
-            scheduler="ray", cfg="fo",
+            scheduler="ray",  # can be determined by type of cluster if more are introduced
+            cfg={
+                "cluster_name": cluster.config.name,
+                "dashboard_address": "localhost:8265",  # dashboard_address,
+                "working_dir": self.config.working_dir,
+                "requirements": self.config.requirements,
+            },
+            workspace=f"file://{Path.cwd()}"
         )
 
-        _, _, self.job_id = parse_app_handle(_app_handle)
+    def submit(self, cluster: "Cluster") -> "TorchXRayJob":
+        """
+        Submit the job definition to a specific cluster, resulting in a Job object.
+        """
+        return TorchXRayJob(self, cluster)
+
+
+class TorchXRayJob(Job):
+    """
+    Active submission of a dist.ddp job to a Ray cluster which can be used to get logs and status.
+    """
+    def __init__(self, job_definition: TorchXJobDefinition, cluster: "Cluster", *script_args):
+        """
+        Creates job which maximizes resource usage on the passed cluster.
+        """
+        self.job_definition: TorchXJobDefinition = job_definition
+        self.cluster: "Cluster" = cluster
+        # dashboard_address = f"{self.cluster.cluster_dashboard_uri(self.cluster.config.namespace).lstrip('http://')}:8265"
+        self._app_handle = torchx_runner.schedule(job_definition._dry_run(cluster, *script_args))
+        _, _, self.job_id = parse_app_handle(self._app_handle)
+        # self.job_id = self.job_id.lstrip(f"{dashboard_address}-")
         all_jobs.append(self)
 
-    def status(self):
+    def status(self) -> str:
         """
-        TODO
+        Get running job status.
         """
-        dashboard_route = self.cluster.cluster_dashboard_uri()
-        client = JobSubmissionClient(dashboard_route)
-        return client.get_job_status(self.job_id)
+        return torchx_runner.status(self._app_handle)
 
-    def logs(self):
+    def logs(self) -> str:
         """
-        TODO
+        Get job logs.
         """
-        dashboard_route = self.cluster_dashboard_uri(namespace=self.config.namespace)
-        client = JobSubmissionClient(dashboard_route)
-        return client.get_job_logs(job_id)
+        return "".join(torchx_runner.log_lines(self._app_handle, None))