antioch-py 2.0.6__py3-none-any.whl → 3.0.12__py3-none-any.whl

antioch/session/ark.py

@@ -1,504 +0,0 @@
-import os
-from typing import TYPE_CHECKING, cast
-
-from antioch.module.token import TokenType
-from antioch.session.error import SessionArkError, SessionAssetError
-from antioch.session.record import NodeOutputRecorder
-from antioch.session.session import SessionContainer
-from antioch.session.views.articulation import Articulation
-from antioch.session.views.camera import Camera
-from antioch.session.views.imu import Imu
-from antioch.session.views.pir_sensor import PirSensor
-from antioch.session.views.radar import Radar
-from common.ark import Ark as ArkDefinition, ArkReference, Environment, HardwareAccessMode
-from common.ark.hardware import ActuatorGroupHardware, CameraHardware, HardwareType, ImuHardware, PirHardware, RadarHardware
-from common.ark.scheduler import NodeCompleteEvent, NodeStartEvent, OnlineScheduler
-from common.ark.sim import SimNodeComplete, SimNodeStart
-from common.core import (
-    Agent,
-    ArkStateResponse,
-    ContainerSource,
-    get_ark_version_reference,
-    list_local_arks,
-    list_remote_arks,
-    load_local_ark,
-    pull_remote_ark,
-)
-from common.message import JointStates, JointTargets, Pose
-from common.utils.comms import CommsAsyncSubscriber, CommsPublisher
-from common.utils.time import now_us, us_to_s
-from common.utils.usd import sanitize_usd_path
-
-if TYPE_CHECKING:
-    from antioch.session.scene import Scene
-
-NODE_START_PUBLISHER_PATH = "_ark/node_start/{module}/{node}"
-NODE_COMPLETE_SUBSCRIBER_PATH = "_ark/node_complete/{module}/{node}"
-
-
-class Ark(SessionContainer):
-    """
-    Complete runtime for an Ark in simulation.
-
-    Manages entire Ark lifecycle from loading through execution. The runtime is designed to be
-    robust and deterministic:
-    - Events/time increase monotonically via the deterministic scheduler
-    - Each execution is processed exactly once via execution tracking
-    - Multiple nodes can execute at the same LET
-
-    Example:
-        scene = Scene()
-
-        # Build and start Ark
-        ark = scene.add_ark(
-            name="my_robot",
-            version="1.0.0",
-            debug=True
-        )
-
-        # Step simulation
-        scene.step(dt=0.01)  # Delegates to ark.step()
-    """
-
-    def __init__(
-        self,
-        path: str,
-        scene: "Scene",
-        name: str,
-        version: str,
-        world_pose: Pose | None = None,
-        local_pose: Pose | None = None,
-        source: ContainerSource = ContainerSource.LOCAL,
-        debug: bool = False,
-        timeout: float = 30.0,
-    ):
-        """
-        Initialize and build complete Ark runtime.
-
-        :param path: USD path where the Ark will be built.
-        :param scene: Scene instance this Ark belongs to.
-        :param name: Name of the Ark.
-        :param version: Version of the Ark.
-        :param world_pose: Optional world pose.
-        :param local_pose: Optional local pose.
-        :param source: Container image source.
-        :param debug: Enable debug mode.
-        :param timeout: Timeout in seconds for the Ark to start.
-        """
-
-        super().__init__()
-        self._path = path
-        self._scene = scene
-        self._agent = Agent()
-
-        # Load Ark definition
-        self._ark_def = load_local_ark(name, version)
-        self._ark_ref = get_ark_version_reference(name, version)
-        self._base_path = f"{path}/{sanitize_usd_path(self._ark_def.name)}"
-
-        # Start Ark via agent
-        self._agent.start_ark(
-            ark=self._ark_def,
-            source=source,
-            environment=Environment.SIM,
-            debug=debug,
-            timeout=timeout,
-        )
-
-        # Initialize scheduler
-        self._scheduler = OnlineScheduler(self._ark_def.edges, self._ark_def.modules)
-        self._next_event = self._scheduler.next()
-        self._node_start_times_us: dict[tuple[str, str], int] = {}
-
-        # Create publishers and subscribers for node coordination
-        self._node_start_publishers: dict[tuple[str, str], CommsPublisher] = {}
-        self._node_complete_subscribers: dict[tuple[str, str], CommsAsyncSubscriber] = {}
-        for module in self._ark_def.modules:
-            for node_name in module.nodes:
-                start_path = NODE_START_PUBLISHER_PATH.format(module=module.name, node=node_name)
-                complete_path = NODE_COMPLETE_SUBSCRIBER_PATH.format(module=module.name, node=node_name)
-                self._node_start_publishers[(module.name, node_name)] = self._session.comms.declare_publisher(start_path)
-                self._node_complete_subscribers[(module.name, node_name)] = self._session.comms.declare_async_subscriber(complete_path)
-
-        # Build in scene
-        self._load_ark_asset(world_pose=world_pose, local_pose=local_pose)
-        self._build_hardware()
-
-    @property
-    def path(self) -> str:
-        """
-        Get the USD path where this Ark is built.
-
-        :return: The USD path of this Ark.
-        """
-
-        return self._path
-
-    @property
-    def definition(self) -> ArkDefinition:
-        """
-        Get the Ark definition.
-
-        :return: The Ark definition.
-        """
-
-        return self._ark_def
-
-    @property
-    def name(self) -> str:
-        """
-        Get the Ark name.
-
-        :return: The Ark name.
-        """
-
-        return self._ark_def.name
-
-    @property
-    def state(self) -> ArkStateResponse:
-        """
-        Get current state of all modules.
-
-        :return: Ark state with module information.
-        """
-
-        return self._agent.get_ark_state()
-
-    @staticmethod
-    def list_local() -> list[ArkReference]:
-        """
-        List all locally available Arks.
-
-        :return: List of ArkReference objects from local storage.
-        """
-
-        return list_local_arks()
-
-    @staticmethod
-    def list_remote() -> list[ArkReference]:
-        """
-        List all Arks from remote registry.
-
-        Requires authentication. Call session.login() first if not authenticated.
-
-        :return: List of ArkReference objects from remote registry.
-        :raises SessionAuthError: If not authenticated.
-        """
-
-        return list_remote_arks()
-
-    @staticmethod
-    def pull(name: str, version: str, overwrite: bool = False) -> ArkDefinition:
-        """
-        Pull an Ark from remote registry to local storage.
-
-        Requires authentication. Call session.login() first if not authenticated.
-        If the Ark already exists locally, returns the existing Ark unless overwrite=True.
-
-        :param name: Name of the Ark.
-        :param version: Version of the Ark.
-        :param overwrite: Overwrite local Ark if it already exists.
-        :return: The loaded Ark definition.
-        :raises SessionAuthError: If not authenticated.
-        """
-
-        return pull_remote_ark(name=name, version=version, overwrite=overwrite)
-
-    def step(self, dt_us: int = 1_000_000) -> None:
-        """
-        Step the Ark simulation forward by exactly dt_us microseconds.
-
-        :param dt_us: Amount of time to step in microseconds (default 1 second).
-        """
-
-        current_sim_time_us = self._scene.time_us
-        target_sim_time_us = current_sim_time_us + dt_us
-        while current_sim_time_us < target_sim_time_us:
-            event = self._next_event
-
-            # Step physics directly to target time and break
-            if event.let_us > target_sim_time_us:
-                self._scene._step_physics(target_sim_time_us - current_sim_time_us)
-                break
-
-            # Step physics to event time if needed
-            if event.let_us > current_sim_time_us:
-                self._scene._step_physics(event.let_us - current_sim_time_us)
-                current_sim_time_us = self._scene.time_us
-
-            # Node start: read hardware for node and send non-blocking start to nodes
-            if isinstance(event, NodeStartEvent):
-                hardware_reads = self._read_node_hardware(event.module, event.node)
-                self._node_start_times_us[(event.module, event.node)] = now_us()
-                self._node_start_publishers[(event.module, event.node)].publish(
-                    SimNodeStart(
-                        module_name=event.module,
-                        node_name=event.node,
-                        start_let_us=event.start_let_us,
-                        start_timestamp_us=self._node_start_times_us[(event.module, event.node)],
-                        input_tokens=event.input_tokens,
-                        hardware_reads=hardware_reads,
-                    )
-                )
-
-            # Node complete: wait for completion message or budget to elapse and write hardware for node
-            elif isinstance(event, NodeCompleteEvent):
-                node_config = next(m for m in self._ark_def.modules if m.name == event.module).nodes[event.node]
-                node_complete_subscriber = self._node_complete_subscribers[(event.module, event.node)]
-                target_time_us = self._node_start_times_us[(event.module, event.node)] + node_config.budget_us
-                while True:
-                    remaining_time_s = us_to_s(max(0, target_time_us - now_us()))
-                    complete = node_complete_subscriber.recv_timeout(SimNodeComplete, timeout=remaining_time_s)
-                    if complete is None and now_us() > target_time_us:
-                        break
-                    if complete is not None and complete.completion_let_us == event.completion_let_us:
-                        if complete.hardware_writes is not None:
-                            self._write_node_hardware(event.module, event.node, complete.hardware_writes)
-                        break
-
-            # Fetch next event
-            self._next_event = self._scheduler.next()
-
-    def get_articulation(self) -> Articulation:
-        """
-        Get the articulation for the Ark.
-
-        :return: The articulation instance.
-        """
-
-        return self._articulation
-
-    def get_camera(self, module_name: str, hardware_name: str) -> Camera:
-        """
-        Get camera hardware view.
-
-        :param module_name: Module name.
-        :param hardware_name: Hardware name.
-        :return: Camera view.
-        """
-
-        if (module_name, hardware_name) not in self._cameras:
-            raise SessionArkError(f"Camera '{hardware_name}' not found in module '{module_name}'")
-        return self._cameras[(module_name, hardware_name)]
-
-    def get_imu(self, module_name: str, hardware_name: str) -> Imu:
-        """
-        Get IMU hardware view.
-
-        :param module_name: Module name.
-        :param hardware_name: Hardware name.
-        :return: IMU view.
-        """
-
-        if (module_name, hardware_name) not in self._imus:
-            raise SessionArkError(f"IMU '{hardware_name}' not found in module '{module_name}'")
-        return self._imus[(module_name, hardware_name)]
-
-    def get_radar(self, module_name: str, hardware_name: str) -> Radar:
-        """
-        Get radar hardware view.
-
-        :param module_name: Module name.
-        :param hardware_name: Hardware name.
-        :return: Radar view.
-        """
-
-        if (module_name, hardware_name) not in self._radars:
-            raise SessionArkError(f"Radar '{hardware_name}' not found in module '{module_name}'")
-        return self._radars[(module_name, hardware_name)]
-
-    def get_pir(self, module_name: str, hardware_name: str) -> PirSensor:
-        """
-        Get PIR sensor hardware view.
-
-        :param module_name: Module name.
-        :param hardware_name: Hardware name.
-        :return: PIR sensor view.
-        """
-
-        if (module_name, hardware_name) not in self._pirs:
-            raise SessionArkError(f"PIR sensor '{hardware_name}' not found in module '{module_name}'")
-        return self._pirs[(module_name, hardware_name)]
-
-    def list_hardware(self) -> dict[str, list[tuple[str, str]]]:
-        """
-        List all hardware in the Ark.
-
-        :return: Dictionary of hardware (module name, hardware name) keyed by hardware type.
-        """
-
-        return {
-            "CAMERA": list(self._cameras.keys()),
-            "IMU": list(self._imus.keys()),
-            "PIR": list(self._pirs.keys()),
-            "RADAR": list(self._radars.keys()),
-        }
-
-    def record_node_output(
-        self,
-        module_name: str,
-        node_name: str,
-        output_name: str,
-        token_type: TokenType | None = TokenType.DATA,
-        last_n: int = 10,
-    ) -> NodeOutputRecorder:
-        """
-        Create a recorder for monitoring node output tokens.
-
-        :param module_name: Module name.
-        :param node_name: Node name within module.
-        :param output_name: Output name within node.
-        :param token_type: Token type to record (default: TokenType.DATA, None = all types).
-        :param last_n: Number of recent tokens to buffer.
-        :return: NodeOutputRecorder instance.
-        """
-
-        return NodeOutputRecorder(
-            comms=self._session.comms,
-            ark_def=self._ark_def,
-            module_name=module_name,
-            node_name=node_name,
-            output_name=output_name,
-            token_type=token_type,
-            last_n=last_n,
-        )
-
-    def _read_node_hardware(self, module_name: str, node_name: str) -> dict[str, bytes]:
-        """
-        Read hardware state for node execution.
-
-        :param module_name: Module name.
-        :param node_name: Node name.
-        :return: Hardware reads keyed by hardware name.
-        """
-
-        hardware_reads = {}
-        node_def = next(m for m in self._ark_def.modules if m.name == module_name).nodes[node_name]
-        for hardware_name, access_mode in node_def.hardware_access.items():
-            if access_mode not in (HardwareAccessMode.READ, HardwareAccessMode.READ_WRITE):
-                continue
-
-            # Read from camera
-            if camera := self._cameras.get((module_name, hardware_name)):
-                camera_frame = camera.get_frame()
-                if camera_frame is not None:
-                    hardware_reads[hardware_name] = camera_frame.pack()
-
-            # Read from IMU
-            elif imu := self._imus.get((module_name, hardware_name)):
-                imu_sample = imu.get_sample()
-                if imu_sample is not None:
-                    hardware_reads[hardware_name] = imu_sample.pack()
-
-            # Read from radar
-            elif radar := self._radars.get((module_name, hardware_name)):
-                radar_scan = radar.get_scan()
-                if radar_scan is not None:
-                    hardware_reads[hardware_name] = radar_scan.pack()
-
-            # Read from PIR sensor
-            elif pir := self._pirs.get((module_name, hardware_name)):
-                pir_status = pir.get_detection_status()
-                if pir_status is not None:
-                    hardware_reads[hardware_name] = pir_status.pack()
-
-            # Read from actuator group
-            else:
-                joint_states = self._articulation.get_joint_states()
-                if joint_states is not None:
-                    hardware_reads[hardware_name] = JointStates(states=joint_states).pack()
-
-        return hardware_reads
-
-    def _write_node_hardware(self, module_name: str, node_name: str, hardware_writes: dict[str, bytes]) -> None:
-        """
-        Write hardware state from node execution.
-
-        :param module_name: Module name.
-        :param node_name: Node name.
-        :param hardware_writes: Hardware writes keyed by hardware name.
-        """
-
-        node_def = next(m for m in self._ark_def.modules if m.name == module_name).nodes[node_name]
-        for hardware_name, write_data in hardware_writes.items():
-            access_mode = node_def.hardware_access.get(hardware_name)
-            if access_mode not in (HardwareAccessMode.WRITE, HardwareAccessMode.READ_WRITE):
-                continue
-
-            # Write to actuator group
-            targets = JointTargets.unpack(write_data)
-            self._articulation.set_joint_targets(joint_targets=targets.targets)
-
-    def _build_hardware(self) -> None:
-        """
-        Create articulation and sensor views from hardware configs.
-        """
-
-        actuator_group_configs = [
-            cast(ActuatorGroupHardware, hardware.config)
-            for hardware in self._ark_def.hardware
-            if hardware.type == HardwareType.ACTUATOR_GROUP
-        ]
-
-        # Create single articulation view with all joint configs from actuator groups
-        self._articulation = self._scene.add_articulation(
-            path=self._base_path,
-            joint_configs=sum([config.config.joint_configs for config in actuator_group_configs], []),
-        )
-
-        # Create sensor views
-        self._cameras: dict[tuple[str, str], Camera] = {}
-        self._imus: dict[tuple[str, str], Imu] = {}
-        self._pirs: dict[tuple[str, str], PirSensor] = {}
-        self._radars: dict[tuple[str, str], Radar] = {}
-        for hardware in self._ark_def.hardware:
-            if hardware.type == HardwareType.CAMERA:
-                camera_hw = cast(CameraHardware, hardware.config)
-                self._cameras[(camera_hw.module, camera_hw.name)] = self._scene.add_camera(
-                    path=f"{self._base_path}{camera_hw.path}",
-                    config=camera_hw.config,
-                    local_pose=camera_hw.pose,
-                )
-
-            elif hardware.type == HardwareType.IMU:
-                imu_hw = cast(ImuHardware, hardware.config)
-                self._imus[(imu_hw.module, imu_hw.name)] = self._scene.add_imu(
-                    path=f"{self._base_path}{imu_hw.path}",
-                    config=imu_hw.config,
-                    local_pose=imu_hw.pose,
-                )
-
-            elif hardware.type == HardwareType.PIR:
-                pir_hw = cast(PirHardware, hardware.config)
-                self._pirs[(pir_hw.module, pir_hw.name)] = self._scene.add_pir_sensor(
-                    path=f"{self._base_path}{pir_hw.path}",
-                    config=pir_hw.config,
-                    local_pose=pir_hw.pose,
-                )
-
-            elif hardware.type == HardwareType.RADAR:
-                radar_hw = cast(RadarHardware, hardware.config)
-                self._radars[(radar_hw.module, radar_hw.name)] = self._scene.add_radar(
-                    path=f"{self._base_path}{radar_hw.path}",
-                    config=radar_hw.config,
-                    local_pose=radar_hw.pose,
-                )
-
-    def _load_ark_asset(self, world_pose: Pose | None = None, local_pose: Pose | None = None) -> None:
-        """
-        Load the asset for the Ark.
-
-        :param world_pose: Optional world pose.
-        :param local_pose: Optional local pose.
-        :raises SessionAssetError: If the asset file is not found.
-        """
-
-        if (asset_path := self._ark_ref.asset_path) is None or not os.path.exists(asset_path):
-            raise SessionAssetError(f"Asset file not found for Ark {self._ark_def.name}")
-
-        self._scene.add_asset(
-            path=self._base_path,
-            asset_file_path=asset_path,
-            world_pose=world_pose,
-            local_pose=local_pose,
-        )

