from __future__ import annotations import hashlib import json import math import os import re import stat import subprocess import threading from dataclasses import dataclass from pathlib import Path from typing import Any, TypeGuard import numpy as np import rerun as rr from k1link.artifacts import write_json_atomic from k1link.sessions import SessionIntegrityError from .jobs import CameraComputeJob, validate_camera_compute_job COMPUTE_RESULT_SCHEMA = "missioncore.compute-result/v1" COMPUTE_RESULT_IDENTITY_SCHEMA = "missioncore.compute-result-identity/v1" OBJECT_DETECTIONS_SCHEMA = "missioncore.object-detections/v1" SESSION_TIMELINE = "session_time" MAX_RESULT_JSON_BYTES = 32 * 1024 * 1024 MAX_RESULT_FRAMES = 10_000 MAX_DETECTIONS_PER_FRAME = 10_000 MAX_OVERLAY_SOURCE_BYTES = 512 * 1024 * 1024 MAX_OVERLAY_DECODED_BYTES = 512 * 1024 * 1024 MAX_OVERLAY_RRD_BYTES = 512 * 1024 * 1024 MAX_JOB_SCAN = 512 _SAFE_RESULT_ID = re.compile(r"^result-[a-f0-9]{64}$") _SAFE_RECORDING_ID = re.compile(r"^[A-Za-z0-9][A-Za-z0-9._:-]{0,127}$") _SHA256 = re.compile(r"^[a-f0-9]{64}$") @dataclass(frozen=True, slots=True) class ObjectDetection: class_id: int class_name: str score: float bbox_xyxy: tuple[float, float, float, float] @dataclass(frozen=True, slots=True) class DetectionFrame: frame_index: int session_time_ns: int detections: tuple[ObjectDetection, ...] @dataclass(frozen=True, slots=True) class RecordedPerceptionResult: result_id: str result_root: Path job: CameraComputeJob created_at_utc: str frames: tuple[DetectionFrame, ...] class RecordedPerceptionOverlayError(RuntimeError): """A validated compute result could not be projected into Rerun.""" def validate_recorded_perception_result( job_root: Path, result_root: Path, ) -> RecordedPerceptionResult: """Validate one content-addressed worker result against its exact job.""" job = validate_camera_compute_job(job_root) root = result_root.expanduser().resolve(strict=True) if not root.is_dir() or _SAFE_RESULT_ID.fullmatch(root.name) is None: raise SessionIntegrityError("compute result root is invalid") result = _read_json_object(root / "result.json", root) detections = _read_json_object(root / "detections.json", root) identity_sha256 = result.get("identity_sha256") pipeline = result.get("pipeline") model = result.get("model") parameters = result.get("parameters") if ( result.get("schema_version") != COMPUTE_RESULT_SCHEMA or result.get("result_id") != root.name or not isinstance(identity_sha256, str) or _SHA256.fullmatch(identity_sha256) is None or root.name != f"result-{identity_sha256}" or result.get("job_id") != job.job_id or result.get("input_sha256") != job.input_sha256 or not isinstance(pipeline, dict) or not isinstance(model, dict) or not isinstance(parameters, dict) ): raise SessionIntegrityError("compute result identity is inconsistent") identity = { "schema_version": COMPUTE_RESULT_IDENTITY_SCHEMA, "job_id": job.job_id, "input_sha256": job.input_sha256, "pipeline": pipeline, "model": { "id": model.get("id"), "version": model.get("version"), "sha256": model.get("sha256"), }, "parameters": parameters, } if hashlib.sha256(_canonical_json(identity)).hexdigest() != identity_sha256: raise SessionIntegrityError("compute result generation is inconsistent") artifacts = result.get("artifacts") detection_path = root / "detections.json" detection_stat = _confined_regular_file(detection_path, root) if not isinstance(artifacts, list) or len(artifacts) != 1: raise SessionIntegrityError("compute result artifact descriptor is invalid") artifact = artifacts[0] if ( not isinstance(artifact, dict) or artifact.get("kind") != "object-detections" or artifact.get("path") != "detections.json" or artifact.get("schema_version") != OBJECT_DETECTIONS_SCHEMA or artifact.get("byte_length") != detection_stat.st_size or artifact.get("sha256") != _sha256_file(detection_path) ): raise SessionIntegrityError("compute result artifact identity changed") if ( detections.get("schema_version") != OBJECT_DETECTIONS_SCHEMA or detections.get("result_id") != root.name or detections.get("job_id") != job.job_id or detections.get("input_sha256") != job.input_sha256 or detections.get("timestamp_basis") != "session-time-seconds" ): raise SessionIntegrityError("object detections are not bound to the result") raw_frames = detections.get("frames") if not isinstance(raw_frames, list) or not 1 <= len(raw_frames) <= MAX_RESULT_FRAMES: raise SessionIntegrityError("object detection frame set is outside bounds") frames: list[DetectionFrame] = [] detection_count = 0 previous_time_ns = -1 for expected_index, value in enumerate(raw_frames): frame = _validate_detection_frame(value, expected_index, job) if frame.session_time_ns <= previous_time_ns: raise SessionIntegrityError("object detection timestamps are not increasing") previous_time_ns = frame.session_time_ns frames.append(frame) detection_count += len(frame.detections) metrics = result.get("metrics") if ( not isinstance(metrics, dict) or metrics.get("frames_processed") != len(frames) or metrics.get("detections") != detection_count ): raise SessionIntegrityError("compute result metrics do not match its artifact") created_at_utc = result.get("created_at_utc") if not isinstance(created_at_utc, str) or len(created_at_utc) > 64: raise SessionIntegrityError("compute result creation time is invalid") return RecordedPerceptionResult( result_id=root.name, result_root=root, job=job, created_at_utc=created_at_utc, frames=tuple(frames), ) class RecordedPerceptionOverlayStore: """Discover admitted local results and cache complete overlay RRD files.""" def __init__( self, *, jobs_root: Path, results_root: Path, cache_root: Path, ffmpeg_path: Path, ffprobe_path: Path, ) -> None: self.jobs_root = jobs_root.expanduser().absolute() self.results_root = results_root.expanduser().absolute() self.cache_root = cache_root.expanduser().absolute() self.ffmpeg_path = ffmpeg_path.expanduser().resolve(strict=True) self.ffprobe_path = ffprobe_path.expanduser().resolve(strict=True) self._lock = threading.Lock() def render( self, session_id: str, *, application_id: str, recording_id: str, ) -> bytes | None: if _SAFE_RECORDING_ID.fullmatch(session_id) is None: raise ValueError("observation session id is invalid") if application_id != "nodedc_mission_core_recorded": raise ValueError("recorded perception application id is invalid") if _SAFE_RECORDING_ID.fullmatch(recording_id) is None: raise ValueError("recorded perception recording id is invalid") with self._lock: result = self._latest_result(session_id) if result is None: return None cache_root = _private_directory(self.cache_root) session_cache = _private_child_directory(cache_root, session_id) cache_dir = _private_child_directory(session_cache, result.result_id) output = cache_dir / f"{recording_id}.rrd" sidecar = output.with_suffix(".rrd.cache.json") cached = _read_cached_overlay(output, sidecar, result) if cached is not None: return cached payload = _render_overlay( result, application_id=application_id, recording_id=recording_id, ffmpeg_path=self.ffmpeg_path, ffprobe_path=self.ffprobe_path, ) temporary = output.with_name(f".{output.name}.{os.getpid()}.tmp") try: with temporary.open("xb") as stream: stream.write(payload) stream.flush() os.fsync(stream.fileno()) os.chmod(temporary, 0o600) os.replace(temporary, output) write_json_atomic( sidecar, { "schema_version": "missioncore.perception-overlay-cache/v1", "result_id": result.result_id, "recording_id": recording_id, "byte_length": len(payload), "sha256": hashlib.sha256(payload).hexdigest(), }, ) finally: temporary.unlink(missing_ok=True) return payload def _latest_result(self, session_id: str) -> RecordedPerceptionResult | None: try: job_roots = sorted(self.jobs_root.iterdir()) except FileNotFoundError: return None if len(job_roots) > MAX_JOB_SCAN: raise RecordedPerceptionOverlayError("compute job catalog is outside bounds") matches: list[RecordedPerceptionResult] = [] for job_root in job_roots: if job_root.is_symlink(): continue try: job = validate_camera_compute_job(job_root) except (OSError, SessionIntegrityError): continue if job.session_id != session_id: continue result_parent = self.results_root / job.job_id try: result_roots = sorted( path for path in result_parent.iterdir() if path.is_dir() and not path.is_symlink() and _SAFE_RESULT_ID.fullmatch(path.name) is not None ) except FileNotFoundError: continue for result_root in result_roots: matches.append(validate_recorded_perception_result(job_root, result_root)) if not matches: return None return max(matches, key=lambda value: (value.created_at_utc, value.result_id)) def _validate_detection_frame( value: object, expected_index: int, job: CameraComputeJob, ) -> DetectionFrame: if not isinstance(value, dict) or value.get("frame_index") != expected_index: raise SessionIntegrityError("object detection frame index is inconsistent") session_seconds = value.get("session_seconds") epoch_seconds = value.get("epoch_seconds") raw_detections = value.get("detections") if ( not _finite_number(session_seconds) or not _finite_number(epoch_seconds) or float(epoch_seconds) < 0 or float(session_seconds) < job.timeline_start_seconds - 0.001 or float(session_seconds) > job.timeline_end_seconds + 0.001 or not isinstance(raw_detections, list) or len(raw_detections) > MAX_DETECTIONS_PER_FRAME ): raise SessionIntegrityError("object detection frame is invalid") found = tuple(_validate_detection(item) for item in raw_detections) return DetectionFrame( frame_index=expected_index, session_time_ns=round(float(session_seconds) * 1_000_000_000), detections=found, ) def _validate_detection(value: object) -> ObjectDetection: if not isinstance(value, dict): raise SessionIntegrityError("object detection is invalid") class_id = value.get("class_id") class_name = value.get("class_name") score = value.get("score") bbox = value.get("bbox_xyxy") if ( not isinstance(class_id, int) or isinstance(class_id, bool) or not 0 <= class_id <= 10_000 or not isinstance(class_name, str) or not class_name or len(class_name) > 128 or not _finite_number(score) or not 0 <= float(score) <= 1 or not isinstance(bbox, list) or len(bbox) != 4 or not all(_finite_number(item) for item in bbox) ): raise SessionIntegrityError("object detection fields are invalid") x1, y1, x2, y2 = (float(item) for item in bbox) if min(x1, y1) < 0 or x2 < x1 or y2 < y1 or max(x2, y2) > 100_000: raise SessionIntegrityError("object detection bounds are invalid") return ObjectDetection(class_id, class_name, float(score), (x1, y1, x2, y2)) def _render_overlay( result: RecordedPerceptionResult, *, application_id: str, recording_id: str, ffmpeg_path: Path, ffprobe_path: Path, ) -> bytes: if validate_camera_compute_job(result.job.job_root) != result.job: raise RecordedPerceptionOverlayError("camera compute job changed before projection") source = _camera_source_bytes(result.job) width, height, frame_count = _probe_video(source, ffprobe_path) if frame_count != len(result.frames): raise RecordedPerceptionOverlayError("camera and detection frame counts differ") _validate_image_bounds(result, width, height) frame_bytes = width * height * 3 expected_bytes = frame_bytes * len(result.frames) if expected_bytes > MAX_OVERLAY_DECODED_BYTES: raise RecordedPerceptionOverlayError("decoded camera overlay is outside bounds") decoded = _decode_rgb(source, ffmpeg_path, frame_count) if len(decoded) != expected_bytes: raise RecordedPerceptionOverlayError("decoded camera frame count changed") recording = rr.RecordingStream( application_id, recording_id=recording_id, send_properties=False, ) stream = rr.binary_stream(recording) try: recording.log( "/perception/camera/metadata", rr.AnyValues( result_id=result.result_id, job_id=result.job.job_id, source_id=result.job.source_id, pipeline="recorded-camera-coco-detection/v1", warning="Generic detections are evidence, not a driving decision.", ), static=True, ) for frame in result.frames: offset = frame.frame_index * frame_bytes image = np.frombuffer( decoded, dtype=np.uint8, count=frame_bytes, offset=offset, ).reshape((height, width, 3)) recording.set_time( SESSION_TIMELINE, duration=np.timedelta64(frame.session_time_ns, "ns"), ) recording.log( "/perception/camera/image", rr.Image(image, color_model="RGB"), ) if frame.detections: recording.log( "/perception/camera/detections", rr.Boxes2D( array=[detection.bbox_xyxy for detection in frame.detections], array_format=rr.Box2DFormat.XYXY, labels=[ f"{detection.class_name} ยท {detection.score:.0%}" for detection in frame.detections ], colors=[ _class_color(detection.class_id) for detection in frame.detections ], show_labels=True, ), ) else: recording.log( "/perception/camera/detections", rr.Clear(recursive=False), ) payload = stream.read(flush=True, flush_timeout_sec=120.0) except Exception as exc: raise RecordedPerceptionOverlayError("failed to serialize perception overlay") from exc finally: recording.disconnect() if payload is None or not payload.startswith(b"RRF2") or len(payload) > MAX_OVERLAY_RRD_BYTES: raise RecordedPerceptionOverlayError("serialized perception overlay is invalid") return payload def _camera_source_bytes(job: CameraComputeJob) -> bytes: epoch = ( job.job_root / "input" / "camera" / job.source_id / f"epoch-{job.codec_epoch}" ) paths = [epoch / "init.mp4"] + [ epoch / "segments" / f"{sequence}.m4s" for sequence in range(1, job.segment_count + 1) ] total = sum(path.stat().st_size for path in paths) if total > MAX_OVERLAY_SOURCE_BYTES: raise RecordedPerceptionOverlayError("camera overlay source is outside bounds") return b"".join(path.read_bytes() for path in paths) def _probe_video(source: bytes, ffprobe_path: Path) -> tuple[int, int, int]: completed = _run_media_tool( [ str(ffprobe_path), "-v", "error", "-count_frames", "-select_streams", "v:0", "-show_entries", "stream=width,height,nb_read_frames", "-of", "json", "-i", "pipe:0", ], source, ) try: document = json.loads(completed.stdout) stream = document["streams"][0] width = int(stream["width"]) height = int(stream["height"]) frame_count = int(stream["nb_read_frames"]) except (IndexError, KeyError, TypeError, ValueError, json.JSONDecodeError) as exc: raise RecordedPerceptionOverlayError("camera dimensions are unavailable") from exc if ( width < 1 or height < 1 or width * height > 4_194_304 or not 1 <= frame_count <= MAX_RESULT_FRAMES ): raise RecordedPerceptionOverlayError("camera dimensions are outside bounds") return width, height, frame_count def _decode_rgb(source: bytes, ffmpeg_path: Path, frame_count: int) -> bytes: return _run_media_tool( [ str(ffmpeg_path), "-v", "error", "-i", "pipe:0", "-map", "0:v:0", "-frames:v", str(frame_count), "-fps_mode", "passthrough", "-pix_fmt", "rgb24", "-f", "rawvideo", "pipe:1", ], source, ).stdout def _run_media_tool(argv: list[str], payload: bytes) -> subprocess.CompletedProcess[bytes]: try: metadata = Path(argv[0]).lstat() if stat.S_ISLNK(metadata.st_mode) or not stat.S_ISREG(metadata.st_mode): raise RecordedPerceptionOverlayError("media tool is not a regular file") completed = subprocess.run( argv, input=payload, capture_output=True, check=False, timeout=120, ) except (OSError, subprocess.SubprocessError) as exc: raise RecordedPerceptionOverlayError("media tool failed") from exc if completed.returncode != 0: raise RecordedPerceptionOverlayError("camera media could not be decoded") return completed def _read_cached_overlay( output: Path, sidecar: Path, result: RecordedPerceptionResult, ) -> bytes | None: try: value = _read_json_object(sidecar, sidecar.parent) _confined_regular_file(output, output.parent) payload = output.read_bytes() except (OSError, SessionIntegrityError): return None if ( value.get("schema_version") != "missioncore.perception-overlay-cache/v1" or value.get("result_id") != result.result_id or value.get("recording_id") != output.stem or value.get("byte_length") != len(payload) or value.get("sha256") != hashlib.sha256(payload).hexdigest() or not payload.startswith(b"RRF2") or len(payload) > MAX_OVERLAY_RRD_BYTES ): return None return payload def _validate_image_bounds( result: RecordedPerceptionResult, width: int, height: int, ) -> None: for frame in result.frames: for detection in frame.detections: x1, y1, x2, y2 = detection.bbox_xyxy if x1 > width or x2 > width or y1 > height or y2 > height: raise RecordedPerceptionOverlayError( "object detection escapes the decoded camera image" ) def _private_directory(path: Path) -> Path: path.mkdir(mode=0o700, parents=True, exist_ok=True) try: metadata = path.lstat() except OSError as exc: raise RecordedPerceptionOverlayError("perception cache is unavailable") from exc if stat.S_ISLNK(metadata.st_mode) or not stat.S_ISDIR(metadata.st_mode): raise RecordedPerceptionOverlayError("perception cache is not a real directory") resolved = path.resolve(strict=True) os.chmod(resolved, 0o700) return resolved def _private_child_directory(parent: Path, name: str) -> Path: child = parent / name child.mkdir(mode=0o700, exist_ok=True) try: metadata = child.lstat() resolved = child.resolve(strict=True) except OSError as exc: raise RecordedPerceptionOverlayError("perception cache child is unavailable") from exc if ( stat.S_ISLNK(metadata.st_mode) or not stat.S_ISDIR(metadata.st_mode) or resolved.parent != parent ): raise RecordedPerceptionOverlayError("perception cache child escapes its root") os.chmod(resolved, 0o700) return resolved def _class_color(class_id: int) -> list[int]: palette = ( (185, 255, 74, 255), (67, 191, 255, 255), (255, 193, 92, 255), (222, 110, 255, 255), (255, 103, 117, 255), ) return list(palette[class_id % len(palette)]) def _read_json_object(path: Path, root: Path) -> dict[str, Any]: metadata = _confined_regular_file(path, root) if not 0 < metadata.st_size <= MAX_RESULT_JSON_BYTES: raise SessionIntegrityError("compute result JSON is outside bounds") try: value = json.loads(path.read_text(encoding="utf-8-sig")) except (OSError, UnicodeDecodeError, json.JSONDecodeError) as exc: raise SessionIntegrityError("compute result JSON is unavailable") from exc if not isinstance(value, dict): raise SessionIntegrityError("compute result JSON is not an object") return value def _confined_regular_file(path: Path, root: Path) -> os.stat_result: try: resolved_root = root.resolve(strict=True) resolved = path.resolve(strict=True) metadata = path.lstat() except OSError as exc: raise SessionIntegrityError("compute result artifact is unavailable") from exc if ( stat.S_ISLNK(metadata.st_mode) or not stat.S_ISREG(metadata.st_mode) or not resolved.is_relative_to(resolved_root) ): raise SessionIntegrityError("compute result artifact is not a confined regular file") return metadata def _sha256_file(path: Path) -> str: digest = hashlib.sha256() with path.open("rb") as stream: while chunk := stream.read(1024 * 1024): digest.update(chunk) return digest.hexdigest() def _canonical_json(value: object) -> bytes: try: return json.dumps( value, ensure_ascii=False, sort_keys=True, separators=(",", ":"), allow_nan=False, ).encode("utf-8") except (TypeError, ValueError) as exc: raise SessionIntegrityError("compute result identity cannot be encoded") from exc def _finite_number(value: object) -> TypeGuard[int | float]: return ( isinstance(value, (int, float)) and not isinstance(value, bool) and math.isfinite(float(value)) )