| name | worktree-env-setup |
| description | Sets up an isolated per-worktree Python environment for attention-gym development using nightly PyTorch and the CI-mirroring uv flow. Use when creating a new git worktree or when a worktree lacks a local .venv. |
Worktree Environment Setup
Each attention-gym worktree gets its own .venv so editable installs, concurrent
agents, and test runs never cross-import another checkout. Never reuse a shared
env's editable install across worktrees, and never ln -s another worktree's
.venv as a shortcut: an editable install is a .pth file naming one checkout,
so a shared env makes every other worktree run that checkout's sources. A fresh
uv venv plus hard-linked wheels costs seconds; a wrong import costs hours. If
.venv already exists as a symlink, rm .venv and rebuild it below.
Setup
From the worktree root (mirrors .github/workflows/test.yml):
uv venv --python 3.13
source .venv/bin/activate
uv pip install --pre torch --index-url https://download.pytorch.org/whl/nightly/cu132
uv pip install --prerelease allow -e '.[tests,linear,dev]'
Notes:
- uv hard-links wheels from its cache, so after the first nightly download this
takes seconds and costs almost no extra disk per worktree.
- Activate
.venv before installing so an already-active foreign environment is not modified.
--prerelease allow is required for the flash-attn-4 beta in [tests], but it also lets
the open [linear] bound resolve to nvidia-cutlass-dsl dev releases (4.8.0.dev0 breaks
tcgen05_mma_ws(..., mma_kind=) in the fused CuTeDSL backward). Pin it afterwards:
uv pip install "nvidia-cutlass-dsl[cu13]==4.7.1".- A
.venv symlink into another worktree is not isolation: its editable .pth still
points at that worktree, so pytest imports the other checkout's attn_gym. Replace it
with a real per-worktree env.
- Do not use
uv sync/uv.lock: nightly torch churns daily and CI uses the
imperative uv pip flow above, not a lockfile.
- Drop
[linear] if CuTeDSL/TVM-FFI kernels are not needed (CPU-only work).
[tests] currently brings FlashAttention's CuTeDSL 4.6 pin and cannot be combined with the
CuTeDSL 4.7+ [mega] extra. For Mega worktrees, install -e '.[mega,dev]' pytest pytest-xdist
instead; Mega tests import-skip optional FlashAttention coverage.
Running commands
Prefer the worktree's own interpreter — either activate .venv first, or use
uv run --no-sync pytest test (matches CI exactly). Never invoke a Python from
another worktree or a shared ~/.venvs/* env for attn_gym imports.
Verifying isolation
cd /tmp && python -c "import attn_gym; print(attn_gym.__file__)"
The printed path must be inside the current worktree. If it points at another
checkout, the editable install is wrong — rerun the -e '.[tests,linear,dev]'
install from this worktree root.
Run the check from outside the repo root. From the root, python -c puts the
current directory first on sys.path and masks a wrong editable install, while
python agent_space/script.py and pytest (whose test/ has no __init__.py)
put the script directory first and silently import the other checkout. A
.venv symlinked to another worktree's env fails exactly this way: edits appear
to have no effect because the kernels compile from the other tree.
