torchx.workspace

Tip

Workspaces handle automatic image patching – copying local code changes into the job’s runtime environment so users don’t need manual image rebuilds. Scheduler authors add workspace support by mixing in a WorkspaceMixin subclass alongside their Scheduler.

Why workspaces exist. Without workspaces, every code change requires rebuilding and pushing a Docker image before submitting a remote job. Workspaces automate this: TorchX overlays local changes onto the base image and submits the patched image in one step.

How It Works

User's workspace directory            Scheduler submission
┌──────────────────────┐
│  src/                │
│  train.py            │     build_workspaces()
│  Dockerfile.torchx   │ ──────────────────────────┐
│  .torchxignore       │                           │
└──────────────────────┘                           ▼
                               ┌──────────────────────────────────┐
                               │  For each Role with a workspace: │
                               │                                  │
                               │  1. Walk workspace               │
                               │     (respecting .torchxignore)   │
                               │                                  │
                               │  2. Build patched artifact       │
                               │     ┌────────────────────────┐   │
                               │     │ DockerWorkspaceMixin:  │   │
                               │     │   docker build + push  │   │
                               │     │ DirWorkspaceMixin:     │   │
                               │     │   copy to shared dir   │   │
                               │     └────────────────────────┘   │
                               │                                  │
                               │  3. Mutate role.image in-place   │
                               │     to reference patched image   │
                               └──────────────┬───────────────────┘
                                              │
                                              ▼
                               ┌──────────────────────────────────┐
                               │  Scheduler._submit_dryrun(app)   │
                               │  uses the patched role.image     │
                               └──────────────────────────────────┘

When workspace= is passed to run() (or --workspace on the CLI), TorchX patches the image before submission:

1. build_workspaces() iterates over each role’s workspace.

2. For each role with a workspace, it calls caching_build_workspace_and_update_role(), which builds the workspace and mutates role.image in-place to reference the patched artifact.

3. For remote schedulers, dryrun_push_images() and push_images() handle pushing the built image to a remote registry.

Important

The role mutated in step 3 belongs to a copy of your AppDef that dryrun() takes before patching – your own AppDef still carries the image you authored. To see what was actually submitted, read app:

info = runner.dryrun(app, "kubernetes", cfg)
info.app.roles[0].image  # patched; app.roles[0].image is not

submit() takes no such copy and patches the AppDef you hand it.

Note

DockerWorkspaceMixin uses Dockerfile.torchx from the workspace root (if present) instead of the default Dockerfile.

Built-in Mixins

Mixin	Strategy
DockerWorkspaceMixin	Builds a Docker image from a Dockerfile.torchx in the workspace, tags it with a content hash, and pushes to the configured image_repo. Used by kubernetes, local_docker, and slurm with workspace_type=docker.
DirWorkspaceMixin	Copies workspace files into a shared job directory on the filesystem. Used by slurm (its default).
MultiWorkspaceMixin	Offers several builders on one scheduler; the workspace_type run option picks one. Used by slurm.

Offering More Than One Builder

A backend that accepts more than one kind of image (a shared directory and a container image, say) offers one builder per kind through MultiWorkspaceMixin. The workspace_type run option picks the builder; the first one is the default, so existing invocations behave as before.

from typing import Any, Mapping

from torchx.schedulers.api import Scheduler
from torchx.specs import CfgVal
from torchx.workspace import MultiWorkspaceMixin, WorkspaceMixin
from torchx.workspace.dir_workspace import DirWorkspaceMixin
from torchx.workspace.docker_workspace import DockerWorkspaceMixin


class MyScheduler(MultiWorkspaceMixin, Scheduler[Mapping[str, CfgVal]]):
    def __init__(self, session_name: str) -> None:
        super().__init__("my_backend", session_name)
        self._builders: dict[str, WorkspaceMixin[Any]] = {
            "dir": DirWorkspaceMixin(),        # the default
            "docker": DockerWorkspaceMixin(),
        }

    def workspace_builders(self) -> Mapping[str, WorkspaceMixin[Any]]:
        return self._builders

