worktree-env-setup

---
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`):

```bash
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]`.
  `[tests]` omits FlashAttention on aarch64, so its transitive CuTeDSL pin does not apply
  there or to linear-only installs. When updating CuTeDSL, run
  `pytest -n 6 test/test_kda_bwd_wy_compile.py` to catch NVVM binding changes without a
  Blackwell GPU, then validate forward/backward numerics on supported hardware.
- 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]` and `[cudnn]` are declared conflicting extras in `pyproject.toml`; keep them in
  separate environments. For cuDNN worktrees, install `-e '.[cudnn,dev]' pytest pytest-xdist`
  instead; cuDNN 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

```bash
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.

More agent context in meta-pytorch/attention-gym

5 other files this repository gives its agents.