492 lines
15 KiB
Python
492 lines
15 KiB
Python
|
from __future__ import annotations
|
||
|
|
|
||
|
|
import asyncio
|
||
|
|
import base64
|
||
|
|
import hashlib
|
||
|
|
import sys
|
||
|
|
from collections.abc import Mapping, Sequence
|
||
|
|
from pathlib import Path
|
||
|
|
from typing import Any
|
||
|
|
|
||
|
|
import pytest
|
||
|
|
from pydantic import ValidationError
|
||
|
|
|
||
|
|
from cmvr_edge_ai.contracts import (
|
||
|
|
ImageArtifact,
|
||
|
|
ImageInput,
|
||
|
|
InferenceRequest,
|
||
|
|
InferenceStatus,
|
||
|
|
JsonArtifact,
|
||
|
|
ScalarOutput,
|
||
|
|
)
|
||
|
|
from cmvr_edge_ai.core import Envelope
|
||
|
|
from cmvr_edge_ai.gauge import (
|
||
|
|
GAUGE_CATEGORY,
|
||
|
|
GAUGE_MODEL_ID,
|
||
|
|
AnalogGaugeReaderOperator,
|
||
|
|
)
|
||
|
|
from cmvr_edge_ai.workers import PROTOCOL_VERSION, FramedWorkerResponse
|
||
|
|
|
||
|
|
|
||
|
|
class FakeWorker:
|
||
|
|
def __init__(
|
||
|
|
self,
|
||
|
|
responses: Sequence[FramedWorkerResponse | Exception],
|
||
|
|
*,
|
||
|
|
ready: Mapping[str, Any] | None = None,
|
||
|
|
) -> None:
|
||
|
|
self.responses = list(responses)
|
||
|
|
self.ready = dict(
|
||
|
|
ready
|
||
|
|
or {
|
||
|
|
"protocol": PROTOCOL_VERSION,
|
||
|
|
"type": "ready",
|
||
|
|
"model_id": GAUGE_MODEL_ID,
|
||
|
|
}
|
||
|
|
)
|
||
|
|
self.calls: list[tuple[Mapping[str, Any], tuple[bytes, ...], float]] = []
|
||
|
|
self.is_alive = False
|
||
|
|
self.last_error = ""
|
||
|
|
self.stopped = False
|
||
|
|
|
||
|
|
async def start(self) -> Mapping[str, Any]:
|
||
|
|
self.is_alive = True
|
||
|
|
return self.ready
|
||
|
|
|
||
|
|
async def submit(
|
||
|
|
self,
|
||
|
|
header: Mapping[str, Any],
|
||
|
|
blobs: Sequence[bytes] = (),
|
||
|
|
*,
|
||
|
|
timeout_s: float,
|
||
|
|
) -> FramedWorkerResponse:
|
||
|
|
self.calls.append((dict(header), tuple(blobs), timeout_s))
|
||
|
|
outcome = self.responses.pop(0)
|
||
|
|
if isinstance(outcome, Exception):
|
||
|
|
raise outcome
|
||
|
|
return outcome
|
||
|
|
|
||
|
|
async def stop(self) -> None:
|
||
|
|
self.is_alive = False
|
||
|
|
self.stopped = True
|
||
|
|
|
||
|
|
|
||
|
|
def _params(tmp_path: Path) -> dict[str, object]:
|
||
|
|
model_dir = tmp_path / "models"
|
||
|
|
model_dir.mkdir(exist_ok=True)
|
||
|
|
for name in ("detection.pt", "keypoints.pt", "segmentation.pt"):
|
||
|
|
(model_dir / name).write_bytes(b"model")
|
||
|
|
return {
|
||
|
|
"python_executable": sys.executable,
|
||
|
|
"project_root": str(tmp_path),
|
||
|
|
"detection_model_path": "models/detection.pt",
|
||
|
|
"key_point_model_path": "models/keypoints.pt",
|
||
|
|
"segmentation_model_path": "models/segmentation.pt",
|
||
|
|
"startup_timeout_s": 12,
|
||
|
|
"request_timeout_s": 34,
|
||
|
|
"shutdown_timeout_s": 5,
|
||
|
|
"max_header_bytes": 4096,
|
||
|
|
"max_blob_bytes": 8192,
|
||
|
|
"max_pixels": 1_234_567,
|
||
|
|
}
|
||
|
|
|
||
|
|
|
||
|
|
def _request(
|
||
|
|
*,
|
||
|
|
request_id: str = "gauge-1",
|
||
|
|
roles: tuple[str, ...] = ("original", "annotated", "diagnostics"),
|
||
|
|
) -> InferenceRequest:
|
||
|
|
encoded = b"encoded-gauge-image"
|
||
|
|
return InferenceRequest(
|
||
|
|
request_id=request_id,
|
||
|
|
category=GAUGE_CATEGORY,
|
||
|
|
source_id="camera-front",
|
||
|
|
inputs=(
|
||
|
|
ImageInput(
|
||
|
|
name="image",
|
||
|
|
media_type="image/jpeg",
|
||
|
|
data=base64.b64encode(encoded).decode("ascii"),
|
||
|
|
width=640,
|
||
|
|
height=480,
|
||
|
|
sha256=hashlib.sha256(encoded).hexdigest(),
|
||
|
|
),
|
||
|
|
),
|
||
|
|
requested_artifact_roles=roles,
|
||
|
|
)
|
||
|
|
|
||
|
|
|
||
|
|
def _envelope(request: InferenceRequest) -> Envelope[InferenceRequest]:
|
||
|
|
return Envelope(
|
||
|
|
request,
|
||
|
|
schema_name="InferenceRequest",
|
||
|
|
trace_id=request.request_id or "fallback",
|
||
|
|
source_id=request.source_id or "",
|
||
|
|
attributes={
|
||
|
|
"invocation_request_id": request.request_id,
|
||
|
|
"invocation_category": request.category,
|
||
|
|
},
|
||
|
|
)
|
||
|
|
|
||
|
|
|
||
|
|
def _worker_response(
|
||
|
|
*,
|
||
|
|
request_id: str = "gauge-1",
|
||
|
|
status: str = "succeeded",
|
||
|
|
value: float | None = 12.5,
|
||
|
|
with_annotated: bool = True,
|
||
|
|
input_image: dict[str, int] | None = None,
|
||
|
|
) -> FramedWorkerResponse:
|
||
|
|
annotated = b"annotated-jpeg"
|
||
|
|
artifacts = (
|
||
|
|
[
|
||
|
|
{
|
||
|
|
"blob_index": 0,
|
||
|
|
"role": "annotated",
|
||
|
|
"media_type": "image/jpeg",
|
||
|
|
"width": 640,
|
||
|
|
"height": 480,
|
||
|
|
"sha256": hashlib.sha256(annotated).hexdigest(),
|
||
|
|
}
|
||
|
|
]
|
||
|
|
if with_annotated
|
||
|
|
else []
|
||
|
|
)
|
||
|
|
error = (
|
||
|
|
{
|
||
|
|
"code": "GAUGE_READER_FAILED",
|
||
|
|
"message": "gauge could not be read",
|
||
|
|
"retryable": False,
|
||
|
|
"stage": "ocr",
|
||
|
|
"details": {"reason": "no labels"},
|
||
|
|
}
|
||
|
|
if status == "failed"
|
||
|
|
else None
|
||
|
|
)
|
||
|
|
return FramedWorkerResponse(
|
||
|
|
header={
|
||
|
|
"protocol": PROTOCOL_VERSION,
|
||
|
|
"type": "response",
|
||
|
|
"request_id": request_id,
|
||
|
|
"category": GAUGE_CATEGORY,
|
||
|
|
"model_id": GAUGE_MODEL_ID,
|
||
|
|
"status": status,
|
||
|
|
"value": value,
|
||
|
|
"unit": "bar" if value is not None else None,
|
||
|
|
"inference_ms": 7.5,
|
||
|
|
"input_image": (
|
||
|
|
{"width": 640, "height": 480}
|
||
|
|
if input_image is None and status != "failed"
|
||
|
|
else input_image
|
||
|
|
),
|
||
|
|
"diagnostics": {"ocr_labels": 4, "fit": {"residual": 0.02}},
|
||
|
|
"warnings": [],
|
||
|
|
"error": error,
|
||
|
|
"artifacts": artifacts,
|
||
|
|
},
|
||
|
|
blobs=(annotated,) if with_annotated else (),
|
||
|
|
)
|
||
|
|
|
||
|
|
|
||
|
|
def test_operator_maps_scalar_and_requested_artifacts(tmp_path: Path) -> None:
|
||
|
|
async def exercise() -> None:
|
||
|
|
fake = FakeWorker((_worker_response(),))
|
||
|
|
constructed: dict[str, object] = {}
|
||
|
|
|
||
|
|
def factory(command, **kwargs): # type: ignore[no-untyped-def]
|
||
|
|
constructed["command"] = tuple(command)
|
||
|
|
constructed["kwargs"] = kwargs
|
||
|
|
return fake
|
||
|
|
|
||
|
|
operator = AnalogGaugeReaderOperator(
|
||
|
|
"reader",
|
||
|
|
_params(tmp_path),
|
||
|
|
worker_factory=factory,
|
||
|
|
)
|
||
|
|
await operator.start()
|
||
|
|
emission = await operator.process(_envelope(_request()))
|
||
|
|
response = emission.envelope.payload
|
||
|
|
|
||
|
|
assert response.status is InferenceStatus.SUCCEEDED
|
||
|
|
assert response.request_id == "gauge-1"
|
||
|
|
assert response.category == GAUGE_CATEGORY
|
||
|
|
assert response.source_id == "camera-front"
|
||
|
|
assert response.model is not None
|
||
|
|
assert response.model.model_id == GAUGE_MODEL_ID
|
||
|
|
assert response.outputs == (
|
||
|
|
ScalarOutput(name="reading", value=12.5, unit="bar"),
|
||
|
|
)
|
||
|
|
assert [artifact.role for artifact in response.artifacts] == [
|
||
|
|
"original",
|
||
|
|
"annotated",
|
||
|
|
"diagnostics",
|
||
|
|
]
|
||
|
|
original = response.artifacts[0]
|
||
|
|
annotated = response.artifacts[1]
|
||
|
|
diagnostics = response.artifacts[2]
|
||
|
|
assert isinstance(original, ImageArtifact)
|
||
|
|
assert base64.b64decode(original.data) == b"encoded-gauge-image"
|
||
|
|
assert isinstance(annotated, ImageArtifact)
|
||
|
|
assert base64.b64decode(annotated.data) == b"annotated-jpeg"
|
||
|
|
assert isinstance(diagnostics, JsonArtifact)
|
||
|
|
assert diagnostics.data == {
|
||
|
|
"ocr_labels": 4,
|
||
|
|
"fit": {"residual": 0.02},
|
||
|
|
}
|
||
|
|
|
||
|
|
command = constructed["command"]
|
||
|
|
assert isinstance(command, tuple)
|
||
|
|
assert Path(command[2]).is_absolute()
|
||
|
|
assert command[2].endswith("cmvr_edge_ai/gauge/analog_gauge_worker.py")
|
||
|
|
max_pixels_index = command.index("--max-pixels")
|
||
|
|
assert command[max_pixels_index + 1] == "1234567"
|
||
|
|
sent_header, sent_blobs, timeout_s = fake.calls[0]
|
||
|
|
assert sent_header["request_id"] == "gauge-1"
|
||
|
|
assert sent_header["model_id"] == GAUGE_MODEL_ID
|
||
|
|
assert sent_header["image"] == {
|
||
|
|
"blob_index": 0,
|
||
|
|
"name": "image",
|
||
|
|
"role": "original",
|
||
|
|
"media_type": "image/jpeg",
|
||
|
|
"sha256": hashlib.sha256(b"encoded-gauge-image").hexdigest(),
|
||
|
|
"width": 640,
|
||
|
|
"height": 480,
|
||
|
|
}
|
||
|
|
assert sent_blobs == (b"encoded-gauge-image",)
|
||
|
|
assert timeout_s == 34
|
||
|
|
await operator.stop()
|
||
|
|
assert fake.stopped is True
|
||
|
|
|
||
|
|
asyncio.run(exercise())
|
||
|
|
|
||
|
|
|
||
|
|
def test_invalid_worker_response_fails_one_request_then_recovers(
|
||
|
|
tmp_path: Path,
|
||
|
|
) -> None:
|
||
|
|
async def exercise() -> None:
|
||
|
|
fake = FakeWorker(
|
||
|
|
(
|
||
|
|
_worker_response(request_id="wrong"),
|
||
|
|
_worker_response(request_id="gauge-2"),
|
||
|
|
)
|
||
|
|
)
|
||
|
|
operator = AnalogGaugeReaderOperator(
|
||
|
|
"reader",
|
||
|
|
_params(tmp_path),
|
||
|
|
worker_factory=lambda *args, **kwargs: fake,
|
||
|
|
)
|
||
|
|
|
||
|
|
failed = (
|
||
|
|
await operator.process(_envelope(_request(request_id="gauge-1")))
|
||
|
|
).envelope.payload
|
||
|
|
succeeded = (
|
||
|
|
await operator.process(_envelope(_request(request_id="gauge-2")))
|
||
|
|
).envelope.payload
|
||
|
|
|
||
|
|
assert failed.status is InferenceStatus.FAILED
|
||
|
|
assert failed.error is not None
|
||
|
|
assert failed.error.code == "GAUGE_INVALID_RESPONSE"
|
||
|
|
assert succeeded.status is InferenceStatus.SUCCEEDED
|
||
|
|
health = await operator.health()
|
||
|
|
assert health.healthy is True
|
||
|
|
assert "completed=1" in health.detail
|
||
|
|
assert "failed=1" in health.detail
|
||
|
|
await operator.stop()
|
||
|
|
|
||
|
|
asyncio.run(exercise())
|
||
|
|
|
||
|
|
|
||
|
|
def test_worker_failure_response_preserves_error_without_killing_operator(
|
||
|
|
tmp_path: Path,
|
||
|
|
) -> None:
|
||
|
|
async def exercise() -> None:
|
||
|
|
fake = FakeWorker(
|
||
|
|
(
|
||
|
|
_worker_response(
|
||
|
|
status="failed",
|
||
|
|
value=None,
|
||
|
|
with_annotated=False,
|
||
|
|
input_image=None,
|
||
|
|
),
|
||
|
|
_worker_response(request_id="gauge-2"),
|
||
|
|
)
|
||
|
|
)
|
||
|
|
operator = AnalogGaugeReaderOperator(
|
||
|
|
"reader",
|
||
|
|
_params(tmp_path),
|
||
|
|
worker_factory=lambda *args, **kwargs: fake,
|
||
|
|
)
|
||
|
|
|
||
|
|
failed = (
|
||
|
|
await operator.process(_envelope(_request(request_id="gauge-1")))
|
||
|
|
).envelope.payload
|
||
|
|
succeeded = (
|
||
|
|
await operator.process(_envelope(_request(request_id="gauge-2")))
|
||
|
|
).envelope.payload
|
||
|
|
|
||
|
|
assert failed.status is InferenceStatus.FAILED
|
||
|
|
assert failed.error is not None
|
||
|
|
assert failed.error.code == "GAUGE_READER_FAILED"
|
||
|
|
assert failed.error.stage == "ocr"
|
||
|
|
assert failed.artifacts == ()
|
||
|
|
assert succeeded.status is InferenceStatus.SUCCEEDED
|
||
|
|
await operator.stop()
|
||
|
|
|
||
|
|
asyncio.run(exercise())
|
||
|
|
|
||
|
|
|
||
|
|
def test_timeout_becomes_retryable_failed_response(tmp_path: Path) -> None:
|
||
|
|
async def exercise() -> None:
|
||
|
|
fake = FakeWorker((asyncio.TimeoutError(),))
|
||
|
|
operator = AnalogGaugeReaderOperator(
|
||
|
|
"reader",
|
||
|
|
_params(tmp_path),
|
||
|
|
worker_factory=lambda *args, **kwargs: fake,
|
||
|
|
)
|
||
|
|
|
||
|
|
response = (
|
||
|
|
await operator.process(_envelope(_request(roles=())))
|
||
|
|
).envelope.payload
|
||
|
|
|
||
|
|
assert response.status is InferenceStatus.FAILED
|
||
|
|
assert response.error is not None
|
||
|
|
assert response.error.code == "GAUGE_TIMEOUT"
|
||
|
|
assert response.error.retryable is True
|
||
|
|
await operator.stop()
|
||
|
|
|
||
|
|
asyncio.run(exercise())
|
||
|
|
|
||
|
|
|
||
|
|
@pytest.mark.parametrize(
|
||
|
|
("status", "value", "expected"),
|
||
|
|
[
|
||
|
|
("no_result", None, InferenceStatus.NO_RESULT),
|
||
|
|
("partial", 3.25, InferenceStatus.PARTIAL),
|
||
|
|
],
|
||
|
|
)
|
||
|
|
def test_operator_maps_non_success_statuses(
|
||
|
|
tmp_path: Path,
|
||
|
|
status: str,
|
||
|
|
value: float | None,
|
||
|
|
expected: InferenceStatus,
|
||
|
|
) -> None:
|
||
|
|
async def exercise() -> None:
|
||
|
|
fake = FakeWorker(
|
||
|
|
(
|
||
|
|
_worker_response(
|
||
|
|
status=status,
|
||
|
|
value=value,
|
||
|
|
with_annotated=False,
|
||
|
|
),
|
||
|
|
)
|
||
|
|
)
|
||
|
|
operator = AnalogGaugeReaderOperator(
|
||
|
|
"reader",
|
||
|
|
_params(tmp_path),
|
||
|
|
worker_factory=lambda *args, **kwargs: fake,
|
||
|
|
)
|
||
|
|
response = (
|
||
|
|
await operator.process(_envelope(_request(roles=("diagnostics",))))
|
||
|
|
).envelope.payload
|
||
|
|
assert response.status is expected
|
||
|
|
assert bool(response.outputs) is (value is not None)
|
||
|
|
await operator.stop()
|
||
|
|
|
||
|
|
asyncio.run(exercise())
|
||
|
|
|
||
|
|
|
||
|
|
def test_missing_unit_is_partial_with_warning(tmp_path: Path) -> None:
|
||
|
|
async def exercise() -> None:
|
||
|
|
worker_response = _worker_response(with_annotated=False)
|
||
|
|
worker_response = FramedWorkerResponse(
|
||
|
|
header={**worker_response.header, "unit": None},
|
||
|
|
blobs=worker_response.blobs,
|
||
|
|
)
|
||
|
|
fake = FakeWorker((worker_response,))
|
||
|
|
operator = AnalogGaugeReaderOperator(
|
||
|
|
"reader",
|
||
|
|
_params(tmp_path),
|
||
|
|
worker_factory=lambda *args, **kwargs: fake,
|
||
|
|
)
|
||
|
|
response = (
|
||
|
|
await operator.process(_envelope(_request(roles=())))
|
||
|
|
).envelope.payload
|
||
|
|
assert response.status is InferenceStatus.PARTIAL
|
||
|
|
assert response.outputs == (
|
||
|
|
ScalarOutput(name="reading", value=12.5, unit=None),
|
||
|
|
)
|
||
|
|
assert "gauge unit was not recognized" in response.warnings
|
||
|
|
await operator.stop()
|
||
|
|
|
||
|
|
asyncio.run(exercise())
|
||
|
|
|
||
|
|
|
||
|
|
def test_operator_configuration_is_strict_and_resolves_model_paths(
|
||
|
|
tmp_path: Path,
|
||
|
|
) -> None:
|
||
|
|
params = _params(tmp_path)
|
||
|
|
fake = FakeWorker((_worker_response(),))
|
||
|
|
operator = AnalogGaugeReaderOperator(
|
||
|
|
"reader",
|
||
|
|
params,
|
||
|
|
worker_factory=lambda *args, **kwargs: fake,
|
||
|
|
)
|
||
|
|
assert operator.params.detection_model_path == (
|
||
|
|
tmp_path / "models/detection.pt"
|
||
|
|
).resolve()
|
||
|
|
assert operator.params.max_pixels == 1_234_567
|
||
|
|
|
||
|
|
default_params = dict(params)
|
||
|
|
default_params.pop("max_pixels")
|
||
|
|
default_operator = AnalogGaugeReaderOperator(
|
||
|
|
"reader-default-pixels",
|
||
|
|
default_params,
|
||
|
|
worker_factory=lambda *args, **kwargs: fake,
|
||
|
|
)
|
||
|
|
assert default_operator.params.max_pixels == 25_000_000
|
||
|
|
|
||
|
|
with pytest.raises(ValidationError, match="Extra inputs are not permitted"):
|
||
|
|
AnalogGaugeReaderOperator(
|
||
|
|
"reader",
|
||
|
|
{**params, "unknown": True},
|
||
|
|
worker_factory=lambda *args, **kwargs: fake,
|
||
|
|
)
|
||
|
|
with pytest.raises(ValidationError, match="absolute path"):
|
||
|
|
AnalogGaugeReaderOperator(
|
||
|
|
"reader",
|
||
|
|
{**params, "python_executable": "python3.8"},
|
||
|
|
worker_factory=lambda *args, **kwargs: fake,
|
||
|
|
)
|
||
|
|
for invalid_max_pixels in (0, 250_000_001):
|
||
|
|
with pytest.raises(ValidationError, match="max_pixels"):
|
||
|
|
AnalogGaugeReaderOperator(
|
||
|
|
"reader",
|
||
|
|
{**params, "max_pixels": invalid_max_pixels},
|
||
|
|
worker_factory=lambda *args, **kwargs: fake,
|
||
|
|
)
|
||
|
|
|
||
|
|
|
||
|
|
def test_operator_command_preserves_virtualenv_python_symlink(
|
||
|
|
tmp_path: Path,
|
||
|
|
) -> None:
|
||
|
|
params = _params(tmp_path)
|
||
|
|
python_link = tmp_path / "worker-python"
|
||
|
|
python_link.symlink_to(sys.executable)
|
||
|
|
params["python_executable"] = str(python_link)
|
||
|
|
fake = FakeWorker((_worker_response(),))
|
||
|
|
constructed: dict[str, object] = {}
|
||
|
|
|
||
|
|
def factory(command, **kwargs): # type: ignore[no-untyped-def]
|
||
|
|
constructed["command"] = tuple(command)
|
||
|
|
return fake
|
||
|
|
|
||
|
|
operator = AnalogGaugeReaderOperator(
|
||
|
|
"reader",
|
||
|
|
params,
|
||
|
|
worker_factory=factory,
|
||
|
|
)
|
||
|
|
|
||
|
|
command = constructed["command"]
|
||
|
|
assert isinstance(command, tuple)
|
||
|
|
assert command[0] == str(python_link)
|
||
|
|
assert operator.params.python_executable == python_link
|