torchx-nightly 2025.10.17__py3-none-any.whl → 2025.10.18__py3-none-any.whl

torchx/runner/api.py

@@ -427,41 +427,25 @@ class Runner:
             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 @@ class Scheduler(abc.ABC, Generic[T, A, D]):
         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 @@ class Scheduler(abc.ABC, Generic[T, A, D]):
         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 logging as logger
 import os
 import pathlib
 import re
+import shutil
 import typing
 import warnings
 from dataclasses import asdict, dataclass, field
@@ -381,6 +382,16 @@ class Workspace:
         """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
@@ -388,6 +399,39 @@ class Workspace:
         """
         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/workspace/api.py

@@ -8,26 +8,17 @@
 
 import abc
 import fnmatch
+import logging
 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,
-    Union,
-)
+from typing import Any, Dict, Generic, Iterable, Mapping, Tuple, TYPE_CHECKING, TypeVar
 
 from torchx.specs import AppDef, CfgVal, Role, runopts, Workspace
 
+logger: logging.Logger = logging.getLogger(__name__)
+
 if TYPE_CHECKING:
     from fsspec import AbstractFileSystem
 
@@ -113,45 +104,72 @@ class WorkspaceMixin(abc.ABC, Generic[T]):
         """
         return runopts()
 
-    def build_workspace_and_update_role2(
+    def build_workspaces(self, roles: list[Role], cfg: Mapping[str, CfgVal]) -> None:
+        """
+        NOTE: this method MUTATES the passed roles!
+
+        Builds the workspaces (if any) for each role and updates the role to reflect the built workspace.
+        Typically ``role.image`` is updated with the newly built image that reflects the local workspace.
+        Some workspace implementations may add extra environment variables to make it easier for other
+        parts of the program to access the workspace. For example a ``WORKSPACE_DIR`` env var may be added
+        to ``role.env`` that scripts can use to refert to the workspace directory in the container.
+        """
+
+        build_cache: dict[object, object] = {}
+
+        for i, role in enumerate(roles):
+            if role.workspace:
+                old_img = role.image
+                self.caching_build_workspace_and_update_role(role, cfg, build_cache)
+
+                if old_img != role.image:
+                    logger.info(
+                        "role[%d]=%s updated with new image to include workspace changes",
+                        i,
+                        role.name,
+                    )
+
+    def caching_build_workspace_and_update_role(
         self,
         role: Role,
-        workspace: Union[Workspace, str],
         cfg: Mapping[str, CfgVal],
+        build_cache: dict[object, object],
     ) -> 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).
+        Same as :py:meth:`build_workspace_and_update_role` but takes
+        a ``build_cache`` that can be used to cache pointers to build artifacts
+        between building workspace for each role.
 
-        If ``workspace`` is a ``str`` this method simply calls
+        This is useful when an appdef has multiple roles where the image and workspace
+        of the roles are the same but other attributes such as entrypoint or args are different.
+
+        NOTE: ``build_cache``'s lifetime is within :py:meth:`build_workspace_and_update_roles`
+        NOTE: the workspace implementation decides what to cache
+
+        Workspace subclasses should prefer implementing this method over
         :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`
+        The default implementation of this method simply calls the (deprecated) non-caching
+        :py:meth:`build_workspace_and_update_role` and deals with multi-dir workspaces by
+        merging them into a single tmpdir before passing it down.
 
-        Subclasses can override this method to customize multi-project
-        workspace building logic.
         """
-        if isinstance(workspace, Workspace):
-            if not workspace.is_unmapped_single_project():
-                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)
-                    return
-            else:  # single project workspace with no target mapping (treat like a str workspace)
-                workspace = str(workspace)
-
-        self.build_workspace_and_update_role(role, workspace, cfg)
 
-    @abc.abstractmethod
+        workspace = role.workspace
+
+        if not workspace:
+            return
+
+        if workspace.is_unmapped_single_project():
+            # single-dir workspace with no target map; no need to copy to a tmp dir
+            self.build_workspace_and_update_role(role, str(workspace), cfg)
+        else:
+            # multi-dirs or single-dir with a target map;
+            # copy all dirs to a tmp dir and treat the tmp dir as a single-dir workspace
+            with tempfile.TemporaryDirectory(suffix="torchx_workspace_") as outdir:
+                workspace.merge_into(outdir)
+                self.build_workspace_and_update_role(role, outdir, cfg)
+
     def build_workspace_and_update_role(
         self,
         role: Role,
@@ -159,6 +177,9 @@ class WorkspaceMixin(abc.ABC, Generic[T]):
         cfg: Mapping[str, CfgVal],
     ) -> None:
         """
+        .. note:: DEPRECATED: Workspace subclasses should implement
+                  :py:meth:`caching_build_workspace_and_update_role` over this method.
+
         Builds the specified ``workspace`` with respect to ``img``
         and updates the ``role`` to reflect the built workspace artifacts.
         In the simplest case, this method builds a new image and updates
@@ -167,7 +188,7 @@ class WorkspaceMixin(abc.ABC, Generic[T]):
 
         Note: this method mutates the passed ``role``.
         """
