248 lines
8.3 KiB
Python
248 lines
8.3 KiB
Python
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import base64
|
|
import hashlib
|
|
from pathlib import Path
|
|
import struct
|
|
import sys
|
|
|
|
from cmvr_edge_ai.workers import FramedSubprocessWorker, PROTOCOL_VERSION
|
|
|
|
|
|
_ONE_PIXEL_PNG = base64.b64decode(
|
|
"iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAQAAAC1HAwCAAAAC0lEQVR42mNk"
|
|
"+A8AAQUBAScY42YAAAAASUVORK5CYII="
|
|
)
|
|
_TWO_PIXEL_HEADER_PNG = (
|
|
_ONE_PIXEL_PNG[:16] + struct.pack("!I", 2) + _ONE_PIXEL_PNG[20:]
|
|
)
|
|
_MODEL_ID = "ethz-analog-gauge-reader@1"
|
|
_CATEGORY = "gauge.analog"
|
|
|
|
|
|
def _request_header(
|
|
request_id: str,
|
|
*,
|
|
sha256: str,
|
|
width: int = 1,
|
|
height: int = 1,
|
|
) -> dict[str, object]:
|
|
return {
|
|
"protocol": PROTOCOL_VERSION,
|
|
"type": "request",
|
|
"request_id": request_id,
|
|
"category": _CATEGORY,
|
|
"model_id": _MODEL_ID,
|
|
"parameters": {},
|
|
"requested_artifact_roles": ["original", "annotated", "diagnostics"],
|
|
"image": {
|
|
"blob_index": 0,
|
|
"name": "image",
|
|
"role": "original",
|
|
"media_type": "image/png",
|
|
"sha256": sha256,
|
|
"width": width,
|
|
"height": height,
|
|
},
|
|
}
|
|
|
|
|
|
def test_analog_gauge_worker_frames_results_and_survives_model_errors(
|
|
tmp_path: Path,
|
|
) -> None:
|
|
project = tmp_path / "fake_gauge_project"
|
|
project.mkdir()
|
|
(project / "pipeline.py").write_text(
|
|
"""
|
|
import json
|
|
import os
|
|
import shutil
|
|
|
|
_calls = 0
|
|
|
|
def process_image(image, detection_model_path, key_point_model_path,
|
|
segmentation_model_path, run_path, debug, eval_mode,
|
|
image_is_raw=False):
|
|
global _calls
|
|
_calls += 1
|
|
print("upstream stdout noise must be redirected")
|
|
if image_is_raw:
|
|
raise AssertionError("worker must pass the encoded image path")
|
|
os.mkdir(run_path)
|
|
with open(os.path.join(run_path, "error.json"), "w") as stream:
|
|
json.dump({"fit_residual": 0.125}, stream)
|
|
if debug:
|
|
shutil.copyfile(image, os.path.join(run_path, "ellipse_results_final.jpg"))
|
|
with open(os.path.join(os.path.dirname(__file__), "seen_paths.txt"), "a") as stream:
|
|
stream.write(run_path + "\\n")
|
|
if _calls == 2:
|
|
raise Exception("No gauge detected in image")
|
|
if _calls == 3:
|
|
return {"value": -3.25, "unit": None}
|
|
if _calls == 4:
|
|
raise RuntimeError("Needle segmentation model checkpoint is incompatible")
|
|
if _calls == 5:
|
|
raise Exception("Segmentation failed, no needle found")
|
|
return {"value": 12.5, "unit": "bar"}
|
|
""".lstrip(),
|
|
encoding="utf-8",
|
|
)
|
|
model_paths = []
|
|
for name in ("detection.pt", "key_point.pt", "segmentation.pt"):
|
|
model_path = project / name
|
|
model_path.write_bytes(b"fake-model")
|
|
model_paths.append(model_path)
|
|
|
|
worker_script = (
|
|
Path(__file__).parents[2]
|
|
/ "src"
|
|
/ "cmvr_edge_ai"
|
|
/ "gauge"
|
|
/ "analog_gauge_worker.py"
|
|
)
|
|
command = (
|
|
sys.executable,
|
|
"-u",
|
|
str(worker_script),
|
|
"--project-root",
|
|
str(project),
|
|
"--detection-model",
|
|
str(model_paths[0]),
|
|
"--key-point-model",
|
|
str(model_paths[1]),
|
|
"--segmentation-model",
|
|
str(model_paths[2]),
|
|
"--model-id",
|
|
_MODEL_ID,
|
|
"--max-pixels",
|
|
"1",
|
|
)
|
|
|
|
async def exercise() -> None:
|
|
worker = FramedSubprocessWorker(
|
|
command,
|
|
cwd=project,
|
|
startup_timeout_s=3,
|
|
shutdown_timeout_s=1,
|
|
)
|
|
try:
|
|
ready = await worker.start()
|
|
assert ready == {
|
|
"protocol": PROTOCOL_VERSION,
|
|
"type": "ready",
|
|
"model_id": _MODEL_ID,
|
|
"category": _CATEGORY,
|
|
}
|
|
|
|
digest = hashlib.sha256(_ONE_PIXEL_PNG).hexdigest()
|
|
response = await worker.submit(
|
|
_request_header("gauge-ok", sha256=digest),
|
|
(_ONE_PIXEL_PNG,),
|
|
timeout_s=3,
|
|
)
|
|
assert response.header["request_id"] == "gauge-ok"
|
|
assert response.header["status"] == "succeeded"
|
|
assert response.header["value"] == 12.5
|
|
assert response.header["unit"] == "bar"
|
|
assert response.header["input_image"] == {"width": 1, "height": 1}
|
|
assert response.header["diagnostics"] == {"fit_residual": 0.125}
|
|
assert response.header["error"] is None
|
|
assert response.header["artifacts"] == [
|
|
{
|
|
"blob_index": 0,
|
|
"role": "annotated",
|
|
"media_type": "image/png",
|
|
"width": 1,
|
|
"height": 1,
|
|
"sha256": digest,
|
|
}
|
|
]
|
|
assert response.blobs == (_ONE_PIXEL_PNG,)
|
|
|
|
failed = await worker.submit(
|
|
_request_header("bad-sha", sha256="0" * 64),
|
|
(_ONE_PIXEL_PNG,),
|
|
timeout_s=3,
|
|
)
|
|
assert failed.header["request_id"] == "bad-sha"
|
|
assert failed.header["status"] == "failed"
|
|
assert failed.header["error"]["code"] == "IMAGE_INTEGRITY_ERROR"
|
|
assert failed.header["input_image"] is None
|
|
assert worker.is_alive is True
|
|
|
|
oversized_digest = hashlib.sha256(_TWO_PIXEL_HEADER_PNG).hexdigest()
|
|
oversized = await worker.submit(
|
|
_request_header(
|
|
"too-many-pixels",
|
|
sha256=oversized_digest,
|
|
width=2,
|
|
),
|
|
(_TWO_PIXEL_HEADER_PNG,),
|
|
timeout_s=3,
|
|
)
|
|
assert oversized.header["status"] == "failed"
|
|
assert oversized.header["error"]["code"] == "IMAGE_TOO_LARGE"
|
|
assert oversized.header["error"]["stage"] == "decode"
|
|
assert oversized.header["input_image"] is None
|
|
|
|
no_result = await worker.submit(
|
|
_request_header("no-gauge", sha256=digest),
|
|
(_ONE_PIXEL_PNG,),
|
|
timeout_s=3,
|
|
)
|
|
assert no_result.header["status"] == "no_result"
|
|
assert no_result.header["value"] is None
|
|
assert no_result.header["input_image"] == {"width": 1, "height": 1}
|
|
assert no_result.header["error"] is None
|
|
assert no_result.header["diagnostics"]["worker"] == {
|
|
"outcome_code": "GAUGE_NOT_FOUND",
|
|
"stage": "detection",
|
|
}
|
|
assert no_result.blobs == ()
|
|
|
|
partial = await worker.submit(
|
|
_request_header("unit-missing", sha256=digest),
|
|
(_ONE_PIXEL_PNG,),
|
|
timeout_s=3,
|
|
)
|
|
assert partial.header["status"] == "partial"
|
|
assert partial.header["value"] == -3.25
|
|
assert partial.header["unit"] is None
|
|
assert "no measurement unit" in partial.header["warnings"][0]
|
|
assert partial.blobs == (_ONE_PIXEL_PNG,)
|
|
|
|
checkpoint_failure = await worker.submit(
|
|
_request_header("bad-checkpoint", sha256=digest),
|
|
(_ONE_PIXEL_PNG,),
|
|
timeout_s=3,
|
|
)
|
|
assert checkpoint_failure.header["status"] == "failed"
|
|
assert checkpoint_failure.header["error"]["code"] == "INFERENCE_FAILED"
|
|
assert checkpoint_failure.header["diagnostics"]["worker"] == {
|
|
"outcome_code": "INFERENCE_FAILED",
|
|
"stage": "inference",
|
|
"exception_type": "RuntimeError",
|
|
}
|
|
|
|
needle_no_result = await worker.submit(
|
|
_request_header("needle-not-found", sha256=digest),
|
|
(_ONE_PIXEL_PNG,),
|
|
timeout_s=3,
|
|
)
|
|
assert needle_no_result.header["status"] == "no_result"
|
|
assert needle_no_result.header["error"] is None
|
|
assert needle_no_result.header["diagnostics"]["worker"] == {
|
|
"outcome_code": "GAUGE_NEEDLE_FAILED",
|
|
"stage": "segmentation",
|
|
}
|
|
assert worker.is_alive is True
|
|
finally:
|
|
await worker.stop()
|
|
|
|
asyncio.run(exercise())
|
|
|
|
recorded = (project / "seen_paths.txt").read_text(encoding="utf-8").splitlines()
|
|
assert len(recorded) == 5
|
|
assert all(not Path(path).exists() for path in recorded)
|