"""Validated calibrated 3D fusion results projected into an opened Rerun recording.""" from __future__ import annotations import hashlib import json import os import re import stat import threading from contextlib import suppress from dataclasses import dataclass from pathlib import Path from typing import Any, Protocol import numpy as np import rerun as rr from rerun.components import FillMode from k1link.artifacts import write_json_atomic from k1link.sessions import SessionIntegrityError from .jobs import CameraComputeJob, validate_camera_compute_job from .perception_epoch import validate_recorded_perception_epoch_result from .results import ( RecordedPerceptionOverlayArtifact, RecordedPerceptionOverlayError, ) FUSION_SCHEMA = "missioncore.recorded-calibrated-fusion/v1" FUSION_IDENTITY_SCHEMA = "missioncore.recorded-calibrated-fusion-identity/v1" SESSION_TIMELINE = "session_time" MAX_JSON_BYTES = 64 * 1024 * 1024 MAX_SCAN = 512 MAX_POINTS_PER_FRAME = 100_000 MAX_BOXES_PER_FRAME = 10_000 _SAFE_FUSION_ID = re.compile(r"^fusion-[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 RecordedCalibratedFusion: fusion_id: str root: Path job: CameraComputeJob perception_result_id: str created_at_utc: str arrays_path: Path labels_path: Path frame_count: int class _OverlayProvider(Protocol): def render( self, session_id: str, *, application_id: str, recording_id: str, ) -> bytes | None: ... class RecordedCalibratedFusionStore: """Discover full-epoch fusion and serialize it for the opened recording ID.""" def __init__( self, *, jobs_root: Path, perception_results_root: Path, fusion_results_root: Path, cache_root: Path, ) -> None: self.jobs_root = jobs_root.expanduser().absolute() self.perception_results_root = perception_results_root.expanduser().absolute() self.fusion_results_root = fusion_results_root.expanduser().absolute() self.cache_root = cache_root.expanduser().absolute() 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 fusion application id is invalid") if _SAFE_RECORDING_ID.fullmatch(recording_id) is None: raise ValueError("recorded fusion recording id is invalid") with self._lock: fusion = self._latest(session_id) if fusion is None: return None cache_root = _private_directory(self.cache_root) session_cache = _private_child_directory(cache_root, session_id) fusion_cache = _private_child_directory(session_cache, fusion.fusion_id) output = fusion_cache / f"{recording_id}.rrd" sidecar = output.with_suffix(".rrd.cache.json") cached = _read_cache(output, sidecar, fusion, recording_id) if cached is not None: return cached payload = _render_fusion( fusion, application_id=application_id, recording_id=recording_id, ) 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.calibrated-fusion-overlay-cache/v1", "fusion_id": fusion.fusion_id, "recording_id": recording_id, "byte_length": len(payload), "sha256": hashlib.sha256(payload).hexdigest(), }, ) finally: temporary.unlink(missing_ok=True) return payload def _latest(self, session_id: str) -> RecordedCalibratedFusion | None: try: jobs = sorted(self.jobs_root.iterdir()) except FileNotFoundError: return None if len(jobs) > MAX_SCAN: raise RecordedPerceptionOverlayError("compute job catalog is outside bounds") matches: list[RecordedCalibratedFusion] = [] for job_root in jobs: 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 parent = self.fusion_results_root / job.job_id try: candidates = sorted( path for path in parent.iterdir() if path.is_dir() and not path.is_symlink() and _SAFE_FUSION_ID.fullmatch(path.name) is not None ) except FileNotFoundError: continue if len(candidates) > MAX_SCAN: raise RecordedPerceptionOverlayError("fusion result catalog is outside bounds") for candidate in candidates: try: matches.append( validate_recorded_calibrated_fusion( job_root, self.perception_results_root, candidate, ) ) except (OSError, SessionIntegrityError): continue if not matches: return None return max(matches, key=lambda value: (value.created_at_utc, value.fusion_id)) class RecordedPerceptionOverlayMux: """Prefer calibrated full-epoch 3D fusion and retain the legacy fallback.""" def __init__(self, primary: _OverlayProvider, fallback: _OverlayProvider | None) -> None: self.primary = primary self.fallback = fallback def render( self, session_id: str, *, application_id: str, recording_id: str, ) -> bytes | None: materialized = self.materialize( session_id, application_id=application_id, recording_id=recording_id, ) if not isinstance(materialized, RecordedPerceptionOverlayArtifact): return materialized try: payload = materialized.path.read_bytes() except OSError as exc: raise RecordedPerceptionOverlayError( "recorded perception cache became unavailable" ) from exc if ( len(payload) != materialized.byte_length or hashlib.sha256(payload).hexdigest() != materialized.sha256 ): raise RecordedPerceptionOverlayError( "recorded perception cache changed after validation" ) return payload def materialize( self, session_id: str, *, application_id: str, recording_id: str, ) -> RecordedPerceptionOverlayArtifact | bytes | None: primary_materializer = getattr(self.primary, "materialize", None) payload = ( primary_materializer( session_id, application_id=application_id, recording_id=recording_id, ) if callable(primary_materializer) else self.primary.render( session_id, application_id=application_id, recording_id=recording_id, ) ) if payload is not None and not isinstance( payload, (bytes, RecordedPerceptionOverlayArtifact) ): raise RecordedPerceptionOverlayError( "recorded perception provider returned an invalid artifact" ) if payload is not None or self.fallback is None: return payload fallback_materializer = getattr(self.fallback, "materialize", None) if callable(fallback_materializer): fallback_payload = fallback_materializer( session_id, application_id=application_id, recording_id=recording_id, ) else: fallback_payload = self.fallback.render( session_id, application_id=application_id, recording_id=recording_id, ) if fallback_payload is not None and not isinstance( fallback_payload, (bytes, RecordedPerceptionOverlayArtifact) ): raise RecordedPerceptionOverlayError( "recorded perception fallback returned an invalid artifact" ) return fallback_payload def status(self, session_id: str, *, recording_id: str) -> dict[str, Any]: primary_status = getattr(self.primary, "status", None) if callable(primary_status): value = primary_status(session_id, recording_id=recording_id) if isinstance(value, dict): return value return { "state": "idle", "phase": "idle", "elapsed_seconds": 0.0, "byte_length": None, } def validate_recorded_calibrated_fusion( job_root: Path, perception_results_root: Path, fusion_root: Path, ) -> RecordedCalibratedFusion: job = validate_camera_compute_job(job_root) root = fusion_root.expanduser().resolve(strict=True) if not root.is_dir() or _SAFE_FUSION_ID.fullmatch(root.name) is None: raise SessionIntegrityError("calibrated fusion root is invalid") manifest = _read_json_object(root / "manifest.json", root) identity = manifest.get("identity") identity_sha256 = manifest.get("identity_sha256") if ( manifest.get("schema_version") != FUSION_SCHEMA or manifest.get("fusion_id") != root.name or not isinstance(identity, dict) or identity.get("schema_version") != FUSION_IDENTITY_SCHEMA or not isinstance(identity_sha256, str) or _SHA256.fullmatch(identity_sha256) is None or root.name != f"fusion-{identity_sha256}" or hashlib.sha256(_canonical_json(identity)).hexdigest() != identity_sha256 or identity.get("job_id") != job.job_id or identity.get("input_sha256") != job.input_sha256 or manifest.get("session_id") != job.session_id or manifest.get("source_id") != job.source_id or manifest.get("frame_count") != job.segment_count or manifest.get("timeline_start_seconds") != job.timeline_start_seconds or manifest.get("timeline_end_seconds") != job.timeline_end_seconds ): raise SessionIntegrityError("calibrated fusion identity is inconsistent") perception_result_id = identity.get("perception_result_id") if not isinstance(perception_result_id, str): raise SessionIntegrityError("calibrated fusion perception binding is invalid") perception = validate_recorded_perception_epoch_result( job_root, perception_results_root / job.job_id / perception_result_id, ) if ( identity.get("calibration_sha256") != perception.calibration_sha256 or identity.get("camera_slot") != perception.calibration_slot ): raise SessionIntegrityError("calibrated fusion calibration binding changed") artifacts = manifest.get("artifacts") expected = {"fusion.npz", "box-labels.json", "fusion-frames.jsonl"} if not isinstance(artifacts, list) or len(artifacts) != len(expected): raise SessionIntegrityError("calibrated fusion artifacts are incomplete") paths: dict[str, Path] = {} for artifact in artifacts: if not isinstance(artifact, dict) or artifact.get("name") not in expected: raise SessionIntegrityError("calibrated fusion artifact is invalid") name = str(artifact["name"]) if name in paths: raise SessionIntegrityError("calibrated fusion artifact is duplicated") path = root / name metadata = _confined_regular_file(path, root) if ( artifact.get("byte_length") != metadata.st_size or not isinstance(artifact.get("sha256"), str) or _SHA256.fullmatch(str(artifact["sha256"])) is None or _sha256(path) != artifact["sha256"] ): raise SessionIntegrityError("calibrated fusion artifact identity changed") paths[name] = path labels = _read_labels(paths["box-labels.json"]) _validate_arrays(paths["fusion.npz"], job, len(labels)) created_at = manifest.get("created_at_utc") if not isinstance(created_at, str) or not 1 <= len(created_at) <= 64: raise SessionIntegrityError("calibrated fusion creation time is invalid") return RecordedCalibratedFusion( fusion_id=root.name, root=root, job=job, perception_result_id=perception_result_id, created_at_utc=created_at, arrays_path=paths["fusion.npz"], labels_path=paths["box-labels.json"], frame_count=job.segment_count, ) def _validate_arrays(path: Path, job: CameraComputeJob, label_count: int) -> None: try: with np.load(path, allow_pickle=False) as arrays: required = { "frame_times_ns", "point_offsets", "points", "point_colors", "box_offsets", "box_centers", "box_half_sizes", "box_colors", } if set(arrays.files) != required: raise SessionIntegrityError("calibrated fusion array set is invalid") frame_times = arrays["frame_times_ns"] point_offsets = arrays["point_offsets"] points = arrays["points"] point_colors = arrays["point_colors"] box_offsets = arrays["box_offsets"] box_centers = arrays["box_centers"] box_half_sizes = arrays["box_half_sizes"] box_colors = arrays["box_colors"] if ( frame_times.dtype != np.int64 or frame_times.shape != (job.segment_count,) or point_offsets.dtype != np.int64 or point_offsets.shape != (job.segment_count + 1,) or box_offsets.dtype != np.int64 or box_offsets.shape != (job.segment_count + 1,) or points.dtype != np.float32 or points.ndim != 2 or points.shape[1:] != (3,) or point_colors.dtype != np.uint8 or point_colors.shape != points.shape or box_centers.dtype != np.float32 or box_centers.ndim != 2 or box_centers.shape[1:] != (3,) or box_half_sizes.dtype != np.float32 or box_half_sizes.shape != box_centers.shape or box_colors.dtype != np.uint8 or box_colors.shape != (box_centers.shape[0], 4) or label_count != box_centers.shape[0] ): raise SessionIntegrityError("calibrated fusion array shapes changed") _validate_offsets(point_offsets, points.shape[0], MAX_POINTS_PER_FRAME) _validate_offsets(box_offsets, box_centers.shape[0], MAX_BOXES_PER_FRAME) if ( np.any(np.diff(frame_times) <= 0) or frame_times[0] < round(job.timeline_start_seconds * 1e9) - 1_000_000 or frame_times[-1] > round(job.timeline_end_seconds * 1e9) + 1_000_000 or not np.isfinite(points).all() or not np.isfinite(box_centers).all() or not np.isfinite(box_half_sizes).all() or np.any(box_half_sizes <= 0) ): raise SessionIntegrityError("calibrated fusion arrays are inconsistent") except (OSError, ValueError, KeyError) as exc: raise SessionIntegrityError("calibrated fusion arrays are unavailable") from exc def _validate_offsets( offsets: np.ndarray[Any, np.dtype[np.int64]], total: int, maximum: int, ) -> None: if ( offsets[0] != 0 or offsets[-1] != total or np.any(np.diff(offsets) < 0) or np.any(np.diff(offsets) > maximum) ): raise SessionIntegrityError("calibrated fusion offsets are inconsistent") def _render_fusion( fusion: RecordedCalibratedFusion, *, application_id: str, recording_id: str, ) -> bytes: labels = _read_labels(fusion.labels_path) recording = rr.RecordingStream(application_id, recording_id=recording_id) stream = rr.binary_stream(recording) try: recording.log("/world", rr.ViewCoordinates.RIGHT_HAND_Z_UP, static=True) recording.log( "/world/perception/contract", rr.TextDocument( "Factory-calibrated KB4 mask-to-LiDAR diagnostic. Distances and support-gated " "boxes are not ground-truthed or safety accepted." ), static=True, ) with np.load(fusion.arrays_path, allow_pickle=False) as arrays: frame_times = arrays["frame_times_ns"] point_offsets = arrays["point_offsets"] points = arrays["points"] point_colors = arrays["point_colors"] box_offsets = arrays["box_offsets"] centers = arrays["box_centers"] half_sizes = arrays["box_half_sizes"] box_colors = arrays["box_colors"] for index, timestamp in enumerate(frame_times): recording.set_time( SESSION_TIMELINE, duration=np.timedelta64(int(timestamp), "ns"), ) point_start, point_end = int(point_offsets[index]), int(point_offsets[index + 1]) if point_end > point_start: recording.log( "/world/perception/semantic_points", rr.Points3D( points[point_start:point_end], colors=point_colors[point_start:point_end], radii=rr.Radius.ui_points(2.5), ), ) else: recording.log( "/world/perception/semantic_points", rr.Clear(recursive=False), ) box_start, box_end = int(box_offsets[index]), int(box_offsets[index + 1]) if box_end > box_start: recording.log( "/world/perception/boxes3d", rr.Boxes3D( centers=centers[box_start:box_end], half_sizes=half_sizes[box_start:box_end], colors=box_colors[box_start:box_end], labels=labels[box_start:box_end], fill_mode=FillMode.Solid, show_labels=True, ), ) else: recording.log( "/world/perception/boxes3d", rr.Clear(recursive=False), ) payload = stream.read(flush=True, flush_timeout_sec=300.0) except Exception as exc: raise RecordedPerceptionOverlayError("failed to serialize calibrated fusion") from exc finally: with suppress(Exception): recording.disconnect() if payload is None or not payload.startswith(b"RRF2"): raise RecordedPerceptionOverlayError("serialized calibrated fusion is invalid") return payload def _read_cache( output: Path, sidecar: Path, fusion: RecordedCalibratedFusion, recording_id: str, ) -> bytes | None: try: value = _read_json_object(sidecar, sidecar.parent) payload = output.read_bytes() except (OSError, SessionIntegrityError): return None if ( value.get("schema_version") != "missioncore.calibrated-fusion-overlay-cache/v1" or value.get("fusion_id") != fusion.fusion_id or value.get("recording_id") != recording_id or value.get("byte_length") != len(payload) or value.get("sha256") != hashlib.sha256(payload).hexdigest() or not payload.startswith(b"RRF2") ): return None return payload def _read_labels(path: Path) -> list[str]: metadata = _confined_regular_file(path, path.parent) if not 1 <= metadata.st_size <= MAX_JSON_BYTES: raise SessionIntegrityError("calibrated fusion labels are outside bounds") try: value = json.loads(path.read_text(encoding="utf-8")) except (OSError, UnicodeDecodeError, json.JSONDecodeError) as exc: raise SessionIntegrityError("calibrated fusion labels are unavailable") from exc if ( not isinstance(value, list) or any(not isinstance(item, str) or not 1 <= len(item) <= 256 for item in value) ): raise SessionIntegrityError("calibrated fusion labels are invalid") return value def _read_json_object(path: Path, root: Path) -> dict[str, Any]: metadata = _confined_regular_file(path, root) if not 0 < metadata.st_size <= MAX_JSON_BYTES: raise SessionIntegrityError("calibrated fusion JSON is outside bounds") try: value = json.loads(path.read_text(encoding="utf-8-sig")) except (OSError, UnicodeDecodeError, json.JSONDecodeError) as exc: raise SessionIntegrityError("calibrated fusion JSON is unavailable") from exc if not isinstance(value, dict): raise SessionIntegrityError("calibrated fusion 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("calibrated fusion 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("calibrated fusion artifact is not confined") return metadata def _sha256(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() except (TypeError, ValueError) as exc: raise SessionIntegrityError("calibrated fusion identity cannot be encoded") from exc def _private_directory(path: Path) -> Path: path.mkdir(mode=0o700, parents=True, exist_ok=True) if path.is_symlink() or not path.is_dir(): raise RecordedPerceptionOverlayError("fusion cache root is invalid") with suppress(OSError): os.chmod(path, 0o700) return path.resolve(strict=True) def _private_child_directory(root: Path, name: str) -> Path: child = root / name child.mkdir(mode=0o700, exist_ok=True) if child.is_symlink() or not child.is_dir() or child.resolve(strict=True).parent != root: raise RecordedPerceptionOverlayError("fusion cache directory is invalid") with suppress(OSError): os.chmod(child, 0o700) return child.resolve(strict=True)