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]. [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

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.