mplang-nightly 0.1.dev182__tar.gz → 0.1.dev184__tar.gz

{mplang_nightly-0.1.dev182 → mplang_nightly-0.1.dev184}/PKG-INFO

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

{mplang_nightly-0.1.dev182 → mplang_nightly-0.1.dev184}/mplang/core/cluster.py

@@ -20,6 +20,7 @@ MPLang cluster configuration.
 from __future__ import annotations
 
 from dataclasses import dataclass, field
+from functools import cached_property
 from typing import Any
 
 
@@ -169,6 +170,16 @@ class ClusterSpec:
             },
         }
 
+    @cached_property
+    def endpoints(self) -> list[str]:
+        eps: list[str] = []
+        for n in sorted(
+            self.nodes.values(),
+            key=lambda x: x.rank,  # type: ignore[attr-defined]
+        ):
+            eps.append(n.endpoint)
+        return eps
+
     @classmethod
     def from_dict(cls, config: dict[str, Any]) -> ClusterSpec:
         """Parses a raw config dictionary and returns a validated ClusterSpec."""

{mplang_nightly-0.1.dev182 → mplang_nightly-0.1.dev184}/mplang/core/comm.py

@@ -48,6 +48,10 @@ class ICommunicator(ABC):
     def recv(self, frm: int, key: str) -> Any:
         """Receive data from peer with the given key"""
 
+    @abstractmethod
+    def onSent(self, frm: int, key: str, data: Any) -> None:
+        """Called when a key is sent to self"""
+
 
 class ICollective(ABC):
     """Interface for collective communication"""

{mplang_nightly-0.1.dev182 → mplang_nightly-0.1.dev184}/mplang/runtime/communicator.py

@@ -29,12 +29,25 @@ from mplang.kernels.value import Value, decode_value, encode_value
 
 class HttpCommunicator(CommunicatorBase):
     def __init__(self, session_name: str, rank: int, endpoints: list[str]):
+        # Validate endpoints
+        if not endpoints:
+            raise ValueError("endpoints cannot be empty")
+
+        if not all(endpoint for endpoint in endpoints):
+            raise ValueError("endpoints cannot contain empty elements")
+
         super().__init__(rank, len(endpoints))
         self.session_name = session_name
-        self.endpoints = endpoints
+        # Ensure all endpoints have protocol prefix
+        self.endpoints = [
+            endpoint
+            if endpoint.startswith(("http://", "https://"))
+            else f"http://{endpoint}"
+            for endpoint in endpoints
+        ]
         self._counter = 0
         logging.info(
-            f"HttpCommunicator initialized: session={session_name}, rank={rank}, endpoints={endpoints}"
+            f"HttpCommunicator initialized: session={session_name}, rank={rank}, endpoints={self.endpoints}"
         )
 
     # override

{mplang_nightly-0.1.dev182 → mplang_nightly-0.1.dev184}/mplang/runtime/server.py

@@ -31,6 +31,7 @@ from fastapi.responses import JSONResponse
 from pydantic import BaseModel
 
 from mplang.core import IrReader, TableType, TensorType
+from mplang.core.cluster import ClusterSpec
 from mplang.kernels.base import KernelContext
 from mplang.kernels.value import Value, decode_value, encode_value
 from mplang.protos.v1alpha1 import mpir_pb2
@@ -40,6 +41,7 @@ from mplang.runtime.session import (
     Computation,
     Session,
     Symbol,
+    create_session_from_spec,
 )
 
 logger = logging.getLogger(__name__)
@@ -260,15 +262,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)
 

{mplang_nightly-0.1.dev182 → mplang_nightly-0.1.dev184}/mplang/runtime/session.py

@@ -34,6 +34,8 @@ from urllib.parse import urlparse
 
 import spu.libspu as libspu
 
+from mplang.core.cluster import ClusterSpec
+from mplang.core.comm import ICommunicator
 from mplang.core.expr.ast import Expr
 from mplang.core.expr.evaluator import IEvaluator, create_evaluator
 from mplang.core.mask import Mask
@@ -96,19 +98,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
@@ -119,18 +127,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
@@ -191,10 +190,11 @@ class Session:
         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):
+            for r, addr in enumerate(self.cluster_spec.endpoints):
                 if r in self.spu_mask:
-                    if "//" not in addr:
-                        addr = f"//{addr}"
+                    # TODO(oeqqwq): addr may contain other schema like grpc://
+                    if not addr.startswith(("http://", "https://")):
+                        addr = f"http://{addr}"
                     parsed = urlparse(addr)
                     assert isinstance(parsed.port, int)
                     new_addr = f"{parsed.hostname}:{parsed.port + SPU_PORT_OFFSET}"
@@ -281,12 +281,17 @@ class Session:
                 )
             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_nightly-0.1.dev182 → mplang_nightly-0.1.dev184}/tests/runtime/test_communicator.py

@@ -24,7 +24,7 @@ import uvicorn
 from mplang.core.cluster import ClusterSpec
 from mplang.runtime.communicator import HttpCommunicator
 from mplang.runtime.server import app, register_session
-from mplang.runtime.session import Session
+from mplang.runtime.session import create_session_from_spec
 from tests.utils.server_fixtures import http_servers  # noqa: F401 (fixture)
 
 
@@ -194,10 +194,12 @@ def run_party_e2e_process(
             runtime_version="test",
             runtime_platform="test",
         ).to_dict()
-        sess = Session.from_cluster_spec_dict(
+        # Create ClusterSpec from dict, then create session
+        cluster_spec = ClusterSpec.from_dict(cluster_spec_dict)
+        sess = create_session_from_spec(
             name=session_name,
             rank=rank,
-            spec_dict=cluster_spec_dict,
+            spec=cluster_spec,
         )
         register_session(sess)