Remove redundant SchedulerBackend in favor of str (#436)

Summary:
Pull Request resolved: #436

Removed SchedulerBackend and replaced with str

Reviewed By: d4l3k

Differential Revision: D35185234

fbshipit-source-id: 8f2e0abbda6eaecc9e8aecd32e3775f2f7ca83be

Author: paksha

torchx/components/component_test_base.py

@@ -27,7 +27,6 @@
     AppHandle,
     AppState,
     CfgVal,
-    SchedulerBackend,
 )
 from torchx.specs.api import _create_args_parser
 from torchx.specs.finder import get_component
@@ -38,7 +37,7 @@ class ComponentUtils:
     def run_appdef_on_scheduler(
         cls,
         app_def: AppDef,
-        scheduler: SchedulerBackend,
+        scheduler: str,
         cfg: Mapping[str, CfgVal],
         dryrun: bool = False,
     ) -> Union[AppHandle, AppDryRunInfo]:

torchx/components/integration_tests/component_provider.py

@@ -11,15 +11,15 @@
 import torchx.components.dist as dist_components
 import torchx.components.serve as serve_components
 import torchx.components.utils as utils_components
-from torchx.specs import AppDef, SchedulerBackend
+from torchx.specs import AppDef
 
 
 class ComponentProvider(ABC):
     """
     Abstract class that represents generic component provider.
     """
 
-    def __init__(self, scheduler: SchedulerBackend, image: str) -> None:
+    def __init__(self, scheduler: str, image: str) -> None:
         self._scheduler = scheduler
         self._image = image
 
@@ -75,7 +75,7 @@ def get_app_def(self) -> AppDef:
 
 
 class TouchComponentProvider(ComponentProvider):
-    def __init__(self, image: str, scheduler: SchedulerBackend) -> None:
+    def __init__(self, image: str, scheduler: str) -> None:
         super(TouchComponentProvider, self).__init__(image, scheduler)
         self._file_path = "<None>"
 
@@ -98,7 +98,7 @@ def tearDown(self) -> None:
 
 
 class CopyComponentProvider(ComponentProvider):
-    def __init__(self, image: str, scheduler: SchedulerBackend) -> None:
+    def __init__(self, image: str, scheduler: str) -> None:
         super(CopyComponentProvider, self).__init__(image, scheduler)
         self._src_path = "<None>"
         self._dst_path = "<None>"

torchx/components/integration_tests/integ_tests.py

@@ -16,12 +16,12 @@
 from pyre_extensions import none_throws
 from torchx.components.component_test_base import ComponentUtils
 from torchx.components.integration_tests.component_provider import ComponentProvider
-from torchx.specs import AppDef, CfgVal, SchedulerBackend
+from torchx.specs import AppDef, CfgVal
 
 
 @dataclass
 class SchedulerInfo:
-    name: SchedulerBackend
+    name: str
     image: str
     cfg: Mapping[str, CfgVal] = field(default_factory=dict)
 

torchx/runner/api.py

@@ -22,7 +22,6 @@
     AppHandle,
     AppStatus,
     CfgVal,
-    SchedulerBackend,
     UnknownAppException,
     from_function,
     make_app_handle,
