diff --git a/docs/ops/distributed-prefill-kv-network.md b/docs/ops/distributed-prefill-kv-network.md index 79f6a9f..7b55fae 100644 --- a/docs/ops/distributed-prefill-kv-network.md +++ b/docs/ops/distributed-prefill-kv-network.md @@ -339,6 +339,8 @@ model's EOS. `--max-response-tokens` defaults to 512 as an explicit safety cap; reaching it marks the stage incomplete instead of presenting a truncated answer as successful. The Critic receives the Generator completion status and must not penalize an honest statement that an open problem has no accepted proof. +The REPL ignores external `SIGTERM`; its shell supervisor restarts signal-based +exits. Only `/quit`, `/exit`, or EOF is treated as approval to stop. ## Rollback diff --git a/scripts/agent_gan_repl.py b/scripts/agent_gan_repl.py index 1eec84f..e9e8adc 100644 --- a/scripts/agent_gan_repl.py +++ b/scripts/agent_gan_repl.py @@ -5,6 +5,7 @@ import argparse import hashlib import json +import signal import time import uuid from pathlib import Path @@ -19,6 +20,17 @@ ) +def install_signal_protection() -> None: + def ignore_sigterm(signum, _frame): + print( + f"\n[protected] ignored external signal {signum}. " + "Type /quit to approve shutdown.", + flush=True, + ) + + signal.signal(signal.SIGTERM, ignore_sigterm) + + class TokenPrinter: def __init__(self, tokenizer, label: str) -> None: self.tokenizer = tokenizer @@ -56,6 +68,7 @@ def _stage(name: str, warm: dict, actual: dict, text: str) -> dict: def main() -> int: + install_signal_protection() parser = argparse.ArgumentParser() parser.add_argument("--worker-ssh", default="allens") parser.add_argument("--address", default="127.0.0.1:51051") diff --git a/scripts/run_agent_gan_repl.sh b/scripts/run_agent_gan_repl.sh index 49c09ca..4748778 100644 --- a/scripts/run_agent_gan_repl.sh +++ b/scripts/run_agent_gan_repl.sh @@ -5,7 +5,23 @@ REPO_ROOT="$(cd "$(dirname "$0")/.." && pwd)" PYTHON="${KAKEYA_BENCH_PYTHON:-$HOME/.venv-distwan/bin/python}" MODEL="${KAKEYA_BENCH_MODEL:-$HOME/kakeya-models/gemma-4-26B-A4B-it-mlx-4bit}" -exec env PYTHONPATH="$REPO_ROOT:$REPO_ROOT/sdks/python" \ - "$PYTHON" "$REPO_ROOT/scripts/agent_gan_repl.py" \ - --tokenizer-id "$MODEL" \ - "$@" +trap 'echo "[supervisor] external termination ignored; use /quit"' TERM HUP + +while true; do + set +e + env PYTHONPATH="$REPO_ROOT:$REPO_ROOT/sdks/python" \ + "$PYTHON" "$REPO_ROOT/scripts/agent_gan_repl.py" \ + --tokenizer-id "$MODEL" \ + "$@" + status=$? + set -e + if [[ "$status" -eq 0 ]]; then + exit 0 + fi + if [[ "$status" -eq 129 || "$status" -eq 137 || "$status" -eq 143 ]]; then + echo "[supervisor] REPL exited from signal ($status); restarting in 2s..." + sleep 2 + continue + fi + exit "$status" +done diff --git a/tests/inference_engine/bridge/test_agent_gan_repl.py b/tests/inference_engine/bridge/test_agent_gan_repl.py index ee34cad..138d4b0 100644 --- a/tests/inference_engine/bridge/test_agent_gan_repl.py +++ b/tests/inference_engine/bridge/test_agent_gan_repl.py @@ -1,4 +1,7 @@ -from scripts.agent_gan_repl import TokenPrinter, _stage +import signal +from pathlib import Path + +from scripts.agent_gan_repl import TokenPrinter, _stage, install_signal_protection class Tokenizer: @@ -51,3 +54,28 @@ def test_repl_stage_is_redacted_and_passes_cache_gate(): {**actual, "complete": False, "stop_reason": "client_safety_limit"}, "cut off", )["ok"] + + +def test_external_sigterm_is_ignored_until_user_quits(monkeypatch, capsys): + installed = {} + monkeypatch.setattr( + signal, + "signal", + lambda number, handler: installed.update({number: handler}), + ) + install_signal_protection() + installed[signal.SIGTERM](signal.SIGTERM, None) + output = capsys.readouterr().out + assert "ignored external signal" in output + assert "Type /quit to approve shutdown" in output + + +def test_shell_supervisor_restarts_signal_exits_only(): + source = ( + Path(__file__).resolve().parents[3] + / "scripts" + / "run_agent_gan_repl.sh" + ).read_text() + assert "trap" in source and "TERM HUP" in source + assert '"$status" -eq 143' in source + assert "restarting in 2s" in source