diff --git a/scripts/gate_slot_lock.py b/scripts/gate_slot_lock.py index c1ab00461e0..9f56ad9d15b 100644 --- a/scripts/gate_slot_lock.py +++ b/scripts/gate_slot_lock.py @@ -30,6 +30,8 @@ from __future__ import annotations import contextlib import fcntl import os +import shlex +import signal import subprocess import sys import time @@ -38,6 +40,7 @@ from typing import IO, TYPE_CHECKING, Final if TYPE_CHECKING: from collections.abc import Iterator + from types import FrameType HELD_MARKER_ENV: Final = "LITELLM_GATE_SLOT_HELD" SLOT_COUNT_ENV: Final = "LITELLM_GATE_SLOTS" @@ -143,6 +146,19 @@ def held_slot() -> Iterator[None]: os.environ.pop(HELD_MARKER_ENV, None) +def _report_never_ran(signal_name: str) -> None: + print( + f"gate_slot_lock: {signal_name} arrived while waiting for a slot, so `{shlex.join(sys.argv[1:])}` never ran", + file=sys.stderr, + flush=True, + ) + + +def _exit_before_the_command_ran(signum: int, _frame: FrameType | None) -> None: + _report_never_ran(signal.Signals(signum).name) + os._exit(128 + signum) + + def _wait_ignoring_interrupts(process: subprocess.Popen[bytes]) -> int: while True: try: @@ -155,10 +171,13 @@ def main() -> int: if len(sys.argv) < 2: print("usage: gate_slot_lock.py [args...]", file=sys.stderr) return 2 + previous_sigterm_handler: Final = signal.signal(signal.SIGTERM, _exit_before_the_command_ran) try: held: Final = acquire_slot() except KeyboardInterrupt: + _report_never_ran("SIGINT") return 130 + signal.signal(signal.SIGTERM, previous_sigterm_handler) try: code: Final = _wait_ignoring_interrupts(subprocess.Popen(sys.argv[1:])) except FileNotFoundError as error: @@ -166,7 +185,10 @@ def main() -> int: return 127 if held is not None: held.close() - return code if code >= 0 else 128 - code + exit_code: Final = code if code >= 0 else 128 - code + if exit_code != 0: + print(f"gate_slot_lock: `{shlex.join(sys.argv[1:])}` exited {exit_code}", file=sys.stderr, flush=True) + return exit_code if __name__ == "__main__": diff --git a/tests/unit/test_gate_slot_lock.py b/tests/unit/test_gate_slot_lock.py index c80e876700b..74d43ec2add 100644 --- a/tests/unit/test_gate_slot_lock.py +++ b/tests/unit/test_gate_slot_lock.py @@ -180,12 +180,64 @@ def test_nested_wrapping_reenters_instead_of_deadlocking(tmp_path: Path) -> None assert b"nested ok" in stdout -def test_wrapped_command_exit_code_is_propagated(tmp_path: Path) -> None: - proc = subprocess.run( - [sys.executable, str(HELPER), sys.executable, "-c", "raise SystemExit(7)"], - env=_env(tmp_path / "locks", "2"), +def test_wrapped_command_exit_code_is_propagated_and_named_on_stderr(tmp_path: Path) -> None: + env = _env(tmp_path / "locks", "2") + failed = subprocess.run( + [sys.executable, str(HELPER), sys.executable, "-c", "raise SystemExit(5)"], + env=env, + capture_output=True, ) - assert proc.returncode == 7 + assert failed.returncode == 5 + assert b"'raise SystemExit(5)'` exited 5" in failed.stderr + succeeded = subprocess.run(_wrapped(["pass"]), env=env, capture_output=True) + assert succeeded.returncode == 0 + assert b"exited" not in succeeded.stderr + + +def test_sigterm_while_queued_says_the_command_never_ran(tmp_path: Path) -> None: + lock_dir = tmp_path / "locks" + holder_started = tmp_path / "holder.started" + release = tmp_path / "release" + done = tmp_path / "done" + env = _env(lock_dir, "1") + holder = subprocess.Popen(_wrapped([START_THEN_WAIT_FOR, str(holder_started), str(release)]), env=env) + try: + assert _wait_until(holder_started.exists, 10) + contender = subprocess.Popen(_wrapped([TOUCH_TARGET, str(done)]), env=env, stderr=subprocess.PIPE) + try: + time.sleep(1.5) + contender.send_signal(signal.SIGTERM) + assert contender.wait(timeout=10) == 128 + signal.SIGTERM + assert contender.stderr is not None + assert b"SIGTERM arrived while waiting for a slot" in contender.stderr.read() + release.touch() + assert holder.wait(timeout=10) == 0 + assert not done.exists() + finally: + _reap(contender) + finally: + release.touch() + _reap(holder) + + +def test_sigterm_after_the_command_started_does_not_claim_it_never_ran(tmp_path: Path) -> None: + started = tmp_path / "started" + release = tmp_path / "release" + wrapper = subprocess.Popen( + _wrapped([START_THEN_WAIT_FOR, str(started), str(release)]), + env=_env(tmp_path / "locks", "1"), + stderr=subprocess.PIPE, + ) + try: + assert _wait_until(started.exists, 10) + wrapper.send_signal(signal.SIGTERM) + wrapper.wait(timeout=10) + release.touch() + assert wrapper.stderr is not None + assert b"never ran" not in wrapper.stderr.read() + finally: + release.touch() + _reap(wrapper) def test_missing_command_exits_127_and_no_command_exits_2(tmp_path: Path) -> None: