marin-iris 0.99__py3-none-any.whl

iris/__init__.py

@@ -0,0 +1,2 @@
+# Copyright The Marin Authors
+# SPDX-License-Identifier: Apache-2.0

iris/_build_info.py

@@ -0,0 +1,3 @@
+# Auto-generated by scripts/python_libs_package.py during wheel builds.
+
+BUILD_DATE = "2026-05-10"

iris/actor/__init__.py

@@ -0,0 +1,35 @@
+# Copyright The Marin Authors
+# SPDX-License-Identifier: Apache-2.0
+
+"""Actor system for distributed RPC.
+
+For ClusterResolver (namespace-aware controller-based resolution),
+see iris.client.resolver.
+"""
+
+from iris.actor.client import ActorClient
+from iris.actor.pool import ActorPool, BroadcastFuture, CallResult
+from iris.actor.resolver import (
+    ACTOR_ENDPOINT_HEADER,
+    FixedResolver,
+    ProxyResolver,
+    ResolvedEndpoint,
+    Resolver,
+    ResolveResult,
+)
+from iris.actor.server import ActorId, ActorServer
+
+__all__ = [
+    "ACTOR_ENDPOINT_HEADER",
+    "ActorClient",
+    "ActorId",
+    "ActorPool",
+    "ActorServer",
+    "BroadcastFuture",
+    "CallResult",
+    "FixedResolver",
+    "ProxyResolver",
+    "ResolveResult",
+    "ResolvedEndpoint",
+    "Resolver",
+]

iris/actor/client.py