antioch/session/asset.py

@@ -1,65 +0,0 @@
-from pathlib import Path
-
-from antioch.session.session import SessionContainer
-from common.ark import AssetReference
-from common.core import list_local_assets, list_remote_assets, pull_remote_asset
-
-
-class Asset(SessionContainer):
-    """
-    Provides methods for interacting with the asset registry.
-
-    All methods handle authentication automatically using the current session's auth context.
-    This is a lightweight wrapper around registry functions that manages auth token complexity.
-
-    Example:
-        # List local assets
-        local_assets = Asset.list_local()
-
-        # List remote assets (requires auth)
-        remote_assets = Asset.list_remote()
-
-        # Pull an asset from remote registry
-        asset_path = Asset.pull(name="my_asset", version="1.0.0")
-    """
-
-    @staticmethod
-    def list_local() -> list[AssetReference]:
-        """
-        List all locally available assets.
-
-        :return: List of AssetReference objects from local storage.
-        """
-
-        return list_local_assets()
-
-    @staticmethod
-    def list_remote() -> list[AssetReference]:
-        """
-        List all assets from remote registry.
-
-        Requires authentication. Call session.login() first if not authenticated.
-
-        :return: List of AssetReference objects from remote registry.
-        :raises SessionAuthError: If not authenticated.
-        """
-
-        return list_remote_assets()
-
-    @staticmethod
-    def pull(name: str, version: str, overwrite: bool = False, show_progress: bool = True) -> Path:
-        """
-        Pull an asset from remote registry to local storage.
-
-        Requires authentication. Call session.login() first if not authenticated.
-        If the asset already exists locally, returns the existing path unless overwrite=True.
-
-        :param name: Name of the asset.
-        :param version: Version of the asset.
-        :param overwrite: Overwrite local asset if it already exists.
-        :param show_progress: Show download progress bar.
-        :return: Path to the downloaded asset file.
-        :raises SessionAuthError: If not authenticated.
-        """
-
-        return pull_remote_asset(name=name, version=version, overwrite=overwrite, show_progress=show_progress)

