Merge b0549b50e7 into e16d79d635

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 <command> [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__":

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: