feat: harden distributed execution and add Python 3.14 support

.github/actions/setup-python-env/action.yml

@@ -26,5 +26,5 @@ runs:
         cache-suffix: ${{ matrix.python-version }}
 
     - name: Install Python dependencies
-      run: uv sync --frozen
+      run: uv sync --frozen --extra pytorch
       shell: bash

.github/workflows/main.yml

@@ -29,7 +29,7 @@ jobs:
     runs-on: self-hosted
     strategy:
       matrix:
-        python-version: ["3.9", "3.10", "3.11", "3.12", "3.13"]
+        python-version: ["3.9", "3.10", "3.11", "3.12", "3.13", "3.14"]
       fail-fast: false
     defaults:
       run:
@@ -44,7 +44,11 @@ jobs:
           python-version: ${{ matrix.python-version }}
 
       - name: Run tests
-        run: uv run python -m pytest tests --cov --cov-config=pyproject.toml --cov-report=xml
+        run: uv run --frozen --extra pytorch python -m pytest tests --cov --cov-config=pyproject.toml --cov-report=xml
+        env:
+          # Ray 2.58+ otherwise creates worker environments from the bare
+          # `uv run` command and drops the optional PyTorch dependencies.
+          RAY_ENABLE_UV_RUN_RUNTIME_ENV: "0"
 
       # - name: Check typing
         # run: uv run mypy

AGENTS.md

