1141 lines
42 KiB
Python
Executable File
1141 lines
42 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
|
|
from datetime import datetime, timezone
|
|
from pathlib import Path
|
|
from typing import Any, Literal, Protocol
|
|
|
|
import yaml
|
|
|
|
|
|
ToggleAction = Literal["start", "stop"]
|
|
|
|
|
|
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 RecorderConfig:
|
|
"""Static configuration for :class:`DataRecorderManager`."""
|
|
|
|
base_directory: Path | str = Path("/home/nvidia/tg3_data_collection")
|
|
topics: Sequence[str] = ()
|
|
required_topics: Sequence[str] = ()
|
|
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)
|
|
if not topics:
|
|
raise ValueError("at least one recording topic is required")
|
|
if any(
|
|
not isinstance(topic, str)
|
|
or not topic.startswith("/")
|
|
or topic.strip() != topic
|
|
or any(character.isspace() for character in 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 isinstance(topic, str)
|
|
or not topic.startswith("/")
|
|
or topic.strip() != topic
|
|
or any(character.isspace() for character in 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 (
|
|
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")
|
|
object.__setattr__(self, "base_directory", base)
|
|
object.__setattr__(self, "topics", topics)
|
|
object.__setattr__(self, "required_topics", required_topics)
|
|
|
|
|
|
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
|
|
required_topic_message_counts: dict[str, 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),
|
|
"required_topics": list(self.config.required_topics),
|
|
"required_topic_message_counts": required_topic_counts,
|
|
"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.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.required_topic_message_counts = required_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))
|
|
return required_counts
|
|
|
|
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"
|
|
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),
|
|
"required_topics": list(self.config.required_topics),
|
|
"required_topic_message_counts": (
|
|
episode.required_topic_message_counts
|
|
),
|
|
"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",
|
|
"RecorderConfig",
|
|
"RecordingError",
|
|
"RecordingToggleGate",
|
|
"ToggleAction",
|
|
"left_joystick_pressed",
|
|
]
|