Use generic TypedDict for runcfg (#482)

Summary:
Pull Request resolved: #482

Pull Request resolved: #481

Use generic TypedDict for runcfg rather than Mapping. This eliminates the need for casts and type asserts to use the values.

Reviewed By: d4l3k

Differential Revision: D36072685

fbshipit-source-id: 3f8d57c385210839fd21c4e2690339055a78a29a

Author: hadarohana

torchx/runner/api.py

@@ -49,6 +49,7 @@ class Runner:
     def __init__(
         self,
         name: str,
+        # pyre-fixme: Scheduler opts
         schedulers: Dict[str, Scheduler],
         component_defaults: Optional[Dict[str, Dict[str, str]]] = None,
     ) -> None:
@@ -539,6 +540,7 @@ def log_lines(
             )
             return log_iter
 
+    # pyre-fixme: Scheduler opts
     def _scheduler(self, scheduler: str) -> Scheduler:
         sched = self._schedulers.get(scheduler)
         if not sched:
@@ -548,7 +550,10 @@ def _scheduler(self, scheduler: str) -> Scheduler:
         return sched
 
     def _scheduler_app_id(
-        self, app_handle: AppHandle, check_session: bool = True
+        self,
+        app_handle: AppHandle,
+        check_session: bool = True
+        # pyre-fixme: Scheduler opts
     ) -> Tuple[Scheduler, str, str]:
         """
         Returns the scheduler and app_id from the app_handle.

torchx/schedulers/__init__.py

@@ -18,6 +18,7 @@
 
 
 class SchedulerFactory(Protocol):
+    # pyre-fixme: Scheduler opts
     def __call__(self, session_name: str, **kwargs: object) -> Scheduler:
         ...
 
@@ -71,7 +72,9 @@ def get_default_scheduler_name() -> str:
 
 
 def get_schedulers(
-    session_name: str, **scheduler_params: object
+    session_name: str,
+    **scheduler_params: object
+    # pyre-fixme: Scheduler opts
 ) -> Dict[str, Scheduler]:
     """
     get_schedulers returns all available schedulers.

torchx/schedulers/api.py

@@ -10,13 +10,12 @@
 from dataclasses import dataclass, field
 from datetime import datetime
 from enum import Enum
-from typing import Iterable, List, Mapping, Optional
+from typing import Generic, Iterable, List, Optional, TypeVar
 
 from torchx.specs import (
     AppDef,
     AppDryRunInfo,
     AppState,
-    CfgVal,
     NONE,
     NULL_RESOURCE,
     Role,
@@ -62,7 +61,10 @@ class DescribeAppResponse:
     roles: List[Role] = field(default_factory=list)
 
 
-class Scheduler(abc.ABC):
+T = TypeVar("T")
+
+
+class Scheduler(abc.ABC, Generic[T]):
     """
     An interface abstracting functionalities of a scheduler.
     Implementors need only implement those methods annotated with
@@ -93,7 +95,7 @@ def close(self) -> None:
     def submit(
         self,
         app: AppDef,
-        cfg: Mapping[str, CfgVal],
+        cfg: T,
         workspace: Optional[str] = None,
     ) -> str:
         """
@@ -129,7 +131,7 @@ def schedule(self, dryrun_info: AppDryRunInfo) -> str:
 
         raise NotImplementedError()
 
-    def submit_dryrun(self, app: AppDef, cfg: Mapping[str, CfgVal]) -> AppDryRunInfo:
+    def submit_dryrun(self, app: AppDef, cfg: T) -> AppDryRunInfo:
         """
         Rather than submitting the request to run the app, returns the
         request object that would have been submitted to the underlying
@@ -138,7 +140,9 @@ def submit_dryrun(self, app: AppDef, cfg: Mapping[str, CfgVal]) -> AppDryRunInfo
         to the scheduler implementation's documentation regarding
         the actual return type.
         """
+        # pyre-fixme: Generic cfg type passed to resolve
         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)
@@ -147,7 +151,7 @@ def submit_dryrun(self, app: AppDef, cfg: Mapping[str, CfgVal]) -> AppDryRunInfo
         return dryrun_info
 
     @abc.abstractmethod
-    def _submit_dryrun(self, app: AppDef, cfg: Mapping[str, CfgVal]) -> AppDryRunInfo:
+    def _submit_dryrun(self, app: AppDef, cfg: T) -> AppDryRunInfo:
         raise NotImplementedError()
 
     def run_opts(self) -> runopts:

torchx/schedulers/aws_batch_scheduler.py

@@ -43,7 +43,6 @@
     Callable,
     Dict,
     Iterable,
-    Mapping,
     Optional,
     Tuple,
     TYPE_CHECKING,
@@ -65,14 +64,14 @@
     AppDef,
     AppState,
     BindMount,
-    CfgVal,
     DeviceMount,
     macros,
     Role,
     runopts,
     VolumeMount,
 )
 from torchx.workspace.docker_workspace import DockerWorkspace
+from typing_extensions import TypedDict
 
 if TYPE_CHECKING:
     from docker import DockerClient
@@ -224,7 +223,12 @@ def _local_session() -> "boto3.session.Session":
     return boto3.session.Session()
 
 
