Files
NODEDC_MISSION_CORE/src/k1link/compute/results.py
T

657 lines
24 KiB
Python

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))
)