diff --git a/AGENTS.md b/AGENTS.md index e515e29..9630bd4 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -1,40 +1,55 @@ # Repository Guidelines -## Project Structure & Module Organization -- `tinyexp/`: main library code (experiment entrypoints, datasets, accelerators, utilities). -- `tinyexp/examples/`: runnable example experiments (e.g., MNIST, ResNet). -- `tests/`: `pytest` suite (mirrors package areas like `tests/utils/` and `tests/tine_engine/`). -- `docs/` + `mkdocs.yml`: MkDocs documentation sources/config. -- `data/`: local datasets used by examples (e.g., `data/MNIST/`). -- `output/`: generated run artifacts/logs—avoid committing large files. +## Agent Quick Reference +Use this index before reading source: the docstrings below are written to be sufficient. Names in backticks are +importable; print a module's with `python -c "import tinyexp; print(tinyexp.__doc__)"` and a class's with +`python -c "from tinyexp.exp_mixins import basic_mixins as m; print(m.CheckpointCfgMixin.CheckpointCfg.__doc__)"`. -## Build, Test, and Development Commands -This repo uses `uv` for dependency management (creates `.venv/` and maintains `uv.lock`). -- `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. `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). +| Task | Read only | +| --- | --- | +| Run or pick a bundled example | `tinyexp/examples/README.md`, then the docstring (run header) of `tinyexp.examples.` | +| Write a new experiment | `tinyexp` (minimal example, fields) and `tinyexp.tiny_engine.accelerator.base_accelerator` (training loop); full training recipe, only if needed: `tinyexp/examples/mnist_exp.py` | +| Config fields, `launcher` pairing, errors; merge rules | `tinyexp`, `tinyexp.exceptions`; `tinyexp.TinyExp.set_cfg` | +| Ray workers (contract, minimal example), multi-node cluster, torchrun, accelerate | `RayCfgMixin` and `RayCfgMixin.RayCfg` in `tinyexp.exp_mixins.basic_mixins`, then one section of `docs/running-modes.md` | +| Checkpoint contents, save/load/resume/eval | `CheckpointCfgMixin.CheckpointCfg` in `tinyexp.exp_mixins.basic_mixins` | +| Logging, W&B, Redis cache | `tinyexp.exp_mixins.basic_mixins` (table), then `LoggerCfgMixin.LoggerCfg`, `WandbCfgMixin.WandbCfg` or `RedisCfgMixin.RedisCfg`; Redis wrapper: "Optional Services" in `docs/running-modes.md` | +| Pick an accelerator; training-loop API | `tinyexp.tiny_engine.accelerator`, `tinyexp.tiny_engine.accelerator.base_accelerator` | +| Infinite sampler, fake dataloader, EMA | `tinyexp.dataset`, `tinyexp.utils` | +| Change one config value | nothing: `python -m mode=help key=value` prints the resolved config and the overridden values, then exits | -## 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; static type checking is currently disabled and may be re-enabled later. +`mode=help` checks key names and value types, not whether a value is supported (e.g. an unknown accelerator name +passes), and it still creates `output_root`. If a run fails with `FileExistsError: ... 'output'`, the local +`./output` is unusable (e.g. a broken symlink): pass `output_root=/tmp/tinyexp` instead of editing code. -## Development Principles -- 如无必要，勿增实体: do not add new concepts, state, helpers, layers, or dependencies unless they are needed to solve the concrete problem. -- Keep It Simple, Stupid (KISS): prefer the smallest direct implementation that preserves behavior and is easy to reason about. +Where to make a change, in order of preference: a CLI override; a new experiment file or a subclass of an existing +`Exp`; the defaults of an example in `tinyexp/examples/`; the library (`tinyexp/` outside `examples/`), only when +the behavior itself must change. -## Testing Guidelines -- Framework: `pytest`; place tests under `tests/` and name files `test_*.py`. -- Keep tests fast and deterministic (CPU-first). If a feature needs external services (e.g., Redis/W&B), gate via env vars and provide safe fallbacks/mocks. +Do not explore: +- `output/`, `data/`: run artifacts and datasets (may be large or broken symlinks). Derive checkpoint paths as + `//{last,best}.ckpt` instead of searching; never glob `**/*.ckpt` or `**/*.pt`. +- `local/`: the maintainer's private experiments, git-ignored and not library code. Do not read, import, modify, + or use it as a reference for new experiments. +- `.venv/`, `site/`, `__pycache__/`, `.ruff_cache/`, `.pytest_cache/`. -## Commit & Pull Request Guidelines -- Commit history follows Conventional Commits (`feat:`, `fix:`) and often includes an emoji. - Example: `feat: 🎸 add new accelerator` -- PRs should include: a clear description, linked issue (if applicable), tests for behavior changes, and doc updates (`README.md`/`docs/`) when user-facing behavior changes. +Before launching GPU or Ray jobs (anything beyond `mode=help` or CPU tests), confirm with the user. -## Agent-Specific Notes (Optional) -- Prefer `make` targets over ad-hoc commands and keep patches focused and reviewable. +## Development +- Layout: `tinyexp/` library, `tinyexp/examples/` runnable experiments, `tests/` mirroring the package + (`tests/utils/`, `tests/tiny_engine/`; the TP accelerator is tested with its example in + `tests/examples/test_vit_tp_exp.py`), `docs/` + `mkdocs.yml` for MkDocs. Do not commit large files from `output/`. +- `uv` manages `.venv/` and `uv.lock`. Prefer `make` targets: `make install` (env + `pre-commit` hooks), + `make check` (lockfile + `pre-commit`: ruff/black/isort), `make test` (`pytest --cov`, writes `coverage.xml`), + `make docs` / `make docs-test` (serve / build MkDocs), `make build` (wheel into `dist/`). `tox` runs the supported + Python versions. `mypy` is disabled; its dependency and configuration are kept for re-enablement. +- Style: Python 3.9+, 4-space indentation, line length 120, explicit names and type hints. +- 如无必要，勿增实体: do not add new concepts, state, helpers, layers, or dependencies unless they are needed to solve + the concrete problem. Keep It Simple, Stupid (KISS): prefer the smallest direct implementation that preserves + behavior and is easy to reason about. Keep patches focused and reviewable. +- Tests: `pytest`, files `tests/**/test_*.py`; fast and deterministic (CPU-first). Gate external services + (Redis/W&B) behind env vars and provide safe fallbacks/mocks. +- Commits follow Conventional Commits and often include an emoji (`feat: 🎸 add new accelerator`). PRs include a + clear description, linked issue (if any), tests for behavior changes, and `README.md`/`docs/` updates when + user-facing behavior changes. +- Changing docstrings, run headers, the docs, or the package layout: first read "Documentation for Coding Agents" + in `docs/index.md` (what the tests check, and the invariants the docs rely on). diff --git a/README.md b/README.md index ccf90a0..969d392 100644 --- a/README.md +++ b/README.md @@ -104,8 +104,8 @@ python -m tinyexp.examples.mnist_exp \ ``` `mode=help` prints the resolved configuration and exits without starting workers. The other mode values are -experiment-specific; for example, MNIST and ResNet provide `train`/`val`, pi provides `run`, and the DeiT-S example -also provides `eval`/`bench`. +experiment-specific; for example, MNIST and ResNet provide `train`/`eval`, pi provides `run`, and the DeiT-S example +provides `train`/`eval`/`bench`. ```bash python -m tinyexp.examples.mnist_exp mode=help @@ -171,13 +171,15 @@ Accelerate, Ray multi-node, network, data, Redis, and W&B requirements. Ray worker resources are explicit. Set `ray_cfg.ray_num_cpus_per_worker` and `ray_cfg.ray_num_gpus_per_worker` to the resources required by one worker; both fields can be overridden from the command line. -Run an ImageNet experiment with TinyExp's Redis helper after installing the package: +Launches without Ray (`launcher=mp`: plain `python`, `torchrun`, `accelerate launch`) can use TinyExp's Redis helper +to own the cache. With `launcher=ray` (the ResNet default) the Ray driver starts and stops Redis itself, and the +helper refuses such a command (exit code 2) before starting anything. ```bash export IMAGENET_HOME=/path/to/imagenet tinyexp-run-with-redis -- \ - python -m tinyexp.examples.resnet_exp \ - redis_cfg.redis_cache_enabled=true + torchrun --nproc-per-node=8 -m tinyexp.examples.resnet_exp \ + launcher=mp redis_cfg.redis_cache_enabled=true ``` `tinyexp-run-with-redis` owns and stops only the Redis processes it starts. If a configured port is already served by @@ -191,45 +193,24 @@ failure coordination. The external launcher or supervisor owns whole-job restart ## Example Experiments -- MNIST baseline: [`tinyexp/examples/mnist_exp.py`](tinyexp/examples/mnist_exp.py) -- ImageNet ResNet-50: [`tinyexp/examples/resnet_exp.py`](tinyexp/examples/resnet_exp.py) -- Distributed Monte Carlo pi (non-DL, `mode=run`): [`tinyexp/examples/pi_exp.py`](tinyexp/examples/pi_exp.py) -- DeiT-S with tensor parallelism: [`tinyexp/examples/vit_tp_exp.py`](tinyexp/examples/vit_tp_exp.py), documented in - [`docs/vit_tp.md`](docs/vit_tp.md) +- MNIST baseline and template for new experiments (CPU or DDP, checkpoints, resume): + [`tinyexp/examples/mnist_exp.py`](tinyexp/examples/mnist_exp.py) +- Distributed Monte Carlo pi, a non-deep-learning workload (`mode=run`) and the fastest smoke test: + [`tinyexp/examples/pi_exp.py`](tinyexp/examples/pi_exp.py) +- ImageNet ResNet-50 with DDP (about 76% top-1 after 90 epochs), FSDP or HF Accelerate, and the Redis byte cache: + [`tinyexp/examples/resnet_exp.py`](tinyexp/examples/resnet_exp.py) +- DeiT-S with tensor parallelism and a TP-vs-DDP benchmark: + [`tinyexp/examples/vit_tp_exp.py`](tinyexp/examples/vit_tp_exp.py), documented in [`docs/vit_tp.md`](docs/vit_tp.md) -Run the ImageNet ResNet-50 example with the dataset root in `IMAGENET_HOME`: +Each example's module docstring starts with a run header — `Modes`, `Run`, `Env`, `Needs`, `Out` — listing its +commands, required data and hardware. The headers are the single source of example commands and are checked by the +test suite. Print one without opening the source: ```bash -export IMAGENET_HOME=/path/to/imagenet -python -m tinyexp.examples.resnet_exp -``` - -The same ResNet training loop also supports FSDP for a parameter/gradient/optimizer-state -sharding smoke test. Run it with `accelerator_cfg.accelerator=fsdp` on a multi-GPU CUDA host; -the recipe and validation target are otherwise the same as DDP. The existing DDP recipe reaches -about 76% ImageNet top-1 after the full 90-epoch run; use the command below to compare FSDP on -the same setup: - -```bash -torchrun --nproc-per-node=2 tinyexp/examples/resnet_exp.py \ - launcher=mp accelerator_cfg.accelerator=fsdp redis_cfg.redis_cache_enabled=false -``` - -The pi example uses Ray workers to all-reduce sample counts and does not use a dataloader: - -```bash -python -m tinyexp.examples.pi_exp \ - pi_cfg.total_samples=100000000 \ - ray_cfg.ray_num_worker=4 -``` - -The DeiT-S example includes a synthetic TP throughput/memory benchmark on two GPUs: - -```bash -python -m tinyexp.examples.vit_tp_exp mode=bench +python -c "import tinyexp.examples.resnet_exp as m; print(m.__doc__)" ``` -See [`docs/vit_tp.md`](docs/vit_tp.md) for ImageNet data preparation, evaluation, and full training commands. +See [`tinyexp/examples/README.md`](tinyexp/examples/README.md) for choosing an example. ## How It Works diff --git a/SKILL.md b/SKILL.md index 1e92a14..7bcf7cb 100644 --- a/SKILL.md +++ b/SKILL.md @@ -1,72 +1,46 @@ --- name: tinyexp-experiments -description: Build, run, debug, and scale PyTorch or general experiments with this TinyExp repository. Use for new experiment files, configuration overrides, CPU/GPU training, Ray jobs, and distributed launches. +description: Write, configure and launch experiments with the TinyExp library (a dataclass Exp per file, CLI overrides, CPU/DDP/FSDP training, Ray workers, torchrun/accelerate). Use when the task involves tinyexp or a class derived from TinyExp. --- # TinyExp experiments -Use TinyExp when the experiment should keep its configuration and execution in one Python entrypoint. Read the nearest example in `tinyexp/examples/` before adding a new abstraction. - -## Core pattern - -1. Define `@dataclass class Exp(TinyExp, ...)`. -2. Put model, data, optimizer, scheduler, and runtime settings in nested dataclasses with `field(default_factory=...)`. -3. Implement `run()` and use an accelerator implementing `AcceleratorProtocol` (`prepare`, `backward`, `reduce_*`, `wait_for_everyone`, `destroy`). -4. End the file with `if __name__ == "__main__": store_and_run_exp(Exp)`. - -Keep training logic in the experiment. Reuse mixins such as `RayCfgMixin`, `CheckpointCfgMixin`, `LoggerCfgMixin`, `WandbCfgMixin`, and `RedisCfgMixin` only when the feature is needed. Save outputs below `output_root` (default `./output`). - -## Configuration and launch - -Hydra accepts dotted CLI overrides without editing code. First inspect a new experiment: - -```bash -python -m tinyexp.examples.mnist_exp mode=help -``` - -`launcher` selects TinyExp's process owner: `mp` runs `Exp.run()` in the current process; `ray` makes the driver create Ray workers from `ray_cfg`. `torchrun`, `accelerate launch`, and `tinyexp-run-with-ray-cluster` are external launchers and must be paired with `launcher=mp`, `launcher=mp`, and `launcher=ray` respectively. - -Typical commands: +An experiment is one `@dataclass class Exp(TinyExp, )` whose `run()` holds the logic; every config field is a +CLI override (`key=value`). The installed package documents itself, so read docstrings instead of source: +`python -c "import tinyexp; print(tinyexp.__doc__)"`, or for a class +`python -c "from tinyexp.exp_mixins import RayCfgMixin as m; print(m.__doc__)"`. + +| Need | Print the `__doc__` of | +| --- | --- | +| New experiment: minimal example, fields, `launcher` pairing | `tinyexp` | +| How overrides are merged and reported | `tinyexp.TinyExp.set_cfg` | +| Training loop API; picking an accelerator | `tinyexp.tiny_engine.accelerator.base_accelerator`; `tinyexp.tiny_engine.accelerator` | +| Code running on Ray workers (minimal example); worker count and resources | `tinyexp.exp_mixins.RayCfgMixin`; `tinyexp.exp_mixins.RayCfgMixin.RayCfg` | +| Checkpoint save / load / resume / eval | `tinyexp.exp_mixins.CheckpointCfgMixin.CheckpointCfg` | +| Logger, W&B, Redis cache mixins | `tinyexp.exp_mixins.basic_mixins` (table, then the nested `XXXCfg`) | +| Infinite sampler, fake dataloader, EMA; errors | `tinyexp.dataset`, `tinyexp.utils`; `tinyexp.exceptions` | +| A bundled example: modes, commands, data, hardware | `tinyexp.examples.` (run header) | + +Bundled examples: `mnist_exp` (template for a training experiment: CPU or DDP, checkpoints, resume), `pi_exp` +(non-deep-learning distributed job, fastest smoke test), `resnet_exp` (ImageNet with DDP, FSDP or HF Accelerate and +a Redis cache), `vit_tp_exp` (tensor parallelism, TP-vs-DDP benchmark). + +Easy to miss: + +- Check a config change before running: `mode=help` prints the resolved config and the overridden values, then + exits. It checks key names and types, not values. +- `torchrun` and `accelerate launch` need `launcher=mp`; plain `python` takes `mp` (one process) or `ray` (TinyExp + starts Ray workers); `tinyexp-run-with-ray-cluster` (multi-node, see `--help`) needs `launcher=ray` and runs + `ray stop --force` on each node, so do not use it where an existing Ray runtime must survive. +- With `launcher=ray`, start with the environment's `python -m package.module`, not `uv run` (or set + `RAY_ENABLE_UV_RUN_RUNTIME_ENV=0`). +- Install with `pip install "tinyexp[pytorch]"`; for a CUDA/ROCm-specific PyTorch, install `tinyexp` without the + extra, then `torch`, `torchvision` and `accelerate` from that backend. ```bash -# Local CPU/debug, one process +python -m tinyexp.examples.pi_exp mode=help pi_cfg.total_samples=100000 python -m tinyexp.examples.pi_exp launcher=mp pi_cfg.total_samples=100000 - -# Local Ray workers (CPU in this example) -python -m tinyexp.examples.pi_exp ray_cfg.ray_num_worker=4 pi_cfg.total_samples=1000000 - -# Single-node GPU/DDP (one process per visible GPU) -torchrun --nproc-per-node=2 -m tinyexp.examples.resnet_exp \ - launcher=mp accelerator_cfg.accelerator=ddp redis_cfg.redis_cache_enabled=false - -# Accelerate owns the processes; TinyExp remains in mp mode -accelerate launch --cpu --num-processes=1 -m tinyexp.examples.pi_exp \ - launcher=mp pi_cfg.total_samples=100000 - -# Existing multi-node Ray cluster -tinyexp-run-with-ray-cluster --node-count=2 --node-rank=0 \ - --head-addr=10.0.0.1 --ray-port=6380 -- \ +torchrun --nproc-per-node=2 -m tinyexp.examples.mnist_exp launcher=mp accelerator_cfg.accelerator=ddp +tinyexp-run-with-ray-cluster --node-count=2 --node-rank=0 --head-addr=10.0.0.1 --ray-port=6380 -- \ python -m tinyexp.examples.pi_exp launcher=ray ``` - -For GPU jobs, verify CUDA visibility and use `DDPAccelerator` with NCCL. For `launcher=ray`, every node needs the same code, Python environment, dependencies, and dataset paths; set `ray_cfg.ray_num_cpus_per_worker` and `ray_cfg.ray_num_gpus_per_worker` for each worker's resources. Do not use the static Ray helper on a node whose existing Ray runtime must be preserved. - -## Examples to reuse - -- `tinyexp/examples/pi_exp.py`: general distributed workload and a fast smoke test (`mode=run`). -- `tinyexp/examples/mnist_exp.py`: CPU or DDP image training, checkpoints, logging, and CLI overrides. -- `tinyexp/examples/resnet_exp.py`: ImageNet training; set `IMAGENET_HOME` or override `dataloader_cfg.data_root`; defaults to DDP and Ray. - -The examples are included in the published package. The bundled examples use PyTorch, so a bare `pip install tinyexp` -does not install all of their dependencies. For a normal install, use the PyTorch extra and run them as modules: - -```bash -python -m pip install "tinyexp[pytorch]" -python -m tinyexp.examples.pi_exp pi_cfg.total_samples=100000 -``` - -If the machine needs a CUDA/ROCm/vendor-specific PyTorch build, install `tinyexp` without the extra and install a compatible `torch`, `torchvision`, and `accelerate` set using that backend's instructions. - -## Development checks - -For repository changes from a source checkout, use `make install-pytorch`, then run `make check` and focused `pytest` tests after changing library behavior. diff --git a/docs/index.md b/docs/index.md index a14c66b..2b97cc9 100644 --- a/docs/index.md +++ b/docs/index.md @@ -22,6 +22,46 @@ different process owners and must be paired with the correct TinyExp `launcher` [Running Modes and Environment Requirements](running-modes.md) for the mode matrix, setup requirements, and verified command patterns. +## Documentation for Coding Agents + +Work that is done by a coding agent reads the docs instead of the source. Three entry points are kept accurate for +that use: + +- `AGENTS.md` at the repository root: a task-to-file index for work inside a checkout; +- `SKILL.md`: an index of the same contracts for an installed package, where `AGENTS.md` and `docs/` are absent; +- the run header at the top of every example module (`Modes`, `Run`, `Env`, `Needs`, `Out`), which lists that + example's commands, data and hardware. Print it with + `python -c "import tinyexp.examples.mnist_exp as m; print(m.__doc__)"` instead of reading the file. + +Their commands are checked by `tests/examples/test_examples_help_smoke.py`: every example must carry a run header, +every command in the run headers, `README.md`, `SKILL.md` and `docs/*.md` must be legal for the config it targets, +and every `Modes` value in a run header must be handled in `run()`. The check is static (key names and types only), +so it catches stale option names, not stale advice. `tests/test_tinyexp_artifacts.py` asserts in a subprocess that +`import tinyexp` loads neither torch nor wandb, and that every public module has a docstring. + +Invariants the docs rely on: + +- keep the `import torch` in `CheckpointCfg` and the `import wandb` in `WandbCfg` lazy (inside methods); +- keep `from tinyexp.exp_mixins.basic_mixins import *Mixin` importable (released path); +- do not move `tinyexp/__init__.py` and do not add a yaml at the package root (Hydra's `pkg://tinyexp` needs the + ConfigStore entry `cfg`); +- code moved out of `examples/` enters coverage, where the patch target is 80%. + +Known gaps, deliberately left open: + +- `RayCfg`, `RedisCfg` and `CheckpointCfg` docstrings do not list every field default; `mode=help` prints them. +- `tinyexp-run-with-redis` takes its ports and memory limit from `redis_cfg.*` overrides on the command after `--`, + and appends the resolved connection settings to that command. +- With `launcher=ray` the driver starts and stops Redis by itself, so `tinyexp-run-with-redis` refuses such a + command; with `launcher=mp` the wrapper owns Redis. +- The MNIST example always trains 3 epochs (not configurable); shorten it with a larger + `dataloader_cfg.train_batch_size_per_device`. +- Reading budget: a fresh agent finishes each multi-feature task without reading the source, but reads about + 10-21k characters of docs; the 12k once targeted is not a bound multi-feature tasks can be held to. The fixed + entry (`AGENTS.md` plus the `tinyexp` module docstring) is about 7.5k, a task spanning several features needs + interface docs beyond the rest, and most of their sentences close a gap an earlier trial hit. Trial-to-trial + noise is large: the same accelerator task read 20.7k in one trial and 12.4k in the next. + ## Design Philosophy TinyExp is intentionally light. diff --git a/docs/modules.md b/docs/modules.md index c3e48b4..679dde8 100644 --- a/docs/modules.md +++ b/docs/modules.md @@ -1 +1,17 @@ +# Modules + +API reference generated from the module docstrings. Each section starts with the module's own summary. + ::: tinyexp + +::: tinyexp.exp_mixins.basic_mixins + +::: tinyexp.tiny_engine.accelerator + +::: tinyexp.tiny_engine.accelerator.base_accelerator + +::: tinyexp.dataset + +::: tinyexp.exceptions + +::: tinyexp.utils diff --git a/docs/running-modes.md b/docs/running-modes.md index e0b5741..1377361 100644 --- a/docs/running-modes.md +++ b/docs/running-modes.md @@ -120,14 +120,11 @@ python -m tinyexp.examples.pi_exp \ pi_cfg.total_samples=100000 ``` -Use the prepared environment's `python` for this mode. When Ray finds `uv run` in the driver process ancestry, its -automatic `uv` integration adds a job-level `working_dir` and a `uv run` worker executable. Ray then packages the -project working directory and starts workers through `uv` instead of directly using the normal Python executable on -each worker node. That behavior is useful when Ray should distribute source and dependencies dynamically, but it is -not the execution model used by TinyExp's examples or static cluster helper, which expect the environment to be -installed on each node already. - -If a surrounding command must use `uv run`, disable Ray's automatic `uv` runtime environment before Ray is imported: +Use the prepared environment's `python` for this mode, not `uv run`. When Ray finds `uv run` among the driver's +ancestors, its `uv` integration packages the working directory and starts workers through `uv` instead of the +environment already installed on each node, which TinyExp's examples and cluster helper expect. How TinyExp was +installed (`pip`, `uv sync`) does not matter; only a `uv run` ancestor does. If the command must go through `uv run`, +disable the integration before Ray is imported: ```bash RAY_ENABLE_UV_RUN_RUNTIME_ENV=0 uv run python -m tinyexp.examples.pi_exp \ @@ -138,12 +135,6 @@ RAY_ENABLE_UV_RUN_RUNTIME_ENV=0 uv run python -m tinyexp.examples.pi_exp \ pi_cfg.total_samples=100000 ``` -The installation tool does not determine whether this hook runs: - -- `pip install "tinyexp[pytorch]"` followed by `python your_exp.py` does not trigger the `uv` integration. -- `uv sync` followed by an activated `.venv` and `python your_exp.py` does not trigger it either. -- A driver started through `uv run` can trigger it even when TinyExp itself was installed with `pip`. - Requirements: - The `ray` Python package must import successfully. @@ -157,14 +148,16 @@ Requirements: 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. +- Ranks are assigned node by node (`RANK` contiguous within a node, `LOCAL_RANK` from 0), so every participating + node must host the same number of workers. +- Multi-worker runs rendezvous through a zero-CPU TCPStore actor in placement-group bundle 0, so the Ray head 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. -Ray forwards worker stdout asynchronously. `RayCfgMixin.get_ray_run_result()` therefore provides a synchronous channel for -one short rank-zero result that must be printed reliably after all workers finish; it returns `None` by default. It is -not intended for streaming logs or large outputs. +Ray forwards worker stdout asynchronously. `RayCfgMixin.get_ray_run_result()` is a synchronous channel for one short +rank-zero result that the driver prints after all workers finish (`None` by default); it is not meant for logs or +large outputs. `RayCfgMixin`'s docstring has the worker-side pattern. The MNIST example uses `seed=42` by default for repeatable model initialization and dropout behavior. Override it explicitly when comparing runs across launchers or cluster topologies. @@ -221,7 +214,7 @@ For multi-node TorchRun, every node needs the same environment, source code, and ## Accelerate Launch -`accelerate launch` is another external process owner. Use `launcher=mp` for the same reason as TorchRun: +`accelerate launch` is another external process owner. Use `launcher=mp` for the same reason as TorchRun. The following portable smoke test verifies the Accelerate entry path with one CPU process: @@ -274,19 +267,12 @@ For multi-machine Accelerate launches, all machines need matching code and envir ### Mixed Precision Every accelerator takes `mixed_precision` (`none` | `fp16` | `bf16`; `CPUAccelerator` is bf16-only) and -validates what it supports. `prepare_model` wraps the module's `forward` in the accelerator's autocast -context and casts the returned tensors back to fp32, matching what `accelerate` does in its own -`prepare_model`. A training loop therefore gets mixed precision whether or not it opens -`accelerator.autocast()` itself, and swapping `HFAccelerator` for -`DDPAccelerator`/`FSDPAccelerator`/`CPUAccelerator` neither silently falls back to fp32 nor changes the -dtypes the loop sees. Opening `accelerator.autocast()` explicitly around a forward pass — as the MNIST and -ResNet examples do — stays correct, since nesting the same dtype is a no-op, and remains the way to cover -work that runs outside `forward`. - -Parameters, gradients, communication, and checkpoints stay fp32; autocast only reduces precision inside -compute kernels. `fp16` additionally enables a `GradScaler`, and `clip_grad_norm_` unscales the gradients -of every optimizer passed through `prepare_optimizer` before measuring the norm, so `max_norm` means the -same thing in all three modes. Under `none` and `bf16` the scaler is disabled and unscaling is a no-op. +validates what it supports. `prepare_model` wraps the module's `forward` in the accelerator's autocast and casts the +returned tensors back to fp32, as `accelerate` does. A loop therefore gets mixed precision with or without its own +`accelerator.autocast()` and sees the same dtypes under every accelerator; opening `accelerator.autocast()` around a +forward pass (as the MNIST and ResNet examples do) is a harmless no-op nesting and still covers work outside `forward`. +Parameters, gradients, communication and checkpoints stay fp32. `fp16` also enables a `GradScaler`; for gradient +clipping under it, see `tinyexp.tiny_engine.accelerator.base_accelerator`. ## Static Ray Cluster @@ -340,7 +326,7 @@ Multi-node requirements: - 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. +When `--node-count=1`, the helper executes the command unchanged: it does not start a static Ray cluster and does not run `ray stop`, so an existing Ray runtime on the node is left alone. A command with `launcher=ray` may still start a local Ray runtime itself. Check every node before launch: @@ -353,7 +339,10 @@ python -c "import tinyexp, torch, ray; print(tinyexp.__file__); print(torch.__ve Redis and W&B requirements are independent of the four launch styles. -- Enabling Redis cache requires a reachable Redis service. +- Enabling Redis cache requires a reachable Redis service. Its owner depends on the launcher: `launcher=ray` starts + and stops Redis itself; launches without Ray (`launcher=mp` under `python`, `torchrun` or `accelerate launch`) + wrap the command with `tinyexp-run-with-redis`. The wrapper refuses a `launcher=ray` command (exit code 2) + before starting anything; it reads the launcher from the command's `launcher=` override or the experiment's default. - 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 diff --git a/docs/vit_tp.md b/docs/vit_tp.md index c6c31f2..727e501 100644 --- a/docs/vit_tp.md +++ b/docs/vit_tp.md @@ -52,6 +52,9 @@ pass. Attention runs per-head locally with no communication inside, using PyTorc bit-identical to `datasets.ImageFolder`; the first epoch fills the cache from disk, later epochs read from memory. `safe_set` failures (e.g. maxmemory full) fall back to disk reads — correctness never depends on the cache. +- **D8** — the train dataloader raises `ValueError` when it has no batch (per-rank samples + < `train_batch_size_per_device` with `drop_last=True`). Upstream would silently train zero + steps and then fail on an empty metric average; valid configurations are unaffected. ## Usage @@ -76,12 +79,14 @@ python -m tinyexp.examples.vit_tp_exp mode=bench python -m tinyexp.examples.vit_tp_exp mode=bench accelerator_cfg.accelerator=ddp ``` -CPU smoke test (no ImageNet, ray 2 workers, gloo TP=2): +CPU smoke test (no ImageNet, ray 2 workers x 2 CPUs, gloo TP=2). The batch overrides are required: +`fake_data` has 64 samples and the default batch of 256 with `drop_last=True` yields no training step: ```bash python -m tinyexp.examples.vit_tp_exp dataloader_cfg.fake_data=true dataloader_cfg.input_size=32 \ + dataloader_cfg.train_batch_size_per_device=8 dataloader_cfg.val_batch_size_per_device=8 \ module_cfg.img_size=32 module_cfg.num_classes=10 epochs=1 max_train_steps=4 \ - ray_cfg.ray_num_gpus_per_worker=0.0 ray_cfg.ray_num_cpus_per_worker=4 + ray_cfg.ray_num_gpus_per_worker=0.0 ray_cfg.ray_num_cpus_per_worker=2 ``` Full 300-epoch training on many GPUs (e.g. 8×H200/H800 single node, DDP — DeiT-S has diff --git a/tests/examples/test_examples_help_smoke.py b/tests/examples/test_examples_help_smoke.py new file mode 100644 index 0000000..7a78d25 --- /dev/null +++ b/tests/examples/test_examples_help_smoke.py @@ -0,0 +1,209 @@ +"""Keep the bundled examples runnable and every documented command truthful. + +- Every example answers ``mode=help`` from a clean directory (first-run reliability). +- Every example starts with a run header (``Modes`` / ``Run`` / ``Env`` / ``Needs``), and every mode it lists is + handled by the experiment. +- Every command in a run header and in the user docs (README, SKILL, docs/) is checked statically: + overrides for a bundled example are composed with Hydra (unknown keys and bad types fail) and + ``tinyexp-run-with-*`` arguments are parsed with the tool's own argument parser. +""" + +import ast +import importlib +import os +import re +import shlex +import subprocess +import sys +from pathlib import Path + +import pytest +from hydra import compose, initialize +from hydra.core.config_store import ConfigStore + +PROJECT_ROOT = Path(__file__).resolve().parents[2] +EXAMPLES_DIR = PROJECT_ROOT / "tinyexp" / "examples" +EXAMPLES = sorted(p.stem for p in EXAMPLES_DIR.glob("*_exp.py")) +HEADER_KEYS = ("Modes", "Run", "Env", "Needs") +# User-facing docs; the refactor plan is a historical record and may quote superseded commands on purpose. +USER_DOCS = [ + PROJECT_ROOT / "README.md", + PROJECT_ROOT / "SKILL.md", + *sorted((PROJECT_ROOT / "docs").glob("*.md")), +] +CLI_TOOLS = { + "tinyexp-run-with-ray-cluster": "tinyexp.cli.run_with_ray_cluster", + "tinyexp-run-with-redis": "tinyexp.cli.run_with_redis", +} +_EXAMPLE_TOKEN = re.compile(r"^tinyexp(?:\.examples\.|/examples/)(\w+_exp)(?:\.py)?$") + + +def _run_help(module: str, cwd: Path, *overrides: str) -> subprocess.CompletedProcess: + env = os.environ.copy() + existing_pythonpath = env.get("PYTHONPATH") + env["PYTHONPATH"] = ( + str(PROJECT_ROOT) if not existing_pythonpath else f"{PROJECT_ROOT}{os.pathsep}{existing_pythonpath}" + ) + env.setdefault("HYDRA_FULL_ERROR", "1") + env.setdefault("WANDB_MODE", "disabled") + env.setdefault("WANDB_SILENT", "true") + return subprocess.run( # noqa: S603 + [sys.executable, "-m", f"tinyexp.examples.{module}", "mode=help", *overrides], + cwd=cwd, + env=env, + capture_output=True, + text=True, + timeout=120, + ) + + +def _header(example: str) -> dict[str, list[str]]: + """Return ``{key: lines}`` for the run header block (``Key : value`` lines plus indented continuations).""" + doc = ast.get_docstring(ast.parse((EXAMPLES_DIR / f"{example}.py").read_text())) or "" + sections: dict[str, list[str]] = {} + current = None + for line in doc.splitlines(): + match = re.match(r"^(\w+)\s*:(.*)$", line) + if match and match.group(1) in HEADER_KEYS: + current = match.group(1) + sections[current] = [match.group(2).strip()] if match.group(2).strip() else [] + elif current and line.startswith(" "): + sections[current].append(line.strip()) + elif sections and not line.strip(): + break # the header is one contiguous block + return sections + + +def _shell_segments(command: str) -> list[list[str]]: + """Split one logical shell line into token lists at ``&&``, ``||`` and ``;``; comments are dropped.""" + return [tokens for part in re.split(r"\s(?:&&|\|\||;)\s", command) if (tokens := shlex.split(part, comments=True))] + + +def _doc_commands() -> list[tuple[str, list[str]]]: + """Return ``(location, tokens)`` for documented commands that use a bundled example or a TinyExp CLI tool.""" + found = [] + for path in USER_DOCS: + in_block, logical, start = False, "", 0 + for number, line in enumerate(path.read_text().splitlines(), 1): + if line.lstrip().startswith("```"): + in_block, logical = not in_block, "" + continue + if not in_block: + continue + start = start if logical else number + if line.rstrip().endswith("\\"): + logical += line.rstrip()[:-1] + " " + continue + command, logical = logical + line, "" + for tokens in _shell_segments(command): + if any(_EXAMPLE_TOKEN.match(t) or t in CLI_TOOLS for t in tokens): + found.append((f"{path.relative_to(PROJECT_ROOT)}:{start}", tokens)) + return found + + +def _example_and_overrides(tokens: list[str]) -> tuple[str, list[str]]: + """Return ``(example, overrides)`` for the bundled example started by ``tokens``, or ``("", [])``.""" + for index, token in enumerate(tokens): + match = _EXAMPLE_TOKEN.match(token) + if match: + overrides = [t for t in tokens[index + 1 :] if "=" in t and not t.startswith("-")] + return match.group(1), overrides + return "", [] + + +def _exp_class(example: str) -> type: + module = importlib.import_module(f"tinyexp.examples.{example}") + class_name = re.search(r"store_and_run_exp\((\w+)\)", Path(module.__file__).read_text()).group(1) + return getattr(module, class_name) + + +def _compose(example: str, overrides: list[str]) -> None: + """Compose the example's config with ``overrides`` exactly as its CLI would; raises on unknown keys/bad types.""" + config_name = f"_command_check_{example}" + ConfigStore.instance().store(name=config_name, node=_exp_class(example)()) + with initialize(version_base=None): + compose(config_name=config_name, overrides=overrides) + + +def _check_command(tokens: list[str]) -> None: + for position, token in enumerate(tokens): + if token in CLI_TOOLS: + importlib.import_module(CLI_TOOLS[token]).parse_args(tokens[position + 1 :]) + example, overrides = _example_and_overrides(tokens) + if example: + _compose(example, overrides) + + +@pytest.mark.parametrize("module", EXAMPLES) +def test_example_help_succeeds_in_clean_directory(module: str, tmp_path: Path) -> None: + result = _run_help(module, tmp_path) + + assert result.returncode == 0, result.stderr + assert "-------- Configurations --------" in f"{result.stdout}\n{result.stderr}" + + +@pytest.mark.parametrize("example", EXAMPLES) +def test_example_has_run_header(example: str) -> None: + header = _header(example) + + assert set(HEADER_KEYS) <= set(header), f"{example} header has {sorted(header)}" + assert all(header[key] for key in HEADER_KEYS) + assert any(f"-m tinyexp.examples.{example} mode=help" in line for line in header["Run"]) + + +@pytest.mark.parametrize("example", EXAMPLES) +def test_example_header_commands_compose(example: str) -> None: + commands = [tokens for line in _header(example)["Run"] for tokens in _shell_segments(line)] + assert commands + + for tokens in commands: + assert _example_and_overrides(tokens)[0] == example, f"{example} header runs something else: {tokens}" + _check_command(tokens) + + +@pytest.mark.parametrize("example", EXAMPLES) +def test_example_header_modes_are_handled(example: str) -> None: + source = (EXAMPLES_DIR / f"{example}.py").read_text() + modes = [m.strip() for m in _header(example)["Modes"][0].split("(")[0].split("|")] + + assert "help" in modes + for mode in modes: + if mode != "help": + assert re.search(rf'self\.mode\s*[!=]=\s*"{mode}"', source), f"{example}: mode {mode!r} not handled" + + +def test_user_docs_contain_checked_commands() -> None: + """Guard against the doc scan silently finding nothing (e.g. after a docs restructure).""" + locations = {location.split(":")[0] for location, _ in _doc_commands()} + + assert {"README.md", "SKILL.md", "docs/running-modes.md", "docs/vit_tp.md"} <= locations + + +@pytest.mark.parametrize("location, tokens", _doc_commands(), ids=lambda value: value if isinstance(value, str) else "") +def test_documented_command_is_valid(location: str, tokens: list[str]) -> None: + try: + _check_command(tokens) + except SystemExit as error: # argparse reports invalid CLI arguments by exiting + pytest.fail(f"{location}: invalid arguments for {tokens[0]!r}: exit {error.code}") + + +def test_help_writes_only_under_output_root(tmp_path: Path) -> None: + cwd = tmp_path / "cwd" + cwd.mkdir() + output_root = tmp_path / "out" + + result = _run_help("mnist_exp", cwd, f"output_root={output_root}") + + assert result.returncode == 0, result.stderr + assert list(cwd.iterdir()) == [] + + +def test_help_ignores_unusable_default_output_in_cwd(tmp_path: Path) -> None: + """A stale ``./output`` (e.g. a broken symlink) must not break runs that set ``output_root``.""" + cwd = tmp_path / "cwd" + cwd.mkdir() + (cwd / "output").symlink_to(tmp_path / "missing-target") + + result = _run_help("mnist_exp", cwd, f"output_root={tmp_path / 'out'}") + + assert result.returncode == 0, result.stderr diff --git a/tests/examples/test_mnist_exp_run.py b/tests/examples/test_mnist_exp_run.py index 667004f..0f4ac58 100644 --- a/tests/examples/test_mnist_exp_run.py +++ b/tests/examples/test_mnist_exp_run.py @@ -9,8 +9,8 @@ from tinyexp.examples.mnist_exp import Exp -def test_mnist_run_val_mode_requires_resume_from(tmp_path: Path, monkeypatch) -> None: - exp = Exp(output_root=str(tmp_path), exp_name="mnist_val", mode="val", resume_from="") +def test_mnist_run_eval_mode_requires_resume_from(tmp_path: Path, monkeypatch) -> None: + exp = Exp(output_root=str(tmp_path), exp_name="mnist_val", mode="eval", resume_from="") dummy_accelerator = SimpleNamespace( rank=0, device="cpu", @@ -26,7 +26,7 @@ def test_mnist_run_val_mode_requires_resume_from(tmp_path: Path, monkeypatch) -> exp.run() -def test_mnist_run_val_mode_uses_checkpoint(tmp_path: Path, monkeypatch) -> None: +def test_mnist_run_eval_mode_uses_checkpoint(tmp_path: Path, monkeypatch) -> None: exp_for_ckpt = Exp(output_root=str(tmp_path), exp_name="mnist_val") checkpoint_path = exp_for_ckpt.checkpoint_cfg.save_checkpoint( run_dir=str(tmp_path / "mnist_val"), @@ -36,7 +36,7 @@ def test_mnist_run_val_mode_uses_checkpoint(tmp_path: Path, monkeypatch) -> None exp_class=exp_for_ckpt.exp_class, ) - exp = Exp(output_root=str(tmp_path), exp_name="mnist_val", mode="val", resume_from=checkpoint_path) + exp = Exp(output_root=str(tmp_path), exp_name="mnist_val", mode="eval", resume_from=checkpoint_path) destroy_calls = [] dummy_accelerator = SimpleNamespace( rank=0, @@ -113,3 +113,14 @@ def wait_for_everyone(self) -> None: assert partitions == [[0, 3, 6, 9], [1, 4, 7], [2, 5, 8]] assert sorted(index for partition in partitions for index in partition) == list(range(10)) + + +def test_mnist_run_unknown_mode_lists_supported_modes(tmp_path: Path, monkeypatch) -> None: + from tinyexp.exceptions import UnknownExperimentModeError + + monkeypatch.setenv("RANK", "0") + monkeypatch.setenv("WORLD_SIZE", "1") + exp = Exp(output_root=str(tmp_path), exp_name="mnist_bad_mode", mode="val") + + with pytest.raises(UnknownExperimentModeError, match=r"Unknown mode 'val'.*train, eval, help"): + exp.run() diff --git a/tests/examples/test_pi_exp_run.py b/tests/examples/test_pi_exp_run.py index 22e990b..30dcc4b 100644 --- a/tests/examples/test_pi_exp_run.py +++ b/tests/examples/test_pi_exp_run.py @@ -4,7 +4,6 @@ from types import SimpleNamespace import pytest -import ray from tinyexp.examples import pi_exp @@ -29,7 +28,7 @@ def destroy(self) -> None: accelerator = DummyAccelerator() logger = SimpleNamespace() - exp = pi_exp.Exp(output_root=str(tmp_path), exp_name="pi_test") + exp = pi_exp.Exp(output_root=str(tmp_path), exp_name="pi_test", launcher="mp") monkeypatch.setattr(pi_exp, "torch", SimpleNamespace(pi=3.14)) monkeypatch.setattr("tinyexp.tiny_engine.accelerator.CPUAccelerator", lambda: accelerator) @@ -62,12 +61,6 @@ def destroy(self) -> None: monkeypatch.setattr(exp.logger_cfg, "build_logger", lambda **kwargs: SimpleNamespace()) monkeypatch.setattr(exp, "print_cfg", lambda logger: {}) monkeypatch.setattr(exp, "_estimate_pi", lambda accelerator: 3.14) - monkeypatch.setattr(ray, "is_initialized", lambda: True) - monkeypatch.setattr( - ray, - "get_runtime_context", - lambda: SimpleNamespace(get_actor_id=lambda: "actor-id"), - ) exp.run() @@ -99,3 +92,12 @@ def destroy(self) -> None: exp.run() assert events == ["destroy"] + + +def test_pi_run_unknown_mode_lists_supported_modes(tmp_path: Path) -> None: + from tinyexp.exceptions import UnknownExperimentModeError + + exp = pi_exp.Exp(output_root=str(tmp_path), mode="train") + + with pytest.raises(UnknownExperimentModeError, match=r"Unknown mode 'train'.*run, help"): + exp.run() diff --git a/tests/examples/test_resnet_exp_run.py b/tests/examples/test_resnet_exp_run.py index 1fae93d..4536266 100644 --- a/tests/examples/test_resnet_exp_run.py +++ b/tests/examples/test_resnet_exp_run.py @@ -238,8 +238,8 @@ def __getitem__(self, index: int) -> tuple[int, int]: assert val_dataloader.persistent_workers is expected -def test_resnet_run_val_mode_requires_resume_from(tmp_path: Path, monkeypatch) -> None: - exp = ResNetExp(output_root=str(tmp_path), exp_name="resnet_val", mode="val", resume_from="") +def test_resnet_run_eval_mode_requires_resume_from(tmp_path: Path, monkeypatch) -> None: + exp = ResNetExp(output_root=str(tmp_path), exp_name="resnet_val", mode="eval", resume_from="") dummy_accelerator = SimpleNamespace( rank=0, device="cpu", @@ -256,7 +256,7 @@ def test_resnet_run_val_mode_requires_resume_from(tmp_path: Path, monkeypatch) - def test_resnet_run_logs_overrides_with_experiment_logger(tmp_path: Path, monkeypatch) -> None: - exp = ResNetExp(mode="val", resume_from="") + exp = ResNetExp(mode="eval", resume_from="") exp.set_cfg(OmegaConf.create({"output_root": str(tmp_path), "exp_name": "resnet_val"})) dummy_accelerator = SimpleNamespace( rank=0, @@ -276,7 +276,7 @@ def test_resnet_run_logs_overrides_with_experiment_logger(tmp_path: Path, monkey assert any(message.startswith("-------- Overridden Configurations --------\n") for message in messages) -def test_resnet_run_val_mode_uses_checkpoint(tmp_path: Path, monkeypatch) -> None: +def test_resnet_run_eval_mode_uses_checkpoint(tmp_path: Path, monkeypatch) -> None: exp_for_ckpt = ResNetExp(output_root=str(tmp_path), exp_name="resnet_val") checkpoint_path = exp_for_ckpt.checkpoint_cfg.save_checkpoint( run_dir=str(tmp_path / "resnet_val"), @@ -289,7 +289,7 @@ def test_resnet_run_val_mode_uses_checkpoint(tmp_path: Path, monkeypatch) -> Non exp = ResNetExp( output_root=str(tmp_path), exp_name="resnet_val", - mode="val", + mode="eval", resume_from=checkpoint_path, ) destroy_calls = [] diff --git a/tests/examples/test_vit_tp_exp.py b/tests/examples/test_vit_tp_exp.py index 78c938a..69de549 100644 --- a/tests/examples/test_vit_tp_exp.py +++ b/tests/examples/test_vit_tp_exp.py @@ -5,6 +5,7 @@ import copy import os import socket +import types import pytest @@ -161,6 +162,20 @@ def __init__(self, *args, **kwargs) -> None: ) +def test_train_dataloader_rejects_empty_epoch() -> None: + """docs/vit_tp.md D8: a batch larger than the per-rank data must fail fast, not train zero steps.""" + experiment = VitTpExp() + cfg = experiment.dataloader_cfg + cfg.fake_data, cfg.input_size, cfg.train_num_workers = True, 32, 0 + accelerator = types.SimpleNamespace(rank=0, world_size=1) + + with pytest.raises(ValueError, match=r"no training batch.*train_batch_size_per_device=256"): + cfg.build_train_dataloader(accelerator, replicate_data=True) + + cfg.train_batch_size_per_device = 8 + assert len(cfg.build_train_dataloader(accelerator, replicate_data=True)) == cfg.fake_data_len // 8 + + def _tp_model() -> VisionTransformer: torch.manual_seed(42) return VisionTransformer( diff --git a/tests/scripts/test_run_with_redis.py b/tests/scripts/test_run_with_redis.py index 0f4f156..19dd646 100644 --- a/tests/scripts/test_run_with_redis.py +++ b/tests/scripts/test_run_with_redis.py @@ -578,3 +578,59 @@ def close(self) -> None: assert exc_info.value.code == 128 + signum assert calls == [(child, signum), "close-lifecycle"] + + +def _write_exp_module(directory: Path, name: str, launcher: str) -> None: + (directory / f"{name}.py").write_text( + "from dataclasses import dataclass, field\n" + "from tinyexp import TinyExp\n" + "from tinyexp.exp_mixins import RedisCfgMixin\n" + "@dataclass\n" + "class Exp(TinyExp, RedisCfgMixin):\n" + f" launcher: str = {launcher!r}\n" + " @dataclass\n" + " class RedisCfg(RedisCfgMixin.RedisCfg):\n" + " redis_cache_max_memory: int = 7\n" + " redis_cfg: RedisCfg = field(default_factory=RedisCfg)\n" + ) + + +@pytest.mark.parametrize( + ("launcher_default", "extra", "rejected"), + [("ray", [], True), ("mp", ["launcher=ray"], True), ("ray", ["launcher=mp"], False)], +) +def test_main_rejects_launcher_ray_before_starting_redis( + tmp_path: Path, monkeypatch, capsys, launcher_default: str, extra: list[str], rejected: bool +) -> None: + _write_exp_module(tmp_path, "ray_default_exp", launcher_default) + monkeypatch.chdir(tmp_path) + started: list[str] = [] + monkeypatch.setattr(run_with_redis, "start_local_redis", lambda *a, **k: started.append("redis") or None) + monkeypatch.setattr("tinyexp.cli.run_with_redis.subprocess.Popen", lambda *args, **kwargs: _FinishedProcess()) + + code = main(["python", "-m", "ray_default_exp", *extra]) + + if rejected: + assert code == 2 + assert started == [] + assert "launcher=ray" in capsys.readouterr().err + else: + assert started == ["redis"] + + +@pytest.mark.parametrize( + "command", + [ + ["python", "-m", "cwd_exp"], + ["torchrun", "--nproc-per-node=2", "-m", "cwd_exp", "launcher=mp"], + ["accelerate", "launch", "--cpu", "-m", "cwd_exp", "launcher=mp"], + ["python", "-m", "torch.distributed.run", "--nproc-per-node=2", "-m", "cwd_exp"], + ], +) +def test_build_redis_cfg_resolves_module_like_python_m(tmp_path: Path, monkeypatch, command: list[str]) -> None: + # ``python -m`` (also run by torchrun/accelerate for ``-m``) puts the current directory on sys.path. + _write_exp_module(tmp_path, "cwd_exp", "mp") + monkeypatch.chdir(tmp_path) + monkeypatch.setattr(sys, "path", [p for p in sys.path if p not in ("", ".", str(tmp_path))]) + + assert build_redis_cfg(command).redis_cache_max_memory == 7 diff --git a/tests/test_tinyexp.py b/tests/test_tinyexp.py index 4b01d13..23ddfbf 100644 --- a/tests/test_tinyexp.py +++ b/tests/test_tinyexp.py @@ -109,6 +109,21 @@ def test_set_cfg_overrides_nested(monkeypatch: pytest.MonkeyPatch) -> None: assert exp.b == 4 +def test_set_cfg_keeps_overrides_recorded_before_overrided_cfg_key(monkeypatch: pytest.MonkeyPatch) -> None: + """A full Hydra config carries ``overrided_cfg: {}``; it must not wipe records of earlier fields.""" + monkeypatch.setenv("RANK", "1") + exp = _CfgExp() + + # Field order matches the dataclass: base TinyExp fields (incl. overrided_cfg) come before subclass fields. + cfg = OmegaConf.create({"mode": "help", "overrided_cfg": {}, "sub_cfg": {"a": 3}}) + exp.set_cfg(cfg) + + assert exp.overrided_cfg == { + "mode": {"value": "help", "original": "train"}, + "a": {"value": 3, "original": 1}, + } + + def test_set_cfg_overrides_mapping_field(monkeypatch: pytest.MonkeyPatch) -> None: monkeypatch.setenv("RANK", "1") exp = _MappingCfgExp() @@ -124,6 +139,17 @@ def test_set_cfg_overrides_mapping_field(monkeypatch: pytest.MonkeyPatch) -> Non } +def test_set_cfg_does_not_report_exp_class_as_override(monkeypatch: pytest.MonkeyPatch) -> None: + """``exp_class`` is always filled in by ``store_and_run_exp``; listing it under overrides is noise.""" + monkeypatch.setenv("RANK", "1") + exp = _CfgExp() + + exp.set_cfg(OmegaConf.create({"exp_class": "pkg.mod.Exp", "b": 4})) + + assert exp.exp_class == "pkg.mod.Exp" + assert exp.overrided_cfg == {"b": {"value": 4, "original": 2}} + + def test_set_cfg_unknown_key_raises(monkeypatch: pytest.MonkeyPatch) -> None: monkeypatch.setenv("RANK", "1") exp = _CfgExp() diff --git a/tests/test_tinyexp_artifacts.py b/tests/test_tinyexp_artifacts.py index 425e69d..4699378 100644 --- a/tests/test_tinyexp_artifacts.py +++ b/tests/test_tinyexp_artifacts.py @@ -1,6 +1,8 @@ from __future__ import annotations import random +import subprocess +import sys from dataclasses import dataclass from pathlib import Path @@ -25,6 +27,17 @@ def test_root_package_re_exports_mixins() -> None: assert WandbCfgMixin.__name__ == "WandbCfgMixin" +def test_import_tinyexp_does_not_load_optional_heavy_packages() -> None: + """``import tinyexp`` must work without the ``[pytorch]`` extra; torch/wandb stay lazily imported. + + Runs in a subprocess because this pytest process has already imported torch. + """ + code = "import sys, tinyexp, tinyexp.tiny_engine; print(sorted({'torch', 'wandb'} & set(sys.modules)))" + result = subprocess.run([sys.executable, "-c", code], capture_output=True, text=True, check=True) # noqa: S603 + + assert result.stdout.strip() == "[]" + + def test_get_run_dir(tmp_path: Path) -> None: exp = TinyExp(output_root=str(tmp_path), exp_name="demo_exp") diff --git a/tests/tine_engine/test_accelerator_protocol.py b/tests/tiny_engine/test_accelerator_protocol.py similarity index 100% rename from tests/tine_engine/test_accelerator_protocol.py rename to tests/tiny_engine/test_accelerator_protocol.py diff --git a/tests/tine_engine/test_cpu_accelerator.py b/tests/tiny_engine/test_cpu_accelerator.py similarity index 100% rename from tests/tine_engine/test_cpu_accelerator.py rename to tests/tiny_engine/test_cpu_accelerator.py diff --git a/tests/tine_engine/test_cpu_accelerator_unit.py b/tests/tiny_engine/test_cpu_accelerator_unit.py similarity index 100% rename from tests/tine_engine/test_cpu_accelerator_unit.py rename to tests/tiny_engine/test_cpu_accelerator_unit.py diff --git a/tests/tine_engine/test_ddp_accelerator.py b/tests/tiny_engine/test_ddp_accelerator.py similarity index 100% rename from tests/tine_engine/test_ddp_accelerator.py rename to tests/tiny_engine/test_ddp_accelerator.py diff --git a/tests/tine_engine/test_ddp_accelerator_unit.py b/tests/tiny_engine/test_ddp_accelerator_unit.py similarity index 100% rename from tests/tine_engine/test_ddp_accelerator_unit.py rename to tests/tiny_engine/test_ddp_accelerator_unit.py diff --git a/tests/tine_engine/test_fsdp_accelerator_unit.py b/tests/tiny_engine/test_fsdp_accelerator_unit.py similarity index 100% rename from tests/tine_engine/test_fsdp_accelerator_unit.py rename to tests/tiny_engine/test_fsdp_accelerator_unit.py diff --git a/tests/tine_engine/test_hf_accelerator_unit.py b/tests/tiny_engine/test_hf_accelerator_unit.py similarity index 100% rename from tests/tine_engine/test_hf_accelerator_unit.py rename to tests/tiny_engine/test_hf_accelerator_unit.py diff --git a/tinyexp/__init__.py b/tinyexp/__init__.py index 7ac622e..6753795 100644 --- a/tinyexp/__init__.py +++ b/tinyexp/__init__.py @@ -1,3 +1,59 @@ +"""TinyExp: the configured experiment class is the entrypoint. + +Minimal experiment (see ``tinyexp/examples/mnist_exp.py`` for a full DDP/CPU recipe):: + + from dataclasses import dataclass, field + from tinyexp import LoggerCfgMixin, TinyExp, store_and_run_exp + + @dataclass(repr=False) + class Exp(TinyExp, LoggerCfgMixin): + @dataclass + class TrainCfg: + lr: float = 0.1 + + train_cfg: TrainCfg = field(default_factory=TrainCfg) + + def run(self) -> None: # every mode except ``help`` lands here; dispatch on self.mode + logger = self.logger_cfg.build_logger(save_dir=self.get_run_dir()) + self.print_cfg(logger) + + if __name__ == "__main__": + store_and_run_exp(Exp) + + # python my_exp.py mode=help train_cfg.lr=0.01 -> print resolved config, run nothing + +``TinyExp`` fields (all overridable as ``key=value`` on the CLI): + +| field | default | meaning | +| --- | --- | --- | +| `exp_class` | `""` | import path of the class, filled in by `store_and_run_exp` (not listed as an override); `__main__.Exp` when started as `python file.py` | +| `exp_name` | main file stem | run subdirectory name (`my_exp.py` -> `my_exp`) | +| `launcher` | `"mp"` | `mp`: run `Exp.run()` in this process; `ray`: spawn Ray workers | +| `output_root` | `"./output"` | root of all outputs; run dir `get_run_dir()` = output_root/exp_name; Hydra's own log is `/tinyexp.log` | +| `mode` | `"train"` | `help` is handled by TinyExp (print config, exit; it still creates `output_root`); any other value reaches `run()`, whose meaning the experiment defines | +| `resume_from` | `""` | checkpoint path to resume or evaluate from; empty means none | +| `hydra` | internal | Hydra settings; do not override | +| `overrided_cfg` | internal | filled by `set_cfg`; read it, do not set it | + +Pairing ``launcher`` with the process owner: + +- ``python exp.py launcher=mp``: one process. +- ``python exp.py launcher=ray``: the driver creates ``ray_cfg.ray_num_worker`` workers (``RayCfgMixin``). +- ``torchrun ... exp.py launcher=mp`` / ``accelerate launch ... launcher=mp``: the external tool owns processes. +- ``tinyexp-run-with-ray-cluster ... -- python exp.py launcher=ray``: multi-node static Ray cluster. + +CLI overrides are checked by Hydra against the dataclasses: unknown keys are rejected +(``Key 'x' not in 'DataloaderCfg'``) and values are type-checked, but not otherwise validated (``mode=help`` +accepts ``accelerator_cfg.accelerator=bogus``). For a ``dict`` field, change an existing key with ``tags.owner=x`` +and add a new one with ``+tags.team=y`` (without the ``+``: ``Key 'team' is not in struct``). How the result is +merged onto the experiment and recorded for ``mode=help``: ``TinyExp.set_cfg``. Ray workers and ``python -m`` vs +``python file.py``: ``RayCfgMixin``. + +Exports: ``TinyExp``, ``store_and_run_exp``, ``simple_launch_exp``, ``ConfigStore`` and the mixins +``LoggerCfgMixin``, ``RayCfgMixin``, ``CheckpointCfgMixin``, ``WandbCfgMixin``, ``RedisCfgMixin`` +(see ``tinyexp.exp_mixins``). +""" + __author__ = "LI Zeming" __email__ = "zane.li@connect.ust.hk" __license__ = "MIT" @@ -60,11 +116,12 @@ def get_help_string(self, action): @dataclass class _HydraConfig(HydraConf): """ - To avoid hydra output the config in unexpected directory. + Keep Hydra's run directory (and its ``tinyexp.log``) under ``output_root`` instead of a fixed ``./output``, + and do not let Hydra write its own config copies. """ output_subdir: Optional[str] = None - run: RunDir = field(default_factory=lambda: RunDir("./output")) + run: RunDir = field(default_factory=lambda: RunDir("${output_root}")) def _default_exp_name() -> str: @@ -116,7 +173,7 @@ class TinyExp: # The experiment mode: # - "run": run a general experiment workload. # - "train": train the model. - # - "val": evaluate the model. + # - "eval": evaluate the model. # - "help": print the experiment configuration and exit. mode: str = "train" @@ -134,9 +191,25 @@ def get_run_dir(self) -> str: return os.path.join(self.output_root, self.exp_name) def set_cfg(self, cfg_hydra, cfg_object=None): + """Copy the Hydra config (dataclass defaults plus CLI overrides) onto this instance; returns ``self``. + + Called by the launcher in every process (with ``launcher=ray``, by each worker on its own copy): + + - nested ``XXXCfg`` dataclasses are updated field by field, recursively; + - fields whose current value is a mapping (``dict``) are replaced as a whole; + - any other value is assigned directly; a key the object does not have raises + ``UnknownConfigurationKeyError``. + + Every changed value is recorded in ``overrided_cfg[] = {"value", "original"}`` (keyed by the + leaf name only, not the dotted path) and listed under "Overridden" by ``print_cfg`` / ``mode=help``. + """ if cfg_object is None: cfg_object = self for key, value in cfg_hydra.items(): + if cfg_object is self and key == "overrided_cfg": + # Internal bookkeeping: a full config carries the default ``{}``, which must not + # overwrite the records collected for fields processed before it. + continue if hasattr(cfg_object, key): sub_object = getattr(cfg_object, key) if isinstance(value, (DictConfig, Mapping)): @@ -160,10 +233,12 @@ def set_cfg(self, cfg_hydra, cfg_object=None): # Otherwise, set the attribute directly ori_value = getattr(cfg_object, key, None) if ori_value != value: - self.overrided_cfg[key] = { - "value": value, - "original": ori_value, - } + if not (cfg_object is self and key == "exp_class"): + # exp_class is always filled in by store_and_run_exp; reporting it is noise. + self.overrided_cfg[key] = { + "value": value, + "original": ori_value, + } setattr(cfg_object, key, value) else: raise UnknownConfigurationKeyError(key) diff --git a/tinyexp/cli/run_with_ray_cluster.py b/tinyexp/cli/run_with_ray_cluster.py index e045d61..69e0bfa 100644 --- a/tinyexp/cli/run_with_ray_cluster.py +++ b/tinyexp/cli/run_with_ray_cluster.py @@ -64,7 +64,7 @@ def parse_args(argv: list[str]) -> argparse.Namespace: "--node-count", type=positive_int, default=1, - help="Total cluster nodes. Default: 1.", + help="Total cluster nodes. Default: 1 (runs the command unchanged; does not start or stop Ray).", ) parser.add_argument("--node-rank", type=uint, default=0, help="Current node rank. Default: 0.") parser.add_argument("--head-addr", default="", help="Address workers use to reach Ray head.") diff --git a/tinyexp/cli/run_with_redis.py b/tinyexp/cli/run_with_redis.py index d41e9fb..e94d4e0 100644 --- a/tinyexp/cli/run_with_redis.py +++ b/tinyexp/cli/run_with_redis.py @@ -1,6 +1,10 @@ #!/usr/bin/env python3 """Run a command with a Redis cache lifecycle owned by this wrapper. +Only for launches without Ray (``launcher=mp``: plain ``python``, ``torchrun``, ``accelerate launch``). With +``launcher=ray`` the Ray driver starts and stops Redis itself, so the wrapper refuses such a command (exit code 2) +before starting anything; the launcher is read from a ``launcher=`` override or the experiment's default. + The wrapper reads ``redis_cfg`` overrides from the target command, starts standalone Redis for single-node jobs or a Redis Cluster for multi-node jobs, injects the final Redis connection settings back into the child command as Hydra @@ -134,7 +138,10 @@ def env_uint(name: str, default: int) -> int: def parse_args(argv: list[str]) -> argparse.Namespace: parser = argparse.ArgumentParser( prog="tinyexp-run-with-redis", - description="Start Redis, inject redis_cfg overrides, then run .", + description=( + "Start Redis, inject redis_cfg overrides, then run . Only for launcher=mp commands " + "(python, torchrun, accelerate launch); launcher=ray starts Redis itself and must not be wrapped." + ), ) parser.add_argument( "--node-count", @@ -190,12 +197,23 @@ def _load_python_file(module_path: Path): # type: ignore[no-untyped-def] def _load_experiment_module(argv: list[str]): # type: ignore[no-untyped-def] - if len(argv) >= 3 and Path(argv[0]).name.startswith("python") and argv[1] == "-m": - return importlib.import_module(argv[2]) + """Import the experiment the command starts, resolving it the way the child process will. + + ``-m `` (python, torchrun, accelerate launch; the last one wins, e.g. after + ``python -m torch.distributed.run``) imports with the current directory on ``sys.path``, as ``python -m`` does; + ``.py`` runs with the file's directory on ``sys.path``, as ``python .py`` does. + """ + modules = [argv[i + 1] for i, arg in enumerate(argv[:-1]) if arg == "-m"] + if modules: + if os.getcwd() not in sys.path: + sys.path.insert(0, os.getcwd()) + return importlib.import_module(modules[-1]) for arg in argv: module_path = Path(arg) if module_path.suffix == ".py" and module_path.is_file(): + if str(module_path.resolve().parent) not in sys.path: + sys.path.insert(0, str(module_path.resolve().parent)) return _load_python_file(module_path.resolve()) return None @@ -210,7 +228,14 @@ def main(argv: list[str]) -> int: global CHILD_PROCESS, REDIS_LIFECYCLE args = parse_args(argv) - redis_cfg = build_redis_cfg(args.command) + redis_cfg, launcher = _infer_experiment(args.command) + if launcher == "ray": + print( + "tinyexp-run-with-redis: the command runs with launcher=ray, whose Ray driver starts and stops Redis " + "itself; run it without this wrapper. The wrapper is for launcher=mp (python, torchrun, accelerate launch).", + file=sys.stderr, + ) + return 2 configured_ports = [int(port) for port in redis_cfg.redis_cluster_ports] if not configured_ports: print("redis_cfg.redis_cluster_ports must not be empty", file=sys.stderr) @@ -292,6 +317,11 @@ def main(argv: list[str]) -> int: def build_redis_cfg(argv: list[str]) -> Any: + return _infer_experiment(argv)[0] + + +def _infer_experiment(argv: list[str]) -> tuple[Any, str]: + """Return ``(redis_cfg, launcher)`` of the command's experiment (``launcher`` is ``""`` when unknown).""" exp: Any | None = None module = _load_experiment_module(argv) if module is not None: @@ -307,11 +337,13 @@ def build_redis_cfg(argv: list[str]) -> Any: exp = cast(Any, exp_classes[0]()) exp.exp_class = f"{exp_classes[0].__module__}.{exp_classes[0].__qualname__}" + launcher_overrides = [arg.split("=", 1)[1] for arg in argv if arg.startswith("launcher=")] + launcher = launcher_overrides[-1] if launcher_overrides else (exp.launcher if exp is not None else "") overrides = [arg for arg in argv if arg.startswith("redis_cfg.")] if exp is not None: if overrides: exp.set_cfg(OmegaConf.from_dotlist(overrides)) - return exp.redis_cfg + return exp.redis_cfg, launcher print( "Could not infer experiment config; using RedisCfgMixin defaults", @@ -320,7 +352,7 @@ def build_redis_cfg(argv: list[str]) -> Any: cfg = OmegaConf.structured(RedisCfgMixin()) if overrides: cfg = OmegaConf.merge(cfg, OmegaConf.from_dotlist(overrides)) - return cfg.redis_cfg + return cfg.redis_cfg, launcher def start_local_redis(redis_cfg: Any, lifecycle: _RedisLifecycle) -> tuple[str, list[int]] | None: diff --git a/tinyexp/dataset/__init__.py b/tinyexp/dataset/__init__.py index b928822..9540005 100644 --- a/tinyexp/dataset/__init__.py +++ b/tinyexp/dataset/__init__.py @@ -1,3 +1,9 @@ +"""Data helpers for training loops. + +Exports: ``InfiniteSampler`` (endless, deterministic, rank-sharded index stream; see ``sampler``) and +``HoldOnesampleDataLoader`` (repeats one cached batch, for profiling/debugging). +""" + # flake8: noqa: F401, F403 from .fake_dataloader import * diff --git a/tinyexp/dataset/fake_dataloader.py b/tinyexp/dataset/fake_dataloader.py index 7c3cf82..bf1f44b 100644 --- a/tinyexp/dataset/fake_dataloader.py +++ b/tinyexp/dataset/fake_dataloader.py @@ -1,3 +1,8 @@ +"""``HoldOnesampleDataLoader(dataloader)``: yields a deep copy of the first batch forever. + +Use it to profile or debug the model without data-loading cost; ``len()`` reports the wrapped length. +""" + import copy __all__ = ["HoldOnesampleDataLoader"] @@ -5,7 +10,8 @@ class HoldOnesampleDataLoader: """ - A fake dataloader that holds one sample from the original dataloader. + A fake dataloader that holds the first batch of the original dataloader and returns a deep copy of it + on every ``next()``; iteration never stops, so bound the loop yourself (``len()`` is the wrapped length). This is useful for testing and profiling purposes, where we want to test the model with a single in-memory sample. """ diff --git a/tinyexp/dataset/sampler.py b/tinyexp/dataset/sampler.py index 4eda6e0..88f80cb 100644 --- a/tinyexp/dataset/sampler.py +++ b/tinyexp/dataset/sampler.py @@ -1,3 +1,14 @@ +"""``InfiniteSampler``: an endless, deterministic, rank-sharded index stream for step-based training. + +Usage:: + + sampler = InfiniteSampler(len(dataset), shuffle=True, seed=0, accelerator=accelerator) + loader = torch.utils.data.DataLoader(dataset, batch_size=bs, sampler=sampler) + batches = iter(loader) # create once; call next(batches) for a fixed number of steps + +The class docstring documents the stream model, ``len()`` semantics and ``drop_last``. +""" + import itertools from collections.abc import Iterator from typing import Any, Optional diff --git a/tinyexp/examples/README.md b/tinyexp/examples/README.md new file mode 100644 index 0000000..ab3e034 --- /dev/null +++ b/tinyexp/examples/README.md @@ -0,0 +1,20 @@ +# TinyExp examples + +Each example is one runnable file. Its module docstring starts with a run header — `Modes`, `Run`, `Env`, +`Needs`, `Out` — which is the single source of commands for that example (kept valid by +`tests/examples/test_examples_help_smoke.py`). Read the header, not the code, to run an example: + +```bash +python -c "import tinyexp.examples.mnist_exp as m; print(m.__doc__)" +``` + +| You want to... | Start from | Data | Hardware | +| --- | --- | --- | --- | +| Write a new training experiment (template) | `mnist_exp.py` | auto-download | CPU, or GPUs with DDP | +| Run a non-deep-learning distributed job / fastest smoke test | `pi_exp.py` | none | CPU | +| Train on ImageNet with DDP, FSDP or HF Accelerate; Redis byte cache | `resnet_exp.py` | `IMAGENET_HOME` | GPUs | +| Use tensor parallelism; reproduce DeiT-S; benchmark TP vs DDP | `vit_tp_exp.py` (+ `docs/vit_tp.md`) | `IMAGENET_HOME` | 2+ GPUs (CPU smoke available) | + +All examples share the same shape: nested `XXXCfg` dataclasses (each with `build_*` helpers) followed by `run()`, +which dispatches on `mode`. Check any override with `mode=help` before running; outputs go to +`//`. diff --git a/tinyexp/examples/mnist_exp.py b/tinyexp/examples/mnist_exp.py index 7bfe03f..6eb06ed 100644 --- a/tinyexp/examples/mnist_exp.py +++ b/tinyexp/examples/mnist_exp.py @@ -1,3 +1,21 @@ +"""MNIST baseline: CPU or DDP image classification with checkpoints, resume and logging. + +Start here when writing a new experiment: config sections (``XXXCfg``) on top, ``run()`` below. + +Modes : train | eval | help (eval requires resume_from=) +Run : + python -m tinyexp.examples.mnist_exp mode=help + python -m tinyexp.examples.mnist_exp # default: 2 Ray workers x 4 CPUs (needs 8 CPUs) + python -m tinyexp.examples.mnist_exp launcher=mp # one process, no Ray + torchrun --nproc-per-node=2 -m tinyexp.examples.mnist_exp launcher=mp accelerator_cfg.accelerator=ddp + python -m tinyexp.examples.mnist_exp mode=eval resume_from=./output/mnist_exp/best.ckpt +Env : none; global rank 0 downloads MNIST to dataloader_cfg.data_root (./data/) on first run while the other + ranks wait, then read it; with several nodes, data_root must be shared or already populated +Note : always trains 3 epochs (not configurable); raise dataloader_cfg.train_batch_size_per_device for a faster run +Needs : tinyexp[pytorch]; accelerator_cfg.accelerator=ddp needs CUDA +Out : /mnist_exp/{log.txt,last.ckpt,best.ckpt} +""" + import datetime from dataclasses import dataclass, field @@ -10,7 +28,7 @@ from torchvision import datasets, transforms from tinyexp import TinyExp, store_and_run_exp -from tinyexp.exceptions import UnknownAcceleratorTypeError +from tinyexp.exceptions import UnknownAcceleratorTypeError, UnknownExperimentModeError from tinyexp.exp_mixins import CheckpointCfgMixin, LoggerCfgMixin, RayCfgMixin, WandbCfgMixin from tinyexp.tiny_engine.accelerator import AcceleratorProtocol @@ -177,16 +195,16 @@ def _run(self, accelerator: AcceleratorProtocol) -> None: cfg_dict=cfg_dict, run_dir=run_dir, ) - elif self.mode == "val": + elif self.mode == "eval": if not self.resume_from: - raise ValueError("resume_from is required when mode='val'") # noqa: TRY003 + raise ValueError("resume_from is required when mode='eval'") # noqa: TRY003 self._evaluate( accelerator=accelerator, logger=logger, module_or_module_path=self.resume_from, ) else: - raise NotImplementedError(f"Mode {self.mode} is not implemented") + raise UnknownExperimentModeError(self.mode, allowed=("train", "eval", "help")) def _evaluate(self, accelerator, logger, module_or_module_path, val_dataloader=None) -> float: if isinstance(module_or_module_path, str): diff --git a/tinyexp/examples/pi_exp.py b/tinyexp/examples/pi_exp.py index df5d547..0401e76 100644 --- a/tinyexp/examples/pi_exp.py +++ b/tinyexp/examples/pi_exp.py @@ -1,24 +1,27 @@ -""" -Estimate pi with distributed Monte Carlo sampling, launched by Ray in ``mode="run"``. - -This example shows that tinyexp can drive general distributed programs (no deep -learning, no dataloader): each Ray worker samples random points in the unit square, -counts how many fall inside the unit circle (x*x + y*y <= 1.0), and the counts are -all-reduced to estimate pi = 4 * inside / total. +"""Distributed Monte Carlo estimate of pi: a non-deep-learning workload driven by TinyExp. -Usage:: +Each worker samples points in the unit square, counts those inside the unit circle, and the counts are +all-reduced to estimate pi = 4 * inside / total. No dataloader, no model; also the fastest smoke test. - python -m tinyexp.examples.pi_exp +Modes : run | help +Run : + python -m tinyexp.examples.pi_exp mode=help + python -m tinyexp.examples.pi_exp # default: 2 Ray CPU workers python -m tinyexp.examples.pi_exp pi_cfg.total_samples=100000000 ray_cfg.ray_num_worker=4 + python -m tinyexp.examples.pi_exp launcher=mp pi_cfg.total_samples=100000 + accelerate launch --cpu --num-processes=1 -m tinyexp.examples.pi_exp launcher=mp +Env : none +Needs : tinyexp[pytorch] +Out : ``pi ~= ...`` on stdout (by rank 0, or by the Ray driver); log in /pi_exp/log.txt """ from dataclasses import dataclass, field from typing import Optional -import ray import torch from tinyexp import TinyExp, store_and_run_exp +from tinyexp.exceptions import UnknownExperimentModeError from tinyexp.exp_mixins import LoggerCfgMixin, RayCfgMixin from tinyexp.tiny_engine.accelerator import AcceleratorProtocol @@ -49,7 +52,7 @@ class PiCfg: def run(self) -> None: if self.mode != "run": - raise NotImplementedError(f"Mode {self.mode} is not implemented") + raise UnknownExperimentModeError(self.mode, allowed=("run", "help")) self._run_result: Optional[str] = None from tinyexp.tiny_engine.accelerator import CPUAccelerator @@ -67,9 +70,8 @@ def _run(self, accelerator: AcceleratorProtocol) -> None: pi = self._estimate_pi(accelerator) if accelerator.is_main_process: result = f"pi ~= {pi:.6f} (error={abs(pi - torch.pi):.6f}, samples={self.pi_cfg.total_samples})" - if ray.is_initialized() and ray.get_runtime_context().get_actor_id() is not None: - self._run_result = result - else: + self._run_result = result + if self.launcher != "ray": print(result, flush=True) def get_ray_run_result(self) -> Optional[str]: diff --git a/tinyexp/examples/resnet_exp.py b/tinyexp/examples/resnet_exp.py index 4c13a48..d778376 100644 --- a/tinyexp/examples/resnet_exp.py +++ b/tinyexp/examples/resnet_exp.py @@ -1,3 +1,18 @@ +"""ImageNet ResNet-50 recipe (about 76% top-1 after 90 epochs with DDP); also runs with FSDP or HF Accelerate. + +Modes : train | eval | help (eval requires resume_from=) +Run : + python -m tinyexp.examples.resnet_exp mode=help + python -m tinyexp.examples.resnet_exp # default: Ray, one DDP worker per GPU, Redis cache on + torchrun --nproc-per-node=8 -m tinyexp.examples.resnet_exp launcher=mp redis_cfg.redis_cache_enabled=false + torchrun --nproc-per-node=2 -m tinyexp.examples.resnet_exp launcher=mp accelerator_cfg.accelerator=fsdp redis_cfg.redis_cache_enabled=false + python -m tinyexp.examples.resnet_exp mode=eval resume_from=./output/resnet_exp/best.ckpt +Env : IMAGENET_HOME (or dataloader_cfg.data_root) with ImageFolder train/ and val/ directories +Needs : tinyexp[pytorch] and CUDA GPUs. The Redis cache (default on) needs redis-server: launcher=ray starts it, + otherwise wrap the command with tinyexp-run-with-redis or set redis_cfg.redis_cache_enabled=false +Out : /resnet_exp/{log.txt,last.ckpt,best.ckpt} +""" + import datetime import io import os @@ -15,7 +30,7 @@ from tinyexp import TinyExp, store_and_run_exp from tinyexp.dataset.sampler import InfiniteSampler -from tinyexp.exceptions import UnknownAcceleratorTypeError +from tinyexp.exceptions import UnknownAcceleratorTypeError, UnknownExperimentModeError from tinyexp.exp_mixins import CheckpointCfgMixin, LoggerCfgMixin, RayCfgMixin, RedisCfgMixin, WandbCfgMixin from tinyexp.tiny_engine.accelerator import AcceleratorProtocol @@ -338,16 +353,16 @@ def _run(self, accelerator: AcceleratorProtocol) -> None: cfg_dict=cfg_dict, run_dir=run_dir, ) - elif self.mode == "val": + elif self.mode == "eval": if not self.resume_from: - raise ValueError("resume_from is required when mode='val'") # noqa: TRY003 + raise ValueError("resume_from is required when mode='eval'") # noqa: TRY003 self._evaluate( accelerator=accelerator, logger=logger, module_or_module_path=self.resume_from, ) else: - raise NotImplementedError(f"Mode {self.mode} is not implemented") + raise UnknownExperimentModeError(self.mode, allowed=("train", "eval", "help")) def _evaluate(self, accelerator, logger, module_or_module_path, val_dataloader=None) -> None: if isinstance(module_or_module_path, str): diff --git a/tinyexp/examples/vit_tp_exp.py b/tinyexp/examples/vit_tp_exp.py index dda6805..bbcfc58 100644 --- a/tinyexp/examples/vit_tp_exp.py +++ b/tinyexp/examples/vit_tp_exp.py @@ -1,4 +1,23 @@ -"""DeiT-S experiment with tensor parallelism, ported against the official DeiT. +"""DeiT-S (ViT-S/16) on ImageNet with tensor parallelism, ported against the official DeiT. + +Modes : train | eval | bench | help (eval requires resume_from= or module_cfg.pretrained_from) +Run : + python -m tinyexp.examples.vit_tp_exp mode=help + python -m tinyexp.examples.vit_tp_exp # default: full recipe, 2 Ray GPU workers, TP=2 + # eval-only against the official checkpoint (accuracy cross-check): + python -m tinyexp.examples.vit_tp_exp mode=eval module_cfg.pretrained_from=https://dl.fbaipublicfiles.com/deit/deit_small_patch16_224-cd65a155.pth + # full 300-epoch recipe on 8 GPUs with DDP and the optional Redis byte cache: + python -m tinyexp.examples.vit_tp_exp accelerator_cfg.accelerator=ddp ray_cfg.ray_num_worker=8 redis_cfg.redis_cache_enabled=true redis_cfg.redis_cache_max_memory=300 + # synthetic throughput/memory benchmark, TP vs DDP: + python -m tinyexp.examples.vit_tp_exp mode=bench + python -m tinyexp.examples.vit_tp_exp mode=bench accelerator_cfg.accelerator=ddp + # CPU smoke without ImageNet (2 Ray workers x 2 CPUs, gloo TP=2; batch must fit the 64 fake samples; + # max_train_steps is only checked after each epoch, so this runs the full 8-step epoch): + python -m tinyexp.examples.vit_tp_exp dataloader_cfg.fake_data=true dataloader_cfg.input_size=32 dataloader_cfg.train_batch_size_per_device=8 dataloader_cfg.val_batch_size_per_device=8 module_cfg.img_size=32 module_cfg.num_classes=10 epochs=1 max_train_steps=4 ray_cfg.ray_num_gpus_per_worker=0.0 ray_cfg.ray_num_cpus_per_worker=2 +Env : IMAGENET_HOME (or dataloader_cfg.data_root) with ImageFolder train/ and val/ (not needed for bench/fake_data) +Needs : tinyexp[vit]; 2 CUDA GPUs for the default TP=2 run (8 for the DDP recipe); Redis cache optional +Out : /vit_tp_exp/{deit_train.log,log.txt (JSON metrics per epoch),last.ckpt,best.ckpt}; + details and results in docs/vit_tp.md This module is written by strict cross-checking (对拍) against upstream sources; only formatting-level changes are allowed (see ``docs/vit_tp.md``). The ported sources are: @@ -33,36 +52,6 @@ (8xH200 DDP via rjob, Redis cache on) reached **79.84% top-1 / 94.99% top-5** (paper: 79.8 / ~95.0). -Usage (full recipe, 2 GPUs, TP=2):: - - # ImageNet must use the standard ImageFolder layout with train/ and val/ - # class directories. The same variable is used by the ResNet example. - export IMAGENET_HOME=/path/to/imagenet - - python -m tinyexp.examples.vit_tp_exp - -Eval-only against the official checkpoint (accuracy cross-check):: - - export IMAGENET_HOME=/path/to/imagenet - python -m tinyexp.examples.vit_tp_exp mode=eval module_cfg.pretrained_from= - -Full 300-epoch training on 8 GPUs with DDP and the optional Redis byte cache:: - - export IMAGENET_HOME=/path/to/imagenet - python -m tinyexp.examples.vit_tp_exp accelerator_cfg.accelerator=ddp \\ - ray_cfg.ray_num_worker=8 redis_cfg.redis_cache_enabled=true redis_cfg.redis_cache_max_memory=300 - -Throughput / memory benchmark, e.g. TP vs DDP:: - - python -m tinyexp.examples.vit_tp_exp mode=bench - python -m tinyexp.examples.vit_tp_exp mode=bench accelerator_cfg.accelerator=ddp - -CPU smoke (no ImageNet on disk, ray 2 workers, gloo TP=2):: - - python -m tinyexp.examples.vit_tp_exp dataloader_cfg.fake_data=true dataloader_cfg.input_size=32 \\ - module_cfg.img_size=32 module_cfg.num_classes=10 epochs=1 max_train_steps=4 \\ - ray_cfg.ray_num_gpus_per_worker=0.0 ray_cfg.ray_num_cpus_per_worker=4 - Not ported: ``HybridEmbed``, distilled models, distillation, ``--cosub``, ``--ThreeAugment``, ``--bce-loss``, ``--attn-only``, INAT/CIFAR datasets, submitit, position-embedding interpolation for non-224 finetuning. @@ -101,7 +90,7 @@ from torchvision import datasets, transforms from tinyexp import TinyExp, store_and_run_exp -from tinyexp.exceptions import UnknownAcceleratorTypeError +from tinyexp.exceptions import UnknownAcceleratorTypeError, UnknownExperimentModeError from tinyexp.exp_mixins import CheckpointCfgMixin, LoggerCfgMixin, RayCfgMixin, RedisCfgMixin, WandbCfgMixin from tinyexp.tiny_engine.accelerator import AcceleratorProtocol, TPAccelerator @@ -620,7 +609,9 @@ class RayCfg(RayCfgMixin.RayCfg): ray_num_gpus_per_worker: float = 1.0 ray_num_cpus_per_worker: int = 12 # main process plus 10 dataloader workers - ray_cfg: RayCfg = field(default_factory=RayCfg) + # Field types are qualified: with ``from __future__ import annotations`` they are strings, and Python 3.9's + # ``typing.get_type_hints`` (used by OmegaConf) does not look up names nested in the class. + ray_cfg: VitTpExp.RayCfg = field(default_factory=RayCfg) @dataclass class RedisCfg(RedisCfgMixin.RedisCfg): @@ -629,7 +620,7 @@ class RedisCfg(RedisCfgMixin.RedisCfg): # filesystems (docs/vit_tp.md D7). No effect on fake_data runs. redis_cache_enabled: bool = False - redis_cfg: RedisCfg = field(default_factory=RedisCfg) + redis_cfg: VitTpExp.RedisCfg = field(default_factory=RedisCfg) @dataclass class AcceleratorCfg: @@ -646,7 +637,7 @@ def build_accelerator(self) -> AcceleratorProtocol: return TPAccelerator() raise UnknownAcceleratorTypeError(self.accelerator) - accelerator_cfg: AcceleratorCfg = field(default_factory=AcceleratorCfg) + accelerator_cfg: VitTpExp.AcceleratorCfg = field(default_factory=AcceleratorCfg) @dataclass class ModuleCfg: @@ -685,7 +676,7 @@ def build_module(self) -> VisionTransformer: num_classes=self.num_classes, ) - module_cfg: ModuleCfg = field(default_factory=ModuleCfg) + module_cfg: VitTpExp.ModuleCfg = field(default_factory=ModuleCfg) @dataclass class DataloaderCfg: @@ -782,7 +773,7 @@ def build_train_dataloader( ) else: sampler = torch.utils.data.RandomSampler(dataset) - return torch.utils.data.DataLoader( + loader = torch.utils.data.DataLoader( dataset, sampler=sampler, batch_size=self.train_batch_size_per_device, @@ -791,6 +782,14 @@ def build_train_dataloader( drop_last=True, persistent_workers=self.train_num_workers > 0, ) + if len(loader) == 0: + # docs/vit_tp.md D8: fail fast instead of training zero steps per epoch. + raise ValueError( + f"The train dataloader has no training batch: {len(sampler)} samples per rank < " + f"dataloader_cfg.train_batch_size_per_device={self.train_batch_size_per_device} with " + "drop_last=True. Lower the batch size or provide more data." + ) + return loader def build_val_dataloader( self, accelerator: AcceleratorProtocol, replicate_data: bool @@ -813,7 +812,7 @@ def build_val_dataloader( persistent_workers=self.val_num_workers > 0, ) - dataloader_cfg: DataloaderCfg = field(default_factory=DataloaderCfg) + dataloader_cfg: VitTpExp.DataloaderCfg = field(default_factory=DataloaderCfg) @dataclass class LossCfg: @@ -852,7 +851,7 @@ def build_criterion(self, num_classes: int) -> nn.Module: criterion = nn.CrossEntropyLoss() return criterion - loss_cfg: LossCfg = field(default_factory=LossCfg) + loss_cfg: VitTpExp.LossCfg = field(default_factory=LossCfg) @dataclass class OptimizerCfg: @@ -880,7 +879,7 @@ def build_optimizer( ] return torch.optim.AdamW(parameters, lr=linear_scaled_lr, eps=self.opt_eps) - optimizer_cfg: OptimizerCfg = field(default_factory=OptimizerCfg) + optimizer_cfg: VitTpExp.OptimizerCfg = field(default_factory=OptimizerCfg) @dataclass class LrSchedulerCfg: @@ -919,14 +918,14 @@ def build_lr_scheduler(self, optimizer: torch.optim.Optimizer, epochs: int): lr_scheduler, _ = create_scheduler(sched_args, optimizer) return lr_scheduler - lr_scheduler_cfg: LrSchedulerCfg = field(default_factory=LrSchedulerCfg) + lr_scheduler_cfg: VitTpExp.LrSchedulerCfg = field(default_factory=LrSchedulerCfg) @dataclass class EmaCfg: model_ema: bool = True # official default on; forced off at tp>1 (docs/vit_tp.md D4) model_ema_decay: float = 0.99996 - ema_cfg: EmaCfg = field(default_factory=EmaCfg) + ema_cfg: VitTpExp.EmaCfg = field(default_factory=EmaCfg) @dataclass class BenchCfg: @@ -934,7 +933,7 @@ class BenchCfg: warmup_steps: int = 10 measure_steps: int = 50 - bench_cfg: BenchCfg = field(default_factory=BenchCfg) + bench_cfg: VitTpExp.BenchCfg = field(default_factory=BenchCfg) # ------------------------------ execution part ------------------------------ # @@ -966,7 +965,7 @@ def _run(self, accelerator: AcceleratorProtocol) -> None: elif self.mode == "bench": self._bench(accelerator=accelerator, logger=logger) else: - raise NotImplementedError(f"Mode {self.mode} is not implemented") + raise UnknownExperimentModeError(self.mode, allowed=("train", "eval", "bench", "help")) def _prepare_for_train_or_eval(self, accelerator): """Components shared by train and eval, in official main.py order.""" diff --git a/tinyexp/exceptions.py b/tinyexp/exceptions.py index a3ac0a0..387a5a0 100644 --- a/tinyexp/exceptions.py +++ b/tinyexp/exceptions.py @@ -1,3 +1,18 @@ +"""Errors raised by TinyExp; each keeps the offending value as an attribute. + +- ``UnknownConfigurationKeyError`` (``.key``): ``set_cfg`` got a key the experiment does not define. +- ``UnknownLauncherError`` (``.launcher``): ``launcher`` is not ``mp`` or ``ray``. +- ``UnknownExperimentModeError(mode, allowed=(...))`` (``.mode``): raise from ``run()`` for an unsupported + ``mode``; the message lists ``allowed``, e.g. ``raise UnknownExperimentModeError(self.mode, allowed=("train", "eval", "help"))``. +- ``UnknownAcceleratorTypeError(accelerator)`` (``.accelerator``): raise from ``build_accelerator()`` for an unknown name. +- ``InvalidWorkerCountError`` (``.num_worker``): ``ray_num_worker`` is 0, below -1, or ``-1`` found no resources. +- ``InsufficientCPUError``: the Ray cluster lacks the requested CPU/GPU resources. +- ``CudaNotAvailableError``: a CUDA-only accelerator was built without CUDA. +- ``UnsupportedCheckpointFormatError`` (``.path``): ``load_checkpoint`` got a dict that is not a TinyExp checkpoint. + +Import them from ``tinyexp.exceptions``. +""" + from collections.abc import Sequence from typing import Optional @@ -21,7 +36,7 @@ def __init__(self, num_worker: int) -> None: class UnknownExperimentModeError(ValueError): - def __init__(self, mode: str, allowed: Sequence[str] = ("run", "train", "val", "help")) -> None: + def __init__(self, mode: str, allowed: Sequence[str] = ("run", "train", "eval", "help")) -> None: self.mode = mode self.allowed = tuple(allowed) super().__init__(f"Unknown mode {mode!r}, please set `mode` to one of: {', '.join(self.allowed)}.") diff --git a/tinyexp/exp_mixins/__init__.py b/tinyexp/exp_mixins/__init__.py index ea7874a..4fe61f0 100644 --- a/tinyexp/exp_mixins/__init__.py +++ b/tinyexp/exp_mixins/__init__.py @@ -1,3 +1,9 @@ +"""Opt-in config components for ``TinyExp``; the table of mixins is in ``basic_mixins``. + +Exports: ``LoggerCfgMixin``, ``RayCfgMixin``, ``CheckpointCfgMixin``, ``WandbCfgMixin``, ``RedisCfgMixin`` +(also re-exported from ``tinyexp``). +""" + from .basic_mixins import CheckpointCfgMixin, LoggerCfgMixin, RayCfgMixin, RedisCfgMixin, WandbCfgMixin __all__ = ["CheckpointCfgMixin", "LoggerCfgMixin", "RayCfgMixin", "RedisCfgMixin", "WandbCfgMixin"] diff --git a/tinyexp/exp_mixins/basic_mixins.py b/tinyexp/exp_mixins/basic_mixins.py index be7d048..41588d8 100644 --- a/tinyexp/exp_mixins/basic_mixins.py +++ b/tinyexp/exp_mixins/basic_mixins.py @@ -1,3 +1,27 @@ +"""Opt-in ``XXXCfgMixin`` components; mix only the ones an experiment needs into ``TinyExp``. + +Each mixin adds one config field whose nested dataclass holds both settings and a ``build_*`` helper: + +| mixin | field | use | +| --- | --- | --- | +| `LoggerCfgMixin` | `logger_cfg` | `build_logger(save_dir, distributed_rank)` -> loguru logger | +| `RayCfgMixin` | `ray_cfg` | `launcher=ray`: worker contract and minimal example in `RayCfgMixin`, count/resources in `RayCfg` | +| `CheckpointCfgMixin` | `checkpoint_cfg` | `save_checkpoint` / `load_checkpoint` / RNG state helpers | +| `WandbCfgMixin` | `wandb_cfg` | `build_wandb(accelerator, **wandb_init_kwargs)` if enabled, else `None` | +| `RedisCfgMixin` | `redis_cfg` | Redis byte-cache settings consumed by datasets and launchers | + +Usage: ``class Exp(TinyExp, RayCfgMixin, CheckpointCfgMixin, LoggerCfgMixin)``. To change defaults, +subclass the nested config and redeclare the field, e.g.:: + + @dataclass + class RayCfg(RayCfgMixin.RayCfg): + ray_num_gpus_per_worker: float = 0.0 + + ray_cfg: RayCfg = field(default_factory=RayCfg) + +``torch`` and ``wandb`` are imported lazily inside methods so ``import tinyexp`` stays torch-free. +""" + import os import random import tempfile @@ -33,6 +57,12 @@ class LoggerCfgMixin: @dataclass class LoggerCfg: + """``build_logger(save_dir, distributed_rank=0, filename="log.txt", mode="a")`` -> loguru logger. + + Creates ``save_dir`` (usually ``self.get_run_dir()``). Every rank writes to ``/`` + and stderr: rank 0 at INFO, other ranks at WARNING and above. ``mode="o"`` overwrites the file. + """ + def build_logger( self, save_dir: str, @@ -77,12 +107,65 @@ def _resolve_ray_num_worker( @dataclass class RayCfgMixin: + """Run ``Exp.run()`` on a group of Ray workers when ``launcher=ray``; worker count/resources: ``RayCfg``. + + Every worker calls ``run()`` with ``RANK``/``WORLD_SIZE``/``LOCAL_RANK``/``MASTER_ADDR``/``MASTER_PORT`` set, + so an accelerator built inside ``run()`` is one rank of a ready process group (``rank``, ``world_size``, + ``reduce``). Worker stdout reaches the driver prefixed and interleaved; for one final line, return it from + ``get_ray_run_result`` on rank 0 and the driver prints it after every worker has finished. The same file runs + unchanged with ``launcher=mp`` (one process, ``world_size == 1``), where that method is never called:: + + @dataclass(repr=False) + class Exp(TinyExp, RayCfgMixin): + launcher: str = "ray" + + def run(self) -> None: + from tinyexp.tiny_engine.accelerator import CPUAccelerator + + accelerator = CPUAccelerator() + try: + part = torch.arange(accelerator.rank, 1000, accelerator.world_size).sum().double() + total = accelerator.reduce(part, reduction="sum").item() + self._result = f"total={total:.0f}" if accelerator.is_main_process else None + if self._result and self.launcher != "ray": + print(self._result, flush=True) + finally: + accelerator.destroy() + + def get_ray_run_result(self): + return self._result + + Start it as ``python -m package.module`` rather than ``python path/to/file.py``: the class then reaches the + workers by import path. From ``__main__`` it is pickled by value, which works locally but is fragile + ("not the same object as ..." pickling errors), and its imports must still resolve on every node. + """ + def get_ray_run_result(self) -> Optional[str]: - """Return a short result for the Ray driver to print after the run.""" + """Rank 0's return value is printed by the Ray driver after all workers finish (default ``None``: nothing). + + Only called with ``launcher=ray``; see ``RayCfgMixin`` for the pattern. + """ return None @dataclass class RayCfg: + """Ray worker group used when ``launcher=ray``. + + - ``ray_num_worker``: total workers across the whole Ray cluster, not per node. Each worker is + one rank (``RANK``/``WORLD_SIZE``/``LOCAL_RANK``/``MASTER_ADDR``/``MASTER_PORT`` are set). + ``-1`` fills the currently available resources. + - ``ray_num_cpus_per_worker`` / ``ray_num_gpus_per_worker``: resources of one worker (one + placement-group bundle). Use ``ray_num_gpus_per_worker=0`` for CPU-only runs. + - ``ray_placement_strategy``: Ray placement-group strategy (``PACK``, ``SPREAD``, ...). + - ``ray_placement_timeout_s``: how long to wait for the placement group before failing. + + On a static multi-node cluster (``tinyexp-run-with-ray-cluster``) the same total is spread over + the nodes by Ray; e.g. 2 nodes x 4 GPUs -> ``ray_cfg.ray_num_worker=8 ray_cfg.ray_num_gpus_per_worker=1`` + (GPU examples additionally take ``accelerator_cfg.accelerator=ddp``). Workers may land on any node, + so ``dataloader_cfg.data_root``, ``output_root`` and ``resume_from`` must be paths every node can + reach (shared filesystem or identical per-node layout). + """ + # ---------------- luancher configuration ---------------- # ray_num_worker: int = -1 # Number of Ray workers, -1 means fill available worker resources. ray_num_cpus_per_worker: int = 1 @@ -93,12 +176,21 @@ class RayCfg: @classmethod def run(cls, exp_class: type[Any], experiment_cfg: DictConfig) -> None: # noqa: C901 """ - Run the Ray driver-side orchestration. + Run the Ray driver-side orchestration for ``launcher=ray``. + + Steps (each marked by a ``# ---- n/5`` comment below): + 1. validate ``ray_cfg`` and start Ray if no runtime exists yet; + 2. resolve the worker count (``-1`` -> available resources) and check cluster capacity; + 3. start the Ray-managed Redis cache if ``redis_cfg`` asks for it; + 4. reserve one placement-group bundle per worker and pick the rendezvous address; + 5. create one actor per worker, ``set_cfg`` + ``run`` on all, print rank 0's run result. + Cleanup releases the rendezvous actor, placement group, Redis and the Ray runtime it started. Args: exp_class: Experiment class instantiated by each Ray worker. experiment_cfg: Complete Hydra configuration for exp_class. """ + # ---- 1/5 validate ray_cfg and own the Ray runtime ------------------------------------ # ray_cfg = experiment_cfg.ray_cfg num_cpus_per_worker = int(ray_cfg.ray_num_cpus_per_worker) num_gpus_per_worker = float(ray_cfg.ray_num_gpus_per_worker) @@ -122,6 +214,7 @@ def run(cls, exp_class: type[Any], experiment_cfg: DictConfig) -> None: # noqa: worker_group = [] try: + # ---- 2/5 resolve the worker count and check cluster capacity ---------------------- # cluster_resources = ray.cluster_resources() sizing_resources = ray.available_resources() if ray_cfg.ray_num_worker == -1 else cluster_resources ray_num_worker = _resolve_ray_num_worker( @@ -150,8 +243,9 @@ def run(cls, exp_class: type[Any], experiment_cfg: DictConfig) -> None: # noqa: needed_gpu=needed_gpu, ) + # ---- 3/5 start the Ray-managed Redis cache (no-op unless redis_cfg enables it) ---- # redis_manager = _maybe_start_ray_redis_cache(experiment_cfg) - # -------------------- allocate resources for run ----------------- # + # ---- 4/5 reserve placement-group bundles and choose the rendezvous address -------- # pg = get_placement_group( num_worker=ray_cfg.ray_num_worker, @@ -179,6 +273,7 @@ def run(cls, exp_class: type[Any], experiment_cfg: DictConfig) -> None: # noqa: node_ids=worker_node_ids, ) runtime_envs = [options["runtime_env"]["env_vars"] for options in options_list] + # ---- 5/5 print topology, launch workers, run them, report rank 0's result -------- # print("==> Ray worker topology:", flush=True) node_ranks = {node_id: node_rank for node_rank, node_id in enumerate(dict.fromkeys(worker_node_ids))} topology = sorted( @@ -202,6 +297,7 @@ def run(cls, exp_class: type[Any], experiment_cfg: DictConfig) -> None: # noqa: if result is not None: print(result, flush=True) finally: + # ---- cleanup: release everything this call created, in reverse order -------------- # if rendezvous_actor is not None: with suppress(Exception): ray.kill(rendezvous_actor, no_restart=True) @@ -226,6 +322,18 @@ def run(cls, exp_class: type[Any], experiment_cfg: DictConfig) -> None: # noqa: class WandbCfgMixin: @dataclass class WandbCfg: + """Optional Weights & Biases logging, off by default (``wandb_cfg.enable_wandb=true`` turns it on). + + ``build_wandb(accelerator=None, **kwargs)`` returns ``None`` when disabled. When enabled it returns the + ``wandb`` module on every rank, but calls ``wandb.init(**kwargs)`` (``project=``, ``name=``, ``config=``, ...) + only when ``accelerator.rank == 0`` (or when no accelerator is given), so log from rank 0 only + (``cfg_dict`` below is what ``self.print_cfg(logger)`` returns):: + + wandb = self.wandb_cfg.build_wandb(accelerator=accelerator, project="p", name=self.exp_name, config=cfg_dict) + if wandb is not None and accelerator.is_main_process: + wandb.log({"loss": loss}) + """ + enable_wandb: bool = False def build_wandb(self, accelerator=None, **kwargs): @@ -241,16 +349,32 @@ def build_wandb(self, accelerator=None, **kwargs): @dataclass class RedisCfgMixin: - """Supplies :attr:`redis_cfg` plus thin Ray entrypoints on the *experiment* object (see below).""" + """Adds ``redis_cfg``: settings of the Redis byte cache (see ``RedisCfg``).""" @dataclass class RedisCfg: - """ - Hydra-overridable options live in dataclass fields below. - - Single-machine runs use standalone Redis. Multi-machine runs can use Redis Cluster. - Redis lifecycle is owned by the launcher: ``tinyexp-run-with-redis`` for - wrapper-based runs, or Ray-managed Redis for Ray launches. + """Settings of the Redis byte cache used by datasets (e.g. cached ImageNet files in ``resnet_exp``). + + - ``redis_cache_enabled`` (default ``True``): whether the experiment should use the cache. + - ``redis_cluster_host`` / ``redis_cluster_ports``: where Redis listens; standalone mode runs one + ``redis-server`` per port and shards keys across them. + - ``redis_cache_max_memory``: total memory limit in **GB** (default 160, sized for ImageNet), split + evenly across the Redis processes. + - ``redis_cluster_startup_timeout_s``: how long to wait for Redis to come up. + - ``redis_rendezvous_world_size``: ``1`` standalone Redis (with ``launcher=ray``: one per Ray node); + ``-1`` Ray-managed Redis Cluster over the alive Ray nodes; ``>1`` externally managed Redis Cluster. + + TinyExp never starts Redis from inside ``run()``; the launcher owns its lifecycle: + ``launcher=ray`` starts and stops it automatically, ``tinyexp-run-with-redis -- `` does so for + any other launch (and injects the matching ``redis_cfg.*`` overrides), or point these fields at a + Redis you manage yourself. The wrapper refuses a ``launcher=ray`` command, which already owns its + Redis. Inside the experiment connect with:: + + from tinyexp.utils.redis_utils import RedisClientManager + cache = RedisClientManager(cfg.redis_cluster_host, list(cfg.redis_cluster_ports), + redis_world_size=int(cfg.redis_rendezvous_world_size)) + if cache.safe_get(key) is None: # miss, or Redis unavailable + cache.safe_set(key, value_bytes) # False when it could not be stored """ redis_cache_enabled: bool = True @@ -273,6 +397,60 @@ class RedisCfg: class CheckpointCfgMixin: @dataclass class CheckpointCfg: + """Atomic ``torch.save``/``torch.load`` of training state; the caller decides where and when. + + Payload written by ``save_checkpoint`` to ``/``:: + + { + "epoch": int | None, "global_step": int | None, "best_metric": float | None, + "meta": {"exp_name": str, "exp_class": str, "saved_at": ISO-8601 UTC str}, # names are "" unless passed + # present only when the object was passed: + "model_state_dict", "optimizer_state_dict", "scheduler_state_dict", + "scaler_state_dict", "extra_state": dict, + } + + ``load_checkpoint(path, model=..., optimizer=..., scheduler=..., scaler=...)`` loads every object + you pass (missing key -> ``KeyError``) and returns the whole dict, so read ``epoch``/``best_metric``/ + ``extra_state`` from it. A dict with none of ``epoch``/``global_step``/``best_metric``/``meta``/ + ``extra_state``/``model_state_dict`` raises ``UnsupportedCheckpointFormatError`` + (``tinyexp.exceptions``); a payload that is not a dict raises ``TypeError``. To inspect a file by + hand use ``torch.load(path, weights_only=False)``. + + Conventions used by the examples:: + + if accelerator.is_main_process: # save on rank 0 only + # RNG state is per process: only a single-process run can restore it faithfully + rng = {"rng_state": self.checkpoint_cfg.capture_rng_state()} if accelerator.world_size == 1 else None + self.checkpoint_cfg.save_checkpoint( + run_dir=self.get_run_dir(), name=self.checkpoint_cfg.last_ckpt_name, + model=accelerator.unwrap_model(module), # DDP/CPU: unwrap first + optimizer=optimizer, scheduler=lr_scheduler, epoch=epoch, + exp_name=self.exp_name, exp_class=self.exp_class, # fills ckpt["meta"] + extra_state=rng, + ) + ckpt = self.checkpoint_cfg.load_checkpoint( # every rank loads + self.resume_from, model=accelerator.unwrap_model(module), optimizer=optimizer, + scheduler=lr_scheduler, map_location=accelerator.device, + ) + start_epoch, best_metric = ckpt["epoch"] + 1, ckpt["best_metric"] + rng_state = (ckpt.get("extra_state") or {}).get("rng_state") + if rng_state is not None and accelerator.world_size == 1: + self.checkpoint_cfg.restore_rng_state(rng_state) # python/numpy/torch(/cuda) RNG + + FSDP (``FSDPAccelerator``, see ``examples/resnet_exp.py``): + + - Save: every rank calls the collective ``model_sd, optim_sd = + accelerator.collect_checkpoint_state(module, optimizer)``; rank 0 receives full CPU state dicts + and the other ranks receive empty dicts. Rank 0 then passes them as ``model=``/``optimizer=`` + wrapped in any object whose ``state_dict()`` returns them (only ``state_dict()`` is called). + - Load: every rank calls ``load_checkpoint`` with the *prepared* (still wrapped) module and the + optimizer returned by ``accelerator.prepare(model, optimizer)``; their ``load_state_dict`` is patched + to accept the full, unsharded state dicts, so the same file works for DDP and FSDP. + + The checkpoint path is therefore ``//{last,best}.ckpt``. For an evaluation + mode, require ``resume_from`` yourself (raise if empty) and call ``load_checkpoint`` with ``model=`` only. + """ + last_ckpt_name: str = "last.ckpt" best_ckpt_name: str = "best.ckpt" diff --git a/tinyexp/tiny_engine/__init__.py b/tinyexp/tiny_engine/__init__.py new file mode 100644 index 0000000..72ce991 --- /dev/null +++ b/tinyexp/tiny_engine/__init__.py @@ -0,0 +1,6 @@ +"""Training-engine pieces that experiments call from inside ``Exp.run()``. + +Currently a single subpackage: ``tinyexp.tiny_engine.accelerator`` (CPU / DDP / FSDP / TP / Hugging Face +backends behind one ``AcceleratorProtocol``). Import from there, e.g. +``from tinyexp.tiny_engine.accelerator import AcceleratorProtocol, DDPAccelerator``. +""" diff --git a/tinyexp/tiny_engine/accelerator/__init__.py b/tinyexp/tiny_engine/accelerator/__init__.py index 839a7b1..286ccd5 100644 --- a/tinyexp/tiny_engine/accelerator/__init__.py +++ b/tinyexp/tiny_engine/accelerator/__init__.py @@ -1,3 +1,20 @@ +"""Interchangeable training backends implementing ``AcceleratorProtocol`` (see ``base_accelerator``). + +| class | backend / device | mixed_precision | when to use | +| --- | --- | --- | --- | +| `CPUAccelerator` | gloo, CPU | none, bf16 | debugging, CI, CPU-only training (DDP when world_size > 1) | +| `DDPAccelerator` | nccl, CUDA required | none, fp16, bf16 | standard multi-GPU data parallel | +| `FSDPAccelerator` | nccl, CUDA required | none, fp16, bf16 | shard params/grads/optimizer state (FULL_SHARD) | +| `TPAccelerator` | nccl on CUDA, else gloo | n/a | tensor parallel, tp_size == world_size (`vit_tp_exp`) | +| `HFAccelerator` | Hugging Face Accelerate | no, fp16, bf16 (`no`, not `none`) | when launching with `accelerate launch` | + +All read ``RANK``/``WORLD_SIZE``/``LOCAL_RANK`` from the environment, so construct them inside the +worker (inside ``Exp.run()``), never in the driver. ``HFAccelerator`` exists only when ``accelerate`` is +installed. Experiments usually select one from config, e.g. ``accelerator_cfg.accelerator=ddp``. + +Exports: the five classes above, ``AcceleratorProtocol``, ``BaseAccelerator``. +""" + # flake8: noqa: F401, F403 import contextlib as _contextlib diff --git a/tinyexp/tiny_engine/accelerator/base_accelerator.py b/tinyexp/tiny_engine/accelerator/base_accelerator.py index c9b2e67..57607d4 100644 --- a/tinyexp/tiny_engine/accelerator/base_accelerator.py +++ b/tinyexp/tiny_engine/accelerator/base_accelerator.py @@ -1,3 +1,46 @@ +"""The accelerator contract that training loops are written against. + +Write the loop against ``AcceleratorProtocol`` and pick the backend from config; the same loop then +runs on CPU, DDP, FSDP, TP or Hugging Face Accelerate. Grouped by when you call them: + +- Topology (attributes): ``rank``, ``world_size``, ``local_rank``, ``device``, ``sync_gradients``; + properties ``is_main_process`` (rank 0) and ``is_local_main_process`` (local rank 0). +- Setup, once: ``prepare(model, optimizer=None)`` -> ``model`` or ``(model, optimizer)``; + or separately ``prepare_model`` / ``prepare_optimizer``. Build the optimizer on the *unprepared* + model's parameters, then prepare both. +- Every step: ``with accelerator.autocast(): loss = ...`` (optional, prepared models already + autocast), ``accelerator.backward(loss)``, ``accelerator.optimizer_step(optimizer)`` + (handles fp16 loss scaling; use it instead of ``optimizer.step()``). +- Cross-rank: ``reduce(tensor, reduction="sum"|"mean")``, ``wait_for_everyone()``, ``print(...)`` + (rank 0 only). ``reduce`` returns a new tensor on the input's device and never modifies the input (with + ``world_size == 1`` it may return the input itself). The result keeps the input dtype, except that "mean" + of an integer tensor becomes floating when ``world_size > 1``; pass floating tensors to avoid that. +- Gradient clipping is not part of the protocol: ``DDPAccelerator``/``FSDPAccelerator``/``HFAccelerator`` provide + ``accelerator.clip_grad_norm_(parameters, max_norm)`` (unscales fp16 gradients first; FSDP reduces the norm + across shards) and it must be used there. ``CPUAccelerator``/``TPAccelerator`` do not define it; they never + scale the loss, so ``torch.nn.utils.clip_grad_norm_`` is correct with them. +- Checkpointing: ``unwrap_model(model)`` returns the bare ``nn.Module`` for ``state_dict()``. +- Teardown: ``destroy()`` in a ``finally`` block; it is idempotent. + +Typical loop:: + + accelerator = DDPAccelerator(mixed_precision="bf16") + try: + model, optimizer = accelerator.prepare(model, torch.optim.SGD(model.parameters(), lr=0.1)) + for x, y in loader: + loss = criterion(model(x.to(accelerator.device)), y.to(accelerator.device)) + optimizer.zero_grad() + accelerator.backward(loss) + accelerator.optimizer_step(optimizer) + mean_loss = accelerator.reduce(loss.detach(), reduction="mean") + finally: + accelerator.destroy() + +``BaseAccelerator`` (the in-repo implementations' base class) additionally provides ``gather``, +``unscale_gradients`` and reads ``RANK``/``WORLD_SIZE``/``LOCAL_RANK``/``MASTER_ADDR``/``MASTER_PORT`` +from the environment, which ``torchrun`` and ``launcher=ray`` both set. +""" + from __future__ import annotations import abc diff --git a/tinyexp/tiny_engine/accelerator/cpu_accelerator.py b/tinyexp/tiny_engine/accelerator/cpu_accelerator.py index 71ff008..cea28e6 100644 --- a/tinyexp/tiny_engine/accelerator/cpu_accelerator.py +++ b/tinyexp/tiny_engine/accelerator/cpu_accelerator.py @@ -1,3 +1,8 @@ +"""CPU accelerator: gloo process group when ``WORLD_SIZE > 1``, optional bf16 autocast. + +Usage: ``CPUAccelerator(mixed_precision="none" | "bf16")``. +""" + import os import socket from contextlib import suppress @@ -11,8 +16,11 @@ class CPUAccelerator(BaseAccelerator): - """ - CPU accelerator for distributed training. + """CPU accelerator: no device placement beyond CPU, gloo process group when ``WORLD_SIZE > 1``. + + ``prepare`` wraps the model in DDP only when ``world_size > 1``; ``backward``/``optimizer_step`` are the + plain PyTorch calls; ``reduce`` all-reduces over gloo. ``mixed_precision="bf16"`` enables CPU bf16 + autocast inside the prepared model's forward (outputs are cast back to fp32). """ def __init__(self, mixed_precision: str = "none") -> None: diff --git a/tinyexp/tiny_engine/accelerator/ddp_accelerator.py b/tinyexp/tiny_engine/accelerator/ddp_accelerator.py index 6143aeb..fc3bb69 100644 --- a/tinyexp/tiny_engine/accelerator/ddp_accelerator.py +++ b/tinyexp/tiny_engine/accelerator/ddp_accelerator.py @@ -1,3 +1,10 @@ +"""CUDA DistributedDataParallel accelerator over NCCL, with fp16 (GradScaler) or bf16 autocast. + +Usage: ``DDPAccelerator(mixed_precision="none" | "fp16" | "bf16")``; one process per GPU, launched by +``torchrun`` (``launcher=mp``) or Ray workers (``launcher=ray``). Raises ``CudaNotAvailableError`` without CUDA. +Extra helpers: ``clip_grad_norm_`` (unscales fp16 first), ``dump_model_to_state_dict``. +""" + import torch import torch.distributed as dist from torch import nn diff --git a/tinyexp/tiny_engine/accelerator/hf_accelerator.py b/tinyexp/tiny_engine/accelerator/hf_accelerator.py index b37af2a..73a3177 100644 --- a/tinyexp/tiny_engine/accelerator/hf_accelerator.py +++ b/tinyexp/tiny_engine/accelerator/hf_accelerator.py @@ -1,9 +1,22 @@ +"""``HFAccelerator``: Hugging Face ``accelerate.Accelerator`` adapted to ``AcceleratorProtocol``. + +Defined only when ``accelerate`` is installed (a warning is emitted otherwise). See ``HFAccelerator``. +""" + try: import contextlib from accelerate import Accelerator class HFAccelerator(Accelerator): + """accelerate's ``Accelerator`` plus ``rank``/``world_size``/``local_rank``, ``optimizer_step`` and an + idempotent ``destroy``. Pair with ``accelerate launch ... launcher=mp``. + + Constructor arguments are accelerate's, e.g. ``HFAccelerator(mixed_precision="bf16")``. accelerate spells + full precision ``"no"`` (``"no" | "fp16" | "bf16"``) and rejects ``"none"``, so map a tinyexp-style + ``mixed_precision=none`` config value to ``"no"`` (as ``resnet_exp`` does). + """ + def __init__(self, *args, **kwargs): super().__init__(*args, **kwargs) self._destroyed = False diff --git a/tinyexp/utils/__init__.py b/tinyexp/utils/__init__.py index e69de29..d287189 100644 --- a/tinyexp/utils/__init__.py +++ b/tinyexp/utils/__init__.py @@ -0,0 +1,7 @@ +"""Internal helpers; experiments normally use them through the mixins. + +- ``log_utils.tiny_logger_setup``: loguru file+console logger (used by ``LoggerCfgMixin``). +- ``model_utils.update_ema``: in-place EMA update of a model copy. +- ``ray_utils``: placement groups, worker env vars and rendezvous for ``RayCfgMixin``. +- ``redis_utils``: Redis client/cluster managers behind ``RedisCfgMixin`` and ``tinyexp-run-with-redis``. +""" diff --git a/tinyexp/utils/log_utils.py b/tinyexp/utils/log_utils.py index 0e107f5..7c8011d 100644 --- a/tinyexp/utils/log_utils.py +++ b/tinyexp/utils/log_utils.py @@ -1,3 +1,9 @@ +"""``tiny_logger_setup(save_dir, distributed_rank=0, filename="log.txt", mode="a")`` -> loguru logger. + +Rank 0 logs INFO to ``/`` and stderr; other ranks log WARNING and above only. +``mode="o"`` overwrites the file. Prefer ``LoggerCfgMixin.LoggerCfg.build_logger``, which also creates ``save_dir``. +""" + import os import sys from sys import stderr @@ -10,7 +16,7 @@ def tiny_logger_setup(save_dir: str, distributed_rank: int = 0, filename: str = Args: save_dir(str): loaction to save log file distributed_rank(int): device rank when multi-gpu environment - mode(str): log file write mode, `append` or `override`. default is `a`. + mode(str): log file write mode: `a` appends (default), `o` overwrites an existing file. Return: logger instance. """ diff --git a/tinyexp/utils/model_utils.py b/tinyexp/utils/model_utils.py index 4090c5a..0760093 100644 --- a/tinyexp/utils/model_utils.py +++ b/tinyexp/utils/model_utils.py @@ -1,3 +1,8 @@ +"""``update_ema(ema_model, model, decay=0.9999)``: ``ema = ema * decay + model * (1 - decay)`` in place. + +Parameter names are matched after stripping ``module.``, so ``model`` may be DDP-wrapped. +""" + from collections import OrderedDict import torch diff --git a/tinyexp/utils/ray_utils.py b/tinyexp/utils/ray_utils.py index 7c56f66..1ed2791 100644 --- a/tinyexp/utils/ray_utils.py +++ b/tinyexp/utils/ray_utils.py @@ -1,3 +1,11 @@ +"""Ray placement and rendezvous primitives used by ``RayCfgMixin.RayCfg.run`` (``launcher=ray``). + +Not needed by experiment code. Main pieces: ``get_placement_group`` (one bundle per worker), +``get_placement_group_node_ids``, ``get_num_worker_options`` (per-worker actor options carrying +``RANK``/``WORLD_SIZE``/``LOCAL_RANK``/``MASTER_ADDR``/``MASTER_PORT``), ``start_ray_rendezvous_store`` +(TCPStore actor for multi-worker runs) and ``get_network_config``. +""" + from __future__ import annotations import os diff --git a/tinyexp/utils/redis_utils.py b/tinyexp/utils/redis_utils.py index f6448f0..a7c00d2 100644 --- a/tinyexp/utils/redis_utils.py +++ b/tinyexp/utils/redis_utils.py @@ -1,3 +1,16 @@ +"""Redis lifecycle and client helpers behind ``RedisCfgMixin`` and ``tinyexp-run-with-redis``. + +- ``RedisClientManager(host, ports, redis_world_size)``: clients that shard keys over standalone Redis + ports or connect to a Redis Cluster; used by Redis-cached datasets (see ``examples/resnet_exp.py``). +- ``RedisClusterManager``: starts/stops local ``redis-server`` processes (standalone shards or Redis Cluster); + use it as a context manager (``with RedisClusterManager(...):``). +- ``RayRedisClusterManager``: the same, one Redis per alive Ray node, for ``launcher=ray``. +- ``wait_for_redis_cluster``; errors ``RedisClusterConfigError`` and ``RedisClusterStartupError``. + +Managers only stop the processes they started; a port already served by another process makes startup fail +(``RedisClusterStartupError``) instead of being taken over. +""" + from __future__ import annotations import contextlib @@ -41,6 +54,17 @@ def __init__( class RedisClientManager: + """Best-effort Redis client for caches: ``RedisClientManager(host, ports, redis_world_size=1)``. + + ``redis_world_size <= 1`` connects to one standalone Redis per port and shards keys across them; + ``> 1`` connects to a Redis Cluster. Build it from ``redis_cfg`` (``redis_cluster_host``, + ``list(redis_cluster_ports)``, ``redis_rendezvous_world_size``). Keys and values are ``bytes``/``str``. + + Nothing here raises: a failed connection is printed and leaves the manager without clients, after which + ``safe_get(key)`` returns ``None`` and ``safe_set(key, value)`` returns ``False``. Treat ``None`` as a + cache miss and fall back to the original data source. + """ + def __init__( self, redis_host: str, @@ -102,6 +126,7 @@ def _redis_client_for_key(self, key: Any) -> Any | None: return self.redis_clients[shard_index] def safe_get(self, key: Any) -> Any | None: + """Return the stored bytes, or ``None`` on a miss or any Redis error.""" redis_client = self._redis_client_for_key(key) if redis_client is None: return None @@ -111,6 +136,7 @@ def safe_get(self, key: Any) -> Any | None: return None def safe_set(self, key: Any, value: Any) -> bool: + """Store ``value``; return ``False`` when it could not be stored (e.g. Redis down or out of memory).""" redis_client = self._redis_client_for_key(key) if redis_client is None: return False