mplang-nightly 0.1.dev330__tar.gz → 0.1.dev332__tar.gz

{mplang_nightly-0.1.dev330 → mplang_nightly-0.1.dev332}/PKG-INFO

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

{mplang_nightly-0.1.dev330 → mplang_nightly-0.1.dev332}/mplang/backends/simp_driver/mem.py

@@ -23,9 +23,10 @@ from collections.abc import Callable
 from typing import TYPE_CHECKING, Any, cast
 
 from mplang.backends.simp_driver.state import SimpDriver
-from mplang.backends.simp_worker import WORKER_HANDLERS, SimpWorker
-from mplang.backends.simp_worker.comm_context import CommContext
+from mplang.backends.simp_worker import WORKER_HANDLERS
+from mplang.backends.simp_worker.infra import DEFAULT_ASYNC_OPS, WorkerInfra
 from mplang.backends.simp_worker.mem import LocalMesh
+from mplang.backends.simp_worker.request import create_request_interpreter
 from mplang.runtime.interpreter import ExecutionTracer, Interpreter
 from mplang.runtime.object_store import FileSystemBackend, ObjectStore
 
@@ -37,10 +38,14 @@ if TYPE_CHECKING:
 
 
 class MemCluster:
-    """Orchestrator that creates and manages local worker Interpreters.
+    """Orchestrator that creates and manages local worker infrastructure.
 
     This class handles worker lifecycle management. It does NOT attach to
     an Interpreter - instead, it creates a SimpMemDriver that can be attached.