@@ -50,7 +49,7 @@ class Runner:
     def __init__(
         self,
         name: str,
-        schedulers: Dict[SchedulerBackend, Scheduler],
+        schedulers: Dict[str, Scheduler],
         component_defaults: Optional[Dict[str, Dict[str, str]]] = None,
     ) -> None:
         """
@@ -102,7 +101,7 @@ def run_component(
         self,
         component: str,
         component_args: List[str],
-        scheduler: SchedulerBackend,
+        scheduler: str,
         cfg: Optional[Mapping[str, CfgVal]] = None,
         workspace: Optional[str] = None,
     ) -> AppHandle:
@@ -151,7 +150,7 @@ def dryrun_component(
         self,
         component: str,
         component_args: List[str],
-        scheduler: SchedulerBackend,
+        scheduler: str,
         cfg: Optional[Mapping[str, CfgVal]] = None,
         workspace: Optional[str] = None,
     ) -> AppDryRunInfo:
@@ -170,7 +169,7 @@ def dryrun_component(
     def run(
         self,
         app: AppDef,
-        scheduler: SchedulerBackend,
+        scheduler: str,
         cfg: Optional[Mapping[str, CfgVal]] = None,
         workspace: Optional[str] = None,
     ) -> AppHandle:
@@ -237,7 +236,7 @@ def name(self) -> str:
     def dryrun(
         self,
         app: AppDef,
-        scheduler: SchedulerBackend,
+        scheduler: str,
         cfg: Optional[Mapping[str, CfgVal]] = None,
         workspace: Optional[str] = None,
     ) -> AppDryRunInfo:
@@ -317,7 +316,7 @@ def run_opts(self) -> Dict[str, runopts]:
             for scheduler_backend, scheduler in self._schedulers.items()
         }
 
-    def scheduler_backends(self) -> List[SchedulerBackend]:
+    def scheduler_backends(self) -> List[str]:
         """
         Returns a list of all supported scheduler backends.
         """
@@ -535,7 +534,7 @@ def log_lines(
             )
             return log_iter
 
-    def _scheduler(self, scheduler: SchedulerBackend) -> Scheduler:
+    def _scheduler(self, scheduler: str) -> Scheduler:
         sched = self._schedulers.get(scheduler)
         if not sched:
             raise KeyError(

torchx/schedulers/__init__.py

@@ -13,7 +13,6 @@
 import torchx.schedulers.local_scheduler as local_scheduler
 import torchx.schedulers.slurm_scheduler as slurm_scheduler
 from torchx.schedulers.api import Scheduler
-from torchx.specs.api import SchedulerBackend
 from torchx.util.entrypoints import load_group
 from typing_extensions import Protocol
 
@@ -73,7 +72,7 @@ def get_default_scheduler_name() -> str:
 
 def get_schedulers(
     session_name: str, **scheduler_params: object
-) -> Dict[SchedulerBackend, Scheduler]:
+) -> Dict[str, Scheduler]:
     """
     get_schedulers returns all available schedulers.
     """

torchx/schedulers/api.py

@@ -21,7 +21,6 @@
     CfgVal,
     Role,
     RoleStatus,
-    SchedulerBackend,
     runopts,
 )
 from torchx.workspace.api import Workspace
@@ -70,7 +69,7 @@ class Scheduler(abc.ABC):
     ``@abc.abstractmethod``.
     """
 
-    def __init__(self, backend: SchedulerBackend, session_name: str) -> None:
+    def __init__(self, backend: str, session_name: str) -> None:
         self.backend = backend
         self.session_name = session_name
 
@@ -284,7 +283,7 @@ def log_iter(
             f"{self.__class__.__qualname__} does not support application log iteration"
         )
 