@@ -0,0 +1,223 @@
+# Copyright The Marin Authors
+# SPDX-License-Identifier: Apache-2.0
+
+"""Actor client for making RPC calls to actor servers.
+
+The ActorClient provides transparent actor discovery and invocation with
+automatic retry logic. Both resolution failures (e.g., actor not yet
+registered) and transient RPC errors are retried up to ``max_call_attempts``
+with exponential backoff.
+
+Example:
+    resolver = ClusterResolver("http://controller:8080")
+    client = ActorClient(resolver, "my-actor")
+    result = client.some_method(arg1, arg2)  # Retries until actor found
+
+Custom backoff behavior:
+    client = ActorClient(
+        resolver, "my-actor",
+        backoff=ExponentialBackoff(initial=0.2, maximum=5.0),
+        max_call_attempts=3,
+    )
+"""
+
+import logging
+import time
+from typing import Any
+
+import cloudpickle
+from connectrpc.code import Code
+from connectrpc.errors import ConnectError
+from rigging.timing import ExponentialBackoff
+
+from iris.actor.resolver import Resolver
+from iris.rpc import actor_pb2
+from iris.rpc.actor_connect import ActorServiceClientSync
+from iris.rpc.errors import call_with_retry
+
+logger = logging.getLogger(__name__)
+
+
+def unwrap_actor_response(resp: actor_pb2.ActorResponse) -> Any:
+    """Unwrap an ActorResponse, raising the embedded exception on error."""
+    if resp.HasField("error"):
+        if resp.error.serialized_exception:
+            raise cloudpickle.loads(resp.error.serialized_exception)
+        raise RuntimeError(f"{resp.error.error_type}: {resp.error.message}")
+    return cloudpickle.loads(resp.serialized_value)
+
+
+class ActorClient:
+    """Actor client with resolver-based discovery.
+
+    By default the client waits forever, i.e. there's no timeout in httpx.
+    Specify ``call_timeout`` to apply a timeout to individual RPC calls.
+    """
+
+    def __init__(
+        self,
+        resolver: Resolver,
+        name: str,
+        call_timeout: float | None = None,
+        max_call_attempts: int = 10,
+        backoff: ExponentialBackoff = ExponentialBackoff(initial=0.5, maximum=10.0, factor=2.0, jitter=0.25),
+    ):
+        """Initialize the actor client.
+
+        Args:
+            resolver: Resolver instance for endpoint discovery
+            name: Name of the actor to invoke
+            call_timeout: Timeout in seconds for individual RPC calls.
+                None (default) means no timeout.
+            max_call_attempts: Maximum number of RPC call attempts (including
+                resolution failures) before giving up.
+            backoff: Exponential backoff configuration for retries between attempts.
+        """
+        self._resolver = resolver
+        self._name = name
+        self._call_timeout = call_timeout
+        self._max_call_attempts = max_call_attempts
+        self._backoff = backoff
+
+        self._rpc_client: ActorServiceClientSync | None = None
+        self._rpc_headers: dict[str, str] = {}
+
+    def rpc_client(self) -> ActorServiceClientSync:
+        """Resolve actor name to an RPC client (single attempt).
+
+        Resolution is attempted once. On failure (empty endpoints or RPC error),
+        the exception propagates to the caller. The outer ``call_with_retry`` in
+        ``_RpcMethod.__call__`` is responsible for retrying.
+
+        Returns:
+            ActorServiceClientSync connected to the resolved endpoint.
+
+        Raises:
+            ConnectError(UNAVAILABLE): If no endpoints are found for the actor.
+        """
+        if self._rpc_client:
+            return self._rpc_client
+
+        logger.info("Resolving name %s via %s", self._name, self._resolver)
+        result = self._resolver.resolve(self._name)
+
+        if result.is_empty:
+            raise ConnectError(
+                Code.UNAVAILABLE,
+                f"No endpoints found for actor '{self._name}'",
+            )
+
+        logger.info(
+            "Resolved actor '%s' to %d endpoint(s)",
+            self._name,
+            len(result.endpoints),
+        )
+        endpoint = result.first()
+        logger.info("First endpoint: url=%s, actor_id=%s", endpoint.url, endpoint.actor_id)
+        self._rpc_headers = dict(endpoint.metadata)
+        self._rpc_client = ActorServiceClientSync(
+            address=endpoint.url,
+            timeout_ms=None if self._call_timeout is None else int(self._call_timeout * 1000),
+            accept_compression=[],
+        )
+        return self._rpc_client
+
+    def _clear_connection(self, _exc: Exception) -> None:
+        self._rpc_client = None
+        self._rpc_headers = {}
+
+    def start_operation(self, method_name: str, *args: Any, **kwargs: Any) -> str:
+        """Start a long-running operation. Returns the operation ID."""
+        call = actor_pb2.ActorCall(
+            method_name=method_name,
+            actor_name=self._name,
+            serialized_args=cloudpickle.dumps(args),
+            serialized_kwargs=cloudpickle.dumps(kwargs),
+        )
+
+        def do_call():
+            client = self.rpc_client()
+            return client.start_operation(call, headers=self._rpc_headers)
+
+        op = call_with_retry(
+            f"{self._name}.start_operation({method_name})",
+            do_call,
+            on_retry=self._clear_connection,
+            max_attempts=self._max_call_attempts,
+            backoff=self._backoff,
+        )
+        return op.operation_id
+
+    def poll_operation_status(self, operation_id: str) -> actor_pb2.Operation:
+        """Single-shot poll of a long-running operation's state."""
+        req = actor_pb2.OperationId(operation_id=operation_id)
+
+        def do_call():
+            return self.rpc_client().get_operation(req, headers=self._rpc_headers)
+
+        return call_with_retry(
+            f"{self._name}.poll_operation_status({operation_id[:8]})",
+            do_call,
+            on_retry=self._clear_connection,
+            max_attempts=self._max_call_attempts,
+            backoff=self._backoff,
+        )
+
+    def get_operation(
+        self,
+        operation_id: str,
+        poll_backoff: ExponentialBackoff | None = None,
+    ) -> actor_pb2.Operation:
+        """Poll a long-running operation until it completes, using exponential backoff."""
+        if poll_backoff is None:
+            poll_backoff = ExponentialBackoff(initial=0.1, maximum=10.0, factor=2.0, jitter=0.25)
+        while True:
+            op = self.poll_operation_status(operation_id)
+            if op.state != actor_pb2.Operation.RUNNING:
+                return op
+            time.sleep(poll_backoff.next_interval())
+
+    def cancel_operation(self, operation_id: str) -> actor_pb2.Operation:
+        """Cancel a long-running operation."""
+        req = actor_pb2.OperationId(operation_id=operation_id)
+
+        def do_call():
+            return self.rpc_client().cancel_operation(req, headers=self._rpc_headers)
+
+        return call_with_retry(
+            f"{self._name}.cancel_operation({operation_id[:8]})",
+            do_call,
+            on_retry=self._clear_connection,
+            max_attempts=self._max_call_attempts,
+            backoff=self._backoff,
+        )
+
+    def __getattr__(self, method_name: str) -> "_RpcMethod":
+        return _RpcMethod(self, method_name)
+
+
+class _RpcMethod:
+    def __init__(self, client: ActorClient, method_name: str):
+        self._client = client
+        self._method_name = method_name
+
+    def __call__(self, *args: Any, **kwargs: Any) -> Any:
+        call = actor_pb2.ActorCall(
+            method_name=self._method_name,
+            actor_name=self._client._name,
+            serialized_args=cloudpickle.dumps(args),
+            serialized_kwargs=cloudpickle.dumps(kwargs),
+        )
+
+        def do_call():
+            client = self._client.rpc_client()
+            resp = client.call(call, headers=self._client._rpc_headers)
+            return unwrap_actor_response(resp)
+
+        return call_with_retry(
+            f"{self._client._name}.{self._method_name}",
+            do_call,
+            on_retry=self._client._clear_connection,
+            max_attempts=self._client._max_call_attempts,
+            backoff=self._client._backoff,
+        )

