feat: agent-facing contracts: docstring index, run headers, verified doc commands

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.<name>` |
+| 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 <module> 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
+  `<output_root>/<exp_name>/{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).

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
 

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, <mixins>)` 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.<name>` (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.