mplang-nightly 0.1.dev158__py3-none-any.whl → 0.1.dev268__py3-none-any.whl

mplang/{runtime → v1/runtime}/server.py

@@ -22,22 +22,30 @@ import logging
 import re
 from typing import Any
 
-import cloudpickle as pickle
-from fastapi import FastAPI, HTTPException, Request
+from fastapi import (
+    FastAPI,
+    HTTPException,
+    Request,
+)
 from fastapi.responses import JSONResponse
 from pydantic import BaseModel
 
-from mplang.core.mpir import Reader
-from mplang.core.table import TableType
-from mplang.core.tensor import TensorType
-from mplang.kernels.base import KernelContext
-from mplang.protos.v1alpha1 import mpir_pb2
-from mplang.runtime.data_providers import DataProvider, ResolvedURI, register_provider
-from mplang.runtime.exceptions import InvalidRequestError, ResourceNotFound
-from mplang.runtime.session import (
+from mplang.v1.core import IrReader, TableType, TensorType
+from mplang.v1.core.cluster import ClusterSpec
+from mplang.v1.kernels.base import KernelContext
+from mplang.v1.kernels.value import Value, decode_value, encode_value
+from mplang.v1.protos.v1alpha1 import mpir_pb2
+from mplang.v1.runtime.data_providers import (
+    DataProvider,
+    ResolvedURI,
+    register_provider,
+)
+from mplang.v1.runtime.exceptions import InvalidRequestError, ResourceNotFound
+from mplang.v1.runtime.session import (
     Computation,
     Session,
     Symbol,
+    create_session_from_spec,
 )
 
 logger = logging.getLogger(__name__)
@@ -112,19 +120,11 @@ class _SymbolsProvider(DataProvider):
         ctx: KernelContext,
     ) -> None:  # type: ignore[override]
         name = self._symbol_name(uri)
-        try:
-            data_b64 = base64.b64encode(pickle.dumps(value)).decode("utf-8")
-        except Exception as e:  # pragma: no cover - defensive
+        if not isinstance(value, Value):
             raise InvalidRequestError(
-                f"Failed to encode value for symbols:// write: {e!s}"
-            ) from e
-        try:
-            obj = pickle.loads(base64.b64decode(data_b64))
-        except Exception as e:  # pragma: no cover - defensive
-            raise InvalidRequestError(
-                f"Failed to decode value for symbols:// write: {e!s}"
-            ) from e
-        _global_symbols[name] = Symbol(name=name, mptype={}, data=obj)
+                f"symbols:// write expects Value instance, got {type(value)}"
+            )
+        _global_symbols[name] = Symbol(name=name, mptype={}, data=value)
 
 
 # Register symbols provider explicitly for server runtime
@@ -208,17 +208,18 @@ class ComputationResponse(BaseModel):
 
 class CreateSymbolRequest(BaseModel):
     mptype: dict
-    data: str  # Base64 encoded data
+    data: str  # Base64 encoded Value data
 
 
 class SymbolResponse(BaseModel):
     name: str
     mptype: dict
-    data: str
+    data: str  # Base64 encoded Value data
 
 
 class CommSendRequest(BaseModel):
-    data: str  # Base64 encoded data
+    data: str  # Base64 encoded binary data
+    is_raw_bytes: bool = False  # True for SPU channel raw bytes
 
 
 # Response Models for enhanced status
@@ -233,7 +234,7 @@ class ComputationListResponse(BaseModel):
 class GlobalSymbolResponse(BaseModel):
     name: str
     mptype: dict
-    data: str
+    data: str  # Base64 encoded Value data
 
 
 @app.get("/health")
@@ -266,15 +267,12 @@ def list_session_computations(session_name: str) -> ComputationListResponse:
 def create_session(session_name: str, request: CreateSessionRequest) -> SessionResponse:
     validate_name(session_name, "session")
     # Delegate cluster spec parsing & session construction to resource layer
-    from mplang.core.cluster import ClusterSpec  # local import to avoid cycles
 
     if session_name in _sessions:
         sess = _sessions[session_name]
     else:
         spec = ClusterSpec.from_dict(request.cluster_spec)
-        if len(spec.get_devices_by_kind("SPU")) == 0:
-            raise InvalidRequestError("No SPU device found in cluster_spec for session")
-        sess = Session(name=session_name, rank=request.rank, cluster_spec=spec)
+        sess = create_session_from_spec(name=session_name, rank=request.rank, spec=spec)
         _sessions[session_name] = sess
     return SessionResponse(name=sess.name)
 
@@ -314,7 +312,7 @@ def create_and_execute_computation(
             f"Invalid base64 or protobuf for mpprogram: {e!s}"
         ) from e
 
-    reader = Reader()
+    reader = IrReader()
     expr = reader.loads(graph_proto)
 
     if expr is None:
@@ -359,7 +357,7 @@ def create_session_symbol(
     if not sess:
         raise ResourceNotFound(f"Session '{session_name}' not found.")
     try:
-        obj = pickle.loads(base64.b64decode(request.data))
+        obj = decode_value(base64.b64decode(request.data))
     except Exception as e:
         raise InvalidRequestError(f"Invalid symbol data: {e!s}") from e
     symbol = Symbol(name=symbol_name, mptype=request.mptype, data=obj)
@@ -368,7 +366,7 @@ def create_session_symbol(
     return SymbolResponse(
         name=symbol.name,
         mptype=symbol.mptype,
-        data=request.data,
+        data=base64.b64encode(encode_value(symbol.data)).decode("utf-8"),
     )
 
 
@@ -388,13 +386,16 @@ def get_session_symbol(session_name: str, symbol_name: str) -> SymbolResponse:
                 status_code=404, detail=f"Symbol {symbol_name} not found"
             )
 
-        data_bytes = pickle.dumps(symbol.data)
-        data_b64 = base64.b64encode(data_bytes).decode("utf-8")
+        # symbol data is None means this party does not participate the computation
+        # that produced the symbol.
+        if symbol.data is None:
+            raise ResourceNotFound(f"Symbol '{symbol_name}' has no data on this party")
 
+        # Serialize using Value envelope
         return SymbolResponse(
             name=symbol.name,
             mptype=symbol.mptype,
-            data=data_b64,
+            data=base64.b64encode(encode_value(symbol.data)).decode("utf-8"),
         )
     except ValueError as e:
         raise HTTPException(status_code=404, detail=str(e)) from e
@@ -430,12 +431,16 @@ def create_global_symbol(
 ) -> GlobalSymbolResponse:
     validate_name(symbol_name, "symbol")
     try:
-        obj = pickle.loads(base64.b64decode(request.data))
+        obj = decode_value(base64.b64decode(request.data))
     except Exception as e:
         raise InvalidRequestError(f"Invalid global symbol data: {e!s}") from e
     sym = Symbol(name=symbol_name, mptype=request.mptype, data=obj)
     _global_symbols[symbol_name] = sym
-    return GlobalSymbolResponse(name=sym.name, mptype=sym.mptype, data=request.data)
+    return GlobalSymbolResponse(
+        name=sym.name,
+        mptype=sym.mptype,
+        data=base64.b64encode(encode_value(sym.data)).decode("utf-8"),
+    )
 
 
 @app.get("/api/v1/symbols/{symbol_name}", response_model=GlobalSymbolResponse)
@@ -443,9 +448,12 @@ def get_global_symbol(symbol_name: str) -> GlobalSymbolResponse:  # route handle
     sym = _global_symbols.get(symbol_name)
     if not sym:
         raise ResourceNotFound(f"Global symbol '{symbol_name}' not found")
-    data_bytes = pickle.dumps(sym.data)
-    data_b64 = base64.b64encode(data_bytes).decode("utf-8")
-    return GlobalSymbolResponse(name=sym.name, mptype=sym.mptype, data=data_b64)
+    # Serialize using Value envelope
+    return GlobalSymbolResponse(
+        name=sym.name,
+        mptype=sym.mptype,
+        data=base64.b64encode(encode_value(sym.data)).decode("utf-8"),
+    )
 
 
 @app.get("/api/v1/symbols")
@@ -480,6 +488,14 @@ def comm_send(
     # The receiver rank should be the rank of the server hosting this endpoint
     # We don't need to validate to_rank since the request is coming to this server
 
+    # For raw bytes (SPU channel), pass through as dict with flag
+    # For normal data, pass the base64 string directly
+    data_payload: str | dict[str, object]
+    if request.is_raw_bytes:
+        data_payload = {"data": request.data, "is_raw_bytes": True}
+    else:
+        data_payload = request.data
+
     # Use the proper onSent mechanism from CommunicatorBase
-    sess.communicator.onSent(from_rank, key, request.data)
+    sess.communicator.onSent(from_rank, key, data_payload)
     return {"status": "ok"}

mplang/{runtime → v1/runtime}/session.py

@@ -25,48 +25,28 @@ Process-wide registries (sessions, global symbols) live in the server layer
 
 from __future__ import annotations
 
-import logging
 import time
 from dataclasses import dataclass, field
 from functools import cached_property
 from typing import TYPE_CHECKING, Any, cast
-from urllib.parse import urlparse
 
 import spu.libspu as libspu
 
-from mplang.core.expr.ast import Expr
-from mplang.core.expr.evaluator import IEvaluator, create_evaluator
-from mplang.core.mask import Mask
-from mplang.kernels.context import RuntimeContext
-from mplang.kernels.spu import PFunction  # type: ignore
-from mplang.runtime.communicator import HttpCommunicator
-from mplang.runtime.exceptions import ResourceNotFound
-from mplang.runtime.link_comm import LinkCommunicator
-from mplang.utils.spu_utils import parse_field, parse_protocol
+from mplang.v1.core.cluster import ClusterSpec
+from mplang.v1.core.comm import ICommunicator
+from mplang.v1.core.expr.ast import Expr
+from mplang.v1.core.expr.evaluator import IEvaluator, create_evaluator
+from mplang.v1.core.mask import Mask
+from mplang.v1.kernels.context import RuntimeContext
+from mplang.v1.kernels.spu import PFunction  # type: ignore
+from mplang.v1.kernels.value import Value
+from mplang.v1.runtime.communicator import HttpCommunicator
+from mplang.v1.runtime.exceptions import ResourceNotFound
+from mplang.v1.runtime.link_comm import LinkCommunicator
+from mplang.v1.utils.spu_utils import parse_field, parse_protocol
 
 if TYPE_CHECKING:  # pragma: no cover - import only for type checking
-    from mplang.core.cluster import ClusterSpec, Node, RuntimeInfo
-
-
-class LinkCommFactory:
-    """Factory for creating and caching link communicators."""
-
-    def __init__(self) -> None:
-        self._cache: dict[tuple[int, tuple[str, ...]], LinkCommunicator] = {}
-
-    def create_link(self, rel_rank: int, addrs: list[str]) -> LinkCommunicator:
-        key = (rel_rank, tuple(addrs))
-        link = self._cache.get(key)
-        if link is not None:
-            return link
-        logging.info(f"LinkCommunicator created: rel_rank={rel_rank} addrs={addrs}")
-        link = LinkCommunicator(rel_rank, addrs)
-        self._cache[key] = link
-        return link
-
-
-# Shared link factory (module-local, not global registry of sessions)
-g_link_factory = LinkCommFactory()
+    from mplang.v1.core.cluster import ClusterSpec, Node, RuntimeInfo
 
 
 @dataclass
@@ -95,19 +75,25 @@ class SessionState:
 class Session:
     """Represents the per-rank execution context.
 
-    Immutable config: name, rank, cluster_spec.
+    Immutable config: name, rank, cluster_spec, communicator.
     Derived: node, runtime_info, endpoints, spu_device, spu_mask, protocol/field, is_spu_party.
     Mutable: state (runtime object, symbols, computations, seeded flag).
+
+    Note: communicator is assumed to be initialized with cluster spec info (e.g. endpoints).
     """
 
-    def __init__(self, name: str, rank: int, cluster_spec: ClusterSpec):
+    def __init__(
+        self,
+        name: str,
+        rank: int,
+        cluster_spec: ClusterSpec,
+        communicator: ICommunicator,
+    ):
         self.name = name
         self.rank = rank
         self.cluster_spec = cluster_spec
         self.state = SessionState()
-        self.communicator = HttpCommunicator(
-            session_name=name, rank=rank, endpoints=self.endpoints
-        )
+        self.communicator = communicator
 
     # --- Derived topology ---
     @cached_property
@@ -118,18 +104,9 @@ class Session:
     def runtime_info(self) -> RuntimeInfo:
         return self.node.runtime_info
 
-    @cached_property
+    @property
     def endpoints(self) -> list[str]:
-        eps: list[str] = []
-        for n in sorted(
-            self.cluster_spec.nodes.values(),
-            key=lambda x: x.rank,  # type: ignore[attr-defined]
-        ):
-            ep = n.endpoint
-            if not ep.startswith(("http://", "https://")):
-                ep = f"http://{ep}"
-            eps.append(ep)
-        return eps
+        return self.cluster_spec.endpoints
 
     @cached_property
     def spu_device(self):  # type: ignore
@@ -184,22 +161,19 @@ class Session:
             return
 
         link_ctx = None
-        # TODO(jint): reuse same port for mplang and spu.
-        SPU_PORT_OFFSET = 100
 
         if self.is_spu_party:
-            # Build SPU address list across all endpoints for ranks in mask
-            spu_addrs: list[str] = []
-            for r, addr in enumerate(self.communicator.endpoints):
-                if r in self.spu_mask:
-                    if "//" not in addr:
-                        addr = f"//{addr}"
-                    parsed = urlparse(addr)
-                    assert isinstance(parsed.port, int)
-                    new_addr = f"{parsed.hostname}:{parsed.port + SPU_PORT_OFFSET}"
-                    spu_addrs.append(new_addr)
-            rel_index = sum(1 for r in range(self.rank) if r in self.spu_mask)
-            link_ctx = g_link_factory.create_link(rel_index, spu_addrs)
+            # Use Channels mode to reuse existing HttpCommunicator
+            # This eliminates the need for separate BRPC ports (SPU_PORT_OFFSET)
+            from mplang.v1.core.comm import CommunicatorBase
+
+            # Type assertion: ICommunicator is actually CommunicatorBase
+            comm = cast(CommunicatorBase, self.communicator)
+            link_ctx = LinkCommunicator(
+                rank=self.rank,
+                comm=comm,
+                spu_mask=self.spu_mask,
+            )
 
         spu_config = libspu.RuntimeConfig(
             protocol=parse_protocol(self.spu_protocol),
@@ -271,14 +245,26 @@ class Session:
                 f"Expected {len(output_names)} results, got {len(results)}"
             )
         for name, val in zip(output_names, results, strict=True):
+            # In pure SIMP model, all nodes should have the same symbol table.
+            # Non-participating nodes get None values.
+            if val is not None and not isinstance(val, Value):
+                raise TypeError(
+                    "Session executions must produce kernel Value outputs; "
+                    f"got {type(val).__name__} for symbol '{name}'"
+                )
             self.add_symbol(Symbol(name=name, mptype={}, data=val))
 
-    # --- Convenience constructor ---
-    @classmethod
-    def from_cluster_spec_dict(cls, name: str, rank: int, spec_dict: dict) -> Session:
-        from mplang.core.cluster import ClusterSpec  # local import to avoid cycles
 
-        spec = ClusterSpec.from_dict(spec_dict)
-        if len(spec.get_devices_by_kind("SPU")) == 0:
-            raise RuntimeError("No SPU device found in cluster_spec")
-        return cls(name=name, rank=rank, cluster_spec=spec)
+# --- Convenience constructor use HttpCommunicator---
+def create_session_from_spec(name: str, rank: int, spec: ClusterSpec) -> Session:
+    if len(spec.get_devices_by_kind("SPU")) == 0:
+        raise RuntimeError("No SPU device found in cluster_spec")
+
+    # Create HttpCommunicator for the session
+    communicator = HttpCommunicator(
+        session_name=name,
+        rank=rank,
+        endpoints=spec.endpoints,
+    )
+
+    return Session(name=name, rank=rank, cluster_spec=spec, communicator=communicator)

mplang/{runtime → v1/runtime}/simulation.py

@@ -18,25 +18,32 @@ import concurrent.futures
 import faulthandler
 import logging
 import sys
+import threading
 import traceback
 from collections.abc import Sequence
 from typing import Any, cast
 
 import spu.libspu as libspu
 
-from mplang.core.cluster import ClusterSpec
-from mplang.core.comm import CollectiveMixin, CommunicatorBase
-from mplang.core.expr.ast import Expr
-from mplang.core.expr.evaluator import IEvaluator, create_evaluator
-from mplang.core.interp import InterpContext, InterpVar
-from mplang.core.mask import Mask
-from mplang.core.mpir import Reader, Writer
-from mplang.core.mpobject import MPObject
-from mplang.core.mptype import MPType, TensorLike
-from mplang.core.pfunc import PFunction  # for spu.seed_env kernel seeding
-from mplang.kernels.context import RuntimeContext
-from mplang.runtime.link_comm import LinkCommunicator
-from mplang.utils.spu_utils import parse_field, parse_protocol
+from mplang.v1.core import (
+    ClusterSpec,
+    CollectiveMixin,
+    CommunicatorBase,
+    InterpContext,
+    InterpVar,
+    IrReader,
+    IrWriter,
+    Mask,
+    MPObject,
+    MPType,
+    PFunction,  # for spu.seed_env kernel seeding
+    TensorLike,
+)
+from mplang.v1.core.expr.ast import Expr
+from mplang.v1.core.expr.evaluator import IEvaluator, create_evaluator
+from mplang.v1.kernels.context import RuntimeContext
+from mplang.v1.runtime.link_comm import LinkCommunicator
+from mplang.v1.utils.spu_utils import parse_field, parse_protocol
 
 
 class ThreadCommunicator(CommunicatorBase, CollectiveMixin):
@@ -73,8 +80,8 @@ class SimVar(InterpVar):
 
     @property
     def values(self) -> list[Any]:
-        """The values of this variable across all ranks."""
-        return self._values
+        """Converted values across all ranks for user inspection."""
+        return [v.to_numpy() if hasattr(v, "to_numpy") else v for v in self._values]
 
     def __repr__(self) -> str:
         return f"SimVar({self.mptype})"
@@ -123,16 +130,37 @@ class Simulator(InterpContext):
             comm.set_peers(self._comms)
 
         # Prepare link contexts for SPU parties (store for evaluator-time initialization)
-        spu_addrs = [f"P{spu_rank}" for spu_rank in spu_mask]
+        # Use Channels mode to reuse ThreadCommunicator instead of separate mem_link
         self._spu_link_ctxs: list[LinkCommunicator | None] = [None] * world_size
-        link_ctx_list = [
-            LinkCommunicator(idx, spu_addrs, mem_link=True)
-            for idx in range(spu_mask.num_parties())
+
+        # Create LinkCommunicators in parallel to avoid deadlock
+        # (create_with_channels does handshake via TestSend/TestRecv)
+        exceptions: dict[int, Exception] = {}
+
+        def create_link(g_rank: int) -> None:
+            try:
+                self._spu_link_ctxs[g_rank] = LinkCommunicator(
+                    rank=g_rank,
+                    comm=self._comms[g_rank],
+                    spu_mask=spu_mask,
+                )
+            except Exception as e:
+                exceptions[g_rank] = e
+
+        threads = [
+            threading.Thread(target=create_link, args=(g_rank,)) for g_rank in spu_mask
         ]
-        for g_rank in range(world_size):
-            if g_rank in spu_mask:
-                rel = Mask(spu_mask).global_to_relative_rank(g_rank)
-                self._spu_link_ctxs[g_rank] = link_ctx_list[rel]
+        for t in threads:
+            t.start()
+        for t in threads:
+            t.join()
+
+        # Check for exceptions during link creation
+        if exceptions:
+            first_exc = next(iter(exceptions.values()))
+            raise RuntimeError(
+                f"Failed to create SPU link contexts for ranks {list(exceptions.keys())}"
+            ) from first_exc
 
         self._spu_runtime_cfg = libspu.RuntimeConfig(
             protocol=spu_protocol, field=spu_field
@@ -187,10 +215,10 @@ class Simulator(InterpContext):
         This exposes potential MPIR serialization bugs by forcing expressions
         to go through the full serialize->deserialize cycle.
         """
-        writer = Writer()
+        writer = IrWriter()
         graph_proto = writer.dumps(expr)
 
-        reader = Reader()
+        reader = IrReader()
         deserialized_expr = reader.loads(graph_proto)
 
         if deserialized_expr is None:
@@ -202,8 +230,7 @@ class Simulator(InterpContext):
     def fetch(self, obj: MPObject) -> list[TensorLike]:
         if not isinstance(obj, SimVar):
             raise ValueError(f"Expected SimVar, got {type(obj)}")
-
-        return list(obj.values)
+        return [v.to_numpy() if hasattr(v, "to_numpy") else v for v in obj._values]
 
     # override
     def evaluate(self, expr: Expr, bindings: dict[str, MPObject]) -> Sequence[MPObject]:
@@ -213,7 +240,7 @@ class Simulator(InterpContext):
                 raise ValueError(f"Variable {name} not in this context, got {var.ctx}.")
 
         pts_env = [
-            {name: cast(SimVar, var).values[rank] for name, var in bindings.items()}
+            {name: cast(SimVar, var)._values[rank] for name, var in bindings.items()}
             for rank in range(self.world_size())
         ]