@@ -13,14 +13,15 @@ This repo uses `uv` for dependency management (creates `.venv/` and maintains `u
 - `make install`: create/sync env and install `pre-commit` hooks.
 - `make check`: verify lockfile consistency and run `pre-commit` (ruff/black/isort, etc.).
 - `make test`: run unit tests with coverage (`pytest --cov`, emits `coverage.xml`).
-- `tox`: run tests across supported Python versions and run `mypy` (matches CI expectations).
+- `tox`: run tests across supported Python versions. `mypy` is currently disabled but its dependency and configuration
+  are retained for future re-enablement.
 - `make docs` / `make docs-test`: serve or build MkDocs docs.
 - `make build`: build a wheel into `dist/` (for release preparation).
 
 ## Coding Style & Naming Conventions
 - Python 3.9+; 4-space indentation; line length target is 120.
 - Formatting/linting is enforced via `pre-commit`: Black + isort + Ruff (auto-fix enabled).
-- Prefer explicit names and type hints; `mypy` is part of the `tox` run and untyped defs are disallowed.
+- Prefer explicit names and type hints; static type checking is currently disabled and may be re-enabled later.
 
 ## Development Principles
 - 如无必要，勿增实体: do not add new concepts, state, helpers, layers, or dependencies unless they are needed to solve the concrete problem.

CONTRIBUTING.md

@@ -65,7 +65,7 @@ cd TinyExp
 Then, install and activate the environment with:
 
 ```bash
-uv sync
+make install-pytorch
 ```
 
 4. Install pre-commit to run linters/formatters at commit time:

Makefile

@@ -1,9 +1,17 @@
 .PHONY: install
-install: ## Install the virtual environment and install the pre-commit hooks
+install: ## Install core dependencies while preserving machine-selected optional packages
 	@echo "🚀 Creating virtual environment using uv"
-	@uv sync
+	@uv sync --locked --inexact
 	@uv run pre-commit install
 
+.PHONY: install-pytorch
+install-pytorch: ## Install the default PyPI PyTorch, TorchVision, and Accelerate builds
+	@uv sync --locked --extra pytorch
+
+.PHONY: install-without-pytorch
+install-without-pytorch: ## Install core dependencies while preserving machine-selected accelerator packages
+	@uv sync --locked --no-extra pytorch --inexact
+
 .PHONY: check
 check: ## Run code quality tools.
 	@echo "🚀 Checking lock file consistency with 'pyproject.toml'"
@@ -16,7 +24,7 @@ check: ## Run code quality tools.
 .PHONY: test
 test: ## Test the code with pytest
 	@echo "🚀 Testing code: Running pytest"
-	@uv run python -m pytest --cov --cov-config=pyproject.toml --cov-report=xml
+	@RAY_ENABLE_UV_RUN_RUNTIME_ENV=0 uv run python -m pytest --cov --cov-config=pyproject.toml --cov-report=xml
 
 .PHONY: build
 build: clean-build ## Build wheel file

README.md

@@ -51,7 +51,7 @@ For a longer explanation, see [`docs/philosophy.md`](docs/philosophy.md).
 ### Option A: Install with pip and use import-based entrypoint
 
 ```bash
-pip install tinyexp
+pip install "tinyexp[pytorch]"
 ```
 
 ```python
@@ -71,7 +71,7 @@ python your_exp.py dataloader_cfg.train_batch_size_per_device=16
 ```bash
 git clone https://github.com/HKUST-SAIL/tinyexp.git
 cd tinyexp
-make install
+make install-pytorch
 uv run python tinyexp/examples/mnist_exp.py
 ```
 
@@ -149,6 +149,11 @@ tinyexp-run-with-redis -- python your_exp.py redis_cfg.redis_cache_enabled=true
 another Redis process, startup fails without shutting down or taking ownership of that server. Connect to externally
 managed Redis directly through `redis_cfg` instead of wrapping the command with `tinyexp-run-with-redis`.
 
+For multi-node training, the helper's Redis lifecycle follows the local command: each wrapper stops the Redis
+resources it owns as soon as its child exits. If one node fails, the whole distributed training job is expected to
+fail and restart; the helper does not keep Redis alive for a global finish barrier or implement heartbeat/lease-based
+failure coordination. The external launcher or supervisor owns whole-job restart and termination.
+
 ## Example Experiments
 
 - MNIST baseline: [`tinyexp/examples/mnist_exp.py`](tinyexp/examples/mnist_exp.py)
@@ -181,12 +186,14 @@ model preparation, reduction, synchronization, and cleanup methods.
 
 ## Development
 
-Install environment and hooks:
+Install the core environment and hooks:
 
 ```bash
 make install
 ```
 
+`make install` installs the core environment and hooks without selecting or removing optional accelerator packages. For the default PyPI stack, run `make install-pytorch` before `make test`. On a machine with a preselected CUDA, ROCm, or vendor PyTorch build, `make install` (or the more explicit `make install-without-pytorch`) preserves that environment; install `torch`, `torchvision`, and `accelerate` together according to that machine's package index/backend, then use ordinary `uv run`. Because the PyTorch packages are optional and ordinary `uv run` does not remove extraneous packages by default, no repeated `--no-sync` flag is needed. Avoid `uv sync` without `--inexact` in that environment.
+
 Run checks:
 
 ```bash

docs/running-modes.md

@@ -40,10 +40,32 @@ make install
 Or install the published package with:
 
 ```bash
-pip install tinyexp
+pip install "tinyexp[pytorch]"
 ```
 
-TinyExp declares PyTorch, Ray, Accelerate, and their Python-level dependencies, but it does not choose a CUDA-specific PyTorch wheel or manage GPU drivers. Follow the PyTorch installation guidance for the target machine when a GPU build is required.
+TinyExp keeps Ray and its Python-level dependencies in the core installation. PyTorch, TorchVision, and Accelerate are provided by the optional `pytorch` extra. This avoids pretending that Python package metadata can choose a compatible CPU, CUDA, ROCm, or vendor-specific build for every machine. The universal `uv.lock` records Python- and operating-system-aware resolutions for the extra; accelerator runtime selection remains machine-specific.
+
+For the default PyPI PyTorch build, use:
+
+```bash
+make install-pytorch
+# or: pip install "tinyexp[pytorch]"
+```
+
+This convenience extra resolves the default PyPI builds recorded in `uv.lock`; it is not a universal CUDA/driver compatibility choice. For a machine that needs a machine-specific CUDA, ROCm, or vendor build, install the base project without the extra, then install all accelerator packages together using the target machine's instructions:
+
+```bash
+make install-without-pytorch
+# Install torch/torchvision from the machine-specific index/build, while
+# keeping PyPI available for accelerate and other regular packages.
+uv pip install --python .venv/bin/python torch torchvision accelerate \
+  --index <pytorch-index-url> --default-index https://pypi.org/simple
+# Or use uv's PyTorch backend selector where supported, e.g.:
+# uv pip install --python .venv/bin/python torch torchvision accelerate --torch-backend cu126
+uv run python -c "import torch; print(torch.__version__, torch.cuda.is_available())"
+```
+
+Python 3.14 requires PyTorch/TorchVision builds that publish compatible `cp314` wheels; this is an ABI requirement, not a CUDA-version requirement. The PyTorch packages are optional, so ordinary `uv run` keeps a machine-selected accelerator build in place and does not require a repeated `--no-sync` flag. Use `make install` (which is intentionally non-exact) on such a machine, and avoid exact `uv sync` without `--inexact` afterward.
 
 Check the active environment before launching:
 
@@ -93,13 +115,17 @@ Requirements:
 
 - The `ray` Python package must import successfully.
 - The machine must have enough resources for every worker bundle.
-- `ray_cfg.ray_num_worker` must be `-1` or a positive integer. `-1` fills the available CPU or GPU capacity.
+- `ray_cfg.ray_num_worker` must be `-1` or a positive integer. `-1` sizes workers from the CPU/GPU resources
+  available when the run starts.
 - `ray_cfg.ray_num_cpus_per_worker` must be positive.
 - `ray_cfg.ray_num_gpus_per_worker` must be non-negative.
 - `ray_cfg.ray_placement_timeout_s` controls how long TinyExp waits for the placement group; the default is 120 seconds.
-- Requests that exceed the cluster's total CPU or GPU capacity fail before placement starts. If the total capacity is sufficient but currently busy, placement waits up to the configured timeout.
+- Requests that exceed the cluster's total CPU or GPU capacity fail before placement starts. If the total capacity is
+  sufficient but currently busy, placement waits up to the configured timeout; a timeout reports total and currently
+  available CPU/GPU resources.
 - GPU workers require a CUDA-enabled PyTorch installation and visible GPUs.
 - TinyExp reads the placement-group bundle-to-node topology before creating Ray worker actors, then derives `RANK` and `LOCAL_RANK` from that topology. Ray workers must be homogeneous: every participating node must host the same number of workers, so every node has the same local-rank range. When bundles are interleaved across nodes, global ranks are reassigned so ranks remain contiguous within each node (for example, node-local workers receive `0..N-1`, then the next node receives the following range).
+- For multi-worker runs, TinyExp starts a zero-CPU TCPStore actor in placement-group bundle 0. The actor binds and holds a dynamic port before worker creation, and all ranks connect to it as clients. The Ray head therefore does not need to host rank 0 or provide a GPU.
 
 If `RAY_ADDRESS` already points to a reachable Ray cluster, `ray.init()` can attach to that cluster. Otherwise, Ray starts a local runtime.
 
@@ -241,12 +267,14 @@ Multi-node requirements:
 
 - `ray` and Python executables must be available on every node. Use `--ray-bin` and `--python-bin` when they are not on `PATH`.
 - The helper owns the Ray runtime on each participating node: it runs `ray stop --force` before startup and during cleanup. Do not use it on a node whose existing Ray runtime must remain active.
+- In multi-node mode, `--node-rank` must be in `[0, --node-count)`, `--ray-port` must be a fixed port in `1..65535`, and `--wait-timeout` bounds Ray head/node readiness.
 - Each node must use compatible Python, TinyExp, PyTorch, Ray, and experiment dependency versions.
 - Custom experiment modules must be importable on every node. The helper does not distribute source code or create environments.
 - Dataset paths used by a scheduled worker must exist on that worker, either through a shared filesystem or equivalent per-node data layout.
 - The head address must be reachable from every worker and must not resolve to loopback for a multi-node job.
 - Firewalls and security groups must allow Ray's head port and Ray's node-to-node runtime traffic. The head port defaults to `6379`; choose an unused `--ray-port` when that port is occupied. The selected head port must also be outside Ray's configured worker port range. Dashboard, metrics, client, and other Ray runtime ports may also need explicit network policy.
 - Aggregate Ray cluster resources must satisfy the requested placement group. A job waits if resources exist in theory but cannot be placed with the requested bundle shape or placement strategy.
+- The Ray head may be CPU-only. GPU worker bundles and the distributed TCPStore are placed on eligible worker nodes according to the resolved placement group.
 
 When `--node-count=1`, the helper executes the command unchanged and does not start a static Ray cluster. A command with `launcher=ray` may still start a local Ray runtime itself.
 
@@ -264,9 +292,25 @@ Redis and W&B requirements are independent of the four launch styles.
 - Enabling Redis cache requires a reachable Redis service.
 - TinyExp-managed standalone Redis requires `redis-server` on the host that starts it.
 - TinyExp-managed Redis Cluster also requires `redis-cli` on the coordinating host.
+- Ray-managed Redis Cluster startup and readiness use `redis_cfg.redis_cluster_startup_timeout_s` (30 seconds by
+  default). The `tinyexp-run-with-redis` wrapper uses its `--wait-timeout` option for the multi-node registration,
+  cluster creation, and readiness deadline.
 - The ResNet example enables Redis cache by default. Set `redis_cfg.redis_cache_enabled=false` when Redis is not installed or desired.
 - Enabling W&B requires suitable credentials and network access, or an explicitly configured offline mode.
 
+### Redis helper lifecycle in multi-node jobs
+
+`tinyexp-run-with-redis` is a per-process wrapper, not a long-lived Redis service manager. It starts Redis resources
+for the command it launches and stops the resources owned by that wrapper as soon as the child command exits, whether
+the child succeeds or fails. In multi-node mode, the HTTP rendezvous is used only to register nodes and create the
+Redis Cluster during startup; it is not a finish barrier.
+
+When a multi-node training process fails, the distributed training job is considered failed and must be restarted as
+a whole by the external process supervisor or launcher. The supervisor must terminate the remaining wrappers; their
+signal handlers stop their child processes and destroy the Redis resources they own. TinyExp deliberately does not add
+heartbeat, lease, or a global terminal-state protocol here: Redis exists only for the lifetime of the command on each
+node, and a failed multi-node job is expected to be torn down rather than kept alive for coordination.
+
 Check Redis system commands when cache management is enabled:
 
 ```bash

pyproject.toml

@@ -17,23 +17,42 @@ classifiers = [
     "Programming Language :: Python :: 3.11",
     "Programming Language :: Python :: 3.12",
     "Programming Language :: Python :: 3.13",
+    "Programming Language :: Python :: 3.14",
     "Topic :: Software Development :: Libraries :: Python Modules",
 ]
 dependencies = [
     "loguru>=0.7.3",
-    "numpy",
+    # NumPy releases before 2.3.2 do not publish CPython 3.14 wheels.
+    "numpy; python_version < '3.14'",
+    "numpy>=2.3.2; python_version >= '3.14'",
     "psutil",
     "tabulate",
     "hydra-core",
     "omegaconf",
-    "ray",
+    # Ray 2.53+ does not publish a CPython 3.9 wheel. Keep the older
+    # interpreter resolution installable while using the current line for
+    # Python versions that have compatible wheels, including 3.14.
+    "ray<2.52; python_version < '3.10'",
+    "ray; python_version >= '3.10'",
+    # Ray pulls rpds-py through jsonschema; use a release with CPython 3.14 wheels.
+    "rpds-py>=0.26; python_version >= '3.14'",
     "tqdm",
-    "torch",
-    "torchvision",
     "redis",
     "wandb",
     "loguru",
+]
+
+[project.optional-dependencies]
+pytorch = [
     "accelerate",
+    # Python 3.14 requires releases that publish cp314 wheels. The CPU/CUDA/
+    # ROCm/vendor runtime is selected separately for each target machine.
+    "torch; python_version < '3.14'",
+    "torch>=2.13; python_version >= '3.14'",
+    "torchvision; python_version < '3.14'",
+    "torchvision>=0.28; python_version >= '3.14'",
+    # TorchVision pulls Pillow; use a release with CPython 3.14 wheels.
+    "pillow>=11.3; python_version >= '3.14'",
 ]
 
 [project.urls]
@@ -80,15 +99,16 @@ include = [
 
 
 
-[tool.mypy]
-files = ["tinyexp"]
-disallow_untyped_defs = true
-disallow_any_unimported = true
-no_implicit_optional = true
-check_untyped_defs = true
-warn_return_any = true
-warn_unused_ignores = true
-show_error_codes = true
+# [tool.mypy]
+# Temporarily disabled. Uncomment this section when mypy is enabled again.
+# files = ["tinyexp"]
+# disallow_untyped_defs = true
+# disallow_any_unimported = true
+# no_implicit_optional = true
+# check_untyped_defs = true
+# warn_return_any = true
+# warn_unused_ignores = true
+# show_error_codes = true
 
 [tool.pytest.ini_options]
 testpaths = ["tests"]

tests/conftest.py

@@ -11,10 +11,12 @@ def ray_session():
     if not ray.is_initialized():
         runtime_env = {
             "working_dir": ".",
-            "excludes": ["*.md", "data/", "tests/", ".git/", ".venv/", "output/", "outputs/", "site/"],
+            "excludes": ["data/", "tests/", ".git/", ".venv/", "output/", "outputs/", "site/"],
         }
         try:
-            ray.init(runtime_env=runtime_env)
+            # Force a fresh local runtime instead of auto-connecting to a
+            # stale cluster address left by an interrupted Ray test run.
+            ray.init(address="local", runtime_env=runtime_env)
         except Exception as exc:
             pytest.skip(f"Ray is not available in this environment: {exc}")
 

tests/dataset/test_sampler.py

@@ -3,6 +3,7 @@
 from types import SimpleNamespace
 
 import pytest
+from torch.utils.data import DataLoader, Dataset
 
 from tinyexp.dataset.fake_dataloader import HoldOnesampleDataLoader
 from tinyexp.dataset.sampler import InfiniteSampler
@@ -83,6 +84,64 @@ def test_infinite_sampler_no_shuffle_multi_worker_slices() -> None:
     assert first_six == [1, 3, 1, 3, 1, 3]
 
 
+def test_infinite_sampler_reuses_one_dataloader_iterator_across_epochs() -> None:
+    class IndexDataset(Dataset[int]):
+        def __len__(self) -> int:
+            return 4
+
+        def __getitem__(self, index: int) -> int:
+            return index
+
+    class CountingInfiniteSampler(InfiniteSampler):
+        iterator_count = 0
+
+        def __iter__(self):
+            type(self).iterator_count += 1
+            return super().__iter__()
+
+    sampler = CountingInfiniteSampler(size=4, shuffle=False, seed=0)
+    dataloader = DataLoader(IndexDataset(), batch_size=2, sampler=sampler, num_workers=0)
+    iterator = iter(dataloader)
+
+    epochs = [[next(iterator).tolist() for _ in range(len(dataloader))] for _ in range(3)]
+
+    assert CountingInfiniteSampler.iterator_count == 1
+    assert epochs == [
+        [[0, 1], [2, 3]],
+        [[0, 1], [2, 3]],
+        [[0, 1], [2, 3]],
+    ]
+
+
+def test_infinite_sampler_resumes_at_epoch_before_iterator_creation() -> None:
+    class IndexDataset(Dataset[int]):
+        def __len__(self) -> int:
+            return 4
+
+        def __getitem__(self, index: int) -> int:
+            return index
+
+    sampler = InfiniteSampler(size=4, shuffle=True, seed=17)
+    dataloader = DataLoader(IndexDataset(), batch_size=2, sampler=sampler, num_workers=0)
+    iterator = iter(dataloader)
+    first_epoch = [next(iterator).tolist() for _ in range(len(dataloader))]
+    second_epoch = [next(iterator).tolist() for _ in range(len(dataloader))]
+
+    resumed_sampler = InfiniteSampler(size=4, shuffle=True, seed=17)
+    resumed_sampler.set_epoch(1)
+    resumed_dataloader = DataLoader(
+        IndexDataset(),
+        batch_size=2,
+        sampler=resumed_sampler,
+        num_workers=0,
+    )
+    resumed_iterator = iter(resumed_dataloader)
+    resumed_epoch = [next(resumed_iterator).tolist() for _ in range(len(resumed_dataloader))]
+
+    assert first_epoch != second_epoch
+    assert resumed_epoch == second_epoch
+
+
 def test_infinite_sampler_set_epoch_continues_fixed_stream() -> None:
     expected_sampler = InfiniteSampler(size=10, shuffle=True, seed=17)
     expected_stream = iter(expected_sampler)