mru-sdk 0.1.0a1__py3-none-any.whl

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
mru/__init__.py ADDED
@@ -0,0 +1,21 @@
1
+ """MRU: robot capabilities for Python jobs; model logic belongs to the job."""
2
+
3
+ __version__ = "0.1.0a1"
4
+
5
+ from .auth import Auth
6
+ from .client import Client, LogEntry, LogPage, RemoteJob
7
+ from .errors import APIError, LaunchError, PlatformConnectionError
8
+ from .errors import (AuthenticationError, AuthorizationError, CommandRejected, DeadlineExceeded,
9
+ JobCancelled, MRUError, ProtocolError, SessionClosed)
10
+ from .job import Actuator, Job, run, run_async
11
+ from .models import (ActuatorRequirement, ActuatorSpec, EncodedImage, ExecutionAck, JointState,
12
+ Observation, Requirements, RobotCapabilities)
13
+ from .session import Session
14
+
15
+ __all__ = [
16
+ "APIError", "Client", "LaunchError", "LogEntry", "LogPage", "PlatformConnectionError", "RemoteJob",
17
+ "Auth", "AuthenticationError", "AuthorizationError",
18
+ "Actuator", "ActuatorRequirement", "ActuatorSpec", "CommandRejected", "DeadlineExceeded",
19
+ "EncodedImage", "ExecutionAck", "Job", "JobCancelled", "JointState", "MRUError", "Observation",
20
+ "ProtocolError", "Requirements", "RobotCapabilities", "Session", "SessionClosed", "run", "run_async",
21
+ ]
mru/auth.py ADDED
@@ -0,0 +1,39 @@
1
+ """Job credentials supplied by the platform, separate from task parameters."""
2
+
3
+ from __future__ import annotations
4
+
5
+ import os
6
+ from dataclasses import dataclass, field
7
+
8
+ from .errors import AuthenticationError
9
+
10
+
11
+ @dataclass(frozen=True, slots=True)
12
+ class Auth:
13
+ """An opaque bearer credential. Construction does not validate access with MRU."""
14
+
15
+ _token: str = field(repr=False)
16
+
17
+ def __post_init__(self) -> None:
18
+ if (not isinstance(self._token, str) or not self._token
19
+ or any(not 33 <= ord(character) <= 126 for character in self._token)):
20
+ raise ValueError("Job token must be a nonempty printable ASCII string without whitespace")
21
+
22
+ @classmethod
23
+ def from_token(cls, token: str) -> Auth:
24
+ """Use a job token already issued by the platform; never a password."""
25
+ return cls(token)
26
+
27
+ @classmethod
28
+ def from_env(cls) -> Auth:
29
+ """Read MRU_JOB_TOKEN, failing if it is missing or malformed."""
30
+ token = os.environ.get("MRU_JOB_TOKEN")
31
+ if token is None:
32
+ raise AuthenticationError("unauthenticated", "Set MRU_JOB_TOKEN or pass auth=Auth.from_token(...)")
33
+ try:
34
+ return cls.from_token(token)
35
+ except ValueError:
36
+ raise AuthenticationError("unauthenticated", "MRU_JOB_TOKEN is malformed") from None
37
+
38
+ def _to_wire(self) -> dict[str, str]:
39
+ return {"scheme": "bearer", "token": self._token}
mru/client.py ADDED
@@ -0,0 +1,281 @@
1
+ """Python job launch client for the proposed mru.python-job/v1 backend contract.
2
+
3
+ The backend contract is documented in docs/launch-api.md; it is not deployed
4
+ by this package. No connection to a legacy deployment or CLI API is made.
5
+ """
6
+
7
+ from __future__ import annotations
8
+
9
+ import json
10
+ import math
11
+ import os
12
+ import re
13
+ import time
14
+ import uuid
15
+ from collections.abc import Mapping, Sequence
16
+ from dataclasses import dataclass
17
+ from email.message import Message
18
+ from http.client import HTTPException
19
+ from pathlib import Path
20
+ from typing import IO, Any
21
+ from urllib.error import HTTPError, URLError
22
+ from urllib.parse import quote, urlencode, urlsplit
23
+ from urllib.request import HTTPRedirectHandler, Request, build_opener
24
+
25
+ from .errors import (APIError, AuthenticationError, AuthorizationError, LaunchError,
26
+ PlatformConnectionError, ProtocolError)
27
+ from .models import Requirements
28
+ from .source import bundle_source
29
+
30
+ TERMINAL_STATUSES = frozenset({"completed", "failed", "cancelled"})
31
+
32
+
33
+ def _positive(value: float, name: str) -> float:
34
+ if isinstance(value, bool) or not math.isfinite(value) or value <= 0:
35
+ raise ValueError(f"{name} must be positive and finite")
36
+ return value
37
+
38
+
39
+ def _identifier(value: str) -> str:
40
+ if not isinstance(value, str) or not re.fullmatch(r"[A-Za-z0-9_-]{1,128}", value):
41
+ raise ValueError("Expected a job ID returned by MRU")
42
+ return quote(value, safe="")
43
+
44
+
45
+ class _NoRedirect(HTTPRedirectHandler):
46
+ def redirect_request(self, req: Request, fp: IO[bytes], code: int, msg: str,
47
+ headers: Message, newurl: str) -> Request | None:
48
+ return None # Never forward account credentials to another location.
49
+
50
+
51
+ @dataclass(frozen=True)
52
+ class LogEntry:
53
+ """One stdout, stderr, or system log line from a remote job."""
54
+
55
+ sequence: int
56
+ timestamp: str
57
+ stream: str
58
+ message: str
59
+
60
+
61
+ @dataclass(frozen=True)
62
+ class LogPage:
63
+ """Entries after the requested cursor, in ascending sequence order.
64
+
65
+ Pass ``cursor`` as ``after`` to fetch the next page. ``has_more`` means another
66
+ page is available now; it does not indicate whether the job has finished.
67
+ """
68
+
69
+ entries: tuple[LogEntry, ...]
70
+ cursor: int
71
+ has_more: bool
72
+
73
+ @classmethod
74
+ def _from_wire(cls, value: dict[str, Any], *, after: int) -> LogPage:
75
+ try:
76
+ raw_entries, cursor, has_more = value["entries"], value["cursor"], value["has_more"]
77
+ if not isinstance(raw_entries, list) or type(cursor) is not int or type(has_more) is not bool:
78
+ raise ValueError()
79
+ entries, previous = [], after
80
+ for item in raw_entries:
81
+ entry = LogEntry(item["sequence"], item["timestamp"], item["stream"], item["message"])
82
+ if (type(entry.sequence) is not int or entry.sequence <= previous
83
+ or not all(isinstance(text, str) for text in (entry.timestamp, entry.stream, entry.message))):
84
+ raise ValueError()
85
+ entries.append(entry)
86
+ previous = entry.sequence
87
+ if cursor < previous:
88
+ raise ValueError()
89
+ except (KeyError, TypeError, ValueError):
90
+ raise ProtocolError("MRU returned an invalid log page") from None
91
+ return cls(tuple(entries), cursor, has_more)
92
+
93
+
94
+ class Client:
95
+ """Launch and monitor remote jobs using an MRU account token.
96
+
97
+ This client runs on the customer's computer, independently of Job/Session
98
+ inside the remote runtime. Constructing it makes no network requests.
99
+ """
100
+
101
+ def __init__(self, *, token: str, api_url: str, timeout: float = 30.0):
102
+ if (not isinstance(token, str) or not token
103
+ or any(not 33 <= ord(c) <= 126 for c in token)):
104
+ raise AuthenticationError("unauthenticated", "Supply a nonempty MRU API token without whitespace")
105
+ url = api_url
106
+ if not isinstance(url, str) or any(ord(c) <= 32 or ord(c) == 127 for c in url):
107
+ raise ValueError("Use an HTTPS MRU API origin")
108
+ try:
109
+ parsed = urlsplit(url)
110
+ parsed.port # Reject malformed ports before attempting a connection.
111
+ valid = (parsed.hostname and not parsed.username and not parsed.password
112
+ and not parsed.query and not parsed.fragment and parsed.path in ("", "/")
113
+ and (parsed.scheme == "https" or
114
+ parsed.scheme == "http" and parsed.hostname in ("localhost", "127.0.0.1", "::1")))
115
+ except ValueError:
116
+ valid = False
117
+ if not valid:
118
+ raise ValueError("Use an HTTPS MRU API origin (HTTP is allowed only for localhost)")
119
+ self.api_url = url.rstrip("/")
120
+ self.timeout = _positive(timeout, "timeout")
121
+ self._token = token
122
+ self._opener = build_opener(_NoRedirect)
123
+
124
+ @classmethod
125
+ def from_env(cls, *, timeout: float = 30.0) -> Client:
126
+ """Read MRU_API_TOKEN and MRU_API_URL; no CLI state or job credentials."""
127
+ token, url = os.environ.get("MRU_API_TOKEN"), os.environ.get("MRU_API_URL")
128
+ if not token:
129
+ raise AuthenticationError("unauthenticated", "Set MRU_API_TOKEN")
130
+ if not url:
131
+ raise ValueError("Set MRU_API_URL to a server implementing the Python launch contract")
132
+ return cls(token=token, api_url=url, timeout=timeout)
133
+
134
+ def _request(self, path: str, body: dict[str, Any] | None = None, *, timeout: float | None = None,
135
+ request_id: str | None = None) -> dict[str, Any]:
136
+ data = None if body is None else json.dumps(body, allow_nan=False).encode("utf-8")
137
+ if data is not None and len(data) > 16 * 1024 * 1024:
138
+ raise ValueError("Launch request exceeds the 16 MiB limit")
139
+ request = Request(self.api_url + path, data=data, headers={
140
+ "Authorization": "Bearer " + self._token,
141
+ "Content-Type": "application/json",
142
+ "Accept": "application/json",
143
+ **({"Idempotency-Key": request_id} if request_id else {}),
144
+ })
145
+ try:
146
+ with self._opener.open(request, timeout=self.timeout if timeout is None else timeout) as response:
147
+ result = json.load(response)
148
+ except HTTPError as exc:
149
+ try:
150
+ with exc:
151
+ reply = json.load(exc)
152
+ detail = reply.get("error", {}).get("message") if isinstance(reply, dict) and isinstance(reply.get("error"), dict) else None
153
+ except (ValueError, OSError, HTTPException):
154
+ detail = None
155
+ message = detail if isinstance(detail, str) else f"MRU API returned HTTP {exc.code}"
156
+ message = message.replace(self._token, "[redacted]")
157
+ if exc.code == 401:
158
+ raise AuthenticationError("unauthenticated", message) from None
159
+ if exc.code == 403:
160
+ raise AuthorizationError("permission_denied", message) from None
161
+ raise APIError(exc.code, message) from None
162
+ except (URLError, OSError, HTTPException):
163
+ raise PlatformConnectionError("MRU connection failed; the remote outcome may be unknown") from None
164
+ except (ValueError, UnicodeError):
165
+ raise ProtocolError("MRU API returned invalid JSON") from None
166
+ if not isinstance(result, dict):
167
+ raise ProtocolError("MRU API response must be an object")
168
+ return result
169
+
170
+ def me(self) -> dict[str, Any]:
171
+ """Verify credentials and return the authenticated account response."""
172
+ return self._request("/v1/me")
173
+
174
+ def launch(self, entrypoint: str | Path, *, robot_id: str,
175
+ params: Mapping[str, Any] | None = None, dependencies: Sequence[str] = (),
176
+ requirements: Requirements | None = None, source_dir: str | Path | None = None,
177
+ files: Sequence[str] = (), python_version: str = "3.12",
178
+ timeout_seconds: int = 600, request_id: str | None = None) -> RemoteJob:
179
+ """Upload selected source files and queue a Python job on MRU.
180
+
181
+ Nothing is imported or executed locally. Only entrypoint and explicit
182
+ files are uploaded. Requires the proposed backend contract in launch-api.md.
183
+ """
184
+ request_id = uuid.uuid4().hex if request_id is None else request_id
185
+ if not isinstance(request_id, str) or not re.fullmatch(r"[A-Za-z0-9_-]{16,80}", request_id):
186
+ raise ValueError("request_id must contain 16–80 ASCII letters, digits, underscores, or hyphens")
187
+ if not isinstance(robot_id, str) or not robot_id.strip() or len(robot_id) > 128:
188
+ raise ValueError("robot_id must be a nonempty string of at most 128 characters")
189
+ if not isinstance(dependencies, Sequence) or isinstance(dependencies, (str, bytes)) or any(
190
+ not isinstance(dep, str) or not dep.strip() or dep.startswith("-") or "\n" in dep or "\r" in dep
191
+ for dep in dependencies):
192
+ raise ValueError("dependencies must contain individual package requirements, not pip options")
193
+ if not isinstance(python_version, str) or not re.fullmatch(r"3\.(?:1[1-9]|[2-9][0-9])", python_version):
194
+ raise ValueError("python_version must name Python 3.11 or newer, e.g. 3.12")
195
+ if type(timeout_seconds) is not int or timeout_seconds <= 0:
196
+ raise ValueError("timeout_seconds must be a positive integer")
197
+ if params is not None and (not isinstance(params, Mapping) or any(not isinstance(k, str) for k in params)):
198
+ raise ValueError("params must be an object with string keys")
199
+ entry_name, source = bundle_source(entrypoint, source_dir=source_dir, files=files)
200
+ payload = {
201
+ "protocol": "mru.python-job/v1", "entrypoint": entry_name,
202
+ "robot_id": robot_id, "params": dict(params or {}),
203
+ "dependencies": list(dependencies), "python_version": python_version,
204
+ "requirements": (requirements or Requirements()).to_wire(),
205
+ "timeout_seconds": timeout_seconds, "source": source,
206
+ }
207
+ try:
208
+ view = self._request("/v1/jobs", payload, request_id=request_id)
209
+ return RemoteJob(self, view, request_id=request_id)
210
+ except (PlatformConnectionError, ProtocolError) as exc:
211
+ raise LaunchError(request_id, "Launch outcome is unknown; retry unchanged source and options with this request_id") from exc
212
+ except APIError as exc:
213
+ if exc.status_code >= 500:
214
+ raise LaunchError(request_id, "Launch outcome is unknown; retry unchanged source and options with this request_id") from exc
215
+ raise
216
+
217
+ def get_job(self, job_id: str) -> RemoteJob:
218
+ """Reconnect to a previously launched job without starting another one."""
219
+ identifier = _identifier(job_id)
220
+ return RemoteJob(self, self._request("/v1/jobs/" + identifier), expected_id=job_id)
221
+
222
+
223
+ class RemoteJob:
224
+ """Handle for a remote Python job. Status is cached until refresh/wait/cancel."""
225
+
226
+ def __init__(self, client: Client, view: dict[str, Any], *, request_id: str | None = None,
227
+ expected_id: str | None = None):
228
+ self._client = client
229
+ self.request_id = request_id
230
+ self._update(view, expected_id=expected_id)
231
+
232
+ def _update(self, view: dict[str, Any], *, expected_id: str | None = None) -> None:
233
+ try:
234
+ identifier, status = view["job"]["id"], view["job"]["status"]
235
+ _identifier(identifier)
236
+ if not isinstance(status, str) or not status or expected_id is not None and identifier != expected_id:
237
+ raise ValueError()
238
+ except (KeyError, TypeError, ValueError):
239
+ raise ProtocolError("MRU returned an invalid job response") from None
240
+ self.id: str = identifier
241
+ self.status: str = status
242
+ self.details: dict[str, Any] = dict(view["job"]) # Full job record from the last response.
243
+
244
+ def refresh(self) -> RemoteJob:
245
+ view = self._client._request("/v1/jobs/" + quote(self.id, safe=""))
246
+ self._update(view, expected_id=self.id)
247
+ return self
248
+
249
+ def wait(self, *, timeout: float = 300.0, poll_interval: float = 2.0) -> RemoteJob:
250
+ """Wait for terminal status; timeout leaves the remote job running."""
251
+ deadline = time.monotonic() + _positive(timeout, "timeout")
252
+ _positive(poll_interval, "poll_interval")
253
+ while self.status not in TERMINAL_STATUSES:
254
+ remaining = deadline - time.monotonic()
255
+ if remaining <= 0:
256
+ raise TimeoutError(f"Stopped waiting for {self.id}; the remote job has not been cancelled")
257
+ try:
258
+ view = self._client._request("/v1/jobs/" + quote(self.id, safe=""),
259
+ timeout=min(self._client.timeout, remaining))
260
+ except PlatformConnectionError:
261
+ if time.monotonic() >= deadline:
262
+ raise TimeoutError(f"Stopped waiting for {self.id}; the remote job has not been cancelled") from None
263
+ raise
264
+ self._update(view, expected_id=self.id)
265
+ if self.status not in TERMINAL_STATUSES:
266
+ time.sleep(min(poll_interval, max(0, deadline - time.monotonic())))
267
+ return self
268
+
269
+ def cancel(self) -> RemoteJob:
270
+ """Request cancellation; use refresh/wait to observe the resulting status."""
271
+ view = self._client._request("/v1/jobs/" + quote(self.id, safe="") + "/cancel", {})
272
+ self._update(view, expected_id=self.id)
273
+ return self
274
+
275
+ def logs(self, *, after: int = 0, limit: int = 100) -> LogPage:
276
+ """Fetch one page of log entries with sequence greater than ``after``."""
277
+ if type(after) is not int or after < 0 or type(limit) is not int or not 1 <= limit <= 100:
278
+ raise ValueError("after must be a nonnegative integer; limit must be between 1 and 100")
279
+ view = self._client._request("/v1/jobs/" + quote(self.id, safe="")
280
+ + "/logs?" + urlencode({"after": after, "limit": limit}))
281
+ return LogPage._from_wire(view, after=after)
mru/errors.py ADDED
@@ -0,0 +1,53 @@
1
+ """Errors surfaced to job authors."""
2
+
3
+
4
+ class MRUError(Exception):
5
+ """Base SDK error."""
6
+
7
+
8
+ class SessionClosed(MRUError):
9
+ pass
10
+
11
+
12
+ class JobCancelled(MRUError):
13
+ pass
14
+
15
+
16
+ class DeadlineExceeded(MRUError):
17
+ """The outcome may be unknown; the SDK requests a stop and ends the session."""
18
+
19
+
20
+ class CommandRejected(MRUError):
21
+ def __init__(self, code: str, message: str):
22
+ self.code = code
23
+ super().__init__(message)
24
+
25
+
26
+ class ProtocolError(MRUError):
27
+ pass
28
+
29
+
30
+ class AuthenticationError(CommandRejected):
31
+ """The job credential is missing, invalid, expired, or revoked."""
32
+
33
+
34
+ class AuthorizationError(CommandRejected):
35
+ """The authenticated job lacks permission for this robot or action."""
36
+
37
+
38
+ class APIError(MRUError):
39
+ def __init__(self, status_code: int, message: str):
40
+ self.status_code = status_code
41
+ super().__init__(message)
42
+
43
+
44
+ class PlatformConnectionError(MRUError):
45
+ """The platform could not be reached or its response was lost."""
46
+
47
+
48
+ class LaunchError(PlatformConnectionError):
49
+ """A submission may have succeeded. Preserve request_id before retrying."""
50
+
51
+ def __init__(self, request_id: str, message: str):
52
+ self.request_id = request_id
53
+ super().__init__(f"{message} (request_id={request_id})")
mru/job.py ADDED
@@ -0,0 +1,119 @@
1
+ from __future__ import annotations
2
+
3
+ import asyncio
4
+ import contextlib
5
+ import logging
6
+ import os
7
+ from collections.abc import AsyncIterator, Awaitable, Callable, Mapping, Sequence
8
+ from types import MappingProxyType, TracebackType
9
+ from typing import Any
10
+
11
+ from .auth import Auth
12
+ from .models import ActuatorSpec, ExecutionAck, Observation, Requirements
13
+ from .session import Session
14
+ from .transport import UnixTransport
15
+
16
+
17
+ class Actuator:
18
+ """A named actuator group advertised by the robot. Obtain handles from ``job.actuators``."""
19
+
20
+ def __init__(self, session: Session, name: str, spec: ActuatorSpec):
21
+ self._session, self.name, self.spec = session, name, spec
22
+
23
+ async def move(self, *, positions: Sequence[float], duration_ms: int,
24
+ observation_id: str, timeout: float | None = None) -> ExecutionAck:
25
+ """Execute absolute positions in spec.joints order and spec.unit; await execution."""
26
+ self._session.check_active()
27
+ positions = tuple(positions)
28
+ if "joint_position" not in self.spec.commands:
29
+ raise ValueError("Actuator does not support absolute joint positions")
30
+ self.spec.validate(positions)
31
+ if isinstance(duration_ms, bool) or not isinstance(duration_ms, int) or duration_ms <= 0:
32
+ raise ValueError("duration_ms must be a positive integer")
33
+ if not isinstance(observation_id, str) or not observation_id:
34
+ raise ValueError("An observation_id is required")
35
+ result = await self._session.request(
36
+ "move", timeout=timeout, actuator=self.name, command="joint_position",
37
+ unit=self.spec.unit, positions=positions, duration_ms=duration_ms,
38
+ observation_id=observation_id,
39
+ )
40
+ return ExecutionAck(result["command_id"], result["executed_at_ns"])
41
+
42
+
43
+ class Job:
44
+ """The running job: parameters, observations, actuators, and lifecycle."""
45
+
46
+ def __init__(self, session: Session):
47
+ if session.capabilities is None:
48
+ raise ValueError("Open the session before creating a Job")
49
+ self.session = session
50
+ self.params: Mapping[str, Any] = MappingProxyType(session.params)
51
+ self.actuators: Mapping[str, Actuator] = MappingProxyType({
52
+ name: Actuator(session, name, spec) for name, spec in session.capabilities.actuators.items()
53
+ })
54
+ self.logger = logging.getLogger("mru.job")
55
+ self._finished = False
56
+
57
+ def observations(self) -> AsyncIterator[Observation]:
58
+ """Iterate recent observations; only one consumer is supported per session."""
59
+ return self.session.observations()
60
+
61
+ async def complete(self) -> None:
62
+ """Report success and stop further execution."""
63
+ await self.session.finish()
64
+ self._finished = True
65
+
66
+ async def fail(self, reason: str) -> None:
67
+ """Report failure with a human-readable reason and stop further execution."""
68
+ await self.session.finish(failure=reason)
69
+ self._finished = True
70
+
71
+ async def __aenter__(self) -> Job:
72
+ return self
73
+
74
+ async def __aexit__(self, exc_type: type[BaseException] | None, exc: BaseException | None,
75
+ traceback: TracebackType | None) -> None:
76
+ try:
77
+ if not self._finished:
78
+ if exc is None:
79
+ await self.complete()
80
+ elif isinstance(exc, Exception):
81
+ with contextlib.suppress(Exception):
82
+ await self.fail(str(exc))
83
+ finally:
84
+ await self.session.close()
85
+
86
+
87
+ async def run_async(main: Callable[[Job], Awaitable[None]], *, session: Session | None = None,
88
+ requirements: Requirements | None = None, auth: Auth | None = None) -> None:
89
+ """Run a job; returning normally completes it. Platform cancellation cancels main()."""
90
+ if session is None:
91
+ path = os.environ.get("MRU_SESSION_SOCKET")
92
+ if not path:
93
+ raise RuntimeError("Set MRU_SESSION_SOCKET or provide a Session (e.g. FakeRobot.session())")
94
+ if auth is None:
95
+ auth = Auth.from_env()
96
+ session = Session(await UnixTransport.connect(path))
97
+ await session.open(requirements, auth=auth)
98
+ async with Job(session) as job:
99
+ async def call_main() -> None:
100
+ await main(job) # Accept any awaitable, not only coroutines.
101
+
102
+ task = asyncio.create_task(call_main(), name="mru-job")
103
+ interrupted = asyncio.create_task(session.interrupted.wait())
104
+ try:
105
+ await asyncio.wait({task, interrupted}, return_when=asyncio.FIRST_COMPLETED)
106
+ if interrupted.done():
107
+ session.check_active()
108
+ await task
109
+ finally:
110
+ for pending in (task, interrupted):
111
+ if not pending.done():
112
+ pending.cancel()
113
+ await asyncio.gather(task, interrupted, return_exceptions=True)
114
+
115
+
116
+ def run(main: Callable[[Job], Awaitable[None]], *, session: Session | None = None,
117
+ requirements: Requirements | None = None, auth: Auth | None = None) -> None:
118
+ """Blocking form of :func:`run_async` for scripts without an event loop."""
119
+ asyncio.run(run_async(main, session=session, requirements=requirements, auth=auth))
mru/models.py ADDED
@@ -0,0 +1,136 @@
1
+ """Wire-independent robot data. Position commands use each group's declared unit."""
2
+
3
+ from __future__ import annotations
4
+
5
+ import base64
6
+ import math
7
+ from dataclasses import asdict, dataclass, field
8
+ from typing import Any, Mapping
9
+
10
+
11
+ @dataclass(frozen=True)
12
+ class JointState:
13
+ names: tuple[str, ...]
14
+ positions: tuple[float, ...]
15
+ unit: str = "rad"
16
+
17
+ def __post_init__(self) -> None:
18
+ object.__setattr__(self, "names", tuple(self.names))
19
+ object.__setattr__(self, "positions", tuple(self.positions))
20
+
21
+
22
+ @dataclass(frozen=True)
23
+ class EncodedImage:
24
+ data: bytes
25
+ encoding: str # e.g. image/jpeg or image/png; decoding belongs to the job
26
+
27
+
28
+ @dataclass(frozen=True)
29
+ class Observation:
30
+ id: str
31
+ timestamp_ns: int # Agent capture time, Unix epoch nanoseconds
32
+ joint_states: Mapping[str, JointState] = field(default_factory=dict)
33
+ images: Mapping[str, EncodedImage] = field(default_factory=dict)
34
+ sensors: Mapping[str, float] = field(default_factory=dict)
35
+
36
+ def to_wire(self) -> dict[str, Any]:
37
+ return {
38
+ "id": self.id,
39
+ "timestamp_ns": self.timestamp_ns,
40
+ "joint_states": {name: asdict(value) for name, value in self.joint_states.items()},
41
+ "images": {
42
+ name: {"encoding": value.encoding, "data": base64.b64encode(value.data).decode("ascii")}
43
+ for name, value in self.images.items()
44
+ },
45
+ "sensors": dict(self.sensors),
46
+ }
47
+
48
+ @classmethod
49
+ def from_wire(cls, value: dict[str, Any]) -> Observation:
50
+ return cls(
51
+ id=value["id"], timestamp_ns=value["timestamp_ns"],
52
+ joint_states={
53
+ name: JointState(state["names"], state["positions"], state["unit"])
54
+ for name, state in value.get("joint_states", {}).items()
55
+ },
56
+ images={
57
+ name: EncodedImage(base64.b64decode(image["data"], validate=True), image["encoding"])
58
+ for name, image in value.get("images", {}).items()
59
+ },
60
+ sensors=value.get("sensors", {}),
61
+ )
62
+
63
+
64
+ @dataclass(frozen=True)
65
+ class ActuatorSpec:
66
+ joints: tuple[str, ...]
67
+ unit: str = "rad"
68
+ commands: tuple[str, ...] = ("joint_position",)
69
+ limits: tuple[tuple[float, float], ...] = ()
70
+
71
+ def __post_init__(self) -> None:
72
+ # Accept lists from callers and JSON; comparisons and hashing rely on tuples.
73
+ object.__setattr__(self, "joints", tuple(self.joints))
74
+ object.__setattr__(self, "commands", tuple(self.commands))
75
+ object.__setattr__(self, "limits", tuple((low, high) for low, high in self.limits))
76
+
77
+ def validate(self, positions: tuple[float, ...]) -> None:
78
+ if len(positions) != len(self.joints):
79
+ raise ValueError(f"Expected {len(self.joints)} positions in joint order {self.joints}")
80
+ if not all(math.isfinite(value) for value in positions):
81
+ raise ValueError("Positions must be finite")
82
+ if self.limits and len(self.limits) != len(self.joints):
83
+ raise ValueError("Invalid actuator capability limits")
84
+ if self.limits and any(not low <= value <= high for value, (low, high) in zip(positions, self.limits)):
85
+ raise ValueError("Position exceeds advertised limits")
86
+
87
+
88
+ @dataclass(frozen=True)
89
+ class RobotCapabilities:
90
+ sensors: Mapping[str, str] # name -> joint_state, encoded_image, scalar
91
+ actuators: Mapping[str, ActuatorSpec]
92
+
93
+ def to_wire(self) -> dict[str, Any]:
94
+ return {"sensors": dict(self.sensors), "actuators": {k: asdict(v) for k, v in self.actuators.items()}}
95
+
96
+ @classmethod
97
+ def from_wire(cls, value: dict[str, Any]) -> RobotCapabilities:
98
+ return cls(value["sensors"], {
99
+ name: ActuatorSpec(spec["joints"], spec["unit"], spec["commands"], spec.get("limits", ()))
100
+ for name, spec in value["actuators"].items()
101
+ })
102
+
103
+
104
+ @dataclass(frozen=True)
105
+ class ActuatorRequirement:
106
+ unit: str
107
+ command: str = "joint_position"
108
+ joints: tuple[str, ...] = ()
109
+
110
+ def __post_init__(self) -> None:
111
+ object.__setattr__(self, "joints", tuple(self.joints))
112
+
113
+
114
+ @dataclass(frozen=True)
115
+ class Requirements:
116
+ sensors: Mapping[str, str] = field(default_factory=dict)
117
+ actuators: Mapping[str, ActuatorRequirement] = field(default_factory=dict)
118
+
119
+ def to_wire(self) -> dict[str, Any]:
120
+ return {"sensors": dict(self.sensors), "actuators": {k: asdict(v) for k, v in self.actuators.items()}}
121
+
122
+ def validate(self, capabilities: RobotCapabilities) -> None:
123
+ for name, kind in self.sensors.items():
124
+ if capabilities.sensors.get(name) != kind:
125
+ raise ValueError(f"Missing or incompatible sensor: {name} ({kind})")
126
+ for name, requirement in self.actuators.items():
127
+ spec = capabilities.actuators.get(name)
128
+ if (spec is None or spec.unit != requirement.unit or requirement.command not in spec.commands
129
+ or (requirement.joints and requirement.joints != spec.joints)):
130
+ raise ValueError(f"Missing or incompatible actuator: {name}")
131
+
132
+
133
+ @dataclass(frozen=True)
134
+ class ExecutionAck:
135
+ command_id: str
136
+ executed_at_ns: int
mru/py.typed ADDED
File without changes