expose more options to the user optionally

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

Author: KPostOffice

src/codeflare_sdk/jobs/config.py

@@ -19,17 +19,20 @@
 """
 
 from dataclasses import dataclass, field
-from ..cluster.cluster import Cluster
+from typing import Optional, Dict
 
-
-@dataclass
+@dataclass(slots=True)
 class JobConfiguration:
     """
     This dataclass is used to specify resource requirements and other details, and
     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
+    debug: bool = False

src/codeflare_sdk/jobs/jobs.py

@@ -35,16 +35,12 @@
 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 submit(self, cluster: "Cluster"):
         """
         Method for creating a job on a specific cluster
         """
@@ -76,7 +72,7 @@ def __init__(self, config: JobConfiguration):
         """
         self.config = config
 
-    def submit(self, cluster: Cluster) -> "TorchXRayJob":
+    def submit(self, cluster: "Cluster") -> "TorchXRayJob":
         """
         Submit the job definition to a specific cluster, resulting in a Job object.
         """
@@ -87,48 +83,53 @@ 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):
+    def __init__(self, job_definition: TorchXJobDefinition, cluster: "Cluster", *script_args):
         """
-        TODO
+        Creates job which maximizes resource usage on the passed cluster.
         """
         self.job_definition: TorchXJobDefinition = job_definition
-        self.cluster: Cluster = cluster
+        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(
             app=ddp(
                 *script_args,
                 script = job_definition.config.script,
-                m=None,  # python module to run (might be worth exposing)
-                name = job_definition.config.name,
+                m=job_definition.config.m,
+                name=job_definition.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=job_definition.config.env,
+                debug=job_definition.config.debug  # TODO: expose
+                # 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": self.cluster.config.name,
+                "dashboard_address": self.cluster.cluster_dashboard_uri(self.cluster.config.namespace).split(":")[0],
+                "working_dir": job_definition.config.working_dir,
+                "requirements": job_definition.config.requirements,
+            },
         )
 
         _, _, self.job_id = parse_app_handle(_app_handle)
         all_jobs.append(self)
 
-    def status(self):
+    def status(self) -> str:
         """
-        TODO
+        Get running job status.
         """
-        dashboard_route = self.cluster.cluster_dashboard_uri()
+        dashboard_route = self.cluster.cluster_dashboard_uri(namespace=self.cluster.config.namespace)
         client = JobSubmissionClient(dashboard_route)
         return client.get_job_status(self.job_id)
 
-    def logs(self):
+    def logs(self) -> str:
         """
-        TODO
+        Get job logs.
         """
-        dashboard_route = self.cluster_dashboard_uri(namespace=self.config.namespace)
+        dashboard_route = self.cluster_dashboard_uri(namespace=self.cluster.config.namespace)
         client = JobSubmissionClient(dashboard_route)
-        return client.get_job_logs(job_id)
+        return client.get_job_logs(self.job_id)