antioch/session/error.py

@@ -1,80 +0,0 @@
-class SessionError(Exception):
-    """
-    Base exception for all session-related errors.
-    """
-
-
-class SessionSimRpcNotConnectedError(SessionError):
-    """
-    Raised when the Sim RPC server is not reachable.
-    """
-
-
-class SessionSimRpcInterruptedError(SessionError):
-    """
-    Raised when the Sim RPC is interrupted.
-    """
-
-
-class SessionSimRpcClientError(SessionError):
-    """
-    Raised when the Sim RPC server returns a client error (user error).
-    """
-
-
-class SessionSimRpcInternalError(SessionError):
-    """
-    Raised when the Sim RPC server returns an internal error with traceback.
-    """
-
-    def __init__(self, message: str, traceback: str | None = None):
-        super().__init__(message)
-        self.traceback = traceback
-
-
-class SessionAuthError(SessionError):
-    """
-    Raised when authentication fails.
-    """
-
-
-class SessionAgentError(SessionError):
-    """
-    Raised when an agent operation fails (start/stop ark, telemetry, etc).
-    """
-
-
-class SessionTaskError(SessionError):
-    """
-    Raised when a task operation fails (invalid lifecycle, etc).
-    """
-
-
-class SessionArkError(SessionError):
-    """
-    Raised when an Ark operation fails (loading, stepping, hardware access, etc).
-    """
-
-
-class SessionHardwareError(SessionError):
-    """
-    Raised when a hardware operation fails.
-    """
-
-
-class SessionAssetError(SessionError):
-    """
-    Raised when an asset operation fails.
-    """
-
-
-class SessionRecordError(SessionError):
-    """
-    Raised when a recording operation fails (node output recording, etc).
-    """
-
-
-class SessionValidationError(SessionError):
-    """
-    Raised when input validation fails (user error).
-    """