torchx runopts my_backend then lists workspace_type next to every builder’s own options, and torchx run -s my_backend -cfg workspace_type=docker,image_repo=example.com/repo --workspace . ... builds and pushes an image instead of copying a directory. The slurm scheduler ships this way: dir copies the workspace into job_dir when that option is set; docker builds a patched image, pushes it to image_repo and runs every replica through the cluster’s container plugin (srun --container-image).

Implementing a Custom WorkspaceMixin

Subclass WorkspaceMixin and implement caching_build_workspace_and_update_role():

from typing import Any, Mapping

from torchx.specs import CfgVal, Role, runopts
from torchx.workspace import WorkspaceMixin


class MyWorkspaceMixin(WorkspaceMixin[None]):
    """Patches images by uploading workspace to a custom artifact store."""

    def workspace_opts(self) -> runopts:
        opts = runopts()
        opts.add("artifact_bucket", type_=str, required=True, help="S3 bucket for workspace artifacts")
        return opts

    def caching_build_workspace_and_update_role(
        self,
        role: Role,
        cfg: Mapping[str, CfgVal],
        build_cache: dict[object, object],
    ) -> None:
        workspace = role.workspace
        if not workspace:
            return

        bucket = cfg.get("artifact_bucket")
        # ... upload workspace files to bucket ...
        # ... update role.image or role.env to reference the artifact ...
        role.env["WORKSPACE_ARTIFACT"] = f"s3://{bucket}/{role.name}/workspace.tar.gz"

Then mix it into your scheduler:

from torchx.schedulers.api import Scheduler

class MyScheduler(MyWorkspaceMixin, Scheduler[Mapping[str, CfgVal]]):
    def __init__(self, session_name: str, **kwargs: object) -> None:
        super().__init__("my_backend", session_name)
    # ... scheduler methods ...

The generic parameter T is the type returned by dryrun_push_images and consumed by push_images. Use None if no separate push step is needed.

Note

The build_cache dict is shared across all roles in a single build_workspaces call. Use it to skip redundant builds when roles share the same image and workspace.

Testing Your Workspace Mixin

Create a Role with a workspace and assert the role was mutated correctly:

import unittest
from torchx.specs import Role, Resource, Workspace

class MyWorkspaceMixinTest(unittest.TestCase):
    def test_build_updates_role(self) -> None:
        mixin = MyWorkspaceMixin()
        role = Role(
            name="worker", image="base:latest", entrypoint="echo",
            resource=Resource(cpu=1, gpu=0, memMB=512),
            workspace=Workspace.from_str("/tmp/my_workspace"),
        )
        mixin.caching_build_workspace_and_update_role(
            role, cfg={"artifact_bucket": "my-bucket"}, build_cache={},
        )
        self.assertIn("WORKSPACE_ARTIFACT", role.env)

See torchx/workspace/test/ for the built-in mixin tests.

Common Pitfalls

• Implementing the deprecated method: Override caching_build_workspace_and_update_role (not the older build_workspace_and_update_role).

• MRO: List the mixin first in the base class list: class MyScheduler(MyWorkspaceMixin, Scheduler[...]):. Python’s MRO requires cooperative super().__init__().

• Forgetting to check ``role.workspace``: Guard with if not role.workspace: return – the method is called for every role.

.torchxignore

Place a .torchxignore file (same syntax as .dockerignore) at the workspace root to exclude files from the job image:

# Exclude version control and IDE files
.git
.vscode
__pycache__

# Exclude data directories
data/
*.csv

# But include a specific config file
!data/config.yaml

Lines starting with ! negate a previous pattern (include the file even if a prior rule excluded it). Blank lines and lines starting with # are ignored.

API Reference

class torchx.workspace.WorkspaceMixin(*args: object, **kwargs: object)

Scheduler mix-in that auto-builds a local workspace into a deployable image or patch.

Warning

Prototype – this interface may change without notice.

Attach to a Scheduler so that local code changes in the workspace are automatically reflected at runtime (via a rebuilt image or an overlaid diff patch) without a manual image rebuild.

build_workspace_and_update_role(role: Role, workspace: str, cfg: Mapping[str, str | int | float | bool | list[str] | dict[str, str] | None]) → None

Build workspace and mutate role to reference the resulting artifact.

Deprecated since version Implement: caching_build_workspace_and_update_role() instead.

build_workspaces(roles: list[Role], cfg: Mapping[str, str | int | float | bool | list[str] | dict[str, str] | None]) → None

Builds workspaces for each role and updates role.image in-place.

Important

Mutates the passed roles. May also add env vars (e.g. WORKSPACE_DIR) to role.env. role.workspace is left set, so a role that has been through this method is not distinguishable by inspection from one that has not – compare role.image instead.

Called by dryrun() with the roles of its private copy, so the caller’s AppDef is unaffected there; submit() passes the caller’s own roles.

caching_build_workspace_and_update_role(role: Role, cfg: Mapping[str, str | int | float | bool | list[str] | dict[str, str] | None], build_cache: dict[object, object]) → None

Like build_workspace_and_update_role() but with a per-call build_cache.

Subclasses should implement this method instead of build_workspace_and_update_role(). The cache avoids redundant builds when multiple roles share the same image and workspace.

Important

build_cache lifetime is scoped to a single build_workspaces() call. What gets cached is up to the implementation.

The default implementation delegates to the (deprecated) build_workspace_and_update_role(), merging multi-dir workspaces into a single tmpdir first.

dryrun_push_images(app: AppDef, cfg: Mapping[str, str | int | float | bool | list[str] | dict[str, str] | None]) → T

Dry-run the image push: updates app with final image names.

Only called for remote jobs. push_images() must be called with the return value before scheduling. Leave this method unoverridden when the builder has no push step (its build already put the artifact where the job reads it); MultiWorkspaceMixin then pushes nothing for it.

push_images(images_to_push: T) → None

Pushes images (returned by dryrun_push_images()) to the remote repo.

workspace_opts() → runopts

Returns the runopts accepted by this workspace.

class torchx.workspace.MultiWorkspaceMixin(*args: object, **kwargs: object)

Scheduler mix-in that offers several workspace builders; the workspace_type run option picks one.

A scheduler whose backend accepts more than one kind of image (a shared directory and a container image, say) returns one builder per kind from workspace_builders(). Every WorkspaceMixin call is forwarded to the builder the run config selects, and workspace_opts() is the union of the builders’ options plus the selector, so torchx runopts <scheduler> lists them all.

>>> from typing import Any, Mapping
>>> from torchx.workspace import MultiWorkspaceMixin, WorkspaceMixin
>>> from torchx.workspace.dir_workspace import DirWorkspaceMixin, TmpDirWorkspaceMixin
>>> class TwoWaysMixin(MultiWorkspaceMixin):
...     def __init__(self) -> None:
...         self.builders = {"dir": DirWorkspaceMixin(), "tmpdir": TmpDirWorkspaceMixin()}
...     def workspace_builders(self) -> Mapping[str, WorkspaceMixin[Any]]:
...         return self.builders
>>> mixin = TwoWaysMixin()
>>> mixin.workspace_builder({})[0]              # the first builder is the default
'dir'
>>> mixin.workspace_builder({"workspace_type": "tmpdir"})[0]
'tmpdir'
>>> sorted(mixin.workspace_opts()._opts)
['workspace_type']

The first builder is the default, so adding this mix-in to a scheduler that used one builder changes nothing until a user sets workspace_type.

