(torchx/workspace) Add API for caching workspace build artifacts between roles (#1149)

Summary:

Since we now allow each `Role` to specify its own workspace through the `Role.workspace` attribute, there are cases where all the roles have the same workspace and base image. In this case, we don't have to re-build the workspace image for each role since the ephemerals will all have the same content!

This diff introduces the `caching_build_workspace_and_update_role()` method to the `WorkspaceMixIn`. This is the preferred method (over `build_workspace_and_update_role()`) that subclasses of `WorkspaceMixIn` should implement.

Until we can completely deprecate `build_workspace_and_update_role()`, we bridge the two by having the default implementation of `caching_build_workspace_and_update_role()` call `build_workspace_and_update_role()`.

To support `caching_build_workspace_and_update_role()`, this diff also does the following:

1. Renames `build_workspace_and_udpate_role2()` to `caching_build_workspace_and_update_role()`.

2. Iterating over the roles and building the workspace for each role is now pushed into the `WorkspaceMixIn.build_workspace_and_update_roles()` instead of it being in the runner.

3. Pushes the logic of merging of multi-directory workspace into a single tmp dir to the `Workspace` class's `merge_into()` method.

Reviewed By: AbishekS

Differential Revision: D84466900

Author: kiukchung

torchx/runner/api.py

@@ -427,41 +427,25 @@ def dryrun(
             sched._pre_build_validate(app, scheduler, resolved_cfg)
 
             if isinstance(sched, WorkspaceMixin):
-                for i, role in enumerate(app.roles):
-                    role_workspace = role.workspace
-
-                    if i == 0 and workspace:
-                        # NOTE: torchx originally took workspace as a runner arg and only applied the workspace to role[0]
-                        #   later, torchx added support for the workspace attr in Role
-                        #   for BC, give precedence to the workspace argument over the workspace attr for role[0]
-                        if role_workspace:
-                            logger.info(
-                                f"Using workspace={workspace} over role[{i}].workspace={role_workspace} for role[{i}]={role.name}."
-                                " To use the role's workspace attr pass: --workspace='' from CLI or workspace=None programmatically."  # noqa: B950
-                            )
-                        role_workspace = workspace
-
-                    if role_workspace:
-                        old_img = role.image
+                if workspace:
+                    # NOTE: torchx originally took workspace as a runner arg and only applied the workspace to role[0]
+                    # later, torchx added support for the workspace attr in Role
+                    # for BC, give precedence to the workspace argument over the workspace attr for role[0]
+                    if app.roles[0].workspace:
                         logger.info(
-                            f"Checking for changes in workspace `{role_workspace}` for role[{i}]={role.name}..."
-                        )
-                        # TODO kiuk@ once we deprecate the `workspace` argument in runner APIs we can simplify the signature of
-                        #   build_workspace_and_update_role2() to just taking the role and resolved_cfg
-                        sched.build_workspace_and_update_role2(
-                            role, role_workspace, resolved_cfg
+                            "Overriding role[%d] (%s) workspace to `%s`"
+                            "To use the role's workspace attr pass: --workspace='' from CLI or workspace=None programmatically.",
+                            0,
+                            role.name,
+                            str(app.roles[0].workspace),
                         )
+                    app.roles[0].workspace = (
+                        Workspace.from_str(workspace)
+                        if isinstance(workspace, str)
+                        else workspace
+                    )
 
-                        if old_img != role.image:
-                            logger.info(
-                                f"Built new image `{role.image}` based on original image `{old_img}`"
-                                f" and changes in workspace `{role_workspace}` for role[{i}]={role.name}."
-                            )
-                        else:
-                            logger.info(
-                                f"Reusing original image `{old_img}` for role[{i}]={role.name}."
-                                " Either a patch was built or no changes to workspace was detected."
-                            )
+                sched.build_workspaces(app.roles, resolved_cfg)
 
             sched._validate(app, scheduler, resolved_cfg)
             dryrun_info = sched.submit_dryrun(app, resolved_cfg)

torchx/schedulers/api.py

@@ -131,7 +131,7 @@ def submit(
         self,
         app: A,
         cfg: T,
-        workspace: Optional[Union[Workspace, str]] = None,
+        workspace: str | Workspace | None = None,
     ) -> str:
         """
         Submits the application to be run by the scheduler.
@@ -145,7 +145,12 @@ def submit(
         resolved_cfg = self.run_opts().resolve(cfg)
         if workspace:
             assert isinstance(self, WorkspaceMixin)
-            self.build_workspace_and_update_role2(app.roles[0], workspace, resolved_cfg)
+
+            if isinstance(workspace, str):
+                workspace = Workspace.from_str(workspace)
+
+            app.roles[0].workspace = workspace
+            self.build_workspaces(app.roles, resolved_cfg)
 
         # pyre-fixme: submit_dryrun takes Generic type for resolved_cfg
         dryrun_info = self.submit_dryrun(app, resolved_cfg)

torchx/specs/api.py

@@ -14,6 +14,7 @@
 import os
 import pathlib
 import re
+import shutil
 import typing
 import warnings
 from dataclasses import asdict, dataclass, field
@@ -381,13 +382,56 @@ def __bool__(self) -> bool:
         """False if no projects mapping. Lets us use workspace object in an if-statement"""
         return bool(self.projects)
 
+    def __eq__(self, other: object) -> bool:
+        if not isinstance(other, Workspace):
+            return False
+        return self.projects == other.projects
+
+    def __hash__(self) -> int:
+        # makes it possible to use Workspace as the key in the workspace build cache
+        # see WorkspaceMixin.caching_build_workspace_and_update_role
+        return hash(frozenset(self.projects.items()))
+
     def is_unmapped_single_project(self) -> bool:
         """
         Returns ``True`` if this workspace only has 1 project
         and its target mapping is an empty string.
         """
         return len(self.projects) == 1 and not next(iter(self.projects.values()))
 
+    def merge_into(self, outdir: str | pathlib.Path) -> None:
+        """
+        Copies each project dir of this workspace into the specified ``outdir``.
+        Each project dir is copied into ``{outdir}/{target}`` where ``target`` is
+        the target mapping of the project dir.
+
+        For example:
+
+        .. code-block:: python
+            from os.path import expanduser
+
+            workspace = Workspace(
+                projects={
+                    expanduser("~/workspace/torch"): "torch",
+                    expanduser("~/workspace/my_project": "")
+                }
+            )
+            workspace.merge_into(expanduser("~/tmp"))
+
+        Copies:
+
+            * ``~/workspace/torch/**`` into ``~/tmp/torch/**``
+            * ``~/workspace/my_project/**`` into ``~/tmp/**``
+
+        """
+
+        for src, dst in self.projects.items():
+            dst_path = pathlib.Path(outdir) / dst
+            if pathlib.Path(src).is_file():
+                shutil.copy2(src, dst_path)
+            else:  # src is dir
+                shutil.copytree(src, dst_path, dirs_exist_ok=True)
+
     @staticmethod
     def from_str(workspace: str | None) -> "Workspace":
         import yaml

torchx/specs/test/api_test.py

@@ -46,6 +46,7 @@
     TORCHX_HOME,
     Workspace,
 )
+from torchx.test.fixtures import TestWithTmpDir
 
 
 class TorchXHomeTest(unittest.TestCase):
@@ -75,7 +76,7 @@ def test_TORCHX_HOME_override(self) -> None:
                 self.assertTrue(conda_pack_out.is_dir())
 
 
-class WorkspaceTest(unittest.TestCase):
+class WorkspaceTest(TestWithTmpDir):
 
     def test_bool(self) -> None:
         self.assertFalse(Workspace(projects={}))
@@ -149,6 +150,28 @@ def test_from_str_multi_project(self) -> None:
             ).projects,
         )
 
+    def test_merge(self) -> None:
+        self.touch("workspace/myproj/README.md")
+        self.touch("workspace/myproj/bin/cli")
+
+        self.touch("workspace/torch/setup.py")
+        self.touch("workspace/torch/torch/__init__.py")
+
+        w = Workspace(
+            projects={
+                str(self.tmpdir / "workspace/myproj"): "",
+                str(self.tmpdir / "workspace/torch"): "torch",
+            }
+        )
+
+        outdir = self.tmpdir / "out"
+        w.merge_into(outdir)
+
+        self.assertTrue((outdir / "README.md").is_file())
+        self.assertTrue((outdir / "bin/cli").is_file())
+        self.assertTrue((outdir / "torch/setup.py").is_file())
+        self.assertTrue((outdir / "torch/torch/__init__.py").is_file())
+
 
 class AppDryRunInfoTest(unittest.TestCase):
     def test_repr(self) -> None:

torchx/test/fixtures.py

@@ -76,6 +76,14 @@ def touch(self, filepath: str) -> Path:
         f.touch()
         return f
 
+    def mkdir(self, dirpath: str) -> Path:
+        """
+        Creates a directory in the test's tmpdir and returns the ``Path`` to the created directory
+        """
+        d = self.tmpdir / dirpath
+        d.mkdir(parents=True, exist_ok=True)
+        return d
+
     def write(self, filepath: str, content: Iterable[str]) -> Path:
         """
         Creates a file given the filepath (can be a file name or a relative path) in the test's tmpdir
@@ -136,6 +144,109 @@ def read(self, filepath: Union[str, Path]) -> List[str]:
         with open(self.tmpdir / filepath, "r") as fin:
             return fin.readlines()
 
+    def assertDirIsEmpty(self, directory: str | Path) -> None:
+        """
+        Asserts that the given directory exists but is empty.
+        If ``dir`` is a relative path, it is assumed to be relative to this test's ``tmpdir``.
+        """
+        directory = Path(directory)
+        d = directory if directory.is_absolute() else self.tmpdir / directory
+        self.assertTrue(d.is_dir(), f"{d} is not a directory")
+        self.assertEqual(0, len(list(d.iterdir())), f"{d} is not empty")
+
+    def assertDirTree(self, root: str | Path, tree: dict[str, object]) -> None:
+        """
+        Asserts that the given ``root`` has the directory structure as specified by ``tree``.
+        If ``root`` is a relative path, it is assumed to be relative to this test's ``tmpdir``
+
+        Usage:
+
+        .. code-block:: python
+            self.assertDirTree(
+                "out",
+                {
+                    "setup.py": "",
+                    "torchx": {
+                        "__init__.py": "",
+                        "version.py": "0.8.0dev0",
+                        "specs": {
+                            "__init__.py": "",
+                            "api.py": "",
+                        },
+                    },
+                },
+            )
+        """
+        root = Path(root)
+        d = root if root.is_absolute() else self.tmpdir / root
+        self.assertTrue(d.is_dir(), f"{d} is not a directory")
+
+        def _assert_tree(current_dir: Path, subtree: dict[str, object]) -> None:
+            # Check that the directory contains exactly the keys in subtree
+            actual_entries = [p.name for p in current_dir.iterdir()]
+            expected_entries = list(subtree.keys())
+            self.assertListEqual(
+                expected_entries,
+                actual_entries,
+                f"contents of the dir `{current_dir}` do not match the expected",
+            )
+            for name, value in subtree.items():
+                path = current_dir / name
+                if isinstance(value, dict):
+                    self.assertTrue(path.is_dir(), f"{path} is not a directory")
+                    _assert_tree(path, value)
+                else:
+                    self.assertTrue(path.is_file(), f"{path} is not a file")
+                    if value != "":
+                        with open(path, "r") as f:
+                            content = f.read().strip()
+                        self.assertEqual(
+                            content,
+                            value,
+                            f"file {path} content {content!r} does not match expected {value!r}",
+                        )
+
+        _assert_tree(d, tree)
+
+    def create_dir_tree(self, root: str, tree: dict[str, object]) -> Path:
+        """
+        Creates the directory structure as specified by ``tree`` under ``self.tmpdir / root``
+
+        Usage:
+
+        .. code-block:: python
+            self.createDirTree(
+                "out",
+                {
+                    "README.md": "foobar",
+                    "torchx": {
+                        "__init__.py": "",
+                        "version.py": "0.8.0dev0",
+                        "specs": {
+                            "__init__.py": "",
+                            "api.py": "",
+                        },
+                    },
+                },
+            )
+        """
+        d = self.tmpdir / root
+        d.mkdir(parents=True, exist_ok=True)
+
+        def _create_tree(current_dir: Path, subtree: dict[str, object]) -> None:
+            for name, value in subtree.items():
+                path = current_dir / name
+                if isinstance(value, dict):
+                    path.mkdir(parents=True, exist_ok=True)
+                    _create_tree(path, value)
+                else:
+                    path.parent.mkdir(parents=True, exist_ok=True)
+                    with open(path, "w") as f:
+                        f.write(str(value))
+
+        _create_tree(d, tree)
+        return d
+
 
 Ret = TypeVar("Ret")