Skip to content

Commit d89b51e

Browse files
YWHyukYWHyuk
authored
Merge pull request #229 from PSAL-POSTECH/feature/worktree-setup-script
[Scripts] Add worktree setup script and docs
2 parents a31c756 + 78ec479 commit d89b51e

2 files changed

Lines changed: 211 additions & 0 deletions

File tree

‎docs/worktrees.md‎

Lines changed: 112 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,112 @@
1+
# Parallel worktrees
2+
3+
Quick setup for working on multiple branches at the same time (e.g. one for
4+
feature work, one for code review, one for a bug fix) without thrashing a
5+
single checkout. Container-dedicated: paths assume the
6+
`ghcr.io/psal-postech/torchsim-ci` layout.
7+
8+
## Why a script
9+
10+
Three things have to line up for parallel worktrees to actually work in this
11+
repo:
12+
13+
1. **Worktree-scoped env vars.** `PyTorchSimFrontend/extension_config.py`
14+
anchors output / log / config paths on `TORCHSIM_DIR`. Without an override
15+
every worktree dumps into the same `outputs/` and `togsim_results/`.
16+
2. **`PYTHONPATH` override.** `pip install -e PyTorchSimDevice` writes a
17+
single editable record into conda's site-packages that points at one
18+
worktree. The override forces `import torch_openreg` to resolve to the
19+
active worktree's code and `.so` first.
20+
3. **Branch tracking.** `git worktree add` from a remote ref sets upstream to
21+
that ref, so `git push` would target `develop`. The script unsets it so
22+
first `git push -u origin <branch>` creates the right remote branch.
23+
24+
The script bundles all three.
25+
26+
## Create a worktree
27+
28+
```bash
29+
scripts/setup_worktree.sh <purpose> [base-ref]
30+
```
31+
32+
`<purpose>` becomes the branch suffix and the dir suffix:
33+
34+
| Command | Worktree dir | Branch |
35+
|---|---|---|
36+
| `setup_worktree.sh feature` | `/workspace/PyTorchSim-feature` | `feature/scratch` |
37+
| `setup_worktree.sh review` | `/workspace/PyTorchSim-review` | `refactor/scratch` (rename after) |
38+
| `setup_worktree.sh bugfix/issue-198` | `/workspace/PyTorchSim-bugfix` | `bugfix/issue-198` |
39+
| `setup_worktree.sh feature origin/master` | `/workspace/PyTorchSim-feature` | `feature/scratch` (off master) |
40+
41+
Default base is `origin/develop` (per `CONTRIBUTING.md`).
42+
43+
## Activate
44+
45+
```bash
46+
cd /workspace/PyTorchSim-bugfix
47+
source .envrc
48+
# → Activated worktree: /workspace/PyTorchSim-bugfix
49+
# → prompt: [torchsim:PyTorchSim-bugfix] ...
50+
```
51+
52+
`.envrc` is local to each worktree and not committed. Re-source it whenever
53+
you open a new shell in that worktree.
54+
55+
## Build once per worktree
56+
57+
The compiled `_C.cpython-*.so` lives under `PyTorchSimDevice/torch_openreg/`
58+
and is not shared across worktrees. After activation:
59+
60+
```bash
61+
(cd PyTorchSimDevice && python setup.py build_ext --inplace)
62+
```
63+
64+
Use `build_ext --inplace` instead of `pip install -e` so the editable
65+
record in `/opt/conda/lib/python3.11/site-packages` keeps pointing at
66+
whichever worktree it already pointed at — `PYTHONPATH` from `.envrc` does
67+
the per-worktree routing. (Running `pip install -e` again rewrites that
68+
record and will pin "default" Python to the new worktree.)
69+
70+
## What the env looks like
71+
72+
Worktree-scoped (auto-set by `.envrc`):
73+
74+
| Var | Value |
75+
|---|---|
76+
| `TORCHSIM_DIR` | `$PWD` of the worktree |
77+
| `TORCHSIM_DUMP_PATH` | `$PWD/outputs` |
78+
| `TORCHSIM_LOG_PATH` | `$PWD/togsim_results` |
79+
| `TOGSIM_CONFIG` | `$PWD/configs/systolic_ws_128x128_c1_simple_noc_tpuv3.yml` |
80+
| `PYTHONPATH` | `$PWD/PyTorchSimDevice:$PWD:$PYTHONPATH` |
81+
82+
Shared (container-dedicated, set the same in every `.envrc`):
83+
84+
| Var | Value |
85+
|---|---|
86+
| `GEM5_PATH` | `/gem5/release/gem5.opt` |
87+
| `TORCHSIM_LLVM_PATH` | `/riscv-llvm/bin` |
88+
| `RISCV` | `/workspace/riscv` |
89+
90+
## Cleanup
91+
92+
```bash
93+
git worktree remove /workspace/PyTorchSim-feature
94+
git branch -D feature/scratch # if you do not want to keep the branch
95+
```
96+
97+
`git worktree list` shows the current set.
98+
99+
## Gotchas
100+
101+
- **Do not commit `.envrc`.** It is per-worktree state. Add to your
102+
personal global gitignore if needed.
103+
- **Editable install conflict.** If you run `pip install -e PyTorchSimDevice`
104+
in worktree A, then again in worktree B, Python's default `import
105+
torch_openreg` flips to B. With `PYTHONPATH` from `.envrc` this still
106+
resolves correctly in either worktree's shell, but a shell with no
107+
`.envrc` sourced will see whichever was installed last.
108+
- **TOGSim FIFO files.** `/tmp/togsim_fifo_<pid>` is keyed on PID, not on
109+
worktree — concurrent runs from different worktrees do not collide.
110+
- **`_C.cpython-*.so` rebuild on PyTorch update.** If you `pip install`
111+
a different torch version in the conda env, every worktree's `.so` is
112+
stale; rebuild each.

‎scripts/setup_worktree.sh‎

Lines changed: 99 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,99 @@
1+
#!/usr/bin/env bash
2+
# Create a sibling git worktree for parallel work and wire up per-worktree env.
3+
#
4+
# Container-dedicated: assumes the ghcr.io/psal-postech/torchsim-ci container
5+
# layout. Shared binaries (gem5, LLVM, riscv toolchain) live at known paths
6+
# and are NOT duplicated per worktree.
7+
#
8+
# Usage:
9+
# scripts/setup_worktree.sh <purpose> [base-ref]
10+
#
11+
# Examples:
12+
# scripts/setup_worktree.sh feature # off origin/develop, branch feature/scratch
13+
# scripts/setup_worktree.sh bugfix/issue-198 # off origin/develop, branch bugfix/issue-198
14+
# scripts/setup_worktree.sh review origin/master # off origin/master, branch review/scratch
15+
#
16+
# Result:
17+
# /workspace/PyTorchSim-<basename(purpose)> new worktree
18+
# /workspace/PyTorchSim-<...>/.envrc per-worktree env (source it)
19+
#
20+
# After creation:
21+
# cd /workspace/PyTorchSim-<...>
22+
# source .envrc
23+
# (cd PyTorchSimDevice && python setup.py build_ext --inplace) # build the .so once
24+
set -euo pipefail
25+
26+
if [[ $# -lt 1 || $# -gt 2 ]]; then
27+
sed -n '3,20p' "$0"
28+
exit 1
29+
fi
30+
31+
PURPOSE="$1"
32+
BASE_REF="${2:-origin/develop}"
33+
34+
# Branch name: use the purpose as-is if it already has a slash, else append /scratch.
35+
if [[ "$PURPOSE" == */* ]]; then
36+
BRANCH="$PURPOSE"
37+
SUFFIX="${PURPOSE%%/*}" # before first slash, used in dir name
38+
else
39+
BRANCH="${PURPOSE}/scratch"
40+
SUFFIX="$PURPOSE"
41+
fi
42+
43+
REPO_ROOT="$(git -C "$(dirname "${BASH_SOURCE[0]}")/.." rev-parse --show-toplevel)"
44+
PARENT_DIR="$(dirname "$REPO_ROOT")"
45+
WT_DIR="${PARENT_DIR}/$(basename "$REPO_ROOT")-${SUFFIX}"
46+
47+
if [[ -e "$WT_DIR" ]]; then
48+
echo "error: $WT_DIR already exists" >&2
49+
exit 1
50+
fi
51+
52+
# Make sure base-ref is up to date if it's a remote ref.
53+
if [[ "$BASE_REF" == origin/* ]]; then
54+
git -C "$REPO_ROOT" fetch origin "${BASE_REF#origin/}" --depth=1
55+
fi
56+
57+
git -C "$REPO_ROOT" worktree add "$WT_DIR" -b "$BRANCH" "$BASE_REF"
58+
59+
# Default `worktree add` from a remote ref sets upstream to that remote ref,
60+
# which means `git push` would target develop/master. Unset so the new branch
61+
# pushes to its own name on first `git push -u origin <branch>`.
62+
git -C "$WT_DIR" branch --unset-upstream || true
63+
64+
# Per-worktree env. Container-dedicated paths for shared binaries.
65+
cat > "$WT_DIR/.envrc" <<'ENVRC'
66+
#!/usr/bin/env bash
67+
# Source this from the worktree root: source .envrc
68+
_self="$(cd "$(dirname "${BASH_SOURCE[0]:-$0}")" && pwd)"
69+
70+
# Worktree-scoped: override defaults from PyTorchSimFrontend/extension_config.py
71+
export TORCHSIM_DIR="$_self"
72+
export TORCHSIM_DUMP_PATH="$_self/outputs"
73+
export TORCHSIM_LOG_PATH="$_self/togsim_results"
74+
export TOGSIM_CONFIG="$_self/configs/systolic_ws_128x128_c1_simple_noc_tpuv3.yml"
75+
76+
# Make `import torch_openreg` resolve to THIS worktree's .so first,
77+
# overriding the conda-wide editable install that points at the main worktree.
78+
export PYTHONPATH="$_self/PyTorchSimDevice:$_self:${PYTHONPATH:-}"
79+
80+
# Container-dedicated shared binaries.
81+
export GEM5_PATH="/gem5/release/gem5.opt"
82+
export TORCHSIM_LLVM_PATH="/riscv-llvm/bin"
83+
export RISCV="/workspace/riscv"
84+
85+
# Prompt hint so you do not lose track of which worktree this shell is on.
86+
export PS1="[torchsim:$(basename "$_self")] ${PS1:-\\w\\$ }"
87+
88+
unset _self
89+
echo "Activated worktree: $TORCHSIM_DIR"
90+
ENVRC
91+
92+
echo
93+
echo "Created worktree: $WT_DIR"
94+
echo "Branch: $BRANCH (base: $BASE_REF)"
95+
echo
96+
echo "Next:"
97+
echo " cd $WT_DIR"
98+
echo " source .envrc"
99+
echo " (cd PyTorchSimDevice && python setup.py build_ext --inplace) # build the .so once"

0 commit comments

Comments
 (0)