+
+    Per-request Interpreters are created on-the-fly by ``create_request_interpreter``
+    to isolate mutable state (CommContext, SimpWorker, SPUState) across concurrent
+    requests.
     """
 
     def __init__(
@@ -75,8 +80,9 @@ class MemCluster:
         )
         self.tracer.start()
 
-        # Create Workers
-        self._workers: list[Interpreter] = []
+        # Create shared WorkerInfra per rank (replaces per-rank Interpreters)
+        self._infras: list[WorkerInfra] = []
+        self._stores: list[ObjectStore] = []
         for rank in range(world_size):
             worker_root = cluster_root / f"node{rank}"
             store = ObjectStore(
@@ -84,60 +90,50 @@ class MemCluster:
                     root_path=str(worker_root / "store"),
                 )
             )
+            self._stores.append(store)
 
-            worker_state = SimpWorker(
+            w_handlers: dict[str, Callable[..., Any]] = {**WORKER_HANDLERS}  # type: ignore[dict-item]
+            infra = WorkerInfra(
                 rank=rank,
                 world_size=world_size,
                 communicator=self._mesh.comms[rank],
                 store=store,
-            )
-
-            w_handlers: dict[str, Callable[..., Any]] = {**WORKER_HANDLERS}  # type: ignore[dict-item]
-            comm_ctx = CommContext(
-                self._mesh.comms[rank],
-                context_id="ctx",
-                my_rank=rank,
-            )
-            w_interp = Interpreter(
-                name=f"Worker-{rank}",
+                handlers=w_handlers,
                 tracer=self.tracer,
                 trace_pid=rank,
-                store=store,
                 root_dir=worker_root,
-                handlers=w_handlers,
-                comm_ctx=comm_ctx,
+                # async_ops has no effect when executor is None (MemCluster
+                # runs each request synchronously on the mesh executor thread).
+                # We still set it so that WorkerInfra carries the canonical set
+                # for introspection and consistency with the HTTP Worker path.
+                async_ops=DEFAULT_ASYNC_OPS,
             )
-            w_interp.set_dialect_state("simp", worker_state)
-
-            w_interp.async_ops = {
-                "bfv.add",
-                "bfv.mul",
-                "bfv.rotate",
-                "bfv.batch_encode",
-                "bfv.relinearize",
-                "bfv.encrypt",
-                "bfv.decrypt",
-                "field.solve_okvs",
-                "field.decode_okvs",
-                "field.aes_expand",
-                "field.mul",
-                "simp.shuffle",
-            }
-            self._workers.append(w_interp)
+            self._infras.append(infra)
 
     @property
     def world_size(self) -> int:
         return self._world_size
 
     @property
-    def workers(self) -> list[Interpreter]:
-        return self._workers
+    def infras(self) -> list[WorkerInfra]:
+        return self._infras
+
+    @property
+    def workers(self) -> list[WorkerInfra]:
+        """Backward-compatible alias for ``infras``.
+
+        .. note:: Returns ``WorkerInfra`` objects (not ``Interpreter``).  Only
+           the ``.store`` attribute is guaranteed by this interface.  Callers
+           needing full Interpreter access should use
+           ``create_request_interpreter(infra, job_id)`` instead.
+        """
+        return self._infras
 
     def create_state(self) -> SimpMemDriver:
         """Create a SimpMemDriver that can be attached to a Driver Interpreter."""
         return SimpMemDriver(
             world_size=self._world_size,
-            workers=self._workers,
+            infras=self._infras,
             mesh=self._mesh,
         )
 
@@ -158,11 +154,11 @@ class SimpMemDriver(SimpDriver):
     def __init__(
         self,
         world_size: int,
-        workers: list[Interpreter],
+        infras: list[WorkerInfra],
         mesh: Any,  # LocalMesh from simp_worker.mem
     ) -> None:
         self._world_size = world_size
-        self._workers = workers
+        self._infras = infras
         self._mesh = mesh
 
     def shutdown(self) -> None:
@@ -174,9 +170,18 @@ class SimpMemDriver(SimpDriver):
         return self._world_size
 
     @property
-    def workers(self) -> list[Interpreter]:
-        """Worker interpreters (for backward compatibility)."""
-        return self._workers
+    def infras(self) -> list[WorkerInfra]:
+        """Worker infrastructure (shared, immutable)."""
+        return self._infras
+
+    @property
+    def workers(self) -> list[WorkerInfra]:
+        """Backward-compatible alias for ``infras``.
+
+        .. note:: Returns ``WorkerInfra`` objects (not ``Interpreter``).  Only
+           the ``.store`` attribute is guaranteed by this interface.
+        """
+        return self._infras
 
     def submit(
         self, rank: int, graph: Graph, inputs: list[Any], job_id: str | None = None
@@ -205,31 +210,32 @@ class SimpMemDriver(SimpDriver):
 
     def fetch(self, rank: int, uri: str) -> Future[Any]:
         """Fetch directly from worker store."""
-        worker = self._workers[rank]
-        worker_ctx = cast(SimpWorker, worker.get_dialect_state("simp"))
-        return self._mesh.executor.submit(lambda: worker_ctx.store.get(uri))  # type: ignore[no-any-return]
+        infra = self._infras[rank]
+        return self._mesh.executor.submit(lambda: infra.store.get(uri))  # type: ignore[no-any-return]
 
     def _run_worker(
         self, rank: int, graph: Graph, inputs: list[Any], job_id: str | None = None
     ) -> Any:
-        """Execute on worker interpreter."""
-        worker_interp = self._workers[rank]
-        worker_ctx = cast(SimpWorker, worker_interp.get_dialect_state("simp"))
-
-        # Resolve URI inputs (None means rank has no data)
-        resolved_inputs = [
-            worker_ctx.store.get(inp) if inp is not None else None for inp in inputs
-        ]
-
-        # Execute
-        results = worker_interp.evaluate_graph(graph, resolved_inputs, job_id)
-
-        # Store results (results is always a list)
-        if not graph.outputs:
-            return None
-        return [
-            worker_ctx.store.put(res) if res is not None else None for res in results
-        ]
+        """Execute on a per-request Interpreter."""
+        infra = self._infras[rank]
+        request_interp = create_request_interpreter(infra, job_id or "anonymous")
+        try:
+            # Resolve URI inputs (None means rank has no data)
+            resolved_inputs = [
+                infra.store.get(inp) if inp is not None else None for inp in inputs
+            ]
+
+            # Execute
+            results = request_interp.evaluate_graph(graph, resolved_inputs, job_id)
+
+            # Store results (results is always a list)
+            if not graph.outputs:
+                return None
+            return [
+                infra.store.put(res) if res is not None else None for res in results
+            ]
+        finally:
+            request_interp.shutdown()
 
 
 def make_simulator(

{mplang_nightly-0.1.dev330 → mplang_nightly-0.1.dev332}/mplang/backends/simp_worker/__init__.py

@@ -38,11 +38,14 @@ from mplang.backends.simp_worker.http import (
     RecvTimeoutError,
     SendTimeoutError,
 )
+from mplang.backends.simp_worker.infra import DEFAULT_ASYNC_OPS, WorkerInfra
 from mplang.backends.simp_worker.mem import LocalMesh, ThreadCommunicator
 from mplang.backends.simp_worker.ops import WORKER_HANDLERS
+from mplang.backends.simp_worker.request import create_request_interpreter
 from mplang.backends.simp_worker.state import SimpWorker
 
 __all__ = [
+    "DEFAULT_ASYNC_OPS",
     "WORKER_HANDLERS",
     "CommConfig",
     "CommStats",
@@ -56,6 +59,8 @@ __all__ = [
     "SendTimeoutError",
     "SimpWorker",
     "ThreadCommunicator",
+    "WorkerInfra",
+    "create_request_interpreter",
     "testall",
     "testany",
     "wait_all",

{mplang_nightly-0.1.dev330 → mplang_nightly-0.1.dev332}/mplang/backends/simp_worker/comm_context.py

@@ -62,15 +62,25 @@ class CommContext:
     def world_size(self) -> int:
         return self._comm.world_size
 
-    def spawn(self) -> CommContext:
+    def spawn(self, suffix: str | None = None) -> CommContext:
         """Create a child context with independent counter namespace.
 
