schedulers parameterized

Differential Revision: D71023681

Pull Request resolved: #1022

Author: bobyangyf

torchx/schedulers/api.py

@@ -94,9 +94,11 @@ def __hash__(self) -> int:
 
 
 T = TypeVar("T")
+A = TypeVar("A")
+D = TypeVar("D")
 
 
-class Scheduler(abc.ABC, Generic[T]):
+class Scheduler(abc.ABC, Generic[T, A, D]):
     """
     An interface abstracting functionalities of a scheduler.
     Implementers need only implement those methods annotated with
@@ -126,7 +128,7 @@ def close(self) -> None:
 
     def submit(
         self,
-        app: AppDef,
+        app: A,
         cfg: T,
         workspace: Optional[str] = None,
     ) -> str:
@@ -150,7 +152,7 @@ def submit(
         return self.schedule(dryrun_info)
 
     @abc.abstractmethod
-    def schedule(self, dryrun_info: AppDryRunInfo) -> str:
+    def schedule(self, dryrun_info: D) -> str:
         """
         Same as ``submit`` except that it takes an ``AppDryRunInfo``.
         Implementers are encouraged to implement this method rather than
@@ -166,7 +168,7 @@ def schedule(self, dryrun_info: AppDryRunInfo) -> str:
 
         raise NotImplementedError()
 
-    def submit_dryrun(self, app: AppDef, cfg: T) -> AppDryRunInfo:
+    def submit_dryrun(self, app: A, cfg: T) -> D:
         """
         Rather than submitting the request to run the app, returns the
         request object that would have been submitted to the underlying
@@ -179,14 +181,16 @@ def submit_dryrun(self, app: AppDef, cfg: T) -> AppDryRunInfo:
         resolved_cfg = self.run_opts().resolve(cfg)
         # pyre-fixme: _submit_dryrun takes Generic type for resolved_cfg
         dryrun_info = self._submit_dryrun(app, resolved_cfg)
-        for role in app.roles:
-            dryrun_info = role.pre_proc(self.backend, dryrun_info)
-        dryrun_info._app = app
-        dryrun_info._cfg = resolved_cfg
+
+        if isinstance(app, AppDef):
+            for role in app.roles:
+                dryrun_info = role.pre_proc(self.backend, dryrun_info)
+            dryrun_info._app = app
+            dryrun_info._cfg = resolved_cfg
         return dryrun_info
 
     @abc.abstractmethod
-    def _submit_dryrun(self, app: AppDef, cfg: T) -> AppDryRunInfo:
+    def _submit_dryrun(self, app: A, cfg: T) -> D:
         raise NotImplementedError()
 
     def run_opts(self) -> runopts:
@@ -345,18 +349,19 @@ def _pre_build_validate(self, app: AppDef, scheduler: str, cfg: T) -> None:
         """
         pass
 
-    def _validate(self, app: AppDef, scheduler: str, cfg: T) -> None:
+    def _validate(self, app: A, scheduler: str, cfg: T) -> None:
         """
         Validates after workspace build whether application is consistent with the scheduler.
 
         Raises error if application is not compatible with scheduler
         """
-        for role in app.roles:
-            if role.resource == NULL_RESOURCE:
-                raise ValueError(
-                    f"No resource for role: {role.image}."
-                    f" Did you forget to attach resource to the role"
-                )
+        if isinstance(app, AppDef):
+            for role in app.roles:
+                if role.resource == NULL_RESOURCE:
+                    raise ValueError(
+                        f"No resource for role: {role.image}."
+                        f" Did you forget to attach resource to the role"
+                    )
 
 
 def filter_regex(regex: str, data: Iterable[str]) -> Iterable[str]:

torchx/schedulers/aws_batch_scheduler.py

@@ -363,7 +363,9 @@ class AWSBatchOpts(TypedDict, total=False):
     execution_role_arn: Optional[str]
 
 