WORKSPACE_TYPE_OPT: str = 'workspace_type'

the run option that picks the builder

caching_build_workspace_and_update_role(role: Role, cfg: Mapping[str, str | int | float | bool | list[str] | dict[str, str] | None], build_cache: dict[object, object]) → None

Builds role’s workspace with the selected builder.

dryrun_push_images(app: AppDef, cfg: Mapping[str, str | int | float | bool | list[str] | dict[str, str] | None]) → tuple[str, Any]

Runs the selected builder’s push dry-run.

Returns (name, images) so push_images() can find the same builder again; images is None for a builder with no push step, and push_images() then pushes nothing.

push_images(images_to_push: tuple[str, Any]) → None

Pushes with the builder named by dryrun_push_images().

workspace_builder(cfg: Mapping[str, str | int | float | bool | list[str] | dict[str, str] | None]) → tuple[str, WorkspaceMixin[Any]]

Returns (name, builder) for cfg["workspace_type"], or the default builder when the option is unset.

Raises:

ValueError – the requested type is not one this scheduler offers; the message lists the names it does offer.

abstractmethod workspace_builders() → Mapping[str, WorkspaceMixin[Any]]

Returns {name: builder}. The first entry is the default.

Return the same builder instances on every call (build the mapping once, in __init__): the push step looks the builder up again by name and expects the instance that ran the build.

workspace_opts() → runopts

The selector plus the union of every builder’s options.

Two builders declaring the same option name must mean the same thing by it; the later builder’s declaration wins.

torchx.workspace.walk_workspace(fs: AbstractFileSystem, path: str, ignore_name: str = '.torchxignore') → Iterable[tuple[str, Iterable[str], Mapping[str, Mapping[str, object]]]]

Walks path on fs, filtering entries via .dockerignore-style rules read from ignore_name.

torchx.workspace.docker_workspace

class torchx.workspace.docker_workspace.DockerWorkspaceMixin(*args: object, docker_client: DockerClient | None = None, **kwargs: object)

Bases: WorkspaceMixin[dict[str, tuple[str, str]]]

Builds patched Docker images from the workspace.

Requires a local Docker daemon. For remote jobs, authenticate via docker login and set the image_repo runopt.

If Dockerfile.torchx exists in the workspace it is used as the Dockerfile; otherwise a default COPY . . Dockerfile is generated. Extra --build-arg values available in Dockerfile.torchx:

• IMAGE – the role’s base image

• WORKSPACE – the workspace path

Use .dockerignore to exclude files from the build context.

build_workspace_and_update_role(role: Role, workspace: str, cfg: Mapping[str, str | int | float | bool | list[str] | dict[str, str] | None]) → None

Builds a Docker image from workspace on top of role.image and updates role.image with the resulting image id.

dryrun_push_images(app: AppDef, cfg: Mapping[str, str | int | float | bool | list[str] | dict[str, str] | None]) → dict[str, tuple[str, str]]

Replaces local sha256:... images in app with remote paths and returns a {local_image: (repo, tag)} mapping for push_images().

push_images(images_to_push: dict[str, tuple[str, str]]) → None

Pushes local images to a remote repository.

Requires docker login authentication to the target repo.

workspace_opts() → runopts

Returns the runopts accepted by this workspace.

torchx.workspace.dir_workspace

class torchx.workspace.dir_workspace.DirWorkspaceMixin(*args: object, **kwargs: object)

Bases: WorkspaceMixin[None]

build_workspace_and_update_role(role: Role, workspace: str, cfg: Mapping[str, str | int | float | bool | list[str] | dict[str, str] | None]) → None

Copies workspace into cfg["job_dir"] and sets role.image to it.

No-op if job_dir is not set. Files matching .torchxignore patterns are skipped.

See also

torchx.schedulers

Scheduler API reference and implementation guide (including workspace integration).

Advanced Usage

Registering custom schedulers and workspace mixins via entry points.