armnet-runtime 0.1.0__py3-none-any.whl

armnet_runtime/context.py

@@ -0,0 +1,974 @@
+"""Job context surfaced to ``@main``-decorated functions.
+
+The cell program injects job env vars plus a JSON-encoded cell config when it
+starts the container; :func:`build_context` reads them and constructs the
+:class:`Context` object that the ``armnet-runtime`` entrypoint passes to
+the user's ``@main`` function.
+"""
+
+from __future__ import annotations
+
+import json
+import logging
+import time
+import os
+import shutil
+import socket
+import tempfile
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import Any, NamedTuple, Optional
+from urllib.parse import urlparse
+
+from armnet_core import Embodiment, RobotCellConfig, Task
+
+from armnet_runtime import env as env_keys
+
+logger = logging.getLogger(__name__)
+
+
+class ResetTimeoutException(RuntimeError):
+    """Raised by ``ctx.cell.reset()`` when a manual reset is not actioned in time.
+
+    The cell waits for an operator to confirm the reset (via the Fleet
+    Management System). If no confirmation arrives within the cell's reset
+    timeout, the cell trips its safety interlock (no further robot commands are
+    allowed, as with a safety violation) and this exception is raised into the
+    job code. Catch it to shut down gracefully — e.g. commit a dataset that was
+    being recorded — before letting the job fail::
+
+        try:
+            ctx.cell.reset()
+        except ResetTimeoutException:
+            dataset.push_to_hub()  # save what we collected
+            raise
+    """
+
+
+class CompletionStatus(NamedTuple):
+    """Result of :meth:`Cell.is_complete`: whether the episode is over and how.
+
+    ``complete`` is True once the current rollout/episode should end (the task
+    was scored complete by the automated monitor, or a human operator reported
+    an outcome). ``success`` is the verdict — for automated completion it equals
+    ``complete`` (a task scored complete is a success); for a human-reported
+    outcome it is the operator's success/fail choice, so an operator can end a
+    rollout as a *failure* (e.g. the robot behaved dangerously).
+
+    Backwards-compatible truthiness: ``bool(status)`` is ``status.complete``, so
+    existing ``if ctx.cell.is_complete(): ...`` callers keep working, while new
+    code can unpack ``complete, success = ctx.cell.is_complete()``.
+    """
+
+    complete: bool
+    success: bool
+
+    def __bool__(self) -> bool:  # noqa: D401 - truthy iff the episode is complete
+        return self.complete
+
+
+# Bounded read timeout for polling the teleop register. Unlike reset (which
+# waits, unbounded, for a human), a teleop read must never stall the control
+# loop: the register is in-memory on the cell, so a response that doesn't come
+# back well within this window means the channel is wedged and we should give
+# up, reconnect next time, and hold position.
+_TELEOP_READ_TIMEOUT_S = 1.0
+# Throttle interval for teleop-channel error logs (they'd otherwise repeat at
+# the control-loop rate while the channel is down).
+_TELEOP_ERROR_LOG_INTERVAL_S = 5.0
+
+
+@dataclass
+class Volume:
+    """User volume mounted into the runtime container."""
+
+    root: Optional[Path] = None
+
+    def path(self, relative_path: str | Path) -> Path:
+        if self.root is None:
+            raise RuntimeError("armnet volume is not mounted in this context")
+        rel = Path(relative_path)
+        if rel.is_absolute() or ".." in rel.parts:
+            raise ValueError("volume path must be relative and must not contain '..'")
+        return self.root / rel
+
+    def read_bytes(self, relative_path: str | Path) -> bytes:
+        return self.path(relative_path).read_bytes()
+
+    def read_text(self, relative_path: str | Path) -> str:
+        return self.path(relative_path).read_text()
+
+    def write_bytes(self, relative_path: str | Path, data: bytes) -> Path:
+        path = self.path(relative_path)
+        path.parent.mkdir(parents=True, exist_ok=True)
+        path.write_bytes(data)
+        return path
+
+    def write_text(self, relative_path: str | Path, data: str) -> Path:
+        path = self.path(relative_path)
+        path.parent.mkdir(parents=True, exist_ok=True)
+        path.write_text(data)
+        return path
+
+
+@dataclass(frozen=True)
+class RuntimeArm:
+    """Runtime-facing handle for one named robot arm in a multi-arm cell."""
+
+    name: str
+    robot_port: str
+    robot_id: Optional[str] = None
+    calibration_dir: Optional[Path] = None
+    calibration_file_path: Optional[Path] = None
+    safety_limit: Optional[float] = None
+
+
+@dataclass(frozen=True)
+class BimanualCalibrationLayout:
+    """Temporary calibration layout matching LeRobot's bimanual id convention."""
+
+    robot_id: str
+    calibration_dir: Path
+
+
+@dataclass
+class Cell:
+    """Handle to the physical cell the user code is running on.
+
+    M0.5 stub: there is no real cell yet, so ``robot_port`` is always ``None``
+    and :meth:`reset` is a no-op. The shape is fixed now so the spec
+    example compiles end-to-end and so M2/M3 can fill in the
+    implementation without touching customer-facing imports.
+    """
+
+    robot_port: Optional[str] = None
+    """Robot port value to pass into LeRobot robot configs.
+
+    In container-backed remote execution this is the connector endpoint, not
+    the host's physical serial path. The SDK's import-system swap routes that
+    endpoint through the cell-side connector, which then opens the real robot
+    port configured on the cell host.
+    """
+
+    robot_id: Optional[str] = None
+    """Stable robot id used by LeRobot to find calibration data."""
+
+    calibration_dir: Optional[Path] = None
+    """Calibration store path visible inside the customer container."""
+
+    calibration_file_path: Optional[Path] = None
+    """Exact LeRobot calibration file path visible inside the customer container."""
+
+    language_instruction: Optional[str] = None
+    """Task instruction provided by the cell."""
+
+    local_control_endpoint: Optional[str] = None
+    """Developer local-container control endpoint for keyboard-driven state."""
+
+    operator_call_endpoint: Optional[str] = None
+    """Operator-call endpoint served by the cell program for human-in-the-loop
+    calls (manual reset confirmation). Distinct from ``robot_port``, which is the
+    robot/bus connector (potentially a headless edge device)."""
+
+    is_local_container: bool = False
+    """True when running a Docker image locally for development."""
+    safety_limit: Optional[float] = None
+    """Relative action safety limit exposed by the cell, if applicable."""
+
+    arms: dict[str, RuntimeArm] = field(default_factory=dict)
+    """Named arms for bimanual/multi-arm cells."""
+
+    # Reused connection + log throttle for teleop polling (see get_teleop_action).
+    # Not part of the constructor or the public/comparable surface.
+    _teleop_conn: Any = field(default=None, init=False, repr=False, compare=False)
+    _teleop_last_error_log: float = field(default=0.0, init=False, repr=False, compare=False)
+    # Reused connection for polling the human-reported rollout outcome served by
+    # the cell program's operator-call endpoint (see is_complete).
+    _completion_conn: Any = field(default=None, init=False, repr=False, compare=False)
+    _completion_last_error_log: float = field(default=0.0, init=False, repr=False, compare=False)
+
+    @property
+    def is_bimanual(self) -> bool:
+        return {"left", "right"}.issubset(self.arms)
+
+    def arm(self, name: str) -> RuntimeArm:
+        try:
+            return self.arms[name]
+        except KeyError as exc:
+            raise KeyError(f"cell has no arm named {name!r}") from exc
+
+    def prepare_bimanual_calibration_dir(self) -> BimanualCalibrationLayout:
+        """Create a temp calibration dir using LeRobot's `<base>_<arm>.json` names.
+
+        Each arm's source calibration is resolved (in order) from its own
+        ``calibration_file_path``, its own ``calibration_dir`` keyed by the arm's
+        ``robot_id``, or—when the arm declares neither—the **cell-level**
+        ``calibration_dir`` keyed by the arm's ``robot_id`` (``<robot_id>.json``).
+        This mirrors how the cell's per-arm health check resolves calibration
+        (``arm.calibration_dir or cell.calibration_dir``), so a config that only
+        sets a top-level ``calibration_dir`` (per-arm ``robot_id`` only) works.
+        """
+
+        if not self.is_bimanual:
+            raise RuntimeError("bimanual calibration requires left and right arms")
+        robot_id = self.robot_id
+        if not robot_id:
+            raise RuntimeError("bimanual calibration requires ctx.cell.robot_id")
+        calibration_dir = Path(tempfile.mkdtemp(prefix="armnet-bimanual-calibration-"))
+        for arm_name in ("left", "right"):
+            arm = self.arm(arm_name)
+            source = arm.calibration_file_path
+            if source is None:
+                cal_dir = arm.calibration_dir or self.calibration_dir
+                if cal_dir is not None and arm.robot_id:
+                    source = Path(cal_dir) / f"{arm.robot_id}.json"
+            if source is None or not Path(source).is_file():
+                raise RuntimeError(
+                    f"no calibration file found for {arm_name} arm "
+                    f"(robot_id={arm.robot_id!r}); looked for "
+                    f"{source if source is not None else '<unresolved>'}. Set the "
+                    "cell-level calibration_dir (with per-arm robot_id) or each "
+                    "arm's calibration_dir/calibration_file_path."
+                )
+            shutil.copy2(source, calibration_dir / f"{robot_id}_{arm_name}.json")
+        return BimanualCalibrationLayout(robot_id=robot_id, calibration_dir=calibration_dir)
+
+    def reset(self) -> None:
+        """Return the robot to rest, then block until the operator confirms.
+
+        Two concerns, two endpoints:
+
+        1. Returning the arm to its rest position is a low-level bus operation,
+           so it is sent to the robot connector (``robot_port``), which may be a
+           headless edge device.
+        2. Operator confirmation is a human-in-the-loop concern, so it is sent
+           to the ``operator_call_endpoint`` served by the ``armnet-cell``
+           process, whose stdin is the operator's terminal.
+
+        The operator-facing prompt is owned by the cell, not by job code: job
+        code only signals *that* a reset point has been reached.
+        """
+
+        self._report_progress("Waiting for workspace reset...")
+
+        # 1. Safety: return the arm to rest via the robot connector, if present.
+        if self.robot_port and _looks_like_connector_endpoint(self.robot_port):
+            response = _connector_request(self.robot_port, {"op": "return_to_rest"})
+            if not response.get("ok"):
+                raise RuntimeError(response.get("error", "robot return-to-rest failed"))
+
+        # 2. Operator confirmation on the cell-served operator-call endpoint
+        # (fallback to the dev local-control endpoint).
+        request = {"op": "reset", "request": {"kind": "manual"}}
+        operator_endpoint = self.operator_call_endpoint or self.local_control_endpoint
+        if operator_endpoint:
+            response = _connector_request(operator_endpoint, request)
+            if not response.get("ok"):
+                error = response.get("error", "operator reset confirmation failed")
+                if response.get("error_type") == "ResetTimeoutException":
+                    raise ResetTimeoutException(error)
+                raise RuntimeError(error)
+            return
+
+        # No operator endpoint attached (degenerate in-process dev run): block on
+        # the local terminal with a standard prompt owned by the runtime.
+        input("Reset the cell workspace, then press Enter. ")
+
+    def is_complete(self, *, block: bool = False) -> CompletionStatus:
+        """Return whether the current episode is complete, and its success.
+
+        Returns a :class:`CompletionStatus` ``(complete, success)``:
+
+        1. A human-reported outcome (an operator hitting success/fail in the
+           FMS during a live rollout) takes precedence and ends the episode
+           immediately, with ``success`` set to the operator's choice. This is
+           how an operator stops a dangerous rollout without stopping the job.
+        2. Otherwise the cell's automated completion monitor is consulted; a
+           task scored complete is reported as a success (``success ==
+           complete``). Pass ``block=True`` for a final episode check that waits
+           for the cell to score the latest cached frames before returning.
+
+        ``bool(status)`` is ``status.complete`` for backwards compatibility.
+        """
+
+        # 1. Human-reported outcome (operator/FMS) wins and ends the episode now.
+        reported = self._human_completion()
+        if reported is not None:
+            return CompletionStatus(complete=True, success=reported)
+
+        # 2. Automated completion scoring via the local-dev or cell connector.
+        request = {"op": "is_complete", "block": block}
+        if self.local_control_endpoint:
+            response = _connector_request(self.local_control_endpoint, request)
+            if not response.get("ok"):
+                raise RuntimeError(response.get("error", "local completion check failed"))
+            return _completion_from_response(response)
+        if self.robot_port and _looks_like_connector_endpoint(self.robot_port):
+            response = _connector_request(self.robot_port, request)
+            if not response.get("ok"):
+                raise RuntimeError(response.get("error", "cell completion check failed"))
+            return _completion_from_response(response)
+        return CompletionStatus(complete=False, success=False)
+
+    def _human_completion(self) -> Optional[bool]:
+        """Return the operator's reported success/fail, or None if none pending.
+
+        Polls the cell program's operator-call endpoint (where FMS rollout
+        commands land). A wedged/slow channel must never stall the control
+        loop, so this mirrors get_teleop_action: bounded read, throttled error
+        logging, drop the connection on error, and treat failures as "no
+        outcome" so the episode simply continues.
+        """
+
+        endpoint = self.operator_call_endpoint or self.local_control_endpoint
+        if not endpoint:
+            return None
+
+        if self._completion_conn is None or self._completion_conn.endpoint != endpoint:
+            if self._completion_conn is not None:
+                self._completion_conn.close()
+            self._completion_conn = _TeleopConnection(endpoint, timeout=_TELEOP_READ_TIMEOUT_S)
+
+        try:
+            response = self._completion_conn.request({"op": "get_completion"})
+        except Exception as exc:  # noqa: BLE001
+            now = time.monotonic()
+            if now - self._completion_last_error_log >= _TELEOP_ERROR_LOG_INTERVAL_S:
+                self._completion_last_error_log = now
+                logger.warning(
+                    "completion read from %s failed (treating as not complete): %r",
+                    endpoint,
+                    exc,
+                )
+            return None
+
+        if not response.get("ok") or not response.get("reported"):
+            return None
+        return bool(response.get("success", False))
+
+    def rollout_begin(
+        self,
+        *,
+        index: Optional[int] = None,
+        total: Optional[int] = None,
+        outcome_controls: bool = True,
+    ) -> None:
+        """Tell the platform a rollout/episode in this job's loop has started.
+
+        The cell publishes this to the FMS, which shows the loop progress
+        ("rollout N / M") for the live job. Pass ``index`` (1-based) and, when
+        known, ``total`` so operators see how far along the loop is.
+
+        ``outcome_controls`` controls whether the FMS also shows operator
+        success/fail buttons: keep the default ``True`` for policy evals; pass
+        ``False`` for progress-only loops such as teleop data collection, where
+        a human verdict doesn't apply. Best-effort: a failed notification never
+        breaks the rollout. Pair with :meth:`rollout_end`.
+        """
+        payload: dict[str, Any] = {"outcome_controls": bool(outcome_controls)}
+        if index is not None:
+            payload["index"] = int(index)
+        if total is not None:
+            payload["total"] = int(total)
+        self._rollout_signal("rollout_begin", **payload)
+
+    def rollout_end(self) -> None:
+        """Tell the platform the current rollout has ended (hides FMS buttons)."""
+        self._rollout_signal("rollout_end")
+
+    def _rollout_signal(self, op: str, **payload: Any) -> None:
+        endpoint = self.operator_call_endpoint or self.local_control_endpoint
+        if not endpoint:
+            return
+        try:
+            _connector_request(endpoint, {"op": op, **payload})
+        except Exception:  # noqa: BLE001 - rollout signalling is best-effort
+            logger.warning("rollout signal %s to %s failed", op, endpoint, exc_info=True)
+
+    def is_shutting_down(self) -> bool:
+        """Return True once the cell has entered the job's post-timeout grace window.
+
+        When a job exceeds its ``timeout_seconds`` the cell does not kill the
+        container straight away: it trips the robot interlock (so any further
+        robot-bus calls fail) and opens a short *grace window* during which this
+        returns True, before force-killing the container. Poll it in your loop
+        and break out to finalize gracefully — e.g. save/push a dataset — instead
+        of being killed mid-write::
+
+            for episode in range(n):
+                if ctx.cell.is_shutting_down():
+                    break  # finalize below
+                ...
+
+        Resilient by design: returns False when no cell/operator endpoint is
+        attached or the status can't be read, so it never stalls or crashes the
+        control loop.
+        """
+
+        endpoint = self.operator_call_endpoint or self.local_control_endpoint
+        if not endpoint:
+            return False
+        try:
+            response = _connector_request(
+                endpoint, {"op": "shutdown_status"}, read_timeout=_TELEOP_READ_TIMEOUT_S
+            )
+        except Exception:  # noqa: BLE001 - never let a status poll break the loop
+            return False
+        return bool(response.get("ok") and response.get("shutting_down"))
+
+    def should_stop(self) -> bool:
+        """Return True when local/remote control asks user code to stop safely."""
+
+        request = {"op": "should_stop"}
+        if self.local_control_endpoint:
+            response = _connector_request(self.local_control_endpoint, request)
+            if not response.get("ok"):
+                raise RuntimeError(response.get("error", "local stop check failed"))
+            return bool(response.get("stop", False))
+        return False
+
+    def get_teleop_action(self) -> Optional[dict[str, float]]:
+        """Return the freshest remote-teleoperation action for this job, or None.
+
+        The client samples a local leader arm and pushes actions to the cell,
+        which keeps only the most recent one (older messages are dropped). This
+        reads that most-recent-value register over the operator-call endpoint.
+
+        Returns ``None`` when no teleop has been received yet (or no operator
+        endpoint is attached), so a control loop can hold position until the
+        operator starts driving. The returned dict is keyed for LeRobot's
+        ``send_action`` (e.g. ``{"shoulder_pan.pos": 12.3, ...}``).
+        """
+
+        endpoint = self.operator_call_endpoint or self.local_control_endpoint
+        if not endpoint:
+            return None
+
+        if self._teleop_conn is None or self._teleop_conn.endpoint != endpoint:
+            if self._teleop_conn is not None:
+                self._teleop_conn.close()
+            self._teleop_conn = _TeleopConnection(endpoint, timeout=_TELEOP_READ_TIMEOUT_S)
+
+        try:
+            response = self._teleop_conn.request({"op": "get_teleop"})
+        except Exception as exc:  # noqa: BLE001
+            # A wedged/slow teleop channel must not stall or crash the control
+            # loop: log (throttled) so a recurrence is diagnosable, drop the
+            # connection (already done in request()) so we reconnect next tick,
+            # and hold position by returning None.
+            now = time.monotonic()
+            if now - self._teleop_last_error_log >= _TELEOP_ERROR_LOG_INTERVAL_S:
+                self._teleop_last_error_log = now
+                logger.warning(
+                    "teleop read from %s failed (holding position; will reconnect): %r",
+                    endpoint,
+                    exc,
+                )
+            return None
+
+        if not response.get("ok"):
+            logger.warning("teleop read returned error: %s", response.get("error"))
+            return None
+        action = response.get("action")
+        if not action:
+            return None
+        return {str(key): float(value) for key, value in action.items()}
+
+    def get_teleop_event(self) -> Optional[str]:
+        """Return the next pending recording-control event, or None.
+
+        While teleoperating, the client can send discrete recording-control
+        events alongside the action stream — LeRobot's standard dataset
+        recording shortcuts: ``"next_episode"`` (Right Arrow: save the episode
+        and move on), ``"rerecord_episode"`` (Left Arrow: discard and redo) and
+        ``"stop_recording"`` (Esc: end the session). The cell queues them in
+        arrival order; each call pops at most one.
+
+        Like :meth:`get_teleop_action`, a wedged channel never stalls the
+        control loop: errors log (throttled), drop the connection so the next
+        call reconnects, and return None.
+        """
+
+        endpoint = self.operator_call_endpoint or self.local_control_endpoint
+        if not endpoint:
+            return None
+
+        if self._teleop_conn is None or self._teleop_conn.endpoint != endpoint:
+            if self._teleop_conn is not None:
+                self._teleop_conn.close()
+            self._teleop_conn = _TeleopConnection(endpoint, timeout=_TELEOP_READ_TIMEOUT_S)
+
+        try:
+            response = self._teleop_conn.request({"op": "get_teleop_event"})
+        except Exception as exc:  # noqa: BLE001
+            now = time.monotonic()
+            if now - self._teleop_last_error_log >= _TELEOP_ERROR_LOG_INTERVAL_S:
+                self._teleop_last_error_log = now
+                logger.warning(
+                    "teleop event read from %s failed (will reconnect): %r",
+                    endpoint,
+                    exc,
+                )
+            return None
+
+        if not response.get("ok"):
+            logger.warning("teleop event read returned error: %s", response.get("error"))
+            return None
+        event = response.get("event")
+        return str(event) if event else None
+
+    def _report_progress(self, message: str) -> None:
+        """Surface a progress message back to the platform.
+
+        M0.5: prints to stdout with a discoverable marker so the cell's
+        captured stdout shows progress in order with other prints. M1+
+        will also publish a NATS message so the orchestrator can stream
+        progress back to the client without waiting for the job to
+        terminate.
+        """
+
+        # Imported locally to avoid pulling markers into the public API
+        # surface of `Context`.
+        from armnet_runtime.markers import PROGRESS_MARKER
+        print(f"{PROGRESS_MARKER} {message}", flush=True)
+        time.sleep(0.01)
+
+
+def _completion_from_response(response: dict[str, Any]) -> CompletionStatus:
+    """Build a CompletionStatus from a connector ``is_complete`` response.
+
+    Automated completion has no separate success signal, so ``success``
+    defaults to ``complete`` (a task scored complete is a success) unless the
+    connector explicitly reports one.
+    """
+    complete = bool(response.get("complete", False))
+    success = bool(response.get("success", complete))
+    return CompletionStatus(complete=complete, success=success)
+
+
+def _looks_like_connector_endpoint(endpoint: str) -> bool:
+    if endpoint.startswith("tcp://"):
+        return True
+    path = Path(endpoint)
+    return path.exists() and path.is_socket()
+
+
+def _connector_request(
+    endpoint: str, payload: dict[str, Any], *, read_timeout: float | None = None
+) -> dict[str, Any]:
+    sock = _connect_connector(endpoint)
+    with sock:
+        # ``read_timeout`` bounds the response wait (used by polling calls that
+        # must never stall a control loop). Operator-paced calls such as ``reset``
+        # leave it None so the wait stays unbounded after the connection succeeds.
+        sock.settimeout(read_timeout)
+        sock.sendall(json.dumps(payload).encode("utf-8") + b"\n")
+        data = _read_line(sock)
+    return json.loads(data.decode("utf-8"))
+
+
+def _connect_connector(endpoint: str) -> socket.socket:
+    if endpoint.startswith("tcp://"):
+        parsed = urlparse(endpoint)
+        if not parsed.hostname or not parsed.port:
+            raise RuntimeError(f"invalid TCP connector endpoint: {endpoint!r}")
+        sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
+        sock.settimeout(10.0)
+        sock.connect((parsed.hostname, parsed.port))
+        return sock
+
+    sock = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
+    sock.settimeout(10.0)
+    sock.connect(endpoint)
+    return sock
+
+
+def _read_line(sock: socket.socket) -> bytes:
+    chunks: list[bytes] = []
+    while True:
+        chunk = sock.recv(65536)
+        if not chunk:
+            break
+        chunks.append(chunk)
+        if b"\n" in chunk:
+            break
+    data = b"".join(chunks)
+    return data.split(b"\n", 1)[0]
+
+
+class _TeleopConnection:
+    """Reused, bounded-timeout line connection for polling teleop from the cell.
+
+    Distinct from :func:`_connector_request` (a fresh socket per call with an
+    *unbounded* read, which is correct for operator-paced ``reset`` but would
+    let a wedged response hang the control loop forever). This keeps one socket
+    open across control steps and uses a short read timeout; on any error the
+    socket is dropped so the next request transparently reconnects.
+    """
+
+    def __init__(self, endpoint: str, *, timeout: float) -> None:
+        self.endpoint = endpoint
+        self._timeout = timeout
+        self._sock: Optional[socket.socket] = None
+        self._buf = b""
+
+    def request(self, payload: dict[str, Any]) -> dict[str, Any]:
+        sock = self._ensure()
+        try:
+            sock.settimeout(self._timeout)
+            sock.sendall(json.dumps(payload).encode("utf-8") + b"\n")
+            line = self._read_line(sock)
+        except Exception:
+            # Drop the (possibly half-open) socket so the next call reconnects.
+            self.close()
+            raise
+        return json.loads(line.decode("utf-8"))
+
+    def _ensure(self) -> socket.socket:
+        if self._sock is None:
+            self._sock = _connect_connector(self.endpoint)
+            self._buf = b""
+        return self._sock
+
+    def _read_line(self, sock: socket.socket) -> bytes:
+        while b"\n" not in self._buf:
+            chunk = sock.recv(65536)
+            if not chunk:
+                raise ConnectionError("teleop connection closed by the cell")
+            self._buf += chunk
+        line, _, self._buf = self._buf.partition(b"\n")
+        return line
+
+    def close(self) -> None:
+        if self._sock is not None:
+            try:
+                self._sock.close()
+            except Exception:  # noqa: BLE001
+                pass
+            self._sock = None
+        self._buf = b""
+
+
+@dataclass
+class Context:
+    """Everything a ``@main``-decorated function needs from the platform."""
+
+    job_id: str
+    embodiment: Embodiment
+    task: Task
+    args: dict[str, Any] = field(default_factory=dict)
+    cell: Cell = field(default_factory=Cell)
+    camera_configs: dict[str, Any] = field(default_factory=dict)
+    cache_home: Optional[Path] = None
+    volume: Volume = field(default_factory=Volume)
+    secrets: dict[str, str] = field(default_factory=dict)
+    timeout_seconds: Optional[int] = None
+    # Lazily created background Rerun streamer (see log_rerun_data). Not part of
+    # the constructor or the public/comparable surface.
+    _rerun_streamer: Any = field(default=None, init=False, repr=False, compare=False)
+
+    def report_progress(self, message: str) -> None:
+        """Surface a progress message back to the platform.
+
+        M0.5: prints to stdout with a discoverable marker so the cell's
+        captured stdout shows progress in order with other prints. M1+
+        will also publish a NATS message so the orchestrator can stream
+        progress back to the client without waiting for the job to
+        terminate.
+        """
+
+        # Imported locally to avoid pulling markers into the public API
+        # surface of `Context`.
+        from armnet_runtime.markers import PROGRESS_MARKER
+        print(f"{PROGRESS_MARKER} {message}", flush=True)
+
+    def is_shutting_down(self) -> bool:
+        """Whether the cell has entered the job's post-timeout grace window.
+
+        Convenience delegate for :meth:`Cell.is_shutting_down`. Poll it in long
+        loops and break out to finalize gracefully before the cell kills the
+        container.
+        """
+        return self.cell.is_shutting_down()
+
+    def log_rerun_data(
+        self,
+        observation: dict[str, Any] | None = None,
+        action: dict[str, Any] | None = None,
+        *,
+        compress_images: bool = True,
+        jpeg_quality: int = 75,
+    ) -> None:
+        """Stream observation/action data to a Rerun viewer on the client.
+
+        Mirrors LeRobot's ``log_rerun_data``: scalars are logged as Rerun
+        scalars, image-like arrays as images, and other arrays as per-element
+        scalars. Keys are namespaced with ``observation.`` / ``action.`` when
+        not already.
+
+        Unlike the LeRobot helper, this does not call ``rr.log`` in-process
+        (the cell container has no viewer). Instead it serializes a protobuf
+        packet and emits it on stdout behind a marker; the cell republishes it
+        on ``logs.<job_id>.rerun`` and the client's orchestrate script replays
+        it into the viewer it started with ``rr.init(...)``.
+
+        Images are JPEG-compressed by default to keep the NATS stream light;
+        set ``compress_images=False`` to send raw RGB. opencv is required for
+        compression and numpy for any array handling; both are imported lazily.
+
+        Non-blocking: the snapshot is handed to a background worker thread that
+        does the encoding and stdout write, so the calling control loop never
+        stalls on visualization. The worker's queue is bounded and drops the
+        oldest pending frame under backpressure (tune with
+        ``ARMNET_RERUN_QUEUE_MAXSIZE``), so a slow consumer sheds frames
+        rather than slowing the robot loop.
+        """
+
+        if not observation and not action:
+            return
+
+        from armnet_runtime.rerun import RerunStreamer
+
+        if self._rerun_streamer is None:
+            self._rerun_streamer = RerunStreamer(self.job_id)
+            self._rerun_streamer.start()
+        self._rerun_streamer.submit(
+            observation,
+            action,
+            compress_images=compress_images,
+            jpeg_quality=jpeg_quality,
+        )
+
+
+SO101_EMBODIMENT = "lerobot/so-101"
+BIMANUAL_SO101_EMBODIMENT = "lerobot/bimanual_so101"
+
+
+def require_so101_embodiment(ctx: "Context", runtime_name: str) -> bool:
+    """Validate the job's embodiment is a (single or bimanual) SO-101.
+
+    The embodiment is the source of truth for how many arms the robot has —
+    ``lerobot/so-101`` is a single arm, ``lerobot/bimanual_so101`` is two — and
+    the orchestrator only routes a job to a cell of the matching embodiment.
+    Returns ``True`` for the bimanual embodiment (so the caller builds a two-arm
+    robot), ``False`` for single-arm.
+
+    Raises :class:`NotImplementedError` for any other embodiment, and
+    :class:`RuntimeError` if the embodiment's arm count disagrees with the cell's
+    actual wiring (``ctx.cell.is_bimanual``) — a misrouted or misconfigured cell.
+    """
+    if ctx.embodiment not in (SO101_EMBODIMENT, BIMANUAL_SO101_EMBODIMENT):
+        raise NotImplementedError(
+            f"{runtime_name} supports {SO101_EMBODIMENT!r} and "
+            f"{BIMANUAL_SO101_EMBODIMENT!r}, got {ctx.embodiment!r}"
+        )
+    expect_bimanual = ctx.embodiment == BIMANUAL_SO101_EMBODIMENT
+    if expect_bimanual != ctx.cell.is_bimanual:
+        raise RuntimeError(
+            f"embodiment {ctx.embodiment!r} expects "
+            f"{'two arms (left+right)' if expect_bimanual else 'a single arm'}, "
+            f"but the cell exposes arms={sorted(ctx.cell.arms)}"
+        )
+    return expect_bimanual
+
+
+def build_context() -> Context:
+    """Read the cell-injected env vars and construct a :class:`Context`.
+
+    Called by the ``armnet-runtime`` entrypoint before invoking the
+    user's ``@main`` function. Raises :class:`RuntimeError` with a clear
+    message if any required env var is missing \u2014 that indicates the
+    program is being run outside a armnet cell.
+    """
+
+    job_id = os.environ.get(env_keys.JOB_ID)
+    embodiment_raw = os.environ.get(env_keys.EMBODIMENT)
+    task_raw = os.environ.get(env_keys.TASK)
+    if not job_id or not embodiment_raw or not task_raw:
+        raise RuntimeError(
+            "armnet runtime env vars not set; this program is meant to "
+            "be executed inside a armnet cell container. Missing one of: "
+            f"{env_keys.JOB_ID}, {env_keys.EMBODIMENT}, {env_keys.TASK}."
+        )
+
+    timeout_raw = os.environ.get(env_keys.TIMEOUT_SECONDS)
+    timeout = int(timeout_raw) if timeout_raw and timeout_raw.isdigit() else None
+
+    args_raw = os.environ.get(env_keys.ARGS, "{}")
+    try:
+        args = json.loads(args_raw)
+    except json.JSONDecodeError as exc:
+        raise RuntimeError(
+            f"{env_keys.ARGS} env var is not valid JSON: {exc}"
+        ) from exc
+    if not isinstance(args, dict):
+        raise RuntimeError(
+            f"{env_keys.ARGS} must decode to a JSON object (dict); got {type(args).__name__}."
+        )
+
+    cell_config = _load_cell_config()
+    return Context(
+        job_id=job_id,
+        embodiment=embodiment_raw,
+        task=task_raw,
+        args=args,
+        cell=Cell(
+            robot_port=(
+                os.environ.get(env_keys.CELL_SOCKET)
+                or (cell_config.robot_connector_endpoint if cell_config else None)
+                or (cell_config.connector_socket_path if cell_config else None)
+                or (cell_config.robot_port if cell_config else None)
+            ),
+            operator_call_endpoint=(
+                os.environ.get(env_keys.OPERATOR_CALL_ENDPOINT)
+                or (cell_config.operator_call_endpoint if cell_config else None)
+            ),
+            robot_id=os.environ.get(env_keys.ROBOT_ID)
+            or (cell_config.robot_id if cell_config else None),
+            calibration_dir=_path_from_env_or_config(
+                env_keys.CALIBRATION_DIR,
+                cell_config.calibration_dir if cell_config else None,
+            ),
+            calibration_file_path=_path_from_env_or_config(
+                env_keys.CALIBRATION_FILE_PATH,
+                cell_config.calibration_file_path if cell_config else None,
+            ),
+            language_instruction=cell_config.language_instruction if cell_config else None,
+            local_control_endpoint=os.environ.get(env_keys.LOCAL_CONTROL_ENDPOINT),
+            is_local_container=bool(os.environ.get(env_keys.LOCAL_CONTAINER)),
+            safety_limit=_cell_safety_limit(cell_config, embodiment_raw),
+            arms=_build_runtime_arms(cell_config),
+        ),
+        camera_configs=_build_camera_configs(
+            cell_config.camera_configs if cell_config else {}
+        ),
+        cache_home=_path_from_env_or_config(env_keys.CACHE_HOME, None),
+        volume=Volume(root=_path_from_env_or_config(env_keys.VOLUME_HOME, None)),
+        secrets=_load_secrets(),
+        timeout_seconds=timeout,
+    )
+
+
+def _load_secrets() -> dict[str, str]:
+    raw = os.environ.get(env_keys.SECRETS, "{}")
+    try:
+        value = json.loads(raw)
+    except json.JSONDecodeError as exc:
+        raise RuntimeError(f"{env_keys.SECRETS} env var is not valid JSON: {exc}") from exc
+    if not isinstance(value, dict):
+        raise RuntimeError(f"{env_keys.SECRETS} must decode to a JSON object.")
+    return {str(key): str(item) for key, item in value.items()}
+
+
+def _cell_safety_limit(cell_config: RobotCellConfig | None, embodiment: str) -> float | None:
+    """Resolve the runtime max_relative_target (degrees) for a cell.
+
+    ``safety_limit`` and ``safety_delta_degrees`` are the same concept (max
+    allowed relative joint motion in degrees) under two names: ``safety_limit``
+    is the runtime/LeRobot-facing name (``max_relative_target``), while configs
+    canonically declare ``robot.safety_delta_degrees`` (also pushed to the edge
+    connector for its own safety checks). An explicit top-level ``safety_limit``
+    wins; otherwise the SO-101 family (single-arm and bimanual) derives it from
+    the configured ``safety_delta_degrees``. Embodiments with connector-only
+    safety (e.g. ARX5) leave it unset.
+    """
+    if cell_config and cell_config.safety_limit is not None:
+        return float(cell_config.safety_limit)
+    if embodiment.startswith("lerobot/so-") or embodiment == BIMANUAL_SO101_EMBODIMENT:
+        return float(cell_config.safety_delta_degrees) if cell_config else 30.0
+    return None
+
+
+def _build_runtime_arms(cell_config: RobotCellConfig | None) -> dict[str, RuntimeArm]:
+    if cell_config is None or not cell_config.arms:
+        return {}
+    base_endpoint = cell_config.robot_connector_endpoint or cell_config.connector_socket_path
+    if not base_endpoint:
+        return {}
+    return {
+        name: RuntimeArm(
+            name=name,
+            robot_port=_arm_connector_endpoint(base_endpoint, name),
+            robot_id=arm.robot_id,
+            calibration_dir=Path(arm.calibration_dir) if arm.calibration_dir else None,
+            calibration_file_path=(
+                Path(arm.calibration_file_path) if arm.calibration_file_path else None
+            ),
+            # Per-arm configs declare safety as ``safety_delta_degrees`` (same as
+            # the cell level); fall back to it so the arm's max_relative_target is
+            # applied. ``safety_limit`` is accepted too for symmetry.
+            safety_limit=(
+                arm.safety_limit if arm.safety_limit is not None else arm.safety_delta_degrees
+            ),
+        )
+        for name, arm in cell_config.arms.items()
+    }
+
+
+def _arm_connector_endpoint(base_endpoint: str, arm_name: str) -> str:
+    if not base_endpoint.startswith("tcp://"):
+        return base_endpoint
+    parsed = urlparse(base_endpoint)
+    query = f"arm={arm_name}"
+    if parsed.query:
+        query = f"{parsed.query}&{query}"
+    return parsed._replace(query=query).geturl()
+
+
+def _load_cell_config() -> RobotCellConfig | None:
+    raw = os.environ.get(env_keys.CELL_CONFIG)
+    if not raw:
+        return None
+    try:
+        return RobotCellConfig.model_validate_json(raw)
+    except Exception as exc:  # noqa: BLE001
+        raise RuntimeError(f"{env_keys.CELL_CONFIG} env var is not valid RobotCellConfig JSON: {exc}") from exc
+
+
+def _path_from_env_or_config(env_key: str, configured: str | None) -> Path | None:
+    value = os.environ.get(env_key) or configured
+    return Path(value) if value else None
+
+
+def _build_camera_configs(camera_configs: dict[str, Any] | None) -> dict[str, Any]:
+    if not camera_configs:
+        return {}
+    return {
+        name: _build_camera_config(config)
+        for name, config in camera_configs.items()
+    }
+
+
+def _build_camera_config(config: Any) -> Any:
+    if not isinstance(config, dict):
+        return config
+
+    camera_type = str(config.get("type", "opencv")).lower()
+    # Strip keys that are not OpenCVCameraConfig params (e.g. arducam_mount is
+    # only used by the edge connector for camera mount actuation; id_path is the
+    # USB port used at provisioning time to create the /dev/video_<name> symlink).
+    _NON_CAMERA_KEYS = {"type", "arducam_mount", "id_path"}
+    kwargs = {key: value for key, value in config.items() if key not in _NON_CAMERA_KEYS}
+    if camera_type != "opencv":
+        return config
+
+    try:
+        from lerobot.cameras import ColorMode
+        from lerobot.cameras.opencv import OpenCVCameraConfig
+    except ImportError:
+        return config
+
+    if "color_mode" in kwargs and not isinstance(kwargs["color_mode"], ColorMode):
+        kwargs["color_mode"] = _coerce_enum(ColorMode, kwargs["color_mode"])
+    return OpenCVCameraConfig(**kwargs)
+
+
+def _coerce_enum(enum_cls: Any, value: Any) -> Any:
+    if isinstance(value, enum_cls):
+        return value
+    if isinstance(value, str):
+        try:
+            return enum_cls[value.upper()]
+        except KeyError:
+            return enum_cls(value)
+    return enum_cls(value)