torchx/workspace: workspace_opts + standardize dryrun_push_images and push_images (#619)

Summary:
This makes a few changes intended to make it easier to swap out the workspace implementation for schedulers that support multiple image runtimes. I.e. LSF and Slurm support Docker, Singularity, native, etc

* This adds a new `workspace_opts` so we can keep runopts attached to the workspace
* This standardizes the `dryrun_push_images` and `push_images` that were used in DockerWorkspace so we can swap the workspace without running into any type issues.

Pull Request resolved: #619

Test Plan: Refactor -- no functional changes. Existing tests/CI should suffice

Reviewed By: kurman

Differential Revision: D40446676

Pulled By: d4l3k

fbshipit-source-id: 763eeaddb14282bf8bc189d979de2277fcd3b7d9

Author: d4l3k

CHANGELOG.md

@@ -77,7 +77,7 @@ Milestone: https://github.com/pytorch/torchx/milestones/3
     * Slurm jobs will by default launch in the current working directory to match `local_cwd` and workspace behavior. #372
     * Replicas now have their own log files and can be accessed programmatically. #373
     * Support for `comment`, `mail-user` and `constraint` fields. #391
-    * Workspace support (prototype) - Slurm jobs can now be launched in isolated experiment directories. #416
+    * WorkspaceMixin support (prototype) - Slurm jobs can now be launched in isolated experiment directories. #416
   * Kubernetes
     * Support for running jobs under service accounts. #408
     * Support for specifying instance types. #433

docs/source/workspace.rst

@@ -6,7 +6,7 @@ torchx.workspace
 
 .. currentmodule:: torchx.workspace
 
-.. autoclass:: Workspace
+.. autoclass:: WorkspaceMixin
   :members:
 
 .. autofunction:: walk_workspace
@@ -18,7 +18,7 @@ torchx.workspace.docker_workspace
 .. automodule:: torchx.workspace.docker_workspace
 .. currentmodule:: torchx.workspace.docker_workspace
 
-.. autoclass:: DockerWorkspace
+.. autoclass:: DockerWorkspaceMixin
   :members:
   :private-members: _update_app_images, _push_images
 
@@ -29,7 +29,7 @@ torchx.workspace.dir_workspace
 .. automodule:: torchx.workspace.dir_workspace
 .. currentmodule:: torchx.workspace.dir_workspace
 
-.. autoclass:: DirWorkspace
+.. autoclass:: DirWorkspaceMixin
   :members:
 
 .. fbcode::
@@ -40,6 +40,6 @@ torchx.workspace.dir_workspace
    .. automodule:: torchx.workspace.fb.jetter_workspace
    .. currentmodule:: torchx.workspace.fb.jetter_workspace
 
-   .. autoclass:: JetterWorkspace
+   .. autoclass:: JetterWorkspaceMixin
      :members:
      :show-inheritance:

torchx/runner/api.py

@@ -34,7 +34,7 @@
 from torchx.tracker.api import tracker_config_env_var_name, TRACKER_ENV_VAR_NAME
 
 from torchx.util.types import none_throws
-from torchx.workspace.api import Workspace
+from torchx.workspace.api import WorkspaceMixin
 
 from .config import get_config, get_configs
 
@@ -363,7 +363,7 @@ def dryrun(
         with log_event("dryrun", scheduler, runcfg=json.dumps(cfg) if cfg else None):
             sched = self._scheduler(scheduler)
 
-            if workspace and isinstance(sched, Workspace):
+            if workspace and isinstance(sched, WorkspaceMixin):
                 role = app.roles[0]
                 old_img = role.image
 

torchx/runner/test/api_test.py

@@ -26,7 +26,7 @@
 from torchx.specs.finder import ComponentNotFoundException
 
 from torchx.util.types import none_throws
-from torchx.workspace import Workspace
+from torchx.workspace import WorkspaceMixin
 
 
 GET_SCHEDULER_FACTORIES = "torchx.runner.api.get_scheduler_factories"
@@ -293,9 +293,9 @@ def test_dryrun_setup_trackers_as_env_variable(self, _) -> None:
                 )
 
     def test_dryrun_with_workspace(self, _) -> None:
-        class TestScheduler(Scheduler, Workspace):
+        class TestScheduler(WorkspaceMixin[None], Scheduler):
             def __init__(self, build_new_img: bool):
-                Scheduler.__init__(self, backend="ignored", session_name="ignored")
+                super().__init__(backend="ignored", session_name="ignored")
                 self.build_new_img = build_new_img
 
             def schedule(self, dryrun_info: AppDryRunInfo) -> str:

torchx/runner/test/config_test.py

@@ -62,7 +62,7 @@ def log_iter(
     def list(self) -> List[ListAppResponse]:
         raise NotImplementedError()
 
-    def run_opts(self) -> runopts:
+    def _run_opts(self) -> runopts:
         opts = runopts()
         opts.add(
             "i",

torchx/schedulers/api.py

@@ -22,7 +22,7 @@
     RoleStatus,
     runopts,
 )
-from torchx.workspace.api import Workspace
+from torchx.workspace.api import WorkspaceMixin
 
 
 DAYS_IN_2_WEEKS = 14
@@ -138,7 +138,7 @@ def submit(
         """
         if workspace:
             sched = self
-            assert isinstance(sched, Workspace)
+            assert isinstance(sched, WorkspaceMixin)
             role = app.roles[0]
             sched.build_workspace_and_update_role(role, workspace, cfg)
         dryrun_info = self.submit_dryrun(app, cfg)
@@ -189,6 +189,12 @@ def run_opts(self) -> runopts:
         Returns the run configuration options expected by the scheduler.
         Basically a ``--help`` for the ``run`` API.
         """
+        opts = self._run_opts()
+        if isinstance(self, WorkspaceMixin):
+            opts.update(self.workspace_opts())
+        return opts
+
+    def _run_opts(self) -> runopts:
         return runopts()
 
     @abc.abstractmethod

torchx/schedulers/aws_batch_scheduler.py

@@ -42,9 +42,11 @@
 from typing import (
     Any,
     Callable,
+    cast,
     Dict,
     Iterable,
     List,
+    Mapping,
     Optional,
     Tuple,
     TYPE_CHECKING,
@@ -67,13 +69,14 @@
     AppDef,
     AppState,
     BindMount,
+    CfgVal,
     DeviceMount,
     macros,
     Role,
     runopts,
     VolumeMount,
 )
-from torchx.workspace.docker_workspace import DockerWorkspace
+from torchx.workspace.docker_workspace import DockerWorkspaceMixin
 from typing_extensions import TypedDict
 
 if TYPE_CHECKING:
@@ -246,7 +249,7 @@ class AWSBatchOpts(TypedDict, total=False):
     priority: Optional[int]
 
 
-class AWSBatchScheduler(Scheduler[AWSBatchOpts], DockerWorkspace):
+class AWSBatchScheduler(DockerWorkspaceMixin, Scheduler[AWSBatchOpts]):
     """
     AWSBatchScheduler is a TorchX scheduling interface to AWS Batch.
 
@@ -308,8 +311,7 @@ def __init__(
         log_client: Optional[Any] = None,
         docker_client: Optional["DockerClient"] = None,
     ) -> None:
-        Scheduler.__init__(self, "aws_batch", session_name)
-        DockerWorkspace.__init__(self, docker_client)
+        super().__init__("aws_batch", session_name, docker_client=docker_client)
 
         # pyre-fixme[4]: Attribute annotation cannot be `Any`.
         self.__client = client
@@ -335,7 +337,7 @@ def schedule(self, dryrun_info: AppDryRunInfo[BatchJob]) -> str:
         assert cfg is not None, f"{dryrun_info} missing cfg"
 
         images_to_push = dryrun_info.request.images_to_push
-        self._push_images(images_to_push)
+        self.push_images(images_to_push)
 
         req = dryrun_info.request
         self._client.register_job_definition(**req.job_def)
@@ -370,7 +372,7 @@ def _submit_dryrun(self, app: AppDef, cfg: AWSBatchOpts) -> AppDryRunInfo[BatchJ
         name = make_unique(f"{app.name}{name_suffix}")
 
         # map any local images to the remote image
-        images_to_push = self._update_app_images(app, cfg.get("image_repo"))
+        images_to_push = self.dryrun_push_images(app, cast(Mapping[str, CfgVal], cfg))
 
         nodes = []
 
@@ -450,14 +452,9 @@ def _cancel_existing(self, app_id: str) -> None:
             reason="killed via torchx CLI",
         )
 
-    def run_opts(self) -> runopts:
+    def _run_opts(self) -> runopts:
         opts = runopts()
         opts.add("queue", type_=str, help="queue to schedule job in", required=True)
-        opts.add(
-            "image_repo",
-            type_=str,
-            help="The image repository to use when pushing patched images, must have push access. Ex: example.com/your/container",
-        )
         opts.add(
             "share_id",
             type_=str,

torchx/schedulers/docker_scheduler.py

@@ -38,7 +38,7 @@
     runopts,
     VolumeMount,
 )
-from torchx.workspace.docker_workspace import DockerWorkspace
+from torchx.workspace.docker_workspace import DockerWorkspaceMixin
 from typing_extensions import TypedDict
 
 
@@ -76,7 +76,7 @@ def __repr__(self) -> str:
         return str(self)
 
 
-LABEL_VERSION: str = DockerWorkspace.LABEL_VERSION
+LABEL_VERSION: str = DockerWorkspaceMixin.LABEL_VERSION
 LABEL_APP_ID: str = "torchx.pytorch.org/app-id"
 LABEL_ROLE_NAME: str = "torchx.pytorch.org/role-name"
 LABEL_REPLICA_ID: str = "torchx.pytorch.org/replica-id"
@@ -98,7 +98,7 @@ class DockerOpts(TypedDict, total=False):
     copy_env: Optional[List[str]]
 
 
-class DockerScheduler(Scheduler[DockerOpts], DockerWorkspace):
+class DockerScheduler(DockerWorkspaceMixin, Scheduler[DockerOpts]):
     """
     DockerScheduler is a TorchX scheduling interface to Docker.
 
@@ -143,8 +143,7 @@ class DockerScheduler(Scheduler[DockerOpts], DockerWorkspace):
     """
 
     def __init__(self, session_name: str) -> None:
-        Scheduler.__init__(self, "docker", session_name)
-        DockerWorkspace.__init__(self)
+        super().__init__("docker", session_name)
 
     def _ensure_network(self) -> None:
         import filelock
@@ -346,7 +345,7 @@ def _cancel_existing(self, app_id: str) -> None:
         for container in containers:
             container.stop()
 
-    def run_opts(self) -> runopts:
+    def _run_opts(self) -> runopts:
         opts = runopts()
         opts.add(
             "copy_env",

torchx/schedulers/kubernetes_scheduler.py

@@ -33,7 +33,17 @@
 import warnings
 from dataclasses import dataclass
 from datetime import datetime
-from typing import Any, Dict, Iterable, List, Mapping, Optional, Tuple, TYPE_CHECKING
+from typing import (
+    Any,
+    cast,
+    Dict,
+    Iterable,
+    List,
+    Mapping,
+    Optional,
+    Tuple,
+    TYPE_CHECKING,
+)
 
 import torchx
 import yaml
@@ -51,6 +61,7 @@
     AppDef,
     AppState,
     BindMount,
+    CfgVal,
     DeviceMount,
     macros,
     ReplicaState,
@@ -61,7 +72,7 @@
     runopts,
     VolumeMount,
 )
-from torchx.workspace.docker_workspace import DockerWorkspace
+from torchx.workspace.docker_workspace import DockerWorkspaceMixin
 from typing_extensions import TypedDict
 
 
@@ -441,7 +452,7 @@ class KubernetesOpts(TypedDict, total=False):
     priority_class: Optional[str]
 
 
-class KubernetesScheduler(Scheduler[KubernetesOpts], DockerWorkspace):
+class KubernetesScheduler(DockerWorkspaceMixin, Scheduler[KubernetesOpts]):
     """
     KubernetesScheduler is a TorchX scheduling interface to Kubernetes.
 
@@ -535,8 +546,7 @@ def __init__(
         client: Optional["ApiClient"] = None,
         docker_client: Optional["DockerClient"] = None,
     ) -> None:
-        Scheduler.__init__(self, "kubernetes", session_name)
-        DockerWorkspace.__init__(self, docker_client)
+        super().__init__("kubernetes", session_name, docker_client=docker_client)
 
         self._client = client
 
@@ -575,7 +585,7 @@ def schedule(self, dryrun_info: AppDryRunInfo[KubernetesJob]) -> str:
         namespace = cfg.get("namespace") or "default"
 
         images_to_push = dryrun_info.request.images_to_push
-        self._push_images(images_to_push)
+        self.push_images(images_to_push)
 
         resource = dryrun_info.request.resource
         try:
@@ -605,7 +615,7 @@ def _submit_dryrun(
             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.get("image_repo"))
+        images_to_push = self.dryrun_push_images(app, cast(Mapping[str, CfgVal], cfg))
 
         service_account = cfg.get("service_account")
         assert service_account is None or isinstance(
@@ -642,7 +652,7 @@ def _cancel_existing(self, app_id: str) -> None:
             name=name,
         )
 
-    def run_opts(self) -> runopts:
+    def _run_opts(self) -> runopts:
         opts = runopts()
         opts.add(
             "namespace",
@@ -656,11 +666,6 @@ def run_opts(self) -> runopts:
             help="Volcano queue to schedule job in",
             required=True,
         )
-        opts.add(
-            "image_repo",
-            type_=str,
-            help="The image repository to use when pushing patched images, must have push access. Ex: example.com/your/container",
-        )
         opts.add(
             "service_account",
             type_=str,

torchx/schedulers/local_scheduler.py

@@ -573,7 +573,7 @@ def __init__(
         self._base_log_dir: Optional[str] = None
         self._created_tmp_log_dir: bool = False
 
-    def run_opts(self) -> runopts:
+    def _run_opts(self) -> runopts:
         opts = runopts()
         opts.add(
             "log_dir",