-        Child ``context_id = f"{parent_id}.{spawn_seq}"``.  ``spawn_seq``
-        increments in the calling thread's order, so if all ranks call
-        ``spawn()`` in the same program-order position, children get matching
-        IDs across ranks.
+        Args:
+            suffix: Explicit suffix for the child context ID.  When provided,
+                ``context_id = f"{parent_id}.{suffix}"``.  This is used when
+                the caller needs a deterministic, content-based ID that is
+                stable across ranks regardless of call order (e.g. the async
+                DAG scheduler uses ``f"{graph_exec_key}.{op_idx}"``).
+
+                When *None* (default), an auto-incrementing counter is used:
+                ``context_id = f"{parent_id}.{spawn_seq}"``.  This is suitable
+                for SPMD code where all ranks call ``spawn()`` in the same
+                program-order position.
         """
-        child_id = f"{self._id}.{self._spawn_counter}"
+        if suffix is None:
+            child_id = f"{self._id}.{self._spawn_counter}"
+        else:
+            child_id = f"{self._id}.{self._spawn_counter}.{suffix}"
         self._spawn_counter += 1
         return CommContext(self._comm, child_id, self._rank)
 

{mplang_nightly-0.1.dev330 → mplang_nightly-0.1.dev332}/mplang/backends/simp_worker/http.py

@@ -61,10 +61,10 @@ from mplang.backends.simp_worker.base import (
     wait_all,
     wait_any,
 )
-from mplang.backends.simp_worker.state import SimpWorker
+from mplang.backends.simp_worker.infra import WorkerInfra
 from mplang.edsl import serde
 from mplang.edsl.graph import Graph
-from mplang.runtime.interpreter import ExecutionTracer, Interpreter
+from mplang.runtime.interpreter import ExecutionTracer
 from mplang.runtime.object_store import FileSystemBackend, ObjectStore
 from mplang.utils.logging import get_logger
 
@@ -636,7 +636,7 @@ def register_routes(
     *,
     rank: int,
     world_size: int,
-    worker: Interpreter,
+    infra: WorkerInfra,
     comm: HttpCommunicator,
     store: ObjectStore,
     exec_pool: concurrent.futures.ThreadPoolExecutor,
@@ -648,36 +648,38 @@ def register_routes(
     enabling reuse in custom server setups.
 
     Note:
-        The *worker* ``Interpreter`` must have the ``"simp"`` dialect state
-        registered (via ``worker.set_dialect_state("simp", ctx)``) before
-        any request is served, as the ``/fetch`` and ``/objects`` endpoints
-        rely on it.
+        The *infra* ``WorkerInfra`` contains all shared infrastructure needed
+        to create per-request Interpreters.
 
     Args:
         app: The FastAPI application to register routes on.
         rank: This worker's rank.
         world_size: Total number of workers.
-        worker: The Interpreter instance for graph execution.
+        infra: The WorkerInfra for creating per-request Interpreters.
         comm: The HttpCommunicator for inter-worker communication.
         store: The ObjectStore for data persistence.
         exec_pool: Thread pool for executing graphs.
     """
