antioch-py 2.2.4__py3-none-any.whl → 3.0.1__py3-none-any.whl

common/core/rome.py

@@ -0,0 +1,251 @@
+from typing import overload
+
+import requests
+from requests import Response
+
+from common.core.types import ArkReference, AssetReference, TaskRun
+
+
+class RomeError(Exception):
+    """
+    Base error for Rome API operations.
+    """
+
+
+class RomeAuthError(RomeError):
+    """
+    Authentication error when interacting with Rome API.
+    """
+
+
+class RomeNetworkError(RomeError):
+    """
+    Network error when interacting with Rome API.
+    """
+
+
+class RomeClient:
+    """
+    Client for interacting with Rome (Antioch's cloud API).
+
+    Handles task runs, artifact uploads/downloads, and registry operations for Arks and Assets.
+    """
+
+    def __init__(self, api_url: str, token: str):
+        """
+        Initialize the Rome client.
+
+        :param api_url: Base URL for Rome API.
+        :param token: Authentication token.
+        """
+
+        self._api_url = api_url
+        self._token = token
+
+    def create_task_run(
+        self,
+        task_run: TaskRun,
+        upload_mcap: bool = False,
+        upload_bundle: bool = False,
+    ) -> tuple[str | None, str | None]:
+        """
+        Create a task run and optionally get signed URLs for artifact uploads.
+
+        :param task_run: TaskRun model with run data.
+        :param upload_mcap: Whether client will upload an MCAP file.
+        :param upload_bundle: Whether client will upload a bundle file.
+        :return: Tuple of (mcap_upload_url, bundle_upload_url).
+        """
+
+        response = self._send_request(
+            "POST",
+            "/tasks/runs",
+            json={
+                "task_run": task_run.model_dump(mode="json"),
+                "upload_mcap": upload_mcap,
+                "upload_bundle": upload_bundle,
+            },
+        )
+
+        return response.get("mcap_upload_url"), response.get("bundle_upload_url")
+
+    def list_arks(self) -> list[ArkReference]:
+        """
+        List all Arks from Rome registry.
+
+        :return: List of ArkReference objects from remote registry.
+        """
+
+        response = self._send_request("GET", "/ark/list")
+        return [ArkReference(**ark) for ark in response.get("data", [])]
+
+    def pull_ark(self, name: str, version: str, config_output_path: str, asset_output_path: str | None = None) -> bool:
+        """
+        Pull Ark config and optionally asset from Rome via signed URLs.
+
+        :param name: Name of the Ark.
+        :param version: Version of the Ark.
+        :param config_output_path: Path where ark.json should be saved.
+        :param asset_output_path: Path where asset.usdz should be saved (if present).
+        :return: True if an asset was downloaded, False otherwise.
+        """
+
+        if not self._token:
+            raise RomeAuthError("User not authenticated")
+
+        try:
+            # Get ark data with download URLs
+            response = self._send_request("GET", "/ark/get", json={"name": name, "version": version})
+
+            # Download config
+            config_response = requests.get(response["config_download_url"], timeout=60)
+            config_response.raise_for_status()
+            with open(config_output_path, "wb") as f:
+                f.write(config_response.content)
+
+            # Download asset if present and output path provided
+            has_asset = response.get("asset_download_url") is not None
+            if has_asset and asset_output_path:
+                asset_response = requests.get(response["asset_download_url"], timeout=None)
+                asset_response.raise_for_status()
+                with open(asset_output_path, "wb") as f:
+                    f.write(asset_response.content)
+                return True
+
+            return False
+        except requests.exceptions.RequestException as e:
+            raise RomeNetworkError(f"Network error: {e}") from e
+
+    def list_assets(self) -> list[AssetReference]:
+        """
+        List all assets from Rome registry.
+
+        :return: List of AssetReference objects from remote registry.
+        """
+
+        response = self._send_request("GET", "/asset/list")
+        return [AssetReference(**asset) for asset in response.get("data", [])]
+
+    def pull_asset(self, name: str, version: str, output_path: str) -> dict[str, str]:
+        """
+        Pull asset file from Rome registry via signed URL.
+
+        :param name: Name of the asset.
+        :param version: Version of the asset.
+        :param output_path: Path where the file should be saved.
+        :return: Metadata dictionary containing extension, file_size, and modified_time.
+        """
+
+        if not self._token:
+            raise RomeAuthError("User not authenticated")
+
+        try:
+            response = self._send_request("GET", "/asset/pull", params={"name": name, "version": version})
+            print(f"Downloading {name}:{version}...")
+            download_response = requests.get(response["download_url"], timeout=None)
+            download_response.raise_for_status()
+            with open(output_path, "wb") as f:
+                f.write(download_response.content)
+            return {
+                "extension": response.get("extension", ""),
+                "file_size": response.get("file_size", ""),
+                "modified_time": response.get("modified_time", ""),
+            }
+        except requests.exceptions.RequestException as e:
+            raise RomeNetworkError(f"Network error: {e}") from e
+
+    def get_gar_token(self) -> dict:
+        """
+        Get a GAR (Google Artifact Registry) access token for Docker operations.
+
+        :return: Dictionary with registry_host, repository, access_token, and expires_at.
+        """
+
+        response = self._send_request("GET", "/token/gar")
+        return response["data"]
+
+    @overload
+    def _send_request(
+        self,
+        method: str,
+        endpoint: str,
+        json: dict | None = None,
+        params: dict | None = None,
+        return_content: bool = False,
+    ) -> dict: ...
+
+    @overload
+    def _send_request(
+        self,
+        method: str,
+        endpoint: str,
+        json: dict | None = None,
+        params: dict | None = None,
+        return_content: bool = True,
+    ) -> bytes: ...
+
+    def _send_request(
+        self,
+        method: str,
+        endpoint: str,
+        json: dict | None = None,
+        params: dict | None = None,
+        return_content: bool = False,
+    ) -> dict | bytes:
+        """
+        Send a request to Rome API with standardized error handling.
+
+        :param method: HTTP method (GET, POST, etc.).
+        :param endpoint: API endpoint path.
+        :param json: Optional JSON payload.
+        :param params: Optional query parameters.
+        :param return_content: If True, return raw bytes content instead of JSON.
+        :return: Response JSON data or raw content bytes.
+        """
+
+        if not self._token:
+            raise RomeAuthError("User not authenticated")
+
+        try:
+            url = f"{self._api_url}{endpoint}"
+            headers = {"Authorization": f"Bearer {self._token}", "Content-Type": "application/json"}
+            response = requests.request(method, url, json=json, params=params, headers=headers, timeout=30)
+            self._check_response_errors(response)
+            if return_content:
+                return response.content
+
+            try:
+                return response.json()
+            except requests.exceptions.JSONDecodeError as e:
+                raise RomeError(f"Invalid JSON response: {e}") from e
+        except requests.exceptions.RequestException as e:
+            raise RomeNetworkError(f"Network error: {e}") from e
+
+    def _check_response_errors(self, response: Response) -> None:
+        """
+        Check response for HTTP errors and raise appropriate exceptions.
+
+        :param response: HTTP response object.
+        """
+
+        if response.status_code >= 400:
+            error_message = self._extract_error_message(response)
+            if response.status_code < 500:
+                raise RomeError(error_message)
+            raise RomeNetworkError(f"Server error: {error_message}")
+
+    def _extract_error_message(self, response: Response) -> str:
+        """
+        Extract error message from response JSON or return generic message.
+
+        :param response: HTTP response object.
+        :return: Error message string from response or generic HTTP status message.
+        """
+
+        try:
+            data = response.json()
+            if isinstance(data, dict) and "message" in data:
+                return data["message"]
+        except Exception:
+            pass
+        return f"HTTP {response.status_code}"

