armnet-runtime 0.1.0__py3-none-any.whl

armnet_runtime/__init__.py

@@ -0,0 +1,66 @@
+"""armnet runtime SDK \u2014 used inside customer containers on a cell.
+
+Customer-facing surface: a ``@main`` decorator and a :class:`Context`. The
+container's entrypoint is the ``armnet-runtime`` console script
+(installed by this package); it loads the user's file, finds the
+``@main``-decorated function, builds the context, and calls it.
+
+Typical use::
+
+    from armnet_runtime import main, Context
+
+    @main
+    def run(ctx: Context):
+        seed = ctx.args.get("seed", 0)
+        ctx.report_progress("starting")
+        ...
+        return {"success_rate": 1.0}
+
+Wire types (``Embodiment``, ``Task``, ``JobSpec``, ``JobResult``, ...) are
+re-exported from :mod:`armnet_core` for ergonomics, so customer code
+never needs to know about the core package directly.
+"""
+
+from armnet_core import (
+    Embodiment,
+    Job,
+    JobResult,
+    JobSpec,
+    JobStatus,
+    RemoteExecutionError,
+    Task,
+    TerminalStatus,
+)
+
+from armnet_runtime.context import (
+    BIMANUAL_SO101_EMBODIMENT,
+    SO101_EMBODIMENT,
+    Cell,
+    CompletionStatus,
+    Context,
+    ResetTimeoutException,
+    require_so101_embodiment,
+)
+from armnet_runtime.decorator import MainRegistrationError, main
+
+__all__ = [
+    "BIMANUAL_SO101_EMBODIMENT",
+    "SO101_EMBODIMENT",
+    "Cell",
+    "CompletionStatus",
+    "Context",
+    "require_so101_embodiment",
+    "Embodiment",
+    "Job",
+    "JobResult",
+    "JobSpec",
+    "JobStatus",
+    "MainRegistrationError",
+    "RemoteExecutionError",
+    "ResetTimeoutException",
+    "Task",
+    "TerminalStatus",
+    "main",
+]
+
+__version__ = "0.1.0"

armnet_runtime/_rerun_pb2.py

@@ -0,0 +1,32 @@
+# -*- coding: utf-8 -*-
+# Generated by the protocol buffer compiler.  DO NOT EDIT!
+# source: rerun.proto
+# Protobuf Python Version: 5.26.1
+"""Generated protocol buffer code."""
+from google.protobuf import descriptor as _descriptor
+from google.protobuf import descriptor_pool as _descriptor_pool
+from google.protobuf import symbol_database as _symbol_database
+from google.protobuf.internal import builder as _builder
+# @@protoc_insertion_point(imports)
+
+_sym_db = _symbol_database.Default()
+
+
+
+
+DESCRIPTOR = _descriptor_pool.Default().AddSerializedFile(b'\n\x0brerun.proto\x12\x10remoterobo.rerun\"y\n\x05Image\x12\x0c\n\x04\x64\x61ta\x18\x01 \x01(\x0c\x12\x31\n\x08\x65ncoding\x18\x02 \x01(\x0e\x32\x1f.remoterobo.rerun.ImageEncoding\x12\r\n\x05width\x18\x03 \x01(\r\x12\x0e\n\x06height\x18\x04 \x01(\r\x12\x10\n\x08\x63hannels\x18\x05 \x01(\r\"f\n\nRerunEntry\x12\x13\n\x0b\x65ntity_path\x18\x01 \x01(\t\x12\x10\n\x06scalar\x18\x02 \x01(\x01H\x00\x12(\n\x05image\x18\x03 \x01(\x0b\x32\x17.remoterobo.rerun.ImageH\x00\x42\x07\n\x05value\"\x82\x01\n\x0bRerunPacket\x12\x0e\n\x06job_id\x18\x01 \x01(\t\x12\x11\n\ttimestamp\x18\x02 \x01(\x01\x12-\n\x07\x65ntries\x18\x03 \x03(\x0b\x32\x1c.remoterobo.rerun.RerunEntry\x12\x0e\n\x06obs_ts\x18\x04 \x01(\x01\x12\x11\n\tsubmit_ts\x18\x05 \x01(\x01*d\n\rImageEncoding\x12\x1e\n\x1aIMAGE_ENCODING_UNSPECIFIED\x10\x00\x12\x17\n\x13IMAGE_ENCODING_JPEG\x10\x01\x12\x1a\n\x16IMAGE_ENCODING_RAW_RGB\x10\x02\x62\x06proto3')
+
+_globals = globals()
+_builder.BuildMessageAndEnumDescriptors(DESCRIPTOR, _globals)
+_builder.BuildTopDescriptorsAndMessages(DESCRIPTOR, 'rerun_pb2', _globals)
+if not _descriptor._USE_C_DESCRIPTORS:
+  DESCRIPTOR._loaded_options = None
+  _globals['_IMAGEENCODING']._serialized_start=393
+  _globals['_IMAGEENCODING']._serialized_end=493
+  _globals['_IMAGE']._serialized_start=33
+  _globals['_IMAGE']._serialized_end=154
+  _globals['_RERUNENTRY']._serialized_start=156
+  _globals['_RERUNENTRY']._serialized_end=258
+  _globals['_RERUNPACKET']._serialized_start=261
+  _globals['_RERUNPACKET']._serialized_end=391
+# @@protoc_insertion_point(module_scope)

