mmo-runtime 0.1.3__py3-none-any.whl

mmo_runtime/__init__.py

@@ -0,0 +1,55 @@
+"""mmo-runtime: Runtime library for MMO gRPC tools."""
+
+from .client import ToolClient, consume_stream, ensure_running
+from .config import BaseServerConfig, load_config
+from .context import CancelToken, ProgressReporter, RunContext
+from .errors import (
+    Cancelled,
+    DeadlineExceeded,
+    Internal,
+    InvalidRequest,
+    ProviderError,
+    ResourceExhausted,
+    ToolError,
+    to_grpc_status,
+    wrap_unknown,
+)
+from .events import ErrorEnvelope, ProgressEvent, ResultEnvelope
+from .layout import RunLayout
+from .logging import setup_logging
+from .manifest import Manifest
+from .server import serve
+from .worker import run_in_worker
+
+try:
+    from ._version import __version__
+except ImportError:  # pragma: no cover
+    __version__ = "0.0.0+unknown"
+
+__all__ = [
+    "BaseServerConfig",
+    "CancelToken",
+    "Cancelled",
+    "DeadlineExceeded",
+    "ErrorEnvelope",
+    "Internal",
+    "InvalidRequest",
+    "Manifest",
+    "ProgressEvent",
+    "ProgressReporter",
+    "ProviderError",
+    "ResourceExhausted",
+    "ResultEnvelope",
+    "RunContext",
+    "RunLayout",
+    "ToolClient",
+    "ToolError",
+    "consume_stream",
+    "ensure_running",
+    "load_config",
+    "run_in_worker",
+    "serve",
+    "setup_logging",
+    "to_grpc_status",
+    "wrap_unknown",
+]

mmo_runtime/_version.py

@@ -0,0 +1,24 @@
+# file generated by vcs-versioning
+# don't change, don't track in version control
+from __future__ import annotations
+
+__all__ = [
+    "__version__",
+    "__version_tuple__",
+    "version",
+    "version_tuple",
+    "__commit_id__",
+    "commit_id",
+]
+
+version: str
+__version__: str
+__version_tuple__: tuple[int | str, ...]
+version_tuple: tuple[int | str, ...]
+commit_id: str | None
+__commit_id__: str | None
+
+__version__ = version = '0.1.3'
+__version_tuple__ = version_tuple = (0, 1, 3)
+
+__commit_id__ = commit_id = None

mmo_runtime/client.py

@@ -0,0 +1,294 @@
+"""gRPC client with deadline, stream consumption, and ensure_running."""
+
+from __future__ import annotations
+
+import logging
+import os
+import subprocess
+import sys
+import threading
+import time
+from collections.abc import Callable, Iterator, Sequence
+from pathlib import Path
+from typing import TYPE_CHECKING, Any, NoReturn, Self, TypeVar
+
+import grpc
+from grpc_health.v1 import health_pb2, health_pb2_grpc
+
+from .errors import Cancelled, DeadlineExceeded, ErrorCode, Internal, ToolError
+from .events import ProgressEvent, ResultEnvelope
+from .transport import normalize_target
+
+if TYPE_CHECKING:
+    from types import TracebackType
+
+T = TypeVar("T")
+
+
+def _get_target_address(socket_path: str) -> str:
+    """Determine gRPC target address for the current platform.
+
+    On Windows ``socket_path`` is a TCP ``host:port`` (grpc Python has no
+    ``windows://`` named-pipe resolver). On Unix it is a Unix domain socket path.
+    """
+    return normalize_target(socket_path)
+
+
+def consume_stream(
+    stream: Iterator[Any],
+    on_progress: Callable[[ProgressEvent], None] | None = None,
+) -> ResultEnvelope:
+    """
+    Consume a server-streaming RPC, yielding progress and returning final result.
+    Raises ToolError on error status or error envelope.
+    """
+    result: ResultEnvelope = ResultEnvelope(payload=None)
+    has_result = False
+    for msg in stream:
+        # Check for progress/result/error oneof
+        if hasattr(msg, "WhichOneof"):
+            oneof = msg.WhichOneof("payload")
+            if oneof == "progress":
+                prog = msg.progress
+                event = ProgressEvent(
+                    step=prog.step,
+                    done=prog.done,
+                    total=prog.total if prog.total > 0 else None,
+                    message=prog.message if prog.message else None,
+                    meta=dict(prog.meta) if hasattr(prog, "meta") else {},
+                )
+                if on_progress:
+                    on_progress(event)
+            elif oneof == "result":
+                result = ResultEnvelope(
+                    payload=msg.result.payload if hasattr(msg.result, "payload") else msg.result,
+                    artifacts=list(msg.result.artifacts)
+                    if hasattr(msg.result, "artifacts")
+                    else [],
+                    meta=dict(msg.result.meta) if hasattr(msg.result, "meta") else {},
+                )
+                has_result = True
+            elif oneof == "error":
+                err = msg.error
+                raise ToolError(
+                    code=err.code,
+                    step=err.step if err.step else None,
+                    provider=err.provider if err.provider else None,
+                    message=err.message,
+                    retryable=err.retryable,
+                )
+        else:
+            # Fallback: assume message is result
+            result = ResultEnvelope(payload=msg)
+            has_result = True
+
+    if not has_result:
+        raise Internal("Stream ended without result")
+    return result
+
+
+class ToolClient:
+    """Client for tool gRPC service."""
+
+    def __init__(
+        self,
+        socket_path: str,
+        default_deadline: float = 30.0,
+        health_check_timeout: float = 5.0,
+    ):
+        self.socket_path = socket_path
+        self.default_deadline = default_deadline
+        self.health_check_timeout = health_check_timeout
+        self._channel: grpc.Channel | None = None
+        self._health_stub: health_pb2_grpc.HealthStub | None = None
+        self._log = logging.getLogger(__name__)
+        self._channel_lock = threading.RLock()
+
+    @property
+    def channel(self) -> grpc.Channel:
+        """Public lazily-created channel for generated service stubs."""
+        return self._ensure_channel()
+
+    def _ensure_channel(self) -> grpc.Channel:
+        with self._channel_lock:
+            if self._channel is None:
+                target = _get_target_address(self.socket_path)
+                self._channel = grpc.insecure_channel(target)
+                self._health_stub = health_pb2_grpc.HealthStub(self._channel)
+            return self._channel
+
+    def close(self) -> None:
+        with self._channel_lock:
+            if self._channel is not None:
+                self._channel.close()
+                self._channel = None
+                self._health_stub = None
+
+    def __enter__(self) -> Self:
+        return self
+
+    def __exit__(
+        self,
+        exc_type: type[BaseException] | None,
+        exc_val: BaseException | None,
+        exc_tb: TracebackType | None,
+    ) -> None:
+        self.close()
+
+    def health_check(
+        self,
+        service: str = "",
+        *,
+        timeout: float | None = None,
+    ) -> health_pb2.HealthCheckResponse.ServingStatus:
+        self._ensure_channel()
+        try:
+            resp: health_pb2.HealthCheckResponse = self._health_stub.Check(  # type: ignore[attr-defined]
+                health_pb2.HealthCheckRequest(service=service),
+                timeout=self.health_check_timeout if timeout is None else timeout,
+            )
+            # grpc_health pb2 ships untyped; status enum member is Unknown to pyright
+            return resp.status  # type: ignore[return-value]
+        except grpc.RpcError as e:
+            if e.code() in (grpc.StatusCode.UNAVAILABLE, grpc.StatusCode.DEADLINE_EXCEEDED):
+                return health_pb2.HealthCheckResponse.ServingStatus.SERVICE_UNKNOWN
+            raise
+
+    def wait_healthy(self, timeout: float = 30.0, interval: float = 0.5) -> bool:
+        deadline = time.monotonic() + timeout
+        while time.monotonic() < deadline:
+            remaining = min(timeout, max(0, deadline - time.monotonic()))
+            if (
+                self.health_check(timeout=min(self.health_check_timeout, remaining))
+                == health_pb2.HealthCheckResponse.ServingStatus.SERVING
+            ):
+                return True
+            time.sleep(max(0, min(interval, deadline - time.monotonic())))
+        return False
+
+    def call_unary(
+        self,
+        method: Callable[..., Any],
+        request: Any,
+        deadline: float | None = None,
+        *,
+        metadata: Sequence[tuple[str, str | bytes]] | None = None,
+    ) -> Any:
+        """Call a unary RPC with deadline."""
+        self._ensure_channel()
+        try:
+            kwargs: dict[str, Any] = {
+                "timeout": self.default_deadline if deadline is None else deadline,
+            }
+            if metadata is not None:
+                kwargs["metadata"] = metadata
+            return method(request, **kwargs)
+        except grpc.RpcError as e:
+            self._handle_grpc_error(e)
+
+    def call_stream(
+        self,
+        method: Callable[..., Iterator[Any]],
+        request: Any,
+        on_progress: Callable[[ProgressEvent], None] | None = None,
+        deadline: float | None = None,
+        *,
+        metadata: Sequence[tuple[str, str | bytes]] | None = None,
+    ) -> ResultEnvelope:
+        """Call a server-streaming RPC with deadline."""
+        self._ensure_channel()
+        try:
+            kwargs: dict[str, Any] = {
+                "timeout": self.default_deadline if deadline is None else deadline,
+            }
+            if metadata is not None:
+                kwargs["metadata"] = metadata
+            stream = method(request, **kwargs)
+            try:
+                return consume_stream(stream, on_progress)
+            finally:
+                cancel = getattr(stream, "cancel", None)
+                if callable(cancel):
+                    cancel()
+        except grpc.RpcError as e:
+            self._handle_grpc_error(e)
+
+    def _handle_grpc_error(self, e: grpc.RpcError) -> NoReturn:
+        status = e.code()
+        message = e.details() or f"RPC error: {status}"
+        if status == grpc.StatusCode.DEADLINE_EXCEEDED:
+            raise DeadlineExceeded(message) from e
+        if status == grpc.StatusCode.CANCELLED:
+            raise Cancelled(message) from e
+        code = {
+            grpc.StatusCode.INVALID_ARGUMENT: ErrorCode.INVALID_REQUEST,
+            grpc.StatusCode.RESOURCE_EXHAUSTED: ErrorCode.RESOURCE_EXHAUSTED,
+            grpc.StatusCode.UNKNOWN: ErrorCode.PROVIDER_ERROR,
+            grpc.StatusCode.UNAVAILABLE: ErrorCode.PROVIDER_ERROR,
+        }.get(status, ErrorCode.INTERNAL)
+        raise ToolError(
+            code=code,
+            message=message,
+            retryable=status in (grpc.StatusCode.RESOURCE_EXHAUSTED, grpc.StatusCode.UNAVAILABLE),
+        ) from e
+
+
+# Process management for ensure_running
+_spawn_locks: dict[str, threading.Lock] = {}
+_spawn_locks_lock = threading.Lock()
+
+
+def get_spawn_lock(socket_path: str) -> threading.Lock:
+    with _spawn_locks_lock:
+        if socket_path not in _spawn_locks:
+            _spawn_locks[socket_path] = threading.Lock()
+        return _spawn_locks[socket_path]
+
+
+def ensure_running(
+    tool_name: str,
+    command: list[str],
+    socket_path: str,
+    timeout: float = 30.0,
+    env: dict[str, str] | None = None,
+    cwd: Path | None = None,
+) -> ToolClient:
+    """
+    Ensure tool server is running. Spawns if not healthy.
+    Uses a per-socket lock to avoid duplicate spawns.
+    """
+    client = ToolClient(socket_path)
+    if client.wait_healthy(timeout=2.0):
+        return client
+
+    lock = get_spawn_lock(socket_path)
+    with lock:
+        # Double-check after acquiring lock
+        if client.wait_healthy(timeout=2.0):
+            return client
+
+        # Spawn server
+        proc_env = os.environ.copy()
+        if env:
+            proc_env.update(env)
+        # On Windows, use CREATE_NEW_PROCESS_GROUP for proper termination
+        creationflags = 0
+        if sys.platform == "win32":
+            creationflags = subprocess.CREATE_NEW_PROCESS_GROUP
+        proc = subprocess.Popen(
+            command,
+            env=proc_env,
+            cwd=cwd,
+            creationflags=creationflags,
+        )
+
+        # Wait for healthy
+        if not client.wait_healthy(timeout=timeout):
+            proc.terminate()
+            try:
+                proc.wait(timeout=5)
+            except subprocess.TimeoutExpired:
+                proc.kill()
+            raise RuntimeError(f"Tool {tool_name} failed to start within {timeout}s")
+
+    return client

mmo_runtime/config.py

@@ -0,0 +1,90 @@
+"""Configuration loading with precedence: defaults < yaml < env < CLI."""
+
+from __future__ import annotations
+
+import os
+import sys
+from pathlib import Path
+from typing import Any
+
+from pydantic import Field, field_validator
+from pydantic_settings import BaseSettings, SettingsConfigDict
+
+
+class BaseServerConfig(BaseSettings):
+    """Base configuration for tool servers."""
+
+    model_config = SettingsConfigDict(
+        env_prefix="TOOL_",
+        env_nested_delimiter="__",
+        case_sensitive=False,
+        extra="ignore",
+    )
+
+    # Socket
+    socket_path: str = Field(
+        default_factory=lambda: "127.0.0.1:50051" if sys.platform == "win32" else "/tmp/tool.sock",
+        description="gRPC host:port or Unix domain socket path",
+    )
+    # Concurrency
+    max_concurrent_jobs: int = Field(default=1, ge=1, description="Max concurrent jobs")
+    max_queue_size: int = Field(default=4, ge=0, description="Max queued jobs (0 = no queue)")
+
+    @property
+    def thread_pool_size(self) -> int:
+        """Total thread pool size: workers + queue + buffer for health/gRPC internals."""
+        return self.max_concurrent_jobs + self.max_queue_size + 4
+
+    # Timeouts
+    idle_timeout_sec: int = Field(
+        default=300, ge=0, description="Idle timeout before shutdown (0 = disabled)"
+    )
+    job_timeout_sec: int = Field(default=0, ge=0, description="Per-job timeout (0 = no limit)")
+    # Device
+    device: str = Field(default="cpu", description="Compute device (cpu, cuda, mps, etc.)")
+    # Paths
+    runs_root: Path = Field(default=Path("runs"), description="Root directory for run outputs")
+    # Logging
+    log_level: str = Field(default="INFO", description="Log level")
+    log_json: bool = Field(default=False, description="Output JSON logs")
+
+    @field_validator("runs_root", mode="before")
+    @classmethod
+    def expand_runs_root(cls, v: str | Path) -> Path:
+        return Path(v).expanduser().resolve()
+
+    @field_validator("socket_path", mode="before")
+    @classmethod
+    def expand_socket_path(cls, v: str) -> str:
+        return os.path.expanduser(v)
+
+
+def load_config(
+    config_cls: type[BaseServerConfig],
+    yaml_path: Path | None = None,
+    cli_overrides: dict[str, Any] | None = None,
+) -> BaseServerConfig:
+    """
+    Load config with precedence: defaults < yaml < env < CLI.
+    """
+    yaml_data: dict[str, Any] = {}
+    if yaml_path and yaml_path.exists():
+        import yaml
+
+        with open(yaml_path) as f:
+            yaml_data = yaml.safe_load(f) or {}
+
+    # pydantic-settings ranks init kwargs above env vars, so yaml and CLI
+    # (both passed as init kwargs) would each beat env. To get yaml < env < CLI:
+    # drop yaml keys that the environment also provides; keep CLI keys so they
+    # win over env. Only flat (non-nested) yaml keys are env-matched.
+    cli = {k: v for k, v in (cli_overrides or {}).items() if v is not None}
+    merged: dict[str, Any] = {**yaml_data, **cli}
+
+    prefix = str(config_cls.model_config.get("env_prefix", ""))
+    env_keys = {k.upper() for k in os.environ}
+    for key in yaml_data:
+        if key not in cli and f"{prefix}{key}".upper() in env_keys:
+            merged.pop(key)
+
+    return config_cls(**merged)

mmo_runtime/context.py

@@ -0,0 +1,109 @@
+"""Context management: cancellation, progress reporting, run context."""
+
+from __future__ import annotations
+
+import logging
+import threading
+from pathlib import Path
+from queue import Full, Queue
+from typing import Any
+
+from .events import ProgressEvent
+
+
+class CancelToken:
+    """Thread-safe cancellation token."""
+
+    def __init__(self):
+        self._event = threading.Event()
+
+    def set(self) -> None:
+        self._event.set()
+
+    def is_set(self) -> bool:
+        return self._event.is_set()
+
+    def raise_if_cancelled(self) -> None:
+        if self._event.is_set():
+            from .errors import Cancelled
+
+            raise Cancelled("Operation cancelled")
+
+    def wait(self, timeout: float | None = None) -> bool:
+        return self._event.wait(timeout)
+
+
+class ProgressReporter:
+    """Progress reporter interface. Default implementation is no-op."""
+
+    def update(
+        self,
+        step: str,
+        done: int,
+        total: int | None = None,
+        message: str | None = None,
+        **meta: Any,
+    ) -> None:
+        pass
+
+
+class QueueProgressReporter(ProgressReporter):
+    """Progress reporter that pushes events to a queue for server-streaming."""
+
+    def __init__(self, queue: Queue[ProgressEvent], *, cancel: CancelToken | None = None):
+        self._queue = queue
+        self._cancel = cancel
+        self._closed = threading.Event()
+
+    def close(self) -> None:
+        """Unblock a producer when its consumer has gone away."""
+        self._closed.set()
+
+    def update(
+        self,
+        step: str,
+        done: int,
+        total: int | None = None,
+        message: str | None = None,
+        **meta: Any,
+    ) -> None:
+        event = ProgressEvent(step=step, done=done, total=total, message=message, meta=meta)
+        while not self._closed.is_set():
+            if self._cancel is not None:
+                self._cancel.raise_if_cancelled()
+            try:
+                self._queue.put(event, timeout=0.05)
+                return
+            except Full:
+                continue
+        from .errors import Cancelled
+
+        raise Cancelled("Progress consumer closed")
+
+
+class RunContext:
+    """Execution context for a single tool run."""
+
+    def __init__(
+        self,
+        run_id: str,
+        out_dir: Path,
+        progress: ProgressReporter | None = None,
+        cancel: CancelToken | None = None,
+        log: logging.Logger | None = None,
+    ):
+        self.run_id = run_id
+        self.out_dir = out_dir
+        self.progress = progress or ProgressReporter()
+        self.cancel = cancel or CancelToken()
+        self.log = log or logging.getLogger(__name__)
+
+    def child(self, step: str) -> RunContext:
+        """Create a child context with a step prefix."""
+        return RunContext(
+            run_id=self.run_id,
+            out_dir=self.out_dir / step,
+            progress=self.progress,
+            cancel=self.cancel,
+            log=self.log.getChild(step),
+        )

mmo_runtime/errors.py

@@ -0,0 +1,125 @@
+"""Error definitions for mmo-runtime."""
+
+from __future__ import annotations
+
+from enum import Enum
+
+import grpc
+
+
+class ErrorCode(str, Enum):
+    INVALID_REQUEST = "INVALID_REQUEST"
+    PROVIDER_ERROR = "PROVIDER_ERROR"
+    CANCELLED = "CANCELLED"
+    DEADLINE_EXCEEDED = "DEADLINE_EXCEEDED"
+    INTERNAL = "INTERNAL"
+    RESOURCE_EXHAUSTED = "RESOURCE_EXHAUSTED"
+
+
+_CODE_TO_GRPC = {
+    ErrorCode.INVALID_REQUEST: grpc.StatusCode.INVALID_ARGUMENT,
+    ErrorCode.PROVIDER_ERROR: grpc.StatusCode.UNKNOWN,
+    ErrorCode.CANCELLED: grpc.StatusCode.CANCELLED,
+    ErrorCode.DEADLINE_EXCEEDED: grpc.StatusCode.DEADLINE_EXCEEDED,
+    ErrorCode.INTERNAL: grpc.StatusCode.INTERNAL,
+    ErrorCode.RESOURCE_EXHAUSTED: grpc.StatusCode.RESOURCE_EXHAUSTED,
+}
+
+
+class ToolError(Exception):
+    """Base error for tool operations."""
+
+    def __init__(
+        self,
+        code: ErrorCode | str,
+        step: str | None = None,
+        provider: str | None = None,
+        message: str = "",
+        retryable: bool = False,
+        cause: Exception | None = None,
+    ):
+        super().__init__(message)
+        try:
+            self.code = ErrorCode(code)
+        except ValueError:
+            self.code = ErrorCode.INTERNAL
+        self.step = step
+        self.provider = provider
+        self.retryable = retryable
+        self.cause = cause
+
+    def __str__(self) -> str:
+        parts = [f"[{self.code.value}]"]
+        if self.step:
+            parts.append(f"step={self.step}")
+        if self.provider:
+            parts.append(f"provider={self.provider}")
+        parts.append(super().__str__())
+        return " ".join(parts)
+
+    def to_grpc_status(self) -> tuple[grpc.StatusCode, str]:
+        status = _CODE_TO_GRPC.get(self.code, grpc.StatusCode.INTERNAL)
+        return status, str(self)
+
+
+class InvalidRequest(ToolError):
+    def __init__(
+        self,
+        message: str,
+        step: str | None = None,
+        provider: str | None = None,
+        cause: Exception | None = None,
+    ):
+        super().__init__(
+            ErrorCode.INVALID_REQUEST, step, provider, message, retryable=False, cause=cause
+        )
+
+
+class ProviderError(ToolError):
+    def __init__(
+        self,
+        message: str,
+        step: str | None = None,
+        provider: str | None = None,
+        cause: Exception | None = None,
+        retryable: bool = True,
+    ):
+        super().__init__(
+            ErrorCode.PROVIDER_ERROR, step, provider, message, retryable=retryable, cause=cause
+        )
+
+
+class Cancelled(ToolError):
+    def __init__(self, message: str = "Operation cancelled", step: str | None = None):
+        super().__init__(ErrorCode.CANCELLED, step, None, message, retryable=False)
+
+
+class DeadlineExceeded(ToolError):
+    def __init__(self, message: str = "Deadline exceeded", step: str | None = None):
+        super().__init__(ErrorCode.DEADLINE_EXCEEDED, step, None, message, retryable=True)
+
+
+class Internal(ToolError):
+    def __init__(
+        self,
+        message: str = "Internal error",
+        step: str | None = None,
+        cause: Exception | None = None,
+    ):
+        super().__init__(ErrorCode.INTERNAL, step, None, message, retryable=False, cause=cause)
+
+
+class ResourceExhausted(ToolError):
+    def __init__(
+        self, message: str = "Resource exhausted", step: str | None = None, retryable: bool = True
+    ):
+        super().__init__(ErrorCode.RESOURCE_EXHAUSTED, step, None, message, retryable=retryable)
+
+
+def to_grpc_status(err: ToolError) -> tuple[grpc.StatusCode, str]:
+    return err.to_grpc_status()
+
+
+def wrap_unknown(exc: BaseException, step: str | None = None) -> Internal:
+    """Wrap an unknown exception as Internal, hiding traceback from client."""
+    return Internal(f"Unexpected error: {type(exc).__name__}", step=step, cause=exc)  # type: ignore[arg-type]