-        ...
+        raise NotImplementedError("implement `caching_build_workspace_and_update_role`")
 
     def dryrun_push_images(self, app: AppDef, cfg: Mapping[str, CfgVal]) -> T:
         """

{torchx_nightly-2025.10.17.dist-info → torchx_nightly-2025.10.18.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: torchx-nightly
-Version: 2025.10.17
+Version: 2025.10.18
 Summary: TorchX SDK and Components
 Home-page: https://github.com/meta-pytorch/torchx
 Author: TorchX Devs

{torchx_nightly-2025.10.17.dist-info → torchx_nightly-2025.10.18.dist-info}/RECORD

@@ -48,7 +48,7 @@ torchx/examples/apps/lightning/profiler.py,sha256=SSSihnwjeUTkBoz0E3qn1b-wbkfUIo
 torchx/examples/apps/lightning/train.py,sha256=0wvvshGHvZowePB4LfclXwn40X7i9euM0ReETWBcPSo,6253
 torchx/pipelines/__init__.py,sha256=2MbRVk5xwRjg-d2qPemeXpEhDsocMQumPQ53lsesZAI,606
 torchx/runner/__init__.py,sha256=x8Sz7s_tLxPgJgvWIhK4ju9BNZU61uBFywGwDY6CqJs,315
-torchx/runner/api.py,sha256=jxtgOl7nNOqpzG-sjUJngXhIOachqaVfKu9rF8YqHWI,31271
+torchx/runner/api.py,sha256=s-Fh9FJoZ9JiEFXyGVF0vdJuXMQ-Wa2v4yLjvWJKBX8,30134
 torchx/runner/config.py,sha256=SaKOB50d79WaMFPWK8CC4as6UaNFaRGhrBkfajq3KC4,18311
 torchx/runner/events/__init__.py,sha256=cMiNjnr4eUNQ2Nxxtu4nsvN5lu56b-a6nJ-ct3i7DQk,5536
 torchx/runner/events/api.py,sha256=bvxKBAYK8LzbrBNaNLgL1x0aivtfANmWo1EMGOrSR8k,2668
@@ -57,7 +57,7 @@ torchx/runtime/__init__.py,sha256=Wxje2BryzeQneFu5r6P9JJiEKG-_C9W1CcZ_JNrKT6g,59
 torchx/runtime/tracking/__init__.py,sha256=dYnAPnrXYREfPXkpHhdOFkcYIODWEbA13PdD-wLQYBo,3055
 torchx/runtime/tracking/api.py,sha256=SmUQyUKZqG3KlAhT7CJOGqRz1O274E4m63wQeOVq3CU,5472
 torchx/schedulers/__init__.py,sha256=FQN9boQM4mwOD3sK9LZ3GBgw-gJ7Vx4MFj6z6ATQIrc,2211
-torchx/schedulers/api.py,sha256=5Amli1httEl82XebAqd8vl3dM8zMKwYfRgfd0mEq3is,14538
+torchx/schedulers/api.py,sha256=smoUv1ocfqsBRmesXbz9i1F86zBOixZ8QHxYmI_MzgQ,14649
 torchx/schedulers/aws_batch_scheduler.py,sha256=-HpjNVhSFBDxZo3cebK-3YEguB49dxoaud2gz30cAVM,29437
 torchx/schedulers/aws_sagemaker_scheduler.py,sha256=flN8GumKE2Dz4X_foAt6Jnvt-ZVojWs6pcyrHwB0hz0,20921
 torchx/schedulers/devices.py,sha256=RjVcu22ZRl_9OKtOtmA1A3vNXgu2qD6A9ST0L0Hsg4I,1734
@@ -70,7 +70,7 @@ torchx/schedulers/lsf_scheduler.py,sha256=YS6Yel8tXJqLPxbcGz95lZG2nCi36AQXdNDyuB
 torchx/schedulers/slurm_scheduler.py,sha256=vypGaCZe61bkyNkqRlK4Iwmk_NaAUQi-DsspaWd6BZw,31873
 torchx/schedulers/streams.py,sha256=8_SLezgnWgfv_zXUsJCUM34-h2dtv25NmZuxEwkzmxw,2007
 torchx/specs/__init__.py,sha256=SXS4r_roOkbbAL-p7EY5fl5ou-AG7S9Ck-zKtRBdHOk,6760
-torchx/specs/api.py,sha256=mRX6sqBQOpEVuyR3M-tK3LqSGNNVzqoOSIWM6FolqZQ,51118
+torchx/specs/api.py,sha256=6_Q5SOUqfMRqQdnlvem36dBusHpynp4PCa9uTNwvzwo,52630
 torchx/specs/builders.py,sha256=Ye3of4MupJ-da8vLaX6_-nzGo_FRw1BFpYsX6dAZCNk,13730
 torchx/specs/file_linter.py,sha256=z0c4mKJv47BWiPaWCdUM0A8kHwnj4b1s7oTmESuD9Tc,14407
 torchx/specs/finder.py,sha256=gWQNEFrLYqrZoI0gMMhQ70YAC4sxqS0ZFpoWAmcVi44,17438
@@ -99,12 +99,12 @@ torchx/util/shlex.py,sha256=eXEKu8KC3zIcd8tEy9_s8Ds5oma8BORr-0VGWNpG2dk,463
 torchx/util/strings.py,sha256=7Ef1loz2IYMrzeJ6Lewywi5cBIc3X3g7lSPbT1Tn_z4,664
 torchx/util/types.py,sha256=E9dxAWQnsJkIDuHtg-poeOJ4etucSI_xP_Z5kNJX8uI,9229
 torchx/workspace/__init__.py,sha256=FqN8AN4VhR1C_SBY10MggQvNZmyanbbuPuE-JCjkyUY,798
-torchx/workspace/api.py,sha256=h2SaC-pYPBLuo3XtkXJ0APMoro-C-ry7KucI7r3EUf4,8753
+torchx/workspace/api.py,sha256=UESQ4qgxXjsb6Y1wP9OGv2ixaFgaTs3SqghmNuOJIZM,10235
 torchx/workspace/dir_workspace.py,sha256=npNW_IjUZm_yS5r-8hrRkH46ndDd9a_eApT64m1S1T4,2268
 torchx/workspace/docker_workspace.py,sha256=PFu2KQNVC-0p2aKJ-W_BKA9ZOmXdCY2ABEkCExp3udQ,10269
-torchx_nightly-2025.10.17.dist-info/LICENSE,sha256=WVHfXhFC0Ia8LTKt_nJVYobdqTJVg_4J3Crrfm2A8KQ,1721
-torchx_nightly-2025.10.17.dist-info/METADATA,sha256=qg3D7hiG2TX50v4beQLMAQ8dRszPGKUVYke16hnu87Q,5046
-torchx_nightly-2025.10.17.dist-info/WHEEL,sha256=tZoeGjtWxWRfdplE7E3d45VPlLNQnvbKiYnx7gwAy8A,92
-torchx_nightly-2025.10.17.dist-info/entry_points.txt,sha256=T328AMXeKI3JZnnxfkEew2ZcMN1oQDtkXjMz7lkV-P4,169
-torchx_nightly-2025.10.17.dist-info/top_level.txt,sha256=pxew3bc2gsiViS0zADs0jb6kC5v8o_Yy_85fhHj_J1A,7
-torchx_nightly-2025.10.17.dist-info/RECORD,,
+torchx_nightly-2025.10.18.dist-info/LICENSE,sha256=WVHfXhFC0Ia8LTKt_nJVYobdqTJVg_4J3Crrfm2A8KQ,1721
+torchx_nightly-2025.10.18.dist-info/METADATA,sha256=uFXR5-5aPu-A24rIhwXi5QhcJUfQEWX0Vz5Ohb66JcA,5046
+torchx_nightly-2025.10.18.dist-info/WHEEL,sha256=tZoeGjtWxWRfdplE7E3d45VPlLNQnvbKiYnx7gwAy8A,92
+torchx_nightly-2025.10.18.dist-info/entry_points.txt,sha256=T328AMXeKI3JZnnxfkEew2ZcMN1oQDtkXjMz7lkV-P4,169
+torchx_nightly-2025.10.18.dist-info/top_level.txt,sha256=pxew3bc2gsiViS0zADs0jb6kC5v8o_Yy_85fhHj_J1A,7
+torchx_nightly-2025.10.18.dist-info/RECORD,,