add options to create_scheduler so that the get_runner method is full… (#767)

* add options to create_scheduler so that the get_runner method is fully configurable

I also added a note to each schedulers __init__ method to help with maintainablility

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

* add tests for create_scheduler params

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

---------

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

Author: KPostOffice

torchx/schedulers/aws_batch_scheduler.py

@@ -408,6 +408,7 @@ def __init__(
         log_client: Optional[Any] = None,
         docker_client: Optional["DockerClient"] = None,
     ) -> None:
+        # NOTE: make sure any new init options are supported in create_scheduler(...)
         super().__init__("aws_batch", session_name, docker_client=docker_client)
 
         # pyre-fixme[4]: Attribute annotation cannot be `Any`.
@@ -796,7 +797,18 @@ def _stream_events(
                 yield event["message"] + "\n"
 
 
-def create_scheduler(session_name: str, **kwargs: object) -> AWSBatchScheduler:
+def create_scheduler(
+    session_name: str,
+    # pyre-fixme[2]: Parameter annotation cannot be `Any`.
+    client: Optional[Any] = None,
+    # pyre-fixme[2]: Parameter annotation cannot be `Any`.
+    log_client: Optional[Any] = None,
+    docker_client: Optional["DockerClient"] = None,
+    **kwargs: object,
+) -> AWSBatchScheduler:
     return AWSBatchScheduler(
         session_name=session_name,
+        client=client,
+        log_client=log_client,
+        docker_client=docker_client,
     )

torchx/schedulers/docker_scheduler.py

@@ -168,6 +168,7 @@ class DockerScheduler(DockerWorkspaceMixin, Scheduler[DockerOpts]):
     """
 
     def __init__(self, session_name: str) -> None:
+        # NOTE: make sure any new init options are supported in create_scheduler(...)
         super().__init__("docker", session_name)
 
     def schedule(self, dryrun_info: AppDryRunInfo[DockerJob]) -> str:

torchx/schedulers/gcp_batch_scheduler.py

@@ -142,6 +142,7 @@ def __init__(
         # pyre-fixme[2]: Parameter annotation cannot be `Any`.
         client: Optional[Any] = None,
     ) -> None:
+        # NOTE: make sure any new init options are supported in create_scheduler(...)
         Scheduler.__init__(self, "gcp_batch", session_name)
         # pyre-fixme[4]: Attribute annotation cannot be `Any`.
         self.__client = client
@@ -474,7 +475,13 @@ def _cancel_existing(self, app_id: str) -> None:
         self._client.delete_job(request=request)
 
 
-def create_scheduler(session_name: str, **kwargs: object) -> GCPBatchScheduler:
+def create_scheduler(
+    session_name: str,
+    # pyre-fixme[2]: Parameter annotation cannot be `Any`.
+    client: Optional[Any] = None,
+    **kwargs: object,
+) -> GCPBatchScheduler:
     return GCPBatchScheduler(
         session_name=session_name,
+        client=client,
     )

torchx/schedulers/kubernetes_mcad_scheduler.py

@@ -897,6 +897,7 @@ def __init__(
         client: Optional["ApiClient"] = None,
         docker_client: Optional["DockerClient"] = None,
     ) -> None:
+        # NOTE: make sure any new init options are supported in create_scheduler(...)
         super().__init__("kubernetes_mcad", session_name, docker_client=docker_client)
 
         self._client = client
@@ -1232,9 +1233,16 @@ def list(self) -> List[ListAppResponse]:
         ]
 
 
-def create_scheduler(session_name: str, **kwargs: Any) -> KubernetesMCADScheduler:
+def create_scheduler(
+    session_name: str,
+    client: Optional["ApiClient"] = None,
+    docker_client: Optional["DockerClient"] = None,
+    **kwargs: Any,
+) -> KubernetesMCADScheduler:
     return KubernetesMCADScheduler(
         session_name=session_name,
+        client=client,
+        docker_client=docker_client,
     )
 
 

torchx/schedulers/kubernetes_scheduler.py

@@ -540,6 +540,7 @@ def __init__(
         client: Optional["ApiClient"] = None,
         docker_client: Optional["DockerClient"] = None,
     ) -> None:
+        # NOTE: make sure any new init options are supported in create_scheduler(...)
         super().__init__("kubernetes", session_name, docker_client=docker_client)
 
         self._client = client
@@ -777,9 +778,16 @@ def list(self) -> List[ListAppResponse]:
         ]
 
 
-def create_scheduler(session_name: str, **kwargs: Any) -> KubernetesScheduler:
+def create_scheduler(
+    session_name: str,
+    client: Optional["ApiClient"] = None,
+    docker_client: Optional["DockerClient"] = None,
+    **kwargs: Any,
+) -> KubernetesScheduler:
     return KubernetesScheduler(
         session_name=session_name,
+        client=client,
+        docker_client=docker_client,
     )
 
 

torchx/schedulers/local_scheduler.py

@@ -556,6 +556,7 @@ def __init__(
         cache_size: int = 100,
         extra_paths: Optional[List[str]] = None,
     ) -> None:
+        # NOTE: make sure any new init options are supported in create_scheduler(...)
         super().__init__("local", session_name)
 
         # TODO T72035686 replace dict with a proper LRUCache data structure
@@ -1124,9 +1125,15 @@ def __next__(self) -> str:
                 return line
 
 
-def create_scheduler(session_name: str, **kwargs: Any) -> LocalScheduler:
+def create_scheduler(
+    session_name: str,
+    cache_size: int = 100,
+    extra_paths: Optional[List[str]] = None,
+    **kwargs: Any,
+) -> LocalScheduler:
     return LocalScheduler(
         session_name=session_name,
-        cache_size=kwargs.get("cache_size", 100),
         image_provider_class=CWDImageProvider,
+        cache_size=cache_size,
+        extra_paths=extra_paths,
     )

torchx/schedulers/lsf_scheduler.py

@@ -438,6 +438,7 @@ class LsfScheduler(Scheduler[LsfOpts]):
     """
 
     def __init__(self, session_name: str) -> None:
+        # NOTE: make sure any new init options are supported in create_scheduler(...)
         super().__init__("lsf", session_name)
 
     def _run_opts(self) -> runopts:

torchx/schedulers/ray_scheduler.py

@@ -153,6 +153,7 @@ class RayScheduler(TmpDirWorkspaceMixin, Scheduler[RayOpts]):
         def __init__(
             self, session_name: str, ray_client: Optional[JobSubmissionClient] = None
         ) -> None:
+            # NOTE: make sure any new init options are supported in create_scheduler(...)
             super().__init__("ray", session_name)
 
             # w/o Final None check in _get_ray_client does not work as it pyre assumes mutability
@@ -441,10 +442,12 @@ def list(self) -> List[ListAppResponse]:
             ]
 
 
-def create_scheduler(session_name: str, **kwargs: Any) -> "RayScheduler":
+def create_scheduler(
+    session_name: str, ray_client: Optional[JobSubmissionClient] = None, **kwargs: Any
+) -> "RayScheduler":
     if not has_ray():  # pragma: no cover
         raise ModuleNotFoundError(
             "Ray is not installed in the current Python environment."
         )
 
-    return RayScheduler(session_name=session_name)
+    return RayScheduler(session_name=session_name, ray_client=ray_client)

torchx/schedulers/slurm_scheduler.py

@@ -321,6 +321,7 @@ class SlurmScheduler(DirWorkspaceMixin, Scheduler[SlurmOpts]):
     """
 
     def __init__(self, session_name: str) -> None:
+        # NOTE: make sure any new init options are supported in create_scheduler(...)
         super().__init__("slurm", session_name)
 
     def _run_opts(self) -> runopts:

torchx/schedulers/test/aws_batch_scheduler_test.py

@@ -95,8 +95,16 @@ def paginate(self, *_1: Any, **_2: Any) -> Iterable[Dict[str, Any]]:
 
 class AWSBatchSchedulerTest(unittest.TestCase):
     def test_create_scheduler(self) -> None:
-        scheduler = create_scheduler("foo")
+        client = MagicMock()
+        log_client = MagicMock()
+        docker_client = MagicMock()
+        scheduler = create_scheduler(
+            "foo", client=client, log_client=log_client, docker_client=docker_client
+        )
         self.assertIsInstance(scheduler, AWSBatchScheduler)
+        self.assertEqual(scheduler._client, client)
+        self.assertEqual(scheduler._log_client, log_client)
+        self.assertEqual(scheduler._docker_client, docker_client)
 
     def test_submit_dryrun_with_share_id(self) -> None:
         app = _test_app()