-class AWSBatchScheduler(DockerWorkspaceMixin, Scheduler[AWSBatchOpts]):
+class AWSBatchScheduler(
+    DockerWorkspaceMixin, Scheduler[AWSBatchOpts, AppDef, AppDryRunInfo[BatchJob]]
+):
     """
     AWSBatchScheduler is a TorchX scheduling interface to AWS Batch.
 

torchx/schedulers/aws_sagemaker_scheduler.py

@@ -156,7 +156,10 @@ def _merge_ordered(
     return merged
 
 
-class AWSSageMakerScheduler(DockerWorkspaceMixin, Scheduler[AWSSageMakerOpts]):  # type: ignore[misc]
+class AWSSageMakerScheduler(
+    DockerWorkspaceMixin,
+    Scheduler[AWSSageMakerOpts, AppDef, AppDryRunInfo[AWSSageMakerJob]],
+):
     """
     AWSSageMakerScheduler is a TorchX scheduling interface to AWS SageMaker.
 

torchx/schedulers/docker_scheduler.py

@@ -128,7 +128,9 @@ class DockerOpts(TypedDict, total=False):
     privileged: bool
 
 
-class DockerScheduler(DockerWorkspaceMixin, Scheduler[DockerOpts]):
+class DockerScheduler(
+    DockerWorkspaceMixin, Scheduler[DockerOpts, AppDef, AppDryRunInfo[DockerJob]]
+):
     """
     DockerScheduler is a TorchX scheduling interface to Docker.
 

torchx/schedulers/gcp_batch_scheduler.py

@@ -104,7 +104,7 @@ class GCPBatchOpts(TypedDict, total=False):
     location: Optional[str]
 
 
-class GCPBatchScheduler(Scheduler[GCPBatchOpts]):
+class GCPBatchScheduler(Scheduler[GCPBatchOpts, AppDef, AppDryRunInfo[GCPBatchJob]]):
     """
     GCPBatchScheduler is a TorchX scheduling interface to GCP Batch.
 

torchx/schedulers/kubernetes_mcad_scheduler.py

@@ -796,7 +796,9 @@ class KubernetesMCADOpts(TypedDict, total=False):
     network: Optional[str]
 
 
-class KubernetesMCADScheduler(DockerWorkspaceMixin, Scheduler[KubernetesMCADOpts]):
+class KubernetesMCADScheduler(
+    DockerWorkspaceMixin, Scheduler[KubernetesMCADOpts, AppDef, AppDryRunInfo]
+):
     """
     KubernetesMCADScheduler is a TorchX scheduling interface to Kubernetes.
 

torchx/schedulers/kubernetes_scheduler.py

@@ -472,7 +472,10 @@ class KubernetesOpts(TypedDict, total=False):
     priority_class: Optional[str]
 
 
-class KubernetesScheduler(DockerWorkspaceMixin, Scheduler[KubernetesOpts]):
+class KubernetesScheduler(
+    DockerWorkspaceMixin,
+    Scheduler[KubernetesOpts, AppDef, AppDryRunInfo[KubernetesJob]],
+):
     """
     KubernetesScheduler is a TorchX scheduling interface to Kubernetes.
 

torchx/schedulers/local_scheduler.py

@@ -529,7 +529,7 @@ def _register_termination_signals() -> None:
         signal.signal(signal.SIGINT, _terminate_process_handler)
 
 
-class LocalScheduler(Scheduler[LocalOpts]):
+class LocalScheduler(Scheduler[LocalOpts, AppDef, AppDryRunInfo[PopenRequest]]):
     """
     Schedules on localhost. Containers are modeled as processes and
     certain properties of the container that are either not relevant

torchx/schedulers/lsf_scheduler.py

@@ -395,7 +395,7 @@ def __repr__(self) -> str:
 {self.materialize()}"""
 
 
-class LsfScheduler(Scheduler[LsfOpts]):
+class LsfScheduler(Scheduler[LsfOpts, AppDef, AppDryRunInfo]):
     """
     **Example: hello_world**
 

torchx/schedulers/ray_scheduler.py

@@ -114,7 +114,9 @@ class RayJob:
         requirements: Optional[str] = None
         actors: List[RayActor] = field(default_factory=list)
 
-    class RayScheduler(TmpDirWorkspaceMixin, Scheduler[RayOpts]):
+    class RayScheduler(
+        TmpDirWorkspaceMixin, Scheduler[RayOpts, AppDef, AppDryRunInfo[RayJob]]
+    ):
         """
         RayScheduler is a TorchX scheduling interface to Ray. The job def
         workers will be launched as Ray actors