-    from typing import cast
+    from mplang.backends.simp_worker.request import create_request_interpreter
 
     def _do_execute(graph: Graph, inputs: list[Any], job_id: str | None = None) -> Any:
-        """Execute graph in worker thread."""
-        # Resolve URI inputs (None means rank has no data)
-        resolved_inputs = [
-            store.get(inp) if inp is not None else None for inp in inputs
-        ]
+        """Execute graph in worker thread with per-request isolation."""
+        request_interp = create_request_interpreter(infra, job_id or "anonymous")
+        try:
+            # Resolve URI inputs (None means rank has no data)
+            resolved_inputs = [
+                store.get(inp) if inp is not None else None for inp in inputs
+            ]
 
-        result = worker.evaluate_graph(graph, resolved_inputs)
-        comm.wait_pending_sends()
+            result = request_interp.evaluate_graph(graph, resolved_inputs, job_id)
+            comm.wait_pending_sends()
 
-        # Store results and return URIs (result is always a list)
-        if not graph.outputs:
-            return None
-        return [store.put(res) if res is not None else None for res in result]
+            # Store results and return URIs (result is always a list)
+            if not graph.outputs:
+                return None
+            return [store.put(res) if res is not None else None for res in result]
+        finally:
+            request_interp.shutdown()
 
     # Mutable closure state shared between endpoint handlers and exec_pool threads.
     async_tasks: dict[str, AsyncTaskState] = {}