iris/actor/pool.py

@@ -0,0 +1,281 @@
+# Copyright The Marin Authors
+# SPDX-License-Identifier: Apache-2.0
+
+"""Actor pool for load-balanced and broadcast RPC calls."""
+
+import logging
+import threading
+import time
+from collections.abc import Callable, Iterator
+from concurrent.futures import Future, ThreadPoolExecutor, as_completed
+from dataclasses import dataclass
+from typing import Any, Generic, TypeVar
+
+import cloudpickle
+from rigging.timing import ExponentialBackoff
+
+from iris.actor.client import unwrap_actor_response
+from iris.actor.resolver import ResolvedEndpoint, Resolver, ResolveResult
+from iris.rpc import actor_pb2
+from iris.rpc.actor_connect import ActorServiceClientSync
+from iris.rpc.errors import call_with_retry
+
+logger = logging.getLogger(__name__)
+
+T = TypeVar("T")
+
+
+@dataclass
+class CallResult:
+    """Result of a single call in a broadcast."""
+
+    endpoint: ResolvedEndpoint
+    value: Any | None = None
+    exception: BaseException | None = None
+
+    @property
+    def success(self) -> bool:
+        return self.exception is None
+
+
+class BroadcastFuture(Generic[T]):
+    """Future representing results from a broadcast call to multiple endpoints."""
+
+    def __init__(self, futures: list[tuple[ResolvedEndpoint, Future]]):
+        self._futures = futures
+
+    def wait_all(self, timeout: float | None = None) -> list[CallResult]:
+        results = []
+        for endpoint, future in self._futures:
+            try:
+                value = future.result(timeout=timeout)
+                results.append(CallResult(endpoint=endpoint, value=value))
+            except Exception as e:
+                results.append(CallResult(endpoint=endpoint, exception=e))
+        return results
+
+    def wait_any(self, timeout: float | None = None) -> CallResult:
+        for future in as_completed([f for _, f in self._futures], timeout=timeout):
+            idx = next(i for i, (_, f) in enumerate(self._futures) if f is future)
+            endpoint = self._futures[idx][0]
+            try:
+                value = future.result()
+                return CallResult(endpoint=endpoint, value=value)
+            except Exception as e:
+                return CallResult(endpoint=endpoint, exception=e)
+        raise TimeoutError("No results within timeout")
+
+    def as_completed(self, timeout: float | None = None) -> Iterator[CallResult]:
+        """Iterate over results as they complete."""
+        endpoint_map = {id(f): ep for ep, f in self._futures}
+        for future in as_completed([f for _, f in self._futures], timeout=timeout):
+            endpoint = endpoint_map[id(future)]
+            try:
+                value = future.result()
+                yield CallResult(endpoint=endpoint, value=value)
+            except Exception as e:
+                yield CallResult(endpoint=endpoint, exception=e)
+
+
+class ActorPool(Generic[T]):
+    """Pool of actors for load-balanced and broadcast calls.
+
+    Resolves a pool of endpoints for an actor name and provides methods to
+    distribute calls across them (round-robin) or broadcast to all endpoints.
+
+    Example:
+        >>> pool = ActorPool(resolver, "inference")
+        >>> result = pool.call().predict(data)  # Round-robin to one endpoint
+        >>> broadcast = pool.broadcast().reload_model()  # Send to all endpoints
+        >>> results = broadcast.wait_all()
+    """
+
+    def __init__(
+        self,
+        resolver: Resolver,
+        name: str,
+        timeout: float = 30.0,
+        max_call_attempts: int = 5,
+        backoff: ExponentialBackoff = ExponentialBackoff(initial=0.1, maximum=10.0, factor=2.0, jitter=0.25),
+        resolve_ttl: float = 5.0,
+    ):
+        """Initialize actor pool.
+
+        Args:
+            resolver: Resolver to discover endpoints
+            name: Actor name to resolve
+            timeout: RPC timeout in seconds
+            max_call_attempts: Maximum number of RPC call attempts before giving up.
+            backoff: Exponential backoff configuration for call retries.
+            resolve_ttl: Seconds to cache resolve results before re-querying the resolver
+        """
+        self._resolver = resolver
+        self._name = name
+        self._timeout = timeout
+        self._max_call_attempts = max_call_attempts
+        self._backoff = backoff
+        self._resolve_ttl = resolve_ttl
+        self._endpoint_index = 0
+        self._cached_result: ResolveResult | None = None
+        self._last_resolve_time: float = 0.0
+        self._lock = threading.Lock()
+        self._executor = ThreadPoolExecutor(max_workers=32)
+        self._clients: dict[str, ActorServiceClientSync] = {}
+
+    def _get_client(self, endpoint: ResolvedEndpoint) -> ActorServiceClientSync:
+        """Return a cached client for the endpoint, creating one if needed."""
+        url = endpoint.url
+        with self._lock:
+            client = self._clients.get(url)
+            if client is not None:
+                return client
+            client = ActorServiceClientSync(
+                address=url,
+                timeout_ms=int(self._timeout * 1000),
+                accept_compression=[],
+            )
+            self._clients[url] = client
+            return client
+
+    def _resolve(self) -> ResolveResult:
+        now = time.monotonic()
+        with self._lock:
+            if self._cached_result is not None and (now - self._last_resolve_time) < self._resolve_ttl:
+                return self._cached_result
+
+        result = self._resolver.resolve(self._name)
+        if result.endpoints:
+            with self._lock:
+                self._cached_result = result
+                self._last_resolve_time = time.monotonic()
+        return result
+
+    def _invalidate_resolve_cache(self) -> None:
+        """Force the next _resolve() call to re-query the resolver."""
+        with self._lock:
+            self._last_resolve_time = 0.0
+            self._cached_result = None
+
+    def _evict_client(self, url: str) -> None:
+        """Remove and close a cached client so it is recreated on next use."""
+        with self._lock:
+            client = self._clients.pop(url, None)
+        if client is not None:
+            try:
+                client.close()
+            except Exception:
+                logger.debug("Error closing evicted client for %s", url, exc_info=True)
+
+    def _get_next_endpoint(self) -> ResolvedEndpoint:
+        """Get the next endpoint in round-robin order.
+
+        Thread-safe: uses a lock to protect the endpoint index.
+        """
+        endpoints = self._resolve().endpoints
+        with self._lock:
+            if not endpoints:
+                raise RuntimeError(f"No endpoints for '{self._name}'")
+            endpoint = endpoints[self._endpoint_index % len(endpoints)]
+            self._endpoint_index += 1
+            return endpoint
+
+    def shutdown(self) -> None:
+        self._executor.shutdown(wait=True)
+        with self._lock:
+            clients = list(self._clients.values())
+            self._clients.clear()
+        for client in clients:
+            try:
+                client.close()
+            except Exception:
+                logger.debug("Error closing client during shutdown", exc_info=True)
+
+    def __enter__(self) -> "ActorPool[T]":
+        return self
+
+    def __exit__(self, *args) -> None:
+        self.shutdown()
+
+    @property
+    def size(self) -> int:
+        return len(self._resolve().endpoints)
+
+    @property
+    def endpoints(self) -> list[ResolvedEndpoint]:
+        return list(self._resolve().endpoints)
+
+    def _call_endpoint(
+        self,
+        endpoint: ResolvedEndpoint,
+        method_name: str,
+        args: tuple,
+        kwargs: dict,
+    ) -> Any:
+        client = self._get_client(endpoint)
+
+        call = actor_pb2.ActorCall(
+            method_name=method_name,
+            actor_name=self._name,
+            serialized_args=cloudpickle.dumps(args),
+            serialized_kwargs=cloudpickle.dumps(kwargs),
+        )
+
+        resp = client.call(call)
+        return unwrap_actor_response(resp)
+
+    def call(self) -> "_PoolCallProxy[T]":
+        return _PoolCallProxy(self)
+
+    def broadcast(self) -> "_PoolBroadcastProxy[T]":
+        return _PoolBroadcastProxy(self)
+
+
+class _PoolCallProxy(Generic[T]):
+    def __init__(self, pool: ActorPool[T]):
+        self._pool = pool
+
+    def __getattr__(self, method_name: str) -> Callable[..., Any]:
+        def call(*args, **kwargs):
+            last_url: list[str | None] = [None]
+
+            def do_call():
+                endpoint = self._pool._get_next_endpoint()
+                last_url[0] = endpoint.url
+                return self._pool._call_endpoint(endpoint, method_name, args, kwargs)
+
+            def on_retry(_exc):
+                self._pool._invalidate_resolve_cache()
+                if last_url[0] is not None:
+                    self._pool._evict_client(last_url[0])
+
+            return call_with_retry(
+                f"{self._pool._name}.{method_name}",
+                do_call,
+                on_retry=on_retry,
+                max_attempts=self._pool._max_call_attempts,
+                backoff=self._pool._backoff,
+            )
+
+        return call
+
+
+class _PoolBroadcastProxy(Generic[T]):
+    def __init__(self, pool: ActorPool[T]):
+        self._pool = pool
+
+    def __getattr__(self, method_name: str) -> Callable[..., BroadcastFuture]:
+        def broadcast(*args, **kwargs) -> BroadcastFuture:
+            result = self._pool._resolve()
+            futures = []
+            for endpoint in result.endpoints:
+                future = self._pool._executor.submit(
+                    self._pool._call_endpoint,
+                    endpoint,
+                    method_name,
+                    args,
+                    kwargs,
+                )
+                futures.append((endpoint, future))
+            return BroadcastFuture(futures)
+
+        return broadcast

