mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-08 08:08:37 +00:00
b602209ca6
* test(integration): hand every harness server a port from one per-worker block
About 45 places in the integration harness picked a port by binding port 0,
reading the number back, closing the socket and starting a server on that
number later. Another process could take the port in between, which surfaced
as "Address already in use" flakes. Earlier fixes retried around the race in
single call sites without removing it.
tests/integration/_support/ports.py now owns every port the harness picks
ahead of a server. Each pytest process claims a 200-port block between 10000
and 32767 by taking an exclusive flock on a per-block lock file in the system
temp directory, so parallel runs on one machine never share a block and the
lock goes away with the process. The range is above the fixed service ports
and below the Linux and macOS ephemeral ranges. Ports go out in order, are
never reused in a run, skip anything an outside process already holds, and
a process that exhausts its block claims the next free one.
The three _free_port copies, worker_endpoints' port-0 binds,
free_loopback_port, the owned proxy's lost-port-race retry and the owned
Redis retry loop are gone. Tests that restart a server on the same port now
allocate that port up front instead of reading it back from the first
server. TQ012 in scripts/check_test_quality.py rejects a port-0 bind read
back with getsockname anywhere under tests/integration.
* ci(integration): record listening sockets and the ephemeral port range
* Revert "ci(integration): record listening sockets and the ephemeral port range"
This reverts commit d3fafc3da2.
759 lines
37 KiB
Python
759 lines
37 KiB
Python
from __future__ import annotations
|
|
|
|
import threading
|
|
import time
|
|
import uuid
|
|
from collections.abc import Callable
|
|
from concurrent.futures import ThreadPoolExecutor
|
|
from itertools import accumulate, repeat
|
|
from pathlib import Path
|
|
from queue import SimpleQueue
|
|
from typing import Final
|
|
|
|
import httpx
|
|
import psutil
|
|
import pytest
|
|
from _logging_only_scope_support import (
|
|
JSON_OBJECT,
|
|
CallerResult,
|
|
ChaosCall,
|
|
_assert_response_id,
|
|
_call_client,
|
|
_chaos_calls,
|
|
_chaos_control_configuration,
|
|
_chaos_model_list,
|
|
_chaos_models,
|
|
_chaos_spend_minimums,
|
|
_configuration,
|
|
_direction,
|
|
_directions_for_audit_leg,
|
|
_drain_upstream,
|
|
_guardrail_entries,
|
|
_is_base_audit_leg,
|
|
_json_contains_exact_string,
|
|
_policy_call_id_matches,
|
|
_spend_rows,
|
|
_spend_rows_for_calls,
|
|
_spend_rows_matching_call,
|
|
wire_server,
|
|
)
|
|
from _logging_only_scope_support import (
|
|
_record_audit_properties as _record_audit_properties,
|
|
)
|
|
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 integration._support.ports import allocate_port
|
|
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:
|
|
identity: Final = f"logging-scope-k1-{uuid.uuid4().hex}"
|
|
marker: Final = uuid.uuid4().hex
|
|
|
|
def policy(_request: Request) -> Reply:
|
|
return Reply(body=b'{"action":"NONE"}')
|
|
|
|
policy_port: Final = allocate_port()
|
|
with gateway.scenario() as scenario:
|
|
deployments: Final = _chaos_models(scenario, marker)
|
|
models: Final = tuple(deployment.model_name for deployment in deployments)
|
|
model_list: Final = _chaos_model_list(deployments)
|
|
chaos_calls: Final = _chaos_calls(deployments, marker)
|
|
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",
|
|
)
|
|
for call in chaos_calls
|
|
)
|
|
control_upstream: Final = _drain_upstream(gateway.upstream_url)
|
|
assert len(control_upstream) == 30, control_upstream
|
|
assert (
|
|
tuple(
|
|
sum(_json_contains_exact_string(observation["body"], call.prompt) for observation in control_upstream)
|
|
for call in chaos_calls
|
|
)
|
|
== (1,) * 30
|
|
), control_upstream
|
|
config: Final = _configuration(
|
|
tmp_path,
|
|
identity,
|
|
f"http://127.0.0.1:{policy_port}",
|
|
"output",
|
|
model_list=model_list,
|
|
)
|
|
with owned_proxy(gateway, tmp_path, {}, config=config, workers=2) as candidate:
|
|
starts: Final = tuple(threading.Event() for _ in range(3))
|
|
|
|
def run(index: int) -> tuple[int, CallerResult]:
|
|
call: Final = chaos_calls[index]
|
|
phase: Final = index // 10
|
|
assert starts[phase].wait(timeout=90), (index, phase)
|
|
return index, _call_client(
|
|
call.client_kind,
|
|
call.endpoint,
|
|
candidate,
|
|
call.model,
|
|
call.prompt,
|
|
call.stream,
|
|
call.call_id,
|
|
)
|
|
|
|
with ThreadPoolExecutor(max_workers=30) as pool:
|
|
futures: Final = tuple(pool.submit(run, index) for index in range(30))
|
|
try:
|
|
with wire_server(policy, port=policy_port) as initial_edge:
|
|
starts[0].set()
|
|
first: Final = tuple(futures[index].result(timeout=90) for index in range(10))
|
|
eventually(
|
|
lambda: initial_edge.received.qsize(),
|
|
lambda count: count == 10 * len(expected_directions),
|
|
seconds=30,
|
|
)
|
|
tuple(
|
|
_spend_rows(model, minimum)
|
|
for model, minimum in zip(models, _chaos_spend_minimums(models, chaos_calls[:10], 5))
|
|
)
|
|
starts[1].set()
|
|
middle: Final = tuple(futures[index].result(timeout=90) for index in range(10, 20))
|
|
tuple(
|
|
_spend_rows(model, minimum)
|
|
for model, minimum in zip(models, _chaos_spend_minimums(models, chaos_calls[:20], 5))
|
|
)
|
|
with wire_server(policy, port=policy_port) as recovered_edge:
|
|
starts[2].set()
|
|
recovered: Final = tuple(futures[index].result(timeout=90) for index in range(20, 30))
|
|
eventually(
|
|
lambda: recovered_edge.received.qsize(),
|
|
lambda count: count == 10 * len(expected_directions),
|
|
seconds=30,
|
|
)
|
|
tuple(_spend_rows(model, 10) for model in models)
|
|
finally:
|
|
for start in starts:
|
|
start.set()
|
|
results: Final = first + middle + recovered
|
|
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
|
|
candidate_ids: Final = tuple(result.response_id for _, result in results)
|
|
assert len(set(candidate_ids)) == 30, candidate_ids
|
|
observed_upstream: Final = _drain_upstream(gateway.upstream_url)
|
|
assert len(observed_upstream) == 30, observed_upstream
|
|
assert (
|
|
tuple(
|
|
sum(
|
|
_json_contains_exact_string(observation["body"], call.prompt)
|
|
for observation in observed_upstream
|
|
)
|
|
for call in chaos_calls
|
|
)
|
|
== (1,) * 30
|
|
), observed_upstream
|
|
expected_success_ids: Final = frozenset(
|
|
call.call_id for call in chaos_calls if call.index < 10 or call.index >= 20
|
|
)
|
|
edge_payloads: Final = tuple(
|
|
JSON_OBJECT.validate_json(call.body) for call in initial_edge.drain() + recovered_edge.drain()
|
|
)
|
|
assert len(edge_payloads) == 20 * len(expected_directions), edge_payloads
|
|
successful_calls: Final = tuple(call for call in chaos_calls if call.call_id in expected_success_ids)
|
|
for call in successful_calls:
|
|
payloads_for_call: Final = tuple(
|
|
payload for payload in edge_payloads if _policy_call_id_matches(payload, call.call_id)
|
|
)
|
|
assert tuple(sorted(_direction(payload) for payload in payloads_for_call)) == tuple(
|
|
sorted(expected_directions)
|
|
), (
|
|
call.call_id,
|
|
payloads_for_call,
|
|
)
|
|
assert all(
|
|
payload["texts"]
|
|
== ([call.prompt] if _direction(payload) == "request" else [results[call.index][1].text])
|
|
for payload in payloads_for_call
|
|
), payloads_for_call
|
|
rows: Final = _spend_rows_for_calls(
|
|
models,
|
|
tuple(
|
|
(
|
|
chaos_calls[index].model,
|
|
response_id,
|
|
chaos_calls[index].call_id,
|
|
)
|
|
for index, response_id in enumerate(candidate_ids)
|
|
),
|
|
)
|
|
for index, call in enumerate(chaos_calls):
|
|
matching_rows: Final = _spend_rows_matching_call(rows, call.model, call.call_id)
|
|
assert len(matching_rows) == 1, (call.call_id, matching_rows)
|
|
row: Final = matching_rows[0]
|
|
expected_status: Final = "guardrail_failed_to_respond" if 10 <= index < 20 else "success"
|
|
entries: Final = _guardrail_entries(row)
|
|
assert tuple(
|
|
(entry["guardrail_name"], entry["guardrail_mode"], entry["guardrail_status"]) for entry in entries
|
|
) == tuple((identity, "logging_only", expected_status) for _ in expected_directions), (index, entries)
|
|
|
|
|
|
@pytest.mark.timeout(180)
|
|
def test_K2_policy_edge_delay_does_not_delay_concurrent_callers(gateway: Gateway, tmp_path: Path) -> None:
|
|
identity: Final = f"logging-scope-k2-{uuid.uuid4().hex}"
|
|
marker: Final = uuid.uuid4().hex
|
|
|
|
def policy(_request: Request) -> Reply:
|
|
time.sleep(_EDGE_DELAY_SECONDS)
|
|
return Reply(body=b'{"action":"NONE"}')
|
|
|
|
with gateway.scenario() as scenario:
|
|
deployments: Final = _chaos_models(scenario, marker)
|
|
models: Final = tuple(deployment.model_name for deployment in deployments)
|
|
model_list: Final = _chaos_model_list(deployments)
|
|
chaos_calls: Final = _chaos_calls(deployments, marker)
|
|
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:
|
|
|
|
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",
|
|
)
|
|
)
|
|
|
|
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 (
|
|
tuple(
|
|
sum(_json_contains_exact_string(observation["body"], call.prompt) for observation in control_upstream)
|
|
for call in chaos_calls
|
|
)
|
|
== (1,) * 30
|
|
), control_upstream
|
|
with wire_server(policy) as edge:
|
|
config: Final = _configuration(tmp_path, identity, edge.url, "output", model_list=model_list)
|
|
with owned_proxy(gateway, tmp_path, {}, config=config, workers=2) as candidate:
|
|
|
|
def run(call: ChaosCall) -> tuple[int, CallerResult, float]:
|
|
started: Final = time.monotonic()
|
|
result: Final = _call_client(
|
|
call.client_kind,
|
|
call.endpoint,
|
|
candidate,
|
|
call.model,
|
|
call.prompt,
|
|
call.stream,
|
|
call.call_id,
|
|
)
|
|
return call.index, result, time.monotonic() - started
|
|
|
|
with ThreadPoolExecutor(max_workers=30) as pool:
|
|
results: Final = tuple(pool.map(run, chaos_calls))
|
|
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_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),
|
|
seconds=30,
|
|
)
|
|
edge_calls: Final = edge.drain()
|
|
payloads: Final = tuple(JSON_OBJECT.validate_json(call.body) for call in edge_calls)
|
|
for call in chaos_calls:
|
|
payloads_for_call: Final = tuple(
|
|
payload for payload in payloads if _policy_call_id_matches(payload, call.call_id)
|
|
)
|
|
assert tuple(sorted(_direction(payload) for payload in payloads_for_call)) == tuple(
|
|
sorted(expected_directions)
|
|
), (
|
|
call,
|
|
payloads_for_call,
|
|
)
|
|
assert all(
|
|
payload["texts"]
|
|
== ([call.prompt] if _direction(payload) == "request" else [controls[call.index].text])
|
|
for payload in payloads_for_call
|
|
), payloads_for_call
|
|
response_ids: Final = tuple(result.response_id for _, result, _ in results)
|
|
assert len(set(response_ids)) == 30, response_ids
|
|
upstream: Final = _drain_upstream(gateway.upstream_url)
|
|
assert len(upstream) == 30, upstream
|
|
assert (
|
|
tuple(
|
|
sum(_json_contains_exact_string(observation["body"], call.prompt) for observation in upstream)
|
|
for call in chaos_calls
|
|
)
|
|
== (1,) * 30
|
|
), upstream
|
|
rows: Final = _spend_rows_for_calls(
|
|
models,
|
|
tuple(
|
|
(call.model, result.response_id, call.call_id)
|
|
for call, (_, result, _) in zip(chaos_calls, results)
|
|
),
|
|
)
|
|
for call, (_, result, _) in zip(chaos_calls, results):
|
|
matching_rows: Final = _spend_rows_matching_call(rows, call.model, call.call_id)
|
|
assert len(matching_rows) == 1, (call, result.response_id, matching_rows)
|
|
row: Final = matching_rows[0]
|
|
entries: Final = _guardrail_entries(row)
|
|
assert tuple(
|
|
(entry["guardrail_name"], entry["guardrail_mode"], entry["guardrail_status"])
|
|
for entry in entries
|
|
) == tuple((identity, "logging_only", "success") for _ in expected_directions), entries
|
|
|
|
|
|
@pytest.mark.timeout(300)
|
|
def test_K3_two_worker_sigkill_checks_post_kill_spend_rows(
|
|
gateway: Gateway,
|
|
tmp_path: Path,
|
|
record_property: Callable[[str, object], None],
|
|
) -> None:
|
|
identity: Final = f"logging-scope-k3-{uuid.uuid4().hex}"
|
|
marker: Final = uuid.uuid4().hex
|
|
scan_started: Final = threading.Event()
|
|
release_scans: Final = threading.Event()
|
|
|
|
def policy(_request: Request) -> Reply:
|
|
scan_started.set()
|
|
assert release_scans.wait(timeout=60), identity
|
|
return Reply(body=b'{"action":"NONE"}')
|
|
|
|
with gateway.scenario() as scenario:
|
|
deployments: Final = _chaos_models(scenario, marker)
|
|
models: Final = tuple(deployment.model_name for deployment in deployments)
|
|
model_list: Final = _chaos_model_list(deployments)
|
|
calls: Final = _chaos_calls(deployments, marker)
|
|
expected_directions: Final = _directions_for_audit_leg(("request", "response"), "output")
|
|
expected_entries: Final = tuple((identity, "logging_only", "success") for _ in expected_directions)
|
|
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",
|
|
)
|
|
for call in calls
|
|
)
|
|
assert len(_drain_upstream(gateway.upstream_url)) == 30
|
|
with wire_server(policy) as edge:
|
|
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
|
|
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:
|
|
result: Final = _call_client(
|
|
call.client_kind,
|
|
call.endpoint,
|
|
candidate,
|
|
call.model,
|
|
call.prompt,
|
|
call.stream,
|
|
call.call_id,
|
|
)
|
|
return call.index, result, None
|
|
except (OpenAIAPIConnectionError, AnthropicAPIConnectionError, httpx.RemoteProtocolError) as error:
|
|
return call.index, None, str(error)
|
|
|
|
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)
|
|
kill()
|
|
finally:
|
|
release_scans.set()
|
|
outcomes: Final = tuple(future.result(timeout=90) for future in futures)
|
|
assert owned.process.poll() is None, "Proxy supervisor exited after a worker was killed"
|
|
successful: Final = tuple(
|
|
(calls[index], result) for index, result, error in outcomes if result is not None and error is None
|
|
)
|
|
assert successful, outcomes
|
|
assert all(
|
|
result.status == controls[call.index].status and result.text == controls[call.index].text
|
|
for call, result in successful
|
|
), successful
|
|
pre_kill_expected: Final = tuple(
|
|
(call.model, result.response_id, call.call_id) for call, result in successful
|
|
)
|
|
pre_kill_rows: Final = _spend_rows_for_calls(
|
|
models,
|
|
pre_kill_expected,
|
|
tolerate_missing=True,
|
|
)
|
|
pre_kill_rows_by_call: Final = tuple(
|
|
(call, _spend_rows_matching_call(pre_kill_rows, call.model, call.call_id)) for call in calls
|
|
)
|
|
assert all(len(rows) <= 1 for _, rows in pre_kill_rows_by_call), pre_kill_rows_by_call
|
|
pre_kill_missing_rows: Final = sum(not rows for _, rows in pre_kill_rows_by_call)
|
|
record_property("k3_pre_kill_missing_spend_rows", pre_kill_missing_rows)
|
|
for call, matching_rows in pre_kill_rows_by_call:
|
|
if not matching_rows:
|
|
continue
|
|
entries: Final = _guardrail_entries(matching_rows[0])
|
|
assert tuple(
|
|
sorted(
|
|
(
|
|
entry["guardrail_name"],
|
|
entry["guardrail_mode"],
|
|
entry["guardrail_status"],
|
|
)
|
|
for entry in entries
|
|
)
|
|
) == tuple(sorted(expected_entries)), (call, entries)
|
|
|
|
post_kill_templates: Final = calls[:6]
|
|
post_kill_calls: Final = tuple(
|
|
ChaosCall(
|
|
index=call.index,
|
|
endpoint=call.endpoint,
|
|
client_kind=call.client_kind,
|
|
model=call.model,
|
|
stream=call.stream,
|
|
prompt=f"synthetic K post-kill burst {marker}-{call.index}",
|
|
call_id=f"{marker}-k-post-kill-{call.index}",
|
|
)
|
|
for call in post_kill_templates
|
|
)
|
|
|
|
def run_post_kill(call: ChaosCall) -> CallerResult:
|
|
return _call_client(
|
|
call.client_kind,
|
|
call.endpoint,
|
|
candidate,
|
|
call.model,
|
|
call.prompt,
|
|
call.stream,
|
|
call.call_id,
|
|
)
|
|
|
|
with ThreadPoolExecutor(max_workers=len(post_kill_calls)) as pool:
|
|
post_kill_futures: Final = tuple(pool.submit(run_post_kill, call) for call in post_kill_calls)
|
|
post_kill_results: Final = tuple(future.result(timeout=90) for future in post_kill_futures)
|
|
assert all(
|
|
result.status == controls[call.index].status and result.text == controls[call.index].text
|
|
for call, result in zip(post_kill_calls, post_kill_results)
|
|
), post_kill_results
|
|
served: Final = successful + tuple(zip(post_kill_calls, post_kill_results))
|
|
response_ids: Final = tuple(result.response_id for _, result in served)
|
|
assert len(set(response_ids)) == len(response_ids), response_ids
|
|
served_calls: Final = tuple(call for call, _ in served)
|
|
requested_call_ids: Final = frozenset(call.call_id for call in calls + post_kill_calls)
|
|
upstream: Final = _drain_upstream(gateway.upstream_url)
|
|
assert all(
|
|
sum(_json_contains_exact_string(observation["body"], call.prompt) for observation in upstream) == 1
|
|
for call in served_calls
|
|
), upstream
|
|
|
|
def accumulate_policy_payloads(
|
|
collected: tuple[dict[str, JsonValue], ...], _: None
|
|
) -> tuple[dict[str, JsonValue], ...]:
|
|
return collected + tuple(JSON_OBJECT.validate_json(call.body) for call in edge.drain())
|
|
|
|
payload_batches: Final = accumulate(repeat(None), accumulate_policy_payloads, initial=())
|
|
|
|
def has_expected_post_kill_scans(collected: tuple[dict[str, JsonValue], ...], call: ChaosCall) -> bool:
|
|
payloads_for_call: Final = tuple(
|
|
payload for payload in collected if _policy_call_id_matches(payload, call.call_id)
|
|
)
|
|
return all(
|
|
sum(_direction(payload) == direction for payload in payloads_for_call)
|
|
>= expected_directions.count(direction)
|
|
for direction in expected_directions
|
|
)
|
|
|
|
payloads: Final = eventually(
|
|
lambda: next(payload_batches),
|
|
lambda collected: all(has_expected_post_kill_scans(collected, call) for call in post_kill_calls),
|
|
seconds=30,
|
|
)
|
|
assert all(
|
|
any(_policy_call_id_matches(payload, call_id) for call_id in requested_call_ids)
|
|
and _direction(payload) in expected_directions
|
|
for payload in payloads
|
|
), payloads
|
|
for call, result in zip(post_kill_calls, post_kill_results):
|
|
payloads_for_call: Final = tuple(
|
|
payload for payload in payloads if _policy_call_id_matches(payload, call.call_id)
|
|
)
|
|
assert tuple(sorted(_direction(payload) for payload in payloads_for_call)) == tuple(
|
|
sorted(expected_directions)
|
|
), (call, payloads_for_call)
|
|
assert all(
|
|
payload["texts"]
|
|
== ([call.prompt] if _direction(payload) == "request" else [controls[call.index].text])
|
|
for payload in payloads_for_call
|
|
), payloads_for_call
|
|
|
|
post_kill_expected: Final = tuple(
|
|
(call.model, result.response_id, call.call_id)
|
|
for call, result in zip(post_kill_calls, post_kill_results)
|
|
)
|
|
rows: Final = _spend_rows_for_calls(
|
|
models,
|
|
post_kill_expected,
|
|
tolerate_missing=True,
|
|
)
|
|
all_candidate_calls: Final = calls + post_kill_calls
|
|
rows_by_call: Final = tuple(
|
|
(call, _spend_rows_matching_call(rows, call.model, call.call_id)) for call in all_candidate_calls
|
|
)
|
|
assert all(len(matching_rows) <= 1 for _, matching_rows in rows_by_call), rows_by_call
|
|
for call, matching_rows in rows_by_call:
|
|
if not matching_rows:
|
|
continue
|
|
entries: Final = _guardrail_entries(matching_rows[0])
|
|
assert tuple(
|
|
sorted(
|
|
(
|
|
entry["guardrail_name"],
|
|
entry["guardrail_mode"],
|
|
entry["guardrail_status"],
|
|
)
|
|
for entry in entries
|
|
)
|
|
) == tuple(sorted(expected_entries)), (call, entries)
|
|
for call, result in successful:
|
|
matching_rows: Final = _spend_rows_matching_call(rows, call.model, call.call_id)
|
|
if matching_rows:
|
|
_assert_response_id(
|
|
call.endpoint,
|
|
str(matching_rows[0]["request_id"]),
|
|
result.response_id,
|
|
marker if call.endpoint == "responses" else call.call_id,
|
|
)
|
|
for call, result in zip(post_kill_calls, post_kill_results):
|
|
matching_rows: Final = _spend_rows_matching_call(rows, call.model, call.call_id)
|
|
assert len(matching_rows) == 1, (result.response_id, matching_rows)
|
|
_assert_response_id(
|
|
call.endpoint,
|
|
str(matching_rows[0]["request_id"]),
|
|
result.response_id,
|
|
marker if call.endpoint == "responses" else call.call_id,
|
|
)
|
|
|
|
|
|
@pytest.mark.timeout(180)
|
|
def test_K4_proxy_restart_after_fifteen_responses_records_lost_ids(
|
|
gateway: Gateway,
|
|
tmp_path: Path,
|
|
record_property: Callable[[str, object], None],
|
|
) -> None:
|
|
identity: Final = f"logging-scope-k4-{uuid.uuid4().hex}"
|
|
marker: Final = uuid.uuid4().hex
|
|
|
|
def policy(_request: Request) -> Reply:
|
|
return Reply(body=b'{"action":"NONE"}')
|
|
|
|
with gateway.scenario() as scenario:
|
|
deployments: Final = _chaos_models(scenario, marker)
|
|
models: Final = tuple(deployment.model_name for deployment in deployments)
|
|
model_list: Final = _chaos_model_list(deployments)
|
|
calls: Final = _chaos_calls(deployments, marker)
|
|
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",
|
|
)
|
|
for call in calls
|
|
)
|
|
control_upstream: Final = _drain_upstream(gateway.upstream_url)
|
|
assert len(control_upstream) == 30, control_upstream
|
|
assert (
|
|
tuple(
|
|
sum(_json_contains_exact_string(observation["body"], call.prompt) for observation in control_upstream)
|
|
for call in calls
|
|
)
|
|
== (1,) * 30
|
|
), control_upstream
|
|
with wire_server(policy) as edge:
|
|
config: Final = _configuration(tmp_path, identity, edge.url, "output", model_list=model_list)
|
|
restart_gate: Final = threading.Event()
|
|
second_wave_ready: Final = threading.Event()
|
|
second_wave_barrier: Final = threading.Barrier(15, action=second_wave_ready.set)
|
|
restarted_gateways: Final[SimpleQueue[Gateway]] = SimpleQueue()
|
|
|
|
def gateway_for_call(call: ChaosCall, first_gateway: Gateway) -> Gateway:
|
|
if call.index < 15:
|
|
return first_gateway
|
|
second_wave_barrier.wait(timeout=90)
|
|
assert restart_gate.wait(timeout=90)
|
|
return restarted_gateways.get()
|
|
|
|
def run(call: ChaosCall, first_gateway: Gateway) -> tuple[int, CallerResult]:
|
|
candidate: Final = gateway_for_call(call, first_gateway)
|
|
result: Final = _call_client(
|
|
call.client_kind,
|
|
call.endpoint,
|
|
candidate,
|
|
call.model,
|
|
call.prompt,
|
|
call.stream,
|
|
call.call_id,
|
|
)
|
|
return call.index, result
|
|
|
|
with ThreadPoolExecutor(max_workers=30) as pool:
|
|
with owned_proxy_process(gateway, tmp_path, {}, config=config, workers=2) as first_proxy:
|
|
first_proxy_port: Final = first_proxy.gateway.client.base_url.port
|
|
assert first_proxy_port is not None
|
|
futures: Final = tuple(pool.submit(run, call, first_proxy.gateway) for call in calls)
|
|
first_results: Final = tuple(futures[index].result(timeout=90)[1] for index in range(15))
|
|
assert all(
|
|
result.status == controls[call.index].status and result.text == controls[call.index].text
|
|
for call, result in zip(calls[:15], first_results)
|
|
), first_results
|
|
assert eventually(lambda: second_wave_ready.is_set(), bool, seconds=30)
|
|
eventually(
|
|
lambda: edge.received.qsize(),
|
|
lambda count: count == 15 * len(expected_directions),
|
|
seconds=30,
|
|
)
|
|
first_wave_expected: Final = tuple(
|
|
(call.model, result.response_id, call.call_id)
|
|
for call, result in zip(calls[:15], first_results)
|
|
)
|
|
first_wave_rows: Final = _spend_rows_for_calls(
|
|
models,
|
|
first_wave_expected,
|
|
tolerate_missing=True,
|
|
)
|
|
assert all(
|
|
len(_spend_rows_matching_call(first_wave_rows, model, call_id)) <= 1
|
|
for model, _, call_id in first_wave_expected
|
|
), first_wave_rows
|
|
first_wave_present_response_ids: Final = frozenset(
|
|
response_id
|
|
for model, response_id, call_id in first_wave_expected
|
|
if len(_spend_rows_matching_call(first_wave_rows, model, call_id)) == 1
|
|
)
|
|
first_wave_lost_response_ids: Final = (
|
|
frozenset(response_id for _, response_id, _ in first_wave_expected)
|
|
- first_wave_present_response_ids
|
|
)
|
|
record_property(
|
|
"K4_PRE_RESTART_LOST_RESPONSE_IDS",
|
|
tuple(sorted(first_wave_lost_response_ids)),
|
|
)
|
|
record_property(
|
|
f"K4_PRE_RESTART_LOST_ROW_COUNT_{'base' if _is_base_audit_leg() else 'head'}",
|
|
len(first_wave_lost_response_ids),
|
|
)
|
|
assert first_proxy.process.poll() is not None, first_proxy.process.pid
|
|
eventually(
|
|
lambda: tuple(
|
|
connection
|
|
for connection in psutil.net_connections(kind="tcp")
|
|
if connection.status == psutil.CONN_LISTEN and connection.laddr.port == first_proxy_port
|
|
),
|
|
lambda listeners: not listeners,
|
|
seconds=30,
|
|
)
|
|
with owned_proxy_process(gateway, tmp_path, {}, config=config, workers=2) as restarted_proxy:
|
|
for _ in range(15):
|
|
restarted_gateways.put(restarted_proxy.gateway)
|
|
restart_gate.set()
|
|
second_results: Final = tuple(futures[index].result(timeout=90)[1] for index in range(15, 30))
|
|
assert all(
|
|
result.status == controls[call.index].status and result.text == controls[call.index].text
|
|
for call, result in zip(calls[15:], second_results)
|
|
), second_results
|
|
eventually(
|
|
lambda: edge.received.qsize(),
|
|
lambda count: count == 30 * len(expected_directions),
|
|
seconds=30,
|
|
)
|
|
results: Final = first_results + second_results
|
|
expected_response_ids: Final = frozenset(result.response_id for result in results)
|
|
assert len(expected_response_ids) == 30, results
|
|
upstream: Final = _drain_upstream(gateway.upstream_url)
|
|
assert len(upstream) == 30, upstream
|
|
assert (
|
|
tuple(
|
|
sum(_json_contains_exact_string(observation["body"], call.prompt) for observation in upstream)
|
|
for call in calls
|
|
)
|
|
== (1,) * 30
|
|
), upstream
|
|
edge_payloads: Final = tuple(JSON_OBJECT.validate_json(call.body) for call in edge.drain())
|
|
assert len(edge_payloads) == 30 * len(expected_directions), edge_payloads
|
|
for call, result in zip(calls, results):
|
|
payloads_for_call: Final = tuple(
|
|
payload for payload in edge_payloads if _policy_call_id_matches(payload, call.call_id)
|
|
)
|
|
assert tuple(sorted(_direction(payload) for payload in payloads_for_call)) == tuple(
|
|
sorted(expected_directions)
|
|
), (
|
|
call,
|
|
payloads_for_call,
|
|
)
|
|
assert all(
|
|
payload["texts"] == ([call.prompt] if _direction(payload) == "request" else [result.text])
|
|
for payload in payloads_for_call
|
|
), payloads_for_call
|
|
post_restart_expected: Final = tuple(
|
|
(call.model, result.response_id, call.call_id) for call, result in zip(calls[15:], second_results)
|
|
)
|
|
rows: Final = _spend_rows_for_calls(
|
|
models,
|
|
post_restart_expected,
|
|
tolerate_missing=True,
|
|
)
|
|
record_property("K4_RESPONSE_IDS", tuple(sorted(expected_response_ids)))
|
|
for call, result in zip(calls, results):
|
|
matching_rows: Final = _spend_rows_matching_call(rows, call.model, call.call_id)
|
|
assert len(matching_rows) <= 1, (result.response_id, matching_rows)
|
|
if call.index >= 15:
|
|
assert len(matching_rows) == 1, (call.call_id, result.response_id, matching_rows)
|
|
if matching_rows:
|
|
entries: Final = _guardrail_entries(matching_rows[0])
|
|
assert tuple(
|
|
(entry["guardrail_name"], entry["guardrail_mode"], entry["guardrail_status"])
|
|
for entry in entries
|
|
) == tuple((identity, "logging_only", "success") for _ in expected_directions), (
|
|
call.call_id,
|
|
entries,
|
|
)
|