diff --git a/README.org b/README.org index d698063..52f3c87 100644 --- a/README.org +++ b/README.org @@ -742,6 +742,32 @@ targets: fields are optional --- omitting them keeps the partition default, which is the right behaviour for whole-node partitions. +*** Deadline watchdog and WIP snapshots + +Every SLURM-backed session (=salloc= or =confined= =sbatch=) starts a +small watchdog on the compute node. It warns at 30, 15, and 5 minutes +before the allocation's =--time= via =tmux display-message= and by +writing =$HOME/.cache/sucoder/slurm-deadline-.warn= (the +un-suffixed =slurm-deadline.warn= is also written for prompts that +poll the older path). It never cancels the job; that stays with +=sucoder release=. + +The same script snapshots the mirror's dirty working tree (tracked and +untracked files, not ignored ones) to =refs/sucoder/wip/= on +the mirror's =origin= at each warning and every +=wip_snapshot_minutes= (default 10; =0= disables the periodic run). A +mirror with no =origin= is never snapshotted, so on a shared-filesystem +mirror this is a no-op today; it becomes live with the local-disk +tiering described in [[file:docs/local-disk-tiering.org][docs/local-disk-tiering.org]]. + +#+begin_src yaml + slurm: + partition: savio4_htc + account: co_carleton + confined: true + wip_snapshot_minutes: 10 +#+end_src + *** Reconnecting If the SSH connection drops, the tmux session on the login node diff --git a/config.example.yaml b/config.example.yaml index 7ea6918..d4f52ff 100644 --- a/config.example.yaml +++ b/config.example.yaml @@ -99,6 +99,9 @@ system_prompt: /home//.sucoder/system_prompt.org # cpus_per_task: 4 # 4 cores out of the shared node # mem: 16G # per-job memory (required on shared partitions) # time: "24:00:00" +# # confined: true # sbatch launch inside the job cgroup (see README) +# # wip_snapshot_minutes: 10 # deadline timer snapshots the dirty tree to +# # # refs/sucoder/wip/ every N min; 0 = off # Mirrors are optional when using zero-config repo detection. mirrors: diff --git a/docs/persistent-presence.org b/docs/persistent-presence.org index 3a78cbf..82565e0 100644 --- a/docs/persistent-presence.org +++ b/docs/persistent-presence.org @@ -422,7 +422,9 @@ login-node-daemon controller variant (v2). controller polls =squeue -o %L= / =sacct -o State=, parsing =D-HH:MM:SS= correctly. It does not depend on the in-node =slurm-deadline.warn= warner regardless of the latter's health - (now fixed; see Appendix) -- the two channels stay independent. + (its day-format parsing is fixed, see Appendix; and as of the + shared =slurm_timer= script it actually runs for confined jobs, + which it never did before) -- the two channels stay independent. 2. *Files survive, process does not.* The repo on NFS =$HOME= plus a pushed mirror are durable; the agent's *conversation context* dies on turnover. Every turnover is therefore bracketed by diff --git a/sucoder/cli.py b/sucoder/cli.py index 1c041a1..93e548f 100644 --- a/sucoder/cli.py +++ b/sucoder/cli.py @@ -45,7 +45,9 @@ ) from .executor import CommandError, CommandExecutor from .logging_utils import setup_logger +from .slurm_timer import TIME_LEFT_TO_MINS_SH, build_timer_script from .mirror import ( + _sanitize_session_token, MirrorError, MirrorManager, # Re-exported for backward-compatible imports (e.g. tests) and reuse; @@ -1134,35 +1136,10 @@ def _ensure_slurm_node( return session.compute_node, cn_control -# Bash helper embedded verbatim into the on-node deadline timer (see -# ``_start_slurm_timer``). Converts SLURM ``squeue -o %L`` time-left -# (TIME_LEFT) into whole minutes remaining. ``%L`` renders as -# ``D-HH:MM:SS`` once a day or more remains, ``HH:MM:SS`` under a day, -# and ``MM:SS`` under an hour; a job with no time limit prints -# ``UNLIMITED``. Splitting on ``:`` alone mis-handles the ``D-HH`` -# field (bash reads ``D-HH`` as the arithmetic ``D - HH``), so the day -# component is split off on ``-`` first. Leading zeros are forced to -# base 10 to avoid octal errors (``08``/``09``). Non-numeric values -# (``UNLIMITED``/``INVALID``/empty) return a large sentinel so no -# deadline warning ever fires. Kept as a module constant (not inlined -# in the f-string) so it is unit-testable under bash and free of -# brace-escaping noise. -_SLURM_TIME_LEFT_TO_MINS_SH = r''' -left_to_mins() { - local s="$1" days=0 rest a b c - if [ -z "$s" ]; then echo 999999; return; fi - case "$s" in - *-*) days="${s%%-*}"; rest="${s#*-}" ;; - *) rest="$s" ;; - esac - IFS=: read -r a b c <<< "$rest" - if [ -z "$c" ]; then b="$a"; a=0; fi - case "${days}${a}${b}" in - *[!0-9]*) echo 999999; return ;; - esac - echo $(( 10#${days:-0}*1440 + 10#${a:-0}*60 + 10#${b:-0} )) -} -'''.strip("\n") +# ``left_to_mins`` now lives in ``slurm_timer`` (shared by the confined +# and unconfined launch paths); re-exported here so existing imports and +# tests keep working. +_SLURM_TIME_LEFT_TO_MINS_SH = TIME_LEFT_TO_MINS_SH def _start_slurm_timer( @@ -1193,9 +1170,7 @@ def _start_slurm_timer( (module constant ``_SLURM_TIME_LEFT_TO_MINS_SH``) so the ``D-HH:MM:SS`` day format is handled correctly and is unit-testable. """ - import shlex import subprocess as _sp - import textwrap job_id = session.slurm_job_id if not job_id: @@ -1203,117 +1178,18 @@ def _start_slurm_timer( tmux_session = f"sucoder-{session.mirror_name}" - # Defensive shell-quoting. ``mirror_name`` and therefore - # ``tmux_session`` come from configuration the user controls; if a - # mirror were ever named with shell metacharacters the unquoted - # interpolation below would be a command-injection vector. - # ``job_id`` is an int from ``int(token)`` so it's already safe, but - # we quote it for symmetry and to insulate against future changes. - q_tmux = shlex.quote(tmux_session) - q_job = shlex.quote(str(job_id)) - - # Use a per-user runtime directory rather than world-writable /tmp. - # On a shared HPC compute node, predictable /tmp/slurm-*.warn paths - # are subject to symlink races: a co-resident user can pre-create - # the path as a symlink to a sensitive file and have the timer - # overwrite it. ``$HOME/.cache/sucoder/`` is per-user (NFS-shared - # across nodes, owned by the same uid) and not writable by other - # local users on the compute node, which closes that vector. - # - # The agent reads ``slurm-deadline.warn`` from the same location - # (the agent runs as the same user inside tmux on the compute - # node), so the path remains discoverable to consumers. - - # The script runs on the compute node, querying squeue via the - # login node is unnecessary — SLURM_JOB_ID is in the environment - # and squeue works locally on compute nodes too. - timer_script = textwrap.dedent(f"""\ - #!/bin/bash - set -u - STATE_DIR="${{HOME}}/.cache/sucoder" - mkdir -p "$STATE_DIR" - chmod 700 "$STATE_DIR" 2>/dev/null || true - WARN_FILE="$STATE_DIR/slurm-deadline.warn" - WARN5="$STATE_DIR/.slurm-warn-5" - WARN15="$STATE_DIR/.slurm-warn-15" - WARN30="$STATE_DIR/.slurm-warn-30" - rm -f "$WARN5" "$WARN15" "$WARN30" "$WARN_FILE" - - # Wait for the agent tmux session to appear before monitoring. - # The timer starts before the session is created, so we must - # not treat its absence as "agent exited". - TMUX_READY=0 - for i in $(seq 1 120); do - if tmux has-session -t {q_tmux} 2>/dev/null; then - TMUX_READY=1 - break - fi - sleep 5 - done - if [ "$TMUX_READY" -eq 0 ]; then - # User owns SLURM lifecycle (see `sucoder release`); leave - # the allocation alone even if the agent's tmux session - # never appeared, since the user may want to debug or - # reuse the compute node manually. - echo "Timed out waiting for tmux session {q_tmux}; SLURM job {q_job} kept alive. Run 'sucoder release' or 'scancel {q_job}' to free the allocation." > "$WARN_FILE" - exit 1 - fi - - # Make each deadline warning linger on the status line so a - # full-screen agent TUI doesn't redraw over it before the human - # notices (scoped to our session via -t, not the global -g). - tmux set-option -t {q_tmux} display-time 15000 2>/dev/null || true - - while true; do - left=$(squeue --job {q_job} --noheader -o "%L" 2>/dev/null) - if [ -z "$left" ]; then - msg="SLURM job {q_job} is no longer queued — allocation may have ended." - echo "$msg" > "$WARN_FILE" - tmux display-message "$msg" 2>/dev/null - break - fi - - # If the agent tmux session is gone, write a warning but - # do NOT auto-cancel the SLURM allocation. Users own the - # SLURM lifecycle (use `sucoder release ` for - # explicit cancel); an automatic scancel here would tear - # down the allocation on transient agent failures and - # destroy any chance of reattaching. - if ! tmux has-session -t {q_tmux} 2>/dev/null; then - echo "Agent tmux session is gone; SLURM job {q_job} kept alive. Run 'sucoder release' or 'scancel {q_job}' to free the allocation." > "$WARN_FILE" - break - fi - - mins=$(left_to_mins "$left") - if [ "$mins" -le 5 ] && [ ! -f "$WARN5" ]; then - msg="SLURM: ~${{mins}} min left (job {q_job}). Commit and save NOW." - echo "$msg" > "$WARN_FILE" - tmux display-message -t {q_tmux} "$msg" 2>/dev/null - touch "$WARN5" - elif [ "$mins" -le 15 ] && [ ! -f "$WARN15" ]; then - msg="SLURM: ~${{mins}} min left (job {q_job}). Start wrapping up." - echo "$msg" > "$WARN_FILE" - tmux display-message -t {q_tmux} "$msg" 2>/dev/null - touch "$WARN15" - elif [ "$mins" -le 30 ] && [ ! -f "$WARN30" ]; then - msg="SLURM: ~${{mins}} min left (job {q_job})." - echo "$msg" > "$WARN_FILE" - tmux display-message -t {q_tmux} "$msg" 2>/dev/null - touch "$WARN30" - fi - sleep 60 - done - """) - - # Inject the time-left parser (kept as a module constant so it can - # be unit-tested under bash) ahead of the monitoring loop. Done - # post-dedent so the helper's column-0 body doesn't flatten the - # common-indent prefix and push the ``#!`` off byte 0. - timer_script = timer_script.replace( - 'rm -f "$WARN5" "$WARN15" "$WARN30" "$WARN_FILE"\n', - 'rm -f "$WARN5" "$WARN15" "$WARN30" "$WARN_FILE"\n\n' - + _SLURM_TIME_LEFT_TO_MINS_SH + "\n", - 1, + # The script is shared with the confined (sbatch) launch path; see + # ``slurm_timer.build_timer_script``. It writes its warnings under + # ``$HOME/.cache/sucoder/`` (per-user, NFS-shared, not writable by + # other local users) rather than world-writable /tmp, where a + # predictable path is subject to symlink races on a shared node. + # No snapshot directory: on this path the mirror root is not known + # until after the node is up, and today's shared mirror has no + # ``origin`` to snapshot to. + timer_script = build_timer_script( + mirror_token=_sanitize_session_token(session.mirror_name), + tmux_session=tmux_session, + job_id=job_id, ) # Write the script to the compute node via stdin, then run it. diff --git a/sucoder/config.py b/sucoder/config.py index 3ce2c56..10474b2 100644 --- a/sucoder/config.py +++ b/sucoder/config.py @@ -60,6 +60,9 @@ class SlurmConfig: # via `sbatch` so it runs inside the # job cgroup, confined to the reserved # cores instead of the whole node + wip_snapshot_minutes: int = 10 # deadline timer snapshots the dirty + # working tree to refs/sucoder/wip/ + # every N minutes; 0 disables @dataclass @@ -425,7 +428,7 @@ class ConfigWarning(UserWarning): # which the parser silently drops, so the option appears to do nothing. _VALID_SLURM_KEYS = frozenset({ "partition", "account", "time", "qos", - "cpus_per_task", "mem", "local_disk", "confined", + "cpus_per_task", "mem", "local_disk", "confined", "wip_snapshot_minutes", }) # Target-level options commonly misplaced under ``slurm:``; warned about # with a tailored "move it up a level" hint. @@ -1138,6 +1141,17 @@ def _parse_slurm_config(raw: Any) -> Optional[SlurmConfig]: if not isinstance(confined, bool): raise ConfigError("`slurm.confined` must be a boolean when provided.") + wip_snapshot_minutes = raw.get("wip_snapshot_minutes", 10) + if ( + isinstance(wip_snapshot_minutes, bool) + or not isinstance(wip_snapshot_minutes, int) + or wip_snapshot_minutes < 0 + ): + raise ConfigError( + "`slurm.wip_snapshot_minutes` must be a non-negative integer " + "(0 disables periodic snapshots) when provided." + ) + # Surface keys that the parser will ignore. The common case is a # target-level option (notably ``system_prompt_extra``) indented one # level too deep, under ``slurm:`` instead of beside it -- which @@ -1169,6 +1183,7 @@ def _parse_slurm_config(raw: Any) -> Optional[SlurmConfig]: mem=mem, local_disk=local_disk, confined=confined, + wip_snapshot_minutes=wip_snapshot_minutes, ) diff --git a/sucoder/mirror.py b/sucoder/mirror.py index bbdecbd..9f57d44 100644 --- a/sucoder/mirror.py +++ b/sucoder/mirror.py @@ -73,6 +73,7 @@ def _frontmatter_problem(metadata: Optional[SkillMetadata]) -> Optional[str]: ensure_directory_mode, ) from .skills_version import validate_skills_version +from .slurm_timer import build_timer_script from .workspace_prefs import WorkspacePrefs @@ -2329,6 +2330,7 @@ def _build_batch_script( mirror_path: str, agent_cmd_str: str, env: Optional[Mapping[str, str]] = None, + timer_path: Optional[str] = None, ) -> str: """sbatch script body for a ``confined`` launch (shared partitions). @@ -2347,6 +2349,12 @@ def _build_batch_script( ``srun --overlap --pty tmux -L attach``. The keeper loop holds the job while the session lives; when the agent exits, the session ends, the keeper exits, and the job frees. + + ``timer_path`` (a staged ``slurm_timer.build_timer_script`` output) + is started with ``nohup`` *after* the session is confirmed and + before the keeper loop, so it runs inside the job cgroup and dies + with the job. Without it a confined job has no deadline watchdog + at all: ``cli._start_slurm_timer`` only runs on the salloc path. """ q_sess = shlex.quote(tmux_session) q_sock = shlex.quote(socket) @@ -2372,7 +2380,11 @@ def _build_batch_script( " echo \"SUCODER: tmux new-session failed (rc=$rc)\" >&2\n" " exit 1\n" "fi\n" - f"while tmux -L {q_sock} has-session -t {q_sess} 2>/dev/null; do\n" + + ( + f"nohup {shlex.quote(timer_path)} > /dev/null 2>&1 &\n" + if timer_path else "" + ) + + f"while tmux -L {q_sock} has-session -t {q_sess} 2>/dev/null; do\n" " sleep 15\n" "done\n" ) @@ -2726,9 +2738,11 @@ def _launch_confined( log_path = f"{cache_dir}/job-{safe}-%j.out" mirror_path = self._resolve_remote_path(ctx) + timer_path = f"{cache_dir}/slurm-timer-{safe}.sh" script = self._build_batch_script( tmux_session=session_name, socket=socket, mirror_path=mirror_path, agent_cmd_str=windowed_cmd, env=env, + timer_path=timer_path, ) self.executor.run_agent( [ @@ -2738,6 +2752,26 @@ def _launch_confined( ], input=script, check=True, capture_output=True, ) + # Deadline watchdog + WIP snapshotter, started by the batch body + # inside the job cgroup. The job id is read from $SLURM_JOB_ID at + # run time (unknown until sbatch assigns it); every tmux call + # carries the dedicated socket. Staged AFTER the batch script so + # the first staged file is still the batch script. + timer_script = build_timer_script( + mirror_token=safe, + tmux_session=session_name, + tmux_socket=socket, + snapshot_dir=mirror_path, + snapshot_minutes=slurm.wip_snapshot_minutes, + ) + self.executor.run_agent( + [ + "sh", "-c", + f"umask 077 && cat > {shlex.quote(timer_path)} " + f"&& chmod 700 {shlex.quote(timer_path)}", + ], + input=timer_script, check=True, capture_output=True, + ) # Submit. check=False so a transport drop AFTER the remote sbatch # created the job (rc 255) is reported as a MirrorError with a diff --git a/sucoder/slurm_timer.py b/sucoder/slurm_timer.py new file mode 100644 index 0000000..0bb90df --- /dev/null +++ b/sucoder/slurm_timer.py @@ -0,0 +1,229 @@ +"""Deadline watchdog + WIP snapshotter for SLURM-backed agent sessions. + +One bash script serves both launch modes: + +- *unconfined* (``salloc`` + compute-node SSH): ``cli._start_slurm_timer`` + renders it with a literal job id and starts it over SSH; +- *confined* (``sbatch``): ``MirrorManager._launch_confined`` stages it to + NFS next to the batch script, and the batch body starts it *inside the + job cgroup* once the tmux session is up. The job id is read from + ``$SLURM_JOB_ID`` because it does not exist when the script is rendered. + +The script warns at 30/15/5 minutes before the allocation's ``--time`` +(``tmux display-message`` for the human, a sentinel file for the agent) +and, when ``snapshot_dir`` is set, snapshots that working tree's dirty +state to ``refs/sucoder/wip/`` on its ``origin`` every +``snapshot_minutes`` and at each warning. A tree with no ``origin`` is +never snapshotted: there is nowhere durable for the snapshot to go, and +``git add -A`` on a shared filesystem is not free (see +``docs/local-disk-tiering.org``). + +It lives in its own module because ``cli`` imports ``mirror``, not the +reverse, so a builder both can share cannot live in either. The script +is a plain template with ``@TOKEN@`` placeholders rather than an +f-string: it is dense with ``$`` and ``{}``. +""" + +from __future__ import annotations + +import shlex +from typing import Optional + +# Converts SLURM ``squeue -o %L`` time-left into whole minutes. ``%L`` +# renders as ``D-HH:MM:SS`` once a day or more remains, ``HH:MM:SS`` under +# a day, ``MM:SS`` under an hour; a job with no limit prints ``UNLIMITED``. +# The day component is split off on ``-`` first (bash reads ``D-HH`` as +# arithmetic), leading zeros are forced to base 10 (``08``/``09`` are not +# octal), and non-numeric values return a large sentinel so no warning +# ever fires. Unit-tested under bash in tests/test_slurm_timer.py. +TIME_LEFT_TO_MINS_SH = r''' +left_to_mins() { + local s="$1" days=0 rest a b c + if [ -z "$s" ]; then echo 999999; return; fi + case "$s" in + *-*) days="${s%%-*}"; rest="${s#*-}" ;; + *) rest="$s" ;; + esac + IFS=: read -r a b c <<< "$rest" + if [ -z "$c" ]; then b="$a"; a=0; fi + case "${days}${a}${b}" in + *[!0-9]*) echo 999999; return ;; + esac + echo $(( 10#${days:-0}*1440 + 10#${a:-0}*60 + 10#${b:-0} )) +} +'''.strip("\n") + +# Snapshot the dirty tree of $SNAPSHOT_DIR to refs/sucoder/wip/$MIRROR_TOKEN +# on origin. Every step is best-effort: the timer must never die because +# a snapshot could not be taken. Runs in a subshell so cd/trap/export do +# not leak. A temporary index leaves the agent's real index untouched; +# the "last tree" marker lives under .git/ because ``add -A`` would sweep +# up a marker in the working tree and defeat the unchanged check. +# Behaviour verified by hand on n0036.savio4, 2026-09-09. +WIP_SNAPSHOT_SH = r''' +snapshot_wip() { + [ -n "$SNAPSHOT_DIR" ] || return 0 + [ -e "$SNAPSHOT_DIR/.git" ] || return 0 + ( + cd "$SNAPSHOT_DIR" || exit 0 + git remote get-url origin >/dev/null 2>&1 || exit 0 + marker="$(git rev-parse --git-dir)/sucoder-last-wip-tree" + GIT_INDEX_FILE=$(mktemp) || exit 0 + export GIT_INDEX_FILE + trap 'rm -f "$GIT_INDEX_FILE"' EXIT + git read-tree HEAD 2>/dev/null || exit 0 + git add -A 2>/dev/null || exit 0 + tree=$(git write-tree 2>/dev/null) || exit 0 + [ "$tree" = "$(cat "$marker" 2>/dev/null)" ] && exit 0 + if [ "$tree" = "$(git rev-parse 'HEAD^{tree}')" ]; then + echo "$tree" > "$marker" + exit 0 + fi + wip=$(git -c user.name=sucoder-wip -c user.email=sucoder-wip@localhost \ + commit-tree "$tree" -p HEAD \ + -m "WIP snapshot $(date -Is) job $JOB") || exit 0 + git update-ref "refs/sucoder/wip/$MIRROR_TOKEN" "$wip" || exit 0 + if git push --quiet --force origin "refs/sucoder/wip/$MIRROR_TOKEN" >/dev/null 2>&1; then + echo "$tree" > "$marker" + fi + ) +} +'''.strip("\n") + +# State files are per mirror: several confined mirrors share one $HOME, +# and a second timer's startup ``rm -f`` must not clear the first's +# markers. The un-suffixed ``slurm-deadline.warn`` is still written for +# prompts that poll the legacy path. +_TEMPLATE = r'''#!/bin/bash +# sucoder SLURM deadline timer + WIP snapshotter (generated; do not edit). +set -u +STATE_DIR="${HOME}/.cache/sucoder" +mkdir -p "$STATE_DIR" +chmod 700 "$STATE_DIR" 2>/dev/null || true +MIRROR_TOKEN=@MIRROR_TOKEN@ +TMUX_SESSION=@TMUX_SESSION@ +TMUX_BIN=(@TMUX_CMD@) +SNAPSHOT_DIR=@SNAPSHOT_DIR@ +SNAPSHOT_MINUTES=@SNAPSHOT_MINUTES@ +JOB=@JOB_REF@ +WARN_FILE="$STATE_DIR/slurm-deadline-$MIRROR_TOKEN.warn" +LEGACY_WARN_FILE="$STATE_DIR/slurm-deadline.warn" +WARN5="$STATE_DIR/.slurm-warn-5-$MIRROR_TOKEN" +WARN15="$STATE_DIR/.slurm-warn-15-$MIRROR_TOKEN" +WARN30="$STATE_DIR/.slurm-warn-30-$MIRROR_TOKEN" +rm -f "$WARN5" "$WARN15" "$WARN30" "$WARN_FILE" + +if [ -z "$JOB" ]; then + echo "sucoder timer: no SLURM job id (not inside a job?); exiting." > "$WARN_FILE" + exit 1 +fi + +@LEFT_TO_MINS@ + +@SNAPSHOT_WIP@ + +warn() { + echo "$1" > "$WARN_FILE" + echo "$1" > "$LEGACY_WARN_FILE" + "${TMUX_BIN[@]}" display-message -t "$TMUX_SESSION" "$1" 2>/dev/null +} + +# Wait for the agent tmux session to appear before monitoring. The +# unconfined timer starts before the session is created, so its absence +# must not be read as "agent exited". +TMUX_READY=0 +for i in $(seq 1 120); do + if "${TMUX_BIN[@]}" has-session -t "$TMUX_SESSION" 2>/dev/null; then + TMUX_READY=1 + break + fi + sleep 5 +done +if [ "$TMUX_READY" -eq 0 ]; then + # The user owns the SLURM lifecycle (see `sucoder release`): leave the + # allocation alone even though the agent never appeared. + echo "Timed out waiting for tmux session $TMUX_SESSION; SLURM job $JOB kept alive. Run 'sucoder release' or 'scancel $JOB' to free the allocation." > "$WARN_FILE" + exit 1 +fi + +# Make each warning linger on the status line so a full-screen agent TUI +# does not redraw over it before the human notices. +"${TMUX_BIN[@]}" set-option -t "$TMUX_SESSION" display-time 15000 2>/dev/null || true + +elapsed=0 +while true; do + left=$(squeue --job "$JOB" --noheader -o "%L" 2>/dev/null) + if [ -z "$left" ]; then + warn "SLURM job $JOB is no longer queued -- allocation may have ended." + break + fi + + # Agent gone: record it but do NOT scancel (see above). + if ! "${TMUX_BIN[@]}" has-session -t "$TMUX_SESSION" 2>/dev/null; then + echo "Agent tmux session is gone; SLURM job $JOB kept alive. Run 'sucoder release' or 'scancel $JOB' to free the allocation." > "$WARN_FILE" + break + fi + + mins=$(left_to_mins "$left") + if [ "$mins" -le 5 ] && [ ! -f "$WARN5" ]; then + warn "SLURM: ~${mins} min left (job $JOB). Commit and save NOW." + touch "$WARN5" + snapshot_wip + elif [ "$mins" -le 15 ] && [ ! -f "$WARN15" ]; then + warn "SLURM: ~${mins} min left (job $JOB). Start wrapping up." + touch "$WARN15" + snapshot_wip + elif [ "$mins" -le 30 ] && [ ! -f "$WARN30" ]; then + warn "SLURM: ~${mins} min left (job $JOB)." + touch "$WARN30" + snapshot_wip + fi + + if [ "$SNAPSHOT_MINUTES" -gt 0 ] && [ $((elapsed % SNAPSHOT_MINUTES)) -eq 0 ]; then + snapshot_wip + fi + sleep 60 + elapsed=$((elapsed + 1)) +done +''' + + +def build_timer_script( + *, + mirror_token: str, + tmux_session: str, + job_id: Optional[int] = None, + tmux_socket: Optional[str] = None, + snapshot_dir: Optional[str] = None, + snapshot_minutes: int = 10, +) -> str: + """Render the timer script. + + ``mirror_token`` names the per-mirror state files and the WIP ref; the + caller sanitizes it (``mirror._sanitize_session_token``). ``job_id`` + ``None`` means "read ``$SLURM_JOB_ID`` at run time" (confined launches, + where the id is unknown until sbatch assigns it). ``tmux_socket`` + adds ``-L `` to *every* tmux call, which a confined launch + requires: without it the timer waits on a session it cannot see and + exits as "timed out". ``snapshot_dir`` ``None`` or ``snapshot_minutes`` + ``0`` disables the periodic snapshot (the deadline snapshots still run + when a directory is given). + + Every user-controlled value is shell-quoted. The only unquoted + substitution is the ``$SLURM_JOB_ID`` reference itself. + """ + if snapshot_minutes < 0: + raise ValueError("snapshot_minutes must be >= 0") + tmux_cmd = "tmux" if tmux_socket is None else f"tmux -L {shlex.quote(tmux_socket)}" + job_ref = '"${SLURM_JOB_ID:-}"' if job_id is None else shlex.quote(str(job_id)) + return ( + _TEMPLATE + .replace("@MIRROR_TOKEN@", shlex.quote(mirror_token)) + .replace("@TMUX_SESSION@", shlex.quote(tmux_session)) + .replace("@TMUX_CMD@", tmux_cmd) + .replace("@SNAPSHOT_DIR@", shlex.quote(snapshot_dir or "")) + .replace("@SNAPSHOT_MINUTES@", str(int(snapshot_minutes))) + .replace("@JOB_REF@", job_ref) + .replace("@LEFT_TO_MINS@", TIME_LEFT_TO_MINS_SH) + .replace("@SNAPSHOT_WIP@", WIP_SNAPSHOT_SH) + ) diff --git a/tests/test_batch_script.py b/tests/test_batch_script.py index 92d60d5..566a441 100644 --- a/tests/test_batch_script.py +++ b/tests/test_batch_script.py @@ -112,3 +112,30 @@ def test_quoting_survives_metacharacters(): ) r = _bash_n(MirrorManager._build_batch_script(**nasty, env={"K": "v'1 `x`"})) assert r.returncode == 0, r.stderr + + +# -- deadline timer / WIP snapshotter ----------------------------------------- + +def test_timer_started_after_session_check_before_keeper(): + """The watchdog runs inside the job cgroup: started by the batch body + once ``new-session`` succeeded (so its wait is trivially satisfied) and + before the keeper loop (so it dies with the job).""" + s = MirrorManager._build_batch_script( + **_BASE, timer_path="/global/home/users/ligon/.cache/sucoder/slurm-timer-K-Aggregators.sh", + ) + nohup = "nohup /global/home/users/ligon/.cache/sucoder/slurm-timer-K-Aggregators.sh > /dev/null 2>&1 &\n" + assert nohup in s + rc_check = s.index("SUCODER: tmux new-session failed") + keeper = s.index("while tmux -L sucoder-K-Aggregators has-session") + assert rc_check < s.index(nohup) < keeper + + +def test_timer_omitted_when_no_path(): + assert "nohup" not in MirrorManager._build_batch_script(**_BASE) + + +@_bash_only +def test_timer_path_is_quoted_and_script_parses(tmp_path): + s = MirrorManager._build_batch_script(**_BASE, timer_path="/p q/t.sh") + assert "nohup '/p q/t.sh' > /dev/null 2>&1 &" in s + assert _bash_n(s).returncode == 0 diff --git a/tests/test_cli.py b/tests/test_cli.py index 72b3d2a..5b131e5 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -797,7 +797,7 @@ def _no_real_ssh(*a, **kw): # ------------------------------------------------------------------ -def test_slurm_timer_script_omits_scancel(): +def test_slurm_timer_script_omits_scancel(monkeypatch): """Regression: the backstop timer must NOT auto-scancel. Previously the on-compute-node monitor script ran @@ -813,19 +813,39 @@ def test_slurm_timer_script_omits_scancel(): `scancel {q_job}` to free the allocation") — those are fine because they're inside an ``echo``/string, not a shell statement. """ - import inspect - src = inspect.getsource(cli._start_slurm_timer) - for raw in src.splitlines(): + # The script body now comes from ``slurm_timer.build_timer_script``; + # inspecting ``_start_slurm_timer``'s own source would check nothing. + # Capture what it actually ships to the node instead. + from sucoder import slurm_timer + + rendered = [] + real_build = slurm_timer.build_timer_script + + def capture(**kw): + script = real_build(**kw) + rendered.append(script) + return script + + monkeypatch.setattr(cli, "build_timer_script", capture) + # ``_start_slurm_timer`` imports subprocess locally as ``_sp``; patch + # the module so its ssh write/start calls are swallowed. + monkeypatch.setattr(subprocess, "run", + lambda *a, **k: subprocess.CompletedProcess(a, 0, "", "")) + session = SimpleNamespace(slurm_job_id=7, mirror_name="sample", compute_node="n0") + control = SimpleNamespace(ssh_options=lambda **kw: []) + cli._start_slurm_timer(session, control, control, mock.Mock()) + + assert rendered, "the timer script was not rendered" + for raw in rendered[0].splitlines(): stripped = raw.strip() - # Skip strings that mention scancel for documentation/warnings. - if not stripped.startswith("scancel"): - continue - # If we got here, a bare `scancel ...` shell command remains. - pytest.fail( - "_start_slurm_timer still emits a `scancel` shell " - f"command line: {stripped!r}. User owns SLURM lifecycle " - "now; use `sucoder release` for explicit cancel." + # Strings may *mention* scancel ("Run `scancel N` to free ..."), + # but no line may execute it. + assert not stripped.startswith("scancel"), ( + "the deadline timer still emits a `scancel` shell command: " + f"{stripped!r}. The user owns the SLURM lifecycle; use " + "`sucoder release` for explicit cancel." ) + assert "JOB=7\n" in rendered[0] def _slurm_config(tmp_path: Path, *, with_session_jobid: bool = False) -> Path: diff --git a/tests/test_config.py b/tests/test_config.py index 8f599b8..cd51d85 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -904,3 +904,29 @@ def test_default_base_branch_optional(tmp_path: Path) -> None: assert cfg.mirrors["auto"].default_base_branch is None assert cfg.mirrors["fixed"].default_base_branch == "develop" + + +# -- slurm.wip_snapshot_minutes ------------------------------------------------ + +def _slurm_raw(**extra): + raw = {"partition": "savio4_htc", "account": "co_carleton"} + raw.update(extra) + return raw + + +def test_slurm_wip_snapshot_minutes_defaults_to_ten(): + from sucoder.config import _parse_slurm_config + assert _parse_slurm_config(_slurm_raw()).wip_snapshot_minutes == 10 + + +@pytest.mark.parametrize("value", [0, 5, 120]) +def test_slurm_wip_snapshot_minutes_accepts_non_negative_ints(value): + from sucoder.config import _parse_slurm_config + assert _parse_slurm_config(_slurm_raw(wip_snapshot_minutes=value)).wip_snapshot_minutes == value + + +@pytest.mark.parametrize("value", [-1, True, "10", 2.5]) +def test_slurm_wip_snapshot_minutes_rejects_bad_values(value): + from sucoder.config import ConfigError, _parse_slurm_config + with pytest.raises(ConfigError, match="wip_snapshot_minutes"): + _parse_slurm_config(_slurm_raw(wip_snapshot_minutes=value)) diff --git a/tests/test_mirror.py b/tests/test_mirror.py index 4a80bec..840b1ff 100644 --- a/tests/test_mirror.py +++ b/tests/test_mirror.py @@ -1517,6 +1517,39 @@ def test_launch_confined_wraps_agent_in_bash_lc(tmp_path, monkeypatch): assert "tmux -L sucoder-sample new-session -A -d -s sucoder-sample" in script +def test_launch_confined_stages_and_starts_deadline_timer(tmp_path, monkeypatch): + """A confined job gets its own deadline watchdog: the timer script is + staged next to the batch script and the batch body nohups it inside + the cgroup. Before this, ``_start_slurm_timer`` (salloc path only) + meant confined jobs had no watchdog at all.""" + manager, ctx = _confined_manager(tmp_path, monkeypatch) + calls = [] + manager.executor.run_agent = _confined_responder(calls, sbatch_out="9") + + manager._launch_confined( + ctx, ["claude"], remote_prelude_text=None, + prelude_sentinel="__X__", env=None, detached=True, + ) + writes = [c for c in calls if c["args"][0] == "sh" and c["input"]] + assert len(writes) == 2, "batch script then timer script must both be staged" + batch, timer = writes[0], writes[1] + timer_path = [t for t in timer["args"][2].split() if "slurm-timer-" in t][0] + assert timer_path.endswith("/.cache/sucoder/slurm-timer-sample.sh") + assert "chmod 700" in timer["args"][2] + # The batch body starts exactly that file, after the session check. + assert f"nohup {timer_path} > /dev/null 2>&1 &" in batch["input"] + # Confined specifics threaded through: runtime job id, dedicated socket, + # the mirror as snapshot dir, the configured cadence. + assert 'JOB="${SLURM_JOB_ID:-}"' in timer["input"] + assert "TMUX_BIN=(tmux -L sucoder-sample)" in timer["input"] + assert "MIRROR_TOKEN=sample" in timer["input"] + assert "SNAPSHOT_MINUTES=10" in timer["input"] + assert "SNAPSHOT_DIR=" in timer["input"] and "SNAPSHOT_DIR=''" not in timer["input"] + # sbatch is submitted only after both files are staged. + sbatch_idx = next(i for i, c in enumerate(calls) if c["args"][0] == "sbatch") + assert all(calls.index(w) < sbatch_idx for w in writes) + + def test_launch_confined_session_not_ready_surfaces_log(tmp_path, monkeypatch): """If the job is RUNNING but the tmux session never came up, fail loudly (with the job-log pointer) instead of attaching into nothing.""" diff --git a/tests/test_slurm_timer_script.py b/tests/test_slurm_timer_script.py new file mode 100644 index 0000000..a6d76b4 --- /dev/null +++ b/tests/test_slurm_timer_script.py @@ -0,0 +1,199 @@ +"""The shared deadline timer + WIP snapshotter (``sucoder.slurm_timer``). + +The script is rendered for both launch modes and exercised under bash: +syntax, socket threading, the no-``scancel`` invariant, and the snapshot +function against a real temporary repository (mirroring the manual spike +on n0036.savio4, 2026-09-09). Nothing here execs ``tmux`` or ``squeue``. +""" +from __future__ import annotations + +import os +import shutil +import subprocess +from pathlib import Path + +import pytest + +from sucoder.slurm_timer import WIP_SNAPSHOT_SH, build_timer_script + +_bash = pytest.mark.skipif(shutil.which("bash") is None, reason="bash not available") +_git = pytest.mark.skipif(shutil.which("git") is None, reason="git not available") + + +def _render(**kw) -> str: + base = dict(mirror_token="K-Aggregators", tmux_session="sucoder-K-Aggregators") + base.update(kw) + return build_timer_script(**base) + + +def _bash_n(script: str, tmp_path: Path) -> subprocess.CompletedProcess: + path = tmp_path / "timer.sh" + path.write_text(script) + return subprocess.run(["bash", "-n", str(path)], capture_output=True, text=True) + + +# -- rendering ---------------------------------------------------------------- + +def test_no_unresolved_tokens_in_either_mode(): + import re + for script in (_render(job_id=123), _render(tmux_socket="s", snapshot_dir="/d")): + assert not re.search(r"@[A-Z_]+@", script), script + + +def test_unconfined_mode_literal_job_id_and_plain_tmux(): + s = _render(job_id=38661192) + assert "JOB=38661192\n" in s + assert "TMUX_BIN=(tmux)\n" in s + assert "SNAPSHOT_DIR=''\n" in s + + +def test_confined_mode_reads_job_id_at_runtime_and_threads_socket(): + s = _render(tmux_socket="sucoder-K-Aggregators", snapshot_dir="/global/home/u/mirrors/K") + assert 'JOB="${SLURM_JOB_ID:-}"\n' in s + assert "TMUX_BIN=(tmux -L sucoder-K-Aggregators)\n" in s + assert "SNAPSHOT_DIR=/global/home/u/mirrors/K\n" in s + # Every tmux invocation goes through the array, so the socket cannot be + # dropped from any one call (has-session, display-message, set-option). + import re + bare_tmux = re.compile(r"(?:^|[;&|(]\s*|\$\(\s*)tmux\s") # tmux in command position + for line in s.splitlines(): + stripped = line.strip() + if stripped.startswith("#") or stripped.startswith("TMUX_BIN=("): + continue + assert not bare_tmux.search(stripped), f"tmux called without the socket array: {line}" + assert s.count('"${TMUX_BIN[@]}"') >= 4 + + +def test_state_files_are_per_mirror_and_legacy_warn_kept(): + s = _render(mirror_token="alpha") + assert 'WARN_FILE="$STATE_DIR/slurm-deadline-$MIRROR_TOKEN.warn"' in s + assert 'LEGACY_WARN_FILE="$STATE_DIR/slurm-deadline.warn"' in s + assert "MIRROR_TOKEN=alpha\n" in s + for n in (5, 15, 30): + assert f'WARN{n}="$STATE_DIR/.slurm-warn-{n}-$MIRROR_TOKEN"' in s + + +def test_user_values_are_shell_quoted(): + s = _render(mirror_token="x y", tmux_session="s;rm -rf /", tmux_socket="a b", + snapshot_dir="/p q") + assert "MIRROR_TOKEN='x y'\n" in s + assert "TMUX_SESSION='s;rm -rf /'\n" in s + assert "TMUX_BIN=(tmux -L 'a b')\n" in s + assert "SNAPSHOT_DIR='/p q'\n" in s + + +def test_negative_snapshot_minutes_rejected(): + with pytest.raises(ValueError): + _render(snapshot_minutes=-1) + + +def test_never_emits_a_bare_scancel(): + """The user owns the SLURM lifecycle (``sucoder release``); the timer + may *mention* scancel in a warning string but never run it.""" + for script in (_render(job_id=1), _render(tmux_socket="s")): + for raw in script.splitlines(): + assert not raw.strip().startswith("scancel"), raw + + +@_bash +def test_bash_syntax_both_modes(tmp_path): + assert _bash_n(_render(job_id=1), tmp_path).returncode == 0 + r = _bash_n(_render(tmux_socket="s", snapshot_dir="/d", snapshot_minutes=0), tmp_path) + assert r.returncode == 0, r.stderr + + +# -- snapshot_wip under bash -------------------------------------------------- + +def _git_run(cwd: Path, *args: str) -> str: + return subprocess.run( + ["git", "-c", "user.name=t", "-c", "user.email=t@x", *args], + cwd=cwd, capture_output=True, text=True, check=True, + ).stdout.strip() + + +def _snapshot(work: Path, token: str = "mirror") -> subprocess.CompletedProcess: + script = ( + "set -u\n" + f"SNAPSHOT_DIR={work}\nMIRROR_TOKEN={token}\nJOB=42\n" + + WIP_SNAPSHOT_SH + "\nsnapshot_wip\n" + ) + return subprocess.run(["bash", "-c", script], capture_output=True, text=True) + + +@pytest.fixture +def repo_pair(tmp_path): + """A non-bare ``origin`` with ``updateInstead`` and a working clone of it.""" + origin = tmp_path / "origin" + _git_run(tmp_path, "init", "-q", "-b", "main", str(origin)) + (origin / "README").write_text("hello\n") + (origin / ".gitignore").write_text(".venv/\n") + _git_run(origin, "add", "-A") + _git_run(origin, "commit", "-q", "-m", "init") + _git_run(origin, "config", "receive.denyCurrentBranch", "updateInstead") + work = tmp_path / "work" + _git_run(tmp_path, "clone", "-q", str(origin), str(work)) + return origin, work + + +@_bash +@_git +def test_snapshot_clean_tree_is_a_noop(repo_pair): + origin, work = repo_pair + assert _snapshot(work).returncode == 0 + assert subprocess.run(["git", "show-ref", "refs/sucoder/wip/mirror"], cwd=origin).returncode != 0 + + +@_bash +@_git +def test_snapshot_dirty_tree_lands_on_origin_and_skips_ignored(repo_pair): + origin, work = repo_pair + (work / "README").write_text("hello\nedit\n") # tracked edit + (work / "notes.org").write_text("new\n") # untracked + (work / ".venv").mkdir(); (work / ".venv" / "x").write_text("ignored\n") + r = _snapshot(work) + assert r.returncode == 0, r.stderr + wip = _git_run(origin, "rev-parse", "refs/sucoder/wip/mirror") + assert _git_run(origin, "rev-parse", "refs/sucoder/wip/mirror^") == _git_run(work, "rev-parse", "HEAD") + files = _git_run(origin, "ls-tree", "-r", "--name-only", wip).splitlines() + assert "notes.org" in files and "README" in files + assert not any(f.startswith(".venv") for f in files) + # The agent's own index and HEAD are untouched. + status = subprocess.run(["git", "status", "--porcelain"], cwd=work, + capture_output=True, text=True, check=True).stdout + assert status.splitlines() == [" M README", "?? notes.org"] + # origin's checked-out branch did not move (the WIP ref is not a branch). + assert _git_run(origin, "rev-parse", "main") == _git_run(work, "rev-parse", "HEAD") + + +@_bash +@_git +def test_snapshot_unchanged_tree_skips_via_marker(repo_pair): + origin, work = repo_pair + (work / "notes.org").write_text("new\n") + assert _snapshot(work).returncode == 0 + first = _git_run(origin, "rev-parse", "refs/sucoder/wip/mirror") + marker = work / ".git" / "sucoder-last-wip-tree" + assert marker.exists() and marker.read_text().strip() == _git_run(origin, "rev-parse", f"{first}^{{tree}}") + assert _snapshot(work).returncode == 0 + assert _git_run(origin, "rev-parse", "refs/sucoder/wip/mirror") == first + # The marker lives under .git/, so it never appears in a snapshot. + assert "sucoder-last-wip-tree" not in _git_run(origin, "ls-tree", "-r", "--name-only", first) + + +@_bash +@_git +def test_snapshot_without_origin_is_a_noop(tmp_path): + solo = tmp_path / "solo" + _git_run(tmp_path, "init", "-q", "-b", "main", str(solo)) + (solo / "a").write_text("x\n") + _git_run(solo, "add", "-A"); _git_run(solo, "commit", "-q", "-m", "init") + (solo / "b").write_text("dirty\n") + assert _snapshot(solo).returncode == 0 + assert subprocess.run(["git", "show-ref", "refs/sucoder/wip/mirror"], cwd=solo).returncode != 0 + assert not (solo / ".git" / "sucoder-last-wip-tree").exists() + + +@_bash +def test_snapshot_missing_or_non_git_dir_is_a_noop(tmp_path): + assert _snapshot(tmp_path / "nope").returncode == 0 + assert _snapshot(tmp_path).returncode == 0