Repository navigation
(feature)[torchx][workspace] Schedulers can offer multiple workspace builders - #1458
Merged
Merged
Conversation
kiukchung
force-pushed
the
kiuk/torchx-multi-workspace-20260919
branch
from
September 19, 2026 21:37
c6e0f02 to
58cc504
Compare
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
force-pushed
the
kiuk/torchx-multi-workspace-20260919
branch
from
September 19, 2026 22:18
58cc504 to
3b0ea7d
Compare
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary:
A scheduler can now offer more than one workspace builder because
MultiWorkspaceMixinselects a builder withworkspace_typeinstead of hard-binding a single mix-in.torchx run -s my_backend --workspace . \ -cfg workspace_type=docker,image_repo=example.com/repo utils.echoslurmnow also offers adockerbuilder that builds and pushes a patched image, renderssrun --container-image=<image>, and leaves theimg_rootmacro (the image's root path, substituted into args and env) empty inside the container.What is unchanged:
local_cwdstill has no workspace support.local_dockerandkubernetesstill useDockerWorkspaceMixinand add noworkspace_typeoption.slurmwithworkspace_typeunset still copies intojob_dir; the renderedsbatchscript stays byte-identical.Closes #590
Test Plan:
pytestacrosstorchx/workspace/test,slurm_scheduler_test,api_test, andworkspace_defaults_test-> 130 tests passed.ruff check(import sorting included) andruff format --checkwith the pinned ruff -> clean.python -m doctest torchx/workspace/api.py-> passed;pyrefly check-> no findings in the changed files.