@@ -845,8 +847,7 @@ def register_routes(
         """Fetch data by URI (e.g. ``mem://abc123``, ``fs://ckpt/s100``)."""
         logger.debug(f"Worker {rank} received fetch request for {req.uri}")
         try:
-            state = cast(SimpWorker, worker.get_dialect_state("simp"))
-            val = state.store.get(req.uri)
+            val = store.get(req.uri)
             return {"result": serde.dumps_b64(val)}
         except Exception as e:
             logger.error(f"Worker {rank} fetch failed: {e}")
@@ -856,8 +857,7 @@ def register_routes(
     def list_objects() -> dict[str, list[str]]:
         """List all objects in the worker's store (transient + persistent)."""
         try:
-            state = cast(SimpWorker, worker.get_dialect_state("simp"))
-            return {"objects": state.store.list_keys()}
+            return {"objects": store.list_keys()}
         except Exception as e:
             logger.error(f"Worker {rank} list_objects failed: {e}")
             raise HTTPException(status_code=500, detail=str(e)) from e
@@ -917,22 +917,22 @@ def create_worker_app(
     comm = HttpCommunicator(rank, world_size, endpoints, tracer=tracer)
     if store is None:
         store = ObjectStore(persistent=FileSystemBackend(root_path=str(root_dir)))
-    ctx = SimpWorker(rank, world_size, comm, store, spu_endpoints)
 
     handlers: dict[str, Callable[..., Any]] = {**WORKER_HANDLERS}  # type: ignore[dict-item]
 
-    from mplang.backends.simp_worker.comm_context import CommContext
+    from mplang.backends.simp_worker.infra import DEFAULT_ASYNC_OPS, WorkerInfra
 
-    comm_ctx = CommContext(comm, context_id="ctx", my_rank=rank)
-    worker = Interpreter(
+    infra = WorkerInfra(
+        rank=rank,
+        world_size=world_size,
+        communicator=comm,
+        store=store,
+        handlers=handlers,
+        spu_endpoints=spu_endpoints,
         tracer=tracer,
         root_dir=root_dir,
-        handlers=handlers,
-        store=store,
-        comm_ctx=comm_ctx,
+        async_ops=DEFAULT_ASYNC_OPS,
     )
-    # Register SimpWorker context as 'simp' dialect state
-    worker.set_dialect_state("simp", ctx)
 
     exec_pool = concurrent.futures.ThreadPoolExecutor(
         max_workers=2, thread_name_prefix=f"exec_{rank}"
@@ -941,7 +941,7 @@ def create_worker_app(
     @asynccontextmanager
     async def lifespan(app: FastAPI) -> AsyncIterator[None]:
         yield
-        await asyncio.to_thread(worker.shutdown)
+        await asyncio.to_thread(infra.shutdown)
         await asyncio.to_thread(comm.shutdown)
         await asyncio.to_thread(exec_pool.shutdown, wait=True)
 
@@ -951,7 +951,7 @@ def create_worker_app(
         app,
         rank=rank,
         world_size=world_size,
-        worker=worker,
+        infra=infra,
         comm=comm,
         store=store,
         exec_pool=exec_pool,

mplang_nightly-0.1.dev332/mplang/backends/simp_worker/infra.py

@@ -0,0 +1,108 @@
+# Copyright 2026 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.
+
+"""WorkerInfra: shared infrastructure container for per-request Interpreter creation.
+
+Created once at Worker startup. Passed to ``create_request_interpreter()`` for
+each incoming request. All fields are either immutable or thread-safe.
+"""
+
+from __future__ import annotations
+
+import concurrent.futures
+import pathlib
+import threading
+from collections.abc import Callable
+from dataclasses import dataclass, field
+from typing import Any
+
+from mplang.backends.simp_worker.base import CommunicatorProtocol
+from mplang.runtime.interpreter import ExecutionTracer
+from mplang.runtime.object_store import ObjectStore
+
+# Opcodes eligible for async DAG scheduling when an Executor is present.
+# Kept as a module-level constant so that both MemCluster and HTTP Worker
+# share the same set without duplication.
+DEFAULT_ASYNC_OPS: frozenset[str] = frozenset({
+    "bfv.add",
+    "bfv.mul",
+    "bfv.rotate",
+    "bfv.batch_encode",
+    "bfv.relinearize",
+    "bfv.encrypt",
+    "bfv.decrypt",
+    "field.solve_okvs",
+    "field.decode_okvs",
+    "field.aes_expand",
+    "field.mul",
+    "simp.shuffle",
+})
+
+
+@dataclass
+class WorkerInfra:
+    """Shared infrastructure for a Worker process.
+
+    Created once at Worker startup. Passed to each per-request Interpreter
+    factory. All fields are either immutable or thread-safe.
+    """
+
+    rank: int
+    world_size: int
+    communicator: CommunicatorProtocol
+    store: ObjectStore
+    handlers: dict[str, Callable[..., Any]]
+    spu_endpoints: dict[int, str] | None = None
+    tracer: ExecutionTracer | None = None
+    trace_pid: int | None = None
+    root_dir: pathlib.Path | None = None
+    executor: concurrent.futures.Executor | None = None
+    async_ops: frozenset[str] = field(default_factory=frozenset)
+
+    # SPU template links (lazily populated, protected by lock).
+    # Cache keys are (local_rank, spu_world_size, protocol, field, link_mode).
+    # The number of distinct keys is bounded by the Cartesian product of SPU
+    # configurations actually used at runtime -- typically 1-3 entries per
+    # worker process (one per distinct SPU device declaration).
+    _spu_lock: threading.Lock = field(default_factory=threading.Lock, repr=False)
+    _spu_template_links: dict[tuple, Any] = field(default_factory=dict, repr=False)
+
+    def get_or_create_spu_link(
+        self,
+        cache_key: tuple,
+        create_fn: Callable[[], Any],
+    ) -> Any:
+        """Thread-safe lazy creation of template SPU links.
+
+        Args:
+            cache_key: Tuple identifying the SPU configuration.
+            create_fn: Factory function to create a new link context.
+
+        Returns:
+            A libspu.link.Context (template link) for the given configuration.
+        """
+        with self._spu_lock:
+            if cache_key not in self._spu_template_links:
+                self._spu_template_links[cache_key] = create_fn()
+            return self._spu_template_links[cache_key]
+
+    def shutdown(self) -> None:
+        """Release cached SPU template links.
+
+        Safe to call multiple times.  Should be called during process
+        shutdown to eagerly close BRPC connections rather than waiting
+        for GC / process exit.
+        """
+        with self._spu_lock:
+            self._spu_template_links.clear()

mplang_nightly-0.1.dev332/mplang/backends/simp_worker/request.py

@@ -0,0 +1,80 @@
+# Copyright 2026 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.
+
+"""Per-request Interpreter factory.
+
+Creates lightweight Interpreter instances for each incoming request,
+providing full isolation of mutable state (CommContext, SimpWorker,
+SPUState) while sharing immutable infrastructure (Communicator, Store,
+handlers).
+"""
+
+from __future__ import annotations
+
+from mplang.backends.simp_worker.comm_context import CommContext
+from mplang.backends.simp_worker.infra import WorkerInfra
+from mplang.backends.simp_worker.state import SimpWorker
+from mplang.backends.spu_state import SPUState
+from mplang.runtime.interpreter import Interpreter
+
+
+def create_request_interpreter(
+    infra: WorkerInfra,
+    job_id: str,
+) -> Interpreter:
+    """Create a lightweight Interpreter for a single request.
+
+    Cost: ~2μs (dict/TLS allocation only).
+    SPU Runtime created on-demand via link.spawn() (~120μs).
+
+    Args:
+        infra: Shared WorkerInfra (process-lifetime).
+        job_id: Unique request identifier (used as CommContext context_id).
+
+    Returns:
+        Per-request Interpreter with isolated state.
+    """
+    # Per-request CommContext with unique context_id
+    comm_ctx = CommContext(infra.communicator, context_id=job_id, my_rank=infra.rank)
+
+    # Per-request Interpreter (does not own shared executor/tracer)
+    interp = Interpreter(
+        name=f"Worker-{infra.rank}-{job_id}",
+        tracer=infra.tracer,
+        trace_pid=infra.trace_pid,
+        store=infra.store,
+        root_dir=infra.root_dir,
+        handlers=infra.handlers,
+        executor=infra.executor,
+        comm_ctx=comm_ctx,
+        owns_executor=False,
+        owns_tracer=False,
+    )
+    interp.async_ops = set(infra.async_ops)
+
+    # Per-request SimpWorker: isolates current_parties
+    worker_state = SimpWorker(
+        rank=infra.rank,
+        world_size=infra.world_size,
+        communicator=infra.communicator,  # raw comm kept for SPU BaseChannel
+        store=infra.store,
+        spu_endpoints=infra.spu_endpoints,
+    )
+    interp.set_dialect_state("simp", worker_state)
+
+    # Per-request SPUState: will use link.spawn() for Runtime isolation
+    spu_state = SPUState(infra=infra)
+    interp.set_dialect_state("spu", spu_state)
+
+    return interp