Batch jobs from a worktree run on the main checkout's environment

README.md

@@ -28,7 +28,7 @@ claude --worktree <task>
    answer to the same question gets the next version, in a new session.
 2. It asks you which working object is current, then works in scripts; outputs go to the task directory.
 3. Results reach a shared working object only after you sign off the dry run.
-4. It opens a PR. Once it is merged, leave with `/exit` and remove the worktree.
+4. It opens a PR. Once it is merged and its batch jobs have finished, leave with `/exit` and remove the worktree.
 
 Notebooks and environment changes (`pixi install`, `pixi add`) stay with you, in the main checkout. With no
 analysis repo yet, ask Claude to start an analysis project: the skill sets one up.

hooks/_repo.py

@@ -6,11 +6,22 @@
 """
 
 import json
+import shlex
 import subprocess
 from pathlib import Path
 
 PLUGIN = "analysis-workflow"
 
+#: `pixi` subcommands that change the workspace environment, or use it and so install it first if it is missing.
+PIXI_ENV_COMMANDS = frozenset(
+    {"install", "add", "lock", "upgrade", "update", "reinstall", "remove", "run", "shell", "shell-hook"}
+)
+#: The first of these after `pixi` is the command, whatever global options precede it.
+PIXI_SUBCOMMANDS = PIXI_ENV_COMMANDS | {
+    "auth", "build", "clean", "completion", "config", "exec", "global", "import", "info", "init", "list",
+    "search", "self-update", "task", "tree", "upload", "workspace",
+}  # fmt: skip
+
 
 def checkouts(directory: Path) -> tuple[Path, Path] | None:
     """The checkout containing ``directory`` and the repository's main checkout; None outside git."""
@@ -37,3 +48,74 @@ def adopted(checkout: Path) -> bool:
         return any(key.split("@")[0] == PLUGIN and on is True for key, on in enabled.items())
     except (OSError, ValueError, AttributeError):
         return False
+
+
+def nearest_manifest(directory: Path) -> Path:
+    """The `pixi.toml` pixi itself finds from ``directory``: the nearest at or above it."""
+    return next(
+        (d / "pixi.toml" for d in (directory, *directory.parents) if (d / "pixi.toml").exists()),
+        directory / "pixi.toml",
+    )
+
+
+def is_root_of_worktree(manifest: Path) -> bool:
+    """Whether ``manifest`` is the root workspace of a worktree of an adopted repo.
+
+    Its editable self-dependency (`path = "."`) would rebase the repo's package, and every path it
+    resolves, into the worktree. A task-local manifest is its own workspace.
+    """
+    if ".claude/worktrees" not in str(manifest):
+        return False
+    if manifest.exists() and 'path = "."' not in manifest.read_text(errors="ignore"):
+        return False
+    found = checkouts(manifest.parent)
+    return found is not None and adopted(found[0])
+
+
+def _is_separator(tok: str) -> bool:
+    """A token made only of command separators (`;`, `&&`, `|`, newlines, parentheses)."""
+    return bool(tok) and set(tok) <= set(";&|()\n")
+
+
+def pixi_calls(command: str, comments: bool = False) -> list[list[str]] | None:
+    """Argument lists of the `pixi` calls in ``command`` that change or use an environment; None if unparsable.
+
+    With ``comments``, an unquoted `#` starts a comment to the end of its line, as in a script. Quotes are
+    honoured either way, so quoted text (a PR body, an echoed message) is never read as a call.
+    """
+    tokens = []
+    for text in command.splitlines(keepends=True) if comments else [command]:
+        lexer = shlex.shlex(text, posix=True, punctuation_chars="();<>|&\n")
+        lexer.whitespace = " \t\r"  # newlines separate commands, so keep them as tokens
+        lexer.commenters = "#" if comments else ""  # one line at a time: a comment swallows its newline
+        lexer.whitespace_split = True
+        try:
+            tokens += [*lexer, "\n"]
+        except ValueError:
+            return None
+    calls = []
+    for i, tok in enumerate(tokens):
+        if Path(tok).name != "pixi":
+            continue
+        args = []
+        for t in tokens[i + 1 :]:
+            if _is_separator(t):
+                break
+            args.append(t)
+        if next((a for a in args if a in PIXI_SUBCOMMANDS), None) in PIXI_ENV_COMMANDS:
+            calls.append(args)
+    return calls
+
+
+def manifest_arg(args: list[str]) -> str | None:
+    """The `--manifest-path` a pixi call passes, if any. In `pixi run`, arguments after the task are the task's."""
+    command = next((a for a in args if a in PIXI_SUBCOMMANDS), None)
+    rest = args[args.index(command) + 1 :] if command else args
+    for i, a in enumerate(rest):
+        if a in ("--manifest-path", "-m"):
+            return rest[i + 1] if i + 1 < len(rest) else None
+        if a.startswith("--manifest-path="):
+            return a.split("=", 1)[1]
+        if command == "run" and not a.startswith("-"):
+            return None
+    return None

