305 lines
9.9 KiB
Python
305 lines
9.9 KiB
Python
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import base64
|
|
from io import BytesIO
|
|
|
|
import pytest
|
|
from PIL import Image
|
|
|
|
from cmvr_edge_ai.contracts import BoundingBox, Detection, DetectionResult, ImageFrame
|
|
from cmvr_edge_ai.contracts.inference import (
|
|
ImageInput,
|
|
InferenceRequest,
|
|
InferenceResponse,
|
|
InferenceStatus,
|
|
)
|
|
from cmvr_edge_ai.core import ComponentContext
|
|
from cmvr_edge_ai.core.envelope import Envelope
|
|
from cmvr_edge_ai.server import InvocationBroker
|
|
from cmvr_edge_ai.server.components import (
|
|
DetectionInferenceResponseOperator,
|
|
InferenceImageDecodeOperator,
|
|
InvocationRequestSource,
|
|
InvocationResponseSink,
|
|
)
|
|
|
|
|
|
def _jpeg_bytes(color: tuple[int, int, int] = (255, 0, 0)) -> bytes:
|
|
image = Image.new("RGB", (4, 3), color)
|
|
output = BytesIO()
|
|
image.save(output, format="JPEG", quality=95)
|
|
return output.getvalue()
|
|
|
|
|
|
def _request(
|
|
image: bytes,
|
|
*,
|
|
request_id: str = "request-1",
|
|
category: str = "detect.mobile_phone",
|
|
roles: tuple[str, ...] = (),
|
|
) -> InferenceRequest:
|
|
return InferenceRequest(
|
|
request_id=request_id,
|
|
category=category,
|
|
inputs=(
|
|
ImageInput(
|
|
media_type="image/jpeg",
|
|
data=base64.b64encode(image).decode("ascii"),
|
|
),
|
|
),
|
|
requested_artifact_roles=roles,
|
|
)
|
|
|
|
|
|
def _context(broker: InvocationBroker[object, object]) -> ComponentContext:
|
|
return ComponentContext(
|
|
pipeline_id="remote",
|
|
node_id="boundary",
|
|
shutdown_event=asyncio.Event(),
|
|
metadata={"invocation_broker": broker},
|
|
)
|
|
|
|
|
|
def test_request_source_consumes_only_its_registered_category() -> None:
|
|
async def exercise() -> None:
|
|
broker: InvocationBroker[InferenceRequest, InferenceResponse] = (
|
|
InvocationBroker()
|
|
)
|
|
broker.register_route("detect.mobile_phone", capacity=1)
|
|
source = InvocationRequestSource(
|
|
"requests", {"category": "detect.mobile_phone"}
|
|
)
|
|
await source.setup(_context(broker))
|
|
request = _request(_jpeg_bytes())
|
|
submitted = await broker.submit(
|
|
request.category, request.request_id or "", request
|
|
)
|
|
|
|
stream = source.messages()
|
|
emission = await anext(stream)
|
|
assert emission.port == "requests"
|
|
assert emission.envelope.payload is request
|
|
assert emission.envelope.trace_id == "request-1"
|
|
assert emission.envelope.attributes["invocation_category"] == request.category
|
|
assert (
|
|
emission.envelope.attributes["invocation_token"]
|
|
== submitted.invocation_token
|
|
)
|
|
await broker.cancel("request-1")
|
|
await broker.close()
|
|
await stream.aclose()
|
|
|
|
asyncio.run(exercise())
|
|
|
|
|
|
def test_image_decode_operator_produces_a_decoded_bgr_frame() -> None:
|
|
async def exercise() -> ImageFrame:
|
|
request = _request(_jpeg_bytes())
|
|
operator = InferenceImageDecodeOperator("decode", {})
|
|
output = await operator.process(
|
|
Envelope(
|
|
request,
|
|
schema_name="InferenceRequest",
|
|
attributes={
|
|
"invocation_request_id": "request-1",
|
|
"invocation_category": request.category,
|
|
"invocation_submitted_at_ns": 1,
|
|
},
|
|
)
|
|
)
|
|
return output.envelope.payload
|
|
|
|
frame = asyncio.run(exercise())
|
|
assert frame.width == 4
|
|
assert frame.height == 3
|
|
assert frame.pixel_format == "BGR8"
|
|
assert frame.codec == "none"
|
|
assert len(frame.data) == 4 * 3 * 3
|
|
# JPEG is lossy, but red remains dominant in the BGR channel order.
|
|
assert frame.data[2] > frame.data[1]
|
|
assert frame.data[2] > frame.data[0]
|
|
|
|
|
|
def test_image_decode_rejects_media_type_mismatch() -> None:
|
|
request = InferenceRequest(
|
|
request_id="request-1",
|
|
category="detect.mobile_phone",
|
|
inputs=(
|
|
ImageInput(
|
|
media_type="image/png",
|
|
data=base64.b64encode(_jpeg_bytes()).decode("ascii"),
|
|
),
|
|
),
|
|
)
|
|
|
|
with pytest.raises(ValueError, match="does not match decoded image type"):
|
|
asyncio.run(
|
|
InferenceImageDecodeOperator("decode", {}).process(
|
|
Envelope(request, schema_name="InferenceRequest")
|
|
)
|
|
)
|
|
|
|
|
|
def test_detection_response_and_sink_complete_the_original_waiter() -> None:
|
|
async def exercise() -> InferenceResponse:
|
|
broker: InvocationBroker[InferenceRequest, InferenceResponse] = (
|
|
InvocationBroker()
|
|
)
|
|
broker.register_route("detect.mobile_phone", capacity=1)
|
|
request = _request(_jpeg_bytes(), roles=("annotated",))
|
|
submitted = await broker.submit(
|
|
request.category, request.request_id or "", request
|
|
)
|
|
await broker.receive(request.category)
|
|
|
|
frame = ImageFrame(
|
|
data=bytes((0, 0, 255)) * 12,
|
|
width=4,
|
|
height=3,
|
|
pixel_format="BGR8",
|
|
)
|
|
result = DetectionResult(
|
|
detections=(
|
|
Detection(
|
|
label="mobile_phone",
|
|
confidence=0.9,
|
|
box=BoundingBox(0, 0, 3, 2),
|
|
),
|
|
),
|
|
model_id="yolov8n-mobile-phone@1",
|
|
model_name="YOLOv8n Mobile Phone",
|
|
inference_ms=4.5,
|
|
source_frame=frame,
|
|
)
|
|
response_operator = DetectionInferenceResponseOperator(
|
|
"response",
|
|
{
|
|
"category": "detect.mobile_phone",
|
|
"backend": "ultralytics-yolo",
|
|
},
|
|
)
|
|
response_emission = await response_operator.process(
|
|
Envelope(
|
|
result,
|
|
schema_name="DetectionResult",
|
|
trace_id="request-1",
|
|
attributes={
|
|
"invocation_request_id": "request-1",
|
|
"invocation_token": submitted.invocation_token,
|
|
"invocation_category": "detect.mobile_phone",
|
|
"invocation_submitted_at_ns": 1,
|
|
"requested_artifact_roles": ("annotated",),
|
|
},
|
|
)
|
|
)
|
|
sink = InvocationResponseSink("responses", {})
|
|
await sink.setup(_context(broker))
|
|
await sink.consume(response_emission.envelope)
|
|
response = await broker.wait("request-1", timeout_s=0.1)
|
|
await broker.close()
|
|
return response
|
|
|
|
response = asyncio.run(exercise())
|
|
assert response.status is InferenceStatus.SUCCEEDED
|
|
assert response.category == "detect.mobile_phone"
|
|
assert response.model is not None
|
|
assert response.model.model_id == "yolov8n-mobile-phone@1"
|
|
assert response.outputs[0].kind == "detections"
|
|
assert response.artifacts[0].role == "annotated"
|
|
assert base64.b64decode(response.artifacts[0].data).startswith(b"\xff\xd8")
|
|
|
|
|
|
def test_response_sink_rejects_response_for_a_different_request_id() -> None:
|
|
async def exercise() -> None:
|
|
broker: InvocationBroker[InferenceRequest, InferenceResponse] = (
|
|
InvocationBroker()
|
|
)
|
|
broker.register_route("detect.mobile_phone", capacity=1)
|
|
request = _request(_jpeg_bytes())
|
|
submitted = await broker.submit(
|
|
request.category, request.request_id or "", request
|
|
)
|
|
await broker.receive(request.category)
|
|
response = InferenceResponse(
|
|
request_id="wrong-request",
|
|
category=request.category,
|
|
status=InferenceStatus.NO_RESULT,
|
|
)
|
|
envelope = Envelope(
|
|
response,
|
|
schema_name="InferenceResponse",
|
|
attributes={
|
|
"invocation_request_id": submitted.request_id,
|
|
"invocation_token": submitted.invocation_token,
|
|
"invocation_category": submitted.category,
|
|
},
|
|
)
|
|
sink = InvocationResponseSink("responses", {})
|
|
await sink.setup(_context(broker))
|
|
|
|
with pytest.raises(ValueError, match="request_id does not match"):
|
|
await sink.consume(envelope)
|
|
|
|
assert await broker.cancel(submitted.request_id) is True
|
|
await broker.close()
|
|
|
|
asyncio.run(exercise())
|
|
|
|
|
|
def test_detection_response_uses_registry_backend_propagated_by_detector() -> None:
|
|
result = DetectionResult(
|
|
detections=(),
|
|
model_id="yolov8n-mobile-phone@2",
|
|
model_name="YOLOv8n Mobile Phone ONNX",
|
|
inference_ms=1.0,
|
|
)
|
|
operator = DetectionInferenceResponseOperator(
|
|
"response",
|
|
{"category": "detect.mobile_phone"},
|
|
)
|
|
envelope = Envelope(
|
|
result,
|
|
schema_name="DetectionResult",
|
|
trace_id="request-onnx",
|
|
attributes={
|
|
"invocation_request_id": "request-onnx",
|
|
"invocation_category": "detect.mobile_phone",
|
|
"invocation_submitted_at_ns": 1,
|
|
"detection_model_backend": "onnxruntime-yolov8",
|
|
},
|
|
)
|
|
|
|
emission = asyncio.run(operator.process(envelope))
|
|
|
|
assert emission.envelope.payload.model is not None
|
|
assert emission.envelope.payload.model.backend == "onnxruntime-yolov8"
|
|
|
|
|
|
def test_detection_response_rejects_stale_backend_override() -> None:
|
|
result = DetectionResult(
|
|
detections=(),
|
|
model_id="yolov8n-mobile-phone@2",
|
|
inference_ms=1.0,
|
|
)
|
|
operator = DetectionInferenceResponseOperator(
|
|
"response",
|
|
{
|
|
"category": "detect.mobile_phone",
|
|
"backend": "ultralytics-yolo",
|
|
},
|
|
)
|
|
envelope = Envelope(
|
|
result,
|
|
schema_name="DetectionResult",
|
|
attributes={
|
|
"invocation_request_id": "request-onnx",
|
|
"invocation_category": "detect.mobile_phone",
|
|
"invocation_submitted_at_ns": 1,
|
|
"detection_model_backend": "onnxruntime-yolov8",
|
|
},
|
|
)
|
|
|
|
with pytest.raises(ValueError, match="does not match"):
|
|
asyncio.run(operator.process(envelope))
|