From 02feb67c20041937afd948d27ec207b0585e7a18 Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Mon, 12 Oct 2026 09:48:42 -0700 Subject: [PATCH] test(integration): cut CircleCI proxy boot time with concurrent worker stacks and single-worker test proxies (#46037) * test(integration): build worker stacks concurrently and record proxy boot phases The controller built every xdist worker's database, Redis, upstream and gateway proxy one after another before any test ran, about 95 s per extensions shard, and since every multi-worker suite is now isolated the same serial wait hits mcp, security and cost too. The stacks now start on a thread pool bounded by the CPU count. Each stack reads its endpoints and results directory from an explicit environment instead of a patched os.environ, every stack that started is registered for cleanup, and a failure in one stops the others before the error is raised Owned proxy startup.json now records the wall-clock spawn time, the worker count and when the Prisma CLI child ran, when the port started accepting and when the first worker answered readiness, and exit.json records when the stop started and ended, so the boot time can be split and overlapping gateway boots can be seen in the artifacts * test(integration): boot one uvicorn worker where a test does not assert cross-worker behaviour A 2-worker owned proxy takes about twice as long to become ready as a 1-worker one, because the uvicorn parent imports the app and each spawned worker imports it again. The passthrough upstream error visibility cases, the single-request S3 dedupe cases and the single-request spend log payload cases send one or two requests and assert on the response, the proxy log and the spend row, none of which depends on a second worker, so they now boot the default single worker Cases that do rely on two workers keep them, the concurrent S3 burst through a sink outage, previous_response_id replay, response cache hits, the concurrent spend bursts and upstream stop and recovery * test(observability): settle OTel excluded-services traces by a quiet period instead of a full window _trace_spans slept through its whole 15 or 30 s window on every call even when the trace had long stopped growing. It now polls and returns once the trace has spans and has not gained one for five tenant export intervals, and the module pins the tenant batch exporter's schedule delay to 1 s so that bound is set by the test rather than by the SDK default. Absence checks still run only after the trace stopped growing, and against a build that leaks postgres spans to the tenant they still fail These cases also boot one uvicorn worker now. None of them asserts anything about a second worker --------- Co-authored-by: Yuneng Jiang --- tests/integration/_support/process.py | 173 ++++++++++++++---- tests/integration/_support/redis_process.py | 10 +- tests/integration/_support/workers.py | 64 ++++++- tests/integration/conftest.py | 26 +-- .../test_otel_excluded_services.py | 78 +++++--- ...t_passthrough_upstream_error_visibility.py | 36 ++-- .../test_s3_v2_streaming_failure_dedupe.py | 3 - .../test_spend_log_tool_payload_content.py | 29 ++- .../unit/integration_support/test_process.py | 40 ++++ .../unit/integration_support/test_workers.py | 51 +++++- 10 files changed, 379 insertions(+), 131 deletions(-) diff --git a/tests/integration/_support/process.py b/tests/integration/_support/process.py index 6ec8dd7447a..a1e1b649202 100644 --- a/tests/integration/_support/process.py +++ b/tests/integration/_support/process.py @@ -40,9 +40,9 @@ CONFIGURATION: Final = TypeAdapter(dict[str, object]) MCP_COVERAGE_CONFIG: Final = Path(__file__).resolve().parents[1] / "mcp_coverage.toml" -def proxy_database_environment() -> Mapping[str, str]: - writer: Final = os.environ.get("INTEGRATION_PROXY_DATABASE_URL", "") - reader: Final = os.environ.get("INTEGRATION_PROXY_READ_REPLICA_URL", "") +def proxy_database_environment(inherited: Mapping[str, str] = os.environ) -> Mapping[str, str]: + writer: Final = inherited.get("INTEGRATION_PROXY_DATABASE_URL", "") + reader: Final = inherited.get("INTEGRATION_PROXY_READ_REPLICA_URL", "") return MappingProxyType( { **({"DATABASE_URL": writer} if writer else {}), @@ -326,6 +326,7 @@ def owned_proxy( workers: int = 1, database_setup: tuple[str, ...] = DB_PUSH, coverage_file: Path | None = None, + inherited: Mapping[str, str] = os.environ, ) -> Iterator[Gateway]: with owned_proxy_process( gateway, @@ -336,6 +337,7 @@ def owned_proxy( workers=workers, database_setup=database_setup, coverage_file=coverage_file, + inherited=inherited, ) as owned: yield owned.gateway @@ -381,8 +383,13 @@ class _Launch: started: float +def _wall_clock(monotonic: float) -> float: + return time.time() - (time.monotonic() - monotonic) + + def _stop_launch(launch: _Launch) -> None: started: Final = time.monotonic() + stop_started_at: Final = time.time() exit_before: Final = launch.process.poll() exception_before: Final = sys.exc_info()[0] try: @@ -399,6 +406,9 @@ def _stop_launch(launch: _Launch) -> None: "exception_at_record": exception_at_record.__name__ if exception_at_record else None, "cleanup_seconds": time.monotonic() - started, "process_lifetime_seconds": time.monotonic() - launch.started, + "spawned_at": _wall_clock(launch.started), + "stop_started_at": stop_started_at, + "stopped_at": time.time(), }, indent=2, ) @@ -439,7 +449,51 @@ def _owns_listener(launch: _Launch) -> bool: return False -def _readiness_pids(launch: _Launch, workers: int) -> Iterator[int]: +@dataclass(frozen=True, slots=True) +class BootPhases: + prisma_started: float | None = None + prisma_finished: float | None = None + listening: float | None = None + first_ready: float | None = None + ready_pids: frozenset[int] = frozenset() + + +@dataclass(frozen=True, slots=True) +class _Probe: + elapsed: float + listening: bool + prisma: bool + ready_pid: int | None + + +def _runs_prisma(process: psutil.Process) -> bool: + try: + return any(Path(part).name == "prisma" for part in process.cmdline()) + except psutil.Error: + return False + + +def _prisma_running(launch: _Launch) -> bool: + try: + children: Final = psutil.Process(launch.process.pid).children() + except psutil.Error: + return False + return any(_runs_prisma(child) for child in children) + + +def _answer(client: httpx.Client, launch: _Launch) -> tuple[bool, int | None]: + try: + response: Final = client.get("/health/integration-readiness", timeout=2) + except httpx.ConnectError: + return False, None + except httpx.TransportError: + return True, None + identity: Final = response.headers.get("x-integration-worker-pid", "") + ready: Final = response.status_code == 200 and identity.isdecimal() and _owns_listener(launch) + return True, int(identity) if ready else None + + +def _probes(launch: _Launch, workers: int) -> Iterator[_Probe]: with httpx.Client( base_url=f"http://127.0.0.1:{launch.port}", timeout=15, @@ -448,19 +502,27 @@ def _readiness_pids(launch: _Launch, workers: int) -> Iterator[int]: ) as client: deadline: Final = time.monotonic() + readiness_seconds(workers) while launch.process.poll() is None: - try: - response: Final = client.get("/health/integration-readiness", timeout=2) - identity: Final = response.headers.get("x-integration-worker-pid", "") - if response.status_code == 200 and identity.isdecimal() and _owns_listener(launch): - yield int(identity) - except httpx.TransportError: - pass + elapsed: Final = time.monotonic() - launch.started + listening, ready_pid = _answer(client, launch) + yield _Probe(elapsed, listening, not listening and _prisma_running(launch), ready_pid) assert time.monotonic() < deadline, "Owned proxy readiness deadline exceeded" time.sleep(0.1) -def _ready_workers(pids: frozenset[int], pid: int) -> frozenset[int]: - return pids | frozenset((pid,)) +def _first(recorded: float | None, observed: bool, elapsed: float) -> float | None: + return recorded if recorded is not None or not observed else elapsed + + +def _observe(phases: BootPhases, probe: _Probe) -> BootPhases: + prisma_started: Final = _first(phases.prisma_started, probe.prisma, probe.elapsed) + prisma_finished: Final = _first(phases.prisma_finished, prisma_started is not None, probe.elapsed) + return BootPhases( + prisma_started=prisma_started, + prisma_finished=None if probe.prisma else prisma_finished, + listening=_first(phases.listening, probe.listening, probe.elapsed), + first_ready=_first(phases.first_ready, probe.ready_pid is not None, probe.elapsed), + ready_pids=phases.ready_pids | frozenset(() if probe.ready_pid is None else (probe.ready_pid,)), + ) def _ready_worker_pid(member: psutil.Process) -> int | None: @@ -475,13 +537,24 @@ def _ready_worker_pid(member: psutil.Process) -> int | None: return None -def _wait_until_ready(launch: _Launch, workers: int = 1) -> None: - 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 - ) - if len(ready & members) >= max(1, workers): - return +def _serves_all(launch: _Launch, phases: BootPhases, workers: int) -> bool: + if len(phases.ready_pids) < max(1, workers): + return False + members: Final = frozenset( + pid for member in group_members(launch.process.pid) if (pid := _ready_worker_pid(member)) is not None + ) + return len(phases.ready_pids & members) >= max(1, workers) + + +def _wait_until_ready(launch: _Launch, workers: int = 1) -> BootPhases | None: + return next( + ( + phases + for phases in accumulate(_probes(launch, workers), _observe, initial=BootPhases()) + if _serves_all(launch, phases, workers) + ), + None, + ) def _launch_until_bound( @@ -489,7 +562,7 @@ def _launch_until_bound( ) -> _Launch: launch: Final = _launch(command, root, environment, output) try: - _wait_until_ready(launch, workers) + phases: Final = _wait_until_ready(launch, workers) exit_code: Final = launch.process.poll() assert exit_code is None or (attempts > 1 and _lost_port_race(exit_code, launch.log)), ( "Owned proxy exited before readiness" @@ -498,12 +571,21 @@ def _launch_until_bound( _stop_launch(launch) raise if exit_code is None: + assert phases is not None launch.log.with_suffix(".startup.json").write_text( json.dumps( { "root_pid": launch.process.pid, "readiness_seconds": time.monotonic() - launch.started, "logical_cpus": os.cpu_count(), + "workers": workers, + "spawned_at": _wall_clock(launch.started), + "phases": { + "prisma_started": phases.prisma_started, + "prisma_finished": phases.prisma_finished, + "listening": phases.listening, + "first_ready": phases.first_ready, + }, } ) + "\n" @@ -517,14 +599,14 @@ def _proxy_root() -> Path: return Path(os.environ.get("INTEGRATION_PROXY_ROOT") or Path(__file__).resolve().parents[3]) -def worker_configuration(config: Path, directory: Path) -> Path: - if os.environ.get("INTEGRATION_WORKER_ISOLATION") != "1": +def worker_configuration(config: Path, directory: Path, inherited: Mapping[str, str] = os.environ) -> Path: + if inherited.get("INTEGRATION_WORKER_ISOLATION") != "1": return config def endpoint(match: re.Match[str]) -> str: if match.group(1) == "8190": - return os.environ["INTEGRATION_UPSTREAM_URL"] - return f"http://127.0.0.1:{os.environ['INTEGRATION_PROVIDER_PORT']}" + return inherited["INTEGRATION_UPSTREAM_URL"] + return f"http://127.0.0.1:{inherited['INTEGRATION_PROVIDER_PORT']}" contents: Final = re.sub(r"http://127[.]0[.]0[.]1:(8190|8191)(?![0-9])", endpoint, config.read_text()) directory.mkdir(parents=True, exist_ok=True) @@ -547,9 +629,12 @@ def worker_configuration(config: Path, directory: Path) -> Path: def _proxy_environment( - gateway: Gateway, overrides: Mapping[str, str], remove_environment: tuple[str, ...] + gateway: Gateway, + overrides: Mapping[str, str], + remove_environment: tuple[str, ...], + source: Mapping[str, str] = os.environ, ) -> Mapping[str, str]: - inherited: Final = {**os.environ, **proxy_database_environment()} + inherited: Final = {**source, **proxy_database_environment(source)} writer: Final = overrides.get("DATABASE_URL") reader: Final = inherited.get("DATABASE_URL_READ_REPLICA") unpaired_reader: Final = ( @@ -567,7 +652,7 @@ def _proxy_environment( { **{name: value for name, value in inherited.items() if name not in removed}, "LITELLM_MASTER_KEY": gateway.key, - "LITELLM_SALT_KEY": os.environ.get("LITELLM_SALT_KEY", "sk-integration-salt"), + "LITELLM_SALT_KEY": source.get("LITELLM_SALT_KEY", "sk-integration-salt"), "STORE_MODEL_IN_DB": "True", "PYTHON_DOTENV_DISABLED": "1", **overrides, @@ -628,13 +713,17 @@ def owned_proxy_process( extra_arguments: tuple[str, ...] = (), coverage_file: Path | None = None, healthcheck: WorkerHealthcheck = WorkerHealthcheck.STRICT, + inherited: Mapping[str, str] = os.environ, ) -> 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 + gateway, + {**overrides, **({"COVERAGE_FILE": str(coverage_file)} if coverage_file else {})}, + remove_environment, + inherited, ) - output: Final = Path(os.environ.get("INTEGRATION_RESULTS_DIR", str(directory))) + output: Final = Path(inherited.get("INTEGRATION_RESULTS_DIR", str(directory))) output.mkdir(parents=True, exist_ok=True) command: Final = ( sys.executable, @@ -642,7 +731,7 @@ def owned_proxy_process( *(("coverage", "run", f"--rcfile={MCP_COVERAGE_CONFIG}", "-m") if coverage_file else ()), "integration._support.proxy", "--config", - str(worker_configuration(config or root / "tests/integration/proxy_config.yaml", directory)), + str(worker_configuration(config or root / "tests/integration/proxy_config.yaml", directory, inherited)), "--host", "127.0.0.1", "--num_workers", @@ -786,13 +875,21 @@ class UpstreamSlot: It runs from the harness's own checkout, never ``INTEGRATION_PROXY_ROOT``: the double belongs to the tests, the proxy root only names the litellm under test.""" - __slots__ = ("certificate", "directory", "port", "process", "root") + __slots__ = ("certificate", "directory", "inherited", "port", "process", "root") - def __init__(self, directory: Path, port: int, root: Path, certificate: UpstreamCertificate | None = None) -> None: + def __init__( + self, + directory: Path, + port: int, + root: Path, + certificate: UpstreamCertificate | None = None, + inherited: Mapping[str, str] = os.environ, + ) -> None: self.directory = directory self.port = port self.root = root self.certificate = certificate + self.inherited = inherited self.process: subprocess.Popen[bytes] | None = None @property @@ -807,7 +904,7 @@ class UpstreamSlot: def start(self) -> None: assert self.process is None, "Owned upstream is already running" - output: Final = Path(os.environ.get("INTEGRATION_RESULTS_DIR") or self.directory) + output: Final = Path(self.inherited.get("INTEGRATION_RESULTS_DIR") or self.directory) log_path: Final = output / f"owned-upstream-{self.port}-{uuid.uuid4().hex}.log" with log_path.open("w") as log: process: Final = subprocess.Popen( @@ -820,7 +917,7 @@ class UpstreamSlot: *self._tls_arguments(), ], cwd=self.root, - env=dict(os.environ), + env=dict(self.inherited), stdout=log, stderr=subprocess.STDOUT, start_new_session=True, @@ -856,9 +953,11 @@ def _harness_root() -> Path: @contextmanager -def owned_upstream(directory: Path, *, tls: bool = False, port: int | None = None) -> Generator[UpstreamSlot]: +def owned_upstream( + directory: Path, *, tls: bool = False, port: int | None = None, inherited: Mapping[str, str] = os.environ +) -> Generator[UpstreamSlot]: certificate: Final = self_signed_certificate(directory) if tls else None - slot: Final = UpstreamSlot(directory, port or _free_port(), _harness_root(), certificate) + slot: Final = UpstreamSlot(directory, port or _free_port(), _harness_root(), certificate, inherited) try: slot.start() yield slot diff --git a/tests/integration/_support/redis_process.py b/tests/integration/_support/redis_process.py index cfdbf041eba..f5e7f509eb5 100644 --- a/tests/integration/_support/redis_process.py +++ b/tests/integration/_support/redis_process.py @@ -5,7 +5,7 @@ import socket import subprocess import time import uuid -from collections.abc import Callable, Iterator +from collections.abc import Callable, Iterator, Mapping from contextlib import contextmanager from dataclasses import dataclass from pathlib import Path @@ -172,7 +172,11 @@ def _started( @contextmanager -def owned_redis(directory: Path, choose_port: Callable[[], int] = free_loopback_port) -> Iterator[OwnedRedis]: +def owned_redis( + directory: Path, + choose_port: Callable[[], int] = free_loopback_port, + inherited: Mapping[str, str] = os.environ, +) -> Iterator[OwnedRedis]: binary: Final = shutil.which("redis-server") if binary: host: Final = "127.0.0.1" @@ -190,7 +194,7 @@ def owned_redis(directory: Path, choose_port: Callable[[], int] = free_loopback_ ).strip() assert host, "CircleCI owned Redis container has no address" prefix: Final = ("docker", "exec", "redis-cache", "redis-server") - output: Final = Path(os.environ.get("INTEGRATION_RESULTS_DIR", str(directory))) + output: Final = Path(inherited.get("INTEGRATION_RESULTS_DIR", str(directory))) output.mkdir(parents=True, exist_ok=True) directory.mkdir(parents=True, exist_ok=True) with (output / f"owned-redis-{uuid.uuid4().hex}.log").open("w") as log: diff --git a/tests/integration/_support/workers.py b/tests/integration/_support/workers.py index 843f5dec4ac..33e086576c5 100644 --- a/tests/integration/_support/workers.py +++ b/tests/integration/_support/workers.py @@ -4,8 +4,9 @@ import os import socket import time import uuid -from collections.abc import Iterator, Mapping -from contextlib import contextmanager, nullcontext +from collections.abc import Callable, Iterator, Mapping +from concurrent.futures import ThreadPoolExecutor +from contextlib import ExitStack, contextmanager, nullcontext from dataclasses import dataclass from pathlib import Path from types import MappingProxyType @@ -151,11 +152,13 @@ def cost_map_environment(upstream: str) -> Mapping[str, str]: @contextmanager -def worker_services(directory: Path, *, replica: bool = False) -> Iterator[Mapping[str, str]]: +def worker_services( + directory: Path, *, replica: bool = False, inherited: Mapping[str, str] = os.environ +) -> Iterator[Mapping[str, str]]: directory.mkdir(parents=True, exist_ok=True) with ( scratch_database(WORKER_DATABASE_PREFIX) as database, - owned_redis(directory) as cache, + owned_redis(directory, inherited=inherited) as cache, ): yield MappingProxyType( { @@ -170,15 +173,17 @@ def worker_services(directory: Path, *, replica: bool = False) -> Iterator[Mappi @contextmanager -def worker_gateway(directory: Path, profile: GatewayProfile) -> Iterator[Mapping[str, str]]: +def worker_gateway( + directory: Path, profile: GatewayProfile, inherited: Mapping[str, str] +) -> Iterator[Mapping[str, str]]: from tests.integration._support.process import owned_proxy, owned_upstream - upstream: Final = os.environ["INTEGRATION_UPSTREAM_URL"] + upstream: Final = inherited["INTEGRATION_UPSTREAM_URL"] port: Final = urlsplit(upstream).port assert port is not None key: Final = f"sk-{uuid.uuid4().hex}" removed: Final = ("LITELLM_LOCAL_MODEL_COST_MAP",) if profile.cost_map else () - with worker_services(directory, replica=profile.replica) as services: + with worker_services(directory, replica=profile.replica, inherited=inherited) as services: overrides: Final = { **services, "DATABASE_URL": services["INTEGRATION_PROXY_DATABASE_URL"], @@ -189,7 +194,7 @@ def worker_gateway(directory: Path, profile: GatewayProfile) -> Iterator[Mapping **(cost_map_environment(upstream) if profile.cost_map else {}), } with ( - owned_upstream(directory, port=port), + owned_upstream(directory, port=port, inherited=inherited), httpx.Client(base_url="http://127.0.0.1:1", trust_env=False) as client, owned_proxy( Gateway(client, key, upstream), @@ -197,8 +202,16 @@ def worker_gateway(directory: Path, profile: GatewayProfile) -> Iterator[Mapping overrides, remove_environment=removed, coverage_file=profile.coverage_file, + inherited=inherited, ) as primary, - owned_proxy(primary, directory, overrides, remove_environment=removed, coverage_file=profile.coverage_file) + owned_proxy( + primary, + directory, + overrides, + remove_environment=removed, + coverage_file=profile.coverage_file, + inherited=inherited, + ) if profile.peer else nullcontext(None) as secondary, ): @@ -215,3 +228,36 @@ def worker_gateway(directory: Path, profile: GatewayProfile) -> Iterator[Mapping "INTEGRATION_UPSTREAM_URL": upstream, } ) + + +@dataclass(frozen=True, slots=True) +class WorkerStack: + resources: ExitStack + environment: Mapping[str, str] + + +def worker_stack(results: Path, profile: GatewayProfile) -> WorkerStack: + endpoints: Final = worker_endpoints() + inherited: Final = MappingProxyType({**os.environ, **endpoints, "INTEGRATION_RESULTS_DIR": str(results)}) + resources: Final = ExitStack() + try: + services: Final = resources.enter_context(worker_gateway(results / "runtime", profile, inherited)) + except BaseException: + resources.close() + raise + return WorkerStack(resources, MappingProxyType({**endpoints, **services, "INTEGRATION_RESULTS_DIR": str(results)})) + + +def start_worker_stacks( + directories: tuple[Path, ...], build: Callable[[Path], WorkerStack], resources: ExitStack, concurrency: int +) -> tuple[Mapping[str, str], ...]: + with ThreadPoolExecutor(max_workers=max(1, concurrency), thread_name_prefix="integration-worker-stack") as pool: + futures: Final = tuple(pool.submit(build, directory) for directory in directories) + for future in futures: + if future.exception() is None: + resources.enter_context(future.result().resources) + failure: Final = next((error for future in futures if (error := future.exception()) is not None), None) + if failure is not None: + resources.close() + raise failure + return tuple(future.result().environment for future in futures) diff --git a/tests/integration/conftest.py b/tests/integration/conftest.py index 8bbf4e71842..cb23713e6b2 100644 --- a/tests/integration/conftest.py +++ b/tests/integration/conftest.py @@ -7,6 +7,7 @@ import re from collections.abc import AsyncIterator, Generator, Iterator, Mapping, Sequence from contextlib import ExitStack from copy import copy +from functools import partial from importlib.metadata import version from pathlib import Path from typing import TYPE_CHECKING, Final, Protocol, runtime_checkable @@ -24,7 +25,7 @@ from tests.integration._support.manifest import OWNED_DIRECTORIES from tests.integration._support.ordering import collected_case, fixture_groups, scheduling_groups from tests.integration._support.provider import SharedProvider, shared_provider from tests.integration._support.routing import RoutingPlugin -from tests.integration._support.workers import gateway_profile, worker_endpoints, worker_gateway +from tests.integration._support.workers import gateway_profile, start_worker_stacks, worker_stack from tests.integration.run import GITHUB_FILES if TYPE_CHECKING: @@ -84,21 +85,14 @@ def pytest_xdist_setupnodes(config: pytest.Config, specs: Sequence[XSpec]) -> No resources: Final = ExitStack() config.add_cleanup(resources.close) destination: Final = Path(os.environ["INTEGRATION_RESULTS_DIR"]) - profile: Final = gateway_profile(os.environ) - for index, spec in enumerate(specs): - results: Final = destination / f"gw{index}" - endpoints: Final = worker_endpoints() - with pytest.MonkeyPatch.context() as environment: - for name, value in endpoints.items(): - environment.setenv(name, value) - environment.setenv("INTEGRATION_RESULTS_DIR", str(results)) - services: Final = resources.enter_context(worker_gateway(results / "runtime", profile)) - spec.env = { - **STRING_ENVIRONMENT.validate_python(spec.env), - **endpoints, - **services, - "INTEGRATION_RESULTS_DIR": str(results), - } + stacks: Final = start_worker_stacks( + tuple(destination / f"gw{index}" for index in range(len(specs))), + partial(worker_stack, profile=gateway_profile(os.environ)), + resources, + os.cpu_count() or 1, + ) + for spec, stack in zip(specs, stacks, strict=True): + spec.env = {**STRING_ENVIRONMENT.validate_python(spec.env), **stack} class IntegrationReportPlugin: diff --git a/tests/integration/observability/test_otel_excluded_services.py b/tests/integration/observability/test_otel_excluded_services.py index 1de071e7103..d62537e2317 100644 --- a/tests/integration/observability/test_otel_excluded_services.py +++ b/tests/integration/observability/test_otel_excluded_services.py @@ -3,6 +3,8 @@ from __future__ import annotations import time import uuid from collections.abc import Callable, Iterator, Mapping +from dataclasses import dataclass, replace +from itertools import accumulate from pathlib import Path from types import MappingProxyType from typing import Final @@ -29,6 +31,13 @@ from integration._support.process import owned_proxy, owned_proxy_process from pydantic import JsonValue DB_SYSTEM_KEYS: Final = frozenset({"db.system.name", "db.system"}) +TENANT_EXPORT_INTERVAL_MILLISECONDS: Final = 1000 +TAIL_QUIET_SECONDS: Final = 5 * TENANT_EXPORT_INTERVAL_MILLISECONDS / 1000 + + +@pytest.fixture(autouse=True) +def tenant_export_interval(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv("OTEL_BSP_SCHEDULE_DELAY", str(TENANT_EXPORT_INTERVAL_MILLISECONDS)) @pytest.fixture @@ -125,21 +134,33 @@ def _trace_id(sink_url: str, response: httpx.Response, seconds: float = 40) -> s return found -def _trace_spans(sink_url: str, trace_id: str, seconds: float = 30) -> tuple[Span, ...]: - """The trace's spans once the post-call tail has landed. +@dataclass(frozen=True, slots=True) +class _TraceSample: + at: float + changed_at: float + spans: tuple[Span, ...] - The spend-writer and other post-response spans flush after the request - answers, so absence assertions poll for the whole window instead of - settling at the first glimpse of the root span. - """ - deadline: Final = time.monotonic() + seconds - group: tuple[Span, ...] = () # rebind-ok: drains samples until the post-call tail lands - while time.monotonic() < deadline: - _, spans = recorded_spans(sink_url) - group = spans_for_trace(spans, trace_id) + +def _trace_samples(sink_url: str, trace_id: str) -> Iterator[_TraceSample]: + while True: + now: Final = time.monotonic() + yield _TraceSample(now, now, spans_for_trace(recorded_spans(sink_url)[1], trace_id)) time.sleep(0.5) - assert group, f"trace {trace_id} never reached {sink_url}" - return group + + +def _keep_unchanged_since(previous: _TraceSample, current: _TraceSample) -> _TraceSample: + return current if len(current.spans) != len(previous.spans) else replace(current, changed_at=previous.changed_at) + + +def _trace_spans(sink_url: str, trace_id: str, seconds: float = 30) -> tuple[Span, ...]: + deadline: Final = time.monotonic() + seconds + settled: Final = next( + sample + for sample in accumulate(_trace_samples(sink_url, trace_id), _keep_unchanged_since) + if (sample.spans and sample.at - sample.changed_at >= TAIL_QUIET_SECONDS) or sample.at >= deadline + ) + assert settled.spans, f"trace {trace_id} never reached {sink_url}" + return settled.spans def _trace_spans_when( @@ -231,7 +252,7 @@ def test_excluded_services_drops_db_spans_at_tenant_only( config: Final = _config_with( tmp_path, otel_audit_config, otel={"excluded_services": ["redis", "postgres"]}, extra=_guardrail_block ) - with owned_proxy(gateway, tmp_path, {"LITELLM_OTEL_V2": "1"}, config=config, workers=2) as candidate: + with owned_proxy(gateway, tmp_path, {"LITELLM_OTEL_V2": "1"}, config=config) as candidate: ten_start, _ = recorded_spans(audit_sinks.tenant) op_start, _ = recorded_spans(audit_sinks.operator) traffic: Final = _drive(candidate, langfuse_vars) @@ -266,7 +287,7 @@ def test_without_excluded_services_the_tenant_still_gets_redis_and_postgres_span tmp_path: Path, ) -> None: config: Final = _config_with(tmp_path, otel_audit_config, extra=_guardrail_block) - with owned_proxy(gateway, tmp_path, {"LITELLM_OTEL_V2": "1"}, config=config, workers=2) as candidate: + with owned_proxy(gateway, tmp_path, {"LITELLM_OTEL_V2": "1"}, config=config) as candidate: tenant_start, _ = recorded_spans(audit_sinks.tenant) traffic: Final = _drive(candidate, langfuse_vars) tenant_trace: Final = _trace_id(audit_sinks.tenant, traffic) @@ -294,7 +315,7 @@ def test_a_non_mapping_otel_block_still_publishes_the_tenant_fan_out( config: Final = _config_with(tmp_path, otel_audit_config, extra=with_callback_settings) overrides: Final = {"LITELLM_OTEL_V2": "1", **_operator_langfuse(audit_sinks)} - with owned_proxy(gateway, tmp_path, overrides, config=config, workers=2) as candidate: + with owned_proxy(gateway, tmp_path, overrides, config=config) as candidate: traffic: Final = _drive(candidate, langfuse_vars) tenant_trace: Final = _trace_id(audit_sinks.tenant, traffic) _await_db_span(audit_sinks.tenant, tenant_trace, "redis") @@ -328,7 +349,7 @@ def test_a_bare_excluded_services_env_var_is_ignored( ) -> None: config: Final = _config_with(tmp_path, otel_audit_config) overrides: Final = {"LITELLM_OTEL_V2": "1", name: "redis,postgres"} - with owned_proxy(gateway, tmp_path, overrides, config=config, workers=2) as candidate: + with owned_proxy(gateway, tmp_path, overrides, config=config) as candidate: tenant_start, _ = recorded_spans(audit_sinks.tenant) traffic: Final = _drive(candidate, langfuse_vars) tenant_trace: Final = _trace_id(audit_sinks.tenant, traffic) @@ -360,7 +381,7 @@ def test_the_documented_env_var_wins_over_a_bare_excluded_services( "LITELLM_OTEL_EXCLUDED_SERVICES": "redis", "EXCLUDED_SERVICES": "postgres", } - with owned_proxy(gateway, tmp_path, overrides, config=config, workers=2) as candidate: + with owned_proxy(gateway, tmp_path, overrides, config=config) as candidate: tenant_start, _ = recorded_spans(audit_sinks.tenant) traffic: Final = _drive(candidate, langfuse_vars) tenant_trace: Final = _trace_id(audit_sinks.tenant, traffic) @@ -397,7 +418,6 @@ def test_case_sensitive_otel_settings_read_only_the_exact_env_name( overrides, config=config, remove_environment=("LITELLM_OTEL_EXCLUDED_SERVICES", "litellm_otel_excluded_services"), - workers=2, ) as candidate: tenant_start, _ = recorded_spans(audit_sinks.tenant) traffic: Final = _drive(candidate, langfuse_vars) @@ -431,7 +451,7 @@ def test_env_ignore_empty_keeps_the_default_service_name( ) -> None: config: Final = _config_with(tmp_path, otel_audit_config, otel={"_env_ignore_empty": True}) overrides: Final = {"LITELLM_OTEL_V2": "1", "OTEL_SERVICE_NAME": ""} - with owned_proxy(gateway, tmp_path, overrides, config=config, workers=2) as candidate: + with owned_proxy(gateway, tmp_path, overrides, config=config) as candidate: traffic: Final = _drive(candidate, langfuse_vars) operator_trace: Final = _trace_id(audit_sinks.operator, traffic) operator_spans: Final = _trace_spans_when( @@ -456,7 +476,7 @@ def test_env_parse_none_str_reads_a_null_traces_endpoint_as_unset( ) -> None: config: Final = _config_with(tmp_path, otel_audit_config, otel={"_env_parse_none_str": "null"}) overrides: Final = {"LITELLM_OTEL_V2": "1", "OTEL_TRACES_ENDPOINT": "null"} - with owned_proxy(gateway, tmp_path, overrides, config=config, workers=2) as candidate: + with owned_proxy(gateway, tmp_path, overrides, config=config) as candidate: traffic: Final = _drive(candidate, langfuse_vars) operator_trace: Final = _trace_id(audit_sinks.operator, traffic) operator_spans: Final = _trace_spans_when( @@ -477,7 +497,7 @@ def test_env_excluded_services_drops_only_redis( ) -> None: config: Final = _config_with(tmp_path, otel_audit_config) with owned_proxy( - gateway, tmp_path, {"LITELLM_OTEL_V2": "1", "LITELLM_OTEL_EXCLUDED_SERVICES": "redis"}, config=config, workers=2 + gateway, tmp_path, {"LITELLM_OTEL_V2": "1", "LITELLM_OTEL_EXCLUDED_SERVICES": "redis"}, config=config ) as candidate: start, _ = recorded_spans(audit_sinks.tenant) _drive(candidate, langfuse_vars) @@ -503,7 +523,7 @@ def test_config_excluded_services_wins_over_env( tmp_path, otel_audit_config, otel={"excluded_services": ["postgres"]}, extra=with_langfuse_otel ) with owned_proxy( - gateway, tmp_path, {"LITELLM_OTEL_V2": "1", "LITELLM_OTEL_EXCLUDED_SERVICES": "redis"}, config=config, workers=2 + gateway, tmp_path, {"LITELLM_OTEL_V2": "1", "LITELLM_OTEL_EXCLUDED_SERVICES": "redis"}, config=config ) as candidate: _assert_tenant_keeps_redis_without_postgres(candidate, audit_sinks, langfuse_vars) @@ -523,7 +543,7 @@ def test_excluded_services_applies_with_preset_ordered_first( tmp_path, otel_audit_config, otel={"excluded_services": ["postgres"]}, extra=preset_first ) overrides: Final = {"LITELLM_OTEL_V2": "1", **_operator_langfuse(audit_sinks)} - with owned_proxy(gateway, tmp_path, overrides, config=config, workers=2) as candidate: + with owned_proxy(gateway, tmp_path, overrides, config=config) as candidate: _assert_tenant_keeps_redis_without_postgres(candidate, audit_sinks, langfuse_vars) @@ -536,7 +556,7 @@ def test_bogus_excluded_service_logs_error_and_drops_at_proxy_start( tmp_path: Path, ) -> None: config: Final = _config_with(tmp_path, otel_audit_config, otel={"excluded_services": ["auth", "postgres"]}) - with owned_proxy_process(gateway, tmp_path, {"LITELLM_OTEL_V2": "1"}, config=config, workers=2) as owned: + with owned_proxy_process(gateway, tmp_path, {"LITELLM_OTEL_V2": "1"}, config=config) as owned: assert "'auth' is not a datastore service; ignored" in owned.log.read_text(), owned.log.read_text()[-3000:] _assert_tenant_keeps_redis_without_postgres(owned.gateway, audit_sinks, langfuse_vars) @@ -551,7 +571,7 @@ def test_valid_config_excluded_services_tolerates_bogus_env( ) -> None: config: Final = _config_with(tmp_path, otel_audit_config, otel={"excluded_services": ["postgres"]}) with owned_proxy( - gateway, tmp_path, {"LITELLM_OTEL_V2": "1", "LITELLM_OTEL_EXCLUDED_SERVICES": "auth"}, config=config, workers=2 + gateway, tmp_path, {"LITELLM_OTEL_V2": "1", "LITELLM_OTEL_EXCLUDED_SERVICES": "auth"}, config=config ) as candidate: _assert_tenant_keeps_redis_without_postgres(candidate, audit_sinks, langfuse_vars) @@ -569,7 +589,7 @@ def test_bogus_excluded_services_env_logs_and_drops_with_preset_alongside_otel( config: Final = _config_with(tmp_path, otel_audit_config, extra=with_langfuse_otel) overrides: Final = {"LITELLM_OTEL_V2": "1", "LITELLM_OTEL_EXCLUDED_SERVICES": "auth,postgres"} - with owned_proxy_process(gateway, tmp_path, overrides, config=config, workers=2) as owned: + with owned_proxy_process(gateway, tmp_path, overrides, config=config) as owned: assert "'auth' is not a datastore service; ignored" in owned.log.read_text(), owned.log.read_text()[-3000:] _assert_tenant_keeps_redis_without_postgres(owned.gateway, audit_sinks, langfuse_vars) @@ -591,7 +611,7 @@ def test_bogus_excluded_services_env_logs_and_drops_without_otel_callback( "LITELLM_OTEL_EXCLUDED_SERVICES": "auth,postgres", **_operator_langfuse(audit_sinks), } - with owned_proxy_process(gateway, tmp_path, overrides, config=config, workers=2) as owned: + with owned_proxy_process(gateway, tmp_path, overrides, config=config) as owned: assert "'auth' is not a datastore service; ignored" in owned.log.read_text(), owned.log.read_text()[-3000:] _assert_tenant_keeps_redis_without_postgres(owned.gateway, audit_sinks, langfuse_vars) @@ -604,7 +624,7 @@ def test_postgres_exclusion_covers_spend_flush( tmp_path: Path, ) -> None: config: Final = _config_with(tmp_path, otel_audit_config, otel={"excluded_services": ["postgres"]}) - with owned_proxy(gateway, tmp_path, {"LITELLM_OTEL_V2": "1"}, config=config, workers=2) as candidate: + with owned_proxy(gateway, tmp_path, {"LITELLM_OTEL_V2": "1"}, config=config) as candidate: op_start, _ = recorded_spans(audit_sinks.operator) ten_start, _ = recorded_spans(audit_sinks.tenant) traffic: Final = _drive(candidate, langfuse_vars) diff --git a/tests/integration/observability/test_passthrough_upstream_error_visibility.py b/tests/integration/observability/test_passthrough_upstream_error_visibility.py index bc6747844ba..0d86440045b 100644 --- a/tests/integration/observability/test_passthrough_upstream_error_visibility.py +++ b/tests/integration/observability/test_passthrough_upstream_error_visibility.py @@ -126,7 +126,7 @@ async def test_gemini_passthrough_async_client_404_body_reaches_proxy_log_and_sp path: Final = tmp_path / "gemini-async.yaml" with wire_server(respond) as wire: _gemini_config(path, wire.url) - with owned_proxy_process(gateway, tmp_path, {}, config=path, workers=2) as owned: + with owned_proxy_process(gateway, tmp_path, {}, config=path) as owned: candidate: Final = owned.gateway async with httpx.AsyncClient( base_url=str(candidate.client.base_url), timeout=15, trust_env=False @@ -156,7 +156,7 @@ def test_gemini_passthrough_streaming_500_relays_full_body_and_logs_bounded_prev path: Final = tmp_path / "gemini-stream-500.yaml" with wire_server(respond) as wire: _gemini_config(path, wire.url) - with owned_proxy_process(gateway, tmp_path, {}, config=path, workers=2) as owned: + with owned_proxy_process(gateway, tmp_path, {}, config=path) as owned: candidate: Final = owned.gateway with candidate.client.stream( "POST", @@ -188,7 +188,7 @@ def test_gemini_passthrough_success_logs_nothing_and_spend_row_is_success(gatewa path: Final = tmp_path / "gemini-200.yaml" with wire_server(respond) as wire: _gemini_config(path, wire.url) - with owned_proxy_process(gateway, tmp_path, {}, config=path, workers=2) as owned: + with owned_proxy_process(gateway, tmp_path, {}, config=path) as owned: candidate: Final = owned.gateway response: Final = candidate.request( "POST", _GEMINI_MODEL_PATH, _GENERATE_CONTENT, headers={"x-goog-api-key": candidate.key} @@ -208,7 +208,7 @@ def test_gemini_passthrough_streaming_200_relays_every_chunk(gateway: Gateway, t path: Final = tmp_path / "gemini-stream-200.yaml" with wire_server(respond) as wire: _gemini_config(path, wire.url) - with owned_proxy_process(gateway, tmp_path, {}, config=path, workers=2) as owned: + with owned_proxy_process(gateway, tmp_path, {}, config=path) as owned: candidate: Final = owned.gateway with candidate.client.stream( "POST", @@ -239,7 +239,7 @@ def test_gemini_streaming_passthrough_request_tags_reach_spend_log(gateway: Gate config_path: Final = tmp_path / "gemini-stream-tags.yaml" with wire_server(respond) as wire: _gemini_config(config_path, wire.url) - with owned_proxy_process(gateway, tmp_path, {}, config=config_path, workers=2) as owned: + with owned_proxy_process(gateway, tmp_path, {}, config=config_path) as owned: candidate: Final = owned.gateway with candidate.client.stream( "POST", @@ -293,7 +293,7 @@ def test_config_pass_through_route_logs_body_and_strips_query(gateway: Gateway, } ] 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) as owned: candidate: Final = owned.gateway response: Final = candidate.request("POST", "/audit-pt", _GENERATE_CONTENT) assert response.status_code == 403, response.text @@ -324,7 +324,7 @@ def test_openai_passthrough_sdk_error_body_reaches_proxy_log_and_spend_row(gatew path: Final = tmp_path / "openai-404.yaml" with wire_server(respond) as wire: _openai_config(path, wire.url) - with owned_proxy_process(gateway, tmp_path, {}, config=path, workers=2) as owned: + with owned_proxy_process(gateway, tmp_path, {}, config=path) as owned: candidate: Final = owned.gateway with OpenAI( api_key=candidate.key, @@ -350,7 +350,7 @@ async def test_openai_passthrough_async_sdk_error_body_reaches_proxy_log_and_spe path: Final = tmp_path / "openai-async-404.yaml" with wire_server(respond) as wire: _openai_config(path, wire.url) - with owned_proxy_process(gateway, tmp_path, {}, config=path, workers=2) as owned: + with owned_proxy_process(gateway, tmp_path, {}, config=path) as owned: candidate: Final = owned.gateway async with AsyncOpenAI( api_key=candidate.key, @@ -377,7 +377,7 @@ def test_gemini_passthrough_control_characters_cannot_forge_log_lines(gateway: G path: Final = tmp_path / "gemini-forged.yaml" with wire_server(respond) as wire: _gemini_config(path, wire.url) - with owned_proxy_process(gateway, tmp_path, {}, config=path, workers=2) as owned: + with owned_proxy_process(gateway, tmp_path, {}, config=path) as owned: candidate: Final = owned.gateway response: Final = candidate.request( "POST", _GEMINI_MODEL_PATH, _GENERATE_CONTENT, headers={"x-goog-api-key": candidate.key} @@ -401,7 +401,7 @@ def test_gemini_passthrough_empty_error_body_still_logged_and_proxy_serves(gatew path: Final = tmp_path / "gemini-empty.yaml" with wire_server(respond) as wire: _gemini_config(path, wire.url) - with owned_proxy_process(gateway, tmp_path, {}, config=path, workers=2) as owned: + with owned_proxy_process(gateway, tmp_path, {}, config=path) as owned: candidate: Final = owned.gateway response: Final = candidate.request( "POST", _GEMINI_MODEL_PATH, _GENERATE_CONTENT, headers={"x-goog-api-key": candidate.key} @@ -434,7 +434,7 @@ def test_gemini_passthrough_gzip_error_body_decoded_for_log_and_client(gateway: path: Final = tmp_path / "gemini-gzip.yaml" with wire_server(respond) as wire: _gemini_config(path, wire.url) - with owned_proxy_process(gateway, tmp_path, {}, config=path, workers=2) as owned: + with owned_proxy_process(gateway, tmp_path, {}, config=path) as owned: candidate: Final = owned.gateway response: Final = candidate.request( "POST", _GEMINI_MODEL_PATH, _GENERATE_CONTENT, headers={"x-goog-api-key": candidate.key} @@ -462,7 +462,7 @@ def test_gemini_passthrough_streaming_gzip_error_body_decoded_for_log_and_client path: Final = tmp_path / "gemini-stream-gzip.yaml" with wire_server(respond) as wire: _gemini_config(path, wire.url) - with owned_proxy_process(gateway, tmp_path, {}, config=path, workers=2) as owned: + with owned_proxy_process(gateway, tmp_path, {}, config=path) as owned: candidate: Final = owned.gateway with candidate.client.stream( "POST", @@ -490,7 +490,7 @@ def test_gemini_passthrough_error_body_redacted_when_message_logging_off(gateway config["environment_variables"] = {"GEMINI_API_BASE": wire.url, "GEMINI_API_KEY": "scripted"} config["litellm_settings"]["turn_off_message_logging"] = True 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) as owned: candidate: Final = owned.gateway response: Final = candidate.request( "POST", _GEMINI_MODEL_PATH, _GENERATE_CONTENT, headers={"x-goog-api-key": candidate.key} @@ -516,7 +516,7 @@ def test_gemini_passthrough_exact_4096_byte_body_logged_without_marker(gateway: path: Final = tmp_path / "gemini-exact.yaml" with wire_server(respond) as wire: _gemini_config(path, wire.url) - with owned_proxy_process(gateway, tmp_path, {}, config=path, workers=2) as owned: + with owned_proxy_process(gateway, tmp_path, {}, config=path) as owned: candidate: Final = owned.gateway response: Final = candidate.request( "POST", _GEMINI_MODEL_PATH, _GENERATE_CONTENT, headers={"x-goog-api-key": candidate.key} @@ -537,7 +537,7 @@ def test_gemini_passthrough_4097_byte_body_truncated_with_marker(gateway: Gatewa path: Final = tmp_path / "gemini-over.yaml" with wire_server(respond) as wire: _gemini_config(path, wire.url) - with owned_proxy_process(gateway, tmp_path, {}, config=path, workers=2) as owned: + with owned_proxy_process(gateway, tmp_path, {}, config=path) as owned: candidate: Final = owned.gateway response: Final = candidate.request( "POST", _GEMINI_MODEL_PATH, _GENERATE_CONTENT, headers={"x-goog-api-key": candidate.key} @@ -557,7 +557,7 @@ def test_gemini_passthrough_one_byte_stream_chunks_reassembled_and_logged(gatewa path: Final = tmp_path / "gemini-one-byte.yaml" with wire_server(respond) as wire: _gemini_config(path, wire.url) - with owned_proxy_process(gateway, tmp_path, {}, config=path, workers=2) as owned: + with owned_proxy_process(gateway, tmp_path, {}, config=path) as owned: candidate: Final = owned.gateway with candidate.client.stream( "POST", @@ -580,7 +580,7 @@ def test_gemini_passthrough_repeated_errors_each_get_row_and_log_line(gateway: G path: Final = tmp_path / "gemini-twice.yaml" with wire_server(respond) as wire: _gemini_config(path, wire.url) - with owned_proxy_process(gateway, tmp_path, {}, config=path, workers=2) as owned: + with owned_proxy_process(gateway, tmp_path, {}, config=path) as owned: candidate: Final = owned.gateway responses: Final = tuple( candidate.request( @@ -604,7 +604,7 @@ def test_gemini_passthrough_repeated_errors_each_get_row_and_log_line(gateway: G def test_budget_rejected_call_keeps_budget_normalized_error(gateway: Gateway, tmp_path: Path) -> None: path: Final = tmp_path / "budget.yaml" path.write_text(Path("tests/integration/proxy_config.yaml").read_text()) - with owned_proxy_process(gateway, tmp_path, {}, config=path, workers=2) as owned: + with owned_proxy_process(gateway, tmp_path, {}, config=path) as owned: candidate: Final = owned.gateway with candidate.scenario() as scenario: model: Final = scenario.model() diff --git a/tests/integration/observability/test_s3_v2_streaming_failure_dedupe.py b/tests/integration/observability/test_s3_v2_streaming_failure_dedupe.py index 003181c72d8..d35a644ea3d 100644 --- a/tests/integration/observability/test_s3_v2_streaming_failure_dedupe.py +++ b/tests/integration/observability/test_s3_v2_streaming_failure_dedupe.py @@ -174,7 +174,6 @@ def test_retried_failure_uploads_one_s3_object( tmp_path, {"DEFAULT_S3_FLUSH_INTERVAL_SECONDS": "2"}, config=config, - workers=2, ) as candidate, candidate.scenario() as scenario, ): @@ -216,7 +215,6 @@ def test_fallback_chain_failure_uploads_one_s3_object( tmp_path, {"DEFAULT_S3_FLUSH_INTERVAL_SECONDS": "2"}, config=config, - workers=2, ) as candidate, candidate.scenario() as scenario, ): @@ -268,7 +266,6 @@ def test_streaming_success_uploads_one_s3_object( tmp_path, {"DEFAULT_S3_FLUSH_INTERVAL_SECONDS": "2"}, config=config, - workers=2, ) as candidate, candidate.scenario() as scenario, ): diff --git a/tests/integration/spend/test_spend_log_tool_payload_content.py b/tests/integration/spend/test_spend_log_tool_payload_content.py index 22ed6c74756..672c0a38407 100644 --- a/tests/integration/spend/test_spend_log_tool_payload_content.py +++ b/tests/integration/spend/test_spend_log_tool_payload_content.py @@ -325,7 +325,7 @@ def test_stored_chat_response_keeps_logprob_tokens(gateway: Gateway, tmp_path: P with ( wire_server(_answering_model_listing(respond)) as wire, - owned_proxy(gateway, tmp_path, {}, config=_prompt_storage_config(tmp_path), workers=2) as isolated, + owned_proxy(gateway, tmp_path, {}, config=_prompt_storage_config(tmp_path)) as isolated, isolated.scenario() as scenario, ): model: Final = scenario.model( @@ -418,7 +418,7 @@ def test_stored_messages_keep_tool_use_input(gateway: Gateway, tmp_path: Path) - with ( wire_server(_answering_model_listing(respond)) as wire, - owned_proxy(gateway, tmp_path, {}, config=_prompt_storage_config(tmp_path), workers=2) as isolated, + owned_proxy(gateway, tmp_path, {}, config=_prompt_storage_config(tmp_path)) as isolated, isolated.scenario() as scenario, ): model: Final = scenario.model( @@ -650,7 +650,7 @@ def test_chat_sdk_keeps_logprob_tokens(gateway: Gateway, tmp_path: Path, client_ with ( wire_server(_answering_model_listing(respond)) as wire, - owned_proxy(gateway, tmp_path, {}, config=_prompt_storage_config(tmp_path), workers=2) as isolated, + owned_proxy(gateway, tmp_path, {}, config=_prompt_storage_config(tmp_path)) as isolated, isolated.scenario() as scenario, ): model: Final = scenario.model( @@ -702,7 +702,7 @@ def test_chat_sdk_stream_include_usage_masks_request_fields(gateway: Gateway, tm with ( wire_server(_answering_model_listing(respond)) as wire, - owned_proxy(gateway, tmp_path, {}, config=_prompt_storage_config(tmp_path), workers=2) as isolated, + owned_proxy(gateway, tmp_path, {}, config=_prompt_storage_config(tmp_path)) as isolated, isolated.scenario() as scenario, ): model: Final = scenario.model( @@ -747,7 +747,7 @@ def test_chat_history_keeps_string_tool_arguments_and_tool_content(gateway: Gate with ( wire_server(_answering_model_listing(respond)) as wire, - owned_proxy(gateway, tmp_path, {}, config=_prompt_storage_config(tmp_path), workers=2) as isolated, + owned_proxy(gateway, tmp_path, {}, config=_prompt_storage_config(tmp_path)) as isolated, isolated.scenario() as scenario, ): model: Final = scenario.model(model="openai/gpt-4o-mini", api_base=wire.url, api_key="synthetic-openai-key") @@ -833,7 +833,7 @@ def test_messages_sdk_keeps_tool_use_input(gateway: Gateway, tmp_path: Path, cli with ( wire_server(_answering_model_listing(respond)) as wire, - owned_proxy(gateway, tmp_path, {}, config=_prompt_storage_config(tmp_path), workers=2) as isolated, + owned_proxy(gateway, tmp_path, {}, config=_prompt_storage_config(tmp_path)) as isolated, isolated.scenario() as scenario, ): model: Final = scenario.model(model=ANTHROPIC_MODEL, api_base=wire.url, api_key="synthetic-anthropic-key") @@ -879,7 +879,7 @@ def test_messages_stream_keeps_tool_use_input(gateway: Gateway, tmp_path: Path) with ( wire_server(_answering_model_listing(respond)) as wire, - owned_proxy(gateway, tmp_path, {}, config=_prompt_storage_config(tmp_path), workers=2) as isolated, + owned_proxy(gateway, tmp_path, {}, config=_prompt_storage_config(tmp_path)) as isolated, isolated.scenario() as scenario, ): model: Final = scenario.model(model=ANTHROPIC_MODEL, api_base=wire.url, api_key="synthetic-anthropic-key") @@ -925,7 +925,7 @@ def test_responses_stream_keeps_function_call_arguments(gateway: Gateway, tmp_pa with ( wire_server(_answering_model_listing(respond)) as wire, - owned_proxy(gateway, tmp_path, {}, config=_prompt_storage_config(tmp_path), workers=2) as isolated, + owned_proxy(gateway, tmp_path, {}, config=_prompt_storage_config(tmp_path)) as isolated, isolated.scenario() as scenario, ): model: Final = scenario.model(model=ANTHROPIC_MODEL, api_base=wire.url, api_key="synthetic-anthropic-key") @@ -1006,7 +1006,7 @@ def test_native_responses_keeps_logprob_tokens(gateway: Gateway, tmp_path: Path) with ( wire_server(_answering_model_listing(respond)) as wire, - owned_proxy(gateway, tmp_path, {}, config=_prompt_storage_config(tmp_path), workers=2) as isolated, + owned_proxy(gateway, tmp_path, {}, config=_prompt_storage_config(tmp_path)) as isolated, isolated.scenario() as scenario, ): model: Final = scenario.model(model="openai/gpt-4o-mini", api_base=wire.url, api_key="synthetic-openai-key") @@ -1052,7 +1052,7 @@ def test_malformed_tool_blocks_keep_only_recognized_content(gateway: Gateway, tm with ( wire_server(_answering_model_listing(respond)) as wire, - owned_proxy(gateway, tmp_path, {}, config=_prompt_storage_config(tmp_path), workers=2) as isolated, + owned_proxy(gateway, tmp_path, {}, config=_prompt_storage_config(tmp_path)) as isolated, isolated.scenario() as scenario, ): model: Final = scenario.model(model="openai/gpt-4o-mini", api_base=wire.url, api_key="synthetic-openai-key") @@ -1097,7 +1097,7 @@ def test_messages_tool_input_handles_mixed_values_and_truncation(gateway: Gatewa with ( wire_server(_answering_model_listing(respond)) as wire, - owned_proxy(gateway, tmp_path, {}, config=_prompt_storage_config(tmp_path), workers=2) as isolated, + owned_proxy(gateway, tmp_path, {}, config=_prompt_storage_config(tmp_path)) as isolated, isolated.scenario() as scenario, ): model: Final = scenario.model(model=ANTHROPIC_MODEL, api_base=wire.url, api_key="synthetic-anthropic-key") @@ -1143,7 +1143,7 @@ def test_messages_without_auth_create_no_spend_row(gateway: Gateway, tmp_path: P with ( wire_server(respond) as wire, - owned_proxy(gateway, tmp_path, {}, config=_prompt_storage_config(tmp_path), workers=2) as isolated, + owned_proxy(gateway, tmp_path, {}, config=_prompt_storage_config(tmp_path)) as isolated, ): response: Final = isolated.client.post( "/v1/messages", @@ -1182,7 +1182,7 @@ def test_messages_upstream_error_keeps_tool_input(gateway: Gateway, tmp_path: Pa with ( wire_server(_answering_model_listing(respond)) as wire, - owned_proxy(gateway, tmp_path, {}, config=_prompt_storage_config(tmp_path), workers=2) as isolated, + owned_proxy(gateway, tmp_path, {}, config=_prompt_storage_config(tmp_path)) as isolated, isolated.scenario() as scenario, ): model: Final = scenario.model(model=ANTHROPIC_MODEL, api_base=wire.url, api_key="synthetic-anthropic-key") @@ -1235,7 +1235,6 @@ def test_store_prompts_off_keeps_chat_and_messages_representation_equal(gateway: tmp_path, {}, config=_prompt_storage_config(tmp_path, store_prompts=False), - workers=2, ) as isolated, isolated.scenario() as scenario, ): @@ -1283,7 +1282,7 @@ def test_identical_messages_requests_have_distinct_spend_rows(gateway: Gateway, with ( wire_server(_answering_model_listing(respond)) as wire, - owned_proxy(gateway, tmp_path, {}, config=_prompt_storage_config(tmp_path), workers=2) as isolated, + owned_proxy(gateway, tmp_path, {}, config=_prompt_storage_config(tmp_path)) as isolated, isolated.scenario() as scenario, ): model: Final = scenario.model(model=ANTHROPIC_MODEL, api_base=wire.url, api_key="synthetic-anthropic-key") diff --git a/tests/unit/integration_support/test_process.py b/tests/unit/integration_support/test_process.py index 998f81155a3..b2271a1ed47 100644 --- a/tests/unit/integration_support/test_process.py +++ b/tests/unit/integration_support/test_process.py @@ -11,6 +11,7 @@ import threading import time from collections.abc import Iterator from contextlib import contextmanager +from functools import reduce from pathlib import Path from types import ModuleType from typing import Final @@ -238,6 +239,45 @@ def test_only_live_readable_workers_can_satisfy_readiness( process_module._wait_until_ready(launch) +def test_boot_phases_mark_when_prisma_ran_the_port_bound_and_each_worker_answered( + process_module: ModuleType, +) -> None: + probe: Final = process_module._Probe + timeline: Final = ( + probe(0.5, listening=False, prisma=False, ready_pid=None), + probe(4.0, listening=False, prisma=True, ready_pid=None), + probe(5.0, listening=False, prisma=False, ready_pid=None), + probe(5.5, listening=False, prisma=True, ready_pid=None), + probe(6.5, listening=False, prisma=False, ready_pid=None), + probe(7.0, listening=True, prisma=False, ready_pid=None), + probe(10.0, listening=True, prisma=False, ready_pid=41), + probe(11.0, listening=True, prisma=False, ready_pid=42), + ) + phases: Final = reduce(process_module._observe, timeline, process_module.BootPhases()) + assert phases == process_module.BootPhases( + prisma_started=4.0, + prisma_finished=6.5, + listening=7.0, + first_ready=10.0, + ready_pids=frozenset((41, 42)), + ) + + +def test_boot_phases_leave_prisma_unset_when_no_prisma_child_was_seen(process_module: ModuleType) -> None: + probe: Final = process_module._Probe + timeline: Final = ( + probe(1.0, listening=False, prisma=False, ready_pid=None), + probe(2.0, listening=True, prisma=False, ready_pid=7), + ) + phases: Final = reduce(process_module._observe, timeline, process_module.BootPhases()) + assert (phases.prisma_started, phases.prisma_finished, phases.listening, phases.first_ready) == ( + None, + None, + 2.0, + 2.0, + ) + + @pytest.mark.parametrize( ("writer", "explicit_reader", "expected_reader"), ( diff --git a/tests/unit/integration_support/test_workers.py b/tests/unit/integration_support/test_workers.py index 68c43335e32..30e55a91b11 100644 --- a/tests/unit/integration_support/test_workers.py +++ b/tests/unit/integration_support/test_workers.py @@ -1,12 +1,21 @@ +import threading +from contextlib import ExitStack from dataclasses import dataclass from pathlib import Path +from types import MappingProxyType from typing import Final import httpx import pytest from tests.integration._support.client import Gateway -from tests.integration._support.workers import GatewayProfile, gateway_profile, readiness_failure +from tests.integration._support.workers import ( + GatewayProfile, + WorkerStack, + gateway_profile, + readiness_failure, + start_worker_stacks, +) from tests.integration.run import worker_settings HEALTHY: Final = { @@ -82,3 +91,43 @@ def test_replica_runs_and_coverage_runs_reach_every_worker_gateway() -> None: assert profile == GatewayProfile( peer=True, replica=True, cost_map=False, coverage_file=Path("/results/coverage/data") ) + + +def _stack(name: str, closed: threading.Event) -> WorkerStack: + resources: Final = ExitStack() + resources.callback(closed.set) + return WorkerStack(resources, MappingProxyType({"INTEGRATION_RESULTS_DIR": name})) + + +def test_worker_stacks_start_together_and_return_in_worker_order(tmp_path: Path) -> None: + directories: Final = tuple(tmp_path / f"gw{index}" for index in range(3)) + closed: Final = MappingProxyType({directory: threading.Event() for directory in directories}) + together: Final = threading.Barrier(len(directories), timeout=10) + + def build(directory: Path) -> WorkerStack: + together.wait() + return _stack(directory.name, closed[directory]) + + with ExitStack() as resources: + environments: Final = start_worker_stacks(directories, build, resources, len(directories)) + assert [environment["INTEGRATION_RESULTS_DIR"] for environment in environments] == ["gw0", "gw1", "gw2"] + assert not any(event.is_set() for event in closed.values()) + assert all(event.is_set() for event in closed.values()) + + +def test_a_failed_worker_stack_stops_every_sibling_that_started(tmp_path: Path) -> None: + directories: Final = tuple(tmp_path / f"gw{index}" for index in range(3)) + closed: Final = MappingProxyType({directory: threading.Event() for directory in directories}) + together: Final = threading.Barrier(len(directories), timeout=10) + + def build(directory: Path) -> WorkerStack: + together.wait() + if directory.name == "gw1": + raise RuntimeError("gw1 gateway never became ready") + return _stack(directory.name, closed[directory]) + + with ExitStack() as resources: + with pytest.raises(RuntimeError, match="gw1 gateway never became ready"): + start_worker_stacks(directories, build, resources, len(directories)) + assert closed[directories[0]].is_set() + assert closed[directories[2]].is_set()