Skip to content

(feature)[torchx][workspace] Schedulers can offer multiple workspace builders - #1458

Merged
kiukchung merged 1 commit into
mainfrom
kiuk/torchx-multi-workspace-20260919
Sep 19, 2026
Merged

kiukchung merged 1 commit into
mainfrom
kiuk/torchx-multi-workspace-20260919

Conversation

@kiukchung

@kiukchung kiukchung commented Sep 19, 2026 •

Copy link
Copy Markdown
Contributor

Summary:
A scheduler can now offer more than one workspace builder because MultiWorkspaceMixin selects a builder with workspace_type instead of hard-binding a single mix-in.

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

    def workspace_builders(self) -> Mapping[str, WorkspaceMixin[Any]]:
        return self._builders
torchx run -s my_backend --workspace . \
    -cfg workspace_type=docker,image_repo=example.com/repo utils.echo

slurm now also offers a docker builder that builds and pushes a patched image, renders srun --container-image=<image>, and leaves the img_root macro (the image's root path, substituted into args and env) empty inside the container.

What is unchanged:

  • local_cwd still has no workspace support.
  • local_docker and kubernetes still use DockerWorkspaceMixin and add no workspace_type option.
  • slurm with workspace_type unset still copies into job_dir; the rendered sbatch script stays byte-identical.

Closes #590

Test Plan:

  1. pytest across torchx/workspace/test, slurm_scheduler_test, api_test, and workspace_defaults_test -> 130 tests passed.
  2. ruff check (import sorting included) and ruff format --check with the pinned ruff -> clean.
  3. python -m doctest torchx/workspace/api.py -> passed; pyrefly check -> no findings in the changed files.
  4. Downstream plugin suite with torchx installed editable at this commit: 952 passed; the 3 failures are that suite's own generated-file checks and fail identically with the released torchx.

@meta-cla meta-cla Bot added the CLA Signed This label is managed by the Facebook bot. Authors need to sign the CLA before a PR can be reviewed. label Sep 19, 2026
@kiukchung
kiukchung force-pushed the kiuk/torchx-multi-workspace-20260919 branch from c6e0f02 to 58cc504 Compare September 19, 2026 21:37
@kiukchung
kiukchung marked this pull request as ready for review September 19, 2026 21:37
…builders

Summary:
A scheduler can now offer more than one workspace builder because `MultiWorkspaceMixin` selects a builder with `workspace_type` instead of hard-binding a single mix-in.

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

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

```bash
torchx run -s my_backend --workspace . \
    -cfg workspace_type=docker,image_repo=example.com/repo utils.echo
```

`slurm` now also offers a `docker` builder that builds and pushes a patched image, renders `srun --container-image=<image>`, and leaves the `img_root` macro (the image's root path, substituted into args and env) empty inside the container.

What is unchanged:
- `local_cwd` still has no workspace support.
- `local_docker` and `kubernetes` still use `DockerWorkspaceMixin` and add no `workspace_type` option.
- `slurm` with `workspace_type` unset still copies into `job_dir`; the rendered `sbatch` script stays byte-identical.

Closes #590

Test Plan:
1. `pytest` across `torchx/workspace/test`, `slurm_scheduler_test`, `api_test`, and `workspace_defaults_test` -> 130 tests passed.
2. `ruff check` (import sorting included) and `ruff format --check` with the pinned ruff -> clean.
3. `python -m doctest torchx/workspace/api.py` -> passed; `pyrefly check` -> no findings in the changed files.
4. Downstream plugin suite with torchx installed editable at this commit: see the pull request
@kiukchung
kiukchung force-pushed the kiuk/torchx-multi-workspace-20260919 branch from 58cc504 to 3b0ea7d Compare September 19, 2026 22:18
@kiukchung
kiukchung merged commit 5314a6c into main Sep 19, 2026
10 checks passed
@kiukchung
kiukchung deleted the kiuk/torchx-multi-workspace-20260919 branch September 19, 2026 22:22
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

CLA Signed This label is managed by the Facebook bot. Authors need to sign the CLA before a PR can be reviewed.

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Support multiple workspace types per scheduler

1 participant