mirror of
https://github.com/BerriAI/litellm.git
synced 2026-08-28 05:25:59 +00:00
feat(scripts): queue heavy gates behind a machine-wide slot lock
This commit is contained in:
parent
870a8cf764
commit
eafddaaa12
8 changed files with 573 additions and 16 deletions
21
Makefile
21
Makefile
|
|
@ -8,8 +8,8 @@
|
|||
lint-basedpyright lint-e2e-basedpyright lint-basedpyright-budget-update lint-type-discipline lint-type-discipline-budget-update \
|
||||
lint-ruff-budget lint-ruff-budget-update lint-budget-update lint-gate \
|
||||
install-dev install-proxy-dev install-test-deps install-hooks \
|
||||
install-helm-unittest check-circular-imports check-import-safety check pre-commit \
|
||||
lint-install lint-fetch-base bootstrap
|
||||
install-helm-unittest check-circular-imports check-import-safety check check-inner pre-commit \
|
||||
lint-install lint-fetch-base bootstrap bootstrap-inner
|
||||
|
||||
# Default target
|
||||
help:
|
||||
|
|
@ -52,10 +52,17 @@ help:
|
|||
@echo " make test-proxy-unit-b - Run proxy_unit_tests (p-z, ~28 files)"
|
||||
@echo " make test-integration - Run integration tests"
|
||||
@echo " make test-unit-helm - Run helm unit tests"
|
||||
@echo ""
|
||||
@echo "Heavy targets (check, bootstrap, lint) queue for LITELLM_GATE_SLOTS machine-wide"
|
||||
@echo "slots (default 2; 0 disables) so parallel sessions don't thrash one machine."
|
||||
|
||||
UV := uv
|
||||
UV_RUN := $(UV) run --no-sync
|
||||
|
||||
# Machine-wide slot queue for the heavy targets below; python3 + stdlib only, so
|
||||
# it runs before any venv exists. See scripts/gate_slot_lock.py.
|
||||
GATE_SLOT_LOCK := python3 scripts/gate_slot_lock.py
|
||||
|
||||
LINT_DEP_INSTALL ?= install-dev
|
||||
LINT_E2E_DEP_INSTALL ?= lint-install
|
||||
LINT_DEP_BASE ?= lint-fetch-base
|
||||
|
|
@ -74,6 +81,9 @@ install-dev:
|
|||
$(UV) sync --inexact --frozen
|
||||
|
||||
bootstrap:
|
||||
@$(GATE_SLOT_LOCK) $(MAKE) bootstrap-inner
|
||||
|
||||
bootstrap-inner:
|
||||
$(UV) sync --inexact --frozen --extra proxy --group proxy-dev --group e2e-dev
|
||||
$(UV_RUN) python scripts/prisma_generate_if_needed.py
|
||||
cd ui/litellm-dashboard && ../../scripts/with_dashboard_node.sh npm install --no-audit --no-fund
|
||||
|
|
@ -230,7 +240,7 @@ check-import-safety: $(LINT_DEP_INSTALL)
|
|||
# base fetch) runs once up front; the checks themselves are independent, so a sub-make
|
||||
# fans them out with -j and the fast ones finish under basedpyright's shadow.
|
||||
lint: lint-install lint-fetch-base
|
||||
$(MAKE) -j $(LINT_JOBS) $(LINT_OUTPUT_SYNC) LINT_DEP_INSTALL= LINT_E2E_DEP_INSTALL= LINT_DEP_BASE= lint-checks
|
||||
$(GATE_SLOT_LOCK) $(MAKE) -j $(LINT_JOBS) $(LINT_OUTPUT_SYNC) LINT_DEP_INSTALL= LINT_E2E_DEP_INSTALL= LINT_DEP_BASE= lint-checks
|
||||
|
||||
lint-checks: lint-format-check-changed lint-ruff lint-gate lint-type-discipline lint-basedpyright lint-e2e-basedpyright check-circular-imports check-import-safety
|
||||
|
||||
|
|
@ -244,7 +254,10 @@ lint-dev: lint-format-changed check-circular-imports check-import-safety
|
|||
# test-linting.yml (Python), test-litellm-ui-build.yml's frontend-lint (dashboard), and
|
||||
# check-ui-api-types.yml (API-type drift), skipping any whose files aren't in scope.
|
||||
# Not auto-installed as a git hook so it never slows an unrelated human commit.
|
||||
check: bootstrap
|
||||
check:
|
||||
@$(GATE_SLOT_LOCK) $(MAKE) check-inner
|
||||
|
||||
check-inner: bootstrap
|
||||
./scripts/pre_commit_lint.sh
|
||||
|
||||
pre-commit:
|
||||
|
|
|
|||
173
scripts/gate_slot_lock.py
Normal file
173
scripts/gate_slot_lock.py
Normal file
|
|
@ -0,0 +1,173 @@
|
|||
#!/usr/bin/env python3
|
||||
"""Machine-wide slot lock for this repo's heavy entrypoints.
|
||||
|
||||
`make check`, `make bootstrap`, `make lint`, and the standalone budget gates
|
||||
(scripts/ruff_strict_gate.py, scripts/type_discipline_gate.py,
|
||||
scripts/type_check_gate.py) each hold one of N machine-wide slots while they
|
||||
run, so however many sessions and worktrees share one machine, at most N of
|
||||
them execute a basedpyright/pytest/prettier storm at a time instead of all
|
||||
thrashing it at once. Slots are fcntl.flock files (macOS ships no flock(1)
|
||||
binary, hence python3 + stdlib only, runnable before any venv exists) under a
|
||||
per-user cache directory shared by every worktree and session:
|
||||
~/.cache/litellm/gate-slots by default, $LITELLM_GATE_SLOT_DIR to override.
|
||||
A holder's lock dies with its process, so a crash leaves nothing to clean up.
|
||||
|
||||
$LITELLM_GATE_SLOTS sets the slot count (default 2); 0 disables locking.
|
||||
Waiting is a blocking flock on a turnstile file plus a slow poll of the slots,
|
||||
so contenders queue roughly first-come-first-served without busy-spinning.
|
||||
A process that acquired (or deliberately skipped) a slot exports
|
||||
LITELLM_GATE_SLOT_HELD, and nested acquisitions under that marker are no-ops,
|
||||
so `make check` invoking the gates internally can never deadlock against
|
||||
itself. Any filesystem error fails open and the command runs unlocked: the
|
||||
lock is a courtesy to the machine, never a gate that may break a build (CI
|
||||
runs one job per machine, so there it only ever takes the instant path).
|
||||
|
||||
CLI: python3 scripts/gate_slot_lock.py <command> [args...]
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import contextlib
|
||||
import fcntl
|
||||
import os
|
||||
import subprocess
|
||||
import sys
|
||||
import time
|
||||
from pathlib import Path
|
||||
from typing import IO, TYPE_CHECKING, Final
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from collections.abc import Iterator
|
||||
|
||||
HELD_MARKER_ENV: Final = "LITELLM_GATE_SLOT_HELD"
|
||||
SLOT_COUNT_ENV: Final = "LITELLM_GATE_SLOTS"
|
||||
SLOT_DIR_ENV: Final = "LITELLM_GATE_SLOT_DIR"
|
||||
DEFAULT_SLOT_COUNT: Final = 2
|
||||
POLL_SECONDS: Final = 2.0
|
||||
|
||||
|
||||
def _slot_dir() -> Path:
|
||||
override: Final = os.environ.get(SLOT_DIR_ENV)
|
||||
return Path(override) if override else Path.home() / ".cache" / "litellm" / "gate-slots"
|
||||
|
||||
|
||||
def _slot_count() -> int:
|
||||
raw: Final = os.environ.get(SLOT_COUNT_ENV)
|
||||
if not raw:
|
||||
return DEFAULT_SLOT_COUNT
|
||||
try:
|
||||
return int(raw)
|
||||
except ValueError:
|
||||
print(
|
||||
f"gate_slot_lock: ignoring non-integer {SLOT_COUNT_ENV}={raw!r}; "
|
||||
f"using {DEFAULT_SLOT_COUNT} slots",
|
||||
file=sys.stderr,
|
||||
)
|
||||
return DEFAULT_SLOT_COUNT
|
||||
|
||||
|
||||
def _try_slot(directory: Path, index: int) -> IO[bytes] | None:
|
||||
handle: Final = (directory / f"slot-{index}.lock").open("wb")
|
||||
try:
|
||||
fcntl.flock(handle, fcntl.LOCK_EX | fcntl.LOCK_NB)
|
||||
except BlockingIOError:
|
||||
handle.close()
|
||||
return None
|
||||
except OSError:
|
||||
handle.close()
|
||||
raise
|
||||
return handle
|
||||
|
||||
|
||||
def _wait_for_slot(directory: Path, count: int) -> IO[bytes]:
|
||||
print(
|
||||
f"gate_slot_lock: all {count} machine-wide slots are busy; queueing "
|
||||
f"(set {SLOT_COUNT_ENV}=0 to disable)",
|
||||
file=sys.stderr,
|
||||
flush=True,
|
||||
)
|
||||
with (directory / "turnstile.lock").open("wb") as turnstile:
|
||||
fcntl.flock(turnstile, fcntl.LOCK_EX)
|
||||
while True:
|
||||
for index in range(count):
|
||||
held = _try_slot(directory, index)
|
||||
if held is not None:
|
||||
return held
|
||||
time.sleep(POLL_SECONDS)
|
||||
|
||||
|
||||
def _locked_handle(count: int) -> IO[bytes]:
|
||||
directory: Final = _slot_dir()
|
||||
directory.mkdir(parents=True, exist_ok=True)
|
||||
for index in range(count):
|
||||
immediate = _try_slot(directory, index)
|
||||
if immediate is not None:
|
||||
return immediate
|
||||
return _wait_for_slot(directory, count)
|
||||
|
||||
|
||||
def acquire_slot() -> IO[bytes] | None:
|
||||
"""Hold a machine-wide slot for the life of the returned handle.
|
||||
|
||||
The caller must keep the handle referenced until the process exits;
|
||||
dropping it closes the file and releases the slot. Returns None without
|
||||
locking when this process already runs under a held slot, when locking is
|
||||
disabled, or when the filesystem refuses to cooperate."""
|
||||
if os.environ.get(HELD_MARKER_ENV):
|
||||
return None
|
||||
count: Final = _slot_count()
|
||||
if count <= 0:
|
||||
os.environ[HELD_MARKER_ENV] = "1"
|
||||
return None
|
||||
try:
|
||||
handle: Final = _locked_handle(count)
|
||||
except (OSError, RuntimeError) as error:
|
||||
print(f"gate_slot_lock: locking unavailable ({error}); running unlocked", file=sys.stderr)
|
||||
os.environ[HELD_MARKER_ENV] = "1"
|
||||
return None
|
||||
os.environ[HELD_MARKER_ENV] = "1"
|
||||
return handle
|
||||
|
||||
|
||||
@contextlib.contextmanager
|
||||
def held_slot() -> Iterator[None]:
|
||||
"""Run the with-block while holding a machine-wide slot (or its no-op forms)."""
|
||||
prior_marker: Final = os.environ.get(HELD_MARKER_ENV)
|
||||
handle: Final = acquire_slot()
|
||||
try:
|
||||
yield
|
||||
finally:
|
||||
if handle is not None:
|
||||
handle.close()
|
||||
if not prior_marker:
|
||||
os.environ.pop(HELD_MARKER_ENV, None)
|
||||
|
||||
|
||||
def _wait_ignoring_interrupts(process: subprocess.Popen[bytes]) -> int:
|
||||
while True:
|
||||
try:
|
||||
return process.wait()
|
||||
except KeyboardInterrupt:
|
||||
continue
|
||||
|
||||
|
||||
def main() -> int:
|
||||
if len(sys.argv) < 2:
|
||||
print("usage: gate_slot_lock.py <command> [args...]", file=sys.stderr)
|
||||
return 2
|
||||
try:
|
||||
held: Final = acquire_slot()
|
||||
except KeyboardInterrupt:
|
||||
return 130
|
||||
try:
|
||||
code: Final = _wait_ignoring_interrupts(subprocess.Popen(sys.argv[1:]))
|
||||
except FileNotFoundError as error:
|
||||
print(f"gate_slot_lock: {error}", file=sys.stderr)
|
||||
return 127
|
||||
if held is not None:
|
||||
held.close()
|
||||
return code if code >= 0 else 128 - code
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(main())
|
||||
|
|
@ -24,6 +24,16 @@
|
|||
|
||||
set -eu
|
||||
|
||||
# Queue for one of the machine-wide heavy-work slots (see scripts/gate_slot_lock.py)
|
||||
# before anything else, so N parallel `make check` runs across worktrees execute two
|
||||
# at a time instead of thrashing the machine. The wrapper exports
|
||||
# LITELLM_GATE_SLOT_HELD, so this re-exec happens exactly once and everything this
|
||||
# script spawns (make lint, the budget gates) skips its own acquisition.
|
||||
if [ -z "${LITELLM_GATE_SLOT_HELD:-}" ]; then
|
||||
script_dir=$(python3 -c 'import os, sys; print(os.path.dirname(os.path.realpath(sys.argv[1])))' "$0")
|
||||
exec python3 "$script_dir/gate_slot_lock.py" "$0" "$@"
|
||||
fi
|
||||
|
||||
if [ -z "${PRE_COMMIT_LINT_INNER:-}" ]; then
|
||||
log_file=$(git rev-parse --path-format=absolute --git-path pre_commit_lint.log)
|
||||
if : > "$log_file" 2>/dev/null; then
|
||||
|
|
|
|||
|
|
@ -215,7 +215,10 @@ def main() -> None:
|
|||
parser.add_argument("--base", default=DEFAULT_BASE)
|
||||
parser.add_argument("--update", action="store_true")
|
||||
args = parser.parse_args()
|
||||
cmd_update(args.base) if args.update else cmd_check(args.base)
|
||||
from gate_slot_lock import held_slot
|
||||
|
||||
with held_slot():
|
||||
cmd_update(args.base) if args.update else cmd_check(args.base)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
|
|
|
|||
|
|
@ -670,16 +670,19 @@ def main() -> None:
|
|||
parser.add_argument("--update", action="store_true")
|
||||
parser.add_argument("--emit-counts-dir", type=Path)
|
||||
args = parser.parse_args()
|
||||
ensure_typecheck_env()
|
||||
head = count_basedpyright(run_basedpyright())
|
||||
if args.emit_counts_dir is not None:
|
||||
cmd_emit_counts(
|
||||
head, args.emit_counts_dir, _run(["git", "rev-parse", "HEAD"]).strip()
|
||||
)
|
||||
elif args.update:
|
||||
cmd_update(head, args.base)
|
||||
else:
|
||||
cmd_check(head, args.base)
|
||||
from gate_slot_lock import held_slot
|
||||
|
||||
with held_slot():
|
||||
ensure_typecheck_env()
|
||||
head = count_basedpyright(run_basedpyright())
|
||||
if args.emit_counts_dir is not None:
|
||||
cmd_emit_counts(
|
||||
head, args.emit_counts_dir, _run(["git", "rev-parse", "HEAD"]).strip()
|
||||
)
|
||||
elif args.update:
|
||||
cmd_update(head, args.base)
|
||||
else:
|
||||
cmd_check(head, args.base)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
|
|
|
|||
|
|
@ -267,7 +267,10 @@ def main() -> None:
|
|||
parser.add_argument("--base", default=DEFAULT_BASE)
|
||||
parser.add_argument("--update", action="store_true")
|
||||
args = parser.parse_args()
|
||||
cmd_update(args.base) if args.update else cmd_check(args.base)
|
||||
from gate_slot_lock import held_slot
|
||||
|
||||
with held_slot():
|
||||
cmd_update(args.base) if args.update else cmd_check(args.base)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
|
|
|
|||
311
tests/test_litellm/test_gate_slot_lock.py
Normal file
311
tests/test_litellm/test_gate_slot_lock.py
Normal file
|
|
@ -0,0 +1,311 @@
|
|||
import fcntl
|
||||
import importlib.util
|
||||
import os
|
||||
import signal
|
||||
import subprocess
|
||||
import sys
|
||||
import time
|
||||
from collections.abc import Callable, Sequence
|
||||
from contextlib import suppress
|
||||
from pathlib import Path
|
||||
|
||||
import pytest
|
||||
|
||||
ROOT = Path(__file__).resolve().parents[2]
|
||||
HELPER = ROOT / "scripts" / "gate_slot_lock.py"
|
||||
|
||||
_spec = importlib.util.spec_from_file_location("gate_slot_lock", HELPER)
|
||||
assert _spec is not None and _spec.loader is not None
|
||||
gate_slot_lock = importlib.util.module_from_spec(_spec)
|
||||
_spec.loader.exec_module(gate_slot_lock)
|
||||
|
||||
START_THEN_WAIT_FOR = (
|
||||
"import pathlib, sys, time\n"
|
||||
"pathlib.Path(sys.argv[1]).touch()\n"
|
||||
"deadline = time.monotonic() + 20\n"
|
||||
"while not pathlib.Path(sys.argv[2]).exists():\n"
|
||||
" if time.monotonic() > deadline:\n"
|
||||
" sys.exit(3)\n"
|
||||
" time.sleep(0.05)\n"
|
||||
)
|
||||
|
||||
TOUCH_TARGET = "import pathlib, sys\npathlib.Path(sys.argv[1]).touch()\n"
|
||||
|
||||
RECORD_INTERVAL = (
|
||||
"import sys, time\n"
|
||||
"with open(sys.argv[1], 'a') as events:\n"
|
||||
" events.write(f'start {time.monotonic()}\\n')\n"
|
||||
" events.flush()\n"
|
||||
" time.sleep(0.6)\n"
|
||||
" events.write(f'end {time.monotonic()}\\n')\n"
|
||||
" events.flush()\n"
|
||||
)
|
||||
|
||||
|
||||
def _env(lock_dir: Path, slots: str) -> dict[str, str]:
|
||||
return {
|
||||
"PATH": os.environ["PATH"],
|
||||
"HOME": str(lock_dir.parent),
|
||||
"LITELLM_GATE_SLOT_DIR": str(lock_dir),
|
||||
"LITELLM_GATE_SLOTS": slots,
|
||||
}
|
||||
|
||||
|
||||
def _wrapped(payload: Sequence[str]) -> list[str]:
|
||||
return [sys.executable, str(HELPER), sys.executable, "-c", *payload]
|
||||
|
||||
|
||||
def _wait_until(predicate: Callable[[], bool], timeout_seconds: float) -> bool:
|
||||
deadline = time.monotonic() + timeout_seconds
|
||||
while time.monotonic() < deadline:
|
||||
if predicate():
|
||||
return True
|
||||
time.sleep(0.05)
|
||||
return predicate()
|
||||
|
||||
|
||||
def _terminate_group(process: subprocess.Popen[bytes]) -> None:
|
||||
with suppress(ProcessLookupError, PermissionError):
|
||||
os.killpg(process.pid, signal.SIGKILL)
|
||||
|
||||
|
||||
def _reap(process: subprocess.Popen[bytes]) -> None:
|
||||
with suppress(subprocess.TimeoutExpired):
|
||||
process.wait(timeout=10)
|
||||
if process.poll() is None:
|
||||
process.kill()
|
||||
process.wait(timeout=10)
|
||||
|
||||
|
||||
def test_six_contenders_never_exceed_two_slots_and_all_complete(tmp_path: Path) -> None:
|
||||
lock_dir = tmp_path / "locks"
|
||||
events_file = tmp_path / "events.log"
|
||||
env = _env(lock_dir, "2")
|
||||
procs = [
|
||||
subprocess.Popen(
|
||||
_wrapped([RECORD_INTERVAL, str(events_file)]),
|
||||
env=env,
|
||||
stdout=subprocess.DEVNULL,
|
||||
stderr=subprocess.DEVNULL,
|
||||
)
|
||||
for _ in range(6)
|
||||
]
|
||||
try:
|
||||
assert [proc.wait(timeout=60) for proc in procs] == [0] * 6
|
||||
finally:
|
||||
for proc in procs:
|
||||
if proc.poll() is None:
|
||||
proc.kill()
|
||||
proc.wait(timeout=10)
|
||||
events = sorted(
|
||||
(float(stamp), 1 if kind == "start" else -1)
|
||||
for kind, stamp in (line.split() for line in events_file.read_text().splitlines())
|
||||
)
|
||||
assert len(events) == 12
|
||||
concurrency_peaks = []
|
||||
running = 0
|
||||
for _, delta in events:
|
||||
running += delta
|
||||
concurrency_peaks.append(running)
|
||||
assert max(concurrency_peaks) <= 2
|
||||
|
||||
|
||||
def test_two_slots_admit_two_holders_at_once(tmp_path: Path) -> None:
|
||||
lock_dir = tmp_path / "locks"
|
||||
first_started = tmp_path / "first.started"
|
||||
second_started = tmp_path / "second.started"
|
||||
env = _env(lock_dir, "2")
|
||||
first = subprocess.Popen(
|
||||
_wrapped([START_THEN_WAIT_FOR, str(first_started), str(second_started)]), env=env
|
||||
)
|
||||
second = subprocess.Popen(
|
||||
_wrapped([START_THEN_WAIT_FOR, str(second_started), str(first_started)]), env=env
|
||||
)
|
||||
assert first.wait(timeout=30) == 0
|
||||
assert second.wait(timeout=30) == 0
|
||||
|
||||
|
||||
def test_contender_beyond_capacity_queues_until_the_slot_frees(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)
|
||||
assert not done.exists()
|
||||
release.touch()
|
||||
assert holder.wait(timeout=10) == 0
|
||||
assert contender.wait(timeout=30) == 0
|
||||
assert done.exists()
|
||||
assert contender.stderr is not None
|
||||
assert b"queueing" in contender.stderr.read()
|
||||
finally:
|
||||
release.touch()
|
||||
_reap(contender)
|
||||
finally:
|
||||
release.touch()
|
||||
_reap(holder)
|
||||
|
||||
|
||||
def test_nested_wrapping_reenters_instead_of_deadlocking(tmp_path: Path) -> None:
|
||||
lock_dir = tmp_path / "locks"
|
||||
nested = [
|
||||
sys.executable,
|
||||
str(HELPER),
|
||||
sys.executable,
|
||||
str(HELPER),
|
||||
sys.executable,
|
||||
"-c",
|
||||
"print('nested ok')",
|
||||
]
|
||||
proc = subprocess.Popen(
|
||||
nested,
|
||||
env=_env(lock_dir, "1"),
|
||||
stdout=subprocess.PIPE,
|
||||
start_new_session=True,
|
||||
)
|
||||
try:
|
||||
stdout, _ = proc.communicate(timeout=20)
|
||||
except subprocess.TimeoutExpired:
|
||||
_terminate_group(proc)
|
||||
pytest.fail("nested gate_slot_lock invocations deadlocked")
|
||||
assert proc.returncode == 0
|
||||
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"),
|
||||
)
|
||||
assert proc.returncode == 7
|
||||
|
||||
|
||||
def test_missing_command_exits_127_and_no_command_exits_2(tmp_path: Path) -> None:
|
||||
env = _env(tmp_path / "locks", "2")
|
||||
missing = subprocess.run(
|
||||
[sys.executable, str(HELPER), str(tmp_path / "no-such-binary")],
|
||||
env=env,
|
||||
capture_output=True,
|
||||
)
|
||||
assert missing.returncode == 127
|
||||
bare = subprocess.run([sys.executable, str(HELPER)], env=env, capture_output=True)
|
||||
assert bare.returncode == 2
|
||||
|
||||
|
||||
def test_wrapped_command_killed_by_signal_maps_to_128_plus_signal(tmp_path: Path) -> None:
|
||||
proc = subprocess.run(
|
||||
_wrapped(["import os, signal\nos.kill(os.getpid(), signal.SIGTERM)\n"]),
|
||||
env=_env(tmp_path / "locks", "2"),
|
||||
)
|
||||
assert proc.returncode == 128 + signal.SIGTERM
|
||||
|
||||
|
||||
def test_unusable_lock_dir_fails_open_and_still_runs_the_command(tmp_path: Path) -> None:
|
||||
blocker = tmp_path / "blocker"
|
||||
blocker.write_text("")
|
||||
done = tmp_path / "done"
|
||||
proc = subprocess.run(
|
||||
_wrapped([TOUCH_TARGET, str(done)]),
|
||||
env=_env(blocker / "locks", "2"),
|
||||
capture_output=True,
|
||||
)
|
||||
assert proc.returncode == 0
|
||||
assert done.exists()
|
||||
assert b"running unlocked" in proc.stderr
|
||||
|
||||
|
||||
def test_zero_slots_disables_locking_entirely(tmp_path: Path) -> None:
|
||||
lock_dir = tmp_path / "locks"
|
||||
done = tmp_path / "done"
|
||||
proc = subprocess.run(
|
||||
_wrapped([TOUCH_TARGET, str(done)]),
|
||||
env=_env(lock_dir, "0"),
|
||||
)
|
||||
assert proc.returncode == 0
|
||||
assert done.exists()
|
||||
assert not lock_dir.exists()
|
||||
|
||||
|
||||
def test_non_integer_slot_count_warns_and_falls_back_to_default(tmp_path: Path) -> None:
|
||||
proc = subprocess.run(
|
||||
[sys.executable, str(HELPER), sys.executable, "-c", "print('ran')"],
|
||||
env=_env(tmp_path / "locks", "lots"),
|
||||
capture_output=True,
|
||||
)
|
||||
assert proc.returncode == 0
|
||||
assert b"ran" in proc.stdout
|
||||
assert b"LITELLM_GATE_SLOTS" in proc.stderr
|
||||
|
||||
|
||||
def test_killed_holder_releases_its_slot_for_the_next_contender(tmp_path: Path) -> None:
|
||||
lock_dir = tmp_path / "locks"
|
||||
holder_started = tmp_path / "holder.started"
|
||||
never = tmp_path / "never"
|
||||
env = _env(lock_dir, "1")
|
||||
holder = subprocess.Popen(
|
||||
_wrapped([START_THEN_WAIT_FOR, str(holder_started), str(never)]),
|
||||
env=env,
|
||||
start_new_session=True,
|
||||
)
|
||||
try:
|
||||
assert _wait_until(holder_started.exists, 10)
|
||||
finally:
|
||||
_terminate_group(holder)
|
||||
holder.wait(timeout=10)
|
||||
after = subprocess.run(
|
||||
[sys.executable, str(HELPER), sys.executable, "-c", "print('freed')"],
|
||||
env=env,
|
||||
capture_output=True,
|
||||
timeout=20,
|
||||
)
|
||||
assert after.returncode == 0
|
||||
assert b"freed" in after.stdout
|
||||
|
||||
|
||||
def test_acquire_slot_holds_marks_and_releases_in_process(
|
||||
tmp_path: Path, monkeypatch: pytest.MonkeyPatch
|
||||
) -> None:
|
||||
lock_dir = tmp_path / "locks"
|
||||
monkeypatch.setenv("LITELLM_GATE_SLOT_HELD", "")
|
||||
monkeypatch.setenv("LITELLM_GATE_SLOT_DIR", str(lock_dir))
|
||||
monkeypatch.setenv("LITELLM_GATE_SLOTS", "1")
|
||||
handle = gate_slot_lock.acquire_slot()
|
||||
assert handle is not None
|
||||
assert os.environ["LITELLM_GATE_SLOT_HELD"] == "1"
|
||||
assert gate_slot_lock.acquire_slot() is None
|
||||
with (lock_dir / "slot-0.lock").open("wb") as probe:
|
||||
with pytest.raises(BlockingIOError):
|
||||
fcntl.flock(probe, fcntl.LOCK_EX | fcntl.LOCK_NB)
|
||||
handle.close()
|
||||
fcntl.flock(probe, fcntl.LOCK_EX | fcntl.LOCK_NB)
|
||||
fcntl.flock(probe, fcntl.LOCK_UN)
|
||||
|
||||
|
||||
def test_held_slot_context_manager_releases_on_exit(
|
||||
tmp_path: Path, monkeypatch: pytest.MonkeyPatch
|
||||
) -> None:
|
||||
lock_dir = tmp_path / "locks"
|
||||
monkeypatch.setenv("LITELLM_GATE_SLOT_HELD", "")
|
||||
monkeypatch.setenv("LITELLM_GATE_SLOT_DIR", str(lock_dir))
|
||||
monkeypatch.setenv("LITELLM_GATE_SLOTS", "1")
|
||||
with gate_slot_lock.held_slot():
|
||||
assert os.environ["LITELLM_GATE_SLOT_HELD"] == "1"
|
||||
with (lock_dir / "slot-0.lock").open("wb") as probe:
|
||||
with pytest.raises(BlockingIOError):
|
||||
fcntl.flock(probe, fcntl.LOCK_EX | fcntl.LOCK_NB)
|
||||
assert not os.environ.get("LITELLM_GATE_SLOT_HELD")
|
||||
with (lock_dir / "slot-0.lock").open("wb") as probe:
|
||||
fcntl.flock(probe, fcntl.LOCK_EX | fcntl.LOCK_NB)
|
||||
fcntl.flock(probe, fcntl.LOCK_UN)
|
||||
|
|
@ -1,4 +1,5 @@
|
|||
import os
|
||||
import shutil
|
||||
import signal
|
||||
import subprocess
|
||||
import time
|
||||
|
|
@ -420,6 +421,46 @@ def test_staged_files_matching_no_check_print_an_explicit_noop_note_and_nonempty
|
|||
assert "skipped: Python lint (make lint) (no litellm/ Python files in scope)" in log
|
||||
|
||||
|
||||
def test_run_queues_through_the_machine_wide_gate_slot_lock(tmp_path: Path) -> None:
|
||||
repo, bin_dir = _sandbox(tmp_path)
|
||||
lock_dir = tmp_path / "gate-locks"
|
||||
proc = _run(repo, bin_dir, {"LITELLM_GATE_SLOT_DIR": str(lock_dir)})
|
||||
assert proc.returncode == 0, proc.stdout + proc.stderr
|
||||
assert (lock_dir / "slot-0.lock").exists()
|
||||
|
||||
|
||||
def test_run_under_a_held_slot_skips_reacquiring_the_gate_lock(tmp_path: Path) -> None:
|
||||
repo, bin_dir = _sandbox(tmp_path)
|
||||
lock_dir = tmp_path / "gate-locks"
|
||||
proc = _run(
|
||||
repo,
|
||||
bin_dir,
|
||||
{"LITELLM_GATE_SLOT_DIR": str(lock_dir), "LITELLM_GATE_SLOT_HELD": "1"},
|
||||
)
|
||||
assert proc.returncode == 0, proc.stdout + proc.stderr
|
||||
assert not lock_dir.exists()
|
||||
|
||||
|
||||
def test_hook_symlink_install_still_resolves_the_slot_lock_helper(tmp_path: Path) -> None:
|
||||
repo, bin_dir = _sandbox(tmp_path)
|
||||
scripts_dir = repo / "scripts"
|
||||
scripts_dir.mkdir()
|
||||
shutil.copy(SCRIPT, scripts_dir / "pre_commit_lint.sh")
|
||||
shutil.copy(SCRIPT.parent / "gate_slot_lock.py", scripts_dir / "gate_slot_lock.py")
|
||||
(repo / ".git" / "hooks" / "pre-commit").symlink_to(Path("../../scripts/pre_commit_lint.sh"))
|
||||
lock_dir = tmp_path / "gate-locks"
|
||||
proc = subprocess.run(
|
||||
["git", "-c", "user.email=t@t", "-c", "user.name=t", "commit", "-qm", "hooked"],
|
||||
cwd=repo,
|
||||
capture_output=True,
|
||||
text=True,
|
||||
env=_env(repo, bin_dir, {"LITELLM_GATE_SLOT_DIR": str(lock_dir)}),
|
||||
timeout=120,
|
||||
)
|
||||
assert proc.returncode == 0, proc.stdout + proc.stderr
|
||||
assert (lock_dir / "slot-0.lock").exists()
|
||||
|
||||
|
||||
def test_failing_run_ends_with_a_fail_verdict(tmp_path: Path) -> None:
|
||||
repo, bin_dir = _sandbox(tmp_path)
|
||||
proc = _run(repo, bin_dir, {"STUB_FAIL": "make-lint"})
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue