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