Skip to content

feat(slurm): deadline timer for confined jobs too, plus WIP snapshotter #9

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
ligon merged 1 commit into from
Sep 9, 2026
Merged

feat(slurm): deadline timer for confined jobs too, plus WIP snapshotter #9

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
26 changes: 26 additions & 0 deletions README.org
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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-<mirror>.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/<mirror>= 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
Expand Down
3 changes: 3 additions & 0 deletions config.example.yaml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -99,6 +99,9 @@ system_prompt: /home/<your-username>/.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/<mirror> every N min; 0 = off

# Mirrors are optional when using zero-config repo detection.
mirrors:
Expand Down
4 changes: 3 additions & 1 deletion docs/persistent-presence.org
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
160 changes: 18 additions & 142 deletions sucoder/cli.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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(
Expand Down Expand Up @@ -1193,127 +1170,26 @@ 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:
return

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 <mirror>` 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.
Expand Down
17 changes: 16 additions & 1 deletion sucoder/config.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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.
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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,
)


Expand Down
36 changes: 35 additions & 1 deletion sucoder/mirror.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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


Expand Down Expand Up @@ -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).

Expand All @@ -2347,6 +2349,12 @@ def _build_batch_script(
``srun --overlap --pty tmux -L <socket> 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)
Expand All @@ -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"
)
Expand Down Expand Up @@ -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(
[
Expand All @@ -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
Expand Down
Loading
Loading