diff --git a/scripts/check_test_quality.py b/scripts/check_test_quality.py index 6f40dd6ea84..09b87039048 100644 --- a/scripts/check_test_quality.py +++ b/scripts/check_test_quality.py @@ -63,6 +63,31 @@ TQ009 A child interpreter spawned as `subprocess.run([sys.executable, ...])` w same working-directory hazard seen from the child's side. Use tests.test_litellm_rust.support.child_interpreter.run_child_interpreter, which also asserts the child resolved the same `litellm.__file__` as the parent. +TQ010 An integration test that signals a process itself, outside the harness in + tests/integration/_support: `os.kill`, `os.killpg`, a SIGKILL / SIGSTOP / SIGCONT + constant handed to any call, or a zero-argument `.kill()` / `.terminate()` / + `.suspend()` / `.resume()` / a `.send_signal(...)` method call. A hand-rolled + worker kill returns as soon as a replacement process exists rather than once it + serves, and a hand-rolled suspend outlives the owned proxy's worker healthcheck, + after which uvicorn kills the suspended worker under the test. Use the owned + proxy's `kill_worker` / `replacing_worker` / `suspended` and the process helpers in + tests/integration/_support/process.py. + Only files under tests/integration, outside tests/integration/_support, are + examined. +TQ011 An integration-test assertion that compares an elapsed time against a numeric + literal (`assert elapsed < 2`, `assert time.monotonic() - started <= 5`, + `assert all(duration < 2 for ...)`). A fixed wall-clock bound measures the + machine as much as the code, so it fails whenever sibling workers load the box. + Assert relative to a measured control or to the injected delay instead, or, + when the time is the behaviour under test (a cap, timeout or cooldown the + product enforces), against that configured value plus explicit slack with + `assert_enforced_within` (tests/integration/_support/timing.py). An operand is an elapsed time when it + reads a clock (`time.monotonic()`, `time.perf_counter()`, `time.time()` and + their `_ns` forms), names something called elapsed / duration / latency, or + names a local the enclosing function bound to a clock reading. The other operand + is a literal when it is built only from numbers, or names a module-level number. + Only `assert` statements are examined, so a polling deadline in a `while` loop is + left alone. Same path scope as TQ010. Every rule is suppressible with `# test-quality-ok: ` on the reported line, following the repo's `*-ok: ` convention. A suppression without a @@ -147,9 +172,24 @@ SUBPROCESS_SPAWNS: Final = frozenset(("run", "Popen", "check_output", "check_cal INTERPRETER_ISOLATION_FLAGS: Final = frozenset(("-I", "-P")) RULE_CODES: Final = frozenset(( - "TQ000", "TQ001", "TQ002", "TQ003", "TQ004", "TQ005", "TQ006", "TQ007", "TQ009", + "TQ000", "TQ001", "TQ002", "TQ003", "TQ004", "TQ005", "TQ006", "TQ007", "TQ009", "TQ010", "TQ011", )) +INTEGRATION_ROOT: Final = ("tests", "integration") +INTEGRATION_HARNESS: Final = "_support" +SIGNAL_SENDERS: Final = frozenset(("os.kill", "os.killpg")) +PROCESS_CONTROL_SIGNALS: Final = frozenset(("SIGKILL", "SIGSTOP", "SIGCONT")) +ZERO_ARGUMENT_PROCESS_CONTROLS: Final = frozenset(("kill", "terminate", "suspend", "resume")) +SIGNAL_METHOD: Final = "send_signal" +CLOCK_READERS: Final = frozenset(( + "time.monotonic", "time.perf_counter", "time.time", + "time.monotonic_ns", "time.perf_counter_ns", "time.time_ns", + "monotonic", "perf_counter", "monotonic_ns", "perf_counter_ns", +)) +ELAPSED_NAME_RE: Final = re.compile(r"elapsed|duration|latenc", re.IGNORECASE) +DURATION_READER: Final = "total_seconds" +DURATION_FOLDS: Final = frozenset(("max", "min", "sum", "abs", "round", "float")) + CREDENTIAL_NAME_RE: Final = re.compile( r"(?:API_KEY|_KEY|TOKEN|SECRET|PASSWORD|CREDENTIAL|DATABASE_URL|ACCESS_KEY_ID)$" ) @@ -687,6 +727,156 @@ def iter_child_interpreter_violations(path: Path, tree: ast.Module) -> Iterator[ ) +def is_integration_test_path(path: Path) -> bool: + parts: Final = path.parts + return any( + parts[index:index + 2] == INTEGRATION_ROOT and parts[index + 2] != INTEGRATION_HARNESS + for index in range(len(parts) - 2) + ) + + +def _names_control_signal(node: ast.expr) -> bool: + if isinstance(node, ast.Attribute): + return node.attr in PROCESS_CONTROL_SIGNALS + return isinstance(node, ast.Name) and node.id in PROCESS_CONTROL_SIGNALS + + +def _signals_a_process(call: ast.Call) -> bool: + if _dotted_name(call.func) in SIGNAL_SENDERS: + return True + if any(_names_control_signal(argument) for argument in (*call.args, *(k.value for k in call.keywords))): + return True + func: Final = call.func + if not isinstance(func, ast.Attribute): + return False + if func.attr == SIGNAL_METHOD: + return True + return func.attr in ZERO_ARGUMENT_PROCESS_CONTROLS and not call.args and not call.keywords + + +def iter_process_signal_violations(path: Path, tree: ast.Module) -> Iterator[Violation]: + if not is_integration_test_path(path): + return + for node in ast.walk(tree): + if isinstance(node, ast.Call) and _signals_a_process(node): + yield Violation( + path, + node.lineno, + "TQ010", + "an integration test signals a process itself; a hand-rolled kill returns before the " + "replacement worker serves and a hand-rolled suspend outlives the worker healthcheck, so use " + "the owned proxy's kill_worker / replacing_worker / suspended, or a tests/integration/_support helper " + f"(suppress: `# {SUPPRESSION_TOKEN}: `)", + ) + + +def _reads_clock(node: ast.expr) -> bool: + if isinstance(node, ast.Call): + return _dotted_name(node.func) in CLOCK_READERS + if isinstance(node, ast.BinOp): + return _reads_clock(node.left) or _reads_clock(node.right) + return False + + +def _bound_values(scope: ast.AST) -> Iterator[tuple[ast.expr, ast.expr]]: + for node in ast.walk(scope): + if isinstance(node, ast.Assign): + yield from ((target, node.value) for target in node.targets) + elif isinstance(node, (ast.AnnAssign, ast.NamedExpr)) and node.value is not None: + yield node.target, node.value + + +def _clock_bound_names(scope: ast.AST) -> frozenset[str]: + return frozenset( + target.id for target, value in _bound_values(scope) if isinstance(target, ast.Name) and _reads_clock(value) + ) + + +def _numeric_constant_names(tree: ast.Module) -> frozenset[str]: + return frozenset( + target.id + for statement in tree.body + for target, value in _bound_values(statement) + if isinstance(target, ast.Name) and _is_numeric_literal(value, frozenset()) + ) + + +def _is_numeric_literal(node: ast.expr, constants: frozenset[str]) -> bool: + if isinstance(node, ast.Constant): + return isinstance(node.value, (int, float)) and not isinstance(node.value, bool) + if isinstance(node, ast.UnaryOp): + return _is_numeric_literal(node.operand, constants) + if isinstance(node, ast.BinOp): + return _is_numeric_literal(node.left, constants) and _is_numeric_literal(node.right, constants) + return isinstance(node, ast.Name) and node.id in constants + + +def _is_elapsed(node: ast.expr, clock_names: frozenset[str]) -> bool: + if isinstance(node, ast.Name): + return node.id in clock_names or ELAPSED_NAME_RE.search(node.id) is not None + if isinstance(node, ast.Attribute): + return ELAPSED_NAME_RE.search(node.attr) is not None + if isinstance(node, ast.BinOp): + return _is_elapsed(node.left, clock_names) or _is_elapsed(node.right, clock_names) + if isinstance(node, ast.Subscript): + return _is_elapsed(node.value, clock_names) + if isinstance(node, (ast.GeneratorExp, ast.ListComp)): + return _is_elapsed(node.elt, clock_names) + if not isinstance(node, ast.Call): + return False + if _dotted_name(node.func) in CLOCK_READERS: + return True + if isinstance(node.func, ast.Attribute) and node.func.attr == DURATION_READER: + return _is_elapsed(node.func.value, clock_names) + return _dotted_name(node.func) in DURATION_FOLDS and any(_is_elapsed(arg, clock_names) for arg in node.args) + + +def _upper_bounded_operands(compare: ast.Compare) -> Iterator[tuple[ast.expr, ast.expr]]: + operands: Final = (compare.left, *compare.comparators) + for left, operator, right in zip(operands, compare.ops, operands[1:]): + if isinstance(operator, (ast.Lt, ast.LtE)): + yield left, right + elif isinstance(operator, (ast.Gt, ast.GtE)): + yield right, left + + +def _asserts_in_scope(node: ast.AST, scope: ast.AST) -> Iterator[tuple[ast.Assert, ast.AST]]: + for child in ast.iter_child_nodes(node): + if isinstance(child, ast.Assert): + yield child, scope + inner: Final = child if isinstance(child, (ast.FunctionDef, ast.AsyncFunctionDef)) else scope + yield from _asserts_in_scope(child, inner) + + +def iter_wall_clock_bound_violations(path: Path, tree: ast.Module) -> Iterator[Violation]: + if not is_integration_test_path(path): + return + constants: Final = _numeric_constant_names(tree) + for statement, scope in _asserts_in_scope(tree, tree): + clock_names: Final = _clock_bound_names(scope) + capped: Final = tuple( + compare + for compare in ast.walk(statement.test) + if isinstance(compare, ast.Compare) + and any( + _is_elapsed(bounded, clock_names) and _is_numeric_literal(bound, constants) + for bounded, bound in _upper_bounded_operands(compare) + ) + ) + yield from ( + Violation( + path, + compare.lineno, + "TQ011", + "an elapsed time is capped by a fixed number, which fails whenever sibling workers load the " + "machine; bound it by a measured control or by the injected delay, or a product-enforced limit by " + "its configured value plus slack with assert_enforced_within " + f"(tests/integration/_support/timing.py) (suppress: `# {SUPPRESSION_TOKEN}: `)", + ) + for compare in capped + ) + + def iter_conftest_inventory_violations(path: Path, tree: ast.Module) -> Iterator[Violation]: if path.name != CONFTEST_NAME: return @@ -724,6 +914,8 @@ def check_file(path: Path) -> tuple[Violation, ...]: *iter_credential_skip_violations(path, tree), *iter_conftest_inventory_violations(path, tree), *iter_child_interpreter_violations(path, tree), + *iter_process_signal_violations(path, tree), + *iter_wall_clock_bound_violations(path, tree), ) if violation.line not in skip ) diff --git a/tests/integration/README.md b/tests/integration/README.md index b65c4f36219..572b80ad5c6 100644 --- a/tests/integration/README.md +++ b/tests/integration/README.md @@ -26,7 +26,7 @@ Reuse the existing canned provider handlers through `_support/upstream.py`. It r Owned proxies disable `.env` loading by default. A contract testing dotenv behavior can explicitly set `PYTHON_DOTENV_DISABLED=0` in its owned environment. An explicit scratch writer database drops an inherited reader with a different database name; pass both database URLs when the contract needs a paired reader. Database relays close listeners, active connections and their event loops on teardown. Upstream startup failures reap the child process before returning the failure. Custom subprocesses, direct SQL and background tasks still need an explicit owner in the test -Owned proxy workers configure a 32-thread executor before application startup, preserving additional worker startup hooks. Held synchronous provider calls must not consume a CPU-dependent executor limit before the rest of a fault burst can start. The proxy entrypoint imports LiteLLM inside `main`, so spawned workers begin supervision before importing the application. Worker healthchecks have a five-second deadline; application readiness and graceful teardown retain their separate longer deadlines. A fault that deliberately pauses Python before it can answer worker pings must supply a longer healthcheck through `extra_arguments`. Slow application startup is covered separately by the readiness deadline. Cases with two complete owned proxy lifetimes use a 180-second pytest budget for startup, assertions and teardown; their HTTP, provider and persistence deadlines stay unchanged. Readiness requires responsive worker identities that are still alive and a positively identified owned listener. Unreadable socket records cannot establish ownership. Crash tests own their client tasks with `TaskGroup`, release held peers in `finally` and verify that the intended worker dies. A test requiring traffic on both workers prepares that traffic explicitly instead of assuming OS load balancing +Owned proxy workers configure a 32-thread executor before application startup, preserving additional worker startup hooks. Held synchronous provider calls must not consume a CPU-dependent executor limit before the rest of a fault burst can start. The proxy entrypoint imports LiteLLM inside `main`, so spawned workers begin supervision before importing the application. Worker healthchecks have a five-second deadline by default, so uvicorn kills a worker that stays paused longer than that. Tests never signal proxy processes themselves. `OwnedProxy.kill_worker` SIGKILLs a worker and returns once its replacement answers readiness with its own pid, `replacing_worker` does the same for a kill that has to land mid-burst, and `suspended` pauses a worker or the supervisor, which only works on a proxy launched with `healthcheck=WorkerHealthcheck.LENIENT`. Helper processes that are not the proxy stop through `stop_child`, `stop_processes` and `stop_process_group` and pause through `paused`. `scripts/check_test_quality.py` rule TQ010 rejects a raw `os.kill`, SIGKILL / SIGSTOP / SIGCONT, `.kill()`, `.terminate()`, `.suspend()` or `.send_signal()` anywhere under `tests/integration` outside `_support`, and TQ011 rejects an assertion that caps an elapsed time with a fixed number. Bound timings with `_support/timing.py` instead, relative to a measured control or to the injected delay. `run.py` measures how long this box takes to import the proxy application before it starts pytest, records it in `boot-calibration.json` beside the results and exports it as `INTEGRATION_IMPORT_SECONDS`. Readiness deadlines, graceful stops and the default per-test timeout are multiples of that measurement (`_support/budget.py`), so a slower box gets longer budgets and a two-worker proxy gets twice the readiness deadline of a one-worker one. Readiness requires responsive worker identities that are still alive and a positively identified owned listener. Unreadable socket records cannot establish ownership. Crash tests own their client tasks with `TaskGroup`, release held peers in `finally` and verify that the intended worker dies. A test requiring traffic on both workers prepares that traffic explicitly instead of assuming OS load balancing, for example by sending each half while the other worker is suspended The CircleCI workflow starts its own database and Redis, restricts test-phase egress to its owned services and writes JUnit plus an executed-node manifest. Missing setup, failed cleanup or a selected test with neither a passed call nor a skip fail qualification. Skipped nodes are listed under `skipped` in `execution.json`, so the skip reasons double as the open bug list. `GITHUB_FILES` in `run.py` lists the files excluded from the `run_integration.sh` selection diff --git a/tests/integration/_support/budget.py b/tests/integration/_support/budget.py new file mode 100644 index 00000000000..c9aa4b0802e --- /dev/null +++ b/tests/integration/_support/budget.py @@ -0,0 +1,64 @@ +import json +import os +import subprocess +import sys +import time +from functools import cache +from pathlib import Path +from typing import Final + +IMPORT_SECONDS_VARIABLE: Final = "INTEGRATION_IMPORT_SECONDS" +READINESS_OVERRIDE_VARIABLE: Final = "INTEGRATION_PROXY_READY_SECONDS" +READINESS_IMPORTS: Final = 12 +CALL_IMPORTS: Final = 6 +PROXY_IMPORT: Final = "import litellm.proxy.proxy_server" + + +def measure_import_seconds(root: Path, environment: dict[str, str]) -> float: + started: Final = time.monotonic() + subprocess.run( + (sys.executable, "-P", "-c", PROXY_IMPORT), cwd=root, env=environment, capture_output=True, check=True + ) + return time.monotonic() - started + + +def boot_units(workers: int) -> int: + return 1 if workers <= 1 else 2 + + +def readiness_for(imported: float, workers: int) -> float: + return READINESS_IMPORTS * imported * boot_units(workers) + + +def case_timeout_for(imported: float) -> float: + return readiness_for(imported, 2) + readiness_for(imported, 1) + CALL_IMPORTS * imported + + +def calibration(imported: float) -> str: + return json.dumps( + { + "import_seconds": imported, + "readiness_seconds_1_worker": readiness_for(imported, 1), + "readiness_seconds_2_workers": readiness_for(imported, 2), + "case_timeout_seconds": case_timeout_for(imported), + }, + indent=2, + ) + + +@cache +def import_seconds() -> float: + exported: Final = os.environ.get(IMPORT_SECONDS_VARIABLE) + if exported: + return float(exported) + root: Final = Path(os.environ.get("INTEGRATION_PROXY_ROOT") or Path(__file__).resolve().parents[3]) + return measure_import_seconds(root, dict(os.environ)) + + +def readiness_seconds(workers: int = 1) -> float: + override: Final = os.environ.get(READINESS_OVERRIDE_VARIABLE) + return float(override) if override else readiness_for(import_seconds(), workers) + + +def case_timeout_seconds() -> float: + return case_timeout_for(import_seconds()) diff --git a/tests/integration/_support/client.py b/tests/integration/_support/client.py index 1c58284e42f..b22c1b1da3a 100644 --- a/tests/integration/_support/client.py +++ b/tests/integration/_support/client.py @@ -240,11 +240,20 @@ class Scenario: info: Final = object_value(self.gateway.get("/key/info", {"key": hashed})["info"]) assert info["status"] == "deleted", f"Deleted key still served as live: {info['status']}" - def delete_model(self, identity: str) -> None: - self.gateway.post("/model/delete", {"id": identity}) + def served_model_ids(self) -> frozenset[JsonValue]: entries: Final = self.gateway.get("/model/info")["data"] assert isinstance(entries, list) - assert all(object_value(object_value(entry)["model_info"])["id"] != identity for entry in entries) + return frozenset(object_value(object_value(entry)["model_info"])["id"] for entry in entries) + + def delete_model(self, identity: str) -> None: + from litellm.constants import PROXY_CONFIG_RELOAD_INTERVAL_SECONDS + + self.gateway.post("/model/delete", {"id": identity}) + eventually( + self.served_model_ids, + lambda served: identity not in served, + seconds=2 * PROXY_CONFIG_RELOAD_INTERVAL_SECONDS, + ) assert read_rows('SELECT model_id FROM "LiteLLM_ProxyModelTable" WHERE model_id = %s', (identity,)) == [] def model(self, *, model_info: Mapping[str, JsonValue] | None = None, **parameters: JsonValue) -> str: diff --git a/tests/integration/_support/process.py b/tests/integration/_support/process.py index 597ae2370be..6ec8dd7447a 100644 --- a/tests/integration/_support/process.py +++ b/tests/integration/_support/process.py @@ -1,5 +1,6 @@ import errno import json +import threading import os import re import signal @@ -8,9 +9,10 @@ import subprocess import sys import time import uuid -from collections.abc import Generator, Iterator, Mapping -from contextlib import contextmanager +from collections.abc import Generator, Iterable, Iterator, Mapping +from contextlib import contextmanager, suppress from dataclasses import dataclass +from enum import Enum from datetime import UTC, datetime, timedelta from ipaddress import IPv4Address from itertools import accumulate @@ -28,6 +30,7 @@ from cryptography.hazmat.primitives.asymmetric import rsa from cryptography.x509.oid import NameOID from pydantic import TypeAdapter +from tests.integration._support.budget import import_seconds, readiness_for, readiness_seconds from tests.integration._support.client import GATEWAY_LIMITS, Gateway DB_PUSH: Final = ("--use_prisma_db_push",) @@ -67,7 +70,7 @@ def signal_group(group: int, action: int) -> None: def graceful_stop_seconds() -> float: - return max(30.0, float(os.environ.get("INTEGRATION_PROXY_READY_SECONDS", "70"))) + return readiness_for(import_seconds(), 1) def stop_root_process(process: subprocess.Popen[bytes]) -> bool: @@ -81,11 +84,235 @@ def stop_root_process(process: subprocess.Popen[bytes]) -> bool: return True +class WorkerHealthcheck(Enum): + STRICT = "strict" + LENIENT = "lenient" + + +STRICT_HEALTHCHECK_SECONDS: Final = 5 +WORKER_ENTRYPOINT: Final = "spawn_main" + + +def healthcheck_seconds(healthcheck: WorkerHealthcheck) -> int: + match healthcheck: + case WorkerHealthcheck.STRICT: + return STRICT_HEALTHCHECK_SECONDS + case WorkerHealthcheck.LENIENT: + return int(graceful_stop_seconds()) + + +def _is_serving_worker(process: psutil.Process) -> bool: + try: + return _ready_worker_pid(process) is not None and WORKER_ENTRYPOINT in " ".join(process.cmdline()) + except psutil.Error: + return False + + +def _is_gone(process: psutil.Process) -> bool: + try: + return process.status() == psutil.STATUS_ZOMBIE + except psutil.NoSuchProcess: + return True + + +def _wait_until_gone(process: psutil.Process, seconds: float) -> None: + deadline: Final = time.monotonic() + seconds + while not _is_gone(process): + assert time.monotonic() < deadline, f"Process {process.pid} is still running after {seconds:.0f}s" + time.sleep(0.05) + + +def _is_stopped(process: psutil.Process) -> bool: + try: + return process.status() == psutil.STATUS_STOPPED + except psutil.NoSuchProcess: + return False + + +def _is_reaped(process: psutil.Process) -> bool: + try: + process.status() + except psutil.NoSuchProcess: + return True + return False + + +def _wait_until_reaped(process: psutil.Process, supervisor: psutil.Process, seconds: float) -> None: + deadline: Final = time.monotonic() + seconds + while not _is_reaped(process) and not _is_stopped(supervisor): + assert time.monotonic() < deadline, f"The supervisor did not reap worker {process.pid} within {seconds:.0f}s" + time.sleep(0.05) + + +@dataclass(frozen=True, slots=True) +class WorkerKill: + victim: psutil.Process + supervisor: psutil.Process + reap_seconds: float + fired: threading.Event + + def freeze(self) -> psutil.Process: + self.victim.suspend() + return self.victim + + def __call__(self) -> None: + assert not self.fired.is_set(), f"Worker {self.victim.pid} was already killed" + with suppress(psutil.NoSuchProcess): + self.victim.kill() + _wait_until_gone(self.victim, STRICT_HEALTHCHECK_SECONDS) + _wait_until_reaped(self.victim, self.supervisor, self.reap_seconds) + self.fired.set() + + +def replacement_pid(before: frozenset[int], answered: int | None, serving: frozenset[int]) -> int | None: + return answered if answered is not None and answered not in before and answered in serving else None + + @dataclass(frozen=True, slots=True) class OwnedProxy: gateway: Gateway process: subprocess.Popen[bytes] log: Path + healthcheck: WorkerHealthcheck = WorkerHealthcheck.STRICT + + def worker_pids(self) -> frozenset[int]: + try: + children: Final = psutil.Process(self.process.pid).children() + except psutil.NoSuchProcess: + return frozenset() + return frozenset(child.pid for child in children if _is_serving_worker(child)) + + def _answering_pid(self) -> int | None: + try: + with httpx.Client( + base_url=str(self.gateway.client.base_url), + timeout=2, + trust_env=False, + limits=httpx.Limits(max_keepalive_connections=0), + ) as client: + response: Final = client.get("/health/integration-readiness") + except httpx.TransportError: + return None + identity: Final = response.headers.get("x-integration-worker-pid", "") + return int(identity) if response.status_code == 200 and identity.isdecimal() else None + + def _await_replacement(self, before: frozenset[int]) -> int: + budget: Final = readiness_seconds(1) + deadline: Final = time.monotonic() + budget + while (replacement := replacement_pid(before, self._answering_pid(), self.worker_pids())) is None: + assert self.process.poll() is None, "Owned proxy exited while a worker was being replaced" + assert time.monotonic() < deadline, f"No replacement worker served within {budget:.0f}s" + time.sleep(0.05) + return replacement + + def _victim(self, pid: int) -> tuple[psutil.Process, frozenset[int]]: + before: Final = self.worker_pids() + assert pid in before, f"{pid} is not a serving worker of this proxy: {sorted(before)}" + return psutil.Process(pid), before + + def _worker_kill(self, victim: psutil.Process) -> WorkerKill: + reap_seconds: Final = healthcheck_seconds(self.healthcheck) + STRICT_HEALTHCHECK_SECONDS + return WorkerKill(victim, psutil.Process(self.process.pid), reap_seconds, threading.Event()) + + @contextmanager + def replacing_worker(self, pid: int) -> Iterator[WorkerKill]: + victim, before = self._victim(pid) + kill: Final = self._worker_kill(victim) + yield kill + assert kill.fired.is_set(), f"replacing_worker({pid}) exited without killing the worker" + self._await_replacement(before) + + def kill_worker(self, pid: int) -> int: + victim, before = self._victim(pid) + self._worker_kill(victim)() + return self._await_replacement(before) + + @contextmanager + def suspended(self, pid: int) -> Iterator[psutil.Process]: + assert self.healthcheck is WorkerHealthcheck.LENIENT, ( + "uvicorn kills a suspended worker after the strict worker healthcheck; " + "launch this proxy with healthcheck=WorkerHealthcheck.LENIENT" + ) + assert pid == self.process.pid or pid in self.worker_pids(), f"{pid} is not part of this proxy" + target: Final = psutil.Process(pid) + target.suspend() + try: + yield target + finally: + with suppress(psutil.NoSuchProcess): + target.resume() + assert not _is_gone(target), ( + f"Process {pid} died while suspended under a {healthcheck_seconds(self.healthcheck)}s worker healthcheck" + ) + + def stall_worker(self, pid: int) -> int: + assert self.healthcheck is WorkerHealthcheck.STRICT, "Only the strict worker healthcheck replaces a stall" + victim, before = self._victim(pid) + victim.suspend() + try: + replacement: Final = self._await_replacement(before) + _wait_until_gone(victim, readiness_seconds(1)) + return replacement + finally: + with suppress(psutil.NoSuchProcess): + victim.kill() + + def begin_shutdown(self) -> None: + self.process.terminate() + + def kill_supervisor(self) -> int: + self.process.kill() + return self.process.wait(timeout=STRICT_HEALTHCHECK_SECONDS) + + +def stop_child(process: subprocess.Popen[bytes] | subprocess.Popen[str], grace_seconds: float = 10) -> int: + if process.poll() is None: + process.terminate() + try: + return process.wait(timeout=grace_seconds) + except subprocess.TimeoutExpired: + process.kill() + return process.wait(timeout=grace_seconds) + + +def stop_processes(processes: Iterable[psutil.Process], grace_seconds: float = 5) -> None: + members: Final = tuple(processes) + for member in members: + with suppress(psutil.NoSuchProcess): + member.terminate() + _stopped, alive = psutil.wait_procs(members, timeout=grace_seconds) + for member in alive: + with suppress(psutil.NoSuchProcess): + member.kill() + psutil.wait_procs(alive, timeout=grace_seconds) + + +def stop_process_group(process: subprocess.Popen[bytes] | subprocess.Popen[str], grace_seconds: float) -> None: + if process.poll() is None: + signal_group(process.pid, signal.SIGTERM) + try: + process.wait(timeout=grace_seconds) + except subprocess.TimeoutExpired: + signal_group(process.pid, signal.SIGKILL) + process.wait(timeout=grace_seconds) + + +def kill_stray_process(pid: int) -> None: + with suppress(psutil.NoSuchProcess): + stray: Final = psutil.Process(pid) + stray.kill() + _wait_until_gone(stray, STRICT_HEALTHCHECK_SECONDS) + + +@contextmanager +def paused(process: subprocess.Popen[bytes] | subprocess.Popen[str]) -> Iterator[None]: + target: Final = psutil.Process(process.pid) + target.suspend() + try: + yield + finally: + with suppress(psutil.NoSuchProcess): + target.resume() @contextmanager @@ -212,14 +439,14 @@ def _owns_listener(launch: _Launch) -> bool: return False -def _readiness_pids(launch: _Launch) -> Iterator[int]: +def _readiness_pids(launch: _Launch, workers: int) -> Iterator[int]: with httpx.Client( base_url=f"http://127.0.0.1:{launch.port}", timeout=15, trust_env=False, limits=httpx.Limits(max_keepalive_connections=0), ) as client: - deadline: Final = time.monotonic() + float(os.environ.get("INTEGRATION_PROXY_READY_SECONDS", "70")) + deadline: Final = time.monotonic() + readiness_seconds(workers) while launch.process.poll() is None: try: response: Final = client.get("/health/integration-readiness", timeout=2) @@ -249,7 +476,7 @@ def _ready_worker_pid(member: psutil.Process) -> int | None: def _wait_until_ready(launch: _Launch, workers: int = 1) -> None: - for ready in accumulate(_readiness_pids(launch), _ready_workers, initial=frozenset[int]()): + for ready in accumulate(_readiness_pids(launch, workers), _ready_workers, initial=frozenset[int]()): members: Final = frozenset( pid for member in group_members(launch.process.pid) if (pid := _ready_worker_pid(member)) is not None ) @@ -400,7 +627,9 @@ def owned_proxy_process( database_setup: tuple[str, ...] = DB_PUSH, extra_arguments: tuple[str, ...] = (), coverage_file: Path | None = None, + healthcheck: WorkerHealthcheck = WorkerHealthcheck.STRICT, ) -> Iterator[OwnedProxy]: + assert "--timeout_worker_healthcheck" not in extra_arguments, "Choose the worker healthcheck with `healthcheck`" root: Final = _proxy_root() environment: Final = _proxy_environment( gateway, {**overrides, **({"COVERAGE_FILE": str(coverage_file)} if coverage_file else {})}, remove_environment @@ -419,7 +648,7 @@ def owned_proxy_process( "--num_workers", str(workers), "--timeout_worker_healthcheck", - "5", + str(healthcheck_seconds(healthcheck)), *database_setup, *extra_arguments, ) @@ -429,7 +658,7 @@ def owned_proxy_process( with httpx.Client( base_url=f"http://127.0.0.1:{launch.port}", timeout=15, trust_env=False, limits=GATEWAY_LIMITS ) as client: - yield OwnedProxy(Gateway(client, gateway.key, gateway.upstream_url), process, launch.log) + yield OwnedProxy(Gateway(client, gateway.key, gateway.upstream_url), process, launch.log, healthcheck) finally: _stop_launch(launch) @@ -501,7 +730,7 @@ def refused_boot_log( launch: Final = _launch(command, root, environment, output) try: with httpx.Client(base_url=f"http://127.0.0.1:{launch.port}", timeout=15, trust_env=False) as client: - deadline: Final = time.monotonic() + 70 + deadline: Final = time.monotonic() + readiness_seconds(1) while launch.process.poll() is None: assert not _is_ready(client), ( f"Proxy became ready instead of refusing to boot:\n{launch.log.read_text()}" diff --git a/tests/integration/_support/redis_process.py b/tests/integration/_support/redis_process.py index e5258fbda02..cfdbf041eba 100644 --- a/tests/integration/_support/redis_process.py +++ b/tests/integration/_support/redis_process.py @@ -139,6 +139,14 @@ class OwnedRedis: assert self.pid_file.encode() in command, "Redis process ownership changed" subprocess.run(["docker", "exec", "redis-cache", "kill", f"-{int(action)}", str(pid)], check=True, timeout=2) + @contextmanager + def paused(self) -> Iterator[None]: + self.signal(signal.SIGSTOP) + try: + yield + finally: + self.signal(signal.SIGCONT) + @contextmanager def _started( diff --git a/tests/integration/_support/timing.py b/tests/integration/_support/timing.py new file mode 100644 index 00000000000..5b16c9240cf --- /dev/null +++ b/tests/integration/_support/timing.py @@ -0,0 +1,40 @@ +import time +from collections.abc import Callable, Iterable +from dataclasses import dataclass +from typing import Final + + +@dataclass(frozen=True, slots=True) +class Timed[T]: + value: T + seconds: float + + +def timed[T](operation: Callable[[], T]) -> Timed[T]: + started: Final = time.monotonic() + value: Final = operation() + return Timed(value, time.monotonic() - started) + + +def assert_not_delayed_by(observed: Iterable[float], *, control: Iterable[float], injected_delay: float) -> None: + worst: Final = max(observed) + baseline: Final = max(control) + assert worst < baseline + injected_delay, ( + f"slowest call took {worst:.2f}s, at least the {baseline:.2f}s control plus the {injected_delay}s injected delay" + ) + + +def assert_cut_short(elapsed: float, *, injected_delay: float) -> None: + assert elapsed < injected_delay, ( + f"returned after {elapsed:.2f}s, no sooner than the {injected_delay}s injected delay" + ) + + +def assert_enforced_within(elapsed: float, *, configured: float, slack: float, not_before: float = 0.0) -> None: + assert not_before <= elapsed <= configured + slack, ( + f"took {elapsed:.2f}s, expected between {not_before}s and the configured {configured}s plus {slack}s of slack" + ) + + +def assert_within_factor_of_control(observed: float, *, control: float, factor: float) -> None: + assert observed <= factor * control, f"took {observed:.2f}s, over {factor}x the {control:.2f}s control" diff --git a/tests/integration/authorization/test_alias_chain_budget_waiver.py b/tests/integration/authorization/test_alias_chain_budget_waiver.py index 175d4765a24..0d64415d445 100644 --- a/tests/integration/authorization/test_alias_chain_budget_waiver.py +++ b/tests/integration/authorization/test_alias_chain_budget_waiver.py @@ -2,7 +2,6 @@ import json import math import os import re -import signal import threading import uuid from collections.abc import Callable, Generator, Iterator, Mapping @@ -964,11 +963,11 @@ def test_killing_one_worker_mid_burst_keeps_the_chain_entry_served(rig: ChainRig held_by: Final = MappingProxyType({pid: _open_upstream_connections(pid, wire.url) for pid in workers}) assert sum(held_by.values()) == GATEWAY_BURST, held_by victim_pid, survivor_pid = sorted(workers, key=held_by.__getitem__) - victim: Final = psutil.Process(victim_pid) - victim.suspend() - victim.send_signal(signal.SIGKILL) - release.set() - responses: Final = {marker: future.result() for marker, future in pending.items()} + with owned.replacing_worker(victim_pid) as kill: + kill.freeze() + kill() + release.set() + responses: Final = {marker: future.result() for marker, future in pending.items()} served: Final = frozenset( marker for marker, response in responses.items() if response is not None and response.status_code == 200 ) diff --git a/tests/integration/authorization/test_cli_sso_login_ui_disabled.py b/tests/integration/authorization/test_cli_sso_login_ui_disabled.py index a5e9b2752ea..315a0b20334 100644 --- a/tests/integration/authorization/test_cli_sso_login_ui_disabled.py +++ b/tests/integration/authorization/test_cli_sso_login_ui_disabled.py @@ -2,10 +2,8 @@ from __future__ import annotations import base64 import json -import os import re import secrets -import signal import threading import uuid from collections.abc import Iterator, Mapping, Sequence @@ -523,7 +521,6 @@ def test_login_survives_a_worker_kill(idp: Idp, tmp_path: Path) -> None: link: Final = browser.get(_cli_link(proxy, session.login_id)) at_idp: Final = browser.get(_assert_idp_redirect(proxy, idp, link, session.login_id)) assert at_idp.status_code == 302, f"{at_idp.status_code} {at_idp.text}" - os.kill(victim, signal.SIGKILL) def callback_attempt() -> httpx.Response | None: try: @@ -531,9 +528,11 @@ def test_login_survives_a_worker_kill(idp: Idp, tmp_path: Path) -> None: except httpx.TransportError: return None - callback: Final = eventually( - callback_attempt, lambda response: response is not None and response.status_code == 200, seconds=60 - ) + with owned.replacing_worker(victim) as kill: + kill() + callback: Final = eventually( + callback_attempt, lambda response: response is not None and response.status_code == 200, seconds=60 + ) assert callback is not None done: Final = _complete_in_browser(browser, callback, session.user_code) assert done.status_code == 200 and CLI_SUCCESS_PAGE_TITLE in done.text, f"{done.status_code} {done.text}" diff --git a/tests/integration/authorization/test_hidden_alias_budget_bypass.py b/tests/integration/authorization/test_hidden_alias_budget_bypass.py index eabdb01602e..67ec30cec2a 100644 --- a/tests/integration/authorization/test_hidden_alias_budget_bypass.py +++ b/tests/integration/authorization/test_hidden_alias_budget_bypass.py @@ -1,7 +1,6 @@ import base64 import json import math -import signal import time import uuid from collections.abc import Callable, Iterator, Mapping @@ -1384,7 +1383,7 @@ def test_killing_one_worker_keeps_the_shadowing_alias_refused(rig: AliasRig, tmp settle_candidate(candidate, target, exhausted, BUDGET_EXCEEDED, seconds=120) settle_candidate(candidate, shadow, exhausted, BUDGET_EXCEEDED, seconds=120) victim: Final = before[0] - victim.send_signal(signal.SIGKILL) + owned.kill_worker(victim.pid) eventually(lambda: _workers(root), _replaced(victim), seconds=120) settle_candidate(candidate, shadow, exhausted, BUDGET_EXCEEDED, seconds=120) assert chat_statuses(candidate, shadow, exhausted, GATEWAY_BURST) == {BUDGET_EXCEEDED} diff --git a/tests/integration/authorization/test_hidden_alias_budget_bypass_chaos.py b/tests/integration/authorization/test_hidden_alias_budget_bypass_chaos.py index a7da04f1364..0a9b1cb3c4e 100644 --- a/tests/integration/authorization/test_hidden_alias_budget_bypass_chaos.py +++ b/tests/integration/authorization/test_hidden_alias_budget_bypass_chaos.py @@ -1,6 +1,4 @@ -import os import re -import signal import threading import uuid from collections.abc import Callable, Iterator, Mapping @@ -186,14 +184,15 @@ def test_killed_worker_leaves_the_sibling_serving_hidden_free_aliases(rig: Alias ) settle_candidate(candidate, rig.hidden_free, key, 200, seconds=120) settle_candidate(candidate, rig.hidden_paid, key, BUDGET_EXCEEDED, seconds=120) - os.kill(workers[0], signal.SIGKILL) - eventually( - lambda: _tolerant_status(candidate, rig.hidden_free, key, "kill-probe-" + uuid.uuid4().hex), - lambda found: found == 200, - seconds=60, - ) - assert chat_statuses(candidate, rig.hidden_free, key, 8) == {200} - assert chat_statuses(candidate, rig.hidden_paid, key, 8) == {BUDGET_EXCEEDED} + with owned.replacing_worker(workers[0]) as kill: + kill() + eventually( + lambda: _tolerant_status(candidate, rig.hidden_free, key, "kill-probe-" + uuid.uuid4().hex), + lambda found: found == 200, + seconds=60, + ) + assert chat_statuses(candidate, rig.hidden_free, key, 8) == {200} + assert chat_statuses(candidate, rig.hidden_paid, key, 8) == {BUDGET_EXCEEDED} eventually( lambda: _worker_startups(owned.log), lambda found: len(found[0]) == 3 and found[1] == 3, diff --git a/tests/integration/authorization/test_key_guardrail_opt_out_chaos.py b/tests/integration/authorization/test_key_guardrail_opt_out_chaos.py index f205b3804f6..78b226f619b 100644 --- a/tests/integration/authorization/test_key_guardrail_opt_out_chaos.py +++ b/tests/integration/authorization/test_key_guardrail_opt_out_chaos.py @@ -1,23 +1,20 @@ import json -import os -import signal import socket import threading import time import uuid from concurrent.futures import ThreadPoolExecutor -from hashlib import sha256 from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer from pathlib import Path from queue import SimpleQueue from typing import Final import httpx -import psutil from integration._support.client import Gateway, eventually, string_value from integration._support.database import read_rows from integration._support.process import owned_proxy, owned_proxy_process +from integration._support.timing import assert_cut_short, timed from integration.authorization._guardrail_opt_out import ( chat, guardrail_config, @@ -250,13 +247,8 @@ def test_flag_denial_survives_worker_kill(gateway: Gateway, tmp_path: Path) -> N caller: Final = non_admin_caller(scenario, member, team, model) alias: Final = "audit-kill-" + uuid.uuid4().hex - workers: Final = eventually( - lambda: psutil.Process(owned.process.pid).children(recursive=True), - lambda children: len(children) >= 2, - seconds=30, - ) - victim: Final = workers[0] - os.kill(victim.pid, signal.SIGKILL) + workers: Final = eventually(owned.worker_pids, lambda pids: len(pids) >= 2, seconds=30) + owned.kill_worker(min(workers)) probe: Final = eventually( lambda: candidate.request( @@ -303,17 +295,18 @@ def test_exempt_chats_do_not_wait_on_slow_guardrail_sink(gateway: Gateway, tmp_p model: Final = scenario.model() exempt: Final = scenario.key(models=[model], disable_global_guardrails=True) - started: Final = time.monotonic() with ThreadPoolExecutor(max_workers=10) as pool: - responses: Final = tuple( + calls: Final = tuple( pool.map( - lambda index: chat(candidate, model, exempt, f"slow-sink-{index}-{uuid.uuid4().hex}"), + lambda index: timed( + lambda: chat(candidate, model, exempt, f"slow-sink-{index}-{uuid.uuid4().hex}") + ), range(10), ) ) - elapsed: Final = time.monotonic() - started + responses: Final = tuple(call.value for call in calls) assert all(response.status_code == 200 for response in responses), [ (response.status_code, response.text) for response in responses ] - assert elapsed < sink_delay, f"exempt chats waited on the guardrail sink: {elapsed}s" + assert_cut_short(max(call.seconds for call in calls), injected_delay=sink_delay) assert sink.drain() == () diff --git a/tests/integration/authorization/test_websocket_rejection_logging.py b/tests/integration/authorization/test_websocket_rejection_logging.py index ea0a841a2b5..c12b387bfd0 100644 --- a/tests/integration/authorization/test_websocket_rejection_logging.py +++ b/tests/integration/authorization/test_websocket_rejection_logging.py @@ -304,10 +304,11 @@ async def test_a_rejection_burst_survives_a_killed_worker_without_an_asgi_crash( before: Final = await _burst(owned, shapes, "burst") assert _crash_lines(before) == (), before victim: Final = eventually(lambda: _workers(owned), lambda workers: len(workers) == 2)[0] - victim.kill() - during: Final = await _refused(owned, "/v1/responses", shapes["unknown_key"], f"model={_tag('during')}") - assert during.status == 403, during.window - assert _crash_lines(during.window) == (), during.window + with owned.replacing_worker(victim.pid) as kill: + kill() + during: Final = await _refused(owned, "/v1/responses", shapes["unknown_key"], f"model={_tag('during')}") + assert during.status == 403, during.window + assert _crash_lines(during.window) == (), during.window eventually( lambda: frozenset(worker.pid for worker in _workers(owned)), lambda pids: len(pids) == 2 and victim.pid not in pids, diff --git a/tests/integration/caching/test_response_cache_chaos.py b/tests/integration/caching/test_response_cache_chaos.py index 2e663c97d1f..17637c088a0 100644 --- a/tests/integration/caching/test_response_cache_chaos.py +++ b/tests/integration/caching/test_response_cache_chaos.py @@ -438,14 +438,16 @@ def test_worker_killed_before_burst_leaves_the_survivor_serving_and_never_pins_a before: Final = _workers(root) assert len(before) == WORKERS, [process.pid for process in before] victim: Final = before[0] - victim.kill() - victim.wait(timeout=10) - with httpx.Client(base_url=target, timeout=15, trust_env=False) as fresh: - readiness: Final = fresh.get("/health/readiness") - assert readiness.status_code == 200, f"/health/readiness with worker {victim.pid} dead: {readiness.text}" - burst: Final = Burst(target, key) - burst.start(pool, _burst_calls(fleet, PER_ENDPOINT // 2, STREAMS_PER_ENDPOINT // 2)) - outcomes: Final = burst.outcomes() + with owned.replacing_worker(victim.pid) as kill: + kill() + with httpx.Client(base_url=target, timeout=15, trust_env=False) as fresh: + readiness: Final = fresh.get("/health/readiness") + assert readiness.status_code == 200, ( + f"/health/readiness with worker {victim.pid} dead: {readiness.text}" + ) + burst: Final = Burst(target, key) + burst.start(pool, _burst_calls(fleet, PER_ENDPOINT // 2, STREAMS_PER_ENDPOINT // 2)) + outcomes: Final = burst.outcomes() respawned: Final = eventually( lambda: tuple(process.pid for process in _workers(root)), lambda pids: len(pids) == WORKERS and victim.pid not in pids, diff --git a/tests/integration/compatibility/test_granian_admission_saturation.py b/tests/integration/compatibility/test_granian_admission_saturation.py index 9f7cec44d89..02c349b2c03 100644 --- a/tests/integration/compatibility/test_granian_admission_saturation.py +++ b/tests/integration/compatibility/test_granian_admission_saturation.py @@ -11,6 +11,7 @@ import httpx import pytest from tests._master_key import MASTER_KEY +from tests.integration._support.process import stop_child MAX_IN_FLIGHT: Final = 8 MAX_QUEUED: Final = 8 @@ -102,16 +103,6 @@ general_settings: assert len(rejected) >= REQUESTS - MAX_IN_FLIGHT - MAX_QUEUED, statuses assert all(result[1] for result in rejected), rejected finally: - proxy_process.terminate() - try: - proxy_process.wait(timeout=10) - except subprocess.TimeoutExpired: - proxy_process.kill() - proxy_process.wait() + stop_child(proxy_process) finally: - fake_process.terminate() - try: - fake_process.wait(timeout=10) - except subprocess.TimeoutExpired: - fake_process.kill() - fake_process.wait() + stop_child(fake_process) diff --git a/tests/integration/compatibility/test_integration_worker_cleanup.py b/tests/integration/compatibility/test_integration_worker_cleanup.py index a95927dee5e..c5f19dfa50f 100644 --- a/tests/integration/compatibility/test_integration_worker_cleanup.py +++ b/tests/integration/compatibility/test_integration_worker_cleanup.py @@ -13,6 +13,7 @@ from psycopg import sql from pydantic import BaseModel, ConfigDict from tests.integration._support.database import read_rows +from tests.integration._support.process import stop_processes FAULT_CASE: Final = """import json, os, pathlib, psutil from urllib.parse import urlsplit @@ -44,13 +45,7 @@ def _survivor(identity: tuple[int, float]) -> psutil.Process | None: def _clean_fault_resources(resources: CrashResources) -> None: - processes: Final = tuple(process for identity in resources.processes if (process := _survivor(identity))) - for process in processes: - process.terminate() - _stopped, alive = psutil.wait_procs(processes, timeout=5) - for process in alive: - process.kill() - psutil.wait_procs(alive, timeout=3) + stop_processes(process for identity in resources.processes if (process := _survivor(identity))) with psycopg.connect(os.environ["DATABASE_URL"], autocommit=True) as admin: admin.execute(sql.SQL("DROP DATABASE IF EXISTS {} WITH (FORCE)").format(sql.Identifier(resources.database))) diff --git a/tests/integration/compatibility/test_missing_body_param_status.py b/tests/integration/compatibility/test_missing_body_param_status.py index facf8007910..fbdf7027d4b 100644 --- a/tests/integration/compatibility/test_missing_body_param_status.py +++ b/tests/integration/compatibility/test_missing_body_param_status.py @@ -3,7 +3,6 @@ from __future__ import annotations import asyncio import json import os -import signal import socket import subprocess import sys @@ -20,7 +19,7 @@ import psutil import pytest from integration._support.client import JSON_OBJECT, Gateway, Scenario, eventually, object_value, string_value from integration._support.database import read_rows -from integration._support.process import owned_proxy_process +from integration._support.process import owned_proxy_process, paused, stop_child from integration._support.upstream import ScenarioHandle, delete_scenario, register_scenario from openai import AsyncOpenAI, BadRequestError, OpenAI from pydantic import JsonValue @@ -1455,10 +1454,7 @@ def _upstream(directory: Path) -> Iterator[tuple[subprocess.Popen[bytes], str]]: eventually(lambda: _healthy(url), lambda status: status == 200, seconds=30) yield process, url finally: - if process.poll() is None: - process.send_signal(signal.SIGCONT) - process.terminate() - process.wait(timeout=10) + stop_child(process) def _register_owned(url: str, identity: str, scripted: JsonResponse) -> ScenarioHandle: @@ -1478,10 +1474,6 @@ def _delete_owned(handle: ScenarioHandle) -> None: ).raise_for_status() -def _workers(process: subprocess.Popen[bytes]) -> tuple[psutil.Process, ...]: - return tuple(psutil.Process(process.pid).children(recursive=True)) - - def _process_tree_line(process: psutil.Process) -> str: try: return f"{process.pid} {' '.join(process.cmdline())}" @@ -1489,13 +1481,6 @@ def _process_tree_line(process: psutil.Process) -> str: return f"{process.pid} " -def _worker_alive(worker: psutil.Process) -> bool: - try: - return worker.is_running() and worker.status() != psutil.STATUS_ZOMBIE - except psutil.NoSuchProcess: - return False - - def _chat_body(model: str, marker: str, valid: bool) -> dict[str, JsonValue]: return {"model": model, "user": marker, **({"messages": [{"role": "user", "content": marker}]} if valid else {})} @@ -1556,11 +1541,7 @@ def test_upstream_pause_and_worker_kill_preserve_required_body_status( processes: Final = (process_root, *process_root.children(recursive=True)) process_tree: Final = "\n".join(_process_tree_line(process) for process in processes) record_property("owned_proxy_process_tree", process_tree) - workers: Final = eventually( - lambda: _workers(owned.process), - lambda children: len(children) >= 2, - seconds=30, - ) + workers: Final = eventually(owned.worker_pids, lambda pids: len(pids) >= 2, seconds=30) observations: Final = _Observations(url) missing_routes: Final = ( ("aspeech", "input"), @@ -1585,26 +1566,22 @@ def test_upstream_pause_and_worker_kill_preserve_required_body_status( valid_calls: Final = tuple( ("/v1/chat/completions", _chat_body(chat_model, marker, True), marker) for marker in markers ) - upstream.send_signal(signal.SIGSTOP) - try: - with ThreadPoolExecutor(max_workers=10) as pool: - missing_futures: Final = tuple( - pool.submit(_post, candidate, path, body) for path, body, _expected, _identity in missing_calls - ) - missing: Final = tuple( - (call, future.result(timeout=15)) for call, future in zip(missing_calls, missing_futures) - ) - assert all( - response.status_code == 400 and response.json() == expected - for (_path, _body, expected, _identity), response in missing - ), [response.text for _call, response in missing] - paused_statuses: Final = tuple(response.status_code for _call, response in missing) - record_property( - "chaos_paused_missing_status_counts", - str({status: paused_statuses.count(status) for status in sorted(set(paused_statuses))}), - ) - finally: - upstream.send_signal(signal.SIGCONT) + with paused(upstream), ThreadPoolExecutor(max_workers=10) as pool: + missing_futures: Final = tuple( + pool.submit(_post, candidate, path, body) for path, body, _expected, _identity in missing_calls + ) + missing: Final = tuple( + (call, future.result(timeout=15)) for call, future in zip(missing_calls, missing_futures) + ) + assert all( + response.status_code == 400 and response.json() == expected + for (_path, _body, expected, _identity), response in missing + ), [response.text for _call, response in missing] + paused_statuses: Final = tuple(response.status_code for _call, response in missing) + record_property( + "chaos_paused_missing_status_counts", + str({status: paused_statuses.count(status) for status in sorted(set(paused_statuses))}), + ) with ThreadPoolExecutor(max_workers=20) as pool: valid_futures: Final = tuple( pool.submit(_post, candidate, path, body) for path, body, _marker in valid_calls @@ -1650,14 +1627,14 @@ def test_upstream_pause_and_worker_kill_preserve_required_body_status( record_property("chaos_spend_query", 'SELECT request_id FROM "LiteLLM_SpendLogs" WHERE request_id = %s') record_property("chaos_spend_response_id_count", len(request_ids)) record_property("chaos_spend_row_counts", str(tuple(len(rows) for rows in spend_rows))) - workers[0].kill() - eventually(lambda: _worker_alive(workers[0]), lambda alive: not alive, seconds=10) - post_kill: Final = tuple( - _post(candidate, path, body) for path, body, _expected, _identity in missing_calls[:5] - ) - assert all( - response.status_code == 400 and response.json() == expected - for response, (_path, _body, expected, _identity) in zip(post_kill, missing_calls[:5]) - ), [response.text for response in post_kill] - recovered: Final = _post(candidate, "/v1/chat/completions", _chat_body(chat_model, "recovered", True)) - assert recovered.status_code == 200, recovered.text + with owned.replacing_worker(min(workers)) as kill: + kill() + post_kill: Final = tuple( + _post(candidate, path, body) for path, body, _expected, _identity in missing_calls[:5] + ) + assert all( + response.status_code == 400 and response.json() == expected + for response, (_path, _body, expected, _identity) in zip(post_kill, missing_calls[:5]) + ), [response.text for response in post_kill] + recovered: Final = _post(candidate, "/v1/chat/completions", _chat_body(chat_model, "recovered", True)) + assert recovered.status_code == 200, recovered.text diff --git a/tests/integration/compatibility/test_openai_sdk_client_paths_chaos.py b/tests/integration/compatibility/test_openai_sdk_client_paths_chaos.py index 0bcc9fc4340..e9ee4a323cd 100644 --- a/tests/integration/compatibility/test_openai_sdk_client_paths_chaos.py +++ b/tests/integration/compatibility/test_openai_sdk_client_paths_chaos.py @@ -1,7 +1,6 @@ from __future__ import annotations import json -import signal import subprocess import uuid from collections.abc import Mapping @@ -19,7 +18,7 @@ from anthropic import Anthropic from anthropic import APIError as AnthropicApiError from integration._support.client import JSON_OBJECT, Gateway, Scenario, eventually, object_value from integration._support.database import read_rows -from integration._support.process import graceful_stop_seconds, owned_proxy_process, owned_upstream +from integration._support.process import graceful_stop_seconds, owned_proxy_process, owned_upstream, paused from integration._support.upstream import ScenarioHandle, delete_scenario, register_scenario from openai import APIError as OpenAiApiError from openai import OpenAI @@ -261,25 +260,6 @@ def _landed_once(response_ids: tuple[str, ...]) -> tuple[tuple[Mapping[str, Json return rows -def _is_worker(child: psutil.Process) -> bool: - try: - return "spawn_main" in " ".join(child.cmdline()) - except psutil.Error: - return False - - -def _worker_alive(worker: psutil.Process) -> bool: - try: - return worker.is_running() and worker.status() != psutil.STATUS_ZOMBIE - except psutil.NoSuchProcess: - return False - - -def _alive_workers(process: subprocess.Popen[bytes]) -> tuple[psutil.Process, ...]: - children: Final = tuple(child for child in psutil.Process(process.pid).children() if _is_worker(child)) - return tuple(child for child in children if _worker_alive(child)) - - def _process_tree(process: subprocess.Popen[bytes]) -> str: root: Final = psutil.Process(process.pid) return "\n".join(_process_line(member) for member in (root, *root.children(recursive=True))) @@ -313,12 +293,9 @@ def test_c01_upstream_pause_mid_burst_keeps_liveliness_and_lands_every_id_once( with ThreadPoolExecutor(max_workers=len(calls)) as pool: futures: Final = _submit(pool, clients, calls) eventually(partial(_done_count, futures), lambda done: done >= 1, seconds=60) - upstream.send_signal(signal.SIGSTOP) - try: + with paused(upstream): done_at_pause: Final = _done_count(futures) paused_liveliness: Final = tuple(_liveliness(gateway) for _ in range(3)) - finally: - upstream.send_signal(signal.SIGCONT) done_at_resume: Final = _done_count(futures) outcomes: Final = tuple(future.result(timeout=180) for future in futures) record_property("c01_burst_size", len(calls)) @@ -358,27 +335,25 @@ def test_c02_worker_sigkill_mid_burst_keeps_serving_and_respawns( partial(_every_worker_serves_every_kind, owned.gateway, deployments), lambda served: served, seconds=90 ) clients: Final = _clients(owned.gateway, deployments) - workers: Final = eventually( - partial(_alive_workers, owned.process), lambda found: len(found) == 2, seconds=60 - ) + workers: Final = eventually(owned.worker_pids, lambda found: len(found) == 2, seconds=60) record_property("c02_process_tree_before_kill", _process_tree(owned.process)) - victim: Final = workers[0] + victim: Final = min(workers) calls: Final = _burst_calls(_CHAT_KINDS, 15) observations: Final = _Observations(slot.url) with ThreadPoolExecutor(max_workers=len(calls)) as pool: futures: Final = _submit(pool, clients, calls) eventually(partial(_done_count, futures), lambda done: done >= 1, seconds=60) - victim.kill() - eventually(partial(_worker_alive, victim), lambda alive: not alive, seconds=10) - done_at_kill: Final = _done_count(futures) - survivor: Final = eventually( - lambda: _attempt(clients, _Call("chat", f"after-kill-{uuid.uuid4().hex}")), - lambda outcome: isinstance(outcome, _Completed), - seconds=30, - ) + with owned.replacing_worker(victim) as kill: + kill() + done_at_kill: Final = _done_count(futures) + survivor: Final = eventually( + lambda: _attempt(clients, _Call("chat", f"after-kill-{uuid.uuid4().hex}")), + lambda outcome: isinstance(outcome, _Completed), + seconds=30, + ) outcomes: Final = tuple(future.result(timeout=180) for future in futures) record_property("c02_burst_size", len(calls)) - record_property("c02_killed_worker_pid", victim.pid) + record_property("c02_killed_worker_pid", victim) record_property("c02_done_at_kill", done_at_kill) assert isinstance(survivor, _Completed), survivor completed: Final = tuple(outcome for outcome in outcomes if isinstance(outcome, _Completed)) @@ -388,12 +363,12 @@ def test_c02_worker_sigkill_mid_burst_keeps_serving_and_respawns( record_property("c02_inflight_failure_errors", str(tuple(outcome.error for outcome in failed))) assert completed, outcomes respawned: Final = eventually( - partial(_alive_workers, owned.process), - lambda found: len(found) == 2 and victim.pid not in {worker.pid for worker in found}, + owned.worker_pids, + lambda found: len(found) == 2 and victim not in found, seconds=graceful_stop_seconds() + 60, ) record_property("c02_process_tree_after_respawn", _process_tree(owned.process)) - record_property("c02_respawned_worker_pids", str(tuple(worker.pid for worker in respawned))) + record_property("c02_respawned_worker_pids", str(tuple(sorted(respawned)))) after_respawn: Final = eventually( lambda: _attempt(clients, _Call("chat", f"after-respawn-{uuid.uuid4().hex}")), lambda outcome: isinstance(outcome, _Completed), diff --git a/tests/integration/configuration/test_codex_catalog_boot.py b/tests/integration/configuration/test_codex_catalog_boot.py index 717b9471244..4868e43b855 100644 --- a/tests/integration/configuration/test_codex_catalog_boot.py +++ b/tests/integration/configuration/test_codex_catalog_boot.py @@ -1,6 +1,5 @@ import json import os -import signal import stat import subprocess import sys @@ -11,12 +10,11 @@ from pathlib import Path from typing import Final import httpx -import psutil import pytest from pydantic import JsonValue from tests.integration._support.client import JSON_OBJECT, Gateway, Scenario, eventually, object_value, string_value -from tests.integration._support.process import group_members, owned_proxy, owned_proxy_process +from tests.integration._support.process import owned_proxy, owned_proxy_process CLIENT_VERSION: Final = "0.159.3" FRESH_CONNECTION: Final = {"Connection": "close"} @@ -224,16 +222,6 @@ def test_team_owned_deployments_of_one_name_serve_each_team_its_own_upstream(gat assert outsider[mixed]["service_tiers"] == [_generic_tier("flex")], outsider[mixed] -def _worker_pids(root_pid: int) -> frozenset[int]: - def is_worker(process: psutil.Process) -> bool: - try: - return "spawn_main" in " ".join(process.cmdline()) - except (psutil.NoSuchProcess, psutil.AccessDenied, psutil.ZombieProcess): - return False - - return frozenset(process.pid for process in group_members(root_pid) if is_worker(process)) - - def _catalog_outcome(proxy: Gateway, slugs: Sequence[JsonValue]) -> str: try: response: Final = _catalog_response(proxy) @@ -252,22 +240,24 @@ def test_catalog_survives_a_worker_kill(gateway: Gateway, tmp_path: Path) -> Non config: Final = _write_config(tmp_path, gateway.upstream_url, tiered, plain, alias) with owned_proxy_process(gateway, tmp_path, {"STORE_MODEL_IN_DB": "False"}, config=config, workers=2) as owned: candidate: Final = owned.gateway - workers: Final = eventually(lambda: _worker_pids(owned.process.pid), lambda pids: len(pids) == 2, seconds=30) + workers: Final = eventually(owned.worker_pids, lambda pids: len(pids) == 2, seconds=30) victim: Final = min(workers) slugs: Final = sorted(entry["slug"] for entry in _entries(_catalog_response(candidate))) assert slugs == sorted((tiered, plain, alias)), slugs - def attempt(index: int) -> str: - if index == 2: - os.kill(victim, signal.SIGKILL) - return _catalog_outcome(candidate, slugs) + with owned.replacing_worker(victim) as kill: - with ThreadPoolExecutor(max_workers=BURST) as pool: - outcomes: Final = tuple(pool.map(attempt, range(BURST))) + def attempt(index: int) -> str: + if index == 2: + kill() + return _catalog_outcome(candidate, slugs) + + with ThreadPoolExecutor(max_workers=BURST) as pool: + outcomes: Final = tuple(pool.map(attempt, range(BURST))) assert outcomes.count("ok") >= 1, outcomes assert all(outcome == "ok" or outcome.startswith("transport:") for outcome in outcomes), outcomes respawned: Final = eventually( - lambda: _worker_pids(owned.process.pid), + owned.worker_pids, lambda pids: len(pids) == 2 and victim not in pids, seconds=60, ) diff --git a/tests/integration/configuration/test_lazy_routes_flag.py b/tests/integration/configuration/test_lazy_routes_flag.py index 28930ec29b4..7809b357a7e 100644 --- a/tests/integration/configuration/test_lazy_routes_flag.py +++ b/tests/integration/configuration/test_lazy_routes_flag.py @@ -23,7 +23,6 @@ from typing import Final import anthropic import httpx import openai -import psutil import pytest import yaml from pydantic import JsonValue, TypeAdapter @@ -273,17 +272,6 @@ def _status(candidate: Gateway, path: str) -> int: return candidate.request("GET", path).status_code -def _workers(owned: OwnedProxy) -> tuple[psutil.Process, ...]: - return tuple(child for child in psutil.Process(owned.process.pid).children() if _is_worker(child)) - - -def _is_worker(child: psutil.Process) -> bool: - try: - return "spawn_main" in " ".join(child.cmdline()) and child.status() != psutil.STATUS_ZOMBIE - except psutil.Error: - return False - - @pytest.mark.parametrize("spelling", ("1", "Yes", "ON")) def test_every_truthy_spelling_of_the_flag_registers_the_ticket_features_at_startup( gateway: Gateway, tmp_path: Path, spelling: str @@ -457,18 +445,19 @@ def test_disable_lazy_routes_flag_route_table_survives_a_boot_burst_and_a_killed assert statuses == (200,) * len(probes), statuses assert _routed_features(owned.gateway) == at_boot, "the boot burst changed the route table" - victim: Final = eventually(lambda: _workers(owned), lambda workers: len(workers) == 2)[0] - victim.kill() - with httpx.Client(base_url=owned.gateway.client.base_url, timeout=15, trust_env=False) as fresh: - survivor: Final = Gateway(fresh, owned.gateway.key, owned.gateway.upstream_url) - during: Final = tuple(survivor.request("GET", "/v1/mcp/server").status_code for _ in range(10)) + victim: Final = min(eventually(owned.worker_pids, lambda workers: len(workers) == 2)) + with owned.replacing_worker(victim) as kill: + kill() + with httpx.Client(base_url=owned.gateway.client.base_url, timeout=15, trust_env=False) as fresh: + survivor: Final = Gateway(fresh, owned.gateway.key, owned.gateway.upstream_url) + during: Final = tuple(survivor.request("GET", "/v1/mcp/server").status_code for _ in range(10)) assert during == (200,) * 10, during respawned: Final = eventually( - lambda: frozenset(worker.pid for worker in _workers(owned)), - lambda pids: len(pids) == 2 and victim.pid not in pids, + owned.worker_pids, + lambda pids: len(pids) == 2 and victim not in pids, seconds=30, ) - assert f"Child process [{victim.pid}] died" in owned.log.read_text(), respawned + assert f"Child process [{victim}] died" in owned.log.read_text(), respawned tables: Final = tuple(_routed_features(owned.gateway) for _ in range(4)) assert all(table == at_boot for table in tables), "the respawned worker disagrees on the route table" diff --git a/tests/integration/configuration/test_multi_worker_slow_import_boot.py b/tests/integration/configuration/test_multi_worker_slow_import_boot.py index fb544f37a66..966a4a79b96 100644 --- a/tests/integration/configuration/test_multi_worker_slow_import_boot.py +++ b/tests/integration/configuration/test_multi_worker_slow_import_boot.py @@ -5,7 +5,7 @@ from typing import Final import pytest from integration._support.client import Gateway, eventually -from integration._support.process import graceful_stop_seconds, owned_proxy_process +from integration._support.process import WorkerHealthcheck, graceful_stop_seconds, owned_proxy_process WORKERS: Final = 2 DEPLOYMENT_HEALTHCHECK_SECONDS: Final = 5 @@ -38,7 +38,7 @@ def test_owned_proxy_workers_outlast_the_deployments_healthcheck_default(gateway tmp_path, _slow_worker_environment(tmp_path), workers=WORKERS, - extra_arguments=("--timeout_worker_healthcheck", str(int(graceful_stop_seconds()))), + healthcheck=WorkerHealthcheck.LENIENT, ) as owned: started: Final = eventually( lambda: STARTED_WORKER.findall(owned.log.read_text()), diff --git a/tests/integration/database/migrations/test_autorouter_baseline_state.py b/tests/integration/database/migrations/test_autorouter_baseline_state.py index 629089eda02..729c34e1bb6 100644 --- a/tests/integration/database/migrations/test_autorouter_baseline_state.py +++ b/tests/integration/database/migrations/test_autorouter_baseline_state.py @@ -18,7 +18,9 @@ from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache from litellm.proxy.db.baseline_accounting import BaselineAccountingStore from litellm.proxy.db.routing_prisma_wrapper import RoutingPrismaWrapper from litellm.proxy.utils import PrismaClient, ProxyLogging +from tests.integration._support.timing import assert_enforced_within +_PROJECTION_LOCK_TIMEOUT_SECONDS: Final = 1 _MIGRATION: Final = Path(__file__).parents[4] / ( "litellm-proxy-extras/litellm_proxy_extras/migrations/20260915010000_add_autorouter_baseline_state/migration.sql" ) @@ -95,7 +97,7 @@ async def test_locked_projection_is_bounded_and_cancellation_propagates( connection.execute('LOCK TABLE "LiteLLM_AutoRouterBaselineComparison" IN ACCESS EXCLUSIVE MODE') started: Final = time.monotonic() assert await store.project("scope") == "unavailable" - assert time.monotonic() - started < 2 + assert_enforced_within(time.monotonic() - started, configured=_PROJECTION_LOCK_TIMEOUT_SECONDS, slack=1) pending: Final = asyncio.create_task(store.project("scope")) await asyncio.sleep(0.01) pending.cancel() diff --git a/tests/integration/database/migrations/test_prisma_toolchain.py b/tests/integration/database/migrations/test_prisma_toolchain.py index fecb1fece91..bf111e7d6b9 100644 --- a/tests/integration/database/migrations/test_prisma_toolchain.py +++ b/tests/integration/database/migrations/test_prisma_toolchain.py @@ -17,7 +17,6 @@ the proxy gave up after four identical timeouts. import json import logging import os -import signal import sys import time from pathlib import Path @@ -33,8 +32,13 @@ from litellm_proxy_extras.prisma_toolchain import ( from litellm_proxy_extras.utils import ProxyExtrasDBManager from tests._process_helpers import process_is_gone +from tests.integration._support.process import kill_stray_process +from tests.integration._support.timing import assert_cut_short, assert_enforced_within +HUNG_BOOTSTRAP_SECONDS = 30 +HUNG_DEPLOY_SECONDS = 60 +COMMAND_TIMEOUT_SECONDS = 1 FAKE_PRISMA = """#!{python} import json import os @@ -159,8 +163,8 @@ def test_bootstrap_is_not_bounded_by_the_per_command_timeout( def test_bootstrap_stops_at_its_own_timeout( toolchain_env: tuple[Path, Path], monkeypatch: pytest.MonkeyPatch ) -> None: - monkeypatch.setenv(PRISMA_BOOTSTRAP_TIMEOUT_ENV_VAR, "1") - monkeypatch.setenv("FAKE_PRISMA_SLEEP", "30") + monkeypatch.setenv(PRISMA_BOOTSTRAP_TIMEOUT_ENV_VAR, str(COMMAND_TIMEOUT_SECONDS)) + monkeypatch.setenv("FAKE_PRISMA_SLEEP", str(HUNG_BOOTSTRAP_SECONDS)) started = time.monotonic() result = ensure_prisma_toolchain( @@ -169,7 +173,8 @@ def test_bootstrap_stops_at_its_own_timeout( elapsed = time.monotonic() - started assert result.ready is False - assert elapsed < 15 + assert_cut_short(elapsed, injected_delay=HUNG_BOOTSTRAP_SECONDS) + assert_enforced_within(elapsed, configured=COMMAND_TIMEOUT_SECONDS, slack=14) def test_setup_database_prepares_the_toolchain_before_migrating( @@ -210,8 +215,8 @@ def test_migrate_deploy_stops_at_its_own_timeout( """The deploy budget still bounds a deploy that hangs, so boot cannot wait forever.""" _, log_path = toolchain_env monkeypatch.setenv("DATABASE_URL", "postgresql://u:p@localhost:9/x") - monkeypatch.setenv(PRISMA_MIGRATE_DEPLOY_TIMEOUT_ENV_VAR, "1") - monkeypatch.setenv("FAKE_PRISMA_FIRST_DEPLOY_SLEEP", "60") + monkeypatch.setenv(PRISMA_MIGRATE_DEPLOY_TIMEOUT_ENV_VAR, str(COMMAND_TIMEOUT_SECONDS)) + monkeypatch.setenv("FAKE_PRISMA_FIRST_DEPLOY_SLEEP", str(HUNG_DEPLOY_SECONDS)) monkeypatch.setenv("FAKE_PRISMA_LATER_DEPLOY_STDERR", "Error: P3018 permission denied for schema public") started = time.monotonic() @@ -220,7 +225,8 @@ def test_migrate_deploy_stops_at_its_own_timeout( elapsed = time.monotonic() - started assert len(_deploy_calls(log_path)) == 2 - assert elapsed < 30 + assert_cut_short(elapsed, injected_delay=HUNG_DEPLOY_SECONDS) + assert_enforced_within(elapsed, configured=2 * COMMAND_TIMEOUT_SECONDS, slack=28) def test_a_timed_out_migrate_deploy_takes_its_process_tree_with_it( @@ -230,7 +236,7 @@ def test_a_timed_out_migrate_deploy_takes_its_process_tree_with_it( _, log_path = toolchain_env pidfile = tmp_path / "grandchild.pid" monkeypatch.setenv("DATABASE_URL", "postgresql://u:p@localhost:9/x") - monkeypatch.setenv(PRISMA_MIGRATE_DEPLOY_TIMEOUT_ENV_VAR, "1") + monkeypatch.setenv(PRISMA_MIGRATE_DEPLOY_TIMEOUT_ENV_VAR, str(COMMAND_TIMEOUT_SECONDS)) monkeypatch.setenv("FAKE_PRISMA_FIRST_DEPLOY_SLEEP", "60") monkeypatch.setenv("FAKE_PRISMA_LATER_DEPLOY_STDERR", "Error: P3018 permission denied for schema public") monkeypatch.setenv("FAKE_PRISMA_GRANDCHILD_PIDFILE", str(pidfile)) @@ -243,10 +249,7 @@ def test_a_timed_out_migrate_deploy_takes_its_process_tree_with_it( assert len(_deploy_calls(log_path)) == 2 assert process_is_gone(grandchild_pid, within_seconds=5) finally: - try: - os.kill(grandchild_pid, signal.SIGKILL) - except ProcessLookupError: - pass + kill_stray_process(grandchild_pid) def test_db_push_timeout_hint_names_the_per_command_budget( diff --git a/tests/integration/database/test_harness_lifecycle.py b/tests/integration/database/test_harness_lifecycle.py index 6993342f2e5..0d604e03f03 100644 --- a/tests/integration/database/test_harness_lifecycle.py +++ b/tests/integration/database/test_harness_lifecycle.py @@ -31,6 +31,7 @@ from typing_extensions import ReadOnly, TypedDict from tests.integration._support import process as process_support from tests.integration._support.async_server import LoopbackServer +from tests.integration._support.budget import case_timeout_seconds from tests.integration._support.client import Gateway, eventually from tests.integration._support.database_relay import ( DatabaseRelay, @@ -102,8 +103,7 @@ def test_owned_readiness_rejects_a_healthy_listener_from_another_process( process_support._wait_until_ready(launch) assert process.poll() is None finally: - process.kill() - process.wait(timeout=3) + process_support.stop_child(process, grace_seconds=3) @pytest.mark.parametrize("factory", (database_relay, dropped_connection_relay, held_statement_relay)) @@ -177,7 +177,7 @@ async def test_owned_runtime_starts_the_entire_held_provider_burst_on_a_small_ru release.set() -@pytest.mark.timeout(90) +@pytest.mark.timeout(case_timeout_seconds()) def test_slow_worker_startup_survives_healthchecks_and_a_stalled_worker_is_replaced( gateway: Gateway, tmp_path: Path ) -> None: @@ -205,23 +205,13 @@ def test_slow_worker_startup_survives_healthchecks_and_a_stalled_worker_is_repla "Proxy yielded before both workers finished startup" ) assert owned.log.read_text().count("Application startup complete.") == 2 - victim: Final = psutil.Process(workers[0]) - victim.suspend() - try: - replacement: Final = eventually( - lambda: tuple(int(pid) for pid in STARTED_WORKER.findall(owned.log.read_text())), - lambda pids: len(pids) == 3 and owned.log.read_text().count("Application startup complete.") == 3, - seconds=30, - ) - assert replacement[-1] not in workers - assert not psutil.pid_exists(victim.pid), "Supervisor did not reap the stalled worker" - response: Final = owned.gateway.request("GET", "/health/readiness") - assert response.status_code == 200, response.text - finally: - try: - victim.kill() - except psutil.NoSuchProcess: - pass + replacement: Final = owned.stall_worker(workers[0]) + assert replacement not in workers + assert tuple(int(pid) for pid in STARTED_WORKER.findall(owned.log.read_text()))[-1] == replacement + assert owned.log.read_text().count("Application startup complete.") == 3 + assert not psutil.pid_exists(workers[0]), "Supervisor did not reap the stalled worker" + response: Final = owned.gateway.request("GET", "/health/readiness") + assert response.status_code == 200, response.text def test_interrupted_cleanup_reaps_the_owned_root_and_its_child() -> None: @@ -261,8 +251,7 @@ def test_interrupted_cleanup_reaps_the_owned_root_and_its_child() -> None: assert process.poll() is not None eventually(lambda: process_support.group_members(process.pid), lambda members: not members, seconds=5) finally: - process_support.signal_group(process.pid, signal.SIGKILL) - process.wait(timeout=3) + process_support.stop_process_group(process, 3) @pytest.mark.parametrize("explicit_dotenv", (False, True)) diff --git a/tests/integration/database/test_v1_migration_error_logging.py b/tests/integration/database/test_v1_migration_error_logging.py index c3a3d8c43d6..43879d371b3 100644 --- a/tests/integration/database/test_v1_migration_error_logging.py +++ b/tests/integration/database/test_v1_migration_error_logging.py @@ -3,7 +3,6 @@ from __future__ import annotations import json import os import re -import signal import socket import socketserver import subprocess @@ -24,6 +23,7 @@ from psycopg.types.json import Jsonb from tests.integration._support.client import eventually from tests.integration._support.process import _free_port as free_port +from tests.integration._support.process import stop_process_group REPO_ROOT: Final = Path(__file__).resolve().parents[3] MIGRATIONS_DIR: Final = REPO_ROOT / "litellm-proxy-extras" / "litellm_proxy_extras" / "migrations" @@ -299,18 +299,6 @@ def _gated_postgres_forwarder(port: int, target: tuple[str, int]) -> Iterator[Ca thread.join(timeout=10) -def _stop_process(process: subprocess.Popen[str]) -> None: - if process.poll() is None: - with suppress(ProcessLookupError): - os.killpg(process.pid, signal.SIGTERM) - try: - process.wait(timeout=30) - except subprocess.TimeoutExpired: - with suppress(ProcessLookupError): - os.killpg(process.pid, signal.SIGKILL) - process.wait(timeout=30) - - def _migration_log_has_p1001_or_process_exited(output_path: Path, process: subprocess.Popen[str]) -> tuple[bool, bool]: p1001_logged: Final = any("P1001" in line for line in error_lines(output_path.read_text())) process_exited: Final = process.poll() is not None @@ -422,7 +410,7 @@ def test_unreachable_database_recovers_after_postgres_forwarder_starts( open_forwarder() completed_returncode: Final = process.wait(timeout=240) finally: - _stop_process(process) + stop_process_group(process, 30) output: Final = output_path.read_text() errors: Final = error_lines(output) p1001_errors: Final = tuple(line for line in errors if "P1001" in line) diff --git a/tests/integration/management/test_batch_file_usage_caps_chaos.py b/tests/integration/management/test_batch_file_usage_caps_chaos.py index 0ec3ab597e0..5f31fa94230 100644 --- a/tests/integration/management/test_batch_file_usage_caps_chaos.py +++ b/tests/integration/management/test_batch_file_usage_caps_chaos.py @@ -1,6 +1,5 @@ import asyncio import re -import signal import time from collections.abc import Callable, Coroutine, Iterator, Mapping from contextlib import contextmanager @@ -220,8 +219,7 @@ async def test_downloads_keep_answering_while_redis_hangs(chaos: Chaos) -> None: key: Final = _key(scenario, **{DOWNLOADS: DOWNLOAD_CAP}) file_id: Final = f"file-{marker()}" minute_ends: Final = window_end(MINUTE_SECONDS, MINUTE_ROOM_SECONDS) - chaos.redis.signal(signal.SIGSTOP) - try: + with chaos.redis.paused(): burst: Final = asyncio.create_task( _download_burst(str(chaos.candidate.client.base_url), key, file_id, BURST) ) @@ -229,8 +227,6 @@ async def test_downloads_keep_answering_while_redis_hangs(chaos: Chaos) -> None: timed, partial(chaos.candidate.request, "GET", "/health/liveliness") ) during: Final = await burst - finally: - chaos.redis.signal(signal.SIGCONT) assert time.time() < minute_ends assert liveliness.response.status_code == 200, liveliness.response.text assert liveliness.after - liveliness.before < 5, liveliness @@ -305,21 +301,24 @@ async def test_a_killed_worker_keeps_its_upload_slots_used_and_the_sibling_servi held_by: Final = {pid: _open_upstream_connections(pid, files.url) for pid in workers} assert sum(held_by.values()) == HELD_UPLOADS, held_by victim_pid, survivor_pid = sorted(workers, key=held_by.__getitem__) - victim: Final = psutil.Process(victim_pid) - victim.suspend() - victim.send_signal(signal.SIGKILL) - release.set() - served: Final = await burst - assert _statuses(served) == [200] * held_by[survivor_pid], (held_by, [response.text for response in served]) - assert_upload_limited( - timed(partial(upload, candidate, key, batch_file(marker(), 1))), - day_ends, - HELD_UPLOADS, - THIS_KEY, - IN_KEY, - ) - assert len(uploads_seen(files.drain(), mark)) == HELD_UPLOADS - assert candidate.request("GET", "/health/liveliness").status_code == 200 + with owned.replacing_worker(victim_pid) as kill: + kill.freeze() + await asyncio.to_thread(kill) + release.set() + served: Final = await burst + assert _statuses(served) == [200] * held_by[survivor_pid], ( + held_by, + [response.text for response in served], + ) + assert_upload_limited( + timed(partial(upload, candidate, key, batch_file(marker(), 1))), + day_ends, + HELD_UPLOADS, + THIS_KEY, + IN_KEY, + ) + assert len(uploads_seen(files.drain(), mark)) == HELD_UPLOADS + assert candidate.request("GET", "/health/liveliness").status_code == 200 def test_the_daily_count_survives_a_proxy_restart(chaos: Chaos) -> None: diff --git a/tests/integration/management/test_credential_federation_values.py b/tests/integration/management/test_credential_federation_values.py index c2c80eef732..99be9f0cb8e 100644 --- a/tests/integration/management/test_credential_federation_values.py +++ b/tests/integration/management/test_credential_federation_values.py @@ -862,9 +862,9 @@ def test_worker_kill_mid_credential_burst(gateway: Gateway, tmp_path: Path) -> N return name, response.status_code victim: Final = eventually(lambda: _workers(owned), lambda workers: len(workers) == 2)[0] - with ThreadPoolExecutor(max_workers=8) as pool: + with owned.replacing_worker(victim.pid) as kill, ThreadPoolExecutor(max_workers=8) as pool: futures: Final = tuple(pool.submit(create, name) for name in names) - victim.kill() + kill() outcomes: Final = tuple(future.result() for future in futures) assert all(status == 200 or isinstance(status, str) for _, status in outcomes), outcomes landed: Final = tuple(name for name, status in outcomes if status == 200) diff --git a/tests/integration/management/test_fallback_management_public_names.py b/tests/integration/management/test_fallback_management_public_names.py index 095cc7d1bb4..75dad3ae428 100644 --- a/tests/integration/management/test_fallback_management_public_names.py +++ b/tests/integration/management/test_fallback_management_public_names.py @@ -397,12 +397,15 @@ def test_fallback_write_survives_a_worker_kill(gateway: Gateway, tmp_path: Path) _wait_until_served(owned.gateway, model) assert _create_fallback(owned.gateway, scenario, before, [target]).status_code == 200 victim: Final = eventually(lambda: _workers(owned), lambda workers: len(workers) == 2)[0] - victim.kill() fresh: Final = scenario.cleanups.enter_context( httpx.Client(base_url=owned.gateway.client.base_url, timeout=15, trust_env=False) ) survivor: Final = Gateway(fresh, owned.gateway.key, owned.gateway.upstream_url) - degraded: Final = tuple(_create_fallback(survivor, scenario, model, [target]) for model in (after, last)) + with owned.replacing_worker(victim.pid) as kill: + kill() + degraded: Final = tuple( + _create_fallback(survivor, scenario, model, [target]) for model in (after, last) + ) assert all(response.status_code == 200 for response in degraded), [r.text for r in degraded] covered: Final = _covered_models(_stored_fallbacks()) assert {before, after, last} <= covered, covered diff --git a/tests/integration/management/test_project_lifecycle.py b/tests/integration/management/test_project_lifecycle.py index fafda579162..d503daf6a0f 100644 --- a/tests/integration/management/test_project_lifecycle.py +++ b/tests/integration/management/test_project_lifecycle.py @@ -1,7 +1,6 @@ import json import os import re -import signal import threading from collections.abc import Iterator from concurrent.futures import ThreadPoolExecutor @@ -958,14 +957,12 @@ def test_ownership_enforced_after_worker_kill(ownership_gateway: Gateway, tmp_pa assert len(workers) == 2, tuple(process.pid for process in workers) victim: Final = workers[0] survivor: Final = workers[1] - with ThreadPoolExecutor(max_workers=24) as pool: + with owned.replacing_worker(victim.pid) as kill, ThreadPoolExecutor(max_workers=24) as pool: futures: Final = tuple(pool.submit(serve, index) for index in range(24)) try: assert started.wait(timeout=10) - os.kill(victim.pid, signal.SIGKILL) + kill() release_stream.set() - psutil.wait_procs((victim,), timeout=10) - assert not psutil.pid_exists(victim.pid), victim.pid finally: release_stream.set() burst: Final = tuple(future.result(timeout=90) for future in futures) diff --git a/tests/integration/management/test_scim_group_pathless_patch.py b/tests/integration/management/test_scim_group_pathless_patch.py index 5ed784d2bb8..efd3afba5b3 100644 --- a/tests/integration/management/test_scim_group_pathless_patch.py +++ b/tests/integration/management/test_scim_group_pathless_patch.py @@ -1,5 +1,3 @@ -import os -import signal import threading import uuid from collections.abc import Callable, Mapping, Sequence @@ -11,7 +9,6 @@ from types import MappingProxyType from typing import Final import httpx -import psutil import pytest from integration._support.client import Gateway, Scenario, eventually, object_value, string_value from integration._support.process import owned_proxy_process @@ -522,21 +519,17 @@ def _patch_status(candidate: Gateway, team: str) -> int | None: return None -def _serving_workers(root: int) -> tuple[psutil.Process, ...]: - return tuple(child for child in psutil.Process(root).children() if child.children()) - - @pytest.mark.timeout(240) def test_worker_kill_mid_burst_keeps_serving_and_leaves_no_empty_key(gateway: Gateway, tmp_path: Path) -> None: with owned_proxy_process(gateway, tmp_path, {}, workers=2) as owned, gateway.scenario() as scenario: teams: Final = _burst_teams(scenario) - workers: Final = eventually( - lambda: _serving_workers(owned.process.pid), lambda children: len(children) >= 2, seconds=30 - ) - attempts: Final = _burst(owned.gateway, tuple(teams), lambda: os.kill(workers[0].pid, signal.SIGKILL)) - assert not workers[0].is_running() or workers[0].status() == psutil.STATUS_ZOMBIE, workers[0] - probe: Final = scenario.team() - eventually(lambda: _patch_status(owned.gateway, probe), lambda status_code: status_code == 200, seconds=60) + workers: Final = eventually(owned.worker_pids, lambda pids: len(pids) >= 2, seconds=30) + with owned.replacing_worker(min(workers)) as kill: + attempts: Final = _burst(owned.gateway, tuple(teams), kill) + probe: Final = scenario.team() + eventually( + lambda: _patch_status(owned.gateway, probe), lambda status_code: status_code == 200, seconds=60 + ) _assert_burst_reflected(owned.gateway, teams, attempts) @@ -545,5 +538,5 @@ def test_rolling_restart_mid_burst_drains_without_empty_keys(gateway: Gateway, t with gateway.scenario() as scenario, owned_proxy_process(gateway, tmp_path, {}, workers=2) as replacement: teams: Final = _burst_teams(scenario) with owned_proxy_process(gateway, tmp_path, {}, workers=2) as retiring: - attempts: Final = _burst(retiring.gateway, tuple(teams), retiring.process.terminate) + attempts: Final = _burst(retiring.gateway, tuple(teams), retiring.begin_shutdown) _assert_burst_reflected(replacement.gateway, teams, attempts) diff --git a/tests/integration/management/test_team_delete_chaos.py b/tests/integration/management/test_team_delete_chaos.py index ebe515f59ec..9df4d1bb973 100644 --- a/tests/integration/management/test_team_delete_chaos.py +++ b/tests/integration/management/test_team_delete_chaos.py @@ -382,21 +382,21 @@ def test_worker_killed_mid_burst_leaves_every_team_fully_deleted_after_retry( fleet: Final = _fleet(owned.gateway, scenario) pinned: Final = fleet[0] burst: Final = Burst(owned.gateway) - with _holding_team_lock(pinned.team_id) as holder_pid: - burst.start(pool, fleet) - _await_pinned_delete_blocked(burst, pinned, holder_pid, record_property) - before: Final = _workers(root) - assert len(before) == WORKERS, [process.pid for process in before] - victim: Final = before[0] - victim.kill() # SIGKILL with the pinned delete blocked: the worker cannot finish its in-flight deletes - victim.wait(timeout=10) - with httpx.Client(base_url=str(owned.gateway.client.base_url), timeout=15, trust_env=False) as fresh: - readiness: Final = fresh.get("/health/readiness") - assert readiness.status_code == 200, ( - f"/health/readiness with worker {victim.pid} dead: {readiness.status_code} {readiness.text}" - ) - # The lock is released: the pinned delete finishes on the survivor, or was dropped with the victim. - outcomes: Final = burst.outcomes() + before: Final = _workers(root) + assert len(before) == WORKERS, [process.pid for process in before] + victim: Final = before[0] + with owned.replacing_worker(victim.pid) as kill: + with _holding_team_lock(pinned.team_id) as holder_pid: + burst.start(pool, fleet) + _await_pinned_delete_blocked(burst, pinned, holder_pid, record_property) + kill() # SIGKILL with the pinned delete blocked: the worker cannot finish its in-flight deletes + with httpx.Client(base_url=str(owned.gateway.client.base_url), timeout=15, trust_env=False) as fresh: + readiness: Final = fresh.get("/health/readiness") + assert readiness.status_code == 200, ( + f"/health/readiness with worker {victim.pid} dead: {readiness.status_code} {readiness.text}" + ) + # The lock is released: the pinned delete finishes on the survivor, or was dropped with the victim. + outcomes: Final = burst.outcomes() respawned: Final = eventually( lambda: tuple(process.pid for process in _workers(root)), lambda pids: len(pids) == WORKERS and victim.pid not in pids, @@ -480,7 +480,7 @@ def test_proxy_terminated_mid_burst_then_restarted_leaves_every_team_fully_delet with _holding_team_lock(pinned.team_id) as holder_pid: burst.start(pool, fleet) _await_pinned_delete_blocked(burst, pinned, holder_pid, record_property) - doomed.process.terminate() # SIGTERM with the pinned delete blocked: uvicorn stops accepting and drains + doomed.begin_shutdown() # SIGTERM with the pinned delete blocked: uvicorn stops accepting and drains assert burst.pending(pinned.team_id), f"{pinned.team_id}: delete answered while its team lock was held" # The lock is released: the drain lets the pinned delete finish before the proxy exits. doomed.process.wait(timeout=120) diff --git a/tests/integration/management/test_team_metadata_validation.py b/tests/integration/management/test_team_metadata_validation.py index b1e30a0d86f..f02f1e66a15 100644 --- a/tests/integration/management/test_team_metadata_validation.py +++ b/tests/integration/management/test_team_metadata_validation.py @@ -10,7 +10,7 @@ import pytest from pydantic import JsonValue from tests.integration._support.client import Gateway, eventually, gateway_from_environment, object_value, string_value -from tests.integration._support.process import owned_proxy +from tests.integration._support.process import owned_proxy, stop_child _SERVICE_URL: Final = "http://127.0.0.1:9414" _CONFIG: Final = Path(__file__).with_name("team_metadata_validation_proxy_config.yaml") @@ -87,9 +87,7 @@ def cost_center_service() -> Iterator[None]: assert health is not None and health.status_code == 200 yield finally: - if process.poll() is None: - process.terminate() - process.wait(timeout=5) + stop_child(process, grace_seconds=5) @pytest.fixture(scope="module") diff --git a/tests/integration/management/test_user_updates_wedged_coordination_redis.py b/tests/integration/management/test_user_updates_wedged_coordination_redis.py index 0715059744e..79b1f3166e3 100644 --- a/tests/integration/management/test_user_updates_wedged_coordination_redis.py +++ b/tests/integration/management/test_user_updates_wedged_coordination_redis.py @@ -1,5 +1,4 @@ import os -import signal import time import uuid from collections.abc import Callable, Mapping @@ -8,7 +7,6 @@ from pathlib import Path from typing import Final from urllib.parse import urlsplit, urlunsplit -import psutil import psycopg import pytest from psycopg import sql @@ -16,12 +14,14 @@ from pydantic import JsonValue from redis import Redis from redis.client import PubSub +from tests.integration._support.budget import case_timeout_seconds from tests.integration._support.client import JSON_OBJECT, Gateway, eventually, object_value, string_value -from tests.integration._support.process import owned_proxy +from tests.integration._support.process import owned_proxy_process from tests.integration._support.redis_process import owned_redis _USERS: Final = 60 _BURST: Final = 30 +_WEDGED_PUBLISHES_PER_BURST_USER: Final = 2 _HANDLER_BUDGET_SECONDS: Final = 0.75 _BULK_BUDGET_SECONDS: Final = 2.0 _CHANNEL: Final = "litellm_proxy.auth_cache_invalidation" @@ -52,26 +52,6 @@ def _received(pubsub: PubSub) -> tuple[dict[str, JsonValue], ...]: messages.append(JSON_OBJECT.validate_json(data)) -def _worker_pid(port: int) -> int: - for process in psutil.process_iter(): - parent = process.parent() - if parent is None: - continue - try: - cmdline = parent.cmdline() - own_cmdline = process.cmdline() - except (psutil.NoSuchProcess, psutil.AccessDenied): - continue - if ( - "integration._support.proxy" in cmdline - and "--port" in cmdline - and str(port) in cmdline - and not any("prisma" in part for part in own_cmdline) - ): - return process.pid - raise AssertionError(f"no uvicorn worker found under the owned proxy on port {port}") - - def _burst_call( index: int, users: tuple[str, ...], key: str, team_id: str, customer_id: str ) -> tuple[str, dict[str, JsonValue]]: @@ -88,7 +68,7 @@ def _burst_call( return "/customer/update", {"user_id": customer_id, "max_budget": 7.0 + index} -@pytest.mark.timeout(240) +@pytest.mark.timeout(case_timeout_seconds()) @pytest.mark.covers( "mgmt.user.update.budget_change_returns_promptly_with_wedged_coordination_redis", "mgmt.user.bulk_update.budget_change_returns_promptly_with_wedged_coordination_redis", @@ -112,7 +92,7 @@ def test_user_budget_updates_return_promptly_while_coordination_redis_is_wedged( prior_logs: Final = frozenset(results_dir.glob("owned-proxy-*.log")) with ( owned_redis(tmp_path) as coordination, - owned_proxy( + owned_proxy_process( gateway, tmp_path, { @@ -123,9 +103,10 @@ def test_user_budget_updates_return_promptly_while_coordination_redis_is_wedged( config=Path("tests/integration/coordination_redis_proxy_config.yaml"), remove_environment=("DATABASE_URL_READ_REPLICA",), workers=2, - ) as candidate, + ) as owned, Redis(host=coordination.host, port=coordination.port, socket_timeout=1) as subscriber_client, ): + candidate: Final = owned.gateway pubsub: Final = subscriber_client.pubsub() pubsub.subscribe(_CHANNEL) received: list[dict[str, JsonValue]] = [] @@ -168,8 +149,8 @@ def test_user_budget_updates_return_promptly_while_coordination_redis_is_wedged( assert timings["h2_healthy_control"] < _HANDLER_BUDGET_SECONDS, ( f"healthy control update took {timings['h2_healthy_control']:.3f}s" ) - coordination.signal(signal.SIGSTOP) - try: + with owned.replacing_worker(min(owned.worker_pids())) as kill, coordination.paused(): + kill() timings["s2_wedged_control"] = _timed_post( candidate, "/user/update", {"user_id": users[0], "tpm_limit": 1000} ) @@ -225,22 +206,17 @@ def test_user_budget_updates_return_promptly_while_coordination_redis_is_wedged( f"/user/update with an invalid body took {timings['s8_invalid_body']:.3f}s" ) - def burst_request(path: str, body: Mapping[str, JsonValue]) -> tuple[object, float]: + def burst_request(path: str, body: Mapping[str, JsonValue]) -> tuple[int, float]: started: Final = time.monotonic() - try: - response: Final = candidate.client.request( - "POST", - path, - json=body, - headers={"Authorization": f"Bearer {candidate.key}"}, - timeout=60, - ) - return response.status_code, time.monotonic() - started - except Exception as error: # noqa: BLE001 # the killed worker drops in-flight requests - return error, time.monotonic() - started + response: Final = candidate.client.request( + "POST", + path, + json=body, + headers={"Authorization": f"Bearer {candidate.key}"}, + timeout=60, + ) + return response.status_code, time.monotonic() - started - port: Final = candidate.client.base_url.port - assert port is not None, f"owned proxy client has no port: {candidate.client.base_url}" with ThreadPoolExecutor(_BURST) as pool: futures: Final = [ pool.submit( @@ -249,15 +225,9 @@ def test_user_budget_updates_return_promptly_while_coordination_redis_is_wedged( ) for i in range(_BURST) ] - os.kill(_worker_pid(port), signal.SIGKILL) results: Final = [future.result() for future in futures] - responses: Final = [(status, elapsed) for status, elapsed in results if isinstance(status, int)] - failures: Final = [status for status, _elapsed in responses if status != 200] + failures: Final = [status for status, _elapsed in results if status != 200] assert not failures, f"burst responses that were not 200: {failures}" - transport_errors: Final = [status for status, _elapsed in results if not isinstance(status, int)] - assert len(transport_errors) <= 3, ( - f"{len(transport_errors)} requests raised transport errors: {transport_errors!r}" - ) elapsed_sorted: Final = sorted( elapsed for i, (status, elapsed) in enumerate(results) if i % 5 in (0, 1, 4) and status == 200 ) @@ -276,39 +246,38 @@ def test_user_budget_updates_return_promptly_while_coordination_redis_is_wedged( assert timings["c1_survivor"] < _HANDLER_BUDGET_SECONDS, ( f"control update on the surviving worker took {timings['c1_survivor']:.3f}s" ) - finally: - coordination.signal(signal.SIGCONT) - wedged_keys: Final = {users[i] for i in range(_BURST) if i % 5 == 0 and i != 0} | {f"team_id:{team_id}"} + wedged_keys: Final = {users[i] for i in range(_BURST) if i % 5 == 0 and i != 0} | { + f"team_id:{team_id}" + } - def proxy_log() -> str: - return "".join( - path.read_text() for path in results_dir.glob("owned-proxy-*.log") if path not in prior_logs + def proxy_log() -> str: + return "".join( + path.read_text() for path in results_dir.glob("owned-proxy-*.log") if path not in prior_logs + ) + + team_wedged_key: Final = f"team_id:{team_id}" + eventually( + proxy_log, + lambda text: ( + all( + text.count(f"publish for {wedged_key} failed") >= _WEDGED_PUBLISHES_PER_BURST_USER + for wedged_key in wedged_keys + if wedged_key != team_wedged_key + ) + and ( + f"publish for {team_wedged_key} failed" in text + or f"internal usage cache entry {team_wedged_key}" in text + ) + ), + seconds=45, ) - - team_wedged_key: Final = f"team_id:{team_id}" - eventually( - proxy_log, - lambda text: ( - all( - f"publish for {wedged_key} failed" in text - for wedged_key in wedged_keys - if wedged_key != team_wedged_key - ) - and ( - f"publish for {team_wedged_key} failed" in text - or f"internal usage cache entry {team_wedged_key}" in text - ) - ), - seconds=45, - ) marker: Final = len(received) drained() recovered_keys: Final = {str(message.get("cache_key")) for message in received[marker:]} assert recovered_keys.isdisjoint(wedged_keys), ( f"wedged publishes unexpectedly landed after recovery: {sorted(recovered_keys & wedged_keys)}" ) - coordination.signal(signal.SIGSTOP) - try: + with coordination.paused(): timings["r1b_short_wedge_a"] = _timed_post( candidate, "/user/update", {"user_id": users[4], "max_budget": 15.0} ) @@ -321,8 +290,6 @@ def test_user_budget_updates_return_promptly_while_coordination_redis_is_wedged( assert timings["r1b_short_wedge_b"] < _HANDLER_BUDGET_SECONDS, ( f"/user/update inside a short wedge took {timings['r1b_short_wedge_b']:.3f}s" ) - finally: - coordination.signal(signal.SIGCONT) eventually( drained, lambda messages: {str(message.get("cache_key")) for message in messages} >= {users[4], users[5]}, diff --git a/tests/integration/management/test_vector_store_file_list_managed_ids.py b/tests/integration/management/test_vector_store_file_list_managed_ids.py index c6124519c9e..13a485f59d4 100644 --- a/tests/integration/management/test_vector_store_file_list_managed_ids.py +++ b/tests/integration/management/test_vector_store_file_list_managed_ids.py @@ -2,7 +2,6 @@ import base64 import hashlib import json import re -import signal import threading import uuid from collections.abc import Callable, Generator, Mapping @@ -743,11 +742,11 @@ def test_worker_sigkill_mid_burst_leaves_the_sibling_mapping_ids(gateway: Gatewa held_by: Final = MappingProxyType({pid: _open_connections_to(pid, wire.url) for pid in workers}) assert sum(held_by.values()) == 20, held_by victim_pid, survivor_pid = sorted(workers, key=held_by.__getitem__) - victim: Final = psutil.Process(victim_pid) - victim.suspend() - victim.send_signal(signal.SIGKILL) - release.set() - served: Final = tuple(result for future in burst if (result := future.result()) is not None) + with owned.replacing_worker(victim_pid) as kill: + kill.freeze() + kill() + release.set() + served: Final = tuple(result for future in burst if (result := future.result()) is not None) assert held_by[survivor_pid] >= 1, held_by assert len(served) == held_by[survivor_pid], (held_by, len(served)) for response in served: diff --git a/tests/integration/mcp/test_mcp_toolsets.py b/tests/integration/mcp/test_mcp_toolsets.py index 7ff512a699b..f275bd0e427 100644 --- a/tests/integration/mcp/test_mcp_toolsets.py +++ b/tests/integration/mcp/test_mcp_toolsets.py @@ -524,8 +524,7 @@ def test_owned_proxy_loss_retains_exit_status_and_first_http_failure(gateway: Ga with owned_proxy_process(gateway, tmp_path, {}) as owned: response: Final = owned.gateway.client.get("/health/readiness") assert response.status_code == 200, response.text - owned.process.kill() - owned.process.wait(timeout=3) + owned.kill_supervisor() owned.gateway.client.get("/health/readiness") record: Final = TypeAdapter(dict[str, int | float | str | None]).validate_json( owned.log.with_suffix(".exit.json").read_text() diff --git a/tests/integration/mcp/test_mcp_user_env_vars.py b/tests/integration/mcp/test_mcp_user_env_vars.py index 0170eb77491..75bb9ede182 100644 --- a/tests/integration/mcp/test_mcp_user_env_vars.py +++ b/tests/integration/mcp/test_mcp_user_env_vars.py @@ -1,4 +1,3 @@ -import signal import uuid from collections.abc import Mapping from concurrent.futures import ThreadPoolExecutor @@ -311,8 +310,7 @@ def test_concurrent_users_across_processes_never_leak_and_survive_a_killed_proce outcomes: Final = burst(pool, (gateway, peer, doomed.gateway)) assert all(code in {200, 412} for _, code in outcomes), outcomes assert all(code == 200 for index, code in outcomes if index % 4 != 2), outcomes - doomed.process.send_signal(signal.SIGKILL) - doomed.process.wait(timeout=10) + doomed.kill_supervisor() after_kill: Final = burst(pool, (gateway, peer)) assert all(code in {200, 412} for _, code in after_kill), after_kill assert all(code == 200 for index, code in after_kill if index % 4 != 2), after_kill diff --git a/tests/integration/mcp/test_proxy_mcp_e2e.py b/tests/integration/mcp/test_proxy_mcp_e2e.py index 4cfac6ed784..959d2620f58 100644 --- a/tests/integration/mcp/test_proxy_mcp_e2e.py +++ b/tests/integration/mcp/test_proxy_mcp_e2e.py @@ -41,6 +41,8 @@ from litellm.proxy.proxy_server import ( initialize, ) from tests._master_key import MASTER_KEY +from tests.integration._support.process import stop_child +from tests.integration._support.timing import assert_enforced_within from tests.integration._support.wire import Reply, Wire, wire_server from tests.integration._support.wire import Request as WireRequest @@ -83,7 +85,10 @@ async def drain_logging_worker() -> typing.AsyncIterator[None]: async def test_schema_validation_stops_expensive_work_and_recovers(cancel: bool, concurrency: int) -> None: import psutil - from litellm.proxy._experimental.mcp_server.tool_search import _tool_argument_validation_error + from litellm.proxy._experimental.mcp_server.tool_search import ( + _VALIDATION_TIMEOUT_SECONDS, + _tool_argument_validation_error, + ) existing_children: Final = frozenset(child.pid for child in psutil.Process().children()) warm_count: Final = min(concurrency, 4) @@ -117,7 +122,7 @@ async def test_schema_validation_stops_expensive_work_and_recovers(cancel: bool, assert all(isinstance(result, asyncio.CancelledError) for result in await group) else: assert await group == ["Tool argument validation exceeded its time limit"] * concurrency - assert time.monotonic() - started < 35 + assert_enforced_within(time.monotonic() - started, configured=_VALIDATION_TIMEOUT_SECONDS, slack=5) assert all(not worker.is_running() for worker in workers), "Cancelled validation must terminate worker CPU work" assert await _tool_argument_validation_error({"type": "object"}, {}) is None finally: @@ -391,12 +396,7 @@ def _math_http_server(offset: int) -> typing.Iterator[str]: time.sleep(0.05) yield f"http://{host}:{port}" finally: - process.terminate() - try: - process.wait(timeout=5) - except subprocess.TimeoutExpired: - process.kill() - process.wait(timeout=5) + stop_child(process, grace_seconds=5) @pytest.fixture(scope="module") diff --git a/tests/integration/messages_endpoint/providers/anthropic/test_anthropic_messages_pre_content_retry_chaos.py b/tests/integration/messages_endpoint/providers/anthropic/test_anthropic_messages_pre_content_retry_chaos.py index b6e848bad75..8944d1990d9 100644 --- a/tests/integration/messages_endpoint/providers/anthropic/test_anthropic_messages_pre_content_retry_chaos.py +++ b/tests/integration/messages_endpoint/providers/anthropic/test_anthropic_messages_pre_content_retry_chaos.py @@ -2,7 +2,6 @@ import asyncio import base64 import json import re -import signal import threading import uuid from collections.abc import Callable, Mapping @@ -364,15 +363,15 @@ async def test_worker_sigkill_mid_burst_leaves_the_sibling_retrying_pre_content_ held_by: Final = MappingProxyType({pid: _open_upstream_connections(pid, wire.url) for pid in workers}) assert sum(held_by.values()) == len(calls), held_by victim_pid, survivor_pid = sorted(workers, key=held_by.__getitem__) - victim: Final = psutil.Process(victim_pid) - victim.suspend() - victim.send_signal(signal.SIGKILL) - release.set() - served: Final = await burst - assert held_by[survivor_pid] >= 10, held_by - assert len(served) == held_by[survivor_pid], (held_by, len(served)) - for item in served: - _assert_completed_on_the_second_attempt(item) + with owned.replacing_worker(victim_pid) as kill: + kill.freeze() + await asyncio.to_thread(kill) + release.set() + served: Final = await burst + assert held_by[survivor_pid] >= 10, held_by + assert len(served) == held_by[survivor_pid], (held_by, len(served)) + for item in served: + _assert_completed_on_the_second_attempt(item) follow_up: Final = _Call(endpoint="messages", stream=True, marker=uuid.uuid4().hex) (answered,) = await _burst(base_url, candidate.key, models, (follow_up,)) _assert_completed_on_the_second_attempt(answered) diff --git a/tests/integration/messages_endpoint/providers/anthropic/test_anthropic_messages_timeout_wire.py b/tests/integration/messages_endpoint/providers/anthropic/test_anthropic_messages_timeout_wire.py index 76c29ad7763..fb0a9d24c0f 100644 --- a/tests/integration/messages_endpoint/providers/anthropic/test_anthropic_messages_timeout_wire.py +++ b/tests/integration/messages_endpoint/providers/anthropic/test_anthropic_messages_timeout_wire.py @@ -5,6 +5,7 @@ from typing import Final import pytest from integration._support.client import JSON_OBJECT, Gateway +from integration._support.timing import assert_cut_short, timed from integration._support.wire import Reply, Request, wire_server _UPSTREAM_STALL_SECONDS: Final = 4.0 @@ -46,13 +47,13 @@ def test_messages_endpoint_honors_configured_timeout_against_stalled_upstream(ga api_key="synthetic-anthropic-key", timeout=_CONFIGURED_TIMEOUT_SECONDS, ) - started: Final = time.monotonic() - response: Final = gateway.request( - "POST", - "/v1/messages", - {"model": model, "max_tokens": 16, "messages": [{"role": "user", "content": prompt}]}, + response: Final = timed( + lambda: gateway.request( + "POST", + "/v1/messages", + {"model": model, "max_tokens": 16, "messages": [{"role": "user", "content": prompt}]}, + ) ) - elapsed: Final = time.monotonic() - started - assert response.status_code == 408, response.text - assert elapsed < _UPSTREAM_STALL_SECONDS, f"timed out only after {elapsed:.2f}s: {response.text}" + assert response.value.status_code == 408, response.value.text + assert_cut_short(response.seconds, injected_delay=_UPSTREAM_STALL_SECONDS) assert len(wire.drain()) == 1 diff --git a/tests/integration/messages_endpoint/providers/anthropic/test_pdf_document_count_tokens_chaos.py b/tests/integration/messages_endpoint/providers/anthropic/test_pdf_document_count_tokens_chaos.py index f196018f117..80ace4c4535 100644 --- a/tests/integration/messages_endpoint/providers/anthropic/test_pdf_document_count_tokens_chaos.py +++ b/tests/integration/messages_endpoint/providers/anthropic/test_pdf_document_count_tokens_chaos.py @@ -1,7 +1,5 @@ import asyncio -import os import re -import signal import threading import uuid from collections.abc import Callable @@ -209,11 +207,13 @@ async def test_worker_sigkill_mid_pdf_count_burst_leaves_the_sibling_pricing_pag ) held_by: Final = MappingProxyType({pid: _held_upstream_connections(pid, wire.url) for pid in workers}) victim: Final = max(workers, key=held_by.__getitem__) - os.kill(victim, signal.SIGKILL) + with owned.replacing_worker(victim) as kill: + await asyncio.to_thread(kill) + release.set() + served: Final = await burst + wire.drain() finally: release.set() - served: Final = await burst - wire.drain() baseline: Final = _baseline(owned.gateway) after: Final = await _fire_counts(owned_url, owned.gateway.key) after_received: Final = wire.drain() diff --git a/tests/integration/observability/test_arize_otel_v2_openinference_chaos.py b/tests/integration/observability/test_arize_otel_v2_openinference_chaos.py index 9ba196cb28e..79b0d9084aa 100644 --- a/tests/integration/observability/test_arize_otel_v2_openinference_chaos.py +++ b/tests/integration/observability/test_arize_otel_v2_openinference_chaos.py @@ -9,7 +9,6 @@ from typing import Final from urllib.parse import urlsplit import httpx -import psutil import pytest from _openinference_support import ( CHAT_TOOLS, @@ -334,13 +333,6 @@ def test_arize_otel_v2_f2_slow_sink_does_not_deadlock(gateway: Gateway, tmp_path _matching_marker_span(rig.destination, marker) -def _has_exited(process: psutil.Process) -> bool: - try: - return process.status() == psutil.STATUS_ZOMBIE - except psutil.NoSuchProcess: - return True - - @pytest.mark.timeout(2 * graceful_stop_seconds() + 120) def test_arize_otel_v2_f3_one_proxy_worker_can_die(gateway: Gateway, tmp_path: Path) -> None: markers: Final = tuple("f3-" + uuid.uuid4().hex for _ in range(8)) @@ -365,11 +357,8 @@ def test_arize_otel_v2_f3_one_proxy_worker_can_die(gateway: Gateway, tmp_path: P destination_handler=sink, fresh_client_connections=True, ) as rig: - children: Final = psutil.Process(rig.owned.process.pid).children(recursive=True) - workers: Final = tuple( - child for child in children if child.is_running() and "resource_tracker" not in " ".join(child.cmdline()) - ) - assert len(workers) >= 2, tuple((worker.pid, worker.name()) for worker in workers) + workers: Final = rig.owned.worker_pids() + assert len(workers) >= 2, workers with ThreadPoolExecutor(max_workers=len(markers)) as executor: try: futures: Final = tuple( @@ -388,12 +377,12 @@ def test_arize_otel_v2_f3_one_proxy_worker_can_die(gateway: Gateway, tmp_path: P seconds=10, ) assert observed >= 2, observed - workers[0].kill() - assert eventually(lambda: _has_exited(workers[0]), bool, seconds=10), workers[0] - release.set() - in_flight: Final = tuple( - (marker, future.result(timeout=60)) for marker, future in zip(markers, futures, strict=True) - ) + with rig.owned.replacing_worker(min(workers)) as kill: + kill() + release.set() + in_flight: Final = tuple( + (marker, future.result(timeout=60)) for marker, future in zip(markers, futures, strict=True) + ) finally: release.set() survivor: Final = "f3-survivor-" + uuid.uuid4().hex diff --git a/tests/integration/observability/test_azure_content_safety_audit.py b/tests/integration/observability/test_azure_content_safety_audit.py index 99d868efe6f..85ced5f719b 100644 --- a/tests/integration/observability/test_azure_content_safety_audit.py +++ b/tests/integration/observability/test_azure_content_safety_audit.py @@ -7,7 +7,6 @@ from contextlib import ExitStack from pathlib import Path from typing import Final -import psutil import pytest import yaml from integration._support.client import Gateway, eventually, gateway_from_environment, object_value @@ -953,13 +952,8 @@ def test_killing_a_worker_mid_burst_leaves_no_duplicate_rows( chaos_rig: tuple[OwnedProxy, Wire, Wire, threading.Event], ) -> None: owned, azure, provider, _ = chaos_rig - port: Final = owned.gateway.client.base_url.port - workers: Final = tuple( - child - for child in psutil.Process(owned.process.pid).children(recursive=False) - if any(connection.laddr.port == port for connection in child.net_connections(kind="tcp")) - ) - assert len(workers) == 2, [worker.pid for worker in workers] + workers: Final = owned.worker_pids() + assert len(workers) == 2, workers with owned.gateway.scenario() as scenario: model: Final = scenario.model( model="openai/gpt-4.1-mini", api_base=provider.url + "/v1", api_key="synthetic-provider-key" @@ -972,11 +966,12 @@ def test_killing_a_worker_mid_burst_leaves_no_duplicate_rows( with ThreadPoolExecutor(max_workers=6) as pool: future_map: Final = tuple(pool.submit(call, identity) for identity in identities) - workers[0].kill() - outcomes: Final = tuple( - future.result() if not future.exception() else (identities[index], -1) - for index, future in enumerate(future_map) - ) + with owned.replacing_worker(min(workers)) as kill: + kill() + outcomes: Final = tuple( + future.result() if not future.exception() else (identities[index], -1) + for index, future in enumerate(future_map) + ) survivors: Final = tuple(status for _, status in outcomes if status != -1) assert survivors and {status for status in survivors} == {200}, outcomes scans: Final = _shield_prompts(azure.drain()) diff --git a/tests/integration/observability/test_azure_content_safety_dispatch_resilience.py b/tests/integration/observability/test_azure_content_safety_dispatch_resilience.py index 5d7b5948d96..d443ab9a2d0 100644 --- a/tests/integration/observability/test_azure_content_safety_dispatch_resilience.py +++ b/tests/integration/observability/test_azure_content_safety_dispatch_resilience.py @@ -2,18 +2,15 @@ import concurrent.futures import socket import threading import uuid -from collections.abc import Iterator from contextlib import ExitStack from dataclasses import dataclass from pathlib import Path from typing import Final import httpx -import psutil import pytest from integration._support.client import Gateway, Scenario, eventually from integration._support.database import read_rows -from integration._support.process import OwnedProxy from integration._support.redis_process import owned_redis from integration._support.wire import Wire, wire_server from integration.observability.azure_dispatch_support import ( @@ -25,7 +22,6 @@ from integration.observability.azure_dispatch_support import ( dispatch_rig, provider_handler, provider_texts, - write_dispatch_config, ) from pydantic import JsonValue @@ -104,16 +100,6 @@ def _safe_request(candidate: Gateway, model: str, prompt: str) -> httpx.Response return None -def _worker_processes(owned: OwnedProxy) -> tuple[psutil.Process, ...]: - descendants: Final = tuple(psutil.Process(owned.process.pid).children(recursive=True)) - return tuple( - process - for process in descendants - if process.is_running() - and any("spawn_main" in argument for argument in process.cmdline()) - ) - - def _operation_specs( phase: str, operation: str, @@ -388,8 +374,8 @@ def test_x3_killing_one_worker_leaves_the_other_serving( _, ): candidate: Final = owned.gateway - workers: Final = _worker_processes(owned) - assert len(workers) == 2, tuple(process.cmdline() for process in workers) + workers: Final = owned.worker_pids() + assert len(workers) == 2, workers with candidate.scenario() as scenario: model: Final = _model(scenario, provider) with concurrent.futures.ThreadPoolExecutor(max_workers=12) as executor: @@ -398,17 +384,17 @@ def test_x3_killing_one_worker_leaves_the_other_serving( for _ in range(4) ) assert entered.wait(timeout=30), "No request reached the slow Azure edge" - killed_worker: Final = workers[0] - survivor: Final = workers[1] - killed_worker.kill() - killed_worker.wait(timeout=10) - assert survivor.is_running(), "The second proxy worker exited with the killed worker" - prompts: Final = tuple(f"worker survivor {uuid.uuid4().hex}" for _ in range(8)) - futures: Final = tuple( - executor.submit(_safe_request, candidate, model, prompt) - for prompt in prompts - ) - responses: Final = tuple(future.result(timeout=70) for future in (*inflight, *futures)) + killed_worker: Final = min(workers) + survivor: Final = max(workers) + with owned.replacing_worker(killed_worker) as kill: + kill() + assert survivor in owned.worker_pids(), "The second proxy worker exited with the killed worker" + prompts: Final = tuple(f"worker survivor {uuid.uuid4().hex}" for _ in range(8)) + futures: Final = tuple( + executor.submit(_safe_request, candidate, model, prompt) + for prompt in prompts + ) + responses: Final = tuple(future.result(timeout=70) for future in (*inflight, *futures)) surviving_responses: Final = tuple(response for response in responses[4:] if response is not None) assert all(response.status_code == 200 for response in surviving_responses), tuple( response.text for response in surviving_responses @@ -470,7 +456,7 @@ def test_x4_proxy_restart_preserves_completed_spend_rows( assert arrived_count == len(interrupted_prompts), ( f"Only {arrived_count} of {len(interrupted_prompts)} interrupted requests reached Azure" ) - first.process.terminate() + first.begin_shutdown() release.set() first_proxy_lifetime.close() interrupted_results: Final = tuple(future.result(timeout=70) for future in interrupted) diff --git a/tests/integration/observability/test_azure_storage_chaos.py b/tests/integration/observability/test_azure_storage_chaos.py index 2235dc8fc56..acc2eb1439d 100644 --- a/tests/integration/observability/test_azure_storage_chaos.py +++ b/tests/integration/observability/test_azure_storage_chaos.py @@ -1,5 +1,3 @@ -import os -import signal import uuid from pathlib import Path from typing import Final @@ -149,13 +147,14 @@ def test_killing_one_worker_keeps_the_other_serving_and_uploading(gateway: Gatew key: Final = scenario.key(models=[openai_model, anthropic_model]) first: Final = mixed_burst(candidate, openai_model, anthropic_model, key, f"{marker}-first", per_surface=2) collect_files(sink, len(first)) - workers: Final = tuple( - process for process in group_members(owned.process.pid) if process.pid != owned.process.pid - ) + workers: Final = owned.worker_pids() assert workers, "no uvicorn workers in the owned proxy process group" - os.kill(workers[0].pid, signal.SIGKILL) - eventually(lambda: _readiness_ok(candidate), lambda ok: ok, seconds=30) - rest: Final = mixed_burst(candidate, openai_model, anthropic_model, key, f"{marker}-rest", per_surface=4) + with owned.replacing_worker(min(workers)) as kill: + kill() + eventually(lambda: _readiness_ok(candidate), lambda ok: ok, seconds=30) + rest: Final = mixed_burst( + candidate, openai_model, anthropic_model, key, f"{marker}-rest", per_surface=4 + ) payloads: Final = eventually( lambda: tuple(sink.payloads().values()), lambda stored: _present_count(stored, rest) == len(rest), diff --git a/tests/integration/observability/test_cache_hit_guardrail_metrics.py b/tests/integration/observability/test_cache_hit_guardrail_metrics.py index 888cfd7ab9d..605130acd38 100644 --- a/tests/integration/observability/test_cache_hit_guardrail_metrics.py +++ b/tests/integration/observability/test_cache_hit_guardrail_metrics.py @@ -1,6 +1,5 @@ import asyncio import json -import subprocess import uuid from collections.abc import Callable, Generator, Mapping from contextlib import contextmanager @@ -15,7 +14,7 @@ import pytest import yaml from integration._support.client import Gateway, Scenario, eventually, object_value, string_value from integration._support.database import read_rows -from integration._support.process import owned_proxy_process +from integration._support.process import OwnedProxy, owned_proxy_process from integration._support.wire import Reply, Request, Wire, wire_server from prometheus_client.parser import text_string_to_metric_families @@ -178,7 +177,7 @@ class Rig: guardrail_name: str policy: Wire provider: Wire - process: subprocess.Popen[bytes] + owned: OwnedProxy @contextmanager @@ -228,7 +227,7 @@ def _rig( identity, policy, provider, - owned.process, + owned, ) diff --git a/tests/integration/observability/test_cache_hit_guardrail_metrics_chaos.py b/tests/integration/observability/test_cache_hit_guardrail_metrics_chaos.py index 7e3b479e5f0..d4b62b36e8b 100644 --- a/tests/integration/observability/test_cache_hit_guardrail_metrics_chaos.py +++ b/tests/integration/observability/test_cache_hit_guardrail_metrics_chaos.py @@ -1,5 +1,4 @@ import json -import signal import threading import uuid from collections.abc import Callable @@ -8,7 +7,6 @@ from pathlib import Path from typing import Final import httpx -import psutil from integration._support.client import Gateway, eventually, object_value, string_value from integration._support.process import owned_proxy_process from integration._support.redis_process import owned_redis @@ -239,14 +237,13 @@ def test_worker_kill_mid_burst_keeps_counting(gateway: Gateway, tmp_path: Path) with _rig(gateway, tmp_path, marker, workers=2) as rig: bodies: Final = _burst_bodies(rig, marker, None) _warm(rig, bodies) - children: Final = psutil.Process(rig.process.pid).children(recursive=True) - assert children, "no uvicorn worker children found" - with ThreadPoolExecutor(max_workers=8) as pool: + workers: Final = eventually(rig.owned.worker_pids, lambda pids: len(pids) == 2) + with rig.owned.replacing_worker(min(workers)) as kill, ThreadPoolExecutor(max_workers=8) as pool: futures: Final = tuple( pool.submit(lambda b: rig.candidate.request("POST", b[0], b[1]), body) for body in bodies ) eventually(lambda: rig.policy.received.qsize() >= 3, bool, seconds=30) - children[0].send_signal(signal.SIGKILL) + kill() statuses: list[int] = [] # mutable-ok: collect per-request outcomes from concurrent futures for future in futures: try: diff --git a/tests/integration/observability/test_grayswan_wire_chaos.py b/tests/integration/observability/test_grayswan_wire_chaos.py index 16800d6c235..b57e9ac9f75 100644 --- a/tests/integration/observability/test_grayswan_wire_chaos.py +++ b/tests/integration/observability/test_grayswan_wire_chaos.py @@ -1,6 +1,4 @@ import json -import os -import signal import threading import time import uuid @@ -8,10 +6,9 @@ from concurrent.futures import ThreadPoolExecutor from pathlib import Path from typing import Final -import psutil import yaml from integration._support.client import Gateway -from integration._support.process import group_members, owned_proxy_process +from integration._support.process import owned_proxy_process from integration._support.wire import Reply, Request, wire_server from pydantic import JsonValue, TypeAdapter from test_grayswan_wire import _PROVIDER_KEY, _REQUEST_MESSAGES, _VENDOR_KEY, _monitor_bodies, _serving_model_probe @@ -206,23 +203,13 @@ def test_worker_kill_mid_burst_survivor_keeps_serving(gateway: Gateway, tmp_path model: Final = scenario.model(model="openai/gpt-4o-mini", api_base=upstream.url, api_key=_PROVIDER_KEY) warm: Final = _fire(candidate, model, "marker-warm", False) assert warm == 200 - members: Final = group_members(owned.process.pid) - candidate_port: Final = candidate.client.base_url.port - workers_listening: Final = tuple( - member - for member in members - if member.pid != owned.process.pid - and any( - connection.laddr.port == candidate_port and connection.status == "LISTEN" - for connection in member.net_connections(kind="inet") + workers: Final = owned.worker_pids() + assert len(workers) == 2, workers + with owned.replacing_worker(min(workers)) as kill: + kill() + statuses: Final = tuple( + _fire(candidate, model, f"marker-kill-{index}", False) for index in range(6) ) - ) - assert len(workers_listening) == 2, [member.pid for member in members] - victim: Final = workers_listening[0] - os.kill(victim.pid, signal.SIGKILL) - psutil.wait_procs((victim,), timeout=10) - assert not psutil.pid_exists(victim.pid), victim.pid - statuses: Final = tuple(_fire(candidate, model, f"marker-kill-{index}", False) for index in range(6)) assert all(status == 200 for status in statuses), statuses bodies: Final = _monitor_bodies(vendor_wire, expected=7) kill_bodies: Final = [ diff --git a/tests/integration/observability/test_guardrail_effects.py b/tests/integration/observability/test_guardrail_effects.py index 4acd9c73ec6..0009e21602b 100644 --- a/tests/integration/observability/test_guardrail_effects.py +++ b/tests/integration/observability/test_guardrail_effects.py @@ -1,8 +1,6 @@ import asyncio import json -import os import re -import signal import socket import threading import uuid @@ -17,13 +15,12 @@ from typing import Final import anthropic import httpx -import psutil import pytest import yaml from integration._support.client import Gateway, eventually, object_value from integration._support.database import read_rows from integration._support.mcp import mcp_peer, register_mcp, tool_names -from integration._support.process import OwnedProxy, group_members, owned_proxy, owned_proxy_process +from integration._support.process import OwnedProxy, owned_proxy, owned_proxy_process from integration._support.wire import Reply, Request, Wire, wire_server from openai import AsyncOpenAI, OpenAI from pydantic import JsonValue @@ -1772,27 +1769,9 @@ def test_responses_pre_call_denial_stream_survives_worker_kill(gateway: Gateway, candidate: Final = owned.gateway with candidate.scenario() as scenario: model: Final = scenario.model(api_base=_dead_api_base()) - members: Final = tuple( - member for member in group_members(owned.process.pid) if member.pid != owned.process.pid - ) - children: Final = tuple(member.pid for member in members) - workers: Final = tuple( - member.pid for member in members if any("spawn_main" in part for part in member.cmdline()) - ) + workers: Final = owned.worker_pids() assert len(workers) >= 2, workers - os.kill(workers[0], signal.SIGKILL) - expected: Final = len(children) - eventually( - lambda: tuple( - member.pid - for member in group_members(owned.process.pid) - if member.pid != owned.process.pid - and member.is_running() - and member.status() != psutil.STATUS_ZOMBIE - ), - lambda pids: len(pids) >= expected and any(pid not in children for pid in pids), - seconds=30, - ) + owned.kill_worker(min(workers)) def burst(index: int) -> httpx.Response: return candidate.request( @@ -2737,14 +2716,10 @@ def test_concurrent_disconnects_during_a_guardrail_outage_each_record_exactly_on def test_disconnect_scans_keep_recording_after_a_worker_is_killed(gateway: Gateway, tmp_path: Path) -> None: with _disconnect_rig(gateway, tmp_path, gated=False, pause=3, workers=2) as rig: - members: Final = tuple( - member for member in group_members(rig.owned.process.pid) if member.pid != rig.owned.process.pid - ) - workers: Final = tuple(member for member in members if any("spawn_main" in part for part in member.cmdline())) - assert len(workers) >= 2, members - workers[0].send_signal(signal.SIGKILL) - psutil.wait_procs((workers[0],), timeout=10) - with ThreadPoolExecutor(max_workers=8) as pool: + workers: Final = rig.owned.worker_pids() + assert len(workers) >= 2, workers + with rig.owned.replacing_worker(min(workers)) as kill, ThreadPoolExecutor(max_workers=8) as pool: + kill() received: Final = tuple(pool.map(lambda index: _chat_httpx(rig, index), range(8))) assert all(rig.secret(index) in line for index, line in enumerate(received)), received assert _post_call_statuses(rig, rig.model, rows=8) == (("success",),) * 8 diff --git a/tests/integration/observability/test_guardrail_stream_scope_chaos.py b/tests/integration/observability/test_guardrail_stream_scope_chaos.py index f8a9d0a4bba..6aa705e52b3 100644 --- a/tests/integration/observability/test_guardrail_stream_scope_chaos.py +++ b/tests/integration/observability/test_guardrail_stream_scope_chaos.py @@ -3,7 +3,6 @@ from __future__ import annotations import json import os import shutil -import signal import socket import subprocess import uuid @@ -18,7 +17,6 @@ from types import MappingProxyType from typing import Final, Literal, TypeAlias, cast import httpx -import psutil import pytest import yaml from integration._support.client import Gateway, JsonValue, eventually @@ -596,19 +594,6 @@ def test_h2_stream_sink_gate_does_not_block_out_of_scope_calls(rig: ChaosRig, tm ) -def _worker_processes(owned: OwnedProxy) -> tuple[psutil.Process, ...]: - return tuple( - child for child in psutil.Process(owned.process.pid).children() if "spawn_main" in _process_command(child) - ) - - -def _process_command(process: psutil.Process) -> str: - try: - return " ".join(process.cmdline()) - except psutil.Error: - return "" - - def _safe_response(future: Future[httpx.Response]) -> httpx.Response | None: try: return future.result(timeout=70) @@ -636,14 +621,14 @@ def test_h3_worker_kill_mid_burst_keeps_remaining_worker_serving(rig: ChaosRig, _rail(name_b, sink_b.url, "non_streaming", default_on=True), ) with _owned_proxy(rig, tmp_path, rails, workers=2) as owned, ThreadPoolExecutor(max_workers=30) as pool: - workers: Final = _worker_processes(owned) - assert len(workers) == 2, tuple(worker.pid for worker in workers) + workers: Final = owned.worker_pids() + assert len(workers) == 2, workers futures: Final[tuple[Future[httpx.Response], ...]] = tuple( pool.submit(_request, owned.gateway, plan, ()) for plan in plans ) assert started.wait(timeout=30), "stream requests did not reach the owned sink" - victim: Final = workers[0] - survivor: Final = workers[1] + victim: Final = min(workers) + survivor: Final = max(workers) survivor_plan: Final = CallPlan( f"audit-h3-survivor-{uuid.uuid4().hex}", f"h3-survivor-{uuid.uuid4().hex}", @@ -651,13 +636,14 @@ def test_h3_worker_kill_mid_burst_keeps_remaining_worker_serving(rig: ChaosRig, False, ) try: - victim.send_signal(signal.SIGKILL) - survivor_response: Final = _request(owned.gateway, survivor_plan, ()) - assert survivor_response.status_code == 200 and survivor_plan.marker in survivor_response.text, ( - survivor_response.status_code, - survivor_response.text, - ) - assert survivor.is_running(), survivor.pid + with owned.replacing_worker(victim) as kill: + kill() + survivor_response: Final = _request(owned.gateway, survivor_plan, ()) + assert survivor_response.status_code == 200 and survivor_plan.marker in survivor_response.text, ( + survivor_response.status_code, + survivor_response.text, + ) + assert survivor in owned.worker_pids(), survivor finally: release.set() responses: Final = tuple(_safe_response(future) for future in futures) diff --git a/tests/integration/observability/test_guardrail_timeout_all_providers.py b/tests/integration/observability/test_guardrail_timeout_all_providers.py index 1f33bd19393..8f6c098a3d4 100644 --- a/tests/integration/observability/test_guardrail_timeout_all_providers.py +++ b/tests/integration/observability/test_guardrail_timeout_all_providers.py @@ -13,6 +13,7 @@ from __future__ import annotations import json import re import socket +import sys import threading import time from collections.abc import Iterator, Mapping @@ -33,11 +34,13 @@ from cryptography.hazmat.primitives import serialization from cryptography.hazmat.primitives.asymmetric import rsa from integration._support.client import Gateway, gateway_from_environment from integration._support.process import owned_proxy_process +from integration._support.timing import assert_enforced_within from integration._support.wire import Reply, Request, wire_server SLOW_SECONDS: Final = 20 FAST_SECONDS: Final = 3 -BOUND_SECONDS: Final = 8 +GUARDRAIL_TIMEOUT_SECONDS: Final = 1 +TIMEOUT_SLACK_SECONDS: Final = 7 PROBE_CONCURRENCY: Final = 8 TOKEN_PATH: Final = "/token" TOKEN_REPLY: Final = json.dumps( @@ -173,6 +176,7 @@ class Sink: port: int seen: list[Seen] = field(default_factory=list) lock: threading.Lock = field(default_factory=threading.Lock) + closing: threading.Event = field(default_factory=threading.Event) server: ThreadingHTTPServer | None = None thread: threading.Thread | None = None @@ -194,7 +198,7 @@ class Sink: ) is_token: Final = self.path.startswith(TOKEN_PATH) if not is_token: - time.sleep(SLOW_SECONDS if self.path.startswith("/slow/") else FAST_SECONDS) + sink.closing.wait(SLOW_SECONDS if self.path.startswith("/slow/") else FAST_SECONDS) payload: Final = TOKEN_REPLY if is_token else b"{}" self.send_response(200) self.send_header("content-type", "application/json") @@ -212,9 +216,14 @@ class Sink: class Server(ThreadingHTTPServer): allow_reuse_address = True - daemon_threads = True + daemon_threads = False + block_on_close = True request_queue_size: Final = socket.SOMAXCONN + def handle_error(self, request: object, client_address: object) -> None: + if not isinstance(sys.exc_info()[1], ConnectionError): + super().handle_error(request, client_address) + self.server = Server(("127.0.0.1", self.port), Handler) self.port = self.server.server_port self.thread = threading.Thread(target=self.server.serve_forever, daemon=True) @@ -222,6 +231,7 @@ class Sink: def stop(self) -> None: assert self.server is not None and self.thread is not None + self.closing.set() self.server.shutdown() self.server.server_close() self.thread.join(timeout=5) @@ -323,7 +333,7 @@ def _rig_config(sink_url: str, root: Path) -> Path: config: Final = yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text()) config["litellm_settings"]["cache"] = False config["guardrails"] = [ - _guardrail(name, provider, sink_url, 1, dict(extra), mode) + _guardrail(name, provider, sink_url, GUARDRAIL_TIMEOUT_SECONDS, dict(extra), mode) for name, provider, extra, mode, _ in _provider_values() ] + [ _guardrail("control-generic", "generic_guardrail_api", sink_url, None, {}, "pre_call"), @@ -431,9 +441,7 @@ def test_litellm_params_timeout_bounds_outbound_call( outcome: Final = outcomes[name] calls: Final = rig.sink.calls_for(name) assert calls, f"{name}: sink saw no request for {provider}" - assert outcome.elapsed < BOUND_SECONDS, ( - f"{name}: elapsed {outcome.elapsed:.2f}s, expected under {BOUND_SECONDS}s with timeout=1" - ) + assert_enforced_within(outcome.elapsed, configured=GUARDRAIL_TIMEOUT_SECONDS, slack=TIMEOUT_SLACK_SECONDS) assert isinstance(outcome.response, httpx.Response), f"{name}: client gave up: {outcome.response!r}" assert outcome.response.status_code != 504, outcome.response.text diff --git a/tests/integration/observability/test_langfuse_delivery.py b/tests/integration/observability/test_langfuse_delivery.py index 018c6a35f31..bee754297f0 100644 --- a/tests/integration/observability/test_langfuse_delivery.py +++ b/tests/integration/observability/test_langfuse_delivery.py @@ -1,9 +1,7 @@ import asyncio import base64 import json -import signal import threading -import time import uuid from collections.abc import Iterator, Mapping, Sequence from concurrent.futures import ThreadPoolExecutor @@ -15,7 +13,6 @@ from typing import Final import anthropic import httpx import openai -import psutil import pytest import yaml from _s3_v2_support import _chat_stream_frames, _responses_stream_frames @@ -28,7 +25,8 @@ from integration._support.client import ( string_value, ) from integration._support.database import read_rows, scratch_database -from integration._support.process import group_members, owned_proxy, owned_proxy_process +from integration._support.process import owned_proxy, owned_proxy_process +from integration._support.timing import assert_cut_short, assert_within_factor_of_control, timed from integration._support.wire import Reply, Request, Wire, wire_server from opentelemetry.proto.collector.trace.v1.trace_service_pb2 import ExportTraceServiceRequest from opentelemetry.proto.common.v1.common_pb2 import KeyValue @@ -336,7 +334,10 @@ def test_prompt_fetch_encodes_the_name_retries_a_5xx_once_and_keeps_langfuse_hea flaky_prompt: Final = f"{marker}/what?" encoded_flaky_prompt: Final = f"{marker}%2Fwhat%3F" missing_prompt: Final = marker + "-missing" + steady_prompt: Final = marker + "-steady" + retry_after: Final = 30 + retried_within_control_factor: Final = 3 seen_prompt_gets: Final[list[str]] = [] # mutable-ok: the double counts attempts across requests def upstream(request: Request) -> Reply: @@ -353,8 +354,10 @@ def test_prompt_fetch_encodes_the_name_retries_a_5xx_once_and_keeps_langfuse_hea prior: Final = sum(1 for seen in seen_prompt_gets if seen.startswith(PROMPTS_PATH + encoded_flaky_prompt)) seen_prompt_gets.append(request.target) if prior == 0: - return Reply(status=503, body=b'{"message":"try later"}', headers={"retry-after": "30"}) + return Reply(status=503, body=b'{"message":"try later"}', headers={"retry-after": str(retry_after)}) return _text_prompt(flaky_prompt) + if request.target.startswith(PROMPTS_PATH + steady_prompt): + return _text_prompt(steady_prompt) seen_prompt_gets.append(request.target) return Reply( status=404, @@ -376,15 +379,24 @@ def test_prompt_fetch_encodes_the_name_retries_a_5xx_once_and_keeps_langfuse_hea missing: Final = scenario.model( model="langfuse/gpt-4o-mini", prompt_id=missing_prompt, api_base=provider.url + "/v1", api_key="synthetic" ) - started: Final = time.monotonic() - response: Final = candidate.request( - "POST", - "/v1/chat/completions", - {"model": flaky, "messages": [{"role": "user", "content": marker}], "prompt_variables": {"word": marker}}, + steady: Final = scenario.model( + model="langfuse/gpt-4o-mini", prompt_id=steady_prompt, api_base=provider.url + "/v1", api_key="synthetic" ) - elapsed: Final = time.monotonic() - started + + def cold_prompt_chat(model: str) -> httpx.Response: + return candidate.request( + "POST", + "/v1/chat/completions", + {"model": model, "messages": [{"role": "user", "content": marker}], "prompt_variables": {"word": marker}}, + ) + + control: Final = timed(lambda: cold_prompt_chat(steady)) + assert control.value.status_code == 200, control.value.text + call: Final = timed(lambda: cold_prompt_chat(flaky)) + response: Final = call.value assert response.status_code == 200, response.text - assert elapsed < 5, f"a retried cold prompt miss took {elapsed:.1f}s" + assert_cut_short(call.seconds, injected_delay=retry_after) + assert_within_factor_of_control(call.seconds, control=control.seconds, factor=retried_within_control_factor) attempts: Final = tuple( target for target in seen_prompt_gets if target.startswith(PROMPTS_PATH + encoded_flaky_prompt) ) @@ -1855,12 +1867,8 @@ def test_audit_surviving_worker_keeps_serving_after_kill(gateway: Gateway, tmp_p ): candidate: Final = owned.gateway model: Final = scenario.model(api_base=provider.url + "/v1", api_key=provider_secret) - workers: Final = tuple( - member for member in group_members(owned.process.pid) if member.pid != owned.process.pid - ) + workers: Final = owned.worker_pids() assert workers, "expected worker processes under the owned proxy" - workers[0].send_signal(signal.SIGKILL) - psutil.wait_procs([workers[0]], timeout=5) def fire(index: int) -> tuple[str, httpx.Response]: caller_trace: Final = uuid.uuid4().hex @@ -1878,7 +1886,8 @@ def test_audit_surviving_worker_keeps_serving_after_kill(gateway: Gateway, tmp_p ) return caller_session, response - with ThreadPoolExecutor(max_workers=10) as pool: + with owned.replacing_worker(min(workers)) as kill, ThreadPoolExecutor(max_workers=10) as pool: + kill() answered: Final = tuple(pool.map(fire, range(10))) received: Final[list[Request]] = [] # mutable-ok: drain() consumes the queue, later polls keep earlier ones for caller_session, response in answered: diff --git a/tests/integration/observability/test_langtrace_delivery.py b/tests/integration/observability/test_langtrace_delivery.py index 03e42087dc7..93006cb6c6b 100644 --- a/tests/integration/observability/test_langtrace_delivery.py +++ b/tests/integration/observability/test_langtrace_delivery.py @@ -1,7 +1,7 @@ import asyncio import json import re -import signal +import threading import time import uuid from collections.abc import Callable, Iterator, Sequence @@ -14,12 +14,12 @@ from queue import SimpleQueue from typing import Final import httpx -import psutil import pytest import yaml from anthropic import Anthropic from integration._support.client import Gateway, eventually from integration._support.process import OwnedProxy, owned_proxy, owned_proxy_process +from integration._support.timing import assert_not_delayed_by, timed from integration._support.wire import Reply, Request, Wire, wire_server from openai import AsyncOpenAI, OpenAI from opentelemetry.proto.collector.trace.v1.trace_service_pb2 import ExportTraceServiceRequest @@ -627,15 +627,21 @@ def test_langtrace_sink_outage_mid_burst_recovers_on_the_same_port(gateway: Gate def test_langtrace_slow_sink_does_not_delay_callers_or_duplicate_spans(gateway: Gateway, tmp_path: Path) -> None: + sink_delay: Final = 3 + slowed: Final = threading.Event() + def slow(request: Request) -> Reply: - time.sleep(1) + if slowed.is_set(): + time.sleep(sink_delay) return _accepted(request) with _langtrace_rig(gateway, tmp_path, respond=slow) as rig: - started: Final = time.monotonic() - markers: Final = _burst(rig, 6) - assert time.monotonic() - started < 5 - _assert_each_once(rig.sink, markers, seconds=40) + control: Final = timed(lambda: _burst(rig, 6)) + _assert_each_once(rig.sink, control.value) + slowed.set() + observed: Final = timed(lambda: _burst(rig, 6)) + assert_not_delayed_by((observed.seconds,), control=(control.seconds,), injected_delay=sink_delay) + _assert_each_once(rig.sink, observed.value, seconds=40) def test_langtrace_survives_a_killed_worker(gateway: Gateway, tmp_path: Path) -> None: @@ -659,25 +665,10 @@ def test_langtrace_survives_a_killed_worker(gateway: Gateway, tmp_path: Path) -> _assert_kill_and_recovery(owned, rig) -def _cmdline(process: psutil.Process) -> str: - try: - return " ".join(process.cmdline()) - except psutil.Error: - return "" - - def _assert_kill_and_recovery(owned: OwnedProxy, rig: _Rig) -> None: _assert_delivered(rig, _Surface(_chat_httpx, False), _marker()) - def uvicorn_workers() -> tuple[psutil.Process, ...]: - return tuple(child for child in psutil.Process(owned.process.pid).children() if "spawn_main" in _cmdline(child)) - - workers: Final = uvicorn_workers() + workers: Final = owned.worker_pids() assert len(workers) == 2, workers - workers[0].send_signal(signal.SIGKILL) - eventually( - uvicorn_workers, - lambda value: len(value) == 2 and workers[0].pid not in {child.pid for child in value}, - seconds=30, - ) + owned.kill_worker(min(workers)) _assert_each_once(rig.sink, _burst(rig, 6)) diff --git a/tests/integration/observability/test_logging_only_scope_chaos.py b/tests/integration/observability/test_logging_only_scope_chaos.py index a34b13f048f..afcaab91b5f 100644 --- a/tests/integration/observability/test_logging_only_scope_chaos.py +++ b/tests/integration/observability/test_logging_only_scope_chaos.py @@ -1,6 +1,5 @@ from __future__ import annotations -import signal import socket import threading import time @@ -45,10 +44,13 @@ from _logging_only_scope_support import ( from anthropic import APIConnectionError as AnthropicAPIConnectionError from integration._support.client import Gateway, eventually from integration._support.process import owned_proxy, owned_proxy_process +from integration._support.timing import Timed, assert_not_delayed_by, timed from integration._support.wire import Reply, Request from openai import APIConnectionError as OpenAIAPIConnectionError from pydantic import JsonValue +_EDGE_DELAY_SECONDS: Final = 5 + @pytest.mark.timeout(180) def test_K1_policy_edge_restart_mid_burst_keeps_output_observation_fail_open(gateway: Gateway, tmp_path: Path) -> None: @@ -216,7 +218,7 @@ def test_K2_policy_edge_delay_does_not_delay_concurrent_callers(gateway: Gateway marker: Final = uuid.uuid4().hex def policy(_request: Request) -> Reply: - time.sleep(2) + time.sleep(_EDGE_DELAY_SECONDS) return Reply(body=b'{"action":"NONE"}') with gateway.scenario() as scenario: @@ -227,18 +229,23 @@ def test_K2_policy_edge_delay_does_not_delay_concurrent_callers(gateway: Gateway expected_directions: Final = _directions_for_audit_leg(("request", "response"), "output") control_config: Final = _chaos_control_configuration(tmp_path, identity, model_list) with owned_proxy(gateway, tmp_path, {}, config=control_config, workers=2) as control_proxy: - controls: Final = tuple( - _call_client( - call.client_kind, - call.endpoint, - control_proxy, - call.model, - call.prompt, - call.stream, - f"{call.call_id}-control", + + def run_control(call: ChaosCall) -> Timed[CallerResult]: + return timed( + lambda: _call_client( + call.client_kind, + call.endpoint, + control_proxy, + call.model, + call.prompt, + call.stream, + f"{call.call_id}-control", + ) ) - for call in chaos_calls - ) + + with ThreadPoolExecutor(max_workers=30) as control_pool: + timed_controls: Final = tuple(control_pool.map(run_control, chaos_calls)) + controls: Final = tuple(control.value for control in timed_controls) control_upstream: Final = _drain_upstream(gateway.upstream_url) assert len(control_upstream) == 30, control_upstream assert ( @@ -270,7 +277,11 @@ def test_K2_policy_edge_delay_does_not_delay_concurrent_callers(gateway: Gateway assert tuple(index for index, _, _ in results) == tuple(range(30)), results assert all(result.status == controls[index].status for index, result, _ in results), results assert all(result.text == controls[index].text for index, result, _ in results), results - assert all(duration < 2 for _, _, duration in results), results + assert_not_delayed_by( + (seconds for _, _, seconds in results), + control=(control.seconds for control in timed_controls), + injected_delay=_EDGE_DELAY_SECONDS, + ) eventually( lambda: edge.received.qsize(), lambda count: count == 30 * len(expected_directions), @@ -364,16 +375,7 @@ def test_K3_two_worker_sigkill_checks_post_kill_spend_rows( config: Final = _configuration(tmp_path, identity, edge.url, "output", model_list=model_list) with owned_proxy_process(gateway, tmp_path, {}, config=config, workers=2) as owned: candidate: Final = owned.gateway - root: Final = psutil.Process(owned.process.pid) - workers: Final = eventually( - lambda: tuple( - child - for child in root.children(recursive=True) - if any("spawn_main" in part for part in child.cmdline()) - ), - lambda children: len(children) == 2, - seconds=30, - ) + workers: Final = tuple(sorted(eventually(owned.worker_pids, lambda pids: len(pids) == 2, seconds=30))) def run(call: ChaosCall) -> tuple[int, CallerResult | None, str | None]: try: @@ -390,17 +392,12 @@ def test_K3_two_worker_sigkill_checks_post_kill_spend_rows( except (OpenAIAPIConnectionError, AnthropicAPIConnectionError, httpx.RemoteProtocolError) as error: return call.index, None, str(error) - with ThreadPoolExecutor(max_workers=30) as pool: + with owned.replacing_worker(workers[0]) as kill, ThreadPoolExecutor(max_workers=30) as pool: futures: Final = tuple(pool.submit(run, call) for call in calls) try: assert eventually(lambda: scan_started.is_set(), bool, seconds=30) eventually(lambda: edge.received.qsize(), lambda count: count >= 5, seconds=30) - workers[0].send_signal(signal.SIGKILL) - killed_workers, surviving_workers = psutil.wait_procs((workers[0],), timeout=10) - assert len(killed_workers) == 1 and not surviving_workers, ( - killed_workers, - surviving_workers, - ) + kill() finally: release_scans.set() outcomes: Final = tuple(future.result(timeout=90) for future in futures) diff --git a/tests/integration/observability/test_otel_conversation_id.py b/tests/integration/observability/test_otel_conversation_id.py index e77e5603148..6d0e5ea197d 100644 --- a/tests/integration/observability/test_otel_conversation_id.py +++ b/tests/integration/observability/test_otel_conversation_id.py @@ -1,7 +1,6 @@ import asyncio import base64 import json -import os import re import signal import threading @@ -16,12 +15,11 @@ from typing import Final import anthropic import httpx import openai -import psutil import pytest import yaml from integration._support.client import Gateway, eventually, gateway_from_environment from integration._support.database import read_rows -from integration._support.process import OwnedProxy, owned_proxy_process +from integration._support.process import OwnedProxy, owned_proxy_process, stop_root_process from integration._support.wire import Reply, Request, Wire, wire_server from pydantic import JsonValue @@ -762,29 +760,26 @@ def test_slow_sink_during_a_burst_lands_every_response_exactly_once(rig: Rig) -> @pytest.mark.covers("other.observability.otel.conversation_id_survives_worker_kill") def test_killing_one_of_two_workers_mid_burst_keeps_serving_and_never_duplicates_a_span(two_worker_rig: Rig) -> None: rig: Final = two_worker_rig - root: Final = psutil.Process(rig.process.process.pid) - workers: Final = eventually( - lambda: tuple(child for child in root.children() if "resource_tracker" not in " ".join(child.cmdline())), - lambda found: len(found) == 2, - seconds=30, - ) + workers: Final = eventually(rig.process.worker_pids, lambda found: len(found) == 2, seconds=30) sessions: Final = {index: f"conv-{index}-{uuid.uuid4().hex}" for index in range(24)} markers: Final = tuple(_marker() for _ in range(24)) - def one(index: int) -> tuple[str, str | None]: - if index == 8: - os.kill(workers[0].pid, signal.SIGKILL) - try: - response: Final = rig.chat(markers[index], headers={"x-litellm-session-id": sessions[index]}) - return f"chatcmpl-{markers[index]}", None if response.status_code == 200 else response.text - except httpx.HTTPError as error: - return f"chatcmpl-{markers[index]}", repr(error) + with rig.process.replacing_worker(min(workers)) as kill: - with ThreadPoolExecutor(max_workers=6) as pool: - results: Final = tuple(pool.map(one, range(24))) - assert rig.process.process.poll() is None, "Proxy root exited after a worker was killed" - after: Final = rig.chat(_marker(), headers={"x-litellm-session-id": "conv-after-kill"}) - assert after.status_code == 200, after.text + def one(index: int) -> tuple[str, str | None]: + if index == 8: + kill() + try: + response: Final = rig.chat(markers[index], headers={"x-litellm-session-id": sessions[index]}) + return f"chatcmpl-{markers[index]}", None if response.status_code == 200 else response.text + except httpx.HTTPError as error: + return f"chatcmpl-{markers[index]}", repr(error) + + with ThreadPoolExecutor(max_workers=6) as pool: + results: Final = tuple(pool.map(one, range(24))) + assert rig.process.process.poll() is None, "Proxy root exited after a worker was killed" + after: Final = rig.chat(_marker(), headers={"x-litellm-session-id": "conv-after-kill"}) + assert after.status_code == 200, after.text assert rig.sink.single_span(after.json()["id"]) == "conv-after-kill" failures: Final = tuple(error for _, error in results if error) assert all(error.startswith(("ReadError(", "RemoteProtocolError(", "ConnectError(")) for error in failures), ( @@ -823,8 +818,8 @@ def test_terminating_the_proxy_right_after_a_burst_flushes_every_span_before_exi lambda: _landed(rig, expected), lambda seen: all(len(values) >= 1 for values in seen.values()), seconds=60 ) assert drained == {session: (identity,) for identity, session in expected.items()}, drained - rig.process.process.terminate() - assert rig.process.process.wait(timeout=40) in (0, -signal.SIGTERM) + assert stop_root_process(rig.process.process), "Proxy did not exit within the graceful stop budget" + assert rig.process.process.returncode in (0, -signal.SIGTERM) assert _landed(rig, expected) == drained with pytest.raises(httpx.ConnectError): next(started) diff --git a/tests/integration/observability/test_otel_excluded_services_matrix.py b/tests/integration/observability/test_otel_excluded_services_matrix.py index 7bca7ffc0dc..07646a0b189 100644 --- a/tests/integration/observability/test_otel_excluded_services_matrix.py +++ b/tests/integration/observability/test_otel_excluded_services_matrix.py @@ -1,8 +1,6 @@ import asyncio import json -import os import re -import signal import uuid from collections.abc import Callable, Generator, Iterator, Mapping from concurrent.futures import ThreadPoolExecutor @@ -15,7 +13,6 @@ from typing import Final, Literal, Protocol, cast import anthropic import httpx import openai -import psutil import pytest import yaml from integration._support.client import Gateway, Scenario, eventually, gateway_from_environment, object_value @@ -687,25 +684,22 @@ def test_stalled_tenant_destination_during_a_burst_does_not_block_responses(rig: @pytest.mark.timeout(300) def test_killing_one_of_two_workers_mid_burst_keeps_the_filter_on_the_survivor(rig: Rig) -> None: - root: Final = psutil.Process(rig.owned.process.pid) - workers: Final = eventually( - lambda: tuple(child for child in root.children() if "resource_tracker" not in " ".join(child.cmdline())), - lambda found: len(found) == 2, - seconds=30, - ) + workers: Final = eventually(rig.owned.worker_pids, lambda found: len(found) == 2, seconds=30) cursors: Final = rig.cursors() - def one(index: int) -> Sent | str: - if index == 6: - os.kill(workers[0].pid, signal.SIGKILL) - try: - return rig.raw("chat", _marker(), stream=index % 2 == 0) - except (httpx.HTTPError, AssertionError) as error: - return repr(error) + with rig.owned.replacing_worker(min(workers)) as kill: - with ThreadPoolExecutor(max_workers=6) as pool: - results: Final = tuple(pool.map(one, range(18))) - assert rig.owned.process.poll() is None, "Proxy root exited after a worker was killed" + def one(index: int) -> Sent | str: + if index == 6: + kill() + try: + return rig.raw("chat", _marker(), stream=index % 2 == 0) + except (httpx.HTTPError, AssertionError) as error: + return repr(error) + + with ThreadPoolExecutor(max_workers=6) as pool: + results: Final = tuple(pool.map(one, range(18))) + assert rig.owned.process.poll() is None, "Proxy root exited after a worker was killed" failures: Final = tuple(result for result in results if isinstance(result, str)) assert all(failure.startswith(("ReadError(", "RemoteProtocolError(", "ConnectError(")) for failure in failures), ( failures @@ -918,25 +912,22 @@ def test_unconfigured_tenant_outage_during_a_mixed_burst(unconfigured_rig: Rig) @pytest.mark.timeout(300) def test_unconfigured_killing_one_of_two_workers_keeps_the_fan_out(unconfigured_rig: Rig) -> None: - root: Final = psutil.Process(unconfigured_rig.owned.process.pid) - workers: Final = eventually( - lambda: tuple(child for child in root.children() if "resource_tracker" not in " ".join(child.cmdline())), - lambda found: len(found) == 2, - seconds=30, - ) + workers: Final = eventually(unconfigured_rig.owned.worker_pids, lambda found: len(found) == 2, seconds=30) cursors: Final = unconfigured_rig.cursors() - def one(index: int) -> Sent | str: - if index == 6: - os.kill(workers[0].pid, signal.SIGKILL) - try: - return unconfigured_rig.raw("chat", _marker(), stream=index % 2 == 0) - except (httpx.HTTPError, AssertionError) as error: - return repr(error) + with unconfigured_rig.owned.replacing_worker(min(workers)) as kill: - with ThreadPoolExecutor(max_workers=6) as pool: - results: Final = tuple(pool.map(one, range(18))) - assert unconfigured_rig.owned.process.poll() is None, "Proxy root exited after a worker was killed" + def one(index: int) -> Sent | str: + if index == 6: + kill() + try: + return unconfigured_rig.raw("chat", _marker(), stream=index % 2 == 0) + except (httpx.HTTPError, AssertionError) as error: + return repr(error) + + with ThreadPoolExecutor(max_workers=6) as pool: + results: Final = tuple(pool.map(one, range(18))) + assert unconfigured_rig.owned.process.poll() is None, "Proxy root exited after a worker was killed" failures: Final = tuple(result for result in results if isinstance(result, str)) assert all(failure.startswith(("ReadError(", "RemoteProtocolError(", "ConnectError(")) for failure in failures), ( failures diff --git a/tests/integration/observability/test_passthrough_upstream_error_chaos.py b/tests/integration/observability/test_passthrough_upstream_error_chaos.py index 014942ccfaf..b4a8524d3f8 100644 --- a/tests/integration/observability/test_passthrough_upstream_error_chaos.py +++ b/tests/integration/observability/test_passthrough_upstream_error_chaos.py @@ -1,10 +1,8 @@ import asyncio import json import re -import signal import threading -from collections.abc import Callable, Iterator -from contextlib import contextmanager +from collections.abc import Callable from dataclasses import dataclass from pathlib import Path from typing import Final @@ -15,7 +13,7 @@ import pytest import yaml from integration._support.client import Gateway, eventually, object_value from integration._support.database import read_rows -from integration._support.process import graceful_stop_seconds, owned_proxy_process +from integration._support.process import WorkerHealthcheck, graceful_stop_seconds, owned_proxy_process from integration._support.wire import Reply, Request, wire_server from pydantic import JsonValue @@ -60,16 +58,6 @@ def _held_connections(pid: int, port: int) -> int: ) -@contextmanager -def _suspended_worker(pid: int) -> Iterator[None]: - worker: Final = psutil.Process(pid) - worker.suspend() - try: - yield - finally: - worker.resume() - - def _error_information(call_id: str) -> dict[str, JsonValue]: rows: Final = eventually( lambda: read_rows('SELECT metadata FROM "LiteLLM_SpendLogs" WHERE request_id=%s', (call_id,)), @@ -166,21 +154,19 @@ async def test_passthrough_worker_sigkill_leaves_sibling_serving_and_logging(gat with wire_server(_held_reply(release)) as wire: config["environment_variables"] = {"GEMINI_API_BASE": wire.url, "GEMINI_API_KEY": "scripted"} path.write_text(yaml.safe_dump(config)) - with owned_proxy_process(gateway, tmp_path, {}, config=path, workers=2) as owned: + with owned_proxy_process( + gateway, tmp_path, {}, config=path, workers=2, healthcheck=WorkerHealthcheck.LENIENT + ) as owned: candidate: Final = owned.gateway - workers: Final = eventually( - lambda: tuple(int(pid) for pid in _STARTED_WORKER.findall(owned.log.read_text())), - lambda pids: len(pids) == 2, - seconds=30, - ) + workers: Final = tuple(sorted(eventually(owned.worker_pids, lambda pids: len(pids) == 2, seconds=30))) async with asyncio.TaskGroup() as tasks: try: - with _suspended_worker(workers[0]): + with owned.suspended(workers[0]): seed: Final = tasks.create_task( _fire_burst(str(candidate.client.base_url), candidate.key, 1, tolerate_transport_errors=True) ) await asyncio.to_thread(eventually, lambda: wire.received.qsize(), lambda size: size >= 1, 30) - with _suspended_worker(workers[1]): + with owned.suspended(workers[1]): rest: Final = tasks.create_task( _fire_burst( str(candidate.client.base_url), @@ -194,16 +180,18 @@ async def test_passthrough_worker_sigkill_leaves_sibling_serving_and_logging(gat peer_port: Final = int(wire.url.rsplit(":", 1)[1]) held_by: Final = {pid: _held_connections(pid, peer_port) for pid in workers} assert sum(held_by.values()) == 20 and min(held_by.values()) > 0, held_by - victim: Final = psutil.Process(max(workers, key=held_by.__getitem__)) - victim.suspend() - victim_ports: Final = frozenset( - connection.raddr.port for connection in victim.net_connections(kind="tcp") if connection.raddr - ) - victim.send_signal(signal.SIGKILL) - victim.wait(timeout=5) + with owned.replacing_worker(max(workers, key=held_by.__getitem__)) as kill: + victim: Final = kill.freeze() + victim_ports: Final = frozenset( + connection.raddr.port + for connection in victim.net_connections(kind="tcp") + if connection.raddr + ) + kill() + release.set() + served: Final = (*await seed, *await rest) finally: release.set() - served: Final = (*await seed, *await rest) for item in served: assert item.response.status_code in (200, 404, 500, 502), item.response.status_code follow_up: Final = candidate.request( diff --git a/tests/integration/observability/test_presidio_streaming_output.py b/tests/integration/observability/test_presidio_streaming_output.py index 680617f1652..d784fecd146 100644 --- a/tests/integration/observability/test_presidio_streaming_output.py +++ b/tests/integration/observability/test_presidio_streaming_output.py @@ -1,6 +1,5 @@ import json import re -import signal import threading import uuid from collections.abc import Callable, Iterator, Mapping @@ -10,16 +9,17 @@ from dataclasses import dataclass from pathlib import Path from typing import Final -import psutil +import httpx import yaml from integration._support.client import Gateway, eventually -from integration._support.process import OwnedProxy, group_members, owned_proxy_process +from integration._support.process import OwnedProxy, owned_proxy_process from integration._support.wire import Reply, Request, Wire, wire_server from openai import OpenAI from pydantic import BaseModel PERSON: Final = "John Smith" MASK: Final = "" +CUT_BY_WORKER_DEATH: Final = 0 GEMINI_MODEL: Final = "gemini-2.5-flash" @@ -170,6 +170,12 @@ class Rig: ) as response: return Received(response.status_code, tuple(response.iter_raw())) + def stream_on_a_worker_that_may_die(self, path: str, body: dict[str, object]) -> Received: + try: + return self.stream(path, body) + except httpx.TransportError: + return Received(CUT_BY_WORKER_DEATH, ()) + def gemini_path(self) -> str: return f"/v1beta/models/{self.gemini}:streamGenerateContent?alt=sse" @@ -567,18 +573,13 @@ def test_native_gemini_keeps_streaming_after_one_worker_is_killed(gateway: Gatew frames: Final = (gemini_frame("alive "), gemini_frame("still.")) provider: Final = gemini_provider(Reply(content_type="text/event-stream", chunks=frames, pause_between_chunks=0.5)) with presidio_rig(gateway, tmp_path, provider) as rig: - workers: Final = eventually( - lambda: tuple( - member for member in group_members(rig.proxy.process.pid) if member.pid != rig.proxy.process.pid - ), - lambda members: len(members) >= 2, - seconds=30, - ) - victim: Final = workers[0] - with ThreadPoolExecutor(max_workers=8) as pool: - futures: Final = tuple(pool.submit(rig.stream, rig.gemini_path(), rig.gemini_body()) for _ in range(8)) - victim.send_signal(signal.SIGKILL) - psutil.wait_procs((victim,), timeout=10) + workers: Final = eventually(rig.proxy.worker_pids, lambda members: len(members) >= 2, seconds=30) + with rig.proxy.replacing_worker(min(workers)) as kill, ThreadPoolExecutor(max_workers=8) as pool: + futures: Final = tuple( + pool.submit(rig.stream_on_a_worker_that_may_die, rig.gemini_path(), rig.gemini_body()) + for _ in range(8) + ) + kill() first_wave: Final = tuple(future.result() for future in futures) survivors: Final = tuple(received for received in first_wave if received.status == 200) assert survivors, [received.text[:200] for received in first_wave] diff --git a/tests/integration/observability/test_prometheus_series_cap_chaos.py b/tests/integration/observability/test_prometheus_series_cap_chaos.py index c2f165ca37e..c62c5b4fb0f 100644 --- a/tests/integration/observability/test_prometheus_series_cap_chaos.py +++ b/tests/integration/observability/test_prometheus_series_cap_chaos.py @@ -6,7 +6,6 @@ from __future__ import annotations import json import re -import signal import threading from collections.abc import Iterator, Mapping, Sequence from concurrent.futures import ThreadPoolExecutor @@ -19,7 +18,6 @@ from types import MappingProxyType from typing import Final import httpx -import psutil import pytest from integration._support.client import eventually, object_value, string_value from integration._support.process import owned_gateway_image, setup_only_proxy_run @@ -255,10 +253,10 @@ def test_killed_worker_is_replaced_by_one_that_reads_the_same_admissions(tmp_pat plans: Final = _plans(extra) with ThreadPoolExecutor(max_workers=1) as pool: burst: Final = pool.submit(_burst, rig, plans, True) - victim: Final = psutil.Process(workers[0]) - victim.suspend() - victim.send_signal(signal.SIGKILL) - served: Final = burst.result() + with rig.proxy.replacing_worker(workers[0]) as kill: + kill.freeze() + kill() + served: Final = burst.result() answered: Final = tuple(item for item in served if item.status == 200) assert answered, [(item.status, item.text[:200]) for item in served] assert all(item.call.answer in item.text for item in answered) diff --git a/tests/integration/observability/test_s3_v2_partition_granularity.py b/tests/integration/observability/test_s3_v2_partition_granularity.py index a362ff55245..382463eb46c 100644 --- a/tests/integration/observability/test_s3_v2_partition_granularity.py +++ b/tests/integration/observability/test_s3_v2_partition_granularity.py @@ -13,7 +13,6 @@ from urllib.parse import quote, unquote import httpx import openai -import psutil import pytest import yaml from _s3_v2_support import ( @@ -31,7 +30,7 @@ from _s3_v2_support import ( from integration._support.client import Gateway, JsonValue, Scenario, eventually, object_value from integration._support.database import read_rows, scratch_database from integration._support.database_relay import database_relay -from integration._support.process import OwnedProxy, group_members, owned_proxy_process +from integration._support.process import OwnedProxy, owned_proxy_process from integration._support.wire import Reply, Request, wire_server FLUSH: Final = {"DEFAULT_S3_FLUSH_INTERVAL_SECONDS": "1"} @@ -1105,14 +1104,6 @@ def test_s3_v2_hour_slow_sink_batches_never_duplicate_an_upload(gateway: Gateway assert _batches_outside_layout(objects, "hour") == () -def _worker_processes(owned: OwnedProxy) -> tuple[int, ...]: - return tuple( - process.pid - for process in group_members(owned.process.pid) - if process.pid != owned.process.pid and "spawn_main" in " ".join(process.cmdline()) - ) - - def test_s3_v2_hour_worker_kill_mid_burst_keeps_the_other_worker_logging(gateway: Gateway, tmp_path: Path) -> None: marker: Final = "s3hkill" + uuid.uuid4().hex[:8] upstream: Final = CountingUpstream() @@ -1124,7 +1115,8 @@ def test_s3_v2_hour_worker_kill_mid_burst_keeps_the_other_worker_logging(gateway owned.gateway.scenario() as scenario, ): openai_model, _, key = _models(scenario, provider.url) - workers: Final = _worker_processes(owned) + workers: Final = owned.worker_pids() + assert len(workers) == 2, workers sent: Final = tuple(f"{marker}-{index}" for index in range(40)) def send(prompt: str) -> tuple[str, bool]: @@ -1143,13 +1135,14 @@ def test_s3_v2_hour_worker_kill_mid_burst_keeps_the_other_worker_logging(gateway return prompt, False return prompt, response.status_code == 200 and response.json()["id"] == prompt - with ThreadPoolExecutor(max_workers=16) as pool: - futures: Final = tuple(pool.submit(send, prompt) for prompt in sent) - eventually(lambda: len(upstream.received()), lambda count: count >= 8, seconds=30) - psutil.Process(workers[0]).kill() - results: Final = tuple(future.result() for future in futures) - later: Final = tuple(f"{marker}-later-{index}" for index in range(8)) - later_results: Final = tuple(send(prompt) for prompt in later) + with owned.replacing_worker(min(workers)) as kill: + with ThreadPoolExecutor(max_workers=16) as pool: + futures: Final = tuple(pool.submit(send, prompt) for prompt in sent) + eventually(lambda: len(upstream.received()), lambda count: count >= 8, seconds=30) + kill() + results: Final = tuple(future.result() for future in futures) + later: Final = tuple(f"{marker}-later-{index}" for index in range(8)) + later_results: Final = tuple(send(prompt) for prompt in later) eventually( lambda: frozenset(str(payload["id"]) for payload in sink.payloads()), lambda landed: frozenset(later) <= landed, @@ -1157,7 +1150,6 @@ def test_s3_v2_hour_worker_kill_mid_burst_keeps_the_other_worker_logging(gateway ) payloads: Final = sink.payloads() objects: Final = sink.objects() - assert len(workers) == 2, workers assert all(ok for _, ok in later_results), "the surviving worker must keep serving after the kill" landed: Final = tuple(str(payload["id"]) for payload in payloads) assert frozenset(landed) <= frozenset((*sent, *later)), "only ids this test sent may land" @@ -1215,7 +1207,7 @@ def test_s3_v2_hour_proxy_restart_mid_burst_keeps_the_layout_without_duplicates( with ThreadPoolExecutor(max_workers=16) as pool: futures: Final = tuple(pool.submit(send, first.gateway, first_key, prompt) for prompt in sent) eventually(lambda: len(upstream.received()), lambda count: count >= 8, seconds=30) - first.process.terminate() + first.begin_shutdown() results: Final = tuple(future.result() for future in futures) first.process.wait(timeout=30) answered: Final = frozenset(prompt for prompt, ok in results if ok) diff --git a/tests/integration/observability/test_s3_v2_upload_fanout.py b/tests/integration/observability/test_s3_v2_upload_fanout.py index 0154cce10df..ead1e5b028b 100644 --- a/tests/integration/observability/test_s3_v2_upload_fanout.py +++ b/tests/integration/observability/test_s3_v2_upload_fanout.py @@ -15,7 +15,7 @@ import yaml from _s3_v2_support import RecordingS3Sink, collect_payloads from _s3_v2_support import s3_config as _recording_s3_config from integration._support.client import Gateway, JsonValue, eventually -from integration._support.process import group_members, owned_proxy, owned_proxy_process +from integration._support.process import owned_proxy, owned_proxy_process from integration._support.wire import Reply, Request, Wire, wire_server BUCKET: Final = "integration-bucket" @@ -814,7 +814,8 @@ def test_s3_v2_worker_kill_mid_burst_keeps_surviving_deliveries(gateway: Gateway ): model: Final = scenario.model(api_base=provider.url + "/v1", api_key="synthetic-provider-key") key: Final = scenario.key(models=[model]) - sent: Final = tuple(f"{marker}-{index}" for index in range(REQUESTS)) + burst: Final = tuple(f"{marker}-burst-{index}" for index in range(REQUESTS)) + after: Final = tuple(f"{marker}-after-{index}" for index in range(REQUESTS // 4)) def send(identity: str) -> tuple[str, bool]: try: @@ -832,23 +833,28 @@ def test_s3_v2_worker_kill_mid_burst_keeps_surviving_deliveries(gateway: Gateway return identity, False return identity, response.status_code == 200 - with ThreadPoolExecutor(max_workers=32) as pool: - futures: Final = tuple(pool.submit(send, identity) for identity in sent) - time.sleep(0.5) - children: Final = tuple( - process for process in group_members(owned.process.pid) if process.pid != owned.process.pid - ) - assert children, "no worker children found to kill" - children[0].kill() - results: Final = tuple(future.result() for future in futures) - survivors: Final = frozenset(identity for identity, ok in results if ok) - assert survivors, "no request survived the worker kill" + workers: Final = owned.worker_pids() + assert workers, "no worker children found to kill" + with owned.replacing_worker(min(workers)) as kill: + with ThreadPoolExecutor(max_workers=32) as pool: + futures: Final = tuple(pool.submit(send, identity) for identity in burst) + time.sleep(0.5) + kill() + results: Final = tuple(future.result() for future in futures) readiness: Final = owned.gateway.client.get("/health/readiness") assert readiness.status_code == 200, readiness.text - payloads: Final = collect_payloads(sink, len(survivors), seconds=90) - landed: Final = frozenset(payload["id"] for payload in payloads) - assert survivors <= landed, "an id whose response succeeded never landed" - assert landed <= frozenset(sent), "an id that was never sent landed" + with ThreadPoolExecutor(max_workers=8) as pool: + after_results: Final = tuple(pool.map(send, after)) + burst_survivors: Final = frozenset(identity for identity, ok in results if ok) + after_survivors: Final = frozenset(identity for identity, ok in after_results if ok) + assert burst_survivors, "no request survived the worker kill" + assert after_survivors == frozenset(after), "the proxy did not serve after the worker was replaced" + landed: Final = eventually( + lambda: frozenset(payload["id"] for payload in sink.payloads()), + lambda ids: after_survivors <= ids and bool(ids & burst_survivors), + seconds=90, + ) + assert landed <= frozenset(burst) | frozenset(after), "an id that was never sent landed" @pytest.mark.covers("other.observability.s3_v2.sigterm_mid_burst_loses_only_inflight_without_duplicates") @@ -895,7 +901,7 @@ def test_s3_v2_sigterm_mid_burst_loses_only_inflight_without_duplicates(gateway: with ThreadPoolExecutor(max_workers=32) as pool: futures: Final = tuple(pool.submit(send, identity) for identity in sent) time.sleep(0.5) - candidate_owned.process.terminate() + candidate_owned.begin_shutdown() results: Final = tuple(future.result() for future in futures) candidate_owned.process.wait(timeout=30) finally: diff --git a/tests/integration/observability/test_signoz_delivery.py b/tests/integration/observability/test_signoz_delivery.py index 56a98b3f17c..5d4b66bbdd6 100644 --- a/tests/integration/observability/test_signoz_delivery.py +++ b/tests/integration/observability/test_signoz_delivery.py @@ -1,7 +1,6 @@ import asyncio import base64 import json -import os import re import signal import threading @@ -17,12 +16,11 @@ from typing import Final import anthropic import httpx import openai -import psutil import pytest import yaml from integration._support.client import Gateway, eventually, gateway_from_environment, object_value, string_value from integration._support.database import read_rows -from integration._support.process import OwnedProxy, owned_proxy_process +from integration._support.process import OwnedProxy, owned_proxy_process, stop_root_process from integration._support.wire import Reply, Request, Wire, wire_server from opentelemetry.proto.collector.trace.v1.trace_service_pb2 import ExportTraceServiceRequest from opentelemetry.proto.common.v1.common_pb2 import AnyValue @@ -915,28 +913,25 @@ def test_slow_sink_during_a_burst_lands_every_response_exactly_once(rig: Rig) -> def test_killing_one_of_two_workers_mid_burst_keeps_serving_and_never_duplicates_a_span(rig: Rig) -> None: - root: Final = psutil.Process(rig.process.process.pid) - workers: Final = eventually( - lambda: tuple(child for child in root.children() if "resource_tracker" not in " ".join(child.cmdline())), - lambda found: len(found) == 2, - seconds=30, - ) + workers: Final = eventually(rig.process.worker_pids, lambda found: len(found) == 2, seconds=30) markers: Final = tuple(_marker() for _ in range(24)) - def one(index: int) -> tuple[str, str | None]: - if index == 8: - os.kill(workers[0].pid, signal.SIGKILL) - try: - response: Final = rig.chat(markers[index]) - return f"chatcmpl-{markers[index]}", None if response.status_code == 200 else response.text - except httpx.HTTPError as error: - return f"chatcmpl-{markers[index]}", repr(error) + with rig.process.replacing_worker(min(workers)) as kill: - with ThreadPoolExecutor(max_workers=6) as pool: - results: Final = tuple(pool.map(one, range(24))) - assert rig.process.process.poll() is None, "Proxy root exited after a worker was killed" - after: Final = rig.chat(_marker()) - assert after.status_code == 200, after.text + def one(index: int) -> tuple[str, str | None]: + if index == 8: + kill() + try: + response: Final = rig.chat(markers[index]) + return f"chatcmpl-{markers[index]}", None if response.status_code == 200 else response.text + except httpx.HTTPError as error: + return f"chatcmpl-{markers[index]}", repr(error) + + with ThreadPoolExecutor(max_workers=6) as pool: + results: Final = tuple(pool.map(one, range(24))) + assert rig.process.process.poll() is None, "Proxy root exited after a worker was killed" + after: Final = rig.chat(_marker()) + assert after.status_code == 200, after.text rig.sink.single_span(_body_id(after)) failures: Final = tuple(error for _, error in results if error) assert all(error.startswith(("ReadError(", "RemoteProtocolError(", "ConnectError(")) for error in failures), ( @@ -969,8 +964,8 @@ def test_terminating_the_proxy_right_after_a_burst_flushes_every_span_before_exi assert all(response.status_code == 200 for response in responses), [response.text for response in responses] identities: Final = tuple(_body_id(response) for response in responses) pending_at_signal: Final = rig.sink.landed(identities) - rig.process.process.terminate() - assert rig.process.process.wait(timeout=40) in (0, -signal.SIGTERM) + assert stop_root_process(rig.process.process), "Proxy did not exit within the graceful stop budget" + assert rig.process.process.returncode in (0, -signal.SIGTERM) with pytest.raises(httpx.ConnectError): next(started) landed: Final = rig.sink.landed(identities) diff --git a/tests/integration/observability/test_straiker_v3_platform.py b/tests/integration/observability/test_straiker_v3_platform.py index 45ae562f5f7..a419d343535 100644 --- a/tests/integration/observability/test_straiker_v3_platform.py +++ b/tests/integration/observability/test_straiker_v3_platform.py @@ -10,8 +10,6 @@ from __future__ import annotations import hashlib import itertools import json -import os -import signal import socket import threading import uuid @@ -30,7 +28,7 @@ import pytest import yaml from integration._support.client import Gateway, eventually, gateway_from_environment, object_value from integration._support.database import read_rows -from integration._support.process import OwnedProxy, owned_proxy, owned_proxy_process +from integration._support.process import OwnedProxy, WorkerHealthcheck, owned_proxy, owned_proxy_process from integration._support.wire import Reply, Request, wire_server V3_KEY: Final = "sk_agt_synthetic_integration_key" @@ -436,7 +434,9 @@ def rig(tmp_path_factory: pytest.TempPathFactory) -> Iterator[Rig]: with gateway_from_environment() as gateway, wire_server(_provider) as provider: config: Final = _rig_config(sink.url, root) with ( - owned_proxy_process(gateway, root, {}, config=config, workers=2) as owned, + owned_proxy_process( + gateway, root, {}, config=config, workers=2, healthcheck=WorkerHealthcheck.LENIENT + ) as owned, owned.gateway.scenario() as scenario, ): chat: Final = scenario.model( @@ -1070,31 +1070,17 @@ def test_burst_with_platform_outage_recovers_without_duplicate_spend(rig: Rig) - # C2: one proxy worker is killed during a burst; the other keeps serving and detect still runs for each call -def _is_live_worker(child: psutil.Process, exclude: int) -> bool: - if child.pid == exclude: - return False - try: - return child.status() != psutil.STATUS_ZOMBIE and "spawn_main" in " ".join(child.cmdline()) - except psutil.Error: - return False - - -def _uvicorn_workers(parent: psutil.Process, *, exclude: int = 0) -> tuple[psutil.Process, ...]: - return tuple(c for c in parent.children() if _is_live_worker(c, exclude)) - - def test_an_unreaped_worker_does_not_hide_the_surviving_proxy_worker(rig: Rig) -> None: - parent: Final = psutil.Process(rig.owned.process.pid) - workers: Final = eventually(lambda: _uvicorn_workers(parent), lambda children: len(children) == 2) + workers: Final = tuple(sorted(eventually(rig.owned.worker_pids, lambda pids: len(pids) == 2))) victim: Final = workers[0] survivor: Final = workers[1] marker: Final = rig.marker() - os.kill(parent.pid, signal.SIGSTOP) - try: - os.kill(victim.pid, signal.SIGKILL) - assert eventually(victim.status, lambda status: status == psutil.STATUS_ZOMBIE) == psutil.STATUS_ZOMBIE - assert tuple(child.pid for child in _uvicorn_workers(parent)) == (survivor.pid,) - assert _uvicorn_workers(parent, exclude=survivor.pid) == () + with rig.owned.replacing_worker(victim) as kill, rig.owned.suspended(rig.owned.process.pid): + kill() + assert eventually(psutil.Process(victim).status, lambda status: status == psutil.STATUS_ZOMBIE) == ( + psutil.STATUS_ZOMBIE + ) + assert rig.owned.worker_pids() == frozenset((survivor,)) with httpx.Client(base_url=rig._base(), timeout=15, trust_env=False) as client: response: Final = client.post( "/v1/chat/completions", @@ -1103,15 +1089,11 @@ def test_an_unreaped_worker_does_not_hide_the_surviving_proxy_worker(rig: Rig) - ) assert response.status_code == 200, response.text assert len(_v3_request_calls(rig, marker)) == 1 - finally: - os.kill(parent.pid, signal.SIGCONT) - assert len(eventually(lambda: _uvicorn_workers(parent, exclude=victim.pid), lambda children: len(children) == 2)) == 2 + assert len(rig.owned.worker_pids() - frozenset((victim,))) == 2 def test_burst_survives_one_worker_kill(rig: Rig) -> None: - parent: Final = psutil.Process(rig.owned.process.pid) - workers: Final = eventually(lambda: _uvicorn_workers(parent), lambda c: len(c) >= 2) - victim: Final = workers[0].pid + workers: Final = tuple(sorted(eventually(rig.owned.worker_pids, lambda pids: len(pids) == 2))) markers: Final = tuple(rig.marker() for _ in range(24)) def fresh_chat(text: str) -> tuple[int, str]: @@ -1126,18 +1108,19 @@ def test_burst_survives_one_worker_kill(rig: Rig) -> None: return 0, repr(error) return response.status_code, response.text - def call(index: int) -> tuple[int, str]: - if index == 6: - os.kill(victim, signal.SIGKILL) - return fresh_chat("kill " + markers[index]) + with rig.owned.replacing_worker(workers[0]) as kill: - with ThreadPoolExecutor(max_workers=4) as pool: - results: Final = tuple(pool.map(call, range(24))) + def call(index: int) -> tuple[int, str]: + if index == 6: + kill() + return fresh_chat("kill " + markers[index]) + + with ThreadPoolExecutor(max_workers=4) as pool: + results: Final = tuple(pool.map(call, range(24))) ok: Final = tuple(i for i, (status, _) in enumerate(results) if status == 200) assert len(ok) >= 20, results for index in ok: assert len(_v3_request_calls(rig, markers[index])) >= 1, markers[index] - eventually(lambda: _uvicorn_workers(parent, exclude=victim), lambda c: len(c) >= 2) after: Final = fresh_chat("after-kill " + rig.marker()) assert after[0] == 200, after diff --git a/tests/integration/providers/test_anthropic_thinking_signature_stream_wire.py b/tests/integration/providers/test_anthropic_thinking_signature_stream_wire.py index cf46885b6b5..fc0234410bb 100644 --- a/tests/integration/providers/test_anthropic_thinking_signature_stream_wire.py +++ b/tests/integration/providers/test_anthropic_thinking_signature_stream_wire.py @@ -1,7 +1,6 @@ import asyncio import json import re -import signal import threading import uuid from collections.abc import Mapping, Sequence @@ -763,19 +762,19 @@ async def test_worker_sigkill_mid_burst_leaves_the_sibling_streaming_signed_thin held_by: Final = MappingProxyType({pid: _open_upstream_connections(pid, wire.url) for pid in workers}) assert sum(held_by.values()) == 20, held_by victim_pid, survivor_pid = sorted(workers, key=held_by.__getitem__) - victim: Final = psutil.Process(victim_pid) - victim.suspend() - victim.send_signal(signal.SIGKILL) - release.set() - served: Final = await burst - assert held_by[survivor_pid] >= 10, held_by - completed: Final = tuple(item for item in served if item.status == 200) - assert len(completed) == held_by[survivor_pid], (held_by, [item.status for item in served]) - for item in completed: - if item.call.stream: - _assert_signed_once(deltas_of(chunks_of(item.text)), item.call.marker) - else: - assert answer(item.call.marker) in item.text, item.text + with owned.replacing_worker(victim_pid) as kill: + kill.freeze() + await asyncio.to_thread(kill) + release.set() + served: Final = await burst + assert held_by[survivor_pid] >= 10, held_by + completed: Final = tuple(item for item in served if item.status == 200) + assert len(completed) == held_by[survivor_pid], (held_by, [item.status for item in served]) + for item in completed: + if item.call.stream: + _assert_signed_once(deltas_of(chunks_of(item.text)), item.call.marker) + else: + assert answer(item.call.marker) in item.text, item.text follow_up: Final = _Call(endpoint="chat", stream=True, marker=uuid.uuid4().hex) (answered,) = await _burst(str(candidate.client.base_url), candidate.key, _CONFIG_MODEL, (follow_up,)) assert answered.status == 200, answered.text diff --git a/tests/integration/providers/test_assistants_route_errors_chaos.py b/tests/integration/providers/test_assistants_route_errors_chaos.py index 5e863932050..045c1ed852a 100644 --- a/tests/integration/providers/test_assistants_route_errors_chaos.py +++ b/tests/integration/providers/test_assistants_route_errors_chaos.py @@ -1,39 +1,36 @@ import asyncio -import signal import threading +from contextlib import AsyncExitStack from pathlib import Path from queue import SimpleQueue from types import MappingProxyType from typing import Final from urllib.parse import urlsplit -import psutil import pytest from integration._support.client import Gateway, eventually -from integration._support.process import graceful_stop_seconds, owned_proxy_process +from integration._support.process import WorkerHealthcheck, graceful_stop_seconds, owned_proxy_process from integration.providers._assistants_route_errors_support import ( MAPPED_BY_STATUS, MARKED_ROUTES, ROUTE_BY_NAME, Call, + Held, answer_all, assert_answered_with_its_own_marker, assistants_wire, held_error_peer, held_one_by_one, held_upstream_connections, - live_worker_pids, marker_of, new_marker, - open_upstream_connections, openai_assistants_config, provider_requests, worker_pids, ) _STATUSES: Final = (400, 404, 429, 500, 503) -_MIN_CALLS: Final = 20 -_MAX_CALLS: Final = 60 +_PER_WORKER: Final = 10 _STARTUP_COMPLETE: Final = "Application startup complete." @@ -45,7 +42,9 @@ def _startups(log: Path) -> tuple[int, int]: async def test_worker_sigkill_mid_burst_leaves_the_sibling_answering_each_mapped_status( gateway: Gateway, tmp_path: Path ) -> None: - calls: Final = tuple(Call(MARKED_ROUTES[index % len(MARKED_ROUTES)], new_marker()) for index in range(_MAX_CALLS)) + calls: Final = tuple( + Call(MARKED_ROUTES[index % len(MARKED_ROUTES)], new_marker()) for index in range(2 * _PER_WORKER) + ) follow_up: Final = Call(ROUTE_BY_NAME["get_thread"], new_marker()) status_by_marker: Final = MappingProxyType( { @@ -59,25 +58,30 @@ async def test_worker_sigkill_mid_burst_leaves_the_sibling_answering_each_mapped port: Final = urlsplit(wire.url).port assert port is not None config: Final = openai_assistants_config(tmp_path, port, 60) - with owned_proxy_process(gateway, tmp_path, {}, config=config, workers=2) as owned: + with owned_proxy_process( + gateway, tmp_path, {}, config=config, workers=2, healthcheck=WorkerHealthcheck.LENIENT + ) as owned: eventually(lambda: _startups(owned.log), lambda found: found == (2, 2), seconds=120) - workers: Final = live_worker_pids(owned.log) - assert len(workers) == 2, workers - - def both_workers_hold_enough() -> bool: - held: Final = tuple(open_upstream_connections(pid, port) for pid in workers) - return sum(held) >= _MIN_CALLS and min(held) > 0 - + victim_pid, survivor_pid = sorted(eventually(owned.worker_pids, lambda pids: len(pids) == 2)) try: - async with held_one_by_one(owned.gateway, calls, arrived, both_workers_hold_enough) as held: - held_by: Final = await asyncio.to_thread(held_upstream_connections, workers, port, len(held.calls)) - assert min(held_by.values()) > 0, held_by - victim_pid, survivor_pid = sorted(workers, key=held_by.__getitem__) - victim: Final = psutil.Process(victim_pid) - victim.suspend() - victim.send_signal(signal.SIGKILL) - release.set() - served: Final = await held.answers() + async with AsyncExitStack() as bursts: + with owned.suspended(survivor_pid): + on_victim: Final = await bursts.enter_async_context( + held_one_by_one(owned.gateway, calls[:_PER_WORKER], arrived, lambda: False) + ) + with owned.suspended(victim_pid): + on_survivor: Final = await bursts.enter_async_context( + held_one_by_one(owned.gateway, calls[_PER_WORKER:], arrived, lambda: False) + ) + held: Final = Held(on_victim.calls + on_survivor.calls, on_victim.pending + on_survivor.pending) + held_by: Final = await asyncio.to_thread( + held_upstream_connections, (victim_pid, survivor_pid), port, len(held.calls) + ) + assert held_by == {victim_pid: _PER_WORKER, survivor_pid: _PER_WORKER}, held_by + with owned.replacing_worker(victim_pid) as kill: + await asyncio.to_thread(kill) + release.set() + served: Final = await held.answers() finally: release.set() (answered,) = await answer_all(owned.gateway, (follow_up,)) diff --git a/tests/integration/providers/test_assistants_route_errors_wire.py b/tests/integration/providers/test_assistants_route_errors_wire.py index ec6d58e6670..395a99a8afe 100644 --- a/tests/integration/providers/test_assistants_route_errors_wire.py +++ b/tests/integration/providers/test_assistants_route_errors_wire.py @@ -1,9 +1,9 @@ import asyncio import json import threading +from contextlib import AsyncExitStack from collections.abc import Iterator, Mapping from dataclasses import dataclass -from pathlib import Path from queue import SimpleQueue from types import MappingProxyType from typing import Final @@ -11,8 +11,8 @@ from typing import Final import httpx import openai import pytest -from integration._support.client import Gateway, gateway_from_environment -from integration._support.process import graceful_stop_seconds, owned_proxy_process +from integration._support.client import Gateway, eventually, gateway_from_environment +from integration._support.process import OwnedProxy, WorkerHealthcheck, graceful_stop_seconds, owned_proxy_process from integration._support.wire import Reply, Request from integration.providers._assistants_route_errors_support import ( ASSISTANT_MODEL, @@ -28,6 +28,7 @@ from integration.providers._assistants_route_errors_support import ( STATUS_TABLE, UNMAPPED_500, Call, + Held, Mapped, Route, answer_all, @@ -42,7 +43,6 @@ from integration.providers._assistants_route_errors_support import ( held_error_peer, held_one_by_one, held_upstream_connections, - live_worker_pids, new_marker, openai_assistants_config, owned_config, @@ -71,15 +71,19 @@ _AZURE_CREDENTIALS: Final = ( _PRIVATE_DETAILS: Final = ("10.20.30.40", "/etc/litellm/secrets/db.yaml", "sk-proj-" + "a1B2c3D4" * 5) _URL_MODEL: Final = "http://169.254.169.254/latest" _BURST_STATUSES: Final = (400, 401, 404, 429, 500, 503) +_BURST_PER_WORKER: Final = 24 _SDK_ERRORS: Final = MappingProxyType({404: openai.NotFoundError, 429: openai.RateLimitError}) _BODY_VARIANTS: Final = ("text_502", "empty_503", "hostile_400", "null_message_404", "garbage_200") @dataclass(frozen=True, slots=True) class _Scripted: - gateway: Gateway + owned: OwnedProxy port: int - log: Path + + @property + def gateway(self) -> Gateway: + return self.owned.gateway @pytest.fixture(scope="module") @@ -89,9 +93,11 @@ def scripted(tmp_path_factory: pytest.TempPathFactory) -> Iterator[_Scripted]: config: Final = openai_assistants_config(directory, port, _DEPLOYMENT_TIMEOUT_SECONDS) with ( gateway_from_environment() as shared, - owned_proxy_process(shared, directory, {}, config=config, workers=2) as owned, + owned_proxy_process( + shared, directory, {}, config=config, workers=2, healthcheck=WorkerHealthcheck.LENIENT + ) as owned, ): - yield _Scripted(owned.gateway, port, owned.log) + yield _Scripted(owned, port) @pytest.fixture(scope="module") @@ -372,7 +378,9 @@ def test_repeated_failing_request_answers_the_same_each_time(scripted: _Scripted async def test_mixed_status_burst_answers_each_request_with_its_own_status_and_marker(scripted: _Scripted) -> None: - calls: Final = tuple(Call(MARKED_ROUTES[index % len(MARKED_ROUTES)], new_marker()) for index in range(48)) + calls: Final = tuple( + Call(MARKED_ROUTES[index % len(MARKED_ROUTES)], new_marker()) for index in range(2 * _BURST_PER_WORKER) + ) status_by_marker: Final = MappingProxyType( { call.marker: _BURST_STATUSES[(index // len(MARKED_ROUTES)) % len(_BURST_STATUSES)] @@ -381,18 +389,29 @@ async def test_mixed_status_burst_answers_each_request_with_its_own_status_and_m ) release: Final = threading.Event() arrived: Final[SimpleQueue[str]] = SimpleQueue() - workers: Final = live_worker_pids(scripted.log) - assert len(workers) == 2, workers + owned: Final = scripted.owned + first, second = sorted(eventually(owned.worker_pids, lambda pids: len(pids) == 2)) with assistants_wire(held_error_peer(status_by_marker, arrived, release), port=scripted.port) as wire: try: - async with held_one_by_one(scripted.gateway, calls, arrived, lambda: False) as held: - held_by: Final = await asyncio.to_thread(held_upstream_connections, workers, scripted.port, len(calls)) + async with AsyncExitStack() as bursts: + with owned.suspended(second): + on_first: Final = await bursts.enter_async_context( + held_one_by_one(owned.gateway, calls[:_BURST_PER_WORKER], arrived, lambda: False) + ) + with owned.suspended(first): + on_second: Final = await bursts.enter_async_context( + held_one_by_one(owned.gateway, calls[_BURST_PER_WORKER:], arrived, lambda: False) + ) + held: Final = Held(on_first.calls + on_second.calls, on_first.pending + on_second.pending) + held_by: Final = await asyncio.to_thread( + held_upstream_connections, (first, second), scripted.port, len(calls) + ) release.set() answers: Final = await held.answers() finally: release.set() received: Final = wire.drain() - assert sum(held_by.values()) == len(calls) and min(held_by.values()) > 0, held_by + assert held_by == {first: _BURST_PER_WORKER, second: _BURST_PER_WORKER}, held_by assert len(answers) == len(calls) for answer in answers: assert_answered_with_its_own_marker(answer, MAPPED_BY_STATUS[status_by_marker[answer.call.marker]]) diff --git a/tests/integration/providers/test_bedrock_converse_lookaround_regex_chaos.py b/tests/integration/providers/test_bedrock_converse_lookaround_regex_chaos.py index b92360a2193..2031be5f3be 100644 --- a/tests/integration/providers/test_bedrock_converse_lookaround_regex_chaos.py +++ b/tests/integration/providers/test_bedrock_converse_lookaround_regex_chaos.py @@ -1,7 +1,6 @@ import asyncio import json import re -import signal import threading import uuid from collections.abc import Mapping @@ -331,15 +330,15 @@ async def test_worker_sigkill_mid_burst_leaves_the_sibling_cleaning_schemas(gate held_by: Final = MappingProxyType({pid: _open_upstream_connections(pid, wire.url) for pid in workers}) assert sum(held_by.values()) == 20, held_by victim_pid, survivor_pid = sorted(workers, key=held_by.__getitem__) - victim: Final = psutil.Process(victim_pid) - victim.suspend() - victim.send_signal(signal.SIGKILL) - release.set() - served: Final = await burst - assert held_by[survivor_pid] >= 10, held_by - assert len(served) == held_by[survivor_pid], (held_by, len(served)) - for item in served: - _assert_answered_with_its_own_marker(item) + with owned.replacing_worker(victim_pid) as kill: + kill.freeze() + await asyncio.to_thread(kill) + release.set() + served: Final = await burst + assert held_by[survivor_pid] >= 10, held_by + assert len(served) == held_by[survivor_pid], (held_by, len(served)) + for item in served: + _assert_answered_with_its_own_marker(item) follow_up: Final = _Call(model=fleet.kimi_bare, endpoint="chat", stream=False, marker=uuid.uuid4().hex) (answered,) = await _burst(str(candidate.client.base_url), candidate.key, (follow_up,)) _assert_answered_with_its_own_marker(answered) diff --git a/tests/integration/providers/test_bedrock_converse_lookaround_regex_wire.py b/tests/integration/providers/test_bedrock_converse_lookaround_regex_wire.py index 97b64c663bb..428f0015766 100644 --- a/tests/integration/providers/test_bedrock_converse_lookaround_regex_wire.py +++ b/tests/integration/providers/test_bedrock_converse_lookaround_regex_wire.py @@ -10,6 +10,7 @@ import httpx import openai import pytest from integration._support.client import Gateway, Scenario, eventually, object_value, string_value +from integration._support.timing import assert_within_factor_of_control, timed from integration._support.upstream import _aws_event_frame from integration._support.wire import Reply, Request, Wire, wire_server from pydantic import JsonValue, TypeAdapter @@ -765,6 +766,10 @@ def test_the_worst_case_lookaround_input_scans_in_linear_time(gateway: Gateway) "type": "object", "properties": {"collection": {"type": "string", "pattern": pattern}}, } + benign: Final[dict[str, JsonValue]] = { + "type": "object", + "properties": {"collection": {"type": "string", "pattern": "a" * len(pattern)}}, + } liveliness: Final[list[tuple[float, int]]] = [] stop: Final = threading.Event() @@ -774,16 +779,24 @@ def test_the_worst_case_lookaround_input_scans_in_linear_time(gateway: Gateway) with wire_server(_bedrock_peer) as wire, gateway.scenario() as scenario: model: Final = _deployment(scenario, wire, f"bedrock/{_KIMI}") + control: Final = timed( + lambda: gateway.request("POST", _path("chat"), _body("chat", model, (_openai_tool(_TOOL, benign),))) + ) + _assert_tool_call_relayed("chat", control.value) + assert len(wire.drain()) == 1 poller: Final = threading.Thread(target=poll) poller.start() - started: Final = time.perf_counter() - response: Final = gateway.request("POST", _path("chat"), _body("chat", model, (_openai_tool(_TOOL, schema),))) - elapsed: Final = time.perf_counter() - started + worst: Final = timed( + lambda: gateway.request("POST", _path("chat"), _body("chat", model, (_openai_tool(_TOOL, schema),))) + ) stop.set() poller.join() - _assert_tool_call_relayed("chat", response) - assert elapsed < 30, elapsed - assert liveliness and max(latency for latency, _ in liveliness) < 5, liveliness + _assert_tool_call_relayed("chat", worst.value) + assert_within_factor_of_control(worst.seconds, control=control.seconds, factor=10) + assert liveliness, liveliness + assert_within_factor_of_control( + max(latency for latency, _ in liveliness), control=control.seconds, factor=10 + ) assert {status for _, status in liveliness} == {200}, liveliness assert len(wire.drain()) == 1 diff --git a/tests/integration/providers/test_bedrock_converse_missing_content_wire.py b/tests/integration/providers/test_bedrock_converse_missing_content_wire.py index 057596941a2..b0a19bbe817 100644 --- a/tests/integration/providers/test_bedrock_converse_missing_content_wire.py +++ b/tests/integration/providers/test_bedrock_converse_missing_content_wire.py @@ -3,7 +3,6 @@ import base64 import json import os import re -import signal import threading import uuid from collections.abc import Iterable, Iterator, Mapping, Sequence @@ -904,13 +903,14 @@ async def test_c02_worker_sigkill_mid_burst_leaves_the_sibling_serving_content_l held_by: Final = MappingProxyType({pid: _open_peer_connections(pid, wire.url) for pid in workers}) assert sum(held_by.values()) == 20, held_by victim_pid, survivor_pid = sorted(workers, key=held_by.__getitem__) - psutil.Process(victim_pid).send_signal(signal.SIGKILL) - release.set() - served: Final = await burst - assert held_by[survivor_pid] >= 10, held_by - assert len(served) == held_by[survivor_pid], (held_by, len(served)) - for item in served: - assert item.status == 200 and _ANSWER in item.text, (item.call, item.status, item.text) + with owned.replacing_worker(victim_pid) as kill: + await asyncio.to_thread(kill) + release.set() + served: Final = await burst + assert held_by[survivor_pid] >= 10, held_by + assert len(served) == held_by[survivor_pid], (held_by, len(served)) + for item in served: + assert item.status == 200 and _ANSWER in item.text, (item.call, item.status, item.text) eventually( lambda: len(_STARTED_WORKER.findall(owned.log.read_text())), lambda count: count == 3, seconds=60 ) @@ -947,7 +947,7 @@ async def test_c03_proxy_terminated_mid_burst_lands_every_answered_content_less_ _burst(_proxy_url(candidate), candidate.key, _PLAIN, calls, tolerate_transport_errors=True) ) await asyncio.to_thread(eventually, held_indexes.qsize, lambda size: size == 12, 60) - owned.process.terminate() + owned.begin_shutdown() release.set() served: Final = await burst eventually(owned.process.poll, lambda code: code is not None, seconds=60) diff --git a/tests/integration/providers/test_bedrock_converse_stream_event_frames_wire.py b/tests/integration/providers/test_bedrock_converse_stream_event_frames_wire.py index bf4f611107d..df566ea69e4 100644 --- a/tests/integration/providers/test_bedrock_converse_stream_event_frames_wire.py +++ b/tests/integration/providers/test_bedrock_converse_stream_event_frames_wire.py @@ -3,7 +3,6 @@ import base64 import json import os import re -import signal import threading import uuid from collections.abc import Callable, Iterable, Mapping @@ -873,17 +872,18 @@ async def test_c02_worker_sigkill_mid_burst_leaves_the_sibling_rejecting_the_hel held_by: Final = MappingProxyType({pid: _open_peer_connections(pid, wire.url) for pid in workers}) assert sum(held_by.values()) == 20, held_by victim_pid, survivor_pid = sorted(workers, key=held_by.__getitem__) - psutil.Process(victim_pid).send_signal(signal.SIGKILL) - release.set() - served: Final = await burst - assert held_by[survivor_pid] >= 10, held_by - assert len(served) == held_by[survivor_pid], (held_by, len(served)) - for item in served: - assert item.status in (200, 400) and _is_rejection(item, _REJECTION), ( - item.call, - item.status, - item.text, - ) + with owned.replacing_worker(victim_pid) as kill: + await asyncio.to_thread(kill) + release.set() + served: Final = await burst + assert held_by[survivor_pid] >= 10, held_by + assert len(served) == held_by[survivor_pid], (held_by, len(served)) + for item in served: + assert item.status in (200, 400) and _is_rejection(item, _REJECTION), ( + item.call, + item.status, + item.text, + ) eventually( lambda: len(_STARTED_WORKER.findall(owned.log.read_text())), lambda count: count == 3, seconds=60 ) @@ -921,7 +921,7 @@ async def test_c03_proxy_terminated_mid_burst_lands_every_served_rejection_at_mo _burst(_proxy_url(candidate), candidate.key, _PLAIN, calls, tolerate_transport_errors=True) ) await asyncio.to_thread(eventually, held_indexes.qsize, lambda size: size == 12, 60) - owned.process.terminate() + owned.begin_shutdown() release.set() served: Final = await burst eventually(owned.process.poll, lambda code: code is not None, seconds=60) diff --git a/tests/integration/providers/test_bedrock_kimi_k3_cache_point_wire.py b/tests/integration/providers/test_bedrock_kimi_k3_cache_point_wire.py index 9bf526bf3d1..db4e564b8e5 100644 --- a/tests/integration/providers/test_bedrock_kimi_k3_cache_point_wire.py +++ b/tests/integration/providers/test_bedrock_kimi_k3_cache_point_wire.py @@ -3,7 +3,6 @@ import base64 import json import os import re -import signal import socket import threading import time @@ -25,6 +24,7 @@ import yaml from integration._support.client import Gateway, Scenario, eventually, gateway_from_environment from integration._support.database import read_rows from integration._support.process import OwnedProxy, owned_proxy_process +from integration._support.timing import assert_within_factor_of_control, timed from integration._support.upstream import _aws_event_frame from integration._support.wire import Reply, Request, Wire, wire_server from pydantic import JsonValue, TypeAdapter @@ -1021,16 +1021,20 @@ def test_p02_a_plain_kimi_reply_is_priced_from_its_own_row(rig: _Rig) -> None: @pytest.mark.timeout(600) def test_s07_a_5kb_caller_model_under_the_wildcard_is_a_bounded_4xx_and_liveliness_stays_up(rig: _Rig) -> None: + control: Final = timed( + lambda: _send(rig.gateway, "chat", _chat_body(f"bedrock/{'a' * 8}", _prompt(), _SYSTEM_AND_USER)) + ) rig.wire.drain() - started: Final = time.monotonic() - outcome: Final = _send(rig.gateway, "chat", _chat_body(f"bedrock/{'a' * 5120}", _prompt(), _SYSTEM_AND_USER)) - elapsed: Final = time.monotonic() - started + oversized: Final = timed( + lambda: _send(rig.gateway, "chat", _chat_body(f"bedrock/{'a' * 5120}", _prompt(), _SYSTEM_AND_USER)) + ) + outcome: Final = oversized.value liveliness: Final = rig.gateway.client.get("/health/liveliness", timeout=5) assert liveliness.status_code == 200, liveliness.text assert outcome.status in (400, 404), (outcome.status, outcome.raw[:300]) error: Final = _JSON.validate_python(_JSON.validate_python(json.loads(outcome.raw))["error"]) assert "a" * 5120 in str(error["message"]), outcome.raw[:300] - assert elapsed < 10, elapsed + assert_within_factor_of_control(oversized.seconds, control=control.seconds, factor=20) assert rig.wire.drain() == () _failure_row(outcome.call_id) @@ -1494,14 +1498,15 @@ async def test_c02_a_worker_killed_mid_traffic_leaves_the_survivor_and_the_respa held_by: Final = MappingProxyType({pid: _open_peer_connections(pid, wire.url) for pid in workers}) assert sum(held_by.values()) == 20, held_by victim_pid, survivor_pid = sorted(workers, key=held_by.__getitem__) - psutil.Process(victim_pid).send_signal(signal.SIGKILL) - hold.set() - served: Final = await burst - assert len(served) == held_by[survivor_pid], (held_by, len(served)) - for item in served: - _assert_answered(item.outcome) - received: Final = _received(wire) - assert len(received) == 20 and all(item.cache_points == 0 for item in received), received + with owned.replacing_worker(victim_pid) as kill: + await asyncio.to_thread(kill) + hold.set() + served: Final = await burst + assert len(served) == held_by[survivor_pid], (held_by, len(served)) + for item in served: + _assert_answered(item.outcome) + received: Final = _received(wire) + assert len(received) == 20 and all(item.cache_points == 0 for item in received), received respawned: Final = eventually( lambda: tuple(int(pid) for pid in _STARTED_WORKER.findall(owned.log.read_text())), lambda pids: len(pids) == 3, @@ -1570,8 +1575,14 @@ def test_c03_a_restart_re_registers_a_stored_deployment_flag_from_the_database( @pytest.mark.timeout(900) async def test_c04_a_slow_peer_under_ten_concurrent_streams_completes_every_call_once(rig: _Rig) -> None: - rig.wire.drain() calls: Final = tuple(_Call("chat", index, True) for index in range(10)) + alone_started: Final = time.monotonic() + (alone,) = await _burst( + _proxy_url(rig.gateway), rig.gateway.key, "kimi-us", calls[:1], f"{_prompt()} {_SLOW_MARKER}" + ) + alone_seconds: Final = time.monotonic() - alone_started + _assert_answered(alone.outcome) + rig.wire.drain() started: Final = time.monotonic() served: Final = await _burst( _proxy_url(rig.gateway), rig.gateway.key, "kimi-us", calls, f"{_prompt()} {_SLOW_MARKER}" @@ -1582,5 +1593,5 @@ async def test_c04_a_slow_peer_under_ten_concurrent_streams_completes_every_call _assert_answered(item.outcome) received: Final = _received(rig.wire) assert len(received) == 10 and all(item.cache_points == 0 for item in received), received - assert elapsed < 30, elapsed + assert_within_factor_of_control(elapsed, control=alone_seconds, factor=len(calls)) _assert_rows_once(served) diff --git a/tests/integration/providers/test_bedrock_mantle_claude_chat_chaos.py b/tests/integration/providers/test_bedrock_mantle_claude_chat_chaos.py index d8a3306ecd5..9c7903497e3 100644 --- a/tests/integration/providers/test_bedrock_mantle_claude_chat_chaos.py +++ b/tests/integration/providers/test_bedrock_mantle_claude_chat_chaos.py @@ -1,7 +1,6 @@ import asyncio import json import re -import signal import threading import uuid from collections.abc import Callable @@ -322,15 +321,15 @@ async def test_worker_sigkill_mid_burst_leaves_the_sibling_serving_the_native_ro held_by: Final = MappingProxyType({pid: _open_upstream_connections(pid, wire.url) for pid in workers}) assert sum(held_by.values()) == 20, held_by victim_pid, survivor_pid = sorted(workers, key=held_by.__getitem__) - victim: Final = psutil.Process(victim_pid) - victim.suspend() - victim.send_signal(signal.SIGKILL) - release.set() - served: Final = await burst - assert held_by[survivor_pid] >= 10, held_by - assert len(served) == held_by[survivor_pid], (held_by, len(served)) - for item in served: - _assert_answered_with_its_own_marker(item) + with owned.replacing_worker(victim_pid) as kill: + kill.freeze() + await asyncio.to_thread(kill) + release.set() + served: Final = await burst + assert held_by[survivor_pid] >= 10, held_by + assert len(served) == held_by[survivor_pid], (held_by, len(served)) + for item in served: + _assert_answered_with_its_own_marker(item) follow_up: Final = _Call(endpoint="chat", stream=False, marker=uuid.uuid4().hex) (answered,) = await _burst(str(candidate.client.base_url), candidate.key, _CONFIG_MODEL, (follow_up,)) _assert_answered_with_its_own_marker(answered) diff --git a/tests/integration/providers/test_bedrock_mantle_count_tokens_wire.py b/tests/integration/providers/test_bedrock_mantle_count_tokens_wire.py index 933fc49c149..7f59b818cba 100644 --- a/tests/integration/providers/test_bedrock_mantle_count_tokens_wire.py +++ b/tests/integration/providers/test_bedrock_mantle_count_tokens_wire.py @@ -1,7 +1,6 @@ import asyncio import json import re -import signal import socket import threading import uuid @@ -1270,10 +1269,11 @@ def test_worker_sigkill_mid_burst_leaves_the_sibling_counting( shares: Final = {pid: _accepted_client_ports(pid, proxy_url.port or 0) & ports for pid in workers} assert sum(map(len, shares.values())) == len(clients), shares victim: Final = min((pid for pid in workers if shares[pid]), key=lambda pid: len(shares[pid])) - psutil.Process(victim).send_signal(signal.SIGKILL) - release.set() - for port, result in (future.result(timeout=60) for future in futures): - assert result == (None if port in shares[victim] else (200, _MANTLE_COUNT)), (port, result, shares) + with owned.replacing_worker(victim) as kill: + kill() + release.set() + for port, result in (future.result(timeout=60) for future in futures): + assert result == (None if port in shares[victim] else (200, _MANTLE_COUNT)), (port, result, shares) second_wave: Final = _clients(stack, str(proxy_url), 6) assert tuple(_counted_on(client, owned.gateway.key, body) for client in second_wave) == ( (200, _MANTLE_COUNT), diff --git a/tests/integration/providers/test_bedrock_mantle_output_limit_wire.py b/tests/integration/providers/test_bedrock_mantle_output_limit_wire.py index 7891522a26f..4c06a70070e 100644 --- a/tests/integration/providers/test_bedrock_mantle_output_limit_wire.py +++ b/tests/integration/providers/test_bedrock_mantle_output_limit_wire.py @@ -3,7 +3,6 @@ import base64 import json import os import re -import signal import threading import uuid from collections.abc import Callable, Mapping @@ -939,7 +938,7 @@ def _chaos_budget(boots: int, stops: int) -> float: return (boots + stops) * graceful_stop_seconds() + _CHAOS_TRAFFIC_SECONDS -@pytest.mark.timeout(_chaos_budget(boots=2, stops=1)) +@pytest.mark.timeout(_chaos_budget(boots=3, stops=1)) async def test_worker_sigkill_mid_burst_leaves_the_sibling_renaming_the_yaml_limit( gateway: Gateway, tmp_path: Path ) -> None: @@ -969,18 +968,18 @@ async def test_worker_sigkill_mid_burst_leaves_the_sibling_renaming_the_yaml_lim held_by: Final = MappingProxyType({pid: _open_upstream_connections(pid, wire.url) for pid in workers}) assert sum(held_by.values()) == 20, held_by victim_pid, survivor_pid = sorted(workers, key=held_by.__getitem__) - victim: Final = psutil.Process(victim_pid) - victim.suspend() - victim.send_signal(signal.SIGKILL) - release.set() - outcome: Final = await burst - assert held_by[survivor_pid] >= 10, held_by - assert (len(outcome.served), outcome.lost) == (held_by[survivor_pid], held_by[victim_pid]), ( - held_by, - outcome, - ) - for item in outcome.served: - _assert_answered_with_its_own_marker(item) + with owned.replacing_worker(victim_pid) as kill: + kill.freeze() + await asyncio.to_thread(kill) + release.set() + outcome: Final = await burst + assert held_by[survivor_pid] >= 10, held_by + assert (len(outcome.served), outcome.lost) == (held_by[survivor_pid], held_by[victim_pid]), ( + held_by, + outcome, + ) + for item in outcome.served: + _assert_answered_with_its_own_marker(item) follow_up: Final = _Call(endpoint="chat", stream=False, marker=uuid.uuid4().hex) (answered,) = ( await _burst(_proxy_url(candidate), candidate.key, _CONFIG_MODEL, (follow_up,), limit=None) @@ -1011,7 +1010,7 @@ async def test_proxy_restart_mid_burst_loses_only_in_flight_requests_and_renames ) ) await asyncio.to_thread(eventually, held_markers.qsize, lambda size: size == 20, 60) - first.process.terminate() + first.begin_shutdown() release.set() outcome: Final = await burst assert len(outcome.served) + outcome.lost == len(calls), outcome diff --git a/tests/integration/providers/test_bedrock_model_discovery_chaos.py b/tests/integration/providers/test_bedrock_model_discovery_chaos.py index f1358837b7e..34350e5846c 100644 --- a/tests/integration/providers/test_bedrock_model_discovery_chaos.py +++ b/tests/integration/providers/test_bedrock_model_discovery_chaos.py @@ -1,6 +1,5 @@ from __future__ import annotations -import signal import threading from collections.abc import Iterator from concurrent.futures import ThreadPoolExecutor @@ -32,7 +31,7 @@ from integration._support.bedrock_discovery import ( wait_for_replacement_worker, ) from integration._support.client import Gateway, eventually, gateway_from_environment -from integration._support.process import graceful_stop_seconds +from integration._support.process import OwnedProxy, graceful_stop_seconds from integration._support.wire import Reply pytestmark: Final = pytest.mark.timeout(2 * graceful_stop_seconds() + 120) @@ -43,6 +42,7 @@ class Rig: gateway: Gateway log: Path plane: ControlPlane + owned: OwnedProxy @pytest.fixture(scope="module") @@ -51,7 +51,7 @@ def rig(tmp_path_factory: pytest.TempPathFactory) -> Iterator[Rig]: directory: Final = tmp_path_factory.mktemp("bedrock_discovery_chaos") with control_plane(directory, hosts) as plane, gateway_from_environment() as parent: with discovery_proxy(parent, directory, plane, check_provider_endpoint=True, workers=WORKERS) as owned: - yield Rig(owned.gateway, owned.log, plane) + yield Rig(owned.gateway, owned.log, plane, owned) def _victim(pids: tuple[int, ...], plane_url: str) -> int | None: @@ -85,13 +85,14 @@ def test_worker_killed_mid_listing_leaves_the_sibling_serving_and_a_replacement_ seconds=LISTING_TIMEOUT_SECONDS, ) assert victim is not None - psutil.Process(victim).send_signal(signal.SIGKILL) - sibling: Final = next(pid for pid in original if pid != victim) - for _ in range(5): - assert rig.gateway.chat(CONTROL_MODEL)["choices"], "the sibling worker stopped answering" - assert psutil.pid_exists(sibling) - severed: Final = in_flight.exception(timeout=LISTING_TIMEOUT_SECONDS * 2) - assert severed is not None or in_flight.result().status_code >= 500 + with rig.owned.replacing_worker(victim) as kill: + kill() + sibling: Final = next(pid for pid in original if pid != victim) + for _ in range(5): + assert rig.gateway.chat(CONTROL_MODEL)["choices"], "the sibling worker stopped answering" + assert psutil.pid_exists(sibling) + severed: Final = in_flight.exception(timeout=LISTING_TIMEOUT_SECONDS * 2) + assert severed is not None or in_flight.result().status_code >= 500 wait_for_replacement_worker(rig.log, original) release.set() assert discovered(rig.gateway, catalog, marker) == catalog.invocable_ids() diff --git a/tests/integration/providers/test_bedrock_runtime_chat_completions_chaos.py b/tests/integration/providers/test_bedrock_runtime_chat_completions_chaos.py index 820d033135f..71b63786c27 100644 --- a/tests/integration/providers/test_bedrock_runtime_chat_completions_chaos.py +++ b/tests/integration/providers/test_bedrock_runtime_chat_completions_chaos.py @@ -7,7 +7,6 @@ import itertools import multiprocessing import os import re -import signal import socket import threading import uuid @@ -29,7 +28,7 @@ import yaml from integration._support.bedrock_runtime_peer import MARKER, marker_of, respond, serve_peer from integration._support.client import Gateway, eventually, gateway_from_environment, object_value from integration._support.database import read_rows -from integration._support.process import owned_proxy_process +from integration._support.process import kill_stray_process, owned_proxy_process from integration._support.wire import Reply, Request, wire_server from pydantic import JsonValue, TypeAdapter @@ -216,7 +215,7 @@ async def _burst_killing_the_peer_once_it_answered( await asyncio.to_thread(eventually, lambda: peer.received.value, lambda count: count == len(calls), 60) first: Final = [await finished for finished in itertools.islice(asyncio.as_completed(tasks), answered)] assert all(item.status == 200 for item in first), [(item.call.marker, item.status) for item in first] - peer.process.kill() + kill_stray_process(peer.process.pid) peer.process.join(timeout=10) return tuple(await asyncio.gather(*tasks)) @@ -254,7 +253,8 @@ def _child_peer(port: int, answer_first: int) -> Iterator[_ChildPeer]: eventually(lambda: _accepts_connections(port), bool, seconds=30) yield _ChildPeer(process=process, received=received, url=f"http://127.0.0.1:{port}") finally: - process.kill() + if process.is_alive(): + kill_stray_process(process.pid) process.join(timeout=10) assert not process.is_alive(), "Owned peer survived cleanup" @@ -428,15 +428,15 @@ async def test_worker_sigkill_mid_burst_leaves_the_sibling_serving(gateway: Gate held_by: Final = MappingProxyType({pid: _open_upstream_connections(pid, wire.url) for pid in workers}) assert sum(held_by.values()) == 20, held_by victim_pid, survivor_pid = sorted(workers, key=held_by.__getitem__) - victim: Final = psutil.Process(victim_pid) - victim.suspend() - victim.send_signal(signal.SIGKILL) - release.set() - served: Final = await burst - assert held_by[survivor_pid] >= 10, held_by - assert len(served) == held_by[survivor_pid], (held_by, len(served)) - for item in served: - _assert_answered_with_its_own_marker(item) + with owned.replacing_worker(victim_pid) as kill: + kill.freeze() + await asyncio.to_thread(kill) + release.set() + served: Final = await burst + assert held_by[survivor_pid] >= 10, held_by + assert len(served) == held_by[survivor_pid], (held_by, len(served)) + for item in served: + _assert_answered_with_its_own_marker(item) follow_up: Final = _Call(endpoint="chat", stream=False, marker=uuid.uuid4().hex) (answered,) = await _burst(str(candidate.client.base_url), candidate.key, _CONFIG_MODEL, (follow_up,)) _assert_answered_with_its_own_marker(answered) diff --git a/tests/integration/providers/test_bedrock_runtime_chat_completions_sad_wire.py b/tests/integration/providers/test_bedrock_runtime_chat_completions_sad_wire.py index 8d6193cd542..0860d79459e 100644 --- a/tests/integration/providers/test_bedrock_runtime_chat_completions_sad_wire.py +++ b/tests/integration/providers/test_bedrock_runtime_chat_completions_sad_wire.py @@ -17,6 +17,7 @@ from integration._support.bedrock_runtime_peer import answer, forwarded_effort, from integration._support.client import Gateway, Scenario, eventually, object_value, string_value from integration._support.database import read_rows from integration._support.process import owned_proxy_process +from integration._support.timing import assert_within_factor_of_control, timed from integration._support.wire import Request, Wire, wire_server from pydantic import JsonValue, TypeAdapter @@ -255,16 +256,21 @@ def test_model_group_info_lists_the_native_supported_params(gateway: Gateway) -> def test_thirty_thousand_digit_version_is_classified_quickly_and_served_by_converse(gateway: Gateway) -> None: marker: Final = uuid.uuid4().hex + control_marker: Final = uuid.uuid4().hex with wire_server(respond) as wire, gateway.scenario() as scenario, ThreadPoolExecutor(max_workers=1) as pool: + control_model: Final = _deployment(scenario, wire) model: Final = _deployment(scenario, wire, model=f"bedrock/{LONG_VERSION_GPT}") + control: Final = timed(lambda: _chat(gateway, control_model, control_marker)) + assert _content(control.value) == answer(control_marker), control.value.text + _only_request(wire, control_marker) liveliness: Final = pool.submit(_timed_liveliness, gateway) - started: Final = time.monotonic() - response: Final = _chat(gateway, model, marker) - elapsed: Final = time.monotonic() - started + served: Final = timed(lambda: _chat(gateway, model, marker)) + response: Final = served.value health_status, health_elapsed = liveliness.result() assert _content(response) == answer(marker), response.text - assert elapsed < 10, elapsed - assert (health_status, health_elapsed < 2) == (200, True), (health_status, health_elapsed) + assert_within_factor_of_control(served.seconds, control=control.seconds, factor=20) + assert health_status == 200, (health_status, health_elapsed) + assert_within_factor_of_control(health_elapsed, control=control.seconds, factor=20) request: Final = _only_request(wire, marker) assert (request.method, target_of(request)) == ("POST", f"/model/{LONG_VERSION_GPT}/converse"), request.target _assert_row(string_value(_payload(response)["id"]), model, "success") @@ -274,12 +280,14 @@ def test_bad_key_on_the_long_version_model_is_refused_before_any_route(gateway: marker: Final = uuid.uuid4().hex control_marker: Final = uuid.uuid4().hex with wire_server(respond) as wire, gateway.scenario() as scenario: + short_model: Final = _deployment(scenario, wire) model: Final = _deployment(scenario, wire, model=f"bedrock/{LONG_VERSION_GPT}") - started: Final = time.monotonic() - refused: Final = _chat(gateway, model, marker, key=BAD_KEY) - elapsed: Final = time.monotonic() - started + short_refusal: Final = timed(lambda: _chat(gateway, short_model, uuid.uuid4().hex, key=BAD_KEY)) + assert short_refusal.value.status_code == 401, short_refusal.value.text + long_refusal: Final = timed(lambda: _chat(gateway, model, marker, key=BAD_KEY)) + refused: Final = long_refusal.value assert refused.status_code == 401, refused.text - assert elapsed < 2, elapsed + assert_within_factor_of_control(long_refusal.seconds, control=short_refusal.seconds, factor=10) assert "Authentication Error" in _error_message(refused), refused.text refused_rows: Final = eventually( lambda: read_rows( diff --git a/tests/integration/providers/test_bedrock_stream_timeout_chaos.py b/tests/integration/providers/test_bedrock_stream_timeout_chaos.py index 2632b81b4dc..496b83e0b2e 100644 --- a/tests/integration/providers/test_bedrock_stream_timeout_chaos.py +++ b/tests/integration/providers/test_bedrock_stream_timeout_chaos.py @@ -1,7 +1,6 @@ import asyncio import json import re -import signal import threading import uuid from collections.abc import Iterator, Mapping @@ -353,14 +352,14 @@ async def test_x3_killing_one_worker_mid_burst_leaves_the_sibling_timing_its_str held_by: Final = {pid: _open_upstream_connections(pid, peer.wire.url) for pid in workers} assert sum(held_by.values()) == 20, held_by victim_pid, survivor_pid = sorted(workers, key=held_by.__getitem__) - victim: Final = psutil.Process(victim_pid) - victim.suspend() - victim.send_signal(signal.SIGKILL) - served: Final = await burst - assert held_by[survivor_pid] >= 10, held_by - assert len(served) == held_by[survivor_pid], (held_by, len(served)) - for item in served: - _assert_timed_out(item, _KILL_TIMEOUT_SECONDS) + with owned.replacing_worker(victim_pid) as kill: + kill.freeze() + await asyncio.to_thread(kill) + served: Final = await burst + assert held_by[survivor_pid] >= 10, held_by + assert len(served) == held_by[survivor_pid], (held_by, len(served)) + for item in served: + _assert_timed_out(item, _KILL_TIMEOUT_SECONDS) peer.hold.clear() _assert_answered(await _one(owned.gateway, _KILL_MODEL, _Call("chat"))) assert len(peer.wire.drain()) == 21 diff --git a/tests/integration/providers/test_cache_control_tool_call_marks_owned_proxy.py b/tests/integration/providers/test_cache_control_tool_call_marks_owned_proxy.py index c945eaff38e..bc0a6f31e74 100644 --- a/tests/integration/providers/test_cache_control_tool_call_marks_owned_proxy.py +++ b/tests/integration/providers/test_cache_control_tool_call_marks_owned_proxy.py @@ -1,7 +1,5 @@ import asyncio -import os import re -import signal import socket import threading from collections import Counter @@ -209,10 +207,12 @@ async def test_worker_sigkill_mid_burst_leaves_the_sibling_serving_capped_reques {pid: _held_upstream_connections(pid, wire.url) for pid in workers} ) victim: Final = max(workers, key=held_by.__getitem__) - os.kill(victim, signal.SIGKILL) + with owned.replacing_worker(victim) as kill: + await asyncio.to_thread(kill) + release.set() + served: Final = await burst finally: release.set() - served: Final = await burst during: Final = wire.drain() after: Final = await _fire(owned_url, owned.gateway.key) after_received: Final = wire.drain() @@ -277,7 +277,7 @@ async def test_proxy_restart_mid_burst_serves_capped_requests_after_the_reboot( ) try: await asyncio.to_thread(eventually, lambda: wire.received.qsize(), lambda size: size >= _BURST, 90) - first.process.terminate() + first.begin_shutdown() finally: release.set() served: Final = await burst diff --git a/tests/integration/providers/test_chatgpt_responses_caller_stream_chaos.py b/tests/integration/providers/test_chatgpt_responses_caller_stream_chaos.py index d1e199c949a..40b87ba1b8b 100644 --- a/tests/integration/providers/test_chatgpt_responses_caller_stream_chaos.py +++ b/tests/integration/providers/test_chatgpt_responses_caller_stream_chaos.py @@ -1,6 +1,5 @@ import asyncio import re -import signal import threading import uuid from collections.abc import Iterator, Mapping, Sequence @@ -304,15 +303,15 @@ async def test_worker_sigkill_mid_burst_leaves_the_sibling_answering_json(gatewa held_by: Final = MappingProxyType({pid: _open_upstream_connections(pid, wire.url) for pid in workers}) assert sum(held_by.values()) == 20, held_by victim_pid, survivor_pid = sorted(workers, key=held_by.__getitem__) - victim: Final = psutil.Process(victim_pid) - victim.suspend() - victim.send_signal(signal.SIGKILL) - release.set() - served: Final = await burst - assert held_by[survivor_pid] >= 10, held_by - assert len(served) == held_by[survivor_pid], (held_by, len(served)) - for item in served: - _assert_answered_in_its_own_shape(item) + with owned.replacing_worker(victim_pid) as kill: + kill.freeze() + await asyncio.to_thread(kill) + release.set() + served: Final = await burst + assert held_by[survivor_pid] >= 10, held_by + assert len(served) == held_by[survivor_pid], (held_by, len(served)) + for item in served: + _assert_answered_in_its_own_shape(item) follow_up: Final = _Call("responses", False, uuid.uuid4().hex) (answered,) = await _burst(candidate, _CONFIG_MODEL, (follow_up,)) _assert_answered_in_its_own_shape(answered) diff --git a/tests/integration/providers/test_databricks_decisions_wire.py b/tests/integration/providers/test_databricks_decisions_wire.py index 24db0153bcd..d3b17044207 100644 --- a/tests/integration/providers/test_databricks_decisions_wire.py +++ b/tests/integration/providers/test_databricks_decisions_wire.py @@ -10,6 +10,7 @@ import pytest from integration._support.client import Gateway, Scenario, eventually, object_value, string_value from integration._support.database import read_rows from integration._support.process import graceful_stop_seconds, owned_proxy_process +from integration._support.timing import assert_within_factor_of_control from integration._support.upstream import ScenarioHandle, delete_scenario, register_scenario from integration.cost_calculation.cost_tracking_case import JsonResponse from pydantic import JsonValue @@ -111,10 +112,28 @@ def _upstream_calls(gateway: Gateway, handle: ScenarioHandle) -> list[dict[str, return _calls_to(_observed(gateway), handle) -def _timed_decide(gateway: Gateway, model: str) -> tuple[httpx.Response, float]: - started: Final = time.perf_counter() - response: Final = _decide(gateway, model) - return response, time.perf_counter() - started +def _timed_decide_under_liveliness_probes( + gateway: Gateway, model: str +) -> tuple[httpx.Response, float, tuple[tuple[int, float], ...]]: + liveliness: Final[list[tuple[int, float]]] = [] + stop: Final = threading.Event() + + def poll() -> None: + while not stop.is_set(): + started: Final = time.perf_counter() + probe: Final = gateway.client.get("/health/liveliness") + liveliness.append((probe.status_code, time.perf_counter() - started)) + + poller: Final = threading.Thread(target=poll) + poller.start() + try: + started: Final = time.perf_counter() + response: Final = _decide(gateway, model) + elapsed: Final = time.perf_counter() - started + finally: + stop.set() + poller.join() + return response, elapsed, tuple(liveliness) def _spend_row(call_id: str) -> dict[str, JsonValue]: @@ -623,23 +642,12 @@ def test_a_megabyte_model_name_is_refused_in_linear_time_while_liveliness_stays_ "/v1/systemone", json={"model": model, "state": _STATE, "questions": _QUESTIONS} ) assert anonymous.status_code == 401, anonymous.text[:300] - tenth_response, tenth_elapsed = _timed_decide(gateway, f"{prefix}/{_TENTH_MEGABYTE_NAME}") - liveliness: Final[list[tuple[int, float]]] = [] - stop: Final = threading.Event() - - def poll() -> None: - while not stop.is_set(): - started: Final = time.perf_counter() - probe: Final = gateway.client.get("/health/liveliness") - liveliness.append((probe.status_code, time.perf_counter() - started)) - - poller: Final = threading.Thread(target=poll) - poller.start() - response, elapsed = _timed_decide(gateway, model) - stop.set() - poller.join() + tenth_response, tenth_elapsed, _ = _timed_decide_under_liveliness_probes( + gateway, f"{prefix}/{_TENTH_MEGABYTE_NAME}" + ) + response, elapsed, liveliness = _timed_decide_under_liveliness_probes(gateway, model) assert tenth_response.status_code == 400, tenth_response.text[:300] - assert elapsed < 25 * tenth_elapsed, (elapsed, tenth_elapsed) + assert_within_factor_of_control(elapsed, control=tenth_elapsed, factor=25) assert liveliness and all(status == 200 for status, _ in liveliness), liveliness assert max(seconds for _, seconds in liveliness) < max(1.0, elapsed / 2), (elapsed, liveliness) assert response.status_code == 400, response.text[:300] diff --git a/tests/integration/providers/test_decisions_chaos.py b/tests/integration/providers/test_decisions_chaos.py index 63b4a482e92..d868e46e6a5 100644 --- a/tests/integration/providers/test_decisions_chaos.py +++ b/tests/integration/providers/test_decisions_chaos.py @@ -1,7 +1,6 @@ import asyncio import json import re -import signal import threading import uuid from dataclasses import dataclass @@ -215,12 +214,12 @@ async def test_worker_sigkill_mid_burst_leaves_the_sibling_serving_the_default_m held_by: Final = MappingProxyType({pid: _open_upstream_connections(pid, wire.url) for pid in workers}) assert sum(held_by.values()) == 20, held_by victim_pid, survivor_pid = sorted(workers, key=held_by.__getitem__) - victim: Final = psutil.Process(victim_pid) - victim.suspend() - victim.send_signal(signal.SIGKILL) - release.set() - served: Final = await burst - assert len(served) == held_by[survivor_pid], (held_by, len(served)) + with owned.replacing_worker(victim_pid) as kill: + kill.freeze() + await asyncio.to_thread(kill) + release.set() + served: Final = await burst + assert len(served) == held_by[survivor_pid], (held_by, len(served)) follow_up: Final = _Call(route="/systemone", marker=f"ok-{uuid.uuid4().hex}", fail=False) (answered,) = await _burst(base_url, candidate.key, None, (follow_up,)) await asyncio.to_thread( diff --git a/tests/integration/providers/test_device_code_login_guard_boot.py b/tests/integration/providers/test_device_code_login_guard_boot.py index 687a8f1d313..34ab7f1597f 100644 --- a/tests/integration/providers/test_device_code_login_guard_boot.py +++ b/tests/integration/providers/test_device_code_login_guard_boot.py @@ -2,7 +2,6 @@ import asyncio import hashlib import itertools import re -import signal import threading import uuid from collections.abc import Mapping @@ -422,12 +421,13 @@ async def test_worker_sigkill_mid_burst_leaves_the_sibling_refusing_logins_and_s owned, key, workers, control_port, held_markers, (), tasks ) victim_pid, survivor_pid = sorted(workers, key=held_by.__getitem__) - victim: Final = psutil.Process(victim_pid) - victim.suspend() - victim.send_signal(signal.SIGKILL) + with owned.replacing_worker(victim_pid) as kill: + kill.freeze() + await asyncio.to_thread(kill) + release.set() + answers: Final = await asyncio.gather(*bursts) finally: release.set() - answers: Final = await asyncio.gather(*bursts) served: Final = tuple(item for item in itertools.chain.from_iterable(answers) if item is not None) assert len(served) == held_by[survivor_pid], (held_by, len(served)) for item in served: diff --git a/tests/integration/providers/test_device_code_login_guard_wire.py b/tests/integration/providers/test_device_code_login_guard_wire.py index 072de4f3782..c88bd0b9273 100644 --- a/tests/integration/providers/test_device_code_login_guard_wire.py +++ b/tests/integration/providers/test_device_code_login_guard_wire.py @@ -24,7 +24,9 @@ from integration._support import responses_vendor as rv from integration._support.client import eventually, gateway_from_environment from integration._support.database import read_rows from integration._support.process import OwnedProxy, owned_proxy_process +from integration._support.timing import assert_cut_short, timed from integration._support.wire import Request +from litellm.llms.chatgpt.authenticator import DEVICE_CODE_COOLDOWN_SECONDS from pydantic import JsonValue pytestmark: Final = pytest.mark.timeout(300) @@ -713,16 +715,18 @@ def test_a_token_mounted_after_a_refusal_serves_without_a_restart_and_removing_i _COOLDOWN_SECONDS_LEFT: Final = 45 +_REFUSAL_MARGIN_SECONDS: Final = 15 def test_a_recent_device_code_request_does_not_hold_the_worker_for_the_rest_of_the_cooldown(scene: _Scene) -> None: - dl.write_chatgpt(scene.rig.chatgpt_dir, {"device_code_requested_at": time.time() - (300 - _COOLDOWN_SECONDS_LEFT)}) - started: Final = time.monotonic() - served: Final = _serve(scene.call(_Cell("chatgpt", "chat", False, "httpx"))) - elapsed: Final = time.monotonic() - started + requested_at: Final = time.time() - (DEVICE_CODE_COOLDOWN_SECONDS - _COOLDOWN_SECONDS_LEFT) + dl.write_chatgpt(scene.rig.chatgpt_dir, {"device_code_requested_at": requested_at}) + remaining: Final = requested_at + DEVICE_CODE_COOLDOWN_SECONDS - time.time() + refusal: Final = timed(lambda: _serve(scene.call(_Cell("chatgpt", "chat", False, "httpx")))) + served: Final = refusal.value assert served.status == 400, served.text assert dl.CHATGPT_REFUSAL in _error_message(served.text), served.text - assert elapsed < _COOLDOWN_SECONDS_LEFT - 15, elapsed + assert_cut_short(refusal.seconds, injected_delay=remaining - _REFUSAL_MARGIN_SECONDS) assert scene.rig.peers.auth_connections() == () diff --git a/tests/integration/providers/test_gemini_embedding_file_block_chaos.py b/tests/integration/providers/test_gemini_embedding_file_block_chaos.py index cee0758b433..17483147107 100644 --- a/tests/integration/providers/test_gemini_embedding_file_block_chaos.py +++ b/tests/integration/providers/test_gemini_embedding_file_block_chaos.py @@ -1,8 +1,6 @@ from __future__ import annotations import json -import os -import signal import threading import uuid import zlib @@ -192,15 +190,17 @@ def test_block_embeddings_survive_a_worker_kill(gateway: Gateway, tmp_path: Path victim: Final = min(workers) assert outcome(candidate, name, 0).startswith("ok:") - def attempt_around_the_kill(index: int) -> str: - if index == 2: - os.kill(victim, signal.SIGKILL) - return outcome(candidate, name, index) + with owned.replacing_worker(victim) as kill: - with ThreadPoolExecutor(max_workers=BURST) as pool: - outcomes: Final = tuple(pool.map(attempt_around_the_kill, range(1, BURST + 1))) - assert outcomes.count("ok") == 0 and any(item.startswith("ok:") for item in outcomes), outcomes - assert all(item.startswith("ok:") or item.startswith("transport:") for item in outcomes), outcomes + def attempt_around_the_kill(index: int) -> str: + if index == 2: + kill() + return outcome(candidate, name, index) + + with ThreadPoolExecutor(max_workers=BURST) as pool: + outcomes: Final = tuple(pool.map(attempt_around_the_kill, range(1, BURST + 1))) + assert outcomes.count("ok") == 0 and any(item.startswith("ok:") for item in outcomes), outcomes + assert all(item.startswith("ok:") or item.startswith("transport:") for item in outcomes), outcomes respawned: Final = eventually( lambda: worker_pids(owned.process.pid), lambda pids: len(pids) == 2 and victim not in pids, diff --git a/tests/integration/providers/test_gemini_thinking_replay_chaos.py b/tests/integration/providers/test_gemini_thinking_replay_chaos.py index 57058eb5fbe..f8bb52ef60a 100644 --- a/tests/integration/providers/test_gemini_thinking_replay_chaos.py +++ b/tests/integration/providers/test_gemini_thinking_replay_chaos.py @@ -1,7 +1,6 @@ import asyncio import json import re -import signal import threading import uuid from collections.abc import Callable, Mapping @@ -327,15 +326,15 @@ async def test_worker_sigkill_mid_burst_leaves_the_sibling_replaying_thinking_si held_by: Final = MappingProxyType({pid: _open_upstream_connections(pid, wire.url) for pid in workers}) assert sum(held_by.values()) == 20, held_by victim_pid, survivor_pid = sorted(workers, key=held_by.__getitem__) - victim: Final = psutil.Process(victim_pid) - victim.suspend() - victim.send_signal(signal.SIGKILL) - release.set() - served: Final = await burst - assert held_by[survivor_pid] >= 10, held_by - assert len(served) == held_by[survivor_pid], (held_by, len(served)) - for item in served: - _assert_answered_with_its_own_marker(item) + with owned.replacing_worker(victim_pid) as kill: + kill.freeze() + await asyncio.to_thread(kill) + release.set() + served: Final = await burst + assert held_by[survivor_pid] >= 10, held_by + assert len(served) == held_by[survivor_pid], (held_by, len(served)) + for item in served: + _assert_answered_with_its_own_marker(item) follow_up: Final = _Call(endpoint="messages", stream=False, marker=uuid.uuid4().hex) (answered,) = await _burst(str(candidate.client.base_url), candidate.key, _CONFIG_MODEL, (follow_up,)) _assert_answered_with_its_own_marker(answered) diff --git a/tests/integration/providers/test_gemini_tts_voice_mapping_owned_proxy.py b/tests/integration/providers/test_gemini_tts_voice_mapping_owned_proxy.py index 8a977427741..3a5a9bf5d06 100644 --- a/tests/integration/providers/test_gemini_tts_voice_mapping_owned_proxy.py +++ b/tests/integration/providers/test_gemini_tts_voice_mapping_owned_proxy.py @@ -8,9 +8,7 @@ so these rows boot their own proxy with ``GEMINI_API_BASE`` pointed at the scrip from __future__ import annotations import asyncio -import os import re -import signal import threading from collections.abc import Callable, Iterator from dataclasses import dataclass @@ -306,11 +304,13 @@ async def test_worker_sigkill_mid_speech_burst_leaves_the_sibling_serving(gatewa ) held_by: Final = MappingProxyType({pid: _held_connections(pid, wire.url) for pid in workers}) victim: Final = max(workers, key=held_by.__getitem__) - os.kill(victim, signal.SIGKILL) + with proxy.replacing_worker(victim) as kill: + await asyncio.to_thread(kill) + release.set() + served: Final = await burst + during: Final = wire.drain() finally: release.set() - served: Final = await burst - during: Final = wire.drain() after: Final = await _fire(url, proxy.gateway.key, model) after_received: Final = wire.drain() assert sum(held_by.values()) == BURST, held_by @@ -335,7 +335,7 @@ async def test_proxy_restart_mid_speech_burst_serves_the_next_burst_after_the_re burst: Final = tasks.create_task(_fire(url, first.gateway.key, model, tolerate_transport_errors=True)) try: await asyncio.to_thread(eventually, lambda: wire.received.qsize(), lambda size: size >= BURST, 90) - first.process.terminate() + first.begin_shutdown() finally: release.set() served: Final = await burst diff --git a/tests/integration/providers/test_hosted_vllm_reasoning_content_chaos.py b/tests/integration/providers/test_hosted_vllm_reasoning_content_chaos.py index 84a5d71e8b8..eabbba425df 100644 --- a/tests/integration/providers/test_hosted_vllm_reasoning_content_chaos.py +++ b/tests/integration/providers/test_hosted_vllm_reasoning_content_chaos.py @@ -1,7 +1,6 @@ import asyncio import json import re -import signal import threading import uuid from collections.abc import Callable @@ -374,15 +373,15 @@ async def test_worker_sigkill_mid_burst_leaves_the_sibling_forwarding_reasoning( held_by: Final = MappingProxyType({pid: _open_upstream_connections(pid, wire.url) for pid in workers}) assert sum(held_by.values()) == 20, held_by victim_pid, survivor_pid = sorted(workers, key=held_by.__getitem__) - victim: Final = psutil.Process(victim_pid) - victim.suspend() - victim.send_signal(signal.SIGKILL) - release.set() - served: Final = await burst - assert held_by[survivor_pid] >= 10, held_by - assert len(served) == held_by[survivor_pid], (held_by, len(served)) - for item in served: - _assert_answered_with_its_own_marker(item) + with owned.replacing_worker(victim_pid) as kill: + kill.freeze() + await asyncio.to_thread(kill) + release.set() + served: Final = await burst + assert held_by[survivor_pid] >= 10, held_by + assert len(served) == held_by[survivor_pid], (held_by, len(served)) + for item in served: + _assert_answered_with_its_own_marker(item) follow_up: Final = _Call(endpoint="chat", stream=False, marker=uuid.uuid4().hex) (answered,) = await _burst(str(candidate.client.base_url), candidate.key, _CONFIG_MODEL, (follow_up,)) _assert_answered_with_its_own_marker(answered) diff --git a/tests/integration/providers/test_ollama_prompt_tools_chaos.py b/tests/integration/providers/test_ollama_prompt_tools_chaos.py index 243c85b2a64..37529f9f692 100644 --- a/tests/integration/providers/test_ollama_prompt_tools_chaos.py +++ b/tests/integration/providers/test_ollama_prompt_tools_chaos.py @@ -1,7 +1,6 @@ import asyncio import json import re -import signal import threading import uuid from collections.abc import Callable @@ -389,15 +388,15 @@ async def test_worker_sigkill_mid_burst_leaves_the_sibling_instructing_ollama(ga held_by: Final = MappingProxyType({pid: _open_upstream_connections(pid, wire.url) for pid in workers}) assert sum(held_by.values()) == 20, held_by victim_pid, survivor_pid = sorted(workers, key=held_by.__getitem__) - victim: Final = psutil.Process(victim_pid) - victim.suspend() - victim.send_signal(signal.SIGKILL) - release.set() - served: Final = await burst - assert held_by[survivor_pid] >= 10, held_by - assert len(served) == held_by[survivor_pid], (held_by, len(served)) - for item in served: - _assert_answered_with_its_own_marker(item) + with owned.replacing_worker(victim_pid) as kill: + kill.freeze() + await asyncio.to_thread(kill) + release.set() + served: Final = await burst + assert held_by[survivor_pid] >= 10, held_by + assert len(served) == held_by[survivor_pid], (held_by, len(served)) + for item in served: + _assert_answered_with_its_own_marker(item) follow_up: Final = _Call(endpoint="chat", stream=False, marker=uuid.uuid4().hex) (answered,) = await _burst(str(candidate.client.base_url), candidate.key, _CONFIG_MODEL, (follow_up,)) _assert_answered_with_its_own_marker(answered) diff --git a/tests/integration/providers/test_openai_dialect_prompt_cache_breakpoint_owned_proxy.py b/tests/integration/providers/test_openai_dialect_prompt_cache_breakpoint_owned_proxy.py index cc7ab06da90..920002316bc 100644 --- a/tests/integration/providers/test_openai_dialect_prompt_cache_breakpoint_owned_proxy.py +++ b/tests/integration/providers/test_openai_dialect_prompt_cache_breakpoint_owned_proxy.py @@ -1,6 +1,5 @@ import asyncio import re -import signal import threading from collections.abc import Callable, Iterator from concurrent.futures import ThreadPoolExecutor @@ -304,12 +303,12 @@ async def test_f2_a_worker_killed_mid_burst_leaves_the_survivor_marking_requests held_by: Final = MappingProxyType({pid: _open_peer_connections(pid, wire.url) for pid in workers}) assert sum(held_by.values()) == _HELD_BURST, held_by victim_pid, survivor_pid = sorted(workers, key=held_by.__getitem__) - victim: Final = psutil.Process(victim_pid) - victim.suspend() - victim.send_signal(signal.SIGKILL) - release.set() - served: Final = await burst - assert len(served) == held_by[survivor_pid], (held_by, len(served)) + with owned.replacing_worker(victim_pid) as kill: + kill.freeze() + await asyncio.to_thread(kill) + release.set() + served: Final = await burst + assert len(served) == held_by[survivor_pid], (held_by, len(served)) follow_up: Final = fresh_marker() (answered,) = await _held_burst(url, candidate.key, (follow_up,)) _assert_served_and_landed(wire, (*served, answered), expected_received=_HELD_BURST + 1) @@ -365,7 +364,7 @@ async def test_f3_a_graceful_restart_mid_burst_drains_the_held_requests_and_keep release.clear() burst: Final = asyncio.create_task(_held_burst(url, candidate.key, markers)) await asyncio.to_thread(eventually, held.qsize, lambda size: size == _HELD_BURST, 60) - owned.process.terminate() + owned.begin_shutdown() release.set() served: Final = await burst assert len(served) == _HELD_BURST, len(served) diff --git a/tests/integration/providers/test_openai_realtime_transcription_url.py b/tests/integration/providers/test_openai_realtime_transcription_url.py index eec0dd37baa..81c450d8f70 100644 --- a/tests/integration/providers/test_openai_realtime_transcription_url.py +++ b/tests/integration/providers/test_openai_realtime_transcription_url.py @@ -30,7 +30,7 @@ from websockets.exceptions import ConnectionClosed, InvalidStatus from tests.integration._support.client import JSON_OBJECT, Gateway, Scenario, eventually, object_value, string_value from tests.integration._support.database import read_rows -from tests.integration._support.process import owned_proxy_process, owned_upstream, stop_root_process +from tests.integration._support.process import OwnedProxy, owned_proxy_process, owned_upstream, stop_root_process from tests.integration._support.upstream import ScenarioHandle, delete_scenario, register_scenario from tests.integration.cost_calculation.cost_tracking_case import RealtimeResponse @@ -532,7 +532,7 @@ async def _first_frames(sockets: tuple[ClientConnection, ...]) -> tuple[dict[str return tuple([JSON_OBJECT.validate_json(await socket.recv()) for socket in sockets]) -async def _sessions_through_worker_kill(ws_base: str, model: str, key: str, root: psutil.Process) -> KillOutcome: +async def _sessions_through_worker_kill(ws_base: str, model: str, key: str, owned: OwnedProxy) -> KillOutcome: query: Final = f"model={model}&{TRANSCRIPTION_QUERY}" headers: Final = {"Authorization": f"Bearer {key}"} sockets: Final = tuple( @@ -544,13 +544,15 @@ async def _sessions_through_worker_kill(ws_base: str, model: str, key: str, root try: created: Final = await asyncio.wait_for(_first_frames(sockets), 60) assert [event.get("type") for event in created] == ["session.created"] * OPEN_SESSIONS, created - workers: Final = _workers(root) - assert len(workers) == WORKERS, [process.pid for process in workers] + workers: Final = tuple(sorted(owned.worker_pids())) + assert len(workers) == WORKERS, workers victim: Final = workers[0] - victim.kill() - await asyncio.to_thread(victim.wait, 10) - served: Final = await asyncio.wait_for(asyncio.gather(*(_one_turn_or_close(socket) for socket in sockets)), 60) - return KillOutcome(served.count(True), served.count(False), victim.pid) + with owned.replacing_worker(victim) as kill: + await asyncio.to_thread(kill) + served: Final = await asyncio.wait_for( + asyncio.gather(*(_one_turn_or_close(socket) for socket in sockets)), 60 + ) + return KillOutcome(served.count(True), served.count(False), victim) finally: await asyncio.gather(*(socket.close() for socket in sockets)) @@ -607,7 +609,7 @@ def test_worker_kill_then_proxy_restart_keep_transcription_sessions_serving( handle.scenario_id, ) root: Final = psutil.Process(owned.process.pid) - outcome: Final = asyncio.run(_sessions_through_worker_kill(_ws_base(owned_url), model, key, root)) + outcome: Final = asyncio.run(_sessions_through_worker_kill(_ws_base(owned_url), model, key, owned)) record_property( "worker_kill", {"served": outcome.served, "closed": outcome.closed, "killed": outcome.killed_pid} ) diff --git a/tests/integration/providers/test_realtime_transcription_guardrail_session_update.py b/tests/integration/providers/test_realtime_transcription_guardrail_session_update.py index e39f7e360f6..cda00b5f3fa 100644 --- a/tests/integration/providers/test_realtime_transcription_guardrail_session_update.py +++ b/tests/integration/providers/test_realtime_transcription_guardrail_session_update.py @@ -2067,18 +2067,18 @@ async def _open_sessions(ws_base: str, model: str, key: str) -> tuple[ClientConn return sockets -async def _sessions_through_worker_kill(ws_base: str, model: str, key: str, root: psutil.Process) -> KillOutcome: +async def _sessions_through_worker_kill(ws_base: str, model: str, key: str, owned: OwnedProxy) -> KillOutcome: sockets: Final = await _open_sessions(ws_base, model, key) try: - workers: Final = _workers(root) - assert len(workers) == WORKERS, [process.pid for process in workers] + workers: Final = tuple(sorted(owned.worker_pids())) + assert len(workers) == WORKERS, workers victim: Final = workers[0] - victim.kill() - await asyncio.to_thread(victim.wait, 10) - served: Final = await asyncio.wait_for( - asyncio.gather(*(_one_transcript_or_close(socket) for socket in sockets)), 60 - ) - return KillOutcome(served.count(True), served.count(False), victim.pid) + with owned.replacing_worker(victim) as kill: + await asyncio.to_thread(kill) + served: Final = await asyncio.wait_for( + asyncio.gather(*(_one_transcript_or_close(socket) for socket in sockets)), 60 + ) + return KillOutcome(served.count(True), served.count(False), victim) finally: await asyncio.gather(*(socket.close() for socket in sockets)) @@ -2138,7 +2138,7 @@ def test_worker_kill_then_proxy_restart_keep_transcription_sessions_uninjected( blocked_after_restart.scenario_id, ) root: Final = psutil.Process(owned.process.pid) - outcome: Final = asyncio.run(_sessions_through_worker_kill(_ws_base(owned_url), model, key, root)) + outcome: Final = asyncio.run(_sessions_through_worker_kill(_ws_base(owned_url), model, key, owned)) record_property( "worker_kill", {"served": outcome.served, "closed": outcome.closed, "killed": outcome.killed_pid} ) diff --git a/tests/integration/providers/test_realtime_unguarded_voice_turns.py b/tests/integration/providers/test_realtime_unguarded_voice_turns.py index 396b266f45f..4dd7fd9aeef 100644 --- a/tests/integration/providers/test_realtime_unguarded_voice_turns.py +++ b/tests/integration/providers/test_realtime_unguarded_voice_turns.py @@ -853,13 +853,6 @@ def _spawned_worker(child: psutil.Process) -> bool: return False -def _exited(process: psutil.Process) -> bool: - try: - return process.status() == psutil.STATUS_ZOMBIE - except psutil.NoSuchProcess: - return True - - def _workers(root: psutil.Process) -> tuple[psutil.Process, ...]: return tuple( sorted((child for child in root.children() if _spawned_worker(child)), key=lambda process: process.pid) @@ -913,16 +906,16 @@ async def _sessions_on_every_worker( async def _sessions_through_worker_kill( - url: str, key: str, workers: tuple[psutil.Process, ...] + url: str, key: str, owned: OwnedProxy, workers: tuple[psutil.Process, ...] ) -> tuple[tuple[Session, bool], ...]: opened: Final = await _sessions_on_every_worker(url, key, workers, ()) sockets: Final = tuple(session.socket for session in opened) try: on_killed_worker: Final = tuple(session.pid == workers[0].pid for session in opened) - workers[0].kill() - await asyncio.to_thread(eventually, lambda: _exited(workers[0]), bool, graceful_stop_seconds()) - turns: Final = (_turn_or_close(socket, label) for socket, label in zip(sockets, _labels(len(sockets)))) - sessions: Final = await asyncio.wait_for(asyncio.gather(*turns), 90) + with owned.replacing_worker(workers[0].pid) as kill: + await asyncio.to_thread(kill) + turns: Final = (_turn_or_close(socket, label) for socket, label in zip(sockets, _labels(len(sockets)))) + sessions: Final = await asyncio.wait_for(asyncio.gather(*turns), 90) return tuple(zip(sessions, on_killed_worker)) finally: await asyncio.gather(*(socket.close() for socket in sockets)) @@ -941,7 +934,7 @@ def test_worker_kill_leaves_the_surviving_sessions_with_only_their_own_frames(ga workers: Final = _workers(root) assert len(workers) == WORKERS, [process.pid for process in workers] outcomes: Final = asyncio.run( - _sessions_through_worker_kill(_session_url(_owned_url(owned), f"model={model}"), key, workers) + _sessions_through_worker_kill(_session_url(_owned_url(owned), f"model={model}"), key, owned, workers) ) served: Final = tuple(label for label, (_, killed) in zip(_labels(len(outcomes)), outcomes) if not killed) assert [session.types for session, _ in outcomes] == [ diff --git a/tests/integration/providers/test_responses_bridge_prompt_cache_breakpoint_chaos.py b/tests/integration/providers/test_responses_bridge_prompt_cache_breakpoint_chaos.py index 10d5cfd9646..e03af85d01a 100644 --- a/tests/integration/providers/test_responses_bridge_prompt_cache_breakpoint_chaos.py +++ b/tests/integration/providers/test_responses_bridge_prompt_cache_breakpoint_chaos.py @@ -2,7 +2,6 @@ from __future__ import annotations import asyncio import re -import signal import threading import uuid from contextlib import ExitStack @@ -126,11 +125,11 @@ async def test_worker_and_peer_outages_preserve_markers_and_recover( survivor_pids: Final = tuple(pid for pid in workers if pid != victim_pid) assert active_by_worker[victim_pid] > 0 and len(survivor_pids) == 1, active_by_worker (survivor_pid,) = survivor_pids - victim: Final = psutil.Process(victim_pid) - victim.suspend() - victim.send_signal(signal.SIGKILL) - release.set() - remaining_served: Final = await asyncio.gather(*burst_tasks[len(early_calls) :]) + with owned.replacing_worker(victim_pid) as kill: + kill.freeze() + await asyncio.to_thread(kill) + release.set() + remaining_served: Final = await asyncio.gather(*burst_tasks[len(early_calls) :]) served: Final = (*early_served, *remaining_served) successful: Final = tuple(item for item in served if item.status == 200) connection_errors: Final = tuple(item for item in served if item.status == 0) diff --git a/tests/integration/providers/test_responses_bridge_prompt_cache_breakpoint_wire_chaos.py b/tests/integration/providers/test_responses_bridge_prompt_cache_breakpoint_wire_chaos.py index 6f9a05d8dfa..ee73d780170 100644 --- a/tests/integration/providers/test_responses_bridge_prompt_cache_breakpoint_wire_chaos.py +++ b/tests/integration/providers/test_responses_bridge_prompt_cache_breakpoint_wire_chaos.py @@ -1,6 +1,5 @@ import asyncio import dataclasses -import signal import threading import uuid from collections import Counter @@ -13,7 +12,6 @@ from typing import Final, Literal, TypeAlias from urllib.parse import urlsplit import httpx -import psutil import pytest import yaml from integration._support import prompt_cache_breakpoint as pcb @@ -371,16 +369,16 @@ async def test_worker_sigkill_mid_burst_leaves_the_sibling_answering( held_by: Final = MappingProxyType({pid: pcb.open_upstream_connections(pid, wire.url) for pid in workers}) assert sum(held_by.values()) == 20, held_by victim_pid, survivor_pid = sorted(workers, key=held_by.__getitem__) - victim: Final = psutil.Process(victim_pid) - victim.suspend() - victim.send_signal(signal.SIGKILL) - held.release.set() - served: Final = await burst - assert held_by[survivor_pid] >= 10, held_by - assert len(served) == held_by[survivor_pid], (held_by, len(served)) - for item in served: - _assert_answered_in_its_own_shape(item) - _marked_once(pcb.drained_posts(wire), calls, pcb.EXPLICIT) + with owned.replacing_worker(victim_pid) as kill: + kill.freeze() + await asyncio.to_thread(kill) + held.release.set() + served: Final = await burst + assert held_by[survivor_pid] >= 10, held_by + assert len(served) == held_by[survivor_pid], (held_by, len(served)) + for item in served: + _assert_answered_in_its_own_shape(item) + _marked_once(pcb.drained_posts(wire), calls, pcb.EXPLICIT) follow_up: Final = uuid.uuid4().hex call_id: Final = _completion(_chat(candidate, _GLOBAL_UNSET, follow_up, "yes"), follow_up) pcb.assert_marker(_user_block_on_wire(wire, follow_up), None) @@ -397,7 +395,7 @@ async def test_proxy_restart_mid_burst_never_lands_a_served_call_twice( with owned_proxy_process(gateway, tmp_path, {}, config=config, workers=2) as first: _worker_pids(first) burst: Final = await _hold_burst(held, first.gateway, _GLOBAL_UNSET, calls) - first.process.terminate() + first.begin_shutdown() held.release.set() served: Final = await burst for item in served: diff --git a/tests/integration/providers/test_responses_minted_reasoning_replay_chaos.py b/tests/integration/providers/test_responses_minted_reasoning_replay_chaos.py index d62b67e08d6..f003e5ffa63 100644 --- a/tests/integration/providers/test_responses_minted_reasoning_replay_chaos.py +++ b/tests/integration/providers/test_responses_minted_reasoning_replay_chaos.py @@ -1,7 +1,6 @@ import asyncio import json import re -import signal import socket import threading import time @@ -351,15 +350,15 @@ async def test_worker_sigkill_mid_burst_leaves_the_sibling_dropping_the_minted_i held_by: Final = MappingProxyType({pid: _open_upstream_connections(pid, wire.url) for pid in workers}) assert sum(held_by.values()) == 20, held_by victim_pid, survivor_pid = sorted(workers, key=held_by.__getitem__) - victim: Final = psutil.Process(victim_pid) - victim.suspend() - victim.send_signal(signal.SIGKILL) - release.set() - served: Final = await burst - assert held_by[survivor_pid] >= 10, held_by - assert len(served) == held_by[survivor_pid], (held_by, len(served)) - for item in served: - _assert_answered_with_its_own_marker(item) + with owned.replacing_worker(victim_pid) as kill: + kill.freeze() + await asyncio.to_thread(kill) + release.set() + served: Final = await burst + assert held_by[survivor_pid] >= 10, held_by + assert len(served) == held_by[survivor_pid], (held_by, len(served)) + for item in served: + _assert_answered_with_its_own_marker(item) follow_up: Final = _Call("responses", False, uuid.uuid4().hex) (answered,) = await _burst(str(candidate.client.base_url), candidate.key, models, (follow_up,)) _assert_answered_with_its_own_marker(answered) diff --git a/tests/integration/providers/test_responses_minted_reasoning_replay_wire.py b/tests/integration/providers/test_responses_minted_reasoning_replay_wire.py index 036f96916a6..2253443dfc0 100644 --- a/tests/integration/providers/test_responses_minted_reasoning_replay_wire.py +++ b/tests/integration/providers/test_responses_minted_reasoning_replay_wire.py @@ -17,6 +17,7 @@ from integration._support import claude_code as cc from integration._support import responses_vendor as rv from integration._support.client import Gateway, Scenario, eventually from integration._support.database import read_rows +from integration._support.timing import assert_within_factor_of_control, timed from integration._support.wire import Request, Wire, wire_server from pydantic import JsonValue, TypeAdapter @@ -26,6 +27,7 @@ _CLAUDE: Final = cc.OPUS _OPENAI_KEY: Final = "synthetic-openai-key" _AZURE_KEY: Final = "synthetic-azure-key" _CACHE_BUST: Final[Mapping[str, JsonValue]] = MappingProxyType({"cache": {"no-cache": True}}) +_RESPONSIVE_FACTOR: Final = 10 _JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue]) _ITEMS: Final = TypeAdapter(list[dict[str, JsonValue]]) @@ -564,6 +566,11 @@ def test_a_megabyte_of_minted_thinking_is_dropped_while_the_proxy_stays_responsi "id": f"rs_{uuid.uuid4()}", "encrypted_content": encrypted, } + single: Final[dict[str, JsonValue]] = { + **minted, + "id": f"rs_{uuid.uuid4()}", + "encrypted_content": json.dumps([block]), + } history: Final = rv.agents_sdk_history(marker, minted) latencies: Final[deque[float]] = deque() done: Final = threading.Event() @@ -577,16 +584,22 @@ def test_a_megabyte_of_minted_thinking_is_dropped_while_the_proxy_stays_responsi with wire_server(rv.ResponsesVendor().respond) as wire, gateway.scenario() as scenario: model: Final = _OPENAI.register(scenario, wire) + same_size: Final[list[dict[str, JsonValue]]] = [ + {"role": "user", "content": "t" * len(encrypted)}, + *rv.agents_sdk_history(marker, single), + ] + control: Final = timed(lambda: _raw(gateway, "/v1/responses", {"model": model, "input": same_size})) + assert control.value.status_code == 200, control.value.text[:500] + _only_request(wire) prober: Final = threading.Thread(target=probe) prober.start() - started: Final = time.monotonic() - response: Final = _raw(gateway, "/v1/responses", {"model": model, "input": history}) - elapsed: Final = time.monotonic() - started + dropped: Final = timed(lambda: _raw(gateway, "/v1/responses", {"model": model, "input": history})) done.set() prober.join(timeout=35) - assert response.status_code == 200, response.text[:500] - assert elapsed < 20, elapsed - assert latencies and max(latencies) < 5, (max(latencies), len(latencies)) + assert dropped.value.status_code == 200, dropped.value.text[:500] + assert_within_factor_of_control(dropped.seconds, control=control.seconds, factor=_RESPONSIVE_FACTOR) + assert latencies, "no liveliness probe completed while the megabyte request was in flight" + assert_within_factor_of_control(max(latencies), control=control.seconds, factor=_RESPONSIVE_FACTOR) _, body = _only_request(wire) assert body["input"] == rv.without(history, (minted,)) diff --git a/tests/integration/providers/test_responses_websocket_session_limit.py b/tests/integration/providers/test_responses_websocket_session_limit.py index 32b01d8387a..f2b3a48c6d1 100644 --- a/tests/integration/providers/test_responses_websocket_session_limit.py +++ b/tests/integration/providers/test_responses_websocket_session_limit.py @@ -24,6 +24,7 @@ import websockets from integration._support.client import Gateway, eventually, gateway_from_environment from integration._support.database import read_rows, scratch_database from integration._support.process import OwnedProxy, graceful_stop_seconds, owned_proxy, owned_proxy_process +from integration._support.timing import assert_cut_short, assert_enforced_within from integration._support.tls import server_context, write_self_signed_cert from integration._support.upstream import ( JsonResponse, @@ -44,6 +45,8 @@ STALL_PROVIDER_MODEL: Final = "ws-stall-peer-model" DEAF_PROVIDER_MODEL: Final = "ws-deaf-peer-model" DEAF_RESPONSE_DELAY_SECONDS: Final = 10 DEAF_READ_PAUSE_SECONDS: Final = 20 +MID_TURN_SECONDS: Final = 40 +DEAF_UNTIL_SECONDS: Final = MID_TURN_SECONDS + DEAF_RESPONSE_DELAY_SECONDS + DEAF_READ_PAUSE_SECONDS PEER_TEXT: Final = "responses websocket peer" TERMINAL: Final = frozenset({"response.completed", "response.failed", "error"}) JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue]) @@ -52,6 +55,8 @@ LIMIT_CLOSE_REASON: Final = "Session duration limit reached" RELOAD_INTERVAL_SECONDS: Final = 3 WORKER_SYNC_SECONDS: Final = RELOAD_INTERVAL_SECONDS + 7 CAP_SECONDS: Final = 60 +CAP_CLOSE_SLACK_SECONDS: Final = 15 +DEAF_CLOSE_SLACK_SECONDS: Final = 3 RESTORED_IDLE_SECONDS: Final = 75 CAPPED_PATHS: Final = ("/v1/responses", "/responses") CAPPED_SOCKETS_MINIMUM: Final = 8 @@ -725,7 +730,7 @@ async def _mid_session( additional_headers={"Authorization": f"Bearer {key}"}, open_timeout=10, ) as connection: - await asyncio.sleep(max(0, started + 40 - time.monotonic())) + await asyncio.sleep(max(0, started + MID_TURN_SECONDS - time.monotonic())) created, turn_error = await _stall_turn(connection, stall_model) remaining: Final = max(0, 75 - (time.monotonic() - started)) close: Final = CloseResult(await _wait_for_close(connection, remaining), time.monotonic() - started) @@ -982,19 +987,23 @@ async def _kill_workload(owned: OwnedProxy, key: str, model: str) -> KillResults ) held_by: Final = MappingProxyType({pid: holders.count(pid) for pid in sorted(workers)}) victim: Final = max(sorted(workers), key=held_by.__getitem__) - psutil.Process(victim).kill() victims: Final = tuple( connection for connection, holder in zip(connections, holders, strict=True) if holder == victim ) survivors: Final = tuple( connection for connection, holder in zip(connections, holders, strict=True) if holder != victim ) - victim_closes, survivor_turns = await asyncio.gather( - asyncio.gather(*tuple(_wait_for_close(connection, 30) for connection in victims)), - asyncio.gather( - *tuple(_turn_result(connection, _create(model, f"survivor-{uuid.uuid4().hex}")) for connection in survivors) - ), - ) + with owned.replacing_worker(victim) as kill: + await asyncio.to_thread(kill) + victim_closes, survivor_turns = await asyncio.gather( + asyncio.gather(*tuple(_wait_for_close(connection, 30) for connection in victims)), + asyncio.gather( + *tuple( + _turn_result(connection, _create(model, f"survivor-{uuid.uuid4().hex}")) + for connection in survivors + ) + ), + ) health: Final = await asyncio.to_thread( eventually, lambda: _health_status(candidate), @@ -1327,7 +1336,9 @@ def test_session_cap_closes_never_started_socket(cap_results: CapResults) -> Non assert cap_results.idle.outcome.closed, cap_results.idle assert cap_results.idle.outcome.close_code == 1000, cap_results.idle assert cap_results.idle.outcome.close_reason == "Session duration limit reached", cap_results.idle - assert 59 <= cap_results.idle.elapsed <= 75, cap_results.idle + assert_enforced_within( + cap_results.idle.elapsed, configured=CAP_SECONDS, slack=CAP_CLOSE_SLACK_SECONDS, not_before=CAP_SECONDS - 1 + ) def test_session_cap_closes_active_socket_and_provider(cap_results: CapResults) -> None: @@ -1336,7 +1347,12 @@ def test_session_cap_closes_active_socket_and_provider(cap_results: CapResults) assert cap_results.active.close.outcome.closed, cap_results.active assert cap_results.active.close.outcome.close_code == 1000, cap_results.active assert cap_results.active.close.outcome.close_reason == "Session duration limit reached", cap_results.active - assert 59 <= cap_results.active.close.elapsed <= 75, cap_results.active + assert_enforced_within( + cap_results.active.close.elapsed, + configured=CAP_SECONDS, + slack=CAP_CLOSE_SLACK_SECONDS, + not_before=CAP_SECONDS - 1, + ) assert cap_results.active.provider_closed, cap_results.active assert all(snapshot.closed for snapshot in cap_results.provider), cap_results.provider @@ -1347,7 +1363,9 @@ def test_session_cap_closes_mid_response_and_allows_new_session(cap_results: Cap assert cap_results.mid.close.outcome.closed, cap_results.mid assert cap_results.mid.close.outcome.close_code == 1000, cap_results.mid assert cap_results.mid.close.outcome.close_reason == "Session duration limit reached", cap_results.mid - assert 59 <= cap_results.mid.close.elapsed <= 75, cap_results.mid + assert_enforced_within( + cap_results.mid.close.elapsed, configured=CAP_SECONDS, slack=CAP_CLOSE_SLACK_SECONDS, not_before=CAP_SECONDS - 1 + ) assert cap_results.mid.provider_closed, cap_results.mid assert cap_results.mid.fresh_error is None, cap_results.mid assert cap_results.mid.fresh_completed, cap_results.mid @@ -1359,7 +1377,13 @@ def test_session_cap_closes_client_promptly_when_provider_ignores_close(cap_resu assert cap_results.deaf.close.outcome.closed, cap_results.deaf assert cap_results.deaf.close.outcome.close_code == 1000, cap_results.deaf assert cap_results.deaf.close.outcome.close_reason == "Session duration limit reached", cap_results.deaf - assert 59 <= cap_results.deaf.close.elapsed <= 63, cap_results.deaf + assert_enforced_within( + cap_results.deaf.close.elapsed, + configured=CAP_SECONDS, + slack=DEAF_CLOSE_SLACK_SECONDS, + not_before=CAP_SECONDS - 1, + ) + assert_cut_short(cap_results.deaf.close.elapsed, injected_delay=DEAF_UNTIL_SECONDS) assert cap_results.deaf.provider_closed, cap_results.deaf assert cap_results.deaf.fresh_error is None, cap_results.deaf assert cap_results.deaf.fresh_completed, cap_results.deaf @@ -1420,7 +1444,10 @@ def test_db_override_caps_new_sessions_on_every_worker_without_restart(override_ ) assert all(socket.close.outcome.closed and socket.close.outcome.close_code == 1000 for socket in capped), capped assert all(socket.close.outcome.close_reason == LIMIT_CLOSE_REASON for socket in capped), capped - assert all(CAP_SECONDS - 1 <= socket.close.elapsed <= CAP_SECONDS + 15 for socket in capped), capped + for socket in capped: + assert_enforced_within( + socket.close.elapsed, configured=CAP_SECONDS, slack=CAP_CLOSE_SLACK_SECONDS, not_before=CAP_SECONDS - 1 + ) assert frozenset(socket.worker_pid for socket in capped) == override_results.workers, capped diff --git a/tests/integration/providers/test_sagemaker_chat_credential_sources.py b/tests/integration/providers/test_sagemaker_chat_credential_sources.py index daaec6a7583..b01454acfcf 100644 --- a/tests/integration/providers/test_sagemaker_chat_credential_sources.py +++ b/tests/integration/providers/test_sagemaker_chat_credential_sources.py @@ -1,5 +1,4 @@ import os -import signal from collections.abc import Mapping from pathlib import Path from types import MappingProxyType @@ -186,13 +185,6 @@ def _worker_pids(log: str) -> tuple[int, ...]: ) -def _dead(process: psutil.Process) -> bool: - try: - return process.status() == psutil.STATUS_ZOMBIE - except psutil.NoSuchProcess: - return True - - @pytest.mark.timeout(_OWNED_PROXY_CELL_SECONDS) def test_every_aws_credential_source_keeps_its_credentials_out_of_the_body_when_a_worker_is_killed( gateway: Gateway, tmp_path: Path @@ -212,11 +204,10 @@ def test_every_aws_credential_source_keeps_its_credentials_out_of_the_body_when_ _assert_each_source_signed_its_own_clean_body( container, {source: _ask(owned, source) for source in _SIGNERS} ) - victim: Final = psutil.Process(workers[0]) - victim.send_signal(signal.SIGKILL) - eventually(lambda: _dead(victim), bool, seconds=graceful_stop_seconds()) - _assert_each_source_signed_its_own_clean_body( - container, {source: _ask(owned, source) for source in _SIGNERS} - ) - assert psutil.Process(workers[1]).is_running(), workers + with owned.replacing_worker(workers[0]) as kill: + kill() + _assert_each_source_signed_its_own_clean_body( + container, {source: _ask(owned, source) for source in _SIGNERS} + ) + assert psutil.Process(workers[1]).is_running(), workers assert owned.process.poll() is None, owned.log.read_text() diff --git a/tests/integration/providers/test_strands_decider_agentcore_wire.py b/tests/integration/providers/test_strands_decider_agentcore_wire.py index c7213a66bc5..9b7a9dbbfd0 100644 --- a/tests/integration/providers/test_strands_decider_agentcore_wire.py +++ b/tests/integration/providers/test_strands_decider_agentcore_wire.py @@ -3,7 +3,6 @@ import json import os import re import threading -import time import uuid from collections.abc import Callable, Iterator, Mapping, Sequence from concurrent.futures import ThreadPoolExecutor @@ -21,6 +20,7 @@ from integration._support.database import read_rows from integration._support.forward_proxy import ForwardProxy, tunnelling_forward_proxy from integration._support.process import owned_proxy from integration._support.sigv4 import encoded_path, signature +from integration._support.timing import Timed, assert_within_factor_of_control, timed from integration._support.tls import server_context, write_self_signed_cert from integration._support.wire import Reply, Request, Wire, wire_server from pydantic import JsonValue, TypeAdapter @@ -167,6 +167,8 @@ _load_yaml: Final[Callable[[str], object]] = yaml.safe_load _JSON: Final[TypeAdapter[JsonValue]] = TypeAdapter(JsonValue) _OWNED_PROXY_CELL_SECONDS: Final = 2 * max(30.0, float(os.environ.get("INTEGRATION_PROXY_READY_SECONDS") or 70)) + 120 pytestmark: Final = pytest.mark.timeout(_OWNED_PROXY_CELL_SECONDS) +_LINEAR_FACTOR: Final = 20 +_CLASSIFIER_SAMPLES: Final = 5 def _host(region: str) -> str: @@ -795,6 +797,10 @@ def test_a_plain_api_base_forwards_amazon_shaped_extra_headers_unchanged(gateway assert json.loads(call.body) == _UPSTREAM_BODY, call.body +def _fastest(operation: Callable[[], str]) -> Timed[str]: + return min((timed(operation) for _ in range(_CLASSIFIER_SAMPLES)), key=lambda sample: sample.seconds) + + def _worst_case_api_base() -> str: return "arn:aws" + "-a" * 131072 + "!" @@ -803,31 +809,28 @@ def test_a_worst_case_runtime_arn_pattern_input_is_classified_in_linear_time( gateway: Gateway, record_property: Callable[[str, object], None] ) -> None: worst: Final = _worst_case_api_base() - started: Final = time.perf_counter() - url: Final = StrandsDeciderDecisionsConfig().get_complete_url(worst, _BODY_MODEL) - classified: Final = time.perf_counter() - started - record_property("classifier_seconds", f"{classified:.6f}") - assert url == f"{worst}/v1/systemone" - assert classified < 0.5, classified + matching: Final = f"{worst.removesuffix('!')}:bedrock-agentcore:{_US_EAST}:123456789012:runtime/decider" + config: Final = StrandsDeciderDecisionsConfig() + control: Final = _fastest(lambda: config.get_complete_url(matching, _BODY_MODEL)) + classified: Final = _fastest(lambda: config.get_complete_url(worst, _BODY_MODEL)) + record_property("classifier_seconds", f"{classified.seconds:.6f}") + assert control.value.endswith("/invocations"), control.value[-200:] + assert classified.value == f"{worst}/v1/systemone" + assert_within_factor_of_control(classified.seconds, control=control.seconds, factor=_LINEAR_FACTOR) with wire_server(_answering) as wire, gateway.scenario() as scenario: model: Final = _plain_model(scenario, wire, configurable_clientside_auth_params=["api_base"]) + answered: Final = timed(lambda: _decide(gateway, model)) + _assert_answered(answered.value, model) + assert len(wire.drain()) == 1 with ThreadPoolExecutor(max_workers=2) as pool: - timed: Final = pool.submit(_timed_decide, gateway, model, worst) - liveliness_started: Final = time.perf_counter() - liveliness: Final = gateway.request("GET", "/health/liveliness") - liveliness_seconds: Final = time.perf_counter() - liveliness_started - response, decide_seconds = timed.result() - record_property("decide_seconds", f"{decide_seconds:.6f}") - record_property("liveliness_seconds", f"{liveliness_seconds:.6f}") - assert liveliness.status_code == 200, liveliness.text - assert 500 <= response.status_code < 600, response.text - assert "litellm.APIConnectionError" in response.text, response.text[:2000] - assert decide_seconds < 5, decide_seconds - assert liveliness_seconds < 2, liveliness_seconds + pending: Final = pool.submit(timed, lambda: _decide(gateway, model, api_base=worst)) + liveliness: Final = timed(lambda: gateway.request("GET", "/health/liveliness")) + decided: Final = pending.result() + record_property("decide_seconds", f"{decided.seconds:.6f}") + record_property("liveliness_seconds", f"{liveliness.seconds:.6f}") + assert liveliness.value.status_code == 200, liveliness.value.text + assert 500 <= decided.value.status_code < 600, decided.value.text + assert "litellm.APIConnectionError" in decided.value.text, decided.value.text[:2000] + assert_within_factor_of_control(decided.seconds, control=answered.seconds, factor=_LINEAR_FACTOR) + assert_within_factor_of_control(liveliness.seconds, control=answered.seconds, factor=_LINEAR_FACTOR) assert wire.drain() == () - - -def _timed_decide(gateway: Gateway, model: str, api_base: str) -> tuple[httpx.Response, float]: - started: Final = time.perf_counter() - response: Final = _decide(gateway, model, api_base=api_base) - return response, time.perf_counter() - started diff --git a/tests/integration/providers/test_validated_reply_boundaries_under_faults.py b/tests/integration/providers/test_validated_reply_boundaries_under_faults.py index 0a358e685f6..ae658554668 100644 --- a/tests/integration/providers/test_validated_reply_boundaries_under_faults.py +++ b/tests/integration/providers/test_validated_reply_boundaries_under_faults.py @@ -1,7 +1,6 @@ import base64 import json import re -import signal import threading import uuid from collections import Counter @@ -449,11 +448,11 @@ def test_a_worker_killed_mid_burst_leaves_the_sibling_validating_provider_payloa eventually(held.qsize, lambda size: size == len(calls), seconds=_GATE_SECONDS) held_by: Final = {pid: _open_upstream_connections(pid, wire.url) for pid in workers} victim_pid, survivor_pid = sorted(workers, key=held_by.__getitem__) - victim: Final = psutil.Process(victim_pid) - victim.suspend() - victim.send_signal(signal.SIGKILL) - release.set() - outcomes: Final = tuple(future.result(timeout=60) for future in futures) + with owned.replacing_worker(victim_pid) as kill: + kill.freeze() + kill() + release.set() + outcomes: Final = tuple(future.result(timeout=60) for future in futures) assert sum(held_by.values()) == len(calls), held_by served: Final = tuple( (call, response) for call, response in zip(calls, outcomes, strict=True) if response is not None diff --git a/tests/integration/providers/test_vertex_anthropic_inline_tools_beta_wire.py b/tests/integration/providers/test_vertex_anthropic_inline_tools_beta_wire.py index 5061f0b34a4..a6b039dc8f7 100644 --- a/tests/integration/providers/test_vertex_anthropic_inline_tools_beta_wire.py +++ b/tests/integration/providers/test_vertex_anthropic_inline_tools_beta_wire.py @@ -1,7 +1,6 @@ import asyncio import json import re -import signal import threading import uuid from collections.abc import Callable, Mapping @@ -564,11 +563,12 @@ def test_worker_sigkill_mid_burst_leaves_the_replacement_forwarding_inline_tools shares: Final = {pid: _accepted_client_ports(pid, proxy_url.port or 0) & frozenset(ports) for pid in workers} assert sum(map(len, shares.values())) == len(clients), shares victim: Final = min((pid for pid in workers if shares[pid]), key=lambda pid: len(shares[pid])) - psutil.Process(victim).send_signal(signal.SIGKILL) - release.set() - results: Final = tuple(future.result(timeout=60) for future in futures) - for port, result in zip(ports, results, strict=True): - assert result == (None if port in shares[victim] else 200), (port, result, shares) + with owned.replacing_worker(victim) as kill: + kill() + release.set() + results: Final = tuple(future.result(timeout=60) for future in futures) + for port, result in zip(ports, results, strict=True): + assert result == (None if port in shares[victim] else 200), (port, result, shares) eventually(lambda: len(_STARTED_WORKER.findall(owned.log.read_text())), lambda started: started >= 3, 120) second_wave: Final = tuple(_marker("after kill") for _ in range(6)) for marker in second_wave: diff --git a/tests/integration/providers/test_vertex_gemini_fragmented_stream_wire.py b/tests/integration/providers/test_vertex_gemini_fragmented_stream_wire.py index 449c5c9c105..95a5f881ba1 100644 --- a/tests/integration/providers/test_vertex_gemini_fragmented_stream_wire.py +++ b/tests/integration/providers/test_vertex_gemini_fragmented_stream_wire.py @@ -1,12 +1,14 @@ import json -import time +from collections.abc import Callable +from functools import partial from typing import Final import pytest from cryptography.hazmat.primitives import serialization from cryptography.hazmat.primitives.asymmetric import rsa -from integration._support.client import Gateway -from integration._support.wire import Reply, Request, wire_server +from integration._support.client import Gateway, Scenario +from integration._support.timing import assert_within_factor_of_control, timed +from integration._support.wire import Reply, Request, Wire, wire_server from pydantic import BaseModel, ConfigDict, JsonValue, TypeAdapter _BACKEND: Final = "gemini-3.7-flash" @@ -17,6 +19,7 @@ _PROMPT: Final = "Write a very long numbered list." _PART_COUNT: Final = 8000 _LINES_PER_FRAGMENT: Final = 64 _STREAM_BUDGET_SECONDS: Final = 10.0 +_FRAGMENTED_FACTOR: Final = 10 _JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue]) @@ -64,8 +67,8 @@ def _expected_text() -> str: return "".join(f"{index}. item\n" for index in range(_PART_COUNT)) -def _gemini_response_fragments() -> tuple[bytes, ...]: - document: Final = json.dumps( +def _gemini_document() -> str: + return json.dumps( { "candidates": [ { @@ -81,18 +84,18 @@ def _gemini_response_fragments() -> tuple[bytes, ...]: }, indent=2, ) + + +def _gemini_response_fragments(document: str, lines_per_fragment: int) -> tuple[bytes, ...]: lines: Final = document.split("\n") fragments: Final = tuple( - "\n".join(lines[start : start + _LINES_PER_FRAGMENT]).encode() + b"\n" - for start in range(0, len(lines), _LINES_PER_FRAGMENT) + "\n".join(lines[start : start + lines_per_fragment]).encode() + b"\n" + for start in range(0, len(lines), lines_per_fragment) ) return (b"data: " + fragments[0], *fragments[1:], b"\n") -@pytest.mark.covers("providers.vertex_gemini.fragmented_stream_json_is_parsed_once_and_stays_live") -def test_vertex_gemini_stream_split_across_many_fragments_completes_without_stalling(gateway: Gateway) -> None: - fragments: Final = _gemini_response_fragments() - +def _scripted_vertex(fragments: tuple[bytes, ...]) -> Callable[[Request], Reply]: def respond(request: Request) -> Reply: assert request.method == "POST" assert request.target == f"{_MODEL_PATH}:streamGenerateContent?alt=sse" @@ -102,37 +105,60 @@ def test_vertex_gemini_stream_split_across_many_fragments_completes_without_stal assert body["generationConfig"] == {"temperature": 0.0} return Reply(content_type="text/event-stream", chunks=fragments) - with wire_server(respond) as wire, gateway.scenario() as scenario: - model: Final = scenario.model( - model=f"vertex_ai/{_BACKEND}", - api_base=f"{wire.url}{_MODEL_PATH}", - api_key=None, - vertex_project=_PROJECT, - vertex_location=_LOCATION, - vertex_credentials=_service_account_json(gateway.upstream_url.rstrip("/")), - ) - started: Final = time.monotonic() - with gateway.client.stream( - "POST", - "/v1/chat/completions", - json={ - "model": model, - "messages": [{"role": "user", "content": _PROMPT}], - "stream": True, - "temperature": 0.0, - }, - headers={"Authorization": f"Bearer {gateway.key}"}, - timeout=_STREAM_BUDGET_SECONDS, - ) as response: - assert response.status_code == 200, response.read() - lines: Final = tuple(line for line in response.iter_lines() if line.startswith("data: ")) - elapsed: Final = time.monotonic() - started - assert elapsed < _STREAM_BUDGET_SECONDS, f"stream took {elapsed:.1f}s for {len(fragments)} fragments" - assert lines[-1] == "data: [DONE]", lines[-3:] - chunks: Final = tuple(_Chunk.model_validate_json(line.removeprefix("data: ")) for line in lines[:-1]) - choices: Final = tuple(choice for chunk in chunks for choice in chunk.choices) - assert "".join(choice.delta.content or "" for choice in choices) == _expected_text() - assert tuple(choice.finish_reason for choice in choices if choice.finish_reason) == ("stop",) + return respond + + +def _vertex_model(gateway: Gateway, scenario: Scenario, wire: Wire) -> str: + return scenario.model( + model=f"vertex_ai/{_BACKEND}", + api_base=f"{wire.url}{_MODEL_PATH}", + api_key=None, + vertex_project=_PROJECT, + vertex_location=_LOCATION, + vertex_credentials=_service_account_json(gateway.upstream_url.rstrip("/")), + ) + + +def _streamed_lines(gateway: Gateway, model: str) -> tuple[str, ...]: + with gateway.client.stream( + "POST", + "/v1/chat/completions", + json={ + "model": model, + "messages": [{"role": "user", "content": _PROMPT}], + "stream": True, + "temperature": 0.0, + }, + headers={"Authorization": f"Bearer {gateway.key}"}, + timeout=_STREAM_BUDGET_SECONDS, + ) as response: + assert response.status_code == 200, response.read() + return tuple(line for line in response.iter_lines() if line.startswith("data: ")) + + +def _assert_streamed_the_whole_list(lines: tuple[str, ...]) -> None: + assert lines[-1] == "data: [DONE]", lines[-3:] + chunks: Final = tuple(_Chunk.model_validate_json(line.removeprefix("data: ")) for line in lines[:-1]) + choices: Final = tuple(choice for chunk in chunks for choice in chunk.choices) + assert "".join(choice.delta.content or "" for choice in choices) == _expected_text() + assert tuple(choice.finish_reason for choice in choices if choice.finish_reason) == ("stop",) + + +@pytest.mark.covers("providers.vertex_gemini.fragmented_stream_json_is_parsed_once_and_stays_live") +def test_vertex_gemini_stream_split_across_many_fragments_completes_without_stalling(gateway: Gateway) -> None: + document: Final = _gemini_document() + whole: Final = _gemini_response_fragments(document, document.count("\n") + 1) + fragments: Final = _gemini_response_fragments(document, _LINES_PER_FRAGMENT) + with ( + wire_server(_scripted_vertex(whole)) as control_wire, + wire_server(_scripted_vertex(fragments)) as wire, + gateway.scenario() as scenario, + ): + control: Final = timed(partial(_streamed_lines, gateway, _vertex_model(gateway, scenario, control_wire))) + _assert_streamed_the_whole_list(control.value) + streamed: Final = timed(partial(_streamed_lines, gateway, _vertex_model(gateway, scenario, wire))) + assert_within_factor_of_control(streamed.seconds, control=control.seconds, factor=_FRAGMENTED_FACTOR) + _assert_streamed_the_whole_list(streamed.value) assert [(request.method, request.target) for request in wire.drain()] == [ ("POST", f"{_MODEL_PATH}:streamGenerateContent?alt=sse") ] diff --git a/tests/integration/providers/test_vertex_partner_count_tokens_wire.py b/tests/integration/providers/test_vertex_partner_count_tokens_wire.py index 5df8dd6e32e..558ee3a004e 100644 --- a/tests/integration/providers/test_vertex_partner_count_tokens_wire.py +++ b/tests/integration/providers/test_vertex_partner_count_tokens_wire.py @@ -1,6 +1,5 @@ import json import re -import signal import socket import threading import uuid @@ -663,15 +662,16 @@ def test_worker_sigkill_mid_burst_leaves_the_sibling_counting(gateway: Gateway, shares: Final = {pid: _accepted_client_ports(pid, proxy_url.port or 0) & frozenset(ports) for pid in workers} assert sum(map(len, shares.values())) == len(clients), shares victim: Final = min((pid for pid in workers if shares[pid]), key=lambda pid: len(shares[pid])) - psutil.Process(victim).send_signal(signal.SIGKILL) - release.set() - results: Final = tuple(future.result(timeout=60) for future in futures) - for port, result in zip(ports, results, strict=True): - assert result == (None if port in shares[victim] else (200, _PEER_COUNT)), (port, result, shares) - second_wave: Final = _clients(stack, str(proxy_url), 6) - assert tuple(_counted_on(client, gateway.key, _full(_OWNED_MODEL)) for client in second_wave) == ( - (200, _PEER_COUNT), - ) * len(second_wave) + with owned.replacing_worker(victim) as kill: + kill() + release.set() + results: Final = tuple(future.result(timeout=60) for future in futures) + for port, result in zip(ports, results, strict=True): + assert result == (None if port in shares[victim] else (200, _PEER_COUNT)), (port, result, shares) + second_wave: Final = _clients(stack, str(proxy_url), 6) + assert tuple(_counted_on(client, gateway.key, _full(_OWNED_MODEL)) for client in second_wave) == ( + (200, _PEER_COUNT), + ) * len(second_wave) assert len(_counted_bodies(wire)) == len(clients) + len(second_wave) eventually(lambda: len(_STARTED_WORKER.findall(owned.log.read_text())), lambda started: started >= 3, 120) assert owned.process.poll() is None diff --git a/tests/integration/providers/test_voyage_mongodb_host_wire.py b/tests/integration/providers/test_voyage_mongodb_host_wire.py index 1a6ed907c9b..926bd5a7c69 100644 --- a/tests/integration/providers/test_voyage_mongodb_host_wire.py +++ b/tests/integration/providers/test_voyage_mongodb_host_wire.py @@ -2,7 +2,6 @@ import asyncio import json import os import re -import signal import socket import uuid from collections.abc import Callable, Iterator @@ -11,7 +10,6 @@ from pathlib import Path from typing import Final import httpx -import psutil import pytest import yaml from integration._support.client import JSON_OBJECT, Gateway, eventually, gateway_from_environment, object_value @@ -675,12 +673,12 @@ async def test_worker_sigkill_mid_traffic_leaves_the_sibling_routing_by_key(gate _burst(str(candidate.client.base_url), candidate.key, 20, tolerate_transport_errors=True) ) await asyncio.to_thread(eventually, lambda: forward.received.qsize(), lambda size: size >= 4, 60) - victim: Final = psutil.Process(workers[0]) - victim.suspend() - victim.send_signal(signal.SIGKILL) - served: Final = await burst - for probe in served: - assert probe.response.status_code == REFUSED_STATUS, probe.response.text + with owned.replacing_worker(workers[0]) as kill: + kill.freeze() + await asyncio.to_thread(kill) + served: Final = await burst + for probe in served: + assert probe.response.status_code == REFUSED_STATUS, probe.response.text assert set(forward.targets()) <= {MONGODB_HOST, VOYAGE_HOST} follow_up: Final = _rerank(candidate, MONGODB_RERANK) _assert_dialed(forward, follow_up, MONGODB_HOST) diff --git a/tests/integration/routing/test_audit_any_sweep_chaos.py b/tests/integration/routing/test_audit_any_sweep_chaos.py index 1e46299e662..fe95d528fd5 100644 --- a/tests/integration/routing/test_audit_any_sweep_chaos.py +++ b/tests/integration/routing/test_audit_any_sweep_chaos.py @@ -1,8 +1,6 @@ from __future__ import annotations import json -import re -import signal from collections.abc import Mapping from concurrent.futures import ThreadPoolExecutor from pathlib import Path @@ -10,7 +8,6 @@ from types import MappingProxyType from typing import Final import httpx -import psutil import pytest from integration._support.client import Gateway, eventually, gateway_from_environment from integration._support.process import graceful_stop_seconds, owned_proxy_process @@ -139,13 +136,6 @@ def _fire( return kind, 200, response_id if isinstance(response_id, str) else None -_WORKER_PID: Final = re.compile(r"Started server process \[(\d+)\]") - - -def _worker_pids(log: Path) -> tuple[int, ...]: - return tuple(int(match) for match in _WORKER_PID.findall(log.read_text())) - - @pytest.mark.timeout(2 * graceful_stop_seconds() + 120) def test_proxy_survives_worker_kill_mid_burst(tmp_path: Path) -> None: with gateway_from_environment() as upstream_gateway: @@ -181,16 +171,17 @@ def test_proxy_survives_worker_kill_mid_burst(tmp_path: Path) -> None: "responses-stream": responses_stream_model, } ) - workers: Final = eventually(lambda: _worker_pids(owned.log), lambda pids: len(pids) == 2, seconds=30) + workers: Final = eventually(owned.worker_pids, lambda pids: len(pids) == 2, seconds=30) records: dict[int, tuple[str, int | None, str | None]] = {} with ThreadPoolExecutor(max_workers=30) as pool: futures: Final = [ (index, pool.submit(_fire, owned.gateway, models, key, _burst_kind(index), str(index))) for index in range(30) ] - psutil.Process(workers[0]).send_signal(signal.SIGKILL) - for index, future in futures: - records[index] = future.result() + with owned.replacing_worker(min(workers)) as kill: + kill() + for index, future in futures: + records[index] = future.result() proxy_stamped: Final = [ record[2] for record in records.values() diff --git a/tests/integration/routing/test_complexity_router_databricks_classifier.py b/tests/integration/routing/test_complexity_router_databricks_classifier.py index 419b786bfc1..af4b76e888f 100644 --- a/tests/integration/routing/test_complexity_router_databricks_classifier.py +++ b/tests/integration/routing/test_complexity_router_databricks_classifier.py @@ -2,7 +2,6 @@ import asyncio import json import os import re -import signal import threading import time import uuid @@ -1008,18 +1007,18 @@ async def test_worker_sigkill_mid_burst_leaves_the_sibling_classifying_through_t held_by: Final = MappingProxyType({pid: _open_upstream_connections(pid, judge.url) for pid in workers}) assert sum(held_by.values()) == 20, held_by victim_pid, survivor_pid = sorted(workers, key=held_by.__getitem__) - victim: Final = psutil.Process(victim_pid) - victim.suspend() - victim.send_signal(signal.SIGKILL) - release.set() - served: Final = await burst - assert len(served) == held_by[survivor_pid], (held_by, len(served)) - assert [item.cause for item in served] == ["jev_classifier"] * len(served), served - follow_up_prompt: Final = _prompt() - (answered,) = await _burst( - base_url, candidate.key, router, (("/v1/chat/completions", False),), (follow_up_prompt,) - ) - assert (answered.cause, answered.text) == ("jev_classifier", _COMPLEX_TEXT), answered + with owned.replacing_worker(victim_pid) as kill: + kill.freeze() + await asyncio.to_thread(kill) + release.set() + served: Final = await burst + assert len(served) == held_by[survivor_pid], (held_by, len(served)) + assert [item.cause for item in served] == ["jev_classifier"] * len(served), served + follow_up_prompt: Final = _prompt() + (answered,) = await _burst( + base_url, candidate.key, router, (("/v1/chat/completions", False),), (follow_up_prompt,) + ) + assert (answered.cause, answered.text) == ("jev_classifier", _COMPLEX_TEXT), answered await asyncio.to_thread( eventually, lambda: _worker_startups(owned.log), lambda found: len(found[0]) == 3 and found[1] == 3, 180 ) diff --git a/tests/integration/routing/test_fallback_hop_foreign_encrypted_reasoning.py b/tests/integration/routing/test_fallback_hop_foreign_encrypted_reasoning.py index bea36ee6bf8..32c7b99d5f1 100644 --- a/tests/integration/routing/test_fallback_hop_foreign_encrypted_reasoning.py +++ b/tests/integration/routing/test_fallback_hop_foreign_encrypted_reasoning.py @@ -19,8 +19,6 @@ serving the group with one deployment until the next resync. import asyncio import json -import re -import signal import time import uuid from collections.abc import AsyncIterator, Callable, Iterator, Mapping, Sequence @@ -34,7 +32,6 @@ from typing import Final import anthropic import httpx import openai -import psutil import pytest from integration._support.client import ( JSON_OBJECT, @@ -1249,13 +1246,6 @@ def test_order_one_outage_mid_burst_hops_stripped_and_recovers_after_restart(gat assert [json.loads(request.body)["input"] for request in restarted.drain()] == [history(recovered)] -WORKER_PID: Final = re.compile(r"Started server process \[(\d+)\]") - - -def worker_pids(log: Path) -> tuple[int, ...]: - return tuple(int(match) for match in WORKER_PID.findall(log.read_text())) - - def assert_completed_shots_hopped_stripped( gateway: Gateway, groups: Mapping[str, OrderedGroup], fired: Sequence[Fired] ) -> None: @@ -1283,17 +1273,18 @@ def test_worker_killed_mid_burst_leaves_the_survivor_hopping_stripped(tmp_path: owned.gateway.post("/config/update", {"router_settings": {"num_retries": 0}}) with owned.gateway.scenario() as scenario: groups: Final = burst_groups(scenario) - workers: Final = eventually(lambda: worker_pids(owned.log), lambda pids: len(pids) == 2, seconds=60) + workers: Final = eventually(owned.worker_pids, lambda pids: len(pids) == 2, seconds=60) shots: Final = tuple((BURST_KINDS[index % len(BURST_KINDS)], marker()) for index in range(30)) - with lanes(owned.gateway, [group.target for group in groups.values()], 30) as pinned: - with ThreadPoolExecutor(max_workers=30) as pool: - futures: Final = [ - pool.submit(fire, lane, groups, kind, mark) - for lane, (kind, mark) in zip(pinned, shots, strict=True) - ] - psutil.Process(workers[0]).send_signal(signal.SIGKILL) - fired: Final = tuple(future.result() for future in futures) - assert_completed_shots_hopped_stripped(owned.gateway, groups, fired) - probes: Final = tuple(fire(owned.gateway, groups, kind, marker()) for kind in BURST_KINDS) - assert all(probe.status == 200 for probe in probes), probes - assert_burst_hopped_stripped(owned.gateway, groups, probes) + with owned.replacing_worker(min(workers)) as kill: + with lanes(owned.gateway, [group.target for group in groups.values()], 30) as pinned: + with ThreadPoolExecutor(max_workers=30) as pool: + futures: Final = [ + pool.submit(fire, lane, groups, kind, mark) + for lane, (kind, mark) in zip(pinned, shots, strict=True) + ] + kill() + fired: Final = tuple(future.result() for future in futures) + assert_completed_shots_hopped_stripped(owned.gateway, groups, fired) + probes: Final = tuple(fire(owned.gateway, groups, kind, marker()) for kind in BURST_KINDS) + assert all(probe.status == 200 for probe in probes), probes + assert_burst_hopped_stripped(owned.gateway, groups, probes) diff --git a/tests/integration/routing/test_payment_required_cooldown.py b/tests/integration/routing/test_payment_required_cooldown.py index a7b80392636..12190efee26 100644 --- a/tests/integration/routing/test_payment_required_cooldown.py +++ b/tests/integration/routing/test_payment_required_cooldown.py @@ -2,7 +2,6 @@ from __future__ import annotations import json import re -import signal import uuid from collections.abc import Callable, Iterator, Mapping, Sequence from concurrent.futures import ThreadPoolExecutor @@ -14,7 +13,6 @@ from typing import Final import anthropic import httpx import openai -import psutil import pytest import yaml from integration._support.client import GATEWAY_LIMITS, Gateway, Scenario, eventually, gateway_from_environment @@ -54,7 +52,6 @@ COOLED_DOWN: Final = "litellm_deployment_cooled_down_total" PRIMARY_GROUP: Final = new_group() FALLBACK_GROUP: Final = new_group() MARKER: Final = re.compile(rb"pr402-[0-9a-f]{32}(?:-\d+)?") -WORKER_PID: Final = re.compile(r"Started server process \[(\d+)\]") @dataclass(frozen=True, slots=True) @@ -506,10 +503,6 @@ def burst_call(group: str, index: int) -> Call: return Call(marker, "/v1/responses", {"model": group, "input": marker, "stream": stream}) -def worker_pids(log: Path) -> tuple[int, ...]: - return tuple(int(pid) for pid in WORKER_PID.findall(log.read_text())) - - def ready_workers(log: Path) -> int: return log.read_text().count("Application startup complete.") @@ -552,12 +545,7 @@ def test_chaos_burst_cooldown_holds_across_both_workers_and_a_worker_kill(tmp_pa assert_served(warm, chat_identity) assert markers_received(refusing) == () assert markers_received(healthy) == markers(warm) - victim: Final = next(pid for pid in sorted(worker_pids(owned.log)) if psutil.pid_exists(pid)) - ready_before: Final = ready_workers(owned.log) - psutil.Process(victim).send_signal(signal.SIGKILL) - eventually( - lambda: ready_workers(owned.log), lambda ready: ready > ready_before, seconds=graceful_stop_seconds() - ) + owned.kill_worker(min(owned.worker_pids())) after_kill: Final = tuple(raw_chat(patient, group, new_marker()) for _ in range(20)) assert_served(after_kill, chat_identity) assert markers_received(refusing) == () diff --git a/tests/integration/routing/test_per_model_rate_limit_fallbacks.py b/tests/integration/routing/test_per_model_rate_limit_fallbacks.py index 428a485f397..67abc211900 100644 --- a/tests/integration/routing/test_per_model_rate_limit_fallbacks.py +++ b/tests/integration/routing/test_per_model_rate_limit_fallbacks.py @@ -2,7 +2,6 @@ from __future__ import annotations import json import re -import signal import uuid from collections import Counter from collections.abc import Iterator, Mapping, Sequence @@ -16,20 +15,18 @@ from typing import Final import anthropic import httpx import openai -import psutil import pytest import yaml from integration._support.client import ( JSON_OBJECT, Gateway, Scenario, - eventually, gateway_from_environment, object_value, string_value, ) from integration._support.openai_wire import chat_reply, responses_reply -from integration._support.process import graceful_stop_seconds, owned_proxy_process +from integration._support.process import OwnedProxy, graceful_stop_seconds, owned_proxy_process from integration._support.wire import Reply, Request, Wire, wire_server from integration.routing.test_payment_required_mapping import OPENAI_MODEL, is_model_info_probe, model_list_reply from pydantic import JsonValue, TypeAdapter @@ -39,10 +36,7 @@ pytestmark: Final = pytest.mark.timeout(2 * graceful_stop_seconds() + 120) SETTING: Final = "disable_fallbacks_on_per_model_rate_limits" FRESH_CONNECTION: Final = MappingProxyType({"Connection": "close"}) MARKER: Final = re.compile(rb"pmrl-[0-9a-f]{32}") -WORKER_STARTED: Final = re.compile(r"Started server process \[(\d+)\]") -STARTUP_COMPLETE: Final = "Application startup complete." CONFIG_LIST: Final = TypeAdapter(list[dict[str, JsonValue]]) -WORKER_PIDS: Final = TypeAdapter(tuple[int, ...]) BURST: Final = 8 TOKEN_LIMIT: Final = 100 INPUT_TOKEN_LIMIT: Final = 120 @@ -75,10 +69,14 @@ class Groups: @dataclass(frozen=True, slots=True) class Rig: gateway: Gateway - log: Path + owned: OwnedProxy groups: Groups wires: Mapping[str, Wire] + @property + def log(self) -> Path: + return self.owned.log + def traffic(self) -> dict[str, tuple[str, ...]]: return {group: markers_on(wire) for group, wire in self.wires.items()} @@ -221,7 +219,7 @@ def enforcing(tmp_path_factory: pytest.TempPathFactory) -> Iterator[Rig]: rig_gateway, directory, {}, config=enforcing_config(directory, groups, wires), workers=2 ) ) - yield Rig(owned.gateway, owned.log, groups, wires) + yield Rig(owned.gateway, owned, groups, wires) def chat(gateway: Gateway, key: str, group: str, content: str, **extra: JsonValue) -> httpx.Response: @@ -326,11 +324,6 @@ def served_once_in_a_burst(rig: Rig, markers: Sequence[str], responses: Sequence return served -def worker_startups(log: Path) -> tuple[tuple[int, ...], int]: - text: Final = log.read_text(errors="replace") - return WORKER_PIDS.validate_python(WORKER_STARTED.findall(text)), text.count(STARTUP_COMPLETE) - - def listed_setting(gateway: Gateway) -> tuple[dict[str, JsonValue], ...]: response: Final = gateway.request("GET", "/config/list", params={"config_type": "general_settings"}) assert response.status_code == 200, response.text @@ -604,16 +597,7 @@ def test_yaml_owned_setting_refuses_a_database_write_and_keeps_enforcing(enforci def test_worker_sigkill_and_replacement_keeps_refusing_per_model_overflow_without_fallback(enforcing: Rig) -> None: enforcing.traffic() - pids, ready = eventually( - lambda: worker_startups(enforcing.log), lambda found: len(found[0]) >= 2 and found[1] >= 2, seconds=60 - ) - victim: Final = next(pid for pid in pids if psutil.pid_exists(pid)) - psutil.Process(victim).send_signal(signal.SIGKILL) - eventually( - lambda: worker_startups(enforcing.log), - lambda found: len(found[0]) > len(pids) and found[1] > ready, - seconds=graceful_stop_seconds(), - ) + enforcing.owned.kill_worker(min(enforcing.owned.worker_pids())) primary: Final = enforcing.groups.primary markers: Final = tuple(new_marker() for _ in range(BURST)) with enforcing.gateway.scenario() as scenario: diff --git a/tests/integration/routing/test_per_model_rate_limit_fallbacks_lifecycle.py b/tests/integration/routing/test_per_model_rate_limit_fallbacks_lifecycle.py index 4361843b2c9..ea47b590f9b 100644 --- a/tests/integration/routing/test_per_model_rate_limit_fallbacks_lifecycle.py +++ b/tests/integration/routing/test_per_model_rate_limit_fallbacks_lifecycle.py @@ -67,7 +67,7 @@ def fallback_rig( owned: Final = stack.enter_context( owned_proxy_process(rig_gateway, directory, environment, config=config, workers=2) ) - yield Rig(owned.gateway, owned.log, groups, wires) + yield Rig(owned.gateway, owned, groups, wires) @pytest.fixture(scope="module") @@ -185,7 +185,7 @@ def test_database_setting_survives_a_restart_and_refuses_per_model_overflow(tmp_ ) listed: Final = listed_setting(restarted.gateway) assert [(entry["field_value"], entry["stored_in_db"]) for entry in listed] == [(True, True)], listed - rig: Final = Rig(restarted.gateway, restarted.log, groups, wires) + rig: Final = Rig(restarted.gateway, restarted, groups, wires) rig.traffic() key: Final = string_value(rig.gateway.post("/key/generate", {"model_rpm_limit": {groups.primary: 1}})["key"]) first, second = new_marker(), new_marker() diff --git a/tests/integration/routing/test_priority_scheduler_queue_cleanup.py b/tests/integration/routing/test_priority_scheduler_queue_cleanup.py index 3c38acadaaf..bcfb28f790a 100644 --- a/tests/integration/routing/test_priority_scheduler_queue_cleanup.py +++ b/tests/integration/routing/test_priority_scheduler_queue_cleanup.py @@ -4,7 +4,6 @@ import asyncio import http.client import json import re -import signal import threading import time import uuid @@ -27,6 +26,7 @@ from integration._support.database import read_rows from integration._support.openai_wire import answering_model_discovery, chat_reply, openai_error, responses_reply from integration._support.process import graceful_stop_seconds, owned_proxy_process from integration._support.redis_process import OwnedRedis, owned_redis +from integration._support.timing import assert_cut_short from integration._support.wire import Reply, Request, Wire, wire_server from pydantic import JsonValue, TypeAdapter from redis import Redis @@ -624,7 +624,7 @@ def test_fresh_queue_during_a_cooldown_is_refused_at_once(pair: Pair) -> None: client.completions.create(model=group, prompt=new_marker(), extra_body={"priority": 1}) assert refused.value.status_code == 429, refused.value.message assert NO_DEPLOYMENTS in refused.value.message - assert time.monotonic() - started < PROMPT_SECONDS + assert_cut_short(time.monotonic() - started, injected_delay=ROUTER_TIMEOUT_SECONDS) assert queue_entries(pair.cache, group) == [] @@ -825,18 +825,18 @@ def test_sibling_worker_keeps_serving_prioritized_requests_after_a_worker_is_kil held_by: Final = {pid: established_upstream_connections(pid, wire) for pid in workers} assert sum(held_by.values()) == 20, held_by victim_pid, survivor_pid = sorted(workers, key=held_by.__getitem__) - victim: Final = psutil.Process(victim_pid) - victim.suspend() - victim.send_signal(signal.SIGKILL) - upstream.release.set() - outcomes: Final = pending.result(timeout=90) - answered: Final = tuple(outcome for outcome in outcomes if isinstance(outcome[2], Answer)) - failed: Final = tuple(outcome for outcome in outcomes if not isinstance(outcome[2], Answer)) - assert len(answered) == held_by[survivor_pid], (held_by, len(answered)) - assert len(failed) == held_by[victim_pid], (held_by, len(failed)) - assert_all_served(answered) - follow_up: Final = new_marker() - served_eventually(owned.gateway, KILL_GROUP, follow_up) + with owned.replacing_worker(victim_pid) as kill: + kill.freeze() + kill() + upstream.release.set() + outcomes: Final = pending.result(timeout=90) + answered: Final = tuple(outcome for outcome in outcomes if isinstance(outcome[2], Answer)) + failed: Final = tuple(outcome for outcome in outcomes if not isinstance(outcome[2], Answer)) + assert len(answered) == held_by[survivor_pid], (held_by, len(answered)) + assert len(failed) == held_by[victim_pid], (held_by, len(failed)) + assert_all_served(answered) + follow_up: Final = new_marker() + served_eventually(owned.gateway, KILL_GROUP, follow_up) eventually( lambda: len(STARTED_WORKER.findall(owned.log.read_text())), lambda count: count == 3, diff --git a/tests/integration/run.py b/tests/integration/run.py index 6dcfe3900d7..43e606f46a7 100644 --- a/tests/integration/run.py +++ b/tests/integration/run.py @@ -1,6 +1,7 @@ from __future__ import annotations import argparse +import importlib.util import json import os import subprocess @@ -8,7 +9,7 @@ import sys from collections.abc import Mapping from dataclasses import dataclass from pathlib import Path -from types import MappingProxyType +from types import MappingProxyType, ModuleType from typing import Final GROUPS: Final = MappingProxyType( @@ -77,6 +78,16 @@ def worker_settings(group: str, workers: int) -> Mapping[str, str]: ) +def _budget() -> ModuleType: + specification: Final = importlib.util.spec_from_file_location( + "integration_budget", Path(__file__).with_name("_support") / "budget.py" + ) + assert specification is not None and specification.loader is not None + module: Final = importlib.util.module_from_spec(specification) + specification.loader.exec_module(module) + return module + + def main() -> int: parser: Final = argparse.ArgumentParser() parser.add_argument("group", choices=tuple(GROUPS)) @@ -116,6 +127,12 @@ def main() -> int: "PYTHON_DOTENV_DISABLED": "1", **worker_settings(options.group, options.workers), } + budget: Final = _budget() + proxy_root: Final = Path(os.environ.get("INTEGRATION_PROXY_ROOT") or root) + imported: Final = float( + os.environ.get(budget.IMPORT_SECONDS_VARIABLE) or budget.measure_import_seconds(proxy_root, environment) + ) + (output / "boot-calibration.json").write_text(budget.calibration(imported) + "\n") result: Final = subprocess.call( [ sys.executable, @@ -129,7 +146,7 @@ def main() -> int: "no:pytest-retry", "-p", "no:rerunfailures", - "--timeout=90", + f"--timeout={int(budget.case_timeout_for(imported))}", "--durations=15", "--tb=short", f"--hypothesis-seed={options.seed}", @@ -144,7 +161,7 @@ def main() -> int: *(("--dist=loadgroup",) if options.workers > 1 else ()), ], cwd=root, - env=environment, + env={**environment, budget.IMPORT_SECONDS_VARIABLE: str(imported)}, ) if result != 0: return result diff --git a/tests/integration/sdk/test_redis_cluster_iam_auth.py b/tests/integration/sdk/test_redis_cluster_iam_auth.py index c649d349dae..055b73bb415 100644 --- a/tests/integration/sdk/test_redis_cluster_iam_auth.py +++ b/tests/integration/sdk/test_redis_cluster_iam_auth.py @@ -11,6 +11,7 @@ from unittest.mock import MagicMock, patch import pytest import redis +from tests.integration._support.process import stop_child from litellm._redis import _get_redis_env_kwarg_mapping, get_redis_client from litellm._redis_credential_provider import _token_cache @@ -91,12 +92,7 @@ def authenticated_redis_cluster(tmp_path: Path, unused_tcp_port_factory: Callabl admin.execute_command("ACL", "SETUSER", "default", "resetpass", ">local-fixture-token") yield port finally: - process.terminate() - try: - process.wait(timeout=5) - except subprocess.TimeoutExpired: - process.kill() - process.wait() + stop_child(process, grace_seconds=5) if process.stdout is not None: process.stdout.close() diff --git a/tests/integration/spend/test_background_interaction_settlement.py b/tests/integration/spend/test_background_interaction_settlement.py index 11444396697..75cb2912076 100644 --- a/tests/integration/spend/test_background_interaction_settlement.py +++ b/tests/integration/spend/test_background_interaction_settlement.py @@ -761,12 +761,12 @@ def test_killed_workers_leave_their_polls_to_the_respawned_workers(rig: Rig) -> for item in created: _state(rig, item, InteractionState(status="in_progress")) rates: Final = _rates(resumer.gateway, model) - killed: Final = _worker_pids(resumer.process.pid) + killed: Final = resumer.worker_pids() assert len(killed) == 2, killed - victims: Final = tuple(psutil.Process(pid) for pid in killed) - for victim in victims: - victim.kill() - psutil.wait_procs(victims, timeout=15) + first, second = sorted(killed) + with resumer.replacing_worker(first) as kill_first, resumer.replacing_worker(second) as kill_second: + kill_first() + kill_second() for item in created: _state(rig, item, _completed()) for item in created: diff --git a/tests/integration/spend/test_daily_activity_key_alias_probes.py b/tests/integration/spend/test_daily_activity_key_alias_probes.py index 8d9b8616435..db22d2402fd 100644 --- a/tests/integration/spend/test_daily_activity_key_alias_probes.py +++ b/tests/integration/spend/test_daily_activity_key_alias_probes.py @@ -1,4 +1,3 @@ -import time import uuid from collections.abc import Callable, Iterator from concurrent.futures import ThreadPoolExecutor @@ -36,10 +35,12 @@ from integration._support.daily_activity import ( ) from integration._support.database import read_rows, scratch_database from integration._support.process import OwnedProxy, owned_proxy_process +from integration._support.timing import assert_enforced_within, timed +from litellm.constants import SPEND_LOG_KEY_METADATA_QUERY_TIMEOUT_MS from pydantic import JsonValue DAY_OUTSIDE_THE_WINDOW: Final = "2026-02-10" -GIVES_UP_WITHIN_SECONDS: Final = 10 +LOCKED_LOOKUP_SLACK_SECONDS: Final = 5 CONCURRENT_READS: Final = 20 CACHED_MISS_CLEARS_WITHIN_SECONDS: Final = 45 ALIAS_OF_ONE_SPEND_LOG: Final = ( @@ -450,16 +451,18 @@ def test_alias_lookup_gives_up_while_spend_logs_are_locked_and_answers_once_they spend_logs_of_key(api_key, (named_row(0, alias),), database_url=database_url), ): with locked_table(SPEND_LOGS_TABLE, database_url=database_url): - started: Final = time.monotonic() - locked: Final = activity_of_key(owned.gateway, AGGREGATED_USER_ACTIVITY, api_key) - waited: Final = time.monotonic() - started + locked: Final = timed(lambda: activity_of_key(owned.gateway, AGGREGATED_USER_ACTIVITY, api_key)) unlocked: Final = eventually( lambda: activity_of_key(owned.gateway, AGGREGATED_USER_ACTIVITY, api_key), _names_the_key(api_key, alias), seconds=CACHED_MISS_CLEARS_WITHIN_SECONDS, ) - assert waited < GIVES_UP_WITHIN_SECONDS, waited - assert_key_reported(locked, api_key, DAY, key_metadata(user=owner, email=email), seeded_metrics(1)) + assert_enforced_within( + locked.seconds, + configured=SPEND_LOG_KEY_METADATA_QUERY_TIMEOUT_MS / 1000, + slack=LOCKED_LOOKUP_SLACK_SECONDS, + ) + assert_key_reported(locked.value, api_key, DAY, key_metadata(user=owner, email=email), seeded_metrics(1)) assert_key_reported(unlocked, api_key, DAY, key_metadata(alias=alias, user=owner, email=email), seeded_metrics(1)) diff --git a/tests/integration/spend/test_daily_activity_key_metadata_query_timeout.py b/tests/integration/spend/test_daily_activity_key_metadata_query_timeout.py index 747570a01ec..780ee43fb3f 100644 --- a/tests/integration/spend/test_daily_activity_key_metadata_query_timeout.py +++ b/tests/integration/spend/test_daily_activity_key_metadata_query_timeout.py @@ -1,7 +1,5 @@ import itertools import json -import os -import signal import threading import uuid from bisect import bisect_left @@ -1025,14 +1023,6 @@ def _running_children(owned: OwnedProxy) -> tuple[int, ...]: ) -def _worker_pids(owned: OwnedProxy) -> tuple[int, ...]: - return tuple( - member.pid - for member in group_members(owned.process.pid) - if member.pid != owned.process.pid and any("spawn_main" in part for part in member.cmdline()) - ) - - @pytest.mark.timeout(600) def test_usage_page_keeps_serving_when_a_worker_dies_mid_outage(gateway: Gateway, tmp_path: Path) -> None: with ( @@ -1042,23 +1032,25 @@ def test_usage_page_keeps_serving_when_a_worker_dies_mid_outage(gateway: Gateway ): spender: Final = _spender(owned.gateway, database_url, "worker-death") children: Final = _running_children(owned) - workers: Final = _worker_pids(owned) + workers: Final = owned.worker_pids() assert len(workers) == WORKERS, workers + victim: Final = min(workers) with _locked_spend_logs(database_url): timed_out: Final = _fresh_read(owned, spender.digest) assert timed_out.status_code == 200, timed_out.text assert timed_out.elapsed >= FAILED_LOOKUP_FLOOR, timed_out.elapsed assert _aliases(timed_out, spender.digest) == (None,), timed_out.text - os.kill(workers[0], signal.SIGKILL) - after_kill: Final = _fresh_read(owned, spender.digest) - assert after_kill.status_code == 200, after_kill.text - assert _aliases(after_kill, spender.digest) == (None,), after_kill.text + with owned.replacing_worker(victim) as kill: + kill() + after_kill: Final = _fresh_read(owned, spender.digest) + assert after_kill.status_code == 200, after_kill.text + assert _aliases(after_kill, spender.digest) == (None,), after_kill.text respawned: Final = eventually( lambda: _running_children(owned), lambda pids: len(pids) >= len(children) and any(pid not in children for pid in pids), seconds=30, ) - assert workers[0] not in respawned, respawned + assert victim not in respawned, respawned recovered: Final = eventually( lambda: _fresh_read(owned, spender.digest), lambda response: response.status_code == 200 and _aliases(response, spender.digest) == (spender.alias,), diff --git a/tests/integration/spend/test_daily_activity_key_owner_faults.py b/tests/integration/spend/test_daily_activity_key_owner_faults.py index 2ab56622ec4..1dba97d6ae1 100644 --- a/tests/integration/spend/test_daily_activity_key_owner_faults.py +++ b/tests/integration/spend/test_daily_activity_key_owner_faults.py @@ -1,6 +1,3 @@ -import os -import signal -import time import uuid from collections.abc import Iterator from contextlib import contextmanager @@ -32,12 +29,14 @@ from integration._support.daily_activity import ( ) from integration._support.database import scratch_database from integration._support.process import OwnedProxy, group_members, owned_proxy_process +from integration._support.timing import assert_enforced_within, timed +from litellm.constants import SPEND_LOG_KEY_METADATA_QUERY_TIMEOUT_MS USER_ACTIVITY: Final = "/user/daily/activity" TEAM_ACTIVITY: Final = "/team/daily/activity" AGGREGATED_TEAM_ACTIVITY: Final = "/team/daily/activity/aggregated" KEYS_OF_ONE_TEAM: Final = 300 -GIVES_UP_WITHIN_SECONDS: Final = 10 +LOCKED_LOOKUP_SLACK_SECONDS: Final = 5 READS_AFTER_THE_WORKER_IS_REPLACED: Final = 6 @@ -214,12 +213,14 @@ def test_owner_lookup_gives_up_while_daily_user_spend_is_locked_and_answers_once rows: Final = (user_row(owner, api_key, DAY), seeded_row(TEAM_SPEND, "team_id", team, api_key, DAY)) with daily_rows(rows, database_url=database_url): with locked_table(USER_SPEND, database_url=database_url): - started: Final = time.monotonic() - locked: Final = activity_of_key(owned.gateway, TEAM_ACTIVITY, api_key, team_ids=team) - waited: Final = time.monotonic() - started + locked: Final = timed(lambda: activity_of_key(owned.gateway, TEAM_ACTIVITY, api_key, team_ids=team)) unlocked: Final = activity_of_key(owned.gateway, TEAM_ACTIVITY, api_key, team_ids=team) - assert waited < GIVES_UP_WITHIN_SECONDS, waited - assert_key_reported(locked, api_key, DAY, key_metadata(), seeded_metrics(1)) + assert_enforced_within( + locked.seconds, + configured=SPEND_LOG_KEY_METADATA_QUERY_TIMEOUT_MS / 1000, + slack=LOCKED_LOOKUP_SLACK_SECONDS, + ) + assert_key_reported(locked.value, api_key, DAY, key_metadata(), seeded_metrics(1)) assert_key_reported(unlocked, api_key, DAY, key_metadata(user=owner, email=email), seeded_metrics(1)) @@ -234,12 +235,11 @@ def test_owner_is_reported_while_a_worker_is_killed_and_after_it_is_replaced(gat member for member in group_members(owned.process.pid) if member.pid != owned.process.pid ) children: Final = tuple(member.pid for member in members) - workers: Final = tuple( - member.pid for member in members if any("spawn_main" in part for part in member.cmdline()) - ) + workers: Final = owned.worker_pids() assert len(workers) >= 2, workers - os.kill(workers[0], signal.SIGKILL) - during: Final = _read_on_a_new_connection(owned.gateway, api_key) + with owned.replacing_worker(min(workers)) as kill: + kill() + during: Final = _read_on_a_new_connection(owned.gateway, api_key) eventually( lambda: _running_children(owned), lambda pids: len(pids) >= len(children) and any(pid not in children for pid in pids), diff --git a/tests/integration/spend/test_daily_tag_spend_retention.py b/tests/integration/spend/test_daily_tag_spend_retention.py index fcb624cb188..0bdb981c878 100644 --- a/tests/integration/spend/test_daily_tag_spend_retention.py +++ b/tests/integration/spend/test_daily_tag_spend_retention.py @@ -1,12 +1,10 @@ import json import os -import signal import uuid from datetime import datetime, timedelta, timezone from pathlib import Path from typing import Final -import psutil import psycopg import pytest import yaml @@ -14,7 +12,7 @@ from pydantic import JsonValue, TypeAdapter from tests.integration._support.client import Gateway, eventually, string_value from tests.integration._support.database import read_rows -from tests.integration._support.process import OwnedProxy, owned_proxy, owned_proxy_process +from tests.integration._support.process import owned_proxy, owned_proxy_process CLEANUP_EVERY_MINUTE: Final = "* * * * *" RETENTION_SETTING: Final = "maximum_daily_tag_spend_retention_period" @@ -82,15 +80,6 @@ def _store_retention_setting(value: JsonValue) -> None: ) -def _listening_workers(owned: OwnedProxy) -> tuple[psutil.Process, ...]: - port: Final = owned.gateway.client.base_url.port - return tuple( - child - for child in psutil.Process(owned.process.pid).children(recursive=True) - if any(conn.status == psutil.CONN_LISTEN and conn.laddr.port == port for conn in child.net_connections("inet")) - ) - - def _listed_retention_value(gateway: Gateway) -> JsonValue: listed: Final = _SETTINGS.validate_json( gateway.request("GET", "/config/list", params={"config_type": "general_settings"}).content @@ -216,13 +205,11 @@ def test_daily_tag_spend_cleanup_completes_after_one_of_two_workers_is_killed(ga with owned_proxy_process(gateway, tmp_path, {}, config=config, workers=2) as owned: with owned.gateway.scenario() as scenario: model: Final = scenario.model() - workers: Final = eventually( - lambda: _listening_workers(owned), lambda found: len(found) == 2, seconds=30 - ) - workers[0].send_signal(signal.SIGKILL) - eventually(lambda: workers[0].is_running(), lambda alive: not alive, seconds=10) - ids: Final = tuple(_completion_id(owned.gateway, model) for _ in range(6)) - assert len(set(ids)) == 6 and all(identity.startswith("chatcmpl-") for identity in ids), ids + workers: Final = eventually(owned.worker_pids, lambda found: len(found) == 2, seconds=30) + with owned.replacing_worker(min(workers)) as kill: + kill() + ids: Final = tuple(_completion_id(owned.gateway, model) for _ in range(6)) + assert len(set(ids)) == 6 and all(identity.startswith("chatcmpl-") for identity in ids), ids remaining: Final = eventually( lambda: _remaining_days(tag), lambda days: expired not in days, seconds=150 ) diff --git a/tests/integration/spend/test_normalized_error_long_message.py b/tests/integration/spend/test_normalized_error_long_message.py index accf160b03f..407b40d36c3 100644 --- a/tests/integration/spend/test_normalized_error_long_message.py +++ b/tests/integration/spend/test_normalized_error_long_message.py @@ -1,7 +1,5 @@ import asyncio import json -import os -import signal import threading import time import uuid @@ -16,18 +14,20 @@ from typing import Final import anthropic import httpx import openai -import psutil import pytest from integration._support.client import Gateway, eventually, object_value, string_value from integration._support.database import read_rows -from integration._support.process import OwnedProxy, owned_proxy_process +from integration._support.process import owned_proxy_process +from integration._support.timing import assert_not_delayed_by, assert_within_factor_of_control, timed from integration._support.wire import Reply, Request, wire_server from pydantic import JsonValue CRAFTED_MODEL: Final = ("exceeded " * 32_000)[:288_000] HOSTILE_5KB_MODEL: Final = ("exceeded budget " * 400)[:5_000] +BENIGN_288KB_MODEL: Final = "x" * len(CRAFTED_MODEL) +CRAFTED_SLOWDOWN_FACTOR: Final = 10 FAST_SECONDS: Final = 10.0 -LIVELINESS_MAX_SECONDS: Final = 5.0 +LIVELINESS_STALL_BUDGET_SECONDS: Final = 5.0 ROW_SECONDS: Final = 70 CHAT: Final = "/v1/chat/completions" MESSAGES: Final = "/v1/messages" @@ -129,15 +129,19 @@ def test_crafted_288kb_model_through_async_openai_sdk_fails_fast(gateway: Gatewa def test_crafted_288kb_model_through_anthropic_sdk_fails_fast(gateway: Gateway) -> None: with gateway.scenario() as scenario: client: Final = anthropic.Anthropic(base_url=str(gateway.client.base_url), api_key=scenario.key(), timeout=120) - started: Final = time.perf_counter() - with pytest.raises(anthropic.BadRequestError) as raised: - client.messages.create( - model=CRAFTED_MODEL, max_tokens=8, messages=[{"role": "user", "content": f"audit {uuid.uuid4().hex}"}] - ) - seconds: Final = time.perf_counter() - started - assert seconds < FAST_SECONDS, f"crafted 288 KB model took {seconds:.2f}s" - assert "Invalid model name passed in" in str(raised.value), str(raised.value)[:300] - info: Final = _error_information(raised.value.response.headers["x-litellm-call-id"]) + + def refused(model: str) -> anthropic.BadRequestError: + with pytest.raises(anthropic.BadRequestError) as raised: + client.messages.create( + model=model, max_tokens=8, messages=[{"role": "user", "content": f"audit {uuid.uuid4().hex}"}] + ) + return raised.value + + control: Final = timed(lambda: refused(BENIGN_288KB_MODEL)) + crafted: Final = timed(lambda: refused(CRAFTED_MODEL)) + assert_within_factor_of_control(crafted.seconds, control=control.seconds, factor=CRAFTED_SLOWDOWN_FACTOR) + assert "Invalid model name passed in" in str(crafted.value), str(crafted.value)[:300] + info: Final = _error_information(crafted.value.response.headers["x-litellm-call-id"]) assert info["normalized_error"] == "400_INVALID_REQUEST", info @@ -151,35 +155,31 @@ def _poll_liveliness(client: httpx.Client, stop: threading.Event) -> list[float] return latencies -def _cmdline(process: psutil.Process) -> str: - try: - return " ".join(process.cmdline()) - except psutil.Error: - return "" - - -def _worker_pids(owned: OwnedProxy) -> tuple[int, ...]: - return tuple(child.pid for child in psutil.Process(owned.process.pid).children() if "spawn_main" in _cmdline(child)) +def _liveliness_while_posting( + client: httpx.Client, model: str, key: str +) -> tuple[tuple[_Timed, ...], tuple[float, ...]]: + stop: Final = threading.Event() + with ThreadPoolExecutor(max_workers=3) as pool: + liveliness: Final = pool.submit(_poll_liveliness, client, stop) + posts: Final = tuple( + pool.submit(_timed_post, client, CHAT, _body(CHAT, model, uuid.uuid4().hex), key) for _ in range(2) + ) + results: Final = tuple(future.result() for future in posts) + stop.set() + return results, tuple(liveliness.result()) def test_two_concurrent_crafted_requests_do_not_stall_liveliness_on_a_two_worker_proxy( gateway: Gateway, tmp_path: Path ) -> None: with owned_proxy_process(gateway, tmp_path, {}, workers=2) as owned, _patient_client(owned.gateway) as client: - assert len(_worker_pids(owned)) == 2, _worker_pids(owned) - stop: Final = threading.Event() - with ThreadPoolExecutor(max_workers=3) as pool: - liveliness: Final = pool.submit(_poll_liveliness, client, stop) - crafted: Final = tuple( - pool.submit(_timed_post, client, CHAT, _body(CHAT, CRAFTED_MODEL, uuid.uuid4().hex), gateway.key) - for _ in range(2) - ) - results: Final = tuple(future.result() for future in crafted) - stop.set() - latencies: Final = liveliness.result() - for timed in results: - _assert_crafted_failure(timed) - assert latencies and max(latencies) < LIVELINESS_MAX_SECONDS, f"liveliness max {max(latencies):.2f}s" + assert len(owned.worker_pids()) == 2, owned.worker_pids() + _benign, control = _liveliness_while_posting(client, BENIGN_288KB_MODEL, gateway.key) + results, latencies = _liveliness_while_posting(client, CRAFTED_MODEL, gateway.key) + for posted in results: + _assert_crafted_failure(posted) + assert latencies and control, (latencies, control) + assert_not_delayed_by(latencies, control=control, injected_delay=LIVELINESS_STALL_BUDGET_SECONDS) def _completion(request: Request) -> Reply: @@ -403,25 +403,28 @@ def test_killing_one_worker_mid_burst_leaves_the_other_serving_crafted_and_happy ): model: Final = scenario.model(api_base=wire.url + "/v1", num_retries=0) key: Final = scenario.key(models=[model]) - workers: Final = _worker_pids(owned) + workers: Final = owned.worker_pids() assert len(workers) == 2, workers + victim: Final = min(workers) - def kill_one_worker() -> None: - os.kill(workers[0], signal.SIGKILL) - - results: Final = _burst(client, model, key, scenario.key(), kill_one_worker) - dropped: Final = tuple(result for result in results if result.status is None) - assert len(dropped) < len(results), results - crafted: Final = tuple(result for result in results if result.label.startswith("crafted") and result.status) - assert crafted and all(result.status == 400 and result.seconds < FAST_SECONDS for result in crafted), crafted - late: Final = tuple(result for result in results if result.label.startswith("late")) - assert all(result.status == 200 for result in late), late - _assert_rows_land_exactly_once(results, "late") - survivor: Final = tuple(pid for pid in _worker_pids(owned) if pid != workers[0]) - assert survivor, "no worker left serving" - after: Final = owned.gateway.chat(model, key=key, text=f"after kill {uuid.uuid4().hex}") - assert isinstance(after["id"], str) and after["id"].startswith("chatcmpl-"), after - assert client.get("/health/liveliness").status_code == 200 + with owned.replacing_worker(victim) as kill: + results: Final = _burst(client, model, key, scenario.key(), kill) + dropped: Final = tuple(result for result in results if result.status is None) + assert len(dropped) < len(results), results + crafted: Final = tuple( + result for result in results if result.label.startswith("crafted") and result.status + ) + assert crafted and all( + result.status == 400 and result.seconds < FAST_SECONDS for result in crafted + ), crafted + late: Final = tuple(result for result in results if result.label.startswith("late")) + assert all(result.status == 200 for result in late), late + _assert_rows_land_exactly_once(results, "late") + survivor: Final = tuple(pid for pid in owned.worker_pids() if pid != victim) + assert survivor, "no worker left serving" + after: Final = owned.gateway.chat(model, key=key, text=f"after kill {uuid.uuid4().hex}") + assert isinstance(after["id"], str) and after["id"].startswith("chatcmpl-"), after + assert client.get("/health/liveliness").status_code == 200 def test_upstream_returning_503_mid_burst_logs_every_failure_with_its_own_cluster_key( diff --git a/tests/integration/spend/test_shutdown_flush.py b/tests/integration/spend/test_shutdown_flush.py index 5441f70d59e..e96511511d9 100644 --- a/tests/integration/spend/test_shutdown_flush.py +++ b/tests/integration/spend/test_shutdown_flush.py @@ -1,6 +1,5 @@ import json import os -import signal import threading import uuid from collections.abc import Callable, Iterator @@ -163,7 +162,7 @@ class _Shutdown: def terminate_once(self, blocked: Callable[[], bool], release: Callable[[], None]) -> None: eventually(blocked, lambda state: state, seconds=60) - self.owned.process.send_signal(signal.SIGTERM) + self.owned.begin_shutdown() eventually(lambda: self.logged(CANCEL_LOG_LINE), lambda seen: seen, seconds=60) release() self.owned.process.wait(timeout=120) diff --git a/tests/integration/spend/test_vertex_image_regional_uplift.py b/tests/integration/spend/test_vertex_image_regional_uplift.py index a06c8842a45..9aedb81c3e0 100644 --- a/tests/integration/spend/test_vertex_image_regional_uplift.py +++ b/tests/integration/spend/test_vertex_image_regional_uplift.py @@ -3,7 +3,6 @@ import base64 import json import os import re -import signal import threading import time import uuid @@ -625,27 +624,28 @@ async def test_worker_sigkill_mid_burst_leaves_the_sibling_billing_the_uplift(ga held_by: Final = {pid: _open_upstream_connections(pid, wire.url) for pid in workers} assert sum(held_by.values()) == len(call_ids), held_by victim_pid, survivor_pid = sorted(workers, key=held_by.__getitem__) - psutil.Process(victim_pid).send_signal(signal.SIGKILL) - release.set() - served: Final = await burst - succeeded: Final = tuple(item.call_id for item in served if item.status == 200) - _artifact( - tmp_path, - "worker-sigkill.json", - { - "held_by_killed_worker": held_by[victim_pid], - "held_by_surviving_worker": held_by[survivor_pid], - "served": len(served), - "succeeded": len(succeeded), - "transport_errors": len(call_ids) - len(served), - }, - ) - assert len(succeeded) == held_by[survivor_pid], ( - held_by, - [(item.status, item.text[:80]) for item in served], - ) - follow_up_id, follow_up = _generate(owned.gateway, _CONFIG_MODEL) - _assert_images(follow_up, follow_up_id, _REGIONAL_TOKEN_COST) + with owned.replacing_worker(victim_pid) as kill: + await asyncio.to_thread(kill) + release.set() + served: Final = await burst + succeeded: Final = tuple(item.call_id for item in served if item.status == 200) + _artifact( + tmp_path, + "worker-sigkill.json", + { + "held_by_killed_worker": held_by[victim_pid], + "held_by_surviving_worker": held_by[survivor_pid], + "served": len(served), + "succeeded": len(succeeded), + "transport_errors": len(call_ids) - len(served), + }, + ) + assert len(succeeded) == held_by[survivor_pid], ( + held_by, + [(item.status, item.text[:80]) for item in served], + ) + follow_up_id, follow_up = _generate(owned.gateway, _CONFIG_MODEL) + _assert_images(follow_up, follow_up_id, _REGIONAL_TOKEN_COST) rows: Final = await asyncio.to_thread(_spend_rows, succeeded, seconds=120) assert rows == {call_id: pytest.approx(_REGIONAL_TOKEN_COST) for call_id in succeeded} await asyncio.to_thread( @@ -670,7 +670,7 @@ async def test_proxy_restart_mid_burst_bills_each_landed_call_once(gateway: Gate ) async with _gated_burst(base_url, gateway.key, call_ids, release) as burst: await asyncio.to_thread(eventually, held.qsize, lambda size: size == len(call_ids), 120) - owned.process.terminate() + owned.begin_shutdown() release.set() served: Final = await burst await asyncio.to_thread(owned.process.wait, 240) diff --git a/tests/integration/streaming/test_interception_converted_stream_options.py b/tests/integration/streaming/test_interception_converted_stream_options.py index 46a4211b121..74fad7deaf1 100644 --- a/tests/integration/streaming/test_interception_converted_stream_options.py +++ b/tests/integration/streaming/test_interception_converted_stream_options.py @@ -1,6 +1,5 @@ import asyncio import json -import signal import threading import uuid from collections.abc import Callable, Generator, Iterator, Mapping, Sequence @@ -882,18 +881,21 @@ async def test_worker_sigkill_mid_burst_leaves_the_sibling_restreaming_with_usag held_by: Final = MappingProxyType({pid: _open_upstream_connections(pid, owned.wire.url) for pid in workers}) assert sum(held_by.values()) == 20, held_by victim_pid, survivor_pid = sorted(workers, key=held_by.__getitem__) - victim: Final = psutil.Process(victim_pid) - victim.suspend() - victim.send_signal(signal.SIGKILL) - release.set() - served: Final = await burst - assert held_by[survivor_pid] >= 10, held_by - assert len(served) == held_by[survivor_pid], (held_by, len(served)) - identities: Final = frozenset( - _assert_restreamed(item.stream, iv.answer(_BURST_DEPLOYMENT), usage=iv.FOLLOWUP_USAGE) for item in served - ) - (follow_up,) = await _burst(owned.base_url, owned.gateway.key, (_chat_request(deployment),)) - recovered: Final = _assert_restreamed(follow_up.stream, iv.answer(_BURST_DEPLOYMENT), usage=iv.FOLLOWUP_USAGE) + with owned.proxy.replacing_worker(victim_pid) as kill: + kill.freeze() + await asyncio.to_thread(kill) + release.set() + served: Final = await burst + assert held_by[survivor_pid] >= 10, held_by + assert len(served) == held_by[survivor_pid], (held_by, len(served)) + identities: Final = frozenset( + _assert_restreamed(item.stream, iv.answer(_BURST_DEPLOYMENT), usage=iv.FOLLOWUP_USAGE) + for item in served + ) + (follow_up,) = await _burst(owned.base_url, owned.gateway.key, (_chat_request(deployment),)) + recovered: Final = _assert_restreamed( + follow_up.stream, iv.answer(_BURST_DEPLOYMENT), usage=iv.FOLLOWUP_USAGE + ) rows: Final = _spend_rows(_BURST_MODEL, len(identities) + 1) landed: Final = tuple(string_value(row["request_id"]) for row in rows if row["status"] == "success") assert sorted(landed) == sorted((*identities, recovered)), rows diff --git a/tests/integration/streaming/test_responses_bridge_stream_chaos.py b/tests/integration/streaming/test_responses_bridge_stream_chaos.py index 6b2ea47a849..745be2e34b1 100644 --- a/tests/integration/streaming/test_responses_bridge_stream_chaos.py +++ b/tests/integration/streaming/test_responses_bridge_stream_chaos.py @@ -1,6 +1,5 @@ import asyncio import re -import signal import socket import threading import uuid @@ -322,18 +321,18 @@ async def test_worker_sigkill_mid_burst_leaves_the_sibling_answering_the_bridged held_by: Final = MappingProxyType({pid: _open_upstream_connections(pid, wire.url) for pid in workers}) assert sum(held_by.values()) == 20, held_by victim_pid, survivor_pid = sorted(workers, key=held_by.__getitem__) - victim: Final = psutil.Process(victim_pid) - victim.suspend() - victim.send_signal(signal.SIGKILL) - held.release.set() - served: Final = await burst - assert held_by[survivor_pid] >= 10, held_by - assert len(served) == held_by[survivor_pid], (held_by, len(served)) - for item in served: - _assert_answered_in_its_own_shape(item) - follow_up: Final = _Call("chat_healthy", uuid.uuid4().hex) - (answered,) = await _burst(rig.gateway, (follow_up,)) - _assert_answered_in_its_own_shape(answered) + with rig.proxy.replacing_worker(victim_pid) as kill: + kill.freeze() + await asyncio.to_thread(kill) + held.release.set() + served: Final = await burst + assert held_by[survivor_pid] >= 10, held_by + assert len(served) == held_by[survivor_pid], (held_by, len(served)) + for item in served: + _assert_answered_in_its_own_shape(item) + follow_up: Final = _Call("chat_healthy", uuid.uuid4().hex) + (answered,) = await _burst(rig.gateway, (follow_up,)) + _assert_answered_in_its_own_shape(answered) _assert_forwarded(wire.drain(), (*calls, follow_up)) _assert_each_lands_once((*served, answered)) @@ -349,7 +348,7 @@ async def test_proxy_sigterm_mid_burst_drains_the_spend_log_queue_and_the_restar with wire_server(held.respond, port=port) as wire: with owned_proxy_process(gateway, tmp_path, {}, config=config, workers=2) as owned: burst: Final = await _held_burst(owned.gateway, calls, held) - owned.process.terminate() + owned.begin_shutdown() held.release.set() served: Final = await burst await asyncio.to_thread( diff --git a/tests/integration/test_oci_proxy_integration.py b/tests/integration/test_oci_proxy_integration.py index 06ed4486941..d059a55b1f2 100644 --- a/tests/integration/test_oci_proxy_integration.py +++ b/tests/integration/test_oci_proxy_integration.py @@ -42,6 +42,7 @@ from typing import Iterator import httpx import pytest from tests._master_key import MASTER_KEY +from tests.integration._support.process import stop_child # --------------------------------------------------------------------------- # Skip gate @@ -147,12 +148,7 @@ def _serve(config_path: str) -> Iterator[str]: _wait_for_health(base_url, proc, time.monotonic() + STARTUP_TIMEOUT_S) yield base_url finally: - proc.terminate() - try: - proc.wait(timeout=10) - except subprocess.TimeoutExpired: - proc.kill() - proc.wait(timeout=5) + stop_child(proc, grace_seconds=10) @pytest.fixture(scope="module") diff --git a/tests/unit/integration_support/test_budget.py b/tests/unit/integration_support/test_budget.py new file mode 100644 index 00000000000..e0e14fab1af --- /dev/null +++ b/tests/unit/integration_support/test_budget.py @@ -0,0 +1,58 @@ +from __future__ import annotations + +import json +from typing import Final + +import pytest + +from tests.integration._support.budget import ( + IMPORT_SECONDS_VARIABLE, + READINESS_OVERRIDE_VARIABLE, + calibration, + case_timeout_for, + import_seconds, + readiness_for, + readiness_seconds, +) + + +@pytest.mark.parametrize("workers", (1, 2, 4)) +def test_a_box_that_imports_twice_as_slowly_gets_twice_the_readiness_deadline(workers: int) -> None: + assert readiness_for(14.0, workers) == pytest.approx(2 * readiness_for(7.0, workers)) + + +def test_a_multi_worker_proxy_gets_the_two_imports_its_boot_pays_for() -> None: + assert readiness_for(7.0, 2) == pytest.approx(2 * readiness_for(7.0, 1)) + assert readiness_for(7.0, 4) == readiness_for(7.0, 2) + + +@pytest.mark.parametrize("imported", (0.5, 7.0, 40.0)) +def test_a_case_timeout_covers_a_previous_teardown_a_two_worker_boot_and_the_call(imported: float) -> None: + assert case_timeout_for(imported) > readiness_for(imported, 2) + readiness_for(imported, 1) + + +def test_the_calibration_record_matches_the_budgets_it_reports() -> None: + record: Final = json.loads(calibration(7.0)) + assert record == { + "import_seconds": 7.0, + "readiness_seconds_1_worker": readiness_for(7.0, 1), + "readiness_seconds_2_workers": readiness_for(7.0, 2), + "case_timeout_seconds": case_timeout_for(7.0), + } + + +def test_the_exported_calibration_is_reused_instead_of_measured_again(monkeypatch: pytest.MonkeyPatch) -> None: + import_seconds.cache_clear() + monkeypatch.setenv(IMPORT_SECONDS_VARIABLE, "3.5") + monkeypatch.delenv(READINESS_OVERRIDE_VARIABLE, raising=False) + try: + assert import_seconds() == 3.5 + assert readiness_seconds(2) == readiness_for(3.5, 2) + finally: + import_seconds.cache_clear() + + +def test_an_explicit_readiness_override_wins_over_the_calibration(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv(READINESS_OVERRIDE_VARIABLE, "42") + assert readiness_seconds(1) == 42.0 + assert readiness_seconds(2) == 42.0 diff --git a/tests/unit/integration_support/test_process.py b/tests/unit/integration_support/test_process.py index b65c9fa7ec6..998f81155a3 100644 --- a/tests/unit/integration_support/test_process.py +++ b/tests/unit/integration_support/test_process.py @@ -5,8 +5,12 @@ import errno import importlib import os import socket +import subprocess +import sys +import threading import time from collections.abc import Iterator +from contextlib import contextmanager from pathlib import Path from types import ModuleType from typing import Final @@ -18,8 +22,18 @@ import pytest import yaml from pydantic import TypeAdapter -from tests.integration._support.client import Gateway -from tests.integration._support.process import worker_configuration +from tests.integration._support.client import Gateway, eventually +from tests.integration._support.process import ( + OwnedProxy, + WorkerHealthcheck, + WorkerKill, + healthcheck_seconds, + kill_stray_process, + owned_proxy_process, + replacement_pid, + stop_process_group, + worker_configuration, +) TESTS_DIR: Final = Path(__file__).resolve().parents[2] UNRELATED_CRASH: Final = "Traceback (most recent call last):\nModuleNotFoundError: No module named 'litellm'\n" @@ -278,3 +292,159 @@ def test_owned_runtime_configuration_preserves_custom_worker_startup_hooks( ) if hook ) + + +_SUPERVISOR: Final = """import multiprocessing +import time + +if __name__ == "__main__": + worker = multiprocessing.get_context("spawn").Process(target=time.sleep, args=(60,)) + worker.start() + print(worker.pid, flush=True) + deadline = time.monotonic() + 60 + while time.monotonic() < deadline: + multiprocessing.active_children() + time.sleep(0.05) +""" + + +@contextmanager +def _supervised(tmp_path: Path, healthcheck: WorkerHealthcheck) -> Iterator[tuple[OwnedProxy, int]]: + script: Final = tmp_path / "supervisor.py" + script.write_text(_SUPERVISOR) + with subprocess.Popen( + (sys.executable, "-P", str(script)), stdout=subprocess.PIPE, start_new_session=True + ) as supervisor: + try: + assert supervisor.stdout is not None + worker: Final = int(supervisor.stdout.readline()) + with httpx.Client(base_url="http://127.0.0.1:9") as client: + owned: Final = OwnedProxy( + Gateway(client, "sk-unit", "http://127.0.0.1:9"), supervisor, tmp_path / "proxy.log", healthcheck + ) + eventually(owned.worker_pids, lambda pids: pids == frozenset((worker,)), seconds=30) + yield owned, worker + finally: + stop_process_group(supervisor, 5) + + +def test_worker_pids_are_the_supervisors_spawned_workers_and_not_the_supervisor(tmp_path: Path) -> None: + with _supervised(tmp_path, WorkerHealthcheck.STRICT) as (owned, worker): + assert owned.worker_pids() == frozenset((worker,)) + assert owned.process.pid not in owned.worker_pids() + + +def test_only_a_serving_worker_of_this_proxy_can_be_killed(tmp_path: Path) -> None: + with _supervised(tmp_path, WorkerHealthcheck.STRICT) as (owned, worker): + with pytest.raises(AssertionError, match="not a serving worker"): + owned.kill_worker(owned.process.pid) + assert owned.process.poll() is None + assert psutil.Process(worker).is_running() + + +def _leave_without_killing(owned: OwnedProxy, worker: int) -> None: + with owned.replacing_worker(worker): + pass + + +def test_a_replacing_block_that_never_kills_its_worker_fails_and_leaves_it_serving(tmp_path: Path) -> None: + with _supervised(tmp_path, WorkerHealthcheck.STRICT) as (owned, worker): + with pytest.raises(AssertionError, match="without killing the worker"): + _leave_without_killing(owned, worker) + assert owned.worker_pids() == frozenset((worker,)) + + +def _kill_of(owned: OwnedProxy, worker: int) -> WorkerKill: + return WorkerKill(psutil.Process(worker), psutil.Process(owned.process.pid), 5, threading.Event()) + + +def test_a_worker_kill_returns_once_the_supervisor_reaped_the_victim_and_cannot_fire_twice(tmp_path: Path) -> None: + with _supervised(tmp_path, WorkerHealthcheck.STRICT) as (owned, worker): + kill: Final = _kill_of(owned, worker) + kill() + assert not psutil.pid_exists(worker) + assert kill.fired.is_set() + with pytest.raises(AssertionError, match="already killed"): + kill() + + +def test_a_frozen_victim_stays_stopped_until_the_kill(tmp_path: Path) -> None: + with _supervised(tmp_path, WorkerHealthcheck.STRICT) as (owned, worker): + kill: Final = _kill_of(owned, worker) + assert kill.freeze().status() == psutil.STATUS_STOPPED + kill() + assert not psutil.pid_exists(worker) + + +def test_a_kill_under_a_suspended_supervisor_returns_with_the_victim_left_unreaped(tmp_path: Path) -> None: + with _supervised(tmp_path, WorkerHealthcheck.LENIENT) as (owned, worker): + kill: Final = _kill_of(owned, worker) + with owned.suspended(owned.process.pid): + kill() + assert psutil.Process(worker).status() == psutil.STATUS_ZOMBIE + eventually(lambda: psutil.pid_exists(worker), lambda exists: not exists, seconds=5) + + +def test_a_kill_fails_when_the_supervisor_never_reaps_the_victim(tmp_path: Path) -> None: + with _supervised(tmp_path, WorkerHealthcheck.STRICT) as (owned, worker): + unreaping: Final = psutil.Process(owned.process.pid) + kill: Final = WorkerKill(psutil.Process(worker), psutil.Process(os.getpid()), 0.5, threading.Event()) + unreaping.suspend() + try: + with pytest.raises(AssertionError, match="did not reap"): + kill() + finally: + unreaping.resume() + assert not kill.fired.is_set() + + +def test_a_strict_healthcheck_proxy_refuses_suspension_before_stopping_anything(tmp_path: Path) -> None: + with _supervised(tmp_path, WorkerHealthcheck.STRICT) as (owned, worker): + with pytest.raises(AssertionError, match="LENIENT"): + with owned.suspended(worker): + pass + assert psutil.Process(worker).status() != psutil.STATUS_STOPPED + + +def test_a_lenient_healthcheck_proxy_suspends_for_the_block_and_resumes_after(tmp_path: Path) -> None: + with _supervised(tmp_path, WorkerHealthcheck.LENIENT) as (owned, worker): + with owned.suspended(worker) as suspended: + assert suspended.status() == psutil.STATUS_STOPPED + assert psutil.Process(worker).status() != psutil.STATUS_STOPPED + assert owned.worker_pids() == frozenset((worker,)) + + +def test_a_worker_killed_while_suspended_fails_the_block(tmp_path: Path) -> None: + with _supervised(tmp_path, WorkerHealthcheck.LENIENT) as (owned, worker): + with pytest.raises(AssertionError, match="died while suspended"): + with owned.suspended(worker): + kill_stray_process(worker) + + +def test_the_worker_healthcheck_cannot_be_overridden_through_extra_arguments(tmp_path: Path) -> None: + with pytest.raises(AssertionError, match="healthcheck"): + with owned_proxy_process( + MagicMock(spec=Gateway), tmp_path, {}, extra_arguments=("--timeout_worker_healthcheck", "60") + ): + pass + assert not tuple(tmp_path.glob("owned-proxy-*.log")) + + +def test_a_lenient_healthcheck_outlasts_a_strict_one() -> None: + assert healthcheck_seconds(WorkerHealthcheck.LENIENT) > healthcheck_seconds(WorkerHealthcheck.STRICT) + + +@pytest.mark.parametrize( + ("answered", "serving", "replacement"), + ( + (30, frozenset((20, 30)), 30), + (20, frozenset((20, 30)), None), + (10, frozenset((20, 30)), None), + (40, frozenset((20, 30)), None), + (None, frozenset((20, 30)), None), + ), +) +def test_only_a_new_serving_worker_answering_readiness_counts_as_the_replacement( + answered: int | None, serving: frozenset[int], replacement: int | None +) -> None: + assert replacement_pid(frozenset((10, 20)), answered, serving) == replacement diff --git a/tests/unit/integration_support/test_timing.py b/tests/unit/integration_support/test_timing.py new file mode 100644 index 00000000000..7e098e879b6 --- /dev/null +++ b/tests/unit/integration_support/test_timing.py @@ -0,0 +1,57 @@ +from __future__ import annotations + +import time +from typing import Final + +import pytest + +from tests.integration._support.timing import ( + assert_cut_short, + assert_enforced_within, + assert_not_delayed_by, + assert_within_factor_of_control, + timed, +) + + +def test_timed_returns_the_value_and_at_least_the_time_the_operation_took() -> None: + measured: Final = timed(lambda: time.sleep(0.05) or "done") + assert measured.value == "done" + assert measured.seconds >= 0.05 + + +def test_a_burst_slowed_by_load_but_not_by_the_injected_delay_passes() -> None: + assert_not_delayed_by((1.2, 2.4), control=(1.1, 1.9), injected_delay=2) + + +def test_a_burst_that_waited_out_the_injected_delay_fails() -> None: + with pytest.raises(AssertionError, match="injected delay"): + assert_not_delayed_by((1.2, 4.0), control=(1.1, 1.9), injected_delay=2) + + +def test_a_call_cut_short_before_the_injected_stall_passes_and_one_that_waited_fails() -> None: + assert_cut_short(1.5, injected_delay=20) + with pytest.raises(AssertionError, match="no sooner than"): + assert_cut_short(20.0, injected_delay=20) + + +def test_an_enforced_limit_passes_inside_its_slack_and_fails_on_either_side() -> None: + assert_enforced_within(59.5, configured=60, slack=15, not_before=59) + assert_enforced_within(75.0, configured=60, slack=15, not_before=59) + with pytest.raises(AssertionError, match="configured 60s plus 15s"): + assert_enforced_within(300.0, configured=60, slack=15, not_before=59) + with pytest.raises(AssertionError, match="between 59s"): + assert_enforced_within(58.0, configured=60, slack=15, not_before=59) + + +def test_an_enforced_limit_without_a_floor_only_caps_the_elapsed_time() -> None: + assert_enforced_within(0.0, configured=1, slack=7) + with pytest.raises(AssertionError): + assert_enforced_within(8.5, configured=1, slack=7) + + +def test_a_bound_relative_to_the_control_scales_with_a_slower_machine() -> None: + assert_within_factor_of_control(0.9, control=0.1, factor=10) + assert_within_factor_of_control(9.0, control=1.0, factor=10) + with pytest.raises(AssertionError, match="over 10x"): + assert_within_factor_of_control(1.1, control=0.1, factor=10) diff --git a/tests/unit/test_check_test_quality.py b/tests/unit/test_check_test_quality.py index bf05775d09d..a0b46cdcdaa 100644 --- a/tests/unit/test_check_test_quality.py +++ b/tests/unit/test_check_test_quality.py @@ -505,6 +505,7 @@ def test_a_named_iterable_that_is_not_a_module_constant_is_skipped_quietly(tmp_p def _written(tmp_path, source, name="conftest.py"): path = tmp_path / name + path.parent.mkdir(parents=True, exist_ok=True) path.write_text(source, encoding="utf-8") return path @@ -637,6 +638,8 @@ _VIOLATING_SNIPPETS: Final = MappingProxyType( "test_snippet.py", 'import subprocess, sys\nsubprocess.run([sys.executable, "-c", "pass"])\n', ), + "TQ010": ("tests/integration/chaos/test_snippet.py", "import os, signal\n\nos.kill(1, signal.SIGKILL)\n"), + "TQ011": ("tests/integration/chaos/test_snippet.py", "def test_fast():\n elapsed = run()\n assert elapsed < 2\n"), } ) @@ -675,3 +678,134 @@ def test_non_interpreter_subprocess_call_is_untouched(tmp_path): def test_popen_sys_executable_tuple_is_flagged(tmp_path): source = 'import subprocess, sys\nsubprocess.Popen((sys.executable, "script.py"))\n' assert _codes(tmp_path, source) == ["TQ009"] + + +_CHAOS_TEST: Final = "tests/integration/observability/test_snippet.py" + + +def _chaos_codes(tmp_path, source, name=_CHAOS_TEST): + return [v.code for v in checker.check_file(_written(tmp_path, source, name))] + + +@pytest.mark.parametrize( + "call", + ( + "os.kill(pid, signal.SIGKILL)", + "os.killpg(pid, signal.SIGTERM)", + "os.kill(pid, 0)", + "worker.send_signal(signal.SIGTERM)", + "signal_group(pid, signal.SIGKILL)", + "redis.signal(SIGSTOP)", + "redis.signal(signal.SIGCONT)", + "worker.kill()", + "worker.terminate()", + "worker.suspend()", + "worker.resume()", + ), +) +def test_an_integration_test_that_signals_a_process_itself_is_flagged_once_per_call(tmp_path, call): + assert _chaos_codes(tmp_path, f"import os, signal\nfrom signal import SIGSTOP\n\n{call}\n") == ["TQ010"] + + +@pytest.mark.parametrize( + "source", + ( + "assert record['exit_before_stop'] == -signal.SIGKILL\n", + "signal.signal(signal.SIGTERM, handler)\n", + "task.cancel()\n", + "slot.kill(reason)\n", + "owned.kill_worker(pid)\n", + "stop_child(process)\n", + ), +) +def test_reading_a_signal_number_or_calling_a_harness_helper_is_not_signalling(tmp_path, source): + assert _chaos_codes(tmp_path, f"import signal\n\n{source}") == [] + + +@pytest.mark.parametrize( + "name", + ( + "tests/integration/_support/process.py", + "tests/unit/test_snippet.py", + "tests/e2e/integration_like/test_snippet.py", + ), +) +def test_the_signal_rule_only_applies_to_integration_tests_outside_the_harness(tmp_path, name): + assert _chaos_codes(tmp_path, "import os, signal\n\nos.kill(1, signal.SIGKILL)\nworker.suspend()\n", name) == [] + + +def test_the_signal_rule_is_suppressible_with_a_reason(tmp_path): + source = "import os, signal\n\nos.kill(1, signal.SIGKILL) # test-quality-ok: reaps a leaked fake helper\n" + assert _chaos_codes(tmp_path, source) == [] + + +@pytest.mark.parametrize( + "assertion", + ( + "assert elapsed < 2", + "assert 2 > elapsed", + "assert elapsed <= 2.5", + "assert time.monotonic() - started < 5", + "assert all(duration < 2 for _, duration in results)", + "assert max(latencies) < 1", + "assert response.elapsed.total_seconds() < 5", + "assert outcome.elapsed < BUDGET_SECONDS", + "assert elapsed < BUDGET_SECONDS - 1", + "assert 59 <= elapsed <= 75", + "assert taken < 3", + ), +) +def test_an_elapsed_time_capped_by_a_number_is_flagged(tmp_path, assertion): + source = ( + "import time\n\nBUDGET_SECONDS = 10\n\n\n" + "def test_fast():\n" + " started = time.monotonic()\n" + " taken = time.perf_counter() - started\n" + f" {assertion}\n" + ) + assert _chaos_codes(tmp_path, source) == ["TQ011"] + + +@pytest.mark.parametrize( + "statement", + ( + "assert elapsed >= 2", + "assert 2 <= elapsed", + "assert elapsed < 25 * control_elapsed", + "assert elapsed < control + injected_delay", + "assert_enforced_within(elapsed, configured=CAP_SECONDS, slack=15, not_before=CAP_SECONDS - 1)", + "assert_cut_short(time.monotonic() - started, injected_delay=30)", + "assert after - baseline < BUDGET_SECONDS", + "assert len(durations) == 3", + "assert count < 2", + "while time.monotonic() - started < 30:\n poll()", + ), +) +def test_lower_bounds_relative_bounds_counts_and_polling_deadlines_are_left_alone(tmp_path, statement): + source = ( + "import time\n\nBUDGET_SECONDS = 10\n\n\n" + "def _relative():\n" + " started = time.monotonic()\n" + " baseline = settled_reads(read(), time.monotonic())\n" + f" {statement}\n" + ) + assert _chaos_codes(tmp_path, source) == [] + + +def test_a_clock_bound_name_is_only_an_elapsed_time_in_the_function_that_binds_it(tmp_path): + source = ( + "import time\n\n\n" + "def _measure():\n" + " taken = time.monotonic() - 1\n" + " return taken\n\n\n" + "def test_counts():\n" + " taken = len(rows())\n" + " assert taken < 3\n" + ) + assert _chaos_codes(tmp_path, source) == [] + + +def test_the_wall_clock_rule_only_applies_to_integration_tests_outside_the_harness(tmp_path): + source = "def test_fast():\n elapsed = run()\n assert elapsed < 2\n" + assert _chaos_codes(tmp_path, source, "tests/unit/test_snippet.py") == [] + assert _chaos_codes(tmp_path, source, "tests/integration/_support/timing.py") == []