mplang-nightly 0.1.dev152__tar.gz → 0.1.dev153__tar.gz

{mplang_nightly-0.1.dev152 → mplang_nightly-0.1.dev153}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: mplang-nightly
-Version: 0.1.dev152
+Version: 0.1.dev153
 Summary: Multi-Party Programming Language
 Author-email: SecretFlow Team <secretflow-contact@service.alipay.com>
 License:                                  Apache License

{mplang_nightly-0.1.dev152 → mplang_nightly-0.1.dev153}/mplang/core/cluster.py

@@ -25,23 +25,28 @@ from typing import Any
 
 @dataclass(frozen=True)
 class RuntimeInfo:
-    """
-    Structured representation of a Physical Node's runtime capabilities.
+    """Per-physical-node runtime configuration.
+
+    ``op_bindings`` is a per-node override map (logical_op -> kernel_id) merged
+    into that node's ``RuntimeContext``. Unknown future / auxiliary fields are
+    preserved in ``extra``.
     """
 
     version: str
     platform: str
-    backends: list[str]
+    # Per-node partial override dispatch table (merged over project defaults).
+    op_bindings: dict[str, str] = field(default_factory=dict)
 
-    # A catch-all for any other custom or future properties
+    # A catch-all for any other custom or future properties (must not collide
+    # with reserved keys: version, platform, op_bindings).
     extra: dict[str, Any] = field(default_factory=dict)
 
     def to_dict(self) -> dict[str, Any]:
-        """Convert RuntimeInfo to a dictionary."""
+        """Convert RuntimeInfo to a dictionary (stable field names)."""
         result = {
             "version": self.version,
             "platform": self.platform,
-            "backends": self.backends,
+            "op_bindings": self.op_bindings,
         }
         result.update(self.extra)
         return result
@@ -175,7 +180,8 @@ class ClusterSpec:
 
         # 2. Parse Physical Nodes, using the list index as the rank
         nodes_map: dict[str, Node] = {}
-        known_runtime_fields = {"version", "platform", "backends"}
+        # Reserved runtime info keys we recognize explicitly.
+        known_runtime_fields = {"version", "platform", "op_bindings"}
         for i, node_cfg in enumerate(config["nodes"]):
             if "rank" in node_cfg:
                 # Optionally, we can log a warning that the explicit 'rank' is ignored.
@@ -187,11 +193,12 @@ class ClusterSpec:
                 for k, v in runtime_info_cfg.items()
                 if k not in known_runtime_fields
             }
-
+            # Gracefully ignore legacy 'backends' if present (treated as extra)
+            # for backward compatibility.
             runtime_info = RuntimeInfo(
                 version=runtime_info_cfg.get("version", "N/A"),
                 platform=runtime_info_cfg.get("platform", "N/A"),
-                backends=runtime_info_cfg.get("backends", []),
+                op_bindings=runtime_info_cfg.get("op_bindings", {}) or {},
                 extra=extra_runtime_info,
             )
 
@@ -227,32 +234,96 @@ class ClusterSpec:
         return cls(nodes=nodes_map, devices=devices_map)
 
     @classmethod
-    def simple(cls, world_size: int) -> ClusterSpec:
-        """Creates a simple cluster spec for simulation with the given number of parties."""
-        nodes = {
-            f"node{i}": Node(
+    def simple(
+        cls,
+        world_size: int,
+        *,
+        endpoints: list[str] | None = None,
+        spu_protocol: str = "SEMI2K",
+        spu_field: str = "FM128",
+        runtime_version: str = "simulated",
+        runtime_platform: str = "simulated",
+        op_bindings: list[dict[str, str]] | None = None,
+        enable_local_device: bool = True,
+        enable_spu_device: bool = True,
+    ) -> ClusterSpec:
+        """Convenience constructor used heavily in tests.
+
+        Parameters
+        ----------
+        world_size:
+            Number of parties (physical nodes).
+        endpoints:
+            Optional explicit endpoint list of length ``world_size``. Each element may
+            include scheme (``http://``) or not; stored verbatim. If not provided we
+            synthesize ``localhost:{5000 + i}`` (5000 is a fixed default; pass explicit
+            endpoints for control). Deprecated ``base_port`` legacy kwarg can adjust it.
+        spu_protocol / spu_field:
+            SPU device config values.
+        runtime_version / runtime_platform:
+            Populated into each node's ``RuntimeInfo``.
+        op_bindings:
+            Optional list of length ``world_size`` supplying per-node op_bindings
+            override dicts (defaults to empty dicts).
+        enable_local_device:
+            If True (default), create one ``local_{rank}`` device per node.
+        enable_spu_device:
+            If True (default) create a shared SPU device named ``SP0``.
+        """
+        base_port = 5000
+
+        if endpoints is not None and len(endpoints) != world_size:
+            raise ValueError(
+                "len(endpoints) must equal world_size when provided: "
+                f"{len(endpoints)} != {world_size}"
+            )
+
+        if op_bindings is not None and len(op_bindings) != world_size:
+            raise ValueError(
+                "len(op_bindings) must equal world_size when provided: "
+                f"{len(op_bindings)} != {world_size}"
+            )
+
+        if not enable_local_device and not enable_spu_device:
+            raise ValueError(
+                "At least one of enable_local_device or enable_spu_device must be True"
+            )
+
+        nodes: dict[str, Node] = {}
+        for i in range(world_size):
+            ep = endpoints[i] if endpoints is not None else f"localhost:{base_port + i}"
+            node_op_bindings = op_bindings[i] if op_bindings is not None else {}
+            nodes[f"node{i}"] = Node(
                 name=f"node{i}",
                 rank=i,
-                endpoint=f"localhost:{5000 + i}",
+                endpoint=ep,
                 runtime_info=RuntimeInfo(
-                    version="simulated",
-                    platform="simulated",
-                    backends=["__all__"],
+                    version=runtime_version,
+                    platform=runtime_platform,
+                    op_bindings=node_op_bindings,
                 ),
             )
-            for i in range(world_size)
-        }
 
-        devices = {
-            "SP0": Device(
+        devices: dict[str, Device] = {}
+        # Optional per-node local devices
+        if enable_local_device:
+            for i in range(world_size):
+                devices[f"local_{i}"] = Device(
+                    name=f"local_{i}",
+                    kind="local",
+                    members=[nodes[f"node{i}"]],
+                )
+
+        # Shared SPU device
+        if enable_spu_device:
+            devices["SP0"] = Device(
                 name="SP0",
                 kind="SPU",
                 members=list(nodes.values()),
                 config={
-                    "protocol": "SEMI2K",
-                    "field": "FM128",
+                    "protocol": spu_protocol,
+                    "field": spu_field,
                 },
             )
-        }
 
         return cls(nodes=nodes, devices=devices)

{mplang_nightly-0.1.dev152 → mplang_nightly-0.1.dev153}/mplang/runtime/client.py

@@ -81,21 +81,14 @@ class HttpExecutorClient:
         self,
         name: str,
         rank: int,
-        endpoints: list[str],
-        *,
-        spu_mask: int = 0,
-        spu_protocol: str = "SEMI2K",
-        spu_field: str = "FM64",
+        cluster_spec: dict,
     ) -> str:
         """Create a new session.
 
         Args:
             name: Session name/ID.
-            rank: The rank of this party in the session.
-            endpoints: List of endpoint URLs for all parties, indexed by rank.
-            spu_mask: SPU mask for the session, 0 means no SPU.
-            spu_protocol: SPU protocol for the session (e.g., "SEMI2K", "ABY3").
-            spu_field: SPU field for the session (e.g., "FM64", "FM128").
+            rank: This party's rank.
+            cluster_spec: Full cluster specification dict (ClusterSpec.to_dict()).
 
         Returns:
             The session name/ID
@@ -104,14 +97,7 @@ class HttpExecutorClient:
             RuntimeError: If session creation fails
         """
         url = f"/sessions/{name}"
-
-        payload: dict[str, Any] = {
-            "rank": rank,
-            "endpoints": endpoints,
-            "spu_mask": spu_mask,
-            "spu_protocol": spu_protocol,
-            "spu_field": spu_field,
-        }
+        payload: dict[str, Any] = {"rank": rank, "cluster_spec": cluster_spec}
 
         try:
             response = await self._client.put(url, json=payload)

{mplang_nightly-0.1.dev152 → mplang_nightly-0.1.dev153}/mplang/runtime/driver.py

@@ -145,8 +145,6 @@ class Driver(InterpContext):
         """Get existing session or create a new one across all HTTP servers."""
         if self._session_id is None:
             new_session_id = new_uuid()
-            endpoints_list = list(self.node_addrs.values())
-
             # Create temporary clients for session creation
             clients = self._create_clients()
             try:
@@ -158,10 +156,7 @@ class Driver(InterpContext):
                     task = client.create_session(
                         name=new_session_id,
                         rank=rank,
-                        endpoints=endpoints_list,
-                        spu_mask=self.spu_mask_int,
-                        spu_protocol=self.spu_protocol_str,
-                        spu_field=self.spu_field_str,
+                        cluster_spec=self.cluster_spec.to_dict(),
                     )
                     tasks.append(task)
 

{mplang_nightly-0.1.dev152 → mplang_nightly-0.1.dev153}/mplang/runtime/server.py

@@ -32,14 +32,30 @@ 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 import resource
 from mplang.runtime.data_providers import DataProvider, ResolvedURI, register_provider
 from mplang.runtime.exceptions import InvalidRequestError, ResourceNotFound
+from mplang.runtime.session import (
+    Computation,
+    Session,
+    Symbol,
+)
 
 logger = logging.getLogger(__name__)
 
 app = FastAPI()
 
+# per-server global state
+_sessions: dict[str, Session] = {}
+_global_symbols: dict[str, Symbol] = {}
+
+
+def register_session(session: Session) -> Session:  # pragma: no cover - test helper
+    existing = _sessions.get(session.name)
+    if existing:
+        return existing
+    _sessions[session.name] = session
+    return session
+
 
 class _SymbolsProvider(DataProvider):
     """Server-local symbols provider backed by BackendRuntime.state."""
@@ -83,7 +99,7 @@ class _SymbolsProvider(DataProvider):
         ctx: KernelContext,
     ) -> Any:  # type: ignore[override]
         name = self._symbol_name(uri)
-        sym = resource.get_global_symbol(name)
+        sym = _global_symbols.get(name)
         if sym is None:
             raise ResourceNotFound(f"Global symbol '{name}' not found")
         return sym.data
@@ -102,8 +118,13 @@ class _SymbolsProvider(DataProvider):
             raise InvalidRequestError(
                 f"Failed to encode value for symbols:// write: {e!s}"
             ) from e
-
-        resource.create_global_symbol(name, {}, data_b64)
+        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)
 
 
 # Register symbols provider explicitly for server runtime
@@ -168,11 +189,7 @@ def validate_name(name: str, name_type: str) -> None:
 # Request/Response Models
 class CreateSessionRequest(BaseModel):
     rank: int
-    endpoints: list[str]
-    # SPU related
-    spu_mask: int
-    spu_protocol: str
-    spu_field: str
+    cluster_spec: dict
 
 
 class SessionResponse(BaseModel):
@@ -229,7 +246,7 @@ async def health_check() -> dict[str, str]:
 @app.get("/sessions", response_model=SessionListResponse)
 def list_sessions() -> SessionListResponse:
     """List all session names."""
-    return SessionListResponse(sessions=resource.list_all_sessions())
+    return SessionListResponse(sessions=list(_sessions.keys()))
 
 
 # List all computations in a session
@@ -238,39 +255,44 @@ def list_sessions() -> SessionListResponse:
 )
 def list_session_computations(session_name: str) -> ComputationListResponse:
     """List all computation names in a session."""
-    session = resource.get_session(session_name)
-    if not session:
+    sess = _sessions.get(session_name)
+    if not sess:
         raise ResourceNotFound(f"Session '{session_name}' not found")
-    return ComputationListResponse(computations=list(session.computations.keys()))
+    return ComputationListResponse(computations=sess.list_computations())
 
 
 # Session endpoints
 @app.put("/sessions/{session_name}", response_model=SessionResponse)
 def create_session(session_name: str, request: CreateSessionRequest) -> SessionResponse:
     validate_name(session_name, "session")
-    session = resource.create_session(
-        name=session_name,
-        rank=request.rank,
-        endpoints=request.endpoints,
-        spu_mask=request.spu_mask,
-        spu_protocol=request.spu_protocol,
-        spu_field=request.spu_field,
-    )
-    return SessionResponse(name=session.name)
+    # 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)
+        _sessions[session_name] = sess
+    return SessionResponse(name=sess.name)
 
 
 @app.get("/sessions/{session_name}", response_model=SessionResponse)
 def get_session(session_name: str) -> SessionResponse:
-    session = resource.get_session(session_name)
-    if not session:
+    sess = _sessions.get(session_name)
+    if not sess:
         raise ResourceNotFound(f"Session '{session_name}' not found")
-    return SessionResponse(name=session.name)
+    return SessionResponse(name=sess.name)
 
 
 @app.delete("/sessions/{session_name}")
 def delete_session(session_name: str) -> dict[str, str]:
     """Delete a session and all its associated resources."""
-    if resource.delete_session(session_name):
+    if session_name in _sessions:
+        del _sessions[session_name]
+        logging.info(f"Session {session_name} deleted successfully")
         return {"message": f"Session '{session_name}' deleted successfully"}
     else:
         raise ResourceNotFound(f"Session '{session_name}' not found")
@@ -299,18 +321,25 @@ def create_and_execute_computation(
         raise InvalidRequestError("Failed to parse expression from protobuf")
 
     # Create the computation resource
-    computation = resource.create_computation(session_name, computation_id, expr)
-    # Execute with input/output names
-    resource.execute_computation(
-        session_name, computation.name, request.input_names, request.output_names
-    )
-    return ComputationResponse(name=computation.name)
+    sess = _sessions.get(session_name)
+    if not sess:
+        raise ResourceNotFound(f"Session '{session_name}' not found.")
+    comp = sess.get_computation(computation_id)
+    if not comp:
+        comp = Computation(name=computation_id, expr=expr)
+        sess.add_computation(comp)
+    sess.execute(comp, request.input_names, request.output_names)
+    return ComputationResponse(name=computation_id)
 
 
 @app.delete("/sessions/{session_name}/computations/{computation_id}")
 def delete_computation(session_name: str, computation_id: str) -> dict[str, str]:
     """Delete a specific computation."""
-    if resource.delete_computation(session_name, computation_id):
+    sess = _sessions.get(session_name)
+    if sess and sess.delete_computation(computation_id):
+        logging.info(
+            f"Computation {computation_id} deleted from session {session_name}"
+        )
         return {"message": f"Computation '{computation_id}' deleted successfully"}
     else:
         raise ResourceNotFound(
@@ -326,9 +355,15 @@ def create_session_symbol(
     session_name: str, symbol_name: str, request: CreateSymbolRequest
 ) -> SymbolResponse:
     """Create a symbol in a session."""
-    symbol = resource.create_symbol(
-        session_name, symbol_name, request.mptype, request.data
-    )
+    sess = _sessions.get(session_name)
+    if not sess:
+        raise ResourceNotFound(f"Session '{session_name}' not found.")
+    try:
+        obj = pickle.loads(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)
+    sess.add_symbol(symbol)
     # Return the base64 data back to client; server stores Python object
     return SymbolResponse(
         name=symbol.name,
@@ -346,8 +381,8 @@ def get_session_symbol(session_name: str, symbol_name: str) -> SymbolResponse:
         logger.debug(
             f"Looking for symbol: '{symbol_name}' in session: '{session_name}'"
         )
-
-        symbol = resource.get_symbol(session_name, symbol_name)
+        sess = _sessions.get(session_name)
+        symbol = sess.get_symbol(symbol_name) if sess else None
         if not symbol:
             raise HTTPException(
                 status_code=404, detail=f"Symbol {symbol_name} not found"
@@ -368,14 +403,19 @@ def get_session_symbol(session_name: str, symbol_name: str) -> SymbolResponse:
 @app.get("/sessions/{session_name}/symbols")
 def list_session_symbols(session_name: str) -> dict[str, list[str]]:
     """List all symbols in a session."""
-    symbols = resource.list_symbols(session_name)
+    sess = _sessions.get(session_name)
+    if not sess:
+        raise ResourceNotFound(f"Session '{session_name}' not found.")
+    symbols = sess.list_symbols()
     return {"symbols": symbols}
 
 
 @app.delete("/sessions/{session_name}/symbols/{symbol_name}")
 def delete_symbol(session_name: str, symbol_name: str) -> dict[str, str]:
     """Delete a specific symbol."""
-    if resource.delete_symbol(session_name, symbol_name):
+    sess = _sessions.get(session_name)
+    if sess and sess.delete_symbol(symbol_name):
+        logging.info(f"Symbol {symbol_name} deleted from session {session_name}")
         return {"message": f"Symbol '{symbol_name}' deleted successfully"}
     else:
         raise ResourceNotFound(
@@ -389,13 +429,18 @@ def create_global_symbol(
     symbol_name: str, request: CreateSymbolRequest
 ) -> GlobalSymbolResponse:
     validate_name(symbol_name, "symbol")
-    sym = resource.create_global_symbol(symbol_name, request.mptype, request.data)
+    try:
+        obj = pickle.loads(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)
 
 
 @app.get("/api/v1/symbols/{symbol_name}", response_model=GlobalSymbolResponse)
-def get_global_symbol(symbol_name: str) -> GlobalSymbolResponse:
-    sym = resource.get_global_symbol(symbol_name)
+def get_global_symbol(symbol_name: str) -> GlobalSymbolResponse:  # route handler
+    sym = _global_symbols.get(symbol_name)
     if not sym:
         raise ResourceNotFound(f"Global symbol '{symbol_name}' not found")
     data_bytes = pickle.dumps(sym.data)
@@ -405,12 +450,13 @@ def get_global_symbol(symbol_name: str) -> GlobalSymbolResponse:
 
 @app.get("/api/v1/symbols")
 def list_global_symbols() -> dict[str, list[str]]:
-    return {"symbols": resource.list_global_symbols()}
+    return {"symbols": list(_global_symbols.keys())}
 
 
 @app.delete("/api/v1/symbols/{symbol_name}")
-def delete_global_symbol(symbol_name: str) -> dict[str, str]:
-    if resource.delete_global_symbol(symbol_name):
+def delete_global_symbol(symbol_name: str) -> dict[str, str]:  # route handler
+    if symbol_name in _global_symbols:
+        del _global_symbols[symbol_name]
         return {"message": f"Global symbol '{symbol_name}' deleted successfully"}
     else:
         raise ResourceNotFound(f"Global symbol '{symbol_name}' not found")
@@ -426,8 +472,8 @@ def comm_send(
     Receive a message from another party and deliver it to the session's communicator.
     This endpoint runs on the receiver's server.
     """
-    session = resource.get_session(session_name)
-    if not session or not session.communicator:
+    sess = _sessions.get(session_name)
+    if not sess or not sess.communicator:
         logger.error(f"Session or communicator not found: session={session_name}")
         raise HTTPException(status_code=404, detail="Session or communicator not found")
 
@@ -435,5 +481,5 @@ def comm_send(
     # We don't need to validate to_rank since the request is coming to this server
 
     # Use the proper onSent mechanism from CommunicatorBase
-    session.communicator.onSent(from_rank, key, request.data)
+    sess.communicator.onSent(from_rank, key, request.data)
     return {"status": "ok"}