common/core/telemetry.py

@@ -0,0 +1,176 @@
+import time
+from pathlib import Path
+
+import foxglove
+import zenoh
+from foxglove import Capability, MCAPWriter  # type: ignore[attr-defined]
+
+from common.constants import FOXGLOVE_WEBSOCKET_PORT
+from common.message import (
+    DetectionDistances,
+    FrameTransforms,
+    Image,
+    ImageAnnotations,
+    Log,
+    Message,
+    PlotData,
+    PointCloud,
+    Pose,
+    RadarScan,
+    RangeMap,
+    Vector2,
+    Vector3,
+)
+from common.utils.comms import CommsSession
+
+TYPE_MAP: dict[str, type] = {
+    "antioch/vector2": Vector2,
+    "antioch/vector3": Vector3,
+    "antioch/pose": Pose,
+    "antioch/image": Image,
+    "antioch/image_annotations": ImageAnnotations,
+    "antioch/point_cloud": PointCloud,
+    "antioch/frame_transforms": FrameTransforms,
+    "antioch/log": Log,
+    "antioch/radar_scan": RadarScan,
+    "antioch/range_map": RangeMap,
+    "antioch/plot_data": PlotData,
+    "antioch/detection_distances": DetectionDistances,
+}
+
+
+class TelemetryManager:
+    """
+    Manages Foxglove WebSocket server and MCAP recording.
+
+    The WebSocket server persists across tasks. MCAP recording is per-task.
+    Uses a callback subscriber to process logs from Zenoh and forward to Foxglove.
+    """
+
+    def __init__(self, run_websocket: bool = True) -> None:
+        """
+        Create a new telemetry manager.
+
+        Initializes the Foxglove WebSocket server (if enabled) and starts listening for logs.
+
+        :param run_websocket: Whether to start the WebSocket server. Set to False for
+            headless non-streaming mode to avoid port conflicts.
+        """
+
+        self._comms = CommsSession()
+        self._mcap_writer: MCAPWriter | None = None
+        self._start_timestamp_ns: int = time.time_ns()
+        self._time_offset_ns: int = 0
+        self._last_let_us: int = 0
+        self._run_websocket = run_websocket
+
+        if run_websocket:
+            try:
+                self._server = foxglove.start_server(
+                    name="antioch-telemetry",
+                    host="0.0.0.0",
+                    port=FOXGLOVE_WEBSOCKET_PORT,
+                    capabilities=[Capability.Time],
+                )
+            except RuntimeError as e:
+                if "Address already in use" in str(e):
+                    raise RuntimeError(
+                        f"Foxglove server port {FOXGLOVE_WEBSOCKET_PORT} is already in use. "
+                        f"Another simulation may be running. Stop it first with sim.stop(), "
+                        f"or kill the process: lsof -ti :{FOXGLOVE_WEBSOCKET_PORT} | xargs -r kill"
+                    ) from None
+                raise
+        else:
+            self._server = None
+
+        self._subscriber = self._comms.declare_callback_subscriber("_logs", self._on_log)
+
+    def stop(self) -> None:
+        """
+        Stop the telemetry manager and clean up resources.
+
+        Stops the log subscriber, MCAP recording, and WebSocket server.
+        """
+
+        self._subscriber.undeclare()
+        self.stop_recording()
+        if self._server is not None:
+            self._server.stop()
+        self._comms.close()
+
+    def start_recording(self, mcap_path: str) -> None:
+        """
+        Start recording telemetry to an MCAP file.
+
+        :param mcap_path: Path to save the MCAP file.
+        """
+
+        self.stop_recording()
+        path = Path(mcap_path)
+        path.parent.mkdir(parents=True, exist_ok=True)
+        self._mcap_writer = foxglove.open_mcap(str(path), allow_overwrite=True)
+
+    def stop_recording(self) -> None:
+        """
+        Stop MCAP recording and finalize the file.
+        """
+
+        if self._mcap_writer:
+            self._mcap_writer.close()
+            self._mcap_writer = None
+
+    def reset_time(self) -> None:
+        """
+        Accumulate the current time offset for monotonic timestamps.
+
+        Call this when the simulation is cleared or stopped. This ensures that
+        timestamps continue monotonically increasing across multiple episodes
+        in the same MCAP file.
+        """
+
+        self._time_offset_ns += self._last_let_us * 1000
+        self._last_let_us = 0
+
+    def _on_log(self, sample: zenoh.Sample) -> None:
+        """
+        Callback for incoming log samples.
+        """
+
+        log = Log.unpack(bytes(sample.payload))
+        if log.channel is None:
+            return
+
+        self._last_let_us = log.let_us
+
+        # Calculate monotonically increasing timestamp
+        log_time_ns = self._start_timestamp_ns + self._time_offset_ns + (log.let_us * 1000)
+
+        # Broadcast time if websocket server is running
+        if self._server is not None:
+            self._server.broadcast_time(log_time_ns)
+
+        if log.telemetry is not None:
+            msg_type = Message.extract_type(log.telemetry)
+            self._log_payload(log.channel, log.telemetry, msg_type, log_time_ns)
+            return
+
+        foxglove_log = log.to_foxglove()
+        if foxglove_log is not None:
+            foxglove.log(log.channel, foxglove_log, log_time=log_time_ns)
+
+    def _log_payload(self, channel: str, data: bytes, msg_type: str | None, log_time_ns: int) -> None:
+        """
+        Extract and log a telemetry payload.
+        """
+
+        msg_class = TYPE_MAP.get(msg_type) if msg_type else None
+        if msg_class is None:
+            # Unknown type - log as JSON directly
+            json_data = Message.extract_data_as_json(data)
+            foxglove.log(channel, json_data, log_time=log_time_ns)
+            return
+
+        msg = msg_class.unpack(data)
+        foxglove_msg = msg.to_foxglove()
+        if foxglove_msg is not None:
+            foxglove.log(channel, foxglove_msg, log_time=log_time_ns)

