(torchx/workspace) Support multi-project/directory workspace

Summary:
Usage:

In your `.torchxconfig` file specify multi-project workspace as follows:

```
[cli:run]
workspace = 
  /home/$USER/github/myproj: 
  /home/$USER/github/verl: verl
  /home/$USER/.torchxconfig: verl/.torchxconfig
```

Then the workspace is built as:

1. `/home/$USER/github/myproj/**` -> `$REMOTE_ROOT/**`
1. `/home/$USER/github/verl/**` -> `$REMOTE_ROOT/verl/**`
1. `/home/$USER/.torchxconfig` -> `$REMOTE_ROOT/verl/.torchxconfig`

Notes:

1. The mappings are basically like BUCK `python_library`'s `resoruce` attribute. Mapping to an empty string, or simply leaving the target blank (as we did for `myproj` in the example above) copies the contents of the src dir directly into the remote workspace root

1. Files can be specified as src. The target of a file can be a directory, and if so, the file will be copied into the directory. For example:

```
[cli:run]
workspace = 
  /home/$USER/github/verl: verl
  /home/$USER/.torchxconfig: verl
```

copies `.torchxconfig` into `$REMOTE_ROOT/verl/.torchxconfig`. However if we switching the order of the projects won't work:

```
[cli:run]
workspace = 
  /home/$USER/.torchxconfig: verl
  /home/$USER/github/verl: verl
```

This would copy `.torchxconfig` to the file called `verl` in  `$REMOTE_ROOT/verl`, then `github/verl` would fail to copy.

Differential Revision: D82169554

Author: kiukchung

torchx/cli/cmd_run.py

@@ -36,6 +36,7 @@
 )
 from torchx.util.log_tee_helpers import tee_logs
 from torchx.util.types import none_throws
+from torchx.workspace import Workspace
 
 
 MISSING_COMPONENT_ERROR_MSG = (
@@ -289,12 +290,13 @@ def _run_inner(self, runner: Runner, args: TorchXRunArgs) -> None:
             else args.component_args
         )
         try:
+            workspace = Workspace.from_str(args.workspace) if args.workspace else None
             if args.dryrun:
                 dryrun_info = runner.dryrun_component(
                     args.component_name,
                     component_args,
                     args.scheduler,
-                    workspace=args.workspace,
+                    workspace=workspace,
                     cfg=args.scheduler_cfg,
                     parent_run_id=args.parent_run_id,
                 )

torchx/runner/api.py

@@ -54,7 +54,7 @@
 from torchx.util.session import get_session_id_or_create_new, TORCHX_INTERNAL_SESSION_ID
 
 from torchx.util.types import none_throws
-from torchx.workspace.api import WorkspaceMixin
+from torchx.workspace.api import Workspace, WorkspaceMixin
 
 if TYPE_CHECKING:
     from typing_extensions import Self
@@ -171,7 +171,7 @@ def run_component(
         component_args: Union[list[str], dict[str, Any]],
         scheduler: str,
         cfg: Optional[Mapping[str, CfgVal]] = None,
-        workspace: Optional[str] = None,
+        workspace: Optional[Union[Workspace, str]] = None,
         parent_run_id: Optional[str] = None,
     ) -> AppHandle:
         """
@@ -206,7 +206,7 @@ def run_component(
             ComponentNotFoundException: if the ``component_path`` is failed to resolve.
         """
 
