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)