mplang-nightly 0.1.dev269__py3-none-any.whl → 0.1.dev270__py3-none-any.whl

mplang/v1/runtime/data_providers.py

@@ -1,303 +0,0 @@
-# Copyright 2025 Ant Group Co., Ltd.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#   http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-from __future__ import annotations
-
-import pathlib
-import struct
-from dataclasses import dataclass
-from typing import Any
-from urllib.parse import ParseResult, urlparse
-
-import numpy as np
-
-from mplang.v1.core import TableLike, TableType, TensorType
-from mplang.v1.kernels.base import KernelContext
-from mplang.v1.kernels.value import (
-    TableValue,
-    TensorValue,
-    Value,
-    decode_value,
-    encode_value,
-)
-from mplang.v1.utils import table_utils
-
-
-@dataclass(frozen=True)
-class ResolvedURI:
-    """Result of resolving a resource path into a normalized form.
-
-    Attributes:
-      scheme: The URI scheme (e.g., 'file', 's3', 'mem', 'var', 'secret').
-      raw: The original path string as provided by the user.
-      parsed: The ParseResult if a scheme was present; otherwise None.
-      local_path: For file paths: concrete filesystem path (absolute or as given).
-    """
-
-    scheme: str
-    raw: str
-    parsed: ParseResult | None
-    local_path: str | None
-
-
-def resolve_uri(path: str) -> ResolvedURI:
-    """Resolve a user-provided resource location into a normalized URI form.
-
-    This helper accepts plain filesystem paths and RFC 3986 style URIs. A path
-    is treated as ``file`` when ``urlparse(path).scheme`` is empty. Detection
-    no longer depends on the presence of the literal substring ``"://"`` so
-    that forms like ``mem:foo`` (no slashes) are still recognized as a URI.
-
-    Captured fields
-    - ``scheme``: Lower-cased scheme (``file`` when absent)
-    - ``raw``: Original input
-    - ``parsed``: ``ParseResult`` when a scheme was provided, else ``None``
-    - ``local_path``: Filesystem path for ``file`` scheme, else ``None``
-
-    Supported (pluggable) schemes out-of-the-box:
-      * ``file`` (default)
-      * ``mem``
-      * ``s3`` (stub)
-      * ``secret`` (stub)
-      * ``symbols`` (registered server-side)
-
-    Examples
-    >>> resolve_uri("data/train.npy").scheme
-    'file'
-    >>> resolve_uri("mem:dataset1").scheme
-    'mem'
-    >>> resolve_uri("mem://dataset1").scheme  # both forms acceptable
-    'mem'
-    >>> resolve_uri("symbols://shared_model").scheme
-    'symbols'
-    >>> resolve_uri("file:///tmp/x.npy").local_path
-    '/tmp/x.npy'
-    """
-
-    pr = urlparse(path)
-    if not pr.scheme:
-        return ResolvedURI("file", path, None, path)
-
-    scheme = pr.scheme.lower()
-    local_path: str | None = None
-    if scheme == "file":
-        local_path = pr.path
-        if pr.netloc and not local_path.startswith("/"):
-            local_path = f"//{pr.netloc}/{pr.path}"
-    return ResolvedURI(scheme, path, pr, local_path)
-
-
-class DataProvider:
-    """Abstract base for data providers.
-
-    Minimal contract: read/write by URI and type spec. Providers may ignore the
-    type spec but SHOULD validate when feasible.
-    """
-
-    def read(
-        self, uri: ResolvedURI, out_spec: TensorType | TableType, *, ctx: KernelContext
-    ) -> Any:
-        raise NotImplementedError
-
-    def write(self, uri: ResolvedURI, value: Any, *, ctx: KernelContext) -> None:
-        raise NotImplementedError
-
-
-_REGISTRY: dict[str, DataProvider] = {}
-
-
-def register_provider(
-    scheme: str, provider: DataProvider, *, replace: bool = False, quiet: bool = False
-) -> None:
-    """Register a provider implementation.
-
-    Args:
-        scheme: URI scheme handled (case-insensitive)
-        provider: Implementation
-        replace: If False and scheme exists -> ValueError
-        quiet: If True, suppress duplicate log messages when replacing
-    """
-    import logging
-
-    key = scheme.lower()
-    if not replace and key in _REGISTRY:
-        raise ValueError(f"provider already registered for scheme: {scheme}")
-    if replace and key in _REGISTRY and not quiet:
-        logging.info(f"Replacing existing provider for scheme '{scheme}'")
-    _REGISTRY[key] = provider
-
-
-def get_provider(scheme: str) -> DataProvider | None:
-    return _REGISTRY.get(scheme.lower())
-
-
-# ---------------- Default Providers ----------------
-MAGIC_MPLANG = b"MPLG"
-MAGIC_PARQUET = b"PAR1"
-MAGIC_ORC = b"ORC"
-MAGIC_NUMPY = b"\x93NUMPY"
-VERSION = 0x01
-
-
-class FileProvider(DataProvider):
-    """Local filesystem provider.
-
-    For tables: CSV bytes via table_utils.
-    For tensors: NumPy .npy via np.load/np.save.
-    """
-
-    def read(
-        self, uri: ResolvedURI, out_spec: TensorType | TableType, *, ctx: KernelContext
-    ) -> Any:
-        path = pathlib.Path(uri.local_path or uri.raw)
-        # try load by magic
-        with path.open("rb") as f:
-            # this is the maximum length needed to detect all supported formats
-            # (numpy requires 6 bytes: '\x93NUMPY').
-            MAGIC_BYTES_LEN_MAX = 6
-            magic = f.read(MAGIC_BYTES_LEN_MAX)
-            f.seek(0)
-            if magic.startswith(MAGIC_MPLANG):
-                MPLANG_HEADER_LEN = len(MAGIC_MPLANG) + 1
-                header = f.read(MPLANG_HEADER_LEN)
-                _, version = struct.unpack(">4sB", header)
-                if version != VERSION:
-                    raise ValueError(f"unsupported mplang version {version}")
-                payload = f.read()
-                return decode_value(payload)
-            elif magic.startswith(MAGIC_PARQUET):
-                if not isinstance(out_spec, TableType):
-                    raise ValueError(
-                        f"PARQUET files require TableType output spec, got {type(out_spec).__name__}"
-                    )
-                return table_utils.read_table(
-                    f, format="parquet", columns=list(out_spec.column_names())
-                )
-            elif magic.startswith(MAGIC_ORC):
-                if not isinstance(out_spec, TableType):
-                    raise ValueError(
-                        f"ORC files require TableType output spec, got {type(out_spec).__name__}"
-                    )
-                return table_utils.read_table(
-                    f, format="orc", columns=list(out_spec.column_names())
-                )
-            elif magic.startswith(MAGIC_NUMPY):
-                if not isinstance(out_spec, TensorType):
-                    raise ValueError(
-                        f"NumPy files require TensorType output spec, got {type(out_spec).__name__}"
-                    )
-                return np.load(f)
-
-            # Fallback: open the file for CSV or NumPy loading.
-            if isinstance(out_spec, TableType):
-                return table_utils.read_table(
-                    f, format="csv", columns=list(out_spec.column_names())
-                )
-            else:
-                return np.load(f)
-
-    def write(self, uri: ResolvedURI, value: Any, *, ctx: KernelContext) -> None:
-        import os
-
-        path = uri.local_path or uri.raw
-        dir_name = os.path.dirname(path)
-        if dir_name:
-            os.makedirs(dir_name, exist_ok=True)
-
-        if not isinstance(value, Value):
-            value = (
-                TableValue(value)
-                if isinstance(value, TableLike)
-                else TensorValue(value)
-            )
-
-        if isinstance(value, TableValue):
-            table_utils.write_table(value.to_arrow(), path, format="parquet")
-        elif isinstance(value, TensorValue):
-            with open(path, "wb") as f:
-                np.save(f, value.to_numpy())
-        else:
-            payload = encode_value(value)
-            with open(path, "wb") as f:
-                f.write(struct.pack(">4sB", MAGIC_MPLANG, VERSION))
-                f.write(payload)
-
-
-class MemProvider(DataProvider):
-    """In-memory per-runtime KV provider (per rank, per session/runtime)."""
-
-    STATE_KEY = "resource.providers.mem"
-
-    @staticmethod
-    def _store(ctx: KernelContext) -> dict[str, Any]:
-        # Use ensure_state so creation is atomic & centralized; enforce dict.
-        store = ctx.runtime.ensure_state(MemProvider.STATE_KEY, dict)
-        if not isinstance(store, dict):  # pragma: no cover - defensive
-            raise TypeError(
-                f"runtime state key '{MemProvider.STATE_KEY}' expected dict, got {type(store).__name__}"
-            )
-        return store  # type: ignore[return-value]
-
-    def read(
-        self, uri: ResolvedURI, out_spec: TensorType | TableType, *, ctx: KernelContext
-    ) -> Any:
-        store = self._store(ctx)
-        key = uri.raw
-        if key not in store:
-            raise FileNotFoundError(f"mem resource not found: {key}")
-        return store[key]
-
-    def write(self, uri: ResolvedURI, value: Any, *, ctx: KernelContext) -> None:
-        store = self._store(ctx)
-        store[uri.raw] = value
-
-
-class S3Provider(DataProvider):
-    """Placeholder S3 provider. Install external plugin to enable."""
-
-    def read(
-        self, uri: ResolvedURI, out_spec: TensorType | TableType, *, ctx: KernelContext
-    ) -> Any:
-        raise NotImplementedError(
-            "S3 provider not installed. Provide an external plugin via register_provider('s3', ...) ."
-        )
-
-    def write(self, uri: ResolvedURI, value: Any, *, ctx: KernelContext) -> None:
-        raise NotImplementedError(
-            "S3 provider not installed. Provide an external plugin via register_provider('s3', ...) ."
-        )
-
-
-class SecretProvider(DataProvider):
-    """Placeholder secret provider. Integrate with KMS/secret manager via plugin."""
-
-    def read(
-        self, uri: ResolvedURI, out_spec: TensorType | TableType, *, ctx: KernelContext
-    ) -> Any:
-        raise NotImplementedError(
-            "secret provider not installed. Provide an external plugin via register_provider('secret', ...) ."
-        )
-
-    def write(self, uri: ResolvedURI, value: Any, *, ctx: KernelContext) -> None:
-        raise NotImplementedError(
-            "secret provider not installed. Provide an external plugin via register_provider('secret', ...) ."
-        )
-
-
-# Register default providers
-register_provider("file", FileProvider())
-register_provider("mem", MemProvider())
-# Stubs to signal missing providers explicitly (can be overridden by plugins)
-register_provider("s3", S3Provider())
-register_provider("secret", SecretProvider())

