torchx - expose scheduler opts to scheduler validate function interface

Differential Revision: D67032224

Pull Request resolved: meta-pytorch#989

Author: tonykao8080

torchx/runner/api.py

@@ -433,7 +433,7 @@ def dryrun(
                         " Either a patch was built or no changes to workspace was detected."
                     )
 
-            sched._validate(app, scheduler)
+            sched._validate(app, scheduler, resolved_cfg)
             dryrun_info = sched.submit_dryrun(app, resolved_cfg)
             dryrun_info._scheduler = scheduler
             return dryrun_info

torchx/schedulers/api.py

@@ -337,7 +337,7 @@ def log_iter(
             f"{self.__class__.__qualname__} does not support application log iteration"
         )
 
-    def _validate(self, app: AppDef, scheduler: str) -> None:
+    def _validate(self, app: AppDef, scheduler: str, cfg: T) -> None:
         """
         Validates whether application is consistent with the scheduler.
 

torchx/schedulers/docker_scheduler.py

@@ -327,7 +327,7 @@ def _submit_dryrun(self, app: AppDef, cfg: DockerOpts) -> AppDryRunInfo[DockerJo
 
         return AppDryRunInfo(req, repr)
 
-    def _validate(self, app: AppDef, scheduler: str) -> None:
+    def _validate(self, app: AppDef, scheduler: str, cfg: DockerOpts) -> None:
         # Skip validation step
         pass
 

torchx/schedulers/gcp_batch_scheduler.py

@@ -464,7 +464,7 @@ def list(self) -> List[ListAppResponse]:
             for job in all_jobs
         ]
 
-    def _validate(self, app: AppDef, scheduler: str) -> None:
+    def _validate(self, app: AppDef, scheduler: str, cfg: GCPBatchOpts) -> None:
         # Skip validation step
         pass
 

torchx/schedulers/kubernetes_mcad_scheduler.py

@@ -1033,7 +1033,7 @@ def _submit_dryrun(
         info._cfg = cfg
         return info
 
-    def _validate(self, app: AppDef, scheduler: str) -> None:
+    def _validate(self, app: AppDef, scheduler: str, cfg: KubernetesMCADOpts) -> None:
         # Skip validation step
         pass
 

torchx/schedulers/kubernetes_scheduler.py

@@ -661,7 +661,7 @@ def _submit_dryrun(
         )
         return AppDryRunInfo(req, repr)
 
-    def _validate(self, app: AppDef, scheduler: str) -> None:
+    def _validate(self, app: AppDef, scheduler: str, cfg: KubernetesOpts) -> None:
         # Skip validation step
         pass
 

torchx/schedulers/local_scheduler.py

@@ -630,7 +630,7 @@ def _run_opts(self) -> runopts:
         )
         return opts
 
-    def _validate(self, app: AppDef, scheduler: str) -> None:
+    def _validate(self, app: AppDef, scheduler: str, cfg: LocalOpts) -> None:
         # Skip validation step for local application
         pass
 

torchx/schedulers/lsf_scheduler.py

@@ -488,7 +488,7 @@ def schedule(self, dryrun_info: AppDryRunInfo[LsfBsub]) -> str:
             subprocess.run(req.cmd, stdout=subprocess.PIPE, check=True)
         return req.app_id
 
-    def _validate(self, app: AppDef, scheduler: str) -> None:
+    def _validate(self, app: AppDef, scheduler: str, cfg: LsfOpts) -> None:
         # Skip validation step for lsf
         pass
 

torchx/schedulers/ray_scheduler.py

@@ -318,7 +318,7 @@ def _submit_dryrun(self, app: AppDef, cfg: RayOpts) -> AppDryRunInfo[RayJob]:
 
             return AppDryRunInfo(job, repr)
 
-        def _validate(self, app: AppDef, scheduler: str) -> None:
+        def _validate(self, app: AppDef, scheduler: str, cfg: RayOpts) -> None:
             if scheduler != "ray":
                 raise ValueError(
                     f"An unknown scheduler backend '{scheduler}' has been passed to the Ray scheduler."

torchx/schedulers/slurm_scheduler.py

@@ -472,7 +472,7 @@ def _submit_dryrun(
 
         return AppDryRunInfo(req, repr)
 
-    def _validate(self, app: AppDef, scheduler: str) -> None:
+    def _validate(self, app: AppDef, scheduler: str, cfg: SlurmOpts) -> None:
         # Skip validation step for slurm
         pass