iris/actor/resolver.py

@@ -0,0 +1,108 @@
+# Copyright The Marin Authors
+# SPDX-License-Identifier: Apache-2.0
+
+"""Core types and resolver implementations for the actor system."""
+
+from dataclasses import dataclass, field
+from typing import Protocol
+
+# Header used by ActorProxy to route requests to the correct actor endpoint.
+# Shared constant between ProxyResolver (client-side) and ActorProxy (server-side).
+ACTOR_ENDPOINT_HEADER = "x-iris-actor-endpoint"
+
+
+@dataclass
+class ResolvedEndpoint:
+    """A single resolved endpoint for an actor."""
+
+    url: str
+    actor_id: str
+    metadata: dict[str, str] = field(default_factory=dict)
+
+
+@dataclass
+class ResolveResult:
+    """Result of resolving an actor name to endpoints."""
+
+    name: str
+    endpoints: list[ResolvedEndpoint] = field(default_factory=list)
+
+    @property
+    def is_empty(self) -> bool:
+        return len(self.endpoints) == 0
+
+    def first(self) -> ResolvedEndpoint:
+        """Get the first endpoint.
+
+        Returns:
+            The first resolved endpoint
+
+        Raises:
+            ValueError: If no endpoints are available
+        """
+        if not self.endpoints:
+            raise ValueError(f"No endpoints for '{self.name}'")
+        return self.endpoints[0]
+
+
+class Resolver(Protocol):
+    """Protocol for resolving actor names to endpoints.
+
+    Implementations:
+    - FixedResolver: Static endpoint mapping
+    - ClusterResolver: Resolves via cluster controller (lives in iris.client)
+    """
+
+    def resolve(self, name: str) -> ResolveResult: ...
+
+
+class FixedResolver:
+    """Resolver with statically configured endpoints.
+
+    Used for testing or when endpoints are known ahead of time.
+    Does not use namespace prefixing since endpoints are static.
+    """
+
+    def __init__(self, endpoints: dict[str, str | list[str]]):
+        """Initialize with a mapping of actor names to URLs."""
+        self._endpoints: dict[str, list[str]] = {}
+        for name, urls in endpoints.items():
+            if isinstance(urls, str):
+                self._endpoints[name] = [urls]
+            else:
+                self._endpoints[name] = list(urls)
+
+    def resolve(self, name: str) -> ResolveResult:
+        urls = self._endpoints.get(name, [])
+        endpoints = [ResolvedEndpoint(url=url, actor_id=f"fixed-{name}-{i}") for i, url in enumerate(urls)]
+        return ResolveResult(name=name, endpoints=endpoints)
+
+
+class ProxyResolver:
+    """Resolver that routes actor calls through the controller's actor proxy.
+
+    Instead of resolving to the actor's direct address, returns the controller
+    URL so all RPCs go through the proxy. The proxy uses the
+    ``X-Iris-Actor-Endpoint`` header to resolve the actual actor endpoint.
+
+    The caller passes the full actor name as registered in the endpoint registry
+    (e.g. ``/user/job/coordinator/actor-0``).
+
+    Args:
+        controller_url: Controller URL (e.g., ``http://localhost:8080``)
+    """
+
+    def __init__(self, controller_url: str):
+        self._controller_url = controller_url.rstrip("/")
+
+    def resolve(self, name: str) -> ResolveResult:
+        return ResolveResult(
+            name=name,
+            endpoints=[
+                ResolvedEndpoint(
+                    url=self._controller_url,
+                    actor_id=f"proxy-{name}",
+                    metadata={ACTOR_ENDPOINT_HEADER: name},
+                )
+            ],
+        )