-class AWSBatchScheduler(Scheduler, DockerWorkspace):
+class AWSBatchOpts(TypedDict, total=False):
+    queue: str
+    image_repo: Optional[str]
+
+
+class AWSBatchScheduler(Scheduler[AWSBatchOpts], DockerWorkspace):
     """
     AWSBatchScheduler is a TorchX scheduling interface to AWS Batch.
 
@@ -326,16 +330,14 @@ def schedule(self, dryrun_info: AppDryRunInfo[BatchJob]) -> str:
 
         return f"{req.queue}:{req.name}"
 
-    def _submit_dryrun(
-        self, app: AppDef, cfg: Mapping[str, CfgVal]
-    ) -> AppDryRunInfo[BatchJob]:
+    def _submit_dryrun(self, app: AppDef, cfg: AWSBatchOpts) -> AppDryRunInfo[BatchJob]:
         queue = cfg.get("queue")
         if not isinstance(queue, str):
             raise TypeError(f"config value 'queue' must be a string, got {queue}")
         name = make_unique(app.name)
 
         # map any local images to the remote image
-        images_to_push = self._update_app_images(app, cfg)
+        images_to_push = self._update_app_images(app, cfg.get("image_repo"))
 
         nodes = []
 
@@ -390,6 +392,7 @@ def _submit_dryrun(
         )
         info = AppDryRunInfo(req, repr)
         info._app = app
+        # pyre-fixme: AppDryRunInfo
         info._cfg = cfg
         return info
 

torchx/schedulers/docker_scheduler.py

@@ -10,7 +10,7 @@
 import tempfile
 from dataclasses import dataclass
 from datetime import datetime
-from typing import Any, Dict, Iterable, List, Mapping, Optional, TYPE_CHECKING, Union
+from typing import Any, Dict, Iterable, List, Optional, TYPE_CHECKING, Union
 
 import torchx
 import yaml
@@ -28,7 +28,6 @@
     AppDef,
     AppState,
     BindMount,
-    CfgVal,
     DeviceMount,
     is_terminal,
     macros,
@@ -39,6 +38,7 @@
     VolumeMount,
 )
 from torchx.workspace.docker_workspace import DockerWorkspace
+from typing_extensions import TypedDict
 
 
 if TYPE_CHECKING:
@@ -93,7 +93,11 @@ def has_docker() -> bool:
         return False
 
 
-class DockerScheduler(Scheduler, DockerWorkspace):
+class DockerOpts(TypedDict, total=False):
+    copy_env: Optional[List[str]]
+
+
+class DockerScheduler(Scheduler[DockerOpts], DockerWorkspace):
     """
     DockerScheduler is a TorchX scheduling interface to Docker.
 
@@ -187,9 +191,7 @@ def schedule(self, dryrun_info: AppDryRunInfo[DockerJob]) -> str:
 
         return req.app_id
 
-    def _submit_dryrun(
-        self, app: AppDef, cfg: Mapping[str, CfgVal]
-    ) -> AppDryRunInfo[DockerJob]:
+    def _submit_dryrun(self, app: AppDef, cfg: DockerOpts) -> AppDryRunInfo[DockerJob]:
         from docker.types import DeviceRequest, Mount
 
         default_env = {}
@@ -301,6 +303,7 @@ def _submit_dryrun(
 
         info = AppDryRunInfo(req, repr)
         info._app = app
+        # pyre-fixme: AppDryRunInfo
         info._cfg = cfg
         return info
 

torchx/schedulers/kubernetes_scheduler.py

@@ -56,7 +56,6 @@
     AppDef,
     AppState,
     BindMount,
-    CfgVal,
     DeviceMount,
     macros,
     ReplicaState,
@@ -68,6 +67,7 @@
     VolumeMount,
 )
 from torchx.workspace.docker_workspace import DockerWorkspace
+from typing_extensions import TypedDict
 
 
 if TYPE_CHECKING:
@@ -434,7 +434,15 @@ def __repr__(self) -> str:
         return str(self)
 
 
-class KubernetesScheduler(Scheduler, DockerWorkspace):
+class KubernetesOpts(TypedDict, total=False):
+    namespace: Optional[str]
+    queue: str
+    image_repo: Optional[str]
+    service_account: Optional[str]
+    priority_class: Optional[str]
+
+
+class KubernetesScheduler(Scheduler[KubernetesOpts], DockerWorkspace):
     """
     KubernetesScheduler is a TorchX scheduling interface to Kubernetes.
 
@@ -590,14 +598,14 @@ def schedule(self, dryrun_info: AppDryRunInfo[KubernetesJob]) -> str:
         return f'{namespace}:{resp["metadata"]["name"]}'
 
     def _submit_dryrun(
-        self, app: AppDef, cfg: Mapping[str, CfgVal]
+        self, app: AppDef, cfg: KubernetesOpts
     ) -> AppDryRunInfo[KubernetesJob]:
         queue = cfg.get("queue")
         if not isinstance(queue, str):
             raise TypeError(f"config value 'queue' must be a string, got {queue}")
 
         # map any local images to the remote image
-        images_to_push = self._update_app_images(app, cfg)
+        images_to_push = self._update_app_images(app, cfg.get("image_repo"))
 
         service_account = cfg.get("service_account")
         assert service_account is None or isinstance(
@@ -616,6 +624,7 @@ def _submit_dryrun(
         )
         info = AppDryRunInfo(req, repr)
         info._app = app
+        # pyre-fixme: AppDryRunInfo
         info._cfg = cfg
         return info