564 lines
20 KiB
Python
564 lines
20 KiB
Python
from __future__ import annotations
|
|
|
|
import asyncio
|
|
|
|
import pytest
|
|
|
|
from cmvr_edge_ai.server.invocation import (
|
|
CompletionDisposition,
|
|
DuplicateInvocationError,
|
|
DuplicateInvocationRouteError,
|
|
DuplicateInvocationWaitError,
|
|
InvocationBroker,
|
|
InvocationBrokerClosedError,
|
|
InvocationCancelledError,
|
|
InvocationNotFoundError,
|
|
InvocationQueueFullError,
|
|
InvocationTimeoutError,
|
|
UnknownInvocationRouteError,
|
|
)
|
|
|
|
|
|
def _broker() -> InvocationBroker[dict[str, object], dict[str, object]]:
|
|
broker: InvocationBroker[dict[str, object], dict[str, object]] = (
|
|
InvocationBroker()
|
|
)
|
|
broker.register_route("detect.mobile_phone", capacity=2)
|
|
broker.register_route("detect.analog_gauge", capacity=1)
|
|
return broker
|
|
|
|
|
|
def test_routes_are_registered_explicitly_and_validated() -> None:
|
|
broker: InvocationBroker[object, object] = InvocationBroker()
|
|
broker.register_route(" detect.mobile_phone ", capacity=2)
|
|
|
|
assert broker.categories == ("detect.mobile_phone",)
|
|
assert broker.route_stats("detect.mobile_phone").capacity == 2
|
|
with pytest.raises(DuplicateInvocationRouteError, match="already registered"):
|
|
broker.register_route("detect.mobile_phone", capacity=1)
|
|
with pytest.raises(ValueError, match="positive integer"):
|
|
broker.register_route("bad.zero", capacity=0)
|
|
with pytest.raises(ValueError, match="positive integer"):
|
|
broker.register_route("bad.bool", capacity=True)
|
|
|
|
|
|
def test_submit_receive_complete_and_wait_round_trip() -> None:
|
|
async def exercise(): # type: ignore[no-untyped-def]
|
|
broker = _broker()
|
|
submitted = await broker.submit(
|
|
"detect.mobile_phone",
|
|
"request-1",
|
|
{"image": b"jpeg"},
|
|
)
|
|
|
|
received = await broker.receive("detect.mobile_phone")
|
|
assert received is submitted
|
|
assert received.category == "detect.mobile_phone"
|
|
assert received.request_id == "request-1"
|
|
assert received.payload == {"image": b"jpeg"}
|
|
assert received.submitted_at_ns > 0
|
|
|
|
waiter = asyncio.create_task(broker.wait("request-1", timeout_s=1))
|
|
await asyncio.sleep(0)
|
|
disposition = await broker.complete(
|
|
"request-1",
|
|
{"count": 1},
|
|
invocation_token=submitted.invocation_token,
|
|
)
|
|
response = await waiter
|
|
stats = broker.route_stats("detect.mobile_phone")
|
|
await broker.close()
|
|
return disposition, response, stats
|
|
|
|
disposition, response, stats = asyncio.run(exercise())
|
|
|
|
assert disposition is CompletionDisposition.ACCEPTED
|
|
assert response == {"count": 1}
|
|
assert stats.queued == 0
|
|
assert stats.pending == 0
|
|
|
|
|
|
def test_category_consumers_never_compete_for_another_route_queue() -> None:
|
|
async def exercise(): # type: ignore[no-untyped-def]
|
|
broker = _broker()
|
|
gauge_receiver = asyncio.create_task(
|
|
broker.receive("detect.analog_gauge")
|
|
)
|
|
phone_receiver = asyncio.create_task(
|
|
broker.receive("detect.mobile_phone")
|
|
)
|
|
await asyncio.sleep(0)
|
|
|
|
await broker.submit("detect.mobile_phone", "phone-1", {"frame": 1})
|
|
phone_request = await asyncio.wait_for(phone_receiver, timeout=1)
|
|
await asyncio.sleep(0)
|
|
assert gauge_receiver.done() is False
|
|
|
|
await broker.complete(
|
|
"phone-1",
|
|
{"boxes": []},
|
|
invocation_token=phone_request.invocation_token,
|
|
)
|
|
await broker.wait("phone-1", timeout_s=1)
|
|
await broker.close()
|
|
with pytest.raises(InvocationBrokerClosedError):
|
|
await gauge_receiver
|
|
return phone_request
|
|
|
|
request = asyncio.run(exercise())
|
|
|
|
assert request.category == "detect.mobile_phone"
|
|
assert request.request_id == "phone-1"
|
|
|
|
|
|
def test_route_capacity_bounds_all_outstanding_work_independently() -> None:
|
|
async def exercise(): # type: ignore[no-untyped-def]
|
|
broker = _broker()
|
|
await broker.submit("detect.analog_gauge", "gauge-1", {"frame": 1})
|
|
|
|
with pytest.raises(InvocationQueueFullError) as captured:
|
|
await broker.submit("detect.analog_gauge", "gauge-2", {"frame": 2})
|
|
assert captured.value.category == "detect.analog_gauge"
|
|
assert captured.value.capacity == 1
|
|
|
|
# A full gauge route does not consume the phone route's capacity.
|
|
await broker.submit("detect.mobile_phone", "phone-1", {"frame": 3})
|
|
assert broker.route_stats("detect.analog_gauge").pending == 1
|
|
assert broker.route_stats("detect.mobile_phone").pending == 1
|
|
|
|
assert await broker.cancel("gauge-1") is True
|
|
await broker.submit("detect.analog_gauge", "gauge-2", {"frame": 2})
|
|
next_gauge = await broker.receive("detect.analog_gauge")
|
|
await broker.cancel("gauge-2")
|
|
await broker.cancel("phone-1")
|
|
await broker.close()
|
|
return next_gauge
|
|
|
|
request = asyncio.run(exercise())
|
|
|
|
assert request.request_id == "gauge-2"
|
|
|
|
|
|
def test_request_ids_are_global_across_routes_and_recent_history() -> None:
|
|
async def exercise(): # type: ignore[no-untyped-def]
|
|
broker = _broker()
|
|
submitted = await broker.submit(
|
|
"detect.mobile_phone", "global-id", {"frame": 1}
|
|
)
|
|
with pytest.raises(DuplicateInvocationError, match="global-id"):
|
|
await broker.submit(
|
|
"detect.analog_gauge",
|
|
"global-id",
|
|
{"frame": 2},
|
|
)
|
|
|
|
await broker.complete(
|
|
"global-id",
|
|
{"boxes": []},
|
|
invocation_token=submitted.invocation_token,
|
|
)
|
|
assert await broker.wait("global-id") == {"boxes": []}
|
|
with pytest.raises(DuplicateInvocationError, match="global-id"):
|
|
await broker.submit(
|
|
"detect.analog_gauge",
|
|
"global-id",
|
|
{"frame": 3},
|
|
)
|
|
await broker.close()
|
|
|
|
asyncio.run(exercise())
|
|
|
|
|
|
def test_out_of_order_results_are_correlated_by_request_id() -> None:
|
|
async def exercise(): # type: ignore[no-untyped-def]
|
|
broker = _broker()
|
|
first = await broker.submit(
|
|
"detect.mobile_phone", "first", {"frame": 1}
|
|
)
|
|
second = await broker.submit(
|
|
"detect.mobile_phone", "second", {"frame": 2}
|
|
)
|
|
assert (await broker.receive("detect.mobile_phone")).request_id == "first"
|
|
assert (await broker.receive("detect.mobile_phone")).request_id == "second"
|
|
|
|
first_waiter = asyncio.create_task(broker.wait("first", timeout_s=1))
|
|
second_waiter = asyncio.create_task(broker.wait("second", timeout_s=1))
|
|
await asyncio.sleep(0)
|
|
await broker.complete(
|
|
"second",
|
|
{"owner": "second"},
|
|
invocation_token=second.invocation_token,
|
|
)
|
|
await broker.complete(
|
|
"first",
|
|
{"owner": "first"},
|
|
invocation_token=first.invocation_token,
|
|
)
|
|
results = await asyncio.gather(first_waiter, second_waiter)
|
|
await broker.close()
|
|
return results
|
|
|
|
assert asyncio.run(exercise()) == [
|
|
{"owner": "first"},
|
|
{"owner": "second"},
|
|
]
|
|
|
|
|
|
def test_duplicate_and_unknown_result_deliveries_are_non_fatal() -> None:
|
|
async def exercise(): # type: ignore[no-untyped-def]
|
|
broker = _broker()
|
|
submitted = await broker.submit(
|
|
"detect.mobile_phone", "request-1", {"frame": 1}
|
|
)
|
|
accepted = await broker.complete(
|
|
"request-1",
|
|
{"boxes": [1]},
|
|
invocation_token=submitted.invocation_token,
|
|
)
|
|
duplicate_before_wait = await broker.fail(
|
|
"request-1",
|
|
RuntimeError("duplicate"),
|
|
invocation_token=submitted.invocation_token,
|
|
)
|
|
response = await broker.wait("request-1")
|
|
duplicate_after_wait = await broker.complete(
|
|
"request-1",
|
|
{"boxes": [2]},
|
|
invocation_token=submitted.invocation_token,
|
|
)
|
|
unknown = await broker.complete(
|
|
"not-submitted",
|
|
{"boxes": []},
|
|
invocation_token="unknown-token",
|
|
)
|
|
await broker.close()
|
|
return (
|
|
accepted,
|
|
duplicate_before_wait,
|
|
duplicate_after_wait,
|
|
unknown,
|
|
response,
|
|
)
|
|
|
|
accepted, duplicate_before, duplicate_after, unknown, response = asyncio.run(
|
|
exercise()
|
|
)
|
|
|
|
assert accepted is CompletionDisposition.ACCEPTED
|
|
assert duplicate_before is CompletionDisposition.DUPLICATE
|
|
assert duplicate_after is CompletionDisposition.DUPLICATE
|
|
assert unknown is CompletionDisposition.UNKNOWN
|
|
assert response == {"boxes": [1]}
|
|
|
|
|
|
def test_pipeline_failure_is_delivered_to_the_correct_waiter() -> None:
|
|
async def exercise(): # type: ignore[no-untyped-def]
|
|
broker = _broker()
|
|
expected = ValueError("invalid gauge ellipse")
|
|
submitted = await broker.submit(
|
|
"detect.analog_gauge", "gauge-1", {"frame": 1}
|
|
)
|
|
waiter = asyncio.create_task(broker.wait("gauge-1", timeout_s=1))
|
|
await asyncio.sleep(0)
|
|
assert (
|
|
await broker.fail(
|
|
"gauge-1",
|
|
expected,
|
|
invocation_token=submitted.invocation_token,
|
|
)
|
|
is CompletionDisposition.ACCEPTED
|
|
)
|
|
with pytest.raises(ValueError, match="invalid gauge ellipse") as captured:
|
|
await waiter
|
|
assert captured.value is expected
|
|
duplicate = await broker.fail(
|
|
"gauge-1",
|
|
RuntimeError("again"),
|
|
invocation_token=submitted.invocation_token,
|
|
)
|
|
await broker.close()
|
|
return duplicate
|
|
|
|
assert asyncio.run(exercise()) is CompletionDisposition.DUPLICATE
|
|
|
|
|
|
def test_timeout_removes_queued_request_and_classifies_late_result() -> None:
|
|
async def exercise(): # type: ignore[no-untyped-def]
|
|
broker = _broker()
|
|
submitted = await broker.submit(
|
|
"detect.analog_gauge", "slow", {"frame": 1}
|
|
)
|
|
with pytest.raises(InvocationTimeoutError, match="slow"):
|
|
await broker.wait("slow", timeout_s=0.01)
|
|
|
|
assert broker.route_stats("detect.analog_gauge").queued == 0
|
|
assert broker.route_stats("detect.analog_gauge").pending == 0
|
|
late = await broker.complete(
|
|
"slow",
|
|
{"reading": 4.2},
|
|
invocation_token=submitted.invocation_token,
|
|
)
|
|
|
|
# The cancelled queue entry cannot poison the next request.
|
|
await broker.submit("detect.analog_gauge", "next", {"frame": 2})
|
|
next_request = await asyncio.wait_for(
|
|
broker.receive("detect.analog_gauge"),
|
|
timeout=1,
|
|
)
|
|
await broker.cancel("next")
|
|
await broker.close()
|
|
return late, next_request
|
|
|
|
late, next_request = asyncio.run(exercise())
|
|
|
|
assert late is CompletionDisposition.LATE
|
|
assert next_request.request_id == "next"
|
|
|
|
|
|
def test_stale_generation_cannot_complete_a_reused_request_id() -> None:
|
|
async def exercise(): # type: ignore[no-untyped-def]
|
|
broker: InvocationBroker[dict[str, object], dict[str, object]] = (
|
|
InvocationBroker(terminal_history_size=1)
|
|
)
|
|
broker.register_route("detect.mobile_phone", capacity=1)
|
|
|
|
old_request = await broker.submit(
|
|
"detect.mobile_phone", "reused-id", {"generation": 1}
|
|
)
|
|
assert await broker.receive("detect.mobile_phone") is old_request
|
|
with pytest.raises(InvocationTimeoutError):
|
|
await broker.wait("reused-id", timeout_s=0.001)
|
|
assert broker.route_stats("detect.mobile_phone").pending == 1
|
|
with pytest.raises(InvocationQueueFullError):
|
|
await broker.submit(
|
|
"detect.mobile_phone", "while-model-busy", {"generation": 9}
|
|
)
|
|
assert (
|
|
await broker.complete(
|
|
old_request.request_id,
|
|
{"generation": 1},
|
|
invocation_token=old_request.invocation_token,
|
|
)
|
|
is CompletionDisposition.LATE
|
|
)
|
|
assert broker.route_stats("detect.mobile_phone").pending == 0
|
|
|
|
eviction = await broker.submit(
|
|
"detect.mobile_phone", "evict-history", {"generation": 0}
|
|
)
|
|
assert await broker.receive("detect.mobile_phone") is eviction
|
|
await broker.complete(
|
|
eviction.request_id,
|
|
{"generation": 0},
|
|
invocation_token=eviction.invocation_token,
|
|
)
|
|
assert await broker.wait(eviction.request_id) == {"generation": 0}
|
|
|
|
new_request = await broker.submit(
|
|
"detect.mobile_phone", "reused-id", {"generation": 2}
|
|
)
|
|
assert await broker.receive("detect.mobile_phone") is new_request
|
|
stale = await broker.complete(
|
|
old_request.request_id,
|
|
{"generation": 1},
|
|
invocation_token=old_request.invocation_token,
|
|
)
|
|
assert stale is CompletionDisposition.STALE
|
|
|
|
accepted = await broker.complete(
|
|
new_request.request_id,
|
|
{"generation": 2},
|
|
invocation_token=new_request.invocation_token,
|
|
)
|
|
response = await broker.wait(new_request.request_id)
|
|
await broker.close()
|
|
return accepted, response
|
|
|
|
accepted, response = asyncio.run(exercise())
|
|
|
|
assert accepted is CompletionDisposition.ACCEPTED
|
|
assert response == {"generation": 2}
|
|
|
|
|
|
def test_consumed_timeout_remains_waitable_after_terminal_history_eviction() -> None:
|
|
async def exercise(): # type: ignore[no-untyped-def]
|
|
broker: InvocationBroker[dict[str, object], dict[str, object]] = (
|
|
InvocationBroker(terminal_history_size=1)
|
|
)
|
|
broker.register_route("detect.mobile_phone", capacity=1)
|
|
broker.register_route("detect.analog_gauge", capacity=1)
|
|
|
|
slow = await broker.submit(
|
|
"detect.mobile_phone", "slow", {"generation": 1}
|
|
)
|
|
assert await broker.receive("detect.mobile_phone") is slow
|
|
with pytest.raises(InvocationTimeoutError):
|
|
await broker.wait("slow", timeout_s=0.001)
|
|
|
|
# Churn unrelated terminal history while the consumed model request is
|
|
# deliberately retained to keep its route capacity reserved.
|
|
await broker.submit("detect.analog_gauge", "evict", {"generation": 0})
|
|
assert await broker.cancel("evict") is True
|
|
with pytest.raises(InvocationTimeoutError, match="slow"):
|
|
await broker.wait("slow")
|
|
|
|
assert (
|
|
await broker.complete(
|
|
"slow",
|
|
{"generation": 1},
|
|
invocation_token=slow.invocation_token,
|
|
)
|
|
is CompletionDisposition.LATE
|
|
)
|
|
assert broker.route_stats("detect.mobile_phone").pending == 0
|
|
await broker.close()
|
|
|
|
asyncio.run(exercise())
|
|
|
|
|
|
def test_explicit_cancel_wakes_waiter_and_rejects_late_result() -> None:
|
|
async def exercise(): # type: ignore[no-untyped-def]
|
|
broker = _broker()
|
|
submitted = await broker.submit(
|
|
"detect.mobile_phone", "cancel-me", {"frame": 1}
|
|
)
|
|
waiter = asyncio.create_task(broker.wait("cancel-me"))
|
|
await asyncio.sleep(0)
|
|
assert await broker.cancel("cancel-me") is True
|
|
assert await broker.cancel("cancel-me") is False
|
|
with pytest.raises(InvocationCancelledError, match="cancel-me"):
|
|
await waiter
|
|
late = await broker.fail(
|
|
"cancel-me",
|
|
RuntimeError("too late"),
|
|
invocation_token=submitted.invocation_token,
|
|
)
|
|
await broker.close()
|
|
return late
|
|
|
|
assert asyncio.run(exercise()) is CompletionDisposition.LATE
|
|
|
|
|
|
def test_cancelling_transport_wait_cancels_correlation_but_not_other_routes() -> None:
|
|
async def exercise(): # type: ignore[no-untyped-def]
|
|
broker = _broker()
|
|
disconnected_request = await broker.submit(
|
|
"detect.mobile_phone", "disconnected", {"frame": 1}
|
|
)
|
|
healthy_request = await broker.submit(
|
|
"detect.analog_gauge", "healthy", {"frame": 2}
|
|
)
|
|
assert (
|
|
await broker.receive("detect.mobile_phone")
|
|
is disconnected_request
|
|
)
|
|
disconnected = asyncio.create_task(broker.wait("disconnected"))
|
|
await asyncio.sleep(0)
|
|
disconnected.cancel()
|
|
with pytest.raises(asyncio.CancelledError):
|
|
await disconnected
|
|
|
|
assert (
|
|
await broker.complete(
|
|
"disconnected",
|
|
{"boxes": []},
|
|
invocation_token=disconnected_request.invocation_token,
|
|
)
|
|
is CompletionDisposition.LATE
|
|
)
|
|
assert (
|
|
await broker.complete(
|
|
"healthy",
|
|
{"reading": 2.0},
|
|
invocation_token=healthy_request.invocation_token,
|
|
)
|
|
is CompletionDisposition.ACCEPTED
|
|
)
|
|
healthy = await broker.wait("healthy", timeout_s=1)
|
|
await broker.close()
|
|
return healthy
|
|
|
|
assert asyncio.run(exercise()) == {"reading": 2.0}
|
|
|
|
|
|
def test_only_one_waiter_can_own_a_request() -> None:
|
|
async def exercise(): # type: ignore[no-untyped-def]
|
|
broker = _broker()
|
|
submitted = await broker.submit(
|
|
"detect.mobile_phone", "one-waiter", {"frame": 1}
|
|
)
|
|
first = asyncio.create_task(broker.wait("one-waiter"))
|
|
await asyncio.sleep(0)
|
|
with pytest.raises(DuplicateInvocationWaitError, match="already has"):
|
|
await broker.wait("one-waiter")
|
|
await broker.complete(
|
|
"one-waiter",
|
|
{"boxes": []},
|
|
invocation_token=submitted.invocation_token,
|
|
)
|
|
response = await first
|
|
await broker.close()
|
|
return response
|
|
|
|
assert asyncio.run(exercise()) == {"boxes": []}
|
|
|
|
|
|
def test_close_is_idempotent_and_wakes_receivers_and_waiters() -> None:
|
|
async def exercise(): # type: ignore[no-untyped-def]
|
|
broker = _broker()
|
|
receiver = asyncio.create_task(broker.receive("detect.analog_gauge"))
|
|
submitted = await broker.submit(
|
|
"detect.mobile_phone", "pending", {"frame": 1}
|
|
)
|
|
waiter = asyncio.create_task(broker.wait("pending"))
|
|
await asyncio.sleep(0)
|
|
|
|
await broker.close()
|
|
await broker.close()
|
|
with pytest.raises(InvocationBrokerClosedError):
|
|
await receiver
|
|
with pytest.raises(InvocationBrokerClosedError):
|
|
await waiter
|
|
assert broker.closed is True
|
|
assert broker.pending_count == 0
|
|
assert (
|
|
await broker.complete(
|
|
"pending",
|
|
{"boxes": []},
|
|
invocation_token=submitted.invocation_token,
|
|
)
|
|
is CompletionDisposition.LATE
|
|
)
|
|
with pytest.raises(InvocationBrokerClosedError):
|
|
await broker.submit("detect.mobile_phone", "new", {"frame": 2})
|
|
with pytest.raises(InvocationBrokerClosedError):
|
|
await broker.receive("detect.mobile_phone")
|
|
with pytest.raises(InvocationBrokerClosedError):
|
|
broker.register_route("talk.asr", capacity=1)
|
|
|
|
asyncio.run(exercise())
|
|
|
|
|
|
def test_unknown_routes_requests_and_invalid_timeout_fail_near_the_caller() -> None:
|
|
async def exercise(): # type: ignore[no-untyped-def]
|
|
broker = _broker()
|
|
with pytest.raises(UnknownInvocationRouteError, match="available"):
|
|
await broker.submit("talk.asr", "request-1", {})
|
|
with pytest.raises(UnknownInvocationRouteError, match="available"):
|
|
await broker.receive("talk.asr")
|
|
with pytest.raises(InvocationNotFoundError, match="missing"):
|
|
await broker.wait("missing")
|
|
submitted = await broker.submit("detect.mobile_phone", "request-2", {})
|
|
with pytest.raises(ValueError, match="finite positive"):
|
|
await broker.wait("request-2", timeout_s=float("inf"))
|
|
with pytest.raises(TypeError, match="must be an Exception"):
|
|
await broker.fail(
|
|
"request-2",
|
|
"not-an-error", # type: ignore[arg-type]
|
|
invocation_token=submitted.invocation_token,
|
|
)
|
|
await broker.cancel("request-2")
|
|
await broker.close()
|
|
|
|
asyncio.run(exercise())
|