-        with log_event("run_component", workspace=workspace) as ctx:
+        with log_event("run_component") as ctx:
             dryrun_info = self.dryrun_component(
                 component,
                 component_args,
@@ -217,7 +217,8 @@ def run_component(
             )
             handle = self.schedule(dryrun_info)
             app = none_throws(dryrun_info._app)
-            ctx._torchx_event.workspace = workspace
+
+            ctx._torchx_event.workspace = str(workspace)
             ctx._torchx_event.scheduler = none_throws(dryrun_info._scheduler)
             ctx._torchx_event.app_image = app.roles[0].image
             ctx._torchx_event.app_id = parse_app_handle(handle)[2]
@@ -230,7 +231,7 @@ def dryrun_component(
         component_args: Union[list[str], dict[str, Any]],
         scheduler: str,
         cfg: Optional[Mapping[str, CfgVal]] = None,
-        workspace: Optional[str] = None,
+        workspace: Optional[Union[Workspace, str]] = None,
         parent_run_id: Optional[str] = None,
     ) -> AppDryRunInfo:
         """
@@ -259,7 +260,7 @@ def run(
         app: AppDef,
         scheduler: str,
         cfg: Optional[Mapping[str, CfgVal]] = None,
-        workspace: Optional[str] = None,
+        workspace: Optional[Union[Workspace, str]] = None,
         parent_run_id: Optional[str] = None,
     ) -> AppHandle:
         """
@@ -272,9 +273,7 @@ def run(
             An application handle that is used to call other action APIs on the app.
         """
 
-        with log_event(
-            api="run", runcfg=json.dumps(cfg) if cfg else None, workspace=workspace
-        ) as ctx:
+        with log_event(api="run") as ctx:
             dryrun_info = self.dryrun(
                 app,
                 scheduler,
@@ -283,10 +282,15 @@ def run(
                 parent_run_id=parent_run_id,
             )
             handle = self.schedule(dryrun_info)
-            ctx._torchx_event.scheduler = none_throws(dryrun_info._scheduler)
-            ctx._torchx_event.app_image = none_throws(dryrun_info._app).roles[0].image
-            ctx._torchx_event.app_id = parse_app_handle(handle)[2]
-            ctx._torchx_event.app_metadata = app.metadata
+
+            event = ctx._torchx_event
+            event.scheduler = scheduler
+            event.runcfg = json.dumps(cfg) if cfg else None
+            event.workspace = str(workspace)
+            event.app_id = parse_app_handle(handle)[2]
+            event.app_image = none_throws(dryrun_info._app).roles[0].image
+            event.app_metadata = app.metadata
+
             return handle
 
     def schedule(self, dryrun_info: AppDryRunInfo) -> AppHandle:
@@ -320,21 +324,22 @@ def schedule(self, dryrun_info: AppDryRunInfo) -> AppHandle:
 
         """
         scheduler = none_throws(dryrun_info._scheduler)
-        app_image = none_throws(dryrun_info._app).roles[0].image
         cfg = dryrun_info._cfg
-        with log_event(
-            "schedule",
-            scheduler,
-            app_image=app_image,
-            runcfg=json.dumps(cfg) if cfg else None,
-        ) as ctx:
+        with log_event("schedule") as ctx:
             sched = self._scheduler(scheduler)
             app_id = sched.schedule(dryrun_info)
             app_handle = make_app_handle(scheduler, self._name, app_id)
+
             app = none_throws(dryrun_info._app)
             self._apps[app_handle] = app
-            _, _, app_id = parse_app_handle(app_handle)
-            ctx._torchx_event.app_id = app_id
+
+            event = ctx._torchx_event
+            event.scheduler = scheduler
+            event.runcfg = json.dumps(cfg) if cfg else None
+            event.app_id = app_id
+            event.app_image = none_throws(dryrun_info._app).roles[0].image
+            event.app_metadata = app.metadata
+
             return app_handle
 
     def name(self) -> str:
@@ -345,7 +350,7 @@ def dryrun(
         app: AppDef,
         scheduler: str,
         cfg: Optional[Mapping[str, CfgVal]] = None,
-        workspace: Optional[str] = None,
+        workspace: Optional[Union[Workspace, str]] = None,
         parent_run_id: Optional[str] = None,
     ) -> AppDryRunInfo:
         """
@@ -414,7 +419,7 @@ def dryrun(
             "dryrun",
             scheduler,
             runcfg=json.dumps(cfg) if cfg else None,
-            workspace=workspace,
+            workspace=str(workspace),
         ):
             sched = self._scheduler(scheduler)
             resolved_cfg = sched.run_opts().resolve(cfg)
@@ -429,7 +434,7 @@ def dryrun(
                 logger.info(
                     'To disable workspaces pass: --workspace="" from CLI or workspace=None programmatically.'
                 )
-                sched.build_workspace_and_update_role(role, workspace, resolved_cfg)
+                sched.build_workspace_and_update_role2(role, workspace, resolved_cfg)
 
                 if old_img != role.image:
                     logger.info(

torchx/runner/events/api.py

@@ -29,7 +29,7 @@ class TorchxEvent:
         scheduler: Scheduler that is used to execute request
         api: Api name
         app_id: Unique id that is set by the underlying scheduler
-        image: Image/container bundle that is used to execute request.
+        app_image: Image/container bundle that is used to execute request.
         app_metadata: metadata to the app (treatment of metadata is scheduler dependent)
         runcfg: Run config that was used to schedule app.
         source: Type of source the event is generated.

torchx/runner/test/config_test.py

@@ -27,6 +27,7 @@
 from torchx.schedulers.api import DescribeAppResponse, ListAppResponse, Stream
 from torchx.specs import AppDef, AppDryRunInfo, CfgVal, runopts
 from torchx.test.fixtures import TestWithTmpDir
+from torchx.workspace import Workspace
 
 
 class TestScheduler(Scheduler):
@@ -506,3 +507,29 @@ def test_dump_and_load_all_registered_schedulers(self) -> None:
                     opt_name in cfg,
                     f"missing {opt_name} in {sched} run opts with cfg {cfg}",
                 )
+
+    def test_get_workspace_config(self) -> None:
+        configdir = self.tmpdir
+        self.write(
+            str(configdir / ".torchxconfig"),
+            """#
+[cli:run]
+workspace =
+    /home/foo/third-party/verl: verl
+    /home/foo/bar/scripts/.torchxconfig: verl/.torchxconfig
+    /home/foo/baz:
+""",
+        )
+
+        workspace_config = get_config(
+            prefix="cli", name="run", key="workspace", dirs=[str(configdir)]
+        )
+        workspace = Workspace.from_str(workspace_config)
+        self.assertDictEqual(
+            {
+                "/home/foo/third-party/verl": "verl",
+                "/home/foo/bar/scripts/.torchxconfig": "verl/.torchxconfig",
+                "/home/foo/baz": None,
+            },
+            workspace.projects,
+        )

torchx/schedulers/api.py

@@ -12,7 +12,7 @@
 from dataclasses import dataclass, field
 from datetime import datetime
 from enum import Enum
-from typing import Generic, Iterable, List, Optional, TypeVar
+from typing import Generic, Iterable, List, Optional, TypeVar, Union
 
 from torchx.specs import (
     AppDef,
@@ -23,7 +23,7 @@
     RoleStatus,
     runopts,
 )
-from torchx.workspace.api import WorkspaceMixin
+from torchx.workspace.api import Workspace, WorkspaceMixin
 
 
 DAYS_IN_2_WEEKS = 14
@@ -131,7 +131,7 @@ def submit(
         self,
         app: A,
         cfg: T,
-        workspace: Optional[str] = None,
+        workspace: Optional[Union[Workspace, str]] = None,
     ) -> str:
         """
         Submits the application to be run by the scheduler.
@@ -144,10 +144,9 @@ def submit(
         # pyre-fixme: Generic cfg type passed to resolve
         resolved_cfg = self.run_opts().resolve(cfg)
         if workspace:
-            sched = self
-            assert isinstance(sched, WorkspaceMixin)
-            role = app.roles[0]
-            sched.build_workspace_and_update_role(role, workspace, resolved_cfg)
+            assert isinstance(self, WorkspaceMixin)
+            self.build_workspace_and_update_role2(app.roles[0], workspace, resolved_cfg)
+
         # pyre-fixme: submit_dryrun takes Generic type for resolved_cfg
         dryrun_info = self.submit_dryrun(app, resolved_cfg)
         return self.schedule(dryrun_info)
@@ -356,13 +355,14 @@ def _validate(self, app: A, scheduler: str, cfg: T) -> None:
 
         Raises error if application is not compatible with scheduler
         """
-        if isinstance(app, AppDef):
-            for role in app.roles:
-                if role.resource == NULL_RESOURCE:
-                    raise ValueError(
-                        f"No resource for role: {role.image}."
-                        f" Did you forget to attach resource to the role"
-                    )
+        if not isinstance(app, AppDef):
+            return
+
+        for role in app.roles:
+            if role.resource == NULL_RESOURCE:
+                raise ValueError(
+                    f"No resource for role: {role.image}. Did you forget to attach resource to the role"
+                )
 
 
 def filter_regex(regex: str, data: Iterable[str]) -> Iterable[str]:

torchx/workspace/__init__.py

@@ -22,4 +22,4 @@
     * ``memory://foo-bar/`` an in-memory workspace for notebook/programmatic usage
 """
 
-from torchx.workspace.api import walk_workspace, WorkspaceMixin  # noqa: F401
+from torchx.workspace.api import walk_workspace, Workspace, WorkspaceMixin  # noqa: F401

torchx/workspace/api.py

@@ -9,8 +9,11 @@
 import abc
 import fnmatch
 import posixpath
+import shutil
+import tempfile
 import warnings
 from dataclasses import dataclass
+from pathlib import Path
 from typing import Any, Dict, Generic, Iterable, Mapping, Tuple, TYPE_CHECKING, TypeVar
 
 from torchx.specs import AppDef, CfgVal, Role, runopts
@@ -75,6 +78,64 @@ def build_workspace(self, sync: bool = True) -> PkgInfo[PackageType]:
         pass
 
 
+@dataclass
+class Workspace:
+    """
+    Specifies a local "workspace" (a set of directories). Workspaces are ad-hoc built
+    into an (usually ephemeral) image. This effectively mirrors the local code changes
+    at job submission time.
+
+    For example:
+
+      1. ``projects={"~/github/torch": "torch"}`` copies ``~/github/torch/**`` into ``$REMOTE_WORKSPACE_ROOT/torch/**``
+      2. ``projects={"~/github/torch": ""}`` copies ``~/github/torch/**`` into ``$REMOTE_WORKSPACE_ROOT/**``
+
+    The exact location of ``$REMOTE_WORKSPACE_ROOT`` is implementation dependent and varies between
+    different implementations of :py:class:`~torchx.workspace.api.WorkspaceMixin`.
+    Check the scheduler documentation for details on which workspace it supports.
+
+    Note: ``projects`` maps the location of the local project to a sub-directory in the remote workspace root directory.
+    Typically the local project location is a directory path (e.g. ``/home/foo/github/torch``).
+
+
+    Attributes:
+        projects: mapping of local project to the sub-dir in the remote workspace dir.
+    """
+
+    projects: dict[str, str]
+
+    @staticmethod
+    def from_str(workspace: str) -> "Workspace":
+        import yaml
+
+        projects = yaml.safe_load(workspace)
+        if isinstance(projects, str):  # single project workspace
+            projects = {projects: ""}
+        else:  # multi-project workspace
+            # Replace None mappings with "" (empty string)
+            projects = {k: ("" if v is None else v) for k, v in projects.items()}
+
+        return Workspace(projects)
+
+    def __str__(self) -> str:
+        """
+        Returns a string representation of the Workspace by concatenating
+        the project mappings using ';' as a delimiter and ':' between key and value.
+        If the single-project workspace with no target mapping, then simply
+        returns the src (local project dir)
+
+        NOTE: meant to be used for logging purposes not serde.
+          Therefore not symmetric with :py:func:`Workspace.from_str`.
+
+        """
+        if len(self.projects) == 1 and not next(iter(self.projects.values())):
+            return next(iter(self.projects))
+        else:
+            return ";".join(
+                k if not v else f"{k}:{v}" for k, v in self.projects.items()
+            )
+
+
 class WorkspaceMixin(abc.ABC, Generic[T]):
     """
     Note: (Prototype) this interface may change without notice!
@@ -100,6 +161,37 @@ def workspace_opts(self) -> runopts:
         """
         return runopts()
 
+    def build_workspace_and_update_role2(
+        self, role: Role, workspace: Workspace | str, cfg: Mapping[str, CfgVal]
+    ) -> None:
+        """
+        Same as :py:meth:`build_workspace_and_update_role` but operates
+        on :py:class:`Workspace` (supports multi-project workspaces)
+        as well as ``str`` (for backwards compatibility).
+
+        If ``workspace`` is a ``str`` this method simply calls
+        :py:meth:`build_workspace_and_update_role`.
+
+        If ``workspace`` is :py:class:`Workspace` then the default
+        impl copies all the projects into a tmp directory and passes the tmp dir to
+        :py:meth:`build_workspace_and_update_role`
+
+        Subclasses can override this method to customize multi-project
+        workspace building logic.
+        """
+        if isinstance(workspace, Workspace):
+            with tempfile.TemporaryDirectory(suffix="torchx_workspace_") as outdir:
+                for src, dst in workspace.projects.items():
+                    dst_path = Path(outdir) / dst
+                    if Path(src).is_file():
+                        shutil.copy2(src, dst_path)
+                    else:  # src is dir
+                        shutil.copytree(src, dst_path, dirs_exist_ok=True)
+
+                self.build_workspace_and_update_role(role, outdir, cfg)
+        else:  # workspace is str
+            self.build_workspace_and_update_role(role, workspace, cfg)
+
     @abc.abstractmethod
     def build_workspace_and_update_role(
         self, role: Role, workspace: str, cfg: Mapping[str, CfgVal]