-    def _validate(self, app: AppDef, scheduler: SchedulerBackend) -> None:
+    def _validate(self, app: AppDef, scheduler: str) -> None:
         """
         Validates whether application is consistent with the scheduler.
 

torchx/schedulers/aws_batch_scheduler.py

@@ -64,7 +64,6 @@
     AppDef,
     AppState,
     Role,
-    SchedulerBackend,
     macros,
     runopts,
     CfgVal,
@@ -363,7 +362,7 @@ def _submit_dryrun(
         info._cfg = cfg
         return info
 
-    def _validate(self, app: AppDef, scheduler: SchedulerBackend) -> None:
+    def _validate(self, app: AppDef, scheduler: str) -> None:
         # Skip validation step
         pass
 

torchx/schedulers/docker_scheduler.py

@@ -30,7 +30,6 @@
     ReplicaStatus,
     Role,
     RoleStatus,
-    SchedulerBackend,
     is_terminal,
     macros,
     runopts,
@@ -295,7 +294,7 @@ def _submit_dryrun(
         info._cfg = cfg
         return info
 
-    def _validate(self, app: AppDef, scheduler: SchedulerBackend) -> None:
+    def _validate(self, app: AppDef, scheduler: str) -> None:
         # Skip validation step
         pass
 

torchx/schedulers/kubernetes_scheduler.py

@@ -61,7 +61,6 @@
     RetryPolicy,
     Role,
     RoleStatus,
-    SchedulerBackend,
     macros,
     runopts,
     BindMount,
@@ -572,7 +571,7 @@ def _submit_dryrun(
         info._cfg = cfg
         return info
 
-    def _validate(self, app: AppDef, scheduler: SchedulerBackend) -> None:
+    def _validate(self, app: AppDef, scheduler: str) -> None:
         # Skip validation step
         pass
 

torchx/schedulers/local_scheduler.py

@@ -51,7 +51,6 @@
     AppState,
     CfgVal,
     Role,
-    SchedulerBackend,
     is_terminal,
     macros,
     runopts,
@@ -602,7 +601,7 @@ def run_opts(self) -> runopts:
         )
         return opts
 
-    def _validate(self, app: AppDef, scheduler: SchedulerBackend) -> None:
+    def _validate(self, app: AppDef, scheduler: str) -> None:
         # Skip validation step for local application
         pass
 

torchx/schedulers/ray_scheduler.py

@@ -28,7 +28,6 @@
 from torchx.specs import (
     AppDef,
     CfgVal,
-    SchedulerBackend,
     macros,
     runopts,
     Role,
@@ -274,7 +273,7 @@ def _submit_dryrun(
 
             return AppDryRunInfo(job, repr)
 
-        def _validate(self, app: AppDef, scheduler: SchedulerBackend) -> None:
+        def _validate(self, app: AppDef, scheduler: str) -> None:
             if scheduler != "ray":
                 raise ValueError(
                     f"An unknown scheduler backend '{scheduler}' has been passed to the Ray scheduler."

torchx/schedulers/slurm_scheduler.py

@@ -31,7 +31,6 @@
     ReplicaStatus,
     Role,
     RoleStatus,
-    SchedulerBackend,
     macros,
     runopts,
 )
@@ -383,7 +382,7 @@ def _submit_dryrun(
 
         return AppDryRunInfo(req, repr)
 
-    def _validate(self, app: AppDef, scheduler: SchedulerBackend) -> None:
+    def _validate(self, app: AppDef, scheduler: str) -> None:
         # Skip validation step for slurm
         pass
 

torchx/specs/__init__.py

@@ -37,7 +37,6 @@
     RetryPolicy,
     Role,
     RoleStatus,
-    SchedulerBackend,
     UnknownAppException,
     UnknownSchedulerException,
     from_function,

torchx/specs/api.py

@@ -38,9 +38,6 @@
 )
 
 
-SchedulerBackend = str
-
-
 # ========================================
 # ==== Distributed AppDef API =======
 # ========================================
@@ -87,7 +84,7 @@ def copy(original: "Resource", **capabilities: Any) -> "Resource":
 NULL_RESOURCE: Resource = Resource(cpu=-1, gpu=-1, memMB=-1)
 
 # used as "*" scheduler backend
-ALL: SchedulerBackend = "all"
+ALL: str = "all"
 
 # sentinel value used to represent missing string attributes, such as image or entrypoint
 MISSING: str = "<MISSING>"
@@ -299,7 +296,7 @@ class Role:
 
     def pre_proc(
         self,
-        scheduler: SchedulerBackend,
+        scheduler: str,
         # pyre-fixme[24]: AppDryRunInfo was designed to work with Any request object
         dryrun_info: "AppDryRunInfo",
         # pyre-fixme[24]: AppDryRunInfo was designed to work with Any request object
@@ -501,7 +498,7 @@ def __init__(self, request: T, fmt: Callable[[T], str]) -> None:
         # unless there is a good reason to
         self._app: Optional[AppDef] = None
         self._cfg: Mapping[str, CfgVal] = {}
-        self._scheduler: Optional[SchedulerBackend] = None
+        self._scheduler: Optional[str] = None
 
     def __repr__(self) -> str:
         return self._fmt(self.request)
@@ -712,7 +709,7 @@ def __init__(self, app_handle: str) -> None:
 
 
 class UnknownSchedulerException(Exception):
-    def __init__(self, scheduler_backend: SchedulerBackend) -> None:
+    def __init__(self, scheduler_backend: str) -> None:
         super().__init__(
             f"Scheduler backend: {scheduler_backend} does not exist."
             f" Use session.scheduler_backends() to see all supported schedulers"
@@ -737,13 +734,11 @@ def __init__(self, app_handle: "AppHandle") -> None:
         )
 
 
-def make_app_handle(
-    scheduler_backend: SchedulerBackend, session_name: str, app_id: str
-) -> str:
+def make_app_handle(scheduler_backend: str, session_name: str, app_id: str) -> str:
     return f"{scheduler_backend}://{session_name}/{app_id}"
 
 
-def parse_app_handle(app_handle: AppHandle) -> Tuple[SchedulerBackend, str, str]:
+def parse_app_handle(app_handle: AppHandle) -> Tuple[str, str, str]:
     """
     parses the app handle into ```(scheduler_backend, session_name, and app_id)```
     """