common/core/types.py

@@ -0,0 +1,219 @@
+from datetime import datetime
+from enum import Enum
+
+from common.message import Message
+
+# =============================================================================
+# Task Types
+# =============================================================================
+
+
+class TaskOutcome(str, Enum):
+    """
+    Task outcome status.
+    """
+
+    SUCCESS = "success"
+    FAILURE = "failure"
+    ERROR = "error"
+    TIMEOUT = "timeout"
+    SKIPPED = "skipped"
+
+
+class TaskTriggerSource(str, Enum):
+    """
+    How the task was triggered.
+    """
+
+    MANUAL = "manual"
+    CI = "ci"
+    SCHEDULED = "scheduled"
+    API = "api"
+
+
+class TaskRunner(Message):
+    """
+    Runner/execution environment metadata for a task run.
+    """
+
+    runner_id: str | None = None
+    runner_name: str | None = None
+    environment: dict | None = None
+
+
+class TaskRun(Message):
+    """
+    A single execution of a task with timing, outcome, and results.
+
+    Task runs are grouped by task_name for display.
+    """
+
+    # Identity - run_id is generated server-side, org_id/user_id set from context
+    run_id: str | None = None
+    org_id: str | None = None
+    user_id: str | None = None
+
+    # Timing
+    started_at: datetime
+    completed_at: datetime
+    duration_ms: int | None = None
+
+    # Task identity
+    task_name: str
+    description: str | None = None
+
+    # Outcome and results
+    outcome: TaskOutcome
+    result: dict | None = None
+
+    # User info (display names for search/filtering)
+    user_name: str | None = None
+    user_email: str | None = None
+
+    # Filtering
+    tags: list[str] | None = None
+    parameters: dict | None = None
+
+    # Execution context
+    runner: TaskRunner | None = None
+    trigger_source: TaskTriggerSource | None = None
+
+
+# =============================================================================
+# Registry Types
+# =============================================================================
+
+
+class ArkRegistryMetadata(Message):
+    """
+    Metadata stored in GCS for a published Ark version.
+
+    This is the canonical schema for Ark registry metadata. All fields are
+    populated by Rome during push and validated against this schema.
+    """
+
+    module_count: int
+    image_count: int
+    capability: str
+    has_assets: bool
+    description: str
+    timestamp: str
+    digest: str
+
+    @classmethod
+    def from_dict(cls, data: dict[str, str]) -> "ArkRegistryMetadata":
+        """
+        Parse metadata from GCS blob metadata (all values are strings).
+
+        :param data: Raw metadata dict from GCS.
+        :return: Parsed ArkRegistryMetadata.
+        """
+
+        return cls(
+            module_count=int(data.get("module_count", "0")),
+            image_count=int(data.get("image_count", "0")),
+            capability=data.get("capability", ""),
+            has_assets=data.get("has_assets", "false").lower() == "true",
+            description=data.get("description", ""),
+            timestamp=data.get("timestamp", ""),
+            digest=data.get("digest", ""),
+        )
+
+    def to_gcs_metadata(self) -> dict[str, str]:
+        """
+        Convert to GCS metadata format (all values as strings).
+
+        :return: Dict suitable for GCS blob metadata.
+        """
+
+        return {
+            "module_count": str(self.module_count),
+            "image_count": str(self.image_count),
+            "capability": self.capability,
+            "has_assets": str(self.has_assets).lower(),
+            "description": self.description,
+            "timestamp": self.timestamp,
+            "digest": self.digest,
+        }
+
+
+class ArkVersionReference(Message):
+    """
+    Reference to an Ark version in the remote Ark registry.
+    """
+
+    version: str
+    created_at: str
+    updated_at: str
+    full_path: str
+    size_bytes: int
+    asset_path: str | None = None
+    asset_size_bytes: int | None = None
+    metadata: ArkRegistryMetadata | None = None
+
+
+class ArkReference(Message):
+    """
+    Reference to an Ark in the remote Ark registry.
+    """
+
+    name: str
+    versions: list[ArkVersionReference]
+    created_at: str
+    updated_at: str
+
+    def __str__(self) -> str:
+        """
+        Return a string representation of the ArkReference.
+        """
+
+        versions_str = ", ".join(v.version for v in self.versions)
+        return f"{self.name} [{versions_str}]"
+
+    def __repr__(self) -> str:
+        """
+        Return a detailed representation of the ArkReference.
+        """
+
+        versions_list = [v.version for v in self.versions]
+        return f"ArkReference(name='{self.name}', versions={versions_list})"
+
+
+class AssetVersionReference(Message):
+    """
+    Reference to a specific asset version.
+    """
+
+    version: str
+    full_path: str
+    size_bytes: int
+    created_at: str
+    updated_at: str
+    metadata: dict[str, str] = {}
+
+
+class AssetReference(Message):
+    """
+    Reference to an asset with all its versions.
+    """
+
+    name: str
+    versions: list[AssetVersionReference]
+    created_at: str
+    updated_at: str
+
+    def __str__(self) -> str:
+        """
+        Return a string representation of the AssetReference.
+        """
+
+        versions_str = ", ".join(v.version for v in self.versions)
+        return f"{self.name} [{versions_str}]"
+
+    def __repr__(self) -> str:
+        """
+        Return a detailed representation of the AssetReference.
+        """
+
+        versions_list = [v.version for v in self.versions]
+        return f"AssetReference(name='{self.name}', versions={versions_list})"