mplang/v1/runtime/driver.py

@@ -1,324 +0,0 @@
-# Copyright 2025 Ant Group Co., Ltd.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#   http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-"""
-HTTP-based driver implementation for distributed execution.
-
-This module provides an HTTP-based driver, using REST APIs
-for distributed multi-party computation coordination.
-"""
-
-from __future__ import annotations
-
-import asyncio
-import base64
-import uuid
-from collections.abc import Sequence
-from typing import Any
-
-import numpy as np
-
-from mplang.v1.core import (
-    ClusterSpec,
-    InterpContext,
-    InterpVar,
-    IrWriter,
-    Mask,
-    MPObject,
-    MPType,
-)
-from mplang.v1.core.expr.ast import Expr
-from mplang.v1.kernels.value import TableValue, TensorValue
-from mplang.v1.runtime.client import HttpExecutorClient
-
-
-def new_uuid() -> str:
-    """Generates a short UUID using URL-safe Base64 encoding."""
-    u = uuid.uuid4()
-    # Get the 16 bytes of the UUID
-    uuid_bytes = u.bytes
-    # Encode using URL-safe Base64
-    encoded_bytes = base64.urlsafe_b64encode(uuid_bytes)
-    # Decode to UTF-8 string, remove padding, and take first 8 characters
-    encoded_string = encoded_bytes.decode("utf-8").rstrip("=")[:8]
-    return encoded_string
-
-
-class DriverVar(InterpVar):
-    """A variable that references a value in distributed HTTP executor nodes.
-
-    This represents a symbol stored on remote HTTP servers that can be
-    retrieved via REST API calls.
-    """
-
-    def __init__(
-        self,
-        ctx: Driver,
-        symbol_name: str,
-        mptype: MPType,
-    ) -> None:
-        super().__init__(ctx, mptype)
-        self.symbol_name = symbol_name
-
-    @property
-    def mptype(self) -> MPType:
-        """The type of this variable."""
-        return self._mptype
-
-    def __repr__(self) -> str:
-        return f"HttpDriverVar(symbol_name={self.symbol_name}, mptype={self.mptype})"
-
-
-class Driver(InterpContext):
-    """Driver for distributed execution using HTTP-based services.
-
-    Args:
-        cluster_spec: The cluster specification defining the distributed environment.
-        trace_ranks: List of ranks to trace execution for debugging.
-        timeout: HTTP request timeout in seconds.
-    """
-
-    def __init__(
-        self,
-        cluster_spec: ClusterSpec,
-        *,
-        trace_ranks: list[int] | None = None,
-        timeout: int = 120,
-    ) -> None:
-        """Initialize a driver with the given cluster specification.
-
-        Args:
-            cluster_spec: The cluster specification defining the distributed environment.
-            trace_ranks: List of ranks to trace execution for debugging.
-            timeout: HTTP request timeout in seconds.
-        """
-        super().__init__(cluster_spec)
-        self._trace_ranks = trace_ranks or []
-        self.timeout = timeout
-
-        self._session_id: str | None = None
-        self._counter = 0
-
-        self.node_addrs = {
-            node_id: node.endpoint for node_id, node in cluster_spec.nodes.items()
-        }
-
-        # Get SPU configuration from cluster_spec
-        spu_devices = cluster_spec.get_devices_by_kind("SPU")
-        if not spu_devices:
-            raise ValueError("No SPU device found in the cluster specification")
-        if len(spu_devices) > 1:
-            raise ValueError("Multiple SPU devices found in the cluster specification")
-        spu_device = spu_devices[0]
-
-        # Store SPU configuration as strings for better readability
-        self.spu_protocol_str = spu_device.config["protocol"]
-        self.spu_field_str = spu_device.config["field"]
-
-        # Compute spu_mask from spu_device members
-        spu_mask = Mask.from_ranks([member.rank for member in spu_device.members])
-        self.spu_mask_int = int(spu_mask)
-
-    def _create_clients(self) -> dict[str, HttpExecutorClient]:
-        """Create HTTP clients for all endpoints."""
-        clients = {}
-        for node_id, endpoint in self.node_addrs.items():
-            clients[node_id] = HttpExecutorClient(endpoint, self.timeout)
-        return clients
-
-    async def _close_clients(self, clients: dict[str, HttpExecutorClient]) -> None:
-        """Close all provided HTTP clients."""
-        await asyncio.gather(*[client.close() for client in clients.values()])
-
-    def new_name(self, prefix: str = "var") -> str:
-        """Generate a unique execution name."""
-        name = f"{prefix}_{self._counter}"
-        self._counter += 1
-        return name
-
-    async def _get_or_create_session(self) -> str:
-        """Get existing session or create a new one across all HTTP servers."""
-        if self._session_id is None:
-            new_session_id = new_uuid()
-            # Create temporary clients for session creation
-            clients = self._create_clients()
-            try:
-                # Create session on all HTTP servers concurrently
-                tasks = []
-                for node_id, client in clients.items():
-                    # Convert node_id to rank for the session creation
-                    rank = list(self.node_addrs.keys()).index(node_id)
-                    task = client.create_session(
-                        name=new_session_id,
-                        rank=rank,
-                        cluster_spec=self.cluster_spec.to_dict(),
-                    )
-                    tasks.append(task)
-
-                try:
-                    results = await asyncio.gather(*tasks)
-                    for session_id in results:
-                        assert session_id == new_session_id
-                    self._session_id = new_session_id
-                except RuntimeError as e:
-                    raise RuntimeError(
-                        f"Failed to create session on one or more parties: {e}"
-                    ) from e
-            finally:
-                await self._close_clients(clients)
-
-        assert self._session_id is not None
-        return self._session_id
-
-    async def _evaluate(
-        self, expr: Expr, bindings: dict[str, MPObject]
-    ) -> Sequence[MPObject]:
-        """Async implementation to evaluate an expression."""
-        session_id = await self._get_or_create_session()
-
-        # Prepare input names from bindings
-        var_names = []
-        party_symbol_names = []
-        for name, var in bindings.items():
-            if var.ctx is not self:
-                raise ValueError(f"Variable {name} not in this context, got {var.ctx}.")
-            assert isinstance(var, DriverVar), (
-                f"Expected HttpDriverVar, got {type(var)}"
-            )
-            var_names.append(name)
-            party_symbol_names.append(var.symbol_name)
-
-        var_name_mapping = dict(zip(var_names, party_symbol_names, strict=True))
-
-        writer = IrWriter(var_name_mapping)
-        program_proto = writer.dumps(expr)
-
-        output_symbols = [self.new_name() for _ in range(expr.num_outputs)]
-
-        # Create temporary clients for computation execution
-        clients = self._create_clients()
-        try:
-            # Concurrently create and execute computation on all parties
-            tasks = []
-            computation_id = new_uuid()
-            for _rank, client in clients.items():
-                task = client.create_and_execute_computation(
-                    session_id,
-                    computation_id,
-                    program_proto.SerializeToString(),
-                    party_symbol_names,
-                    output_symbols,
-                )
-                tasks.append(task)
-
-            try:
-                await asyncio.gather(*tasks)
-            except RuntimeError as e:
-                raise RuntimeError(
-                    f"Failed to create and execute computation on one or more parties: {e}"
-                ) from e
-        finally:
-            await self._close_clients(clients)
-
-        # Create HttpDriverVar objects for each output
-        driver_vars = []
-        for symbol_name, mptype in zip(output_symbols, expr.mptypes, strict=True):
-            driver_var = DriverVar(self, symbol_name, mptype)
-            driver_vars.append(driver_var)
-
-        return driver_vars
-
-    def evaluate(self, expr: Expr, bindings: dict[str, MPObject]) -> Sequence[MPObject]:
-        """Evaluate an expression using distributed HTTP execution."""
-        return asyncio.run(self._evaluate(expr, bindings))
-
-    async def _fetch(self, obj: MPObject) -> list[Any]:
-        """Async implementation to fetch results."""
-        if not isinstance(obj, DriverVar):
-            raise ValueError(f"Expected HttpDriverVar, got {type(obj)}")
-
-        session_id = await self._get_or_create_session()
-        symbol_full_name = obj.symbol_name
-
-        # Create temporary clients for fetching
-        clients = self._create_clients()
-        try:
-            # Concurrently fetch symbol from all parties
-            tasks = []
-            for _rank, client in clients.items():
-                task = client.get_symbol(session_id, symbol_full_name)
-                tasks.append(task)
-
-            try:
-                # The results will be in the same order as the clients (ranks)
-                results = await asyncio.gather(*tasks)
-                converted: list[Any] = []
-                for value in results:
-                    if isinstance(value, TensorValue):
-                        arr = value.to_numpy()
-                        if isinstance(arr, np.ndarray) and arr.size == 1:
-                            converted.append(arr.item())
-                        else:
-                            converted.append(arr)
-                    elif isinstance(value, TableValue):
-                        converted.append(value.to_pandas())
-                    else:
-                        converted.append(value)
-                return converted
-            except RuntimeError as e:
-                raise RuntimeError(
-                    f"Failed to fetch symbol from one or more parties: {e}"
-                ) from e
-        finally:
-            await self._close_clients(clients)
-
-    def fetch(self, obj: MPObject) -> list[Any]:
-        """Fetch results from the distributed HTTP execution."""
-        return asyncio.run(self._fetch(obj))
-
-    async def _ping(self, node_id: str) -> bool:
-        """Async implementation to ping a node.
-
-        Args:
-            node_id: The ID of the node to ping
-
-        Returns:
-            True if the node is healthy, False otherwise
-        """
-        # Create a temporary client for the node
-        if node_id not in self.node_addrs:
-            raise ValueError(f"Node {node_id} not found in party addresses")
-
-        endpoint = self.node_addrs[node_id]
-        client = HttpExecutorClient(endpoint, self.timeout)
-
-        try:
-            # Perform health check
-            return await client.health_check()
-        except Exception:
-            # Any exception means the node is not healthy
-            return False
-        finally:
-            await client.close()
-
-    def ping(self, node_id: str) -> bool:
-        """Ping a node to check if it's healthy.
-
-        Args:
-            node_id: The ID of the node to ping
-
-        Returns:
-            True if the node is healthy, False otherwise
-        """
-        return asyncio.run(self._ping(node_id))

mplang/v1/runtime/exceptions.py

@@ -1,27 +0,0 @@
-# Copyright 2025 Ant Group Co., Ltd.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#   http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-"""Custom exception types for the HTTP backend."""
-
-
-class HttpBackendError(Exception):
-    """Base exception for all HTTP backend errors."""
-
-
-class ResourceNotFound(HttpBackendError):
-    """Raised when a resource (session, computation, etc.) is not found."""
-
-
-class InvalidRequestError(HttpBackendError, ValueError):
-    """Raised for invalid requests, e.g., bad parameters or invalid state."""