armnet_runtime/cli.py

@@ -0,0 +1,217 @@
+"""``armnet-runtime`` console script.
+
+This is the container entrypoint for any image built on top of the
+armnet runtime SDK. The container's ``CMD`` looks like::
+
+    CMD ["armnet-runtime", "/app/hello.py"]
+
+and this script:
+
+  1. (M3+) Performs the LeRobot import-system swap so customer code
+     transparently routes hardware operations through the safety-aware
+     robot connector. Currently a marker; nothing is swapped yet.
+  2. Adds the user script's directory to ``sys.path`` so it can import
+     sibling modules.
+  3. Imports the user script \u2014 which executes the ``@main`` decorator as
+     a side effect, registering the entry-point function.
+  4. Builds a :class:`~armnet_runtime.Context` from cell-injected env
+     vars.
+  5. Calls the registered function with the context.
+  6. Prints the function's return value with the
+     :data:`~armnet_runtime.markers.RESULT_MARKER_JSON` prefix; the
+     cell extracts that and surfaces it as
+     :attr:`~armnet_core.JobResult.return_value`.
+  7. Maps exceptions to a non-zero exit so the cell records the job as
+     ``FAILED``; the traceback ends up in the captured stderr.
+"""
+
+from __future__ import annotations
+
+import argparse
+import importlib.util
+import json
+import sys
+import traceback
+from datetime import datetime
+from pathlib import Path
+from typing import TextIO
+from typing import Sequence
+
+from armnet_runtime.context import build_context
+from armnet_runtime.decorator import registered_main
+from armnet_runtime.markers import (
+    MAX_TRACEBACK_BYTES,
+    RESULT_MARKER,
+    RESULT_MARKER_JSON,
+    TRACEBACK_MARKER_JSON,
+)
+
+# Exit codes. Keep stable \u2014 the cell may eventually distinguish
+# runtime-CLI failures from user-code failures, e.g. by mapping these to
+# JobResult.error.
+EXIT_OK = 0
+EXIT_USER_RAISED = 1
+EXIT_USAGE = 2
+EXIT_NO_MAIN = 3
+EXIT_USER_IMPORT_FAILED = 4
+EXIT_BAD_CONTEXT = 5
+
+_USER_MODULE_NAME = "__armnet_user_main__"
+
+
+class _TimestampedStream:
+    """Prefix every stdout/stderr line with a Unix timestamp."""
+
+    def __init__(self, wrapped: TextIO) -> None:
+        self._wrapped = wrapped
+        self._at_line_start = True
+
+    def write(self, text: str) -> int:
+        written = 0
+        for chunk in text.splitlines(keepends=True):
+            if self._at_line_start and chunk:
+                prefix = f"[{datetime.now().isoformat(timespec='milliseconds')}] "
+                self._wrapped.write(prefix)
+            self._wrapped.write(chunk)
+            written += len(chunk)
+            self._at_line_start = chunk.endswith(("\n", "\r"))
+        return written
+
+    def flush(self) -> None:
+        self._wrapped.flush()
+
+    def isatty(self) -> bool:
+        return self._wrapped.isatty()
+
+    def __getattr__(self, name: str) -> object:
+        return getattr(self._wrapped, name)
+
+
+def _install_timestamped_streams() -> None:
+    sys.stdout = _TimestampedStream(sys.stdout)  # type: ignore[assignment]
+    sys.stderr = _TimestampedStream(sys.stderr)  # type: ignore[assignment]
+
+
+def _parse_args(argv: Sequence[str]) -> argparse.Namespace:
+    parser = argparse.ArgumentParser(
+        prog="armnet-runtime",
+        description="Run a armnet @main-decorated script inside a cell container.",
+    )
+    parser.add_argument(
+        "script",
+        type=Path,
+        help="Path to the customer's Python script containing one @main-decorated function.",
+    )
+    return parser.parse_args(argv)
+
+
+def _load_user_script(script_path: Path) -> None:
+    """Import the user's script; the @main side-effect populates the registry."""
+
+    if not script_path.is_file():
+        raise FileNotFoundError(f"script not found: {script_path}")
+
+    sys.path.insert(0, str(script_path.parent))
+
+    spec = importlib.util.spec_from_file_location(_USER_MODULE_NAME, script_path)
+    if spec is None or spec.loader is None:
+        raise ImportError(f"could not build import spec for {script_path}")
+    module = importlib.util.module_from_spec(spec)
+    sys.modules[_USER_MODULE_NAME] = module
+    spec.loader.exec_module(module)
+
+
+def _print_result(value: object) -> None:
+    """Render the user's return value to stdout.
+
+    Two lines: a human-readable line for ergonomic logs, and a marker line
+    with JSON for the cell to parse into ``JobResult.return_value``.
+    """
+
+    print(f"\n{RESULT_MARKER} {value!r}", flush=True)
+    try:
+        encoded = json.dumps(value, default=repr)
+    except (TypeError, ValueError):
+        encoded = json.dumps(repr(value))
+    print(f"{RESULT_MARKER_JSON}{encoded}", flush=True)
+
+
+def run(argv: Sequence[str] | None = None) -> int:
+    args = _parse_args(sys.argv[1:] if argv is None else list(argv))
+    _install_timestamped_streams()
+
+    # Install guarded LeRobot import replacement before loading user code so
+    # `from lerobot...` inside the script sees the remote Feetech proxy.
+    from armnet_runtime.lerobot import install_lerobot_remote_mode
+
+    install_lerobot_remote_mode()
+
+    try:
+        _load_user_script(args.script)
+    except Exception:
+        print("armnet-runtime: failed to import user script:", file=sys.stderr)
+        traceback.print_exc()
+        # An import failure is still a Python traceback in user code (think
+        # syntax error in `hello.py`); surface it the same way as a runtime
+        # raise so the cell lifts it onto JobResult.traceback.
+        _print_traceback_marker(traceback.format_exc())
+        return EXIT_USER_IMPORT_FAILED
+
+    fn = registered_main()
+    if fn is None:
+        print(
+            f"armnet-runtime: no @main-decorated function found in {args.script}. "
+            "Decorate exactly one function with `@main` to mark it as the entry point.",
+            file=sys.stderr,
+        )
+        return EXIT_NO_MAIN
+
+    try:
+        ctx = build_context()
+    except Exception as exc:  # noqa: BLE001
+        print(f"armnet-runtime: failed to build context: {exc}", file=sys.stderr)
+        return EXIT_BAD_CONTEXT
+
+    try:
+        result = fn(ctx)
+    except SystemExit:
+        # Customer code called sys.exit explicitly \u2014 let that propagate.
+        raise
+    except Exception:
+        # Stderr: full, human-readable for `docker logs` users.
+        print(
+            f"armnet-runtime: @main function raised; "
+            f"job {ctx.job_id} will be marked FAILED.",
+            file=sys.stderr,
+        )
+        traceback.print_exc()
+        # Stdout: structured marker the cell parses into JobResult.traceback.
+        _print_traceback_marker(traceback.format_exc())
+        return EXIT_USER_RAISED
+
+    if result is not None:
+        _print_result(result)
+    return EXIT_OK
+
+
+def _print_traceback_marker(tb_text: str) -> None:
+    """Emit the `[armnet:traceback]:json` marker on stdout.
+
+    Truncates at :data:`MAX_TRACEBACK_BYTES` of UTF-8 to keep NATS payloads
+    bounded. The full untruncated traceback is still in stderr (printed by
+    ``traceback.print_exc()`` above), so power users on the cell side can
+    read ``JobResult.stderr`` if the truncated head isn't enough.
+    """
+
+    encoded = tb_text.encode("utf-8")
+    if len(encoded) > MAX_TRACEBACK_BYTES:
+        # Truncate on a UTF-8 boundary, then append a marker visible inside
+        # the truncated payload itself.
+        head = encoded[:MAX_TRACEBACK_BYTES].decode("utf-8", errors="ignore")
+        tb_text = head + "\n... [traceback truncated; full version in stderr]"
+
+    print(f"{TRACEBACK_MARKER_JSON}{json.dumps(tb_text)}", flush=True)
+
+
+if __name__ == "__main__":
+    sys.exit(run())