Files
TG3/tg3_local_teleop/data_collection.py

1475 lines
56 KiB
Python
Executable File

#!/usr/bin/env python3
"""Session-safe, asynchronous ROS 2 bag data collection helpers.
The module intentionally has no ROS Python dependency. The teleoperation
bridge can call the small, non-blocking API from its control tick while a
worker thread owns all filesystem work and the ``ros2 bag`` subprocess.
"""
from __future__ import annotations
import hashlib
import json
import math
import os
import queue
import shutil
import signal
import subprocess
import threading
import time
import uuid
from collections.abc import Callable, Mapping, Sequence
from dataclasses import dataclass, field
from datetime import datetime, timezone
from pathlib import Path
from types import MappingProxyType
from typing import Any, Literal, Protocol
import yaml
ToggleAction = Literal["start", "stop"]
def _valid_ros_topic(topic: Any) -> bool:
return (
isinstance(topic, str)
and topic.startswith("/")
and topic.strip() == topic
and not any(character.isspace() for character in topic)
)
def left_joystick_pressed(data: Mapping[str, Any]) -> bool | None:
"""Strictly parse the live-frame ``button_joystick.left`` value.
xTELE's live 5003 frame carries this button directly as ``bool`` or the
integer ``0``/``1``. Missing containers and all other values return
``None``. In particular, callers must not interpret malformed input as a
button release.
"""
if not isinstance(data, Mapping):
return None
try:
button_joystick = data["button_joystick"]
if not isinstance(button_joystick, Mapping):
return None
value = button_joystick["left"]
except (KeyError, TypeError):
return None
if type(value) is bool:
return value
if type(value) is int and value in (0, 1):
return bool(value)
return None
class RecordingToggleGate:
"""One-button start/stop gate with long-press and release protection.
Each accepted teleoperation session must call :meth:`new_session`. The
gate then requires a continuous, valid release before accepting its first
hold. A hold toggles once; the same physical press can never toggle a
second time. Missing/malformed input cancels timers and re-enters the
release lock.
"""
def __init__(
self, hold_seconds: float = 1.0, release_seconds: float = 0.5
) -> None:
if not math.isfinite(hold_seconds) or hold_seconds <= 0.0:
raise ValueError("recording hold time must be positive and finite")
if not math.isfinite(release_seconds) or release_seconds <= 0.0:
raise ValueError("recording release time must be positive and finite")
self.hold_seconds = float(hold_seconds)
self.release_seconds = float(release_seconds)
self.session_open = False
self.active = False
self.require_release = True
self.hold_started_at: float | None = None
self.release_started_at: float | None = None
self.input_healthy = False
self.button_pressed: bool | None = None
self.toggle_count = 0
self.last_transition = "initialized"
def new_session(self) -> None:
"""Open a fresh session in the release-locked, inactive state."""
self.session_open = True
self.active = False
self.require_release = True
self.hold_started_at = None
self.release_started_at = None
self.input_healthy = False
self.button_pressed = None
self.last_transition = "new_session"
def end_session(self, reason: str = "session_ended") -> ToggleAction | None:
"""Close the gate and report whether an active recorder must stop."""
action: ToggleAction | None = "stop" if self.active else None
self.session_open = False
self.active = False
self.require_release = True
self.hold_started_at = None
self.release_started_at = None
self.input_healthy = False
self.button_pressed = None
self.last_transition = reason
return action
def force_inactive(self, reason: str = "recorder_stopped") -> None:
"""Reconcile the gate after an automatic or failed recorder stop."""
self.active = False
self.require_release = True
self.hold_started_at = None
self.release_started_at = None
self.input_healthy = False
self.button_pressed = None
self.last_transition = reason
def update(
self,
now: float,
*,
input_healthy: bool,
pressed: bool | None,
) -> ToggleAction | None:
"""Advance the gate and return ``start``/``stop`` only on a toggle."""
if not math.isfinite(now):
raise ValueError("recording gate clock must be finite")
if not self.session_open:
return None
self.input_healthy = bool(input_healthy)
self.button_pressed = pressed if type(pressed) is bool else None
if not input_healthy or pressed is None or type(pressed) is not bool:
self.hold_started_at = None
self.release_started_at = None
self.require_release = True
return None
if self.require_release:
self.hold_started_at = None
if pressed:
self.release_started_at = None
return None
if self.release_started_at is None:
self.release_started_at = now
return None
if now - self.release_started_at >= self.release_seconds:
self.require_release = False
self.release_started_at = None
return None
self.release_started_at = None
if not pressed:
if self.hold_started_at is not None:
# A short press must be followed by another stable release.
self.hold_started_at = None
self.release_started_at = now
self.require_release = True
return None
if self.hold_started_at is None:
self.hold_started_at = now
return None
if now - self.hold_started_at < self.hold_seconds:
return None
self.active = not self.active
self.toggle_count += 1
action: ToggleAction = "start" if self.active else "stop"
self.last_transition = action
self.hold_started_at = None
self.release_started_at = None
self.require_release = True
return action
@property
def state(self) -> str:
if not self.session_open:
return "closed"
if not self.input_healthy or self.button_pressed is None:
return "input_unhealthy"
if self.require_release:
return "awaiting_release"
if self.hold_started_at is not None:
return "holding_stop" if self.active else "holding_start"
return "active" if self.active else "idle"
@dataclass(frozen=True)
class OptionalTopicGroupConfig:
"""Non-fatal completeness and rate checks for an optional sensor group."""
topics: Sequence[str]
minimum_topic_rates_hz: Mapping[str, float] = field(default_factory=dict)
def __post_init__(self) -> None:
topics = (
tuple(self.topics)
if isinstance(self.topics, Sequence)
and not isinstance(self.topics, str)
else ()
)
if not topics:
raise ValueError("optional topic group must contain at least one topic")
if any(not _valid_ros_topic(topic) for topic in topics):
raise ValueError(
"optional topic group topics must be absolute ROS topic names"
)
if len(set(topics)) != len(topics):
raise ValueError("optional topic group topics must not contain duplicates")
if not isinstance(self.minimum_topic_rates_hz, Mapping):
raise ValueError("optional minimum topic rates must be a mapping")
rates: dict[str, float] = {}
for topic, rate in self.minimum_topic_rates_hz.items():
if topic not in topics:
raise ValueError(
"optional minimum-rate topics must be a subset of group topics"
)
if (
isinstance(rate, bool)
or not isinstance(rate, (int, float))
or not math.isfinite(rate)
or rate <= 0.0
):
raise ValueError(
f"optional minimum topic rate for {topic!r} must be "
"positive and finite"
)
rates[topic] = float(rate)
object.__setattr__(self, "topics", topics)
object.__setattr__(
self,
"minimum_topic_rates_hz",
MappingProxyType(rates),
)
@dataclass(frozen=True)
class RecorderConfig:
"""Static configuration for :class:`DataRecorderManager`."""
base_directory: Path | str = Path("/home/nvidia/tg3_data_collection")
topics: Sequence[str] = ()
required_topics: Sequence[str] = ()
minimum_topic_rates_hz: Mapping[str, float] = field(default_factory=dict)
optional_topic_groups: Mapping[str, OptionalTopicGroupConfig] = field(
default_factory=dict
)
retain_failed_episodes: bool = True
minimum_free_bytes: int = 5 * 1024**3
max_duration_seconds: float = 30 * 60.0
poll_interval_seconds: float = 0.1
sigint_timeout_seconds: float = 15.0
kill_timeout_seconds: float = 3.0
ros2_executable: str = "ros2"
validate_bag_info: bool = True
bag_info_timeout_seconds: float = 15.0
def __post_init__(self) -> None:
base = Path(self.base_directory).expanduser()
topics = tuple(self.topics)
required_topics = tuple(self.required_topics)
minimum_topic_rates: dict[str, float] = {}
if not topics:
raise ValueError("at least one recording topic is required")
if any(not _valid_ros_topic(topic) for topic in topics):
raise ValueError("recording topics must be absolute ROS topic names")
if len(set(topics)) != len(topics):
raise ValueError("recording topics must not contain duplicates")
if len(set(required_topics)) != len(required_topics):
raise ValueError("required topics must not contain duplicates")
if any(not _valid_ros_topic(topic) for topic in required_topics):
raise ValueError("required topics must be absolute ROS topic names")
if not set(required_topics).issubset(topics):
raise ValueError("required topics must be a subset of recording topics")
if not isinstance(self.minimum_topic_rates_hz, Mapping):
raise ValueError("minimum topic rates must be a mapping")
for topic, rate in self.minimum_topic_rates_hz.items():
if topic not in required_topics:
raise ValueError(
"minimum-rate topics must be a subset of required topics"
)
if (
isinstance(rate, bool)
or not isinstance(rate, (int, float))
or not math.isfinite(rate)
or rate <= 0.0
):
raise ValueError(
f"minimum topic rate for {topic!r} must be positive and finite"
)
minimum_topic_rates[topic] = float(rate)
if not isinstance(self.optional_topic_groups, Mapping):
raise ValueError("optional topic groups must be a mapping")
optional_groups: dict[str, OptionalTopicGroupConfig] = {}
grouped_topics: set[str] = set()
for name, group in self.optional_topic_groups.items():
if (
not isinstance(name, str)
or not name
or len(name) > 64
or any(
character
not in "-_0123456789ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz"
for character in name
)
):
raise ValueError(
"optional topic group names must contain 1-64 safe characters"
)
if not isinstance(group, OptionalTopicGroupConfig):
raise ValueError(
f"optional topic group {name!r} has an invalid configuration"
)
group_topics = set(group.topics)
if not group_topics.issubset(topics):
raise ValueError(
f"optional topic group {name!r} must be a subset of recording topics"
)
if group_topics.intersection(required_topics):
raise ValueError(
f"optional topic group {name!r} must be disjoint from required topics"
)
overlap = group_topics.intersection(grouped_topics)
if overlap:
raise ValueError(
"optional topic groups must be disjoint; repeated topics: "
+ ", ".join(sorted(overlap))
)
grouped_topics.update(group_topics)
optional_groups[name] = group
if (
isinstance(self.minimum_free_bytes, bool)
or not isinstance(self.minimum_free_bytes, int)
or self.minimum_free_bytes < 0
):
raise ValueError("minimum free bytes must be a non-negative integer")
for name, value in (
("maximum duration", self.max_duration_seconds),
("poll interval", self.poll_interval_seconds),
("SIGINT timeout", self.sigint_timeout_seconds),
("kill timeout", self.kill_timeout_seconds),
("bag info timeout", self.bag_info_timeout_seconds),
):
if not math.isfinite(value) or value <= 0.0:
raise ValueError(f"{name} must be positive and finite")
if not isinstance(self.ros2_executable, str) or not self.ros2_executable:
raise ValueError("ROS 2 executable must be a non-empty string")
if type(self.validate_bag_info) is not bool:
raise ValueError("validate bag info must be a bool")
if type(self.retain_failed_episodes) is not bool:
raise ValueError("retain failed episodes must be a bool")
object.__setattr__(self, "base_directory", base)
object.__setattr__(self, "topics", topics)
object.__setattr__(self, "required_topics", required_topics)
object.__setattr__(
self,
"minimum_topic_rates_hz",
MappingProxyType(minimum_topic_rates),
)
object.__setattr__(
self,
"optional_topic_groups",
MappingProxyType(optional_groups),
)
class RecorderProcess(Protocol):
returncode: int | None
def poll(self) -> int | None: ...
def send_signal(self, sig: int) -> None: ...
def wait(self, timeout: float | None = None) -> int: ...
def terminate(self) -> None: ...
def kill(self) -> None: ...
ProcessFactory = Callable[..., RecorderProcess]
@dataclass(frozen=True)
class _StartRequest:
episode_id: str
session_id: str
reason: str
@dataclass
class _Episode:
request: _StartRequest
active_directory: Path
bag_directory: Path
started_monotonic: float
started_at: datetime
node_name: str = ""
command: tuple[str, ...] = ()
bag_info_validation: dict[str, Any] | None = None
topic_message_counts: dict[str, int] | None = None
required_topic_message_counts: dict[str, int] | None = None
minimum_topic_message_counts: dict[str, int] | None = None
observed_topic_rates_hz: dict[str, float] | None = None
optional_topic_groups: dict[str, Any] | None = None
data_quality_warnings: list[str] | None = None
metadata_duration_nanoseconds: int | None = None
process: RecorderProcess | None = None
stdout_stream: Any = None
stderr_stream: Any = None
class RecordingError(RuntimeError):
"""An episode could not be safely finalized as usable data."""
class DataRecorderManager:
"""Asynchronous owner of one ``ros2 bag record`` process at a time.
:meth:`request_start` and :meth:`request_stop` only update in-memory state
and enqueue work. The worker performs every subprocess and filesystem
operation. Completed episodes move atomically from ``active`` to
``ready``; any startup, process, validation, or finalization failure moves
the preserved episode directory to ``failed``.
"""
def __init__(
self,
config: RecorderConfig,
*,
process_factory: ProcessFactory = subprocess.Popen,
command_runner: Callable[..., Any] = subprocess.run,
free_bytes: Callable[[Path], int] | None = None,
monotonic: Callable[[], float] = time.monotonic,
utc_now: Callable[[], datetime] | None = None,
episode_id_factory: Callable[[], str] | None = None,
) -> None:
self.config = config
self._process_factory = process_factory
self._command_runner = command_runner
self._free_bytes = free_bytes or (
lambda path: int(shutil.disk_usage(path).free)
)
self._monotonic = monotonic
self._utc_now = utc_now or (lambda: datetime.now(timezone.utc))
self._episode_id_factory = episode_id_factory or self._default_episode_id
self._requests: queue.Queue[tuple[str, object | None]] = queue.Queue()
self._condition = threading.Condition(threading.RLock())
self._shutdown_requested = False
self._requested_stop_reason: str | None = None
self._status: dict[str, Any] = {
"state": "idle",
"recording": False,
"episode_id": None,
"session_id": None,
"started_at": None,
"stop_reason": None,
"last_result": None,
"last_episode_directory": None,
"last_error": "",
"topics": list(config.topics),
}
self._worker = threading.Thread(
target=self._run,
name="tg3-data-recorder",
daemon=True,
)
self._worker.start()
@staticmethod
def _default_episode_id() -> str:
stamp = datetime.now(timezone.utc).strftime("%Y%m%dT%H%M%S.%fZ")
return f"{stamp}_{uuid.uuid4().hex[:12]}"
def request_start(
self,
session_id: str,
reason: str = "button",
*,
episode_id: str | None = None,
) -> str | None:
"""Queue a start and return its actual episode ID, or ``None`` if busy."""
self._validate_custom_data_value(session_id, "session ID")
self._validate_text(reason, "start reason")
if episode_id is None:
episode_id = self._episode_id_factory()
self._validate_episode_id(episode_id)
request = _StartRequest(episode_id, session_id, reason)
with self._condition:
if self._shutdown_requested or self._status["state"] != "idle":
return None
self._requested_stop_reason = None
self._status.update(
{
"state": "start_pending",
"recording": False,
"episode_id": episode_id,
"session_id": session_id,
"started_at": None,
"stop_reason": None,
"last_error": "",
}
)
self._condition.notify_all()
self._requests.put(("start", request))
return episode_id
def request_stop(self, reason: str = "operator_button") -> bool:
"""Queue SIGINT shutdown of the active/pending episode."""
self._validate_text(reason, "stop reason")
with self._condition:
if self._status["state"] not in (
"start_pending",
"starting",
"recording",
"stop_pending",
):
return False
if self._requested_stop_reason is None:
self._requested_stop_reason = reason
self._status["state"] = "stop_pending"
self._status["stop_reason"] = self._requested_stop_reason
self._condition.notify_all()
self._requests.put(("stop", reason))
return True
def status(self) -> dict[str, Any]:
"""Return a thread-safe, JSON-serializable status snapshot."""
with self._condition:
snapshot = dict(self._status)
snapshot["topics"] = list(self._status["topics"])
return snapshot
def wait_until_idle(self, timeout: float | None = None) -> bool:
"""Wait explicitly for an idle/terminated worker (never call in tick)."""
deadline = None if timeout is None else self._monotonic() + timeout
with self._condition:
while self._status["state"] not in ("idle", "shutdown"):
remaining = (
None if deadline is None else deadline - self._monotonic()
)
if remaining is not None and remaining <= 0.0:
return False
self._condition.wait(remaining)
return True
def shutdown(self, *, wait: bool = True, timeout: float | None = None) -> bool:
"""Request clean SIGINT shutdown and optionally wait for the worker."""
with self._condition:
if not self._shutdown_requested:
self._shutdown_requested = True
if self._requested_stop_reason is None:
self._requested_stop_reason = "shutdown"
if self._status["state"] not in ("idle", "shutdown"):
self._status["state"] = "stop_pending"
self._status["stop_reason"] = self._requested_stop_reason
self._condition.notify_all()
self._requests.put(("shutdown", None))
if not wait:
return not self._worker.is_alive()
self._worker.join(timeout)
return not self._worker.is_alive()
def _run(self) -> None:
while True:
command, payload = self._requests.get()
if command == "shutdown":
break
if command != "start" or not isinstance(payload, _StartRequest):
continue
self._record_episode(payload)
with self._condition:
should_shutdown = self._shutdown_requested
if should_shutdown:
break
with self._condition:
self._status["state"] = "shutdown"
self._status["recording"] = False
self._status["episode_id"] = None
self._status["session_id"] = None
self._condition.notify_all()
def _record_episode(self, request: _StartRequest) -> None:
episode: _Episode | None = None
stop_reason: str | None = None
result_directory: Path | None = None
try:
with self._condition:
self._status["state"] = "starting"
self._condition.notify_all()
base = Path(self.config.base_directory)
active_root = base / "active"
ready_root = base / "ready"
failed_root = base / "failed"
for directory in (active_root, ready_root, failed_root):
directory.mkdir(parents=True, exist_ok=True)
active_directory = active_root / request.episode_id
active_directory.mkdir(mode=0o750)
episode = _Episode(
request=request,
active_directory=active_directory,
bag_directory=active_directory / "bag",
started_monotonic=self._monotonic(),
started_at=self._as_utc(self._utc_now()),
bag_info_validation={
"enabled": self.config.validate_bag_info,
"passed": None,
"result": "not_run",
"timeout_seconds": self.config.bag_info_timeout_seconds,
},
)
initial_free = int(self._free_bytes(base))
if initial_free < self.config.minimum_free_bytes:
raise RecordingError(
"insufficient free space before recording: "
f"{initial_free} < {self.config.minimum_free_bytes} bytes"
)
episode.stdout_stream = (active_directory / "ros2_bag.stdout.log").open(
"ab", buffering=0
)
episode.stderr_stream = (active_directory / "ros2_bag.stderr.log").open(
"ab", buffering=0
)
command = [
self.config.ros2_executable,
"bag",
"record",
"--storage",
"mcap",
"--storage-preset-profile",
"zstd_fast",
"--max-cache-size",
"67108864",
"--max-bag-duration",
"300",
"--disable-keyboard-controls",
"--node-name",
self._ros_node_name(request.episode_id),
"--output",
str(episode.bag_directory),
"--custom-data",
f"capture_id={request.episode_id}",
f"teleop_session_id={request.session_id}",
"--topics",
*self.config.topics,
]
episode.node_name = command[command.index("--node-name") + 1]
episode.command = tuple(command)
episode.process = self._process_factory(
command,
stdin=subprocess.DEVNULL,
stdout=episode.stdout_stream,
stderr=episode.stderr_stream,
start_new_session=True,
)
with self._condition:
self._status.update(
{
"state": (
"stop_pending"
if self._requested_stop_reason is not None
else "recording"
),
"recording": True,
"started_at": self._iso(episode.started_at),
}
)
self._condition.notify_all()
stop_reason = self._monitor_episode(episode)
self._stop_process(episode.process)
ended_at = self._as_utc(self._utc_now())
ended_monotonic = self._monotonic()
self._close_logs(episode)
with self._condition:
self._status["state"] = "finalizing"
self._status["recording"] = False
self._status["stop_reason"] = stop_reason
self._condition.notify_all()
manifest = self._ready_manifest(
episode,
stop_reason,
ended_at,
ended_monotonic,
)
self._write_manifest(episode.active_directory, manifest)
self._write_ready_marker(episode.active_directory)
self._sync_directory(episode.active_directory)
result_directory = ready_root / request.episode_id
os.replace(episode.active_directory, result_directory)
try:
self._sync_directory(ready_root)
except OSError:
# The same-filesystem rename has already published a complete
# manifest+READY directory. Parent fsync is best-effort here.
pass
self._finish_status("ready", result_directory, "", stop_reason)
except Exception as error:
error_text = f"{type(error).__name__}: {error}"
if episode is not None:
self._abort_process(episode.process)
self._close_logs(episode)
result_directory, failure_error = self._preserve_failed_episode(
episode,
stop_reason or self._requested_stop_reason or "error",
error_text,
)
if failure_error:
error_text += f"; failed to archive episode: {failure_error}"
self._finish_status(
"failed",
result_directory,
error_text,
stop_reason or self._requested_stop_reason or "error",
)
def _monitor_episode(self, episode: _Episode) -> str:
assert episode.process is not None
while True:
returncode = episode.process.poll()
if returncode is not None:
raise RecordingError(
f"ros2 bag exited unexpectedly with return code {returncode}"
)
with self._condition:
requested_reason = self._requested_stop_reason
shutting_down = self._shutdown_requested
if requested_reason is not None:
return requested_reason
if shutting_down:
return "shutdown"
elapsed = self._monotonic() - episode.started_monotonic
if elapsed >= self.config.max_duration_seconds:
return "maximum_duration"
free = int(self._free_bytes(Path(self.config.base_directory)))
if free < self.config.minimum_free_bytes:
return "low_disk_space"
try:
command, payload = self._requests.get(
timeout=self.config.poll_interval_seconds
)
except queue.Empty:
continue
if command == "stop":
return str(payload)
if command == "shutdown":
return "shutdown"
def _stop_process(self, process: RecorderProcess) -> None:
if process.poll() is not None:
return
process.send_signal(signal.SIGINT)
try:
process.wait(timeout=self.config.sigint_timeout_seconds)
return
except subprocess.TimeoutExpired:
process.terminate()
try:
process.wait(timeout=self.config.kill_timeout_seconds)
except subprocess.TimeoutExpired:
process.kill()
process.wait(timeout=self.config.kill_timeout_seconds)
raise RecordingError("ros2 bag did not stop cleanly after SIGINT")
def _abort_process(self, process: RecorderProcess | None) -> None:
if process is None or process.poll() is not None:
return
try:
process.send_signal(signal.SIGINT)
process.wait(timeout=self.config.sigint_timeout_seconds)
return
except Exception:
pass
try:
process.kill()
process.wait(timeout=self.config.kill_timeout_seconds)
except Exception:
pass
def _ready_manifest(
self,
episode: _Episode,
stop_reason: str,
ended_at: datetime,
ended_monotonic: float,
) -> dict[str, Any]:
metadata = episode.bag_directory / "metadata.yaml"
if not metadata.is_file() or metadata.stat().st_size <= 0:
raise RecordingError("bag metadata.yaml is missing or empty")
try:
metadata_text = metadata.read_text(encoding="utf-8")
except UnicodeDecodeError as error:
raise RecordingError("bag metadata.yaml is not valid UTF-8") from error
if "mcap" not in metadata_text.lower():
raise RecordingError("bag metadata.yaml does not identify MCAP storage")
required_topic_counts = self._validate_required_topic_counts(
episode, metadata_text
)
mcap_files = sorted(
path
for path in episode.bag_directory.rglob("*.mcap")
if path.is_file()
)
if not mcap_files:
raise RecordingError("bag contains no MCAP file")
empty_mcap = [path.name for path in mcap_files if path.stat().st_size <= 0]
if empty_mcap:
raise RecordingError(
"bag contains empty MCAP file(s): " + ", ".join(empty_mcap)
)
bag_info_validation = self._validate_bag_info(episode)
artifacts = [metadata, *mcap_files]
files = [
self._file_record(episode.active_directory, path)
for path in artifacts
]
return {
"schema_version": 1,
"state": "complete",
"status": "ready",
"episode_id": episode.request.episode_id,
"session_id": episode.request.session_id,
"start_reason": episode.request.reason,
"stop_reason": stop_reason,
"topics": list(self.config.topics),
"topic_message_counts": episode.topic_message_counts,
"required_topics": list(self.config.required_topics),
"required_topic_message_counts": required_topic_counts,
"minimum_topic_rates_hz": dict(
self.config.minimum_topic_rates_hz
),
"minimum_topic_message_counts": (
episode.minimum_topic_message_counts
),
"observed_topic_rates_hz": episode.observed_topic_rates_hz,
"optional_topic_groups": episode.optional_topic_groups,
"data_quality_warnings": episode.data_quality_warnings,
"metadata_duration_nanoseconds": (
episode.metadata_duration_nanoseconds
),
"storage_id": "mcap",
"custom_data": {
"capture_id": episode.request.episode_id,
"teleop_session_id": episode.request.session_id,
},
"bag_info_validation": bag_info_validation,
"rosbag2": {
"node_name": episode.node_name,
"storage_preset_profile": "zstd_fast",
"max_cache_size_bytes": 67_108_864,
"max_bag_duration_seconds": 300,
},
"started_at_utc": self._iso(episode.started_at),
"stopped_at_utc": self._iso(ended_at),
"duration_seconds": max(
0.0, ended_monotonic - episode.started_monotonic
),
"total_size_bytes": sum(record["size_bytes"] for record in files),
"files": files,
}
def _validate_required_topic_counts(
self, episode: _Episode, metadata_text: str
) -> dict[str, int]:
"""Require positive rosbag2 metadata counts for configured topics."""
try:
metadata = yaml.safe_load(metadata_text)
except yaml.YAMLError as error:
raise RecordingError("bag metadata.yaml is not valid YAML") from error
if not isinstance(metadata, Mapping):
raise RecordingError("bag metadata.yaml root must be a mapping")
information = metadata.get("rosbag2_bagfile_information")
if not isinstance(information, Mapping):
raise RecordingError(
"bag metadata.yaml lacks rosbag2_bagfile_information"
)
if information.get("storage_identifier") != "mcap":
raise RecordingError(
"bag metadata.yaml storage_identifier is not mcap"
)
entries = information.get("topics_with_message_count")
if not isinstance(entries, list):
if self.config.required_topics:
raise RecordingError(
"bag metadata.yaml lacks topics_with_message_count"
)
episode.topic_message_counts = {
topic: 0 for topic in self.config.topics
}
episode.required_topic_message_counts = {}
return {}
all_counts: dict[str, int] = {}
for index, entry in enumerate(entries):
if not isinstance(entry, Mapping):
raise RecordingError(
f"metadata topic-count entry {index} is not a mapping"
)
topic_metadata = entry.get("topic_metadata")
if not isinstance(topic_metadata, Mapping):
raise RecordingError(
f"metadata topic-count entry {index} lacks topic_metadata"
)
name = topic_metadata.get("name")
message_count = entry.get("message_count")
if not isinstance(name, str) or not name.startswith("/"):
raise RecordingError(
f"metadata topic-count entry {index} has an invalid name"
)
if (
isinstance(message_count, bool)
or not isinstance(message_count, int)
or message_count < 0
):
raise RecordingError(
f"metadata topic {name!r} has an invalid message_count"
)
all_counts[name] = all_counts.get(name, 0) + message_count
required_counts = {
topic: all_counts.get(topic, 0)
for topic in self.config.required_topics
}
episode.topic_message_counts = {
topic: all_counts.get(topic, 0) for topic in self.config.topics
}
episode.required_topic_message_counts = required_counts
self._observe_optional_topic_groups(episode, information, all_counts)
missing = [
topic for topic in self.config.required_topics if topic not in all_counts
]
empty = [
topic
for topic in self.config.required_topics
if topic in all_counts and all_counts[topic] <= 0
]
problems: list[str] = []
if missing:
problems.append("missing required topics: " + ", ".join(missing))
if empty:
problems.append("zero-message required topics: " + ", ".join(empty))
if problems:
raise RecordingError("; ".join(problems))
self._validate_minimum_topic_rates(episode, information, all_counts)
return required_counts
def _observe_optional_topic_groups(
self,
episode: _Episode,
information: Mapping[str, Any],
all_counts: Mapping[str, int],
) -> None:
"""Classify optional sensor quality without invalidating core data."""
duration_nanoseconds: int | None = None
duration = information.get("duration")
if isinstance(duration, Mapping):
candidate = duration.get("nanoseconds")
if (
type(candidate) is int
and candidate > 0
):
duration_nanoseconds = candidate
episode.metadata_duration_nanoseconds = candidate
duration_seconds = (
None
if duration_nanoseconds is None
else duration_nanoseconds / 1_000_000_000.0
)
observations: dict[str, Any] = {}
warnings: list[str] = []
for name, group in self.config.optional_topic_groups.items():
counts = {
topic: int(all_counts.get(topic, 0)) for topic in group.topics
}
empty_topics = [
topic for topic, count in counts.items() if count <= 0
]
all_absent = len(empty_topics) == len(group.topics)
minimum_counts: dict[str, int | None] = {}
observed_rates: dict[str, float | None] = {}
below_rate: list[str] = []
for topic, minimum_rate in group.minimum_topic_rates_hz.items():
if duration_seconds is None:
minimum_counts[topic] = None
observed_rates[topic] = None
if not all_absent:
below_rate.append(topic)
continue
minimum_count = math.floor(duration_seconds * minimum_rate)
observed_rate = counts[topic] / duration_seconds
minimum_counts[topic] = minimum_count
observed_rates[topic] = observed_rate
if not all_absent and counts[topic] < minimum_count:
below_rate.append(topic)
if all_absent:
state = "absent"
elif empty_topics:
state = "partial"
warnings.append(
f"optional topic group {name!r} is partial; zero-message "
"topics: " + ", ".join(empty_topics)
)
elif below_rate:
state = "low_rate"
if duration_seconds is None:
warnings.append(
f"optional topic group {name!r} rate could not be "
"validated because bag duration is unavailable"
)
else:
details = [
f"{topic}={observed_rates[topic]:.3f}Hz<"
f"{group.minimum_topic_rates_hz[topic]:.3f}Hz"
for topic in below_rate
]
warnings.append(
f"optional topic group {name!r} is below its observed "
"minimum rate: " + "; ".join(details)
)
else:
state = "healthy"
observations[name] = {
"state": state,
"topics": list(group.topics),
"topic_message_counts": counts,
"minimum_topic_rates_hz": dict(
group.minimum_topic_rates_hz
),
"minimum_topic_message_counts": minimum_counts,
"observed_topic_rates_hz": observed_rates,
"zero_message_topics": empty_topics,
"below_minimum_rate_topics": below_rate,
}
episode.optional_topic_groups = observations
episode.data_quality_warnings = warnings
def _validate_minimum_topic_rates(
self,
episode: _Episode,
information: Mapping[str, Any],
all_counts: Mapping[str, int],
) -> None:
configured = self.config.minimum_topic_rates_hz
if not configured:
episode.minimum_topic_message_counts = {}
episode.observed_topic_rates_hz = {}
return
duration = information.get("duration")
if not isinstance(duration, Mapping):
raise RecordingError(
"bag metadata.yaml lacks duration for minimum-rate validation"
)
nanoseconds = duration.get("nanoseconds")
if (
isinstance(nanoseconds, bool)
or not isinstance(nanoseconds, int)
or nanoseconds <= 0
):
raise RecordingError(
"bag metadata.yaml duration.nanoseconds must be a positive integer"
)
episode.metadata_duration_nanoseconds = nanoseconds
minimum_counts: dict[str, int] = {}
observed_rates: dict[str, float] = {}
failures: list[str] = []
duration_seconds = nanoseconds / 1_000_000_000.0
for topic, minimum_rate in configured.items():
count = all_counts[topic]
minimum_count = math.floor(duration_seconds * minimum_rate)
observed_rate = count / duration_seconds
minimum_counts[topic] = minimum_count
observed_rates[topic] = observed_rate
if count < minimum_count:
failures.append(
f"{topic}: {count} messages < {minimum_count} required "
f"over {duration_seconds:.6f}s "
f"({observed_rate:.3f} Hz < {minimum_rate:.3f} Hz)"
)
episode.minimum_topic_message_counts = minimum_counts
episode.observed_topic_rates_hz = observed_rates
if failures:
raise RecordingError(
"required topics below minimum average rate: "
+ "; ".join(failures)
)
def _validate_bag_info(self, episode: _Episode) -> dict[str, Any]:
"""Run rosbag2's reader-level validation before publishing READY."""
if not self.config.validate_bag_info:
validation = {
"enabled": False,
"passed": None,
"result": "disabled",
"timeout_seconds": self.config.bag_info_timeout_seconds,
}
episode.bag_info_validation = validation
return validation
command = [
self.config.ros2_executable,
"bag",
"info",
str(episode.bag_directory),
]
validation: dict[str, Any] = {
"enabled": True,
"passed": False,
"result": "running",
"returncode": None,
"timeout_seconds": self.config.bag_info_timeout_seconds,
"validated_at_utc": self._iso(self._as_utc(self._utc_now())),
}
episode.bag_info_validation = validation
try:
result = self._command_runner(
command,
stdin=subprocess.DEVNULL,
capture_output=True,
text=True,
timeout=self.config.bag_info_timeout_seconds,
check=False,
)
except subprocess.TimeoutExpired as error:
validation["result"] = "timeout"
self._write_bag_info_logs(
episode,
self._output_text(error.stdout),
self._output_text(error.stderr),
)
raise RecordingError(
"ros2 bag info timed out after "
f"{self.config.bag_info_timeout_seconds:.3f}s"
) from error
except Exception as error:
validation["result"] = "invocation_error"
validation["error_type"] = type(error).__name__
raise RecordingError(
f"could not run ros2 bag info: {type(error).__name__}: {error}"
) from error
returncode = getattr(result, "returncode", None)
stdout = self._output_text(getattr(result, "stdout", ""))
stderr = self._output_text(getattr(result, "stderr", ""))
self._write_bag_info_logs(episode, stdout, stderr)
if type(returncode) is not int:
validation["result"] = "invalid_runner_result"
raise RecordingError("ros2 bag info returned no integer return code")
validation["returncode"] = returncode
if returncode != 0:
validation["result"] = "nonzero_exit"
detail = stderr.strip() or stdout.strip() or "no diagnostic output"
if len(detail) > 512:
detail = detail[:512] + "..."
raise RecordingError(
f"ros2 bag info failed with return code {returncode}: {detail}"
)
validation["passed"] = True
validation["result"] = "passed"
return validation
@staticmethod
def _output_text(value: object) -> str:
if value is None:
return ""
if isinstance(value, bytes):
return value.decode("utf-8", errors="replace")
return str(value)
@staticmethod
def _write_bag_info_logs(
episode: _Episode, stdout: str, stderr: str
) -> None:
# These diagnostics are intentionally not included in the data hashes;
# only metadata.yaml and MCAP payloads define a synchronized episode.
try:
(episode.active_directory / "bag_info.stdout.log").write_text(
stdout, encoding="utf-8"
)
(episode.active_directory / "bag_info.stderr.log").write_text(
stderr, encoding="utf-8"
)
except OSError:
pass
def _preserve_failed_episode(
self,
episode: _Episode,
stop_reason: str,
error_text: str,
) -> tuple[Path | None, str]:
if not episode.active_directory.exists():
return None, "active episode directory is missing"
if not self.config.retain_failed_episodes:
try:
active_root = Path(self.config.base_directory) / "active"
target = episode.active_directory
if (
target.parent != active_root
or target.name != episode.request.episode_id
or target.is_symlink()
or active_root.is_symlink()
):
raise RecordingError(
"refusing to discard a failed episode outside its "
"fixed active root"
)
shutil.rmtree(target)
try:
self._sync_directory(active_root)
except OSError:
pass
return None, ""
except Exception as discard_error:
return None, (
"failed to discard project-owned failed episode: "
f"{type(discard_error).__name__}: {discard_error}"
)
try:
ready_marker = episode.active_directory / "READY"
if ready_marker.exists():
# Preserve, but invalidate, a marker if the final directory
# rename itself failed after the marker was written.
os.replace(ready_marker, episode.active_directory / "READY.invalid")
failed_manifest = {
"schema_version": 1,
"state": "failed",
"status": "failed",
"episode_id": episode.request.episode_id,
"session_id": episode.request.session_id,
"start_reason": episode.request.reason,
"stop_reason": stop_reason,
"topics": list(self.config.topics),
"topic_message_counts": episode.topic_message_counts,
"required_topics": list(self.config.required_topics),
"required_topic_message_counts": (
episode.required_topic_message_counts
),
"minimum_topic_rates_hz": dict(
self.config.minimum_topic_rates_hz
),
"minimum_topic_message_counts": (
episode.minimum_topic_message_counts
),
"observed_topic_rates_hz": episode.observed_topic_rates_hz,
"optional_topic_groups": episode.optional_topic_groups,
"data_quality_warnings": episode.data_quality_warnings,
"metadata_duration_nanoseconds": (
episode.metadata_duration_nanoseconds
),
"storage_id": "mcap",
"custom_data": {
"capture_id": episode.request.episode_id,
"teleop_session_id": episode.request.session_id,
},
"bag_info_validation": episode.bag_info_validation,
"started_at_utc": self._iso(episode.started_at),
"failed_at_utc": self._iso(self._as_utc(self._utc_now())),
"duration_seconds": max(
0.0, self._monotonic() - episode.started_monotonic
),
"error": error_text,
}
self._write_manifest(episode.active_directory, failed_manifest)
failed_directory = (
Path(self.config.base_directory)
/ "failed"
/ episode.request.episode_id
)
os.replace(episode.active_directory, failed_directory)
return failed_directory, ""
except Exception as archive_error:
return None, f"{type(archive_error).__name__}: {archive_error}"
def _finish_status(
self,
result: str,
directory: Path | None,
error: str,
stop_reason: str,
) -> None:
with self._condition:
self._status.update(
{
"state": "idle",
"recording": False,
"episode_id": None,
"session_id": None,
"started_at": None,
"stop_reason": stop_reason,
"last_result": result,
"last_episode_directory": (
None if directory is None else str(directory)
),
"last_error": error,
}
)
self._requested_stop_reason = None
self._condition.notify_all()
@staticmethod
def _close_logs(episode: _Episode) -> None:
for name in ("stdout_stream", "stderr_stream"):
stream = getattr(episode, name)
if stream is not None:
try:
stream.close()
except Exception:
# Recording data remains authoritative; a diagnostic-log
# close failure must not strand the episode in active/.
pass
finally:
setattr(episode, name, None)
@staticmethod
def _file_record(root: Path, path: Path) -> dict[str, Any]:
digest = hashlib.sha256()
with path.open("rb") as stream:
for chunk in iter(lambda: stream.read(1024 * 1024), b""):
digest.update(chunk)
size = path.stat().st_size
return {
"path": path.relative_to(root).as_posix(),
# ``size`` is the on-wire sync protocol field. ``size_bytes`` is
# kept explicit for humans and status/manifest consumers.
"size": size,
"size_bytes": size,
"sha256": digest.hexdigest(),
}
@staticmethod
def _write_manifest(directory: Path, manifest: Mapping[str, Any]) -> None:
temporary = directory / f".manifest.{uuid.uuid4().hex}.tmp"
final = directory / "manifest.json"
with temporary.open("w", encoding="utf-8") as stream:
json.dump(manifest, stream, ensure_ascii=False, indent=2, sort_keys=True)
stream.write("\n")
stream.flush()
os.fsync(stream.fileno())
os.replace(temporary, final)
@staticmethod
def _write_ready_marker(directory: Path) -> None:
"""Publish READY last inside active before the atomic directory move."""
temporary = directory / f".READY.{uuid.uuid4().hex}.tmp"
final = directory / "READY"
with temporary.open("w", encoding="ascii") as stream:
stream.write("ready\n")
stream.flush()
os.fsync(stream.fileno())
os.replace(temporary, final)
@staticmethod
def _sync_directory(directory: Path) -> None:
descriptor = os.open(directory, os.O_RDONLY | getattr(os, "O_DIRECTORY", 0))
try:
os.fsync(descriptor)
finally:
os.close(descriptor)
@staticmethod
def _ros_node_name(episode_id: str) -> str:
digest = hashlib.sha256(episode_id.encode("utf-8")).hexdigest()[:16]
return f"tg3_data_recorder_{digest}"
@staticmethod
def _validate_text(value: str, name: str) -> None:
if not isinstance(value, str) or not value.strip() or len(value) > 512:
raise ValueError(
f"{name} must be a non-empty string up to 512 characters"
)
@staticmethod
def _validate_custom_data_value(value: str, name: str) -> None:
allowed = (
"-_.0123456789"
"ABCDEFGHIJKLMNOPQRSTUVWXYZ"
"abcdefghijklmnopqrstuvwxyz"
)
if (
not isinstance(value, str)
or not value
or len(value) > 128
or any(character not in allowed for character in value)
):
raise ValueError(
f"{name} must contain 1-128 safe custom-data characters"
)
@staticmethod
def _validate_episode_id(value: str) -> None:
allowed = (
"-_.0123456789"
"ABCDEFGHIJKLMNOPQRSTUVWXYZ"
"abcdefghijklmnopqrstuvwxyz"
)
first_allowed = (
"0123456789ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz"
)
if (
not isinstance(value, str)
or not value
or len(value) > 128
or value[0] not in first_allowed
or value in (".", "..")
or Path(value).name != value
or any(character not in allowed for character in value)
):
raise ValueError("episode ID must be one safe path component")
@staticmethod
def _as_utc(value: datetime) -> datetime:
if not isinstance(value, datetime):
raise TypeError("UTC clock must return datetime")
if value.tzinfo is None:
raise ValueError("UTC clock must return a timezone-aware datetime")
return value.astimezone(timezone.utc)
@staticmethod
def _iso(value: datetime) -> str:
return value.astimezone(timezone.utc).isoformat().replace("+00:00", "Z")
__all__ = [
"DataRecorderManager",
"OptionalTopicGroupConfig",
"RecorderConfig",
"RecordingError",
"RecordingToggleGate",
"ToggleAction",
"left_joystick_pressed",
]