hooks/guard_task_paths.py

@@ -6,6 +6,9 @@
 usually gitignored, so git never warns. Absolute paths into a worktree, a temp dir or the checkout
 itself outlive what they point at, or break on the next machine.
 
+A job script written in a worktree must not run `pixi` on the workspace it finds itself: the job
+would build the root environment into the worktree, out of reach of the Bash pixi guard.
+
 It acts only in a repo that has adopted the plugin. Only the text being written is inspected, so
 pre-existing violations elsewhere in a file do not trip the hook; they surface when that line is next
 touched, which is when to fix them.
@@ -19,10 +22,23 @@
 import sys
 from pathlib import Path
 
-from _repo import adopted, checkouts
+from _repo import PIXI_ENV_COMMANDS, adopted, checkouts, is_root_of_worktree, manifest_arg, nearest_manifest, pixi_calls
 
 # Code that runs. Prose may legitimately quote any of these paths.
 CODE_SUFFIXES = {".py", ".sh", ".sbatch", ".ipynb", ".R"}
+JOB_SCRIPT_SUFFIXES = {".sh", ".sbatch"}
+
+#: The one workspace a job may name: the main checkout's, as `$MAIN` passed in at submit. Any other path, relative
+#: or derived in the job (`git rev-parse --show-toplevel`), can resolve to the worktree itself.
+MAIN_MANIFEST = re.compile(r"\$\{?MAIN\}?/")
+
+
+def _pixi_on_found_workspace(script: str) -> bool:
+    """Whether any command in ``script`` runs `pixi` on a workspace other than the main checkout's."""
+    calls = pixi_calls(script.replace("\\\n", " "), comments=True)
+    if calls is None:  # unparsable: err towards blocking
+        return bool(re.search(rf"\bpixi\b.*\b({'|'.join(PIXI_ENV_COMMANDS)})\b", script))
+    return any(not MAIN_MANIFEST.match(manifest_arg(args) or "") for args in calls)
 
 
 def _rules(repo: str) -> list[tuple[re.Pattern, str, str]]:
@@ -65,6 +81,18 @@ def main() -> int:
     # Whatever this call would add: Write.content, Edit.new_string, NotebookEdit.new_source
     written = "\n".join(str(tool_input.get(key, "")) for key in ("content", "new_string", "new_source"))
     hits = [(what, fix) for pattern, what, fix in _rules(found[1].name) if pattern.search(written)]
+    if (
+        file_path.suffix in JOB_SCRIPT_SUFFIXES
+        and _pixi_on_found_workspace(written)
+        and is_root_of_worktree(nearest_manifest(file_path.parent))
+    ):
+        hits.append(
+            (
+                "a `pixi` call on the worktree's own workspace",
+                'jobs run on the main checkout\'s environment: pixi run --manifest-path "$MAIN/pixi.toml", '
+                "with MAIN passed in at submit",
+            )
+        )
     if not hits:
         return 0
 

hooks/guard_worktree_pixi.py

@@ -20,77 +20,24 @@
 
 import json
 import re
-import shlex
 import sys
 from pathlib import Path
 
-from _repo import adopted, checkouts
-
-_MUTATING = {"install", "add", "lock", "upgrade", "update", "reinstall", "remove"}
-# These use the workspace environment, installing it first if it is missing.
-_ENV_USING = {"run", "shell", "shell-hook"}
-_BLOCKED = _MUTATING | _ENV_USING
-# The first of these after `pixi` is the command, whatever global options precede it.
-_SUBCOMMANDS = _BLOCKED | {
-    "auth", "build", "clean", "completion", "config", "exec", "global", "import", "info", "init", "list",
-    "search", "self-update", "task", "tree", "upload", "workspace",
-}  # fmt: skip
-
-
-def _is_separator(tok: str) -> bool:
-    """A token made only of command separators (`;`, `&&`, `|`, newlines, parentheses)."""
-    return bool(tok) and set(tok) <= set(";&|()\n")
-
-
-def _pixi_calls(command: str) -> list[list[str]] | None:
-    """Argument lists of blocked `pixi` calls; None if unparsable."""
-    lexer = shlex.shlex(command, posix=True, punctuation_chars="();<>|&\n")
-    lexer.whitespace = " \t\r"  # newlines separate commands, so keep them as tokens
-    lexer.commenters = ""  # a `#` comment must not swallow the newline after it
-    lexer.whitespace_split = True
-    try:
-        tokens = list(lexer)
-    except ValueError:
-        return None
-    calls = []
-    for i, tok in enumerate(tokens):
-        if Path(tok).name != "pixi":
-            continue
-        args = []
-        for t in tokens[i + 1 :]:
-            if _is_separator(t):
-                break
-            args.append(t)
-        if next((a for a in args if a in _SUBCOMMANDS), None) in _BLOCKED:
-            calls.append(args)
-    return calls
+from _repo import PIXI_ENV_COMMANDS, adopted, checkouts, is_root_of_worktree, manifest_arg, nearest_manifest, pixi_calls
 
 
 def _manifest(args: list[str], cwd: Path) -> Path | None:
     """The manifest a call acts on; None if it is only known at run time."""
-    for i, a in enumerate(args):
-        if a in ("--manifest-path", "-m") and i + 1 < len(args):
-            path = Path(args[i + 1])
-            break
-        if a.startswith("--manifest-path="):
-            path = Path(a.split("=", 1)[1])
-            break
-    else:  # pixi's own discovery: the nearest pixi.toml at or above cwd
-        return next((d / "pixi.toml" for d in (cwd, *cwd.parents) if (d / "pixi.toml").exists()), cwd / "pixi.toml")
-    if any(c in str(path) for c in "$`"):
+    named = manifest_arg(args)
+    if named is None:
+        return nearest_manifest(cwd)
+    if any(c in named for c in "$`"):
         return None  # shell-expanded at run time, so not provably outside a worktree
+    path = Path(named)
     path = (cwd / path).resolve() if not path.is_absolute() else path.resolve()
     return path / "pixi.toml" if path.is_dir() else path
 
 
-def _is_root_of_worktree(manifest: Path) -> bool:
-    if ".claude/worktrees" not in str(manifest):
-        return False
-    if manifest.exists() and 'path = "."' not in manifest.read_text(errors="ignore"):
-        return False  # task-local workspace
-    return _in_adopted_repo(manifest.parent)
-
-
 def _in_adopted_repo(directory: Path) -> bool:
     found = checkouts(directory)
     return found is not None and adopted(found[0])
@@ -104,13 +51,13 @@ def main() -> int:
 
     if ".claude/worktrees" not in str(cwd) and "--manifest-path" not in command and " -m " not in command:
         return 0
-    calls = _pixi_calls(command)
+    calls = pixi_calls(command)
     if calls is None:  # unparsable: fall back to the conservative text match
-        pattern = rf"\bpixi\b.*\b({'|'.join(_BLOCKED)})\b"
+        pattern = rf"\bpixi\b.*\b({'|'.join(PIXI_ENV_COMMANDS)})\b"
         blocked = ".claude/worktrees" in str(cwd) and re.search(pattern, command) and _in_adopted_repo(cwd)
     else:
         manifests = [_manifest(args, cwd) for args in calls]
-        blocked = any(_in_adopted_repo(cwd) if m is None else _is_root_of_worktree(m) for m in manifests)
+        blocked = any(_in_adopted_repo(cwd) if m is None else is_root_of_worktree(m) for m in manifests)
     if not blocked:
         return 0
 

hooks/session_context.md

@@ -8,5 +8,5 @@ analysis code, outputs or data. Always:
    newest file name.
 4. Shared objects grow by addition only (`commit_adata`), and only after sign-off on the dry run.
 5. Agents never write to `data/<dataset>/results/` or the central `figures/`; those are the human lane.
-6. A worktree runs on the main checkout's environment (`pixi run --manifest-path <main>/pixi.toml ...`);
-   environment changes happen in the main checkout.
+6. A worktree, and every batch job submitted from it, runs on the main checkout's environment
+   (`pixi run --manifest-path <main>/pixi.toml ...`); environment changes happen in the main checkout.

skills/analysis-workflow/SKILL.md

@@ -62,6 +62,9 @@ into the human lane.
 
 - Another task's untracked outputs: `task_paths(<its _common.py>).outputs`. Read the main checkout by absolute
   path; never `cd` into it.
+- Batch-job scripts sit in `scripts/` beside the code they run. They get the main checkout as `$MAIN` from the
+  submit command, never hardcoded, run `pixi run --manifest-path "$MAIN/pixi.toml"`, and log to `PATHS.logs`. How
+  to submit them is the cluster's skill.
 - A figure that prints a number may carry `scripts/verify.py` re-deriving it. A code package's invariants belong in
   that package's tests.
 
@@ -71,6 +74,7 @@ into the human lane.
 3. An object other tasks will load is promoted to `data/<dataset>/processed/` at sign-off. Ask which subdirectory;
    record the path.
 4. `analysis-workflow check layout`, then push.
+5. The worktree is removed only once no batch job submitted from it is still queued or running.
 
 ## Data
 

skills/analysis-workflow/assets/task_template/README.md

@@ -32,7 +32,8 @@ Anything promoted to `data/<dataset>/processed/` at sign-off is recorded here wi
 
 ## Running
 
-From the task directory, on the main checkout's environment. Batch jobs log to `PATHS.logs`:
+From the task directory, on the main checkout's environment. Batch jobs get `MAIN=<main checkout>` at submit and
+log to `PATHS.logs`:
 
 ```bash
 PY="pixi run --manifest-path <main checkout>/pixi.toml python"

tests/test_hooks.py

@@ -92,6 +92,41 @@ def test_pixi_allows_running_on_the_main_checkouts_environment(repo, pixi_worktr
     assert run("guard_worktree_pixi.py", bash(command, pixi_worktree)) == 0
 
 
+@pytest.mark.parametrize(
+    "line",
+    [
+        "pixi run python scripts/a.py",
+        "srun -n1 pixi run python a.py",
+        "pixi install",
+        "pixi \\\n    run python a.py",  # continued line
+        "pixi run --manifest-path ../../../pixi.toml python a.py",  # resolves against the job's directory
+        'pixi run python a.py && pixi run --manifest-path "$MAIN/pixi.toml" python b.py',  # each call checked
+        "pixi update",
+        'ROOT=$(git rev-parse --show-toplevel); pixi run --manifest-path "$ROOT/pixi.toml" python a.py',
+        'pixi run python a.py  # TODO: --manifest-path "$MAIN/pixi.toml"',
+        'echo "starting # job"; pixi run python a.py',  # a quoted hash is no comment
+        'pixi run python a.py --manifest-path "$MAIN/pixi.toml"',  # after the task, the flag is the task's
+    ],
+)
+def test_task_paths_blocks_job_scripts_running_pixi_on_the_worktree(pixi_worktree, line):
+    job = pixi_worktree / "analysis" / "t_v1" / "scripts" / "run.sh"
+    assert run("guard_task_paths.py", write(job, f"#!/bin/bash\n{line}\n")) == 2
+
+
+@pytest.mark.parametrize(
+    ("rel", "line"),
+    [
+        ("analysis/t_v1/scripts/run.sh", 'pixi run --manifest-path "$MAIN/pixi.toml" python scripts/a.py'),
+        ("analysis/t_v1/scripts/run.sh", "pixi run --manifest-path=${MAIN}/pixi.toml python scripts/a.py"),
+        ("analysis/t_v1/scripts/run.sh", "# not pixi run python a.py, see the README"),
+        ("analysis/bench_v1/run.sh", "pixi run python a.py"),  # task-local workspace
+        ("analysis/t_v1/scripts/run.py", 'print("pixi run python a.py")'),  # not a job script
+    ],
+)
+def test_task_paths_allows_job_scripts(pixi_worktree, rel, line):
+    assert run("guard_task_paths.py", write(pixi_worktree / rel, f"#!/bin/bash\n{line}\n")) == 0
+
+
 def test_pixi_allows_the_main_checkout(repo):
     main, _ = repo
     (main / "pixi.toml").write_text('[pypi-dependencies]\nmyproject = { path = ".", editable = true }\n')