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

mplang/v2/backends/simp_design.md

@@ -0,0 +1,136 @@
+# Simp Dialect Backend Design
+
+## Overview
+
+The `simp` (Simple Multi-Party) dialect implements SPMD (Single Program Multiple Data) distributed execution. A single program is written once and executed across multiple parties, with the runtime handling distribution, communication, and synchronization.
+
+## Why Two Implementations?
+
+The simp dialect requires **two separate backend implementations** because the same primitives (`pcall`, `shuffle`, `converge`) have fundamentally different semantics depending on where they execute:
+
+| Primitive | Driver (Host) | Worker |
+|-----------|---------------|--------|
+| `pcall` | Dispatch work to workers | Execute locally |
+| `shuffle` | Route data between workers | Send/Receive via communicator |
+| `converge` | Merge HostVars | Pick non-null value |
+
+This is the essence of SPMD: the Driver orchestrates, Workers execute.
+
+## Architecture
+
+```
+┌─────────────────────────────────────────────────────────────────┐
+│                      dialects/simp.py                          │
+│                    (Primitive definitions)                      │
+└─────────────────────────────────────────────────────────────────┘
+                              │
+                ┌─────────────┴─────────────┐
+                ▼                           ▼
+┌───────────────────────────┐   ┌───────────────────────────┐
+│     simp_driver/          │   │     simp_worker/          │
+│   (Host/Driver side)      │   │   (Worker side)           │
+├───────────────────────────┤   ├───────────────────────────┤
+│ base.py    SimpDriver     │   │ state.py   SimpWorker     │
+│ ops.py     HOST_HANDLERS  │   │ ops.py     WORKER_HANDLERS│
+│ values.py  HostVar        │   │                           │
+│ mem.py     SimpMemDriver  │   │ mem.py     LocalMesh      │
+│ http.py    SimpHttpDriver │   │ http.py    HTTP Server    │
+└───────────────────────────┘   └───────────────────────────┘
+```
+
+## Directory Structure
+
+```
+backends/
+├── simp_driver/                    # Driver/Host side
+│   ├── __init__.py                 # Exports
+│   ├── base.py                     # SimpDriver (abstract base)
+│   ├── ops.py                      # HOST_HANDLERS
+│   ├── values.py                   # HostVar
+│   ├── mem.py                      # MemCluster + SimpMemDriver + make_simulator
+│   └── http.py                     # SimpHttpDriver + make_driver
+│
+├── simp_worker/                    # Worker side
+│   ├── __init__.py                 # Exports
+│   ├── state.py                    # SimpWorker (DialectState)
+│   ├── ops.py                      # WORKER_HANDLERS
+│   ├── mem.py                      # LocalMesh + ThreadCommunicator
+│   └── http.py                     # HTTP Worker Server
+```
+
+## Key Classes
+
+### Driver Side
+
+```python
+class SimpDriver(DialectState, ABC):
+    """Abstract interface for simp drivers."""
+    dialect_name = "simp"
+    world_size: int
+    
+    @abstractmethod
+    def submit(self, rank, graph, inputs, job_id=None) -> Future: ...
+    @abstractmethod
+    def fetch(self, rank, uri) -> Future: ...
+    @abstractmethod
+    def collect(self, futures) -> list: ...
+
+class SimpMemDriver(SimpDriver):
+    """In-memory IPC via ThreadPoolExecutor."""
+
+class SimpHttpDriver(SimpDriver):
+    """HTTP IPC via httpx."""
+```
+
+### Worker Side
+
+```python
+class SimpWorker(DialectState):
+    """Worker state with communicator and store."""
+    dialect_name = "simp"
+    rank: int
+    world_size: int
+    communicator: Any  # ThreadCommunicator or HTTP client
+    store: ObjectStore
+```
+
+## IPC Symmetry
+
+| IPC Type | Driver | Worker |
+|----------|--------|--------|
+| Memory | `simp_driver/mem.py` | `simp_worker/mem.py` |
+| HTTP | `simp_driver/http.py` | `simp_worker/http.py` |
+
+## Factory Functions
+
+```python
+# Create local simulator (memory IPC)
+interp = simp.make_simulator(world_size=3)
+
+# Create remote driver (HTTP IPC)
+interp = simp.make_driver(["http://w1:8000", "http://w2:8000"])
+```
+
+## Data Flow
+
+```
+User Code
+    │
+    ▼
+simp.pcall(parties=(0,1), fn, args)
+    │
+    ▼ (Driver Interpreter)
+HOST_HANDLERS["simp.pcall"]
+    │
+    ├─► driver.submit(rank=0, graph, inputs)
+    └─► driver.submit(rank=1, graph, inputs)
+              │
+              ▼ (IPC: Memory or HTTP)
+         Worker Interpreters
+              │
+              ▼
+    WORKER_HANDLERS["simp.pcall"]
+              │
+              ▼
+         Local Execution
+```

mplang/v2/backends/simp_driver/__init__.py

@@ -0,0 +1,41 @@
+# 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.
+
+"""Simp Driver package.
+
+Provides Driver-side state, values, and ops for the simp dialect.
+"""
+
+from mplang.v2.backends.simp_driver.http import SimpHttpDriver, make_driver
+from mplang.v2.backends.simp_driver.mem import (
+    LocalCluster,
+    MemCluster,
+    SimpMemDriver,
+    make_simulator,
+)
+from mplang.v2.backends.simp_driver.ops import DRIVER_HANDLERS
+from mplang.v2.backends.simp_driver.state import SimpDriver
+from mplang.v2.backends.simp_driver.values import DriverVar
+
+__all__ = [
+    "DRIVER_HANDLERS",
+    "DriverVar",
+    "LocalCluster",
+    "MemCluster",
+    "SimpDriver",
+    "SimpHttpDriver",
+    "SimpMemDriver",
+    "make_driver",
+    "make_simulator",
+]

mplang/v2/backends/simp_driver/http.py

@@ -0,0 +1,168 @@
+# 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.
+
+"""Simp Driver HTTP IPC (SimpHttpDriver, make_driver)."""
+
+from __future__ import annotations
+
+import concurrent.futures
+import os
+import pathlib
+from typing import TYPE_CHECKING, Any
+
+import httpx
+
+from mplang.v2.backends.simp_driver.state import SimpDriver
+from mplang.v2.edsl import serde
+from mplang.v2.runtime.interpreter import Interpreter
+from mplang.v2.runtime.object_store import ObjectStore
+
+if TYPE_CHECKING:
+    from concurrent.futures import Future
+
+    from mplang.v2.edsl.graph import Graph
+    from mplang.v2.libs.device import ClusterSpec
+
+
+class SimpHttpDriver(SimpDriver):
+    """Simp Driver for remote HTTP IPC.
+
+    Implements submit/fetch/collect interface for dispatching work via HTTP.
+    """
+
+    dialect_name: str = "simp"
+
+    def __init__(
+        self,
+        endpoints: list[str],
+        *,
+        cluster_spec: ClusterSpec | None = None,
+    ) -> None:
+        """Create remote simp driver.
+
+        Args:
+            endpoints: List of HTTP endpoints for workers.
+            cluster_spec: Optional cluster specification for metadata.
+        """
+        # Normalize endpoints
+        self._endpoints = []
+        for ep in endpoints:
+            if not ep.startswith("http://") and not ep.startswith("https://"):
+                ep = f"http://{ep}"
+            self._endpoints.append(ep)
+
+        self._world_size = len(endpoints)
+        self._cluster_spec = cluster_spec
+
+        # Construct driver root
+        data_root = pathlib.Path(os.environ.get("MPLANG_DATA_ROOT", ".mpl"))
+        if cluster_spec:
+            self.driver_root = data_root / cluster_spec.cluster_id / "__driver__"
+        else:
+            self.driver_root = data_root / "__remote__" / "__driver__"
+
+        # HTTP client and executor
+        self._client = httpx.Client(timeout=None)
+        self._executor = concurrent.futures.ThreadPoolExecutor(
+            max_workers=self._world_size
+        )
+
+    @property
+    def world_size(self) -> int:
+        return self._world_size
+
+    def submit(
+        self, rank: int, graph: Graph, inputs: list[Any], job_id: str | None = None
+    ) -> Future[Any]:
+        """Submit execution to remote worker via HTTP."""
+        return self._executor.submit(self._do_request, rank, graph, inputs, job_id)
+
+    def collect(self, futures: list[Future[Any]]) -> list[Any]:
+        """Collect results from futures."""
+        return [f.result() for f in futures]
+
+    def fetch(self, rank: int, uri: str) -> Future[Any]:
+        """Fetch data from remote worker."""
+        return self._executor.submit(self._do_fetch, rank, uri)
+
+    def _do_request(
+        self, rank: int, graph: Graph, inputs: list[Any], job_id: str | None = None
+    ) -> Any:
+        """Execute HTTP request."""
+        url = f"{self._endpoints[rank]}/exec"
+        graph_b64 = serde.dumps_b64(graph)
+        inputs_b64 = serde.dumps_b64(inputs)
+        payload = {"graph": graph_b64, "inputs": inputs_b64}
+        if job_id:
+            payload["job_id"] = job_id
+
+        resp = self._client.post(url, json=payload)
+        resp.raise_for_status()
+        return serde.loads_b64(resp.json()["result"])
+
+    def _do_fetch(self, rank: int, uri: str) -> Any:
+        """Fetch data from remote worker."""
+        url = f"{self._endpoints[rank]}/fetch"
+        resp = self._client.post(url, json={"uri": uri})
+        resp.raise_for_status()
+        return serde.loads_b64(resp.json()["result"])
+
+    def shutdown(self) -> None:
+        """Close HTTP client and executor."""
+        self._client.close()
+        self._executor.shutdown()
+
+
+def make_driver(endpoints: list[str], *, cluster_spec: Any = None) -> Interpreter:
+    """Create an Interpreter configured for remote SIMP execution.
+
+    This factory creates a SimpHttpDriver and returns an Interpreter
+    with the simp dialect state attached.
+
+    Args:
+        endpoints: List of HTTP endpoints for workers.
+        cluster_spec: Optional ClusterSpec for metadata.
+
+    Returns:
+        Configured Interpreter with simp state attached.
+
+    Example:
+        >>> interp = make_driver(["http://worker1:8000", "http://worker2:8000"])
+        >>> with interp:
+        ...     result = my_func()
+    """
+    from mplang.v2.backends.simp_driver.ops import DRIVER_HANDLERS
+
+    if cluster_spec is None:
+        from mplang.v2.libs.device import ClusterSpec
+
+        cluster_spec = ClusterSpec.simple(
+            world_size=len(endpoints), endpoints=endpoints
+        )
+
+    state = SimpHttpDriver(endpoints, cluster_spec=cluster_spec)
+
+    from collections.abc import Callable
+
+    handlers: dict[str, Callable[..., Any]] = {**DRIVER_HANDLERS}  # type: ignore[dict-item]
+    interp = Interpreter(
+        name="DriverInterpreter",
+        root_dir=state.driver_root,
+        handlers=handlers,
+        store=ObjectStore(fs_root=str(state.driver_root)),
+    )
+    interp.set_dialect_state("simp", state)
+    interp._cluster_spec = cluster_spec  # type: ignore[attr-defined]
+
+    return interp

mplang/v2/backends/simp_driver/mem.py

@@ -0,0 +1,280 @@
+# 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.
+
+"""Simp Driver memory IPC (MemCluster, SimpMemDriver, make_simulator)."""
+
+from __future__ import annotations
+
+import concurrent.futures
+import os
+import pathlib
+from collections.abc import Callable
+from typing import TYPE_CHECKING, Any, cast
+
+from mplang.v2.backends.simp_driver.state import SimpDriver
+from mplang.v2.backends.simp_worker import WORKER_HANDLERS, SimpWorker
+from mplang.v2.backends.simp_worker.mem import LocalMesh
+from mplang.v2.runtime.interpreter import ExecutionTracer, Interpreter
+from mplang.v2.runtime.object_store import ObjectStore
+
+if TYPE_CHECKING:
+    from concurrent.futures import Future
+
+    from mplang.v2.edsl.graph import Graph
+    from mplang.v2.libs.device import ClusterSpec
+
+
+class MemCluster:
+    """Orchestrator that creates and manages local worker Interpreters.
+
+    This class handles worker lifecycle management. It does NOT attach to
+    an Interpreter - instead, it creates a SimpMemDriver that can be attached.
+    """
+
+    def __init__(
+        self,
+        world_size: int,
+        *,
+        cluster_spec: ClusterSpec | None = None,
+        enable_tracing: bool = False,
+    ) -> None:
+        """Create a local memory cluster.
+
+        Args:
+            world_size: Number of workers.
+            cluster_spec: Optional cluster specification for metadata.
+            enable_tracing: If True, enable execution tracing.
+        """
+        self._world_size = world_size
+        self._cluster_spec = cluster_spec
+
+        # Construct root_dir from cluster_id
+        data_root = pathlib.Path(os.environ.get("MPLANG_DATA_ROOT", ".mpl"))
+        cluster_id = cluster_spec.cluster_id if cluster_spec else f"local_{world_size}"
+        cluster_root = data_root / cluster_id
+        self.host_root = cluster_root / "__host__"
+
+        # Create Local Mesh (communication mesh for workers)
+        self._mesh = LocalMesh(world_size)
+
+        # Create Execution Tracer
+        self.tracer: ExecutionTracer = ExecutionTracer(
+            enabled=enable_tracing, trace_dir=self.host_root / "trace"
+        )
+        self.tracer.start()
+
+        # Create Workers
+        self._workers: list[Interpreter] = []
+        for rank in range(world_size):
+            worker_root = cluster_root / f"node{rank}"
+            store = ObjectStore(fs_root=str(worker_root / "store"))
+
+            worker_state = SimpWorker(
+                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]
+            w_interp = Interpreter(
+                name=f"Worker-{rank}",
+                tracer=self.tracer,
+                trace_pid=rank,
+                store=store,
+                root_dir=worker_root,
+                handlers=w_handlers,
+            )
+            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)
+
+    @property
+    def world_size(self) -> int:
+        return self._world_size
+
+    @property
+    def workers(self) -> list[Interpreter]:
+        return self._workers
+
+    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,
+            mesh=self._mesh,
+        )
+
+    def shutdown(self, wait: bool = True) -> None:
+        """Stop all workers and release resources."""
+        self._mesh.shutdown(wait=wait)
+
+
+class SimpMemDriver(SimpDriver):
+    """Simp Driver for local memory IPC.
+
+    Implements submit/fetch/collect interface for dispatching work to local workers.
+    This class is created by MemCluster and attached to a Driver Interpreter.
+    """
+
+    dialect_name: str = "simp"
+
+    def __init__(
+        self,
+        world_size: int,
+        workers: list[Interpreter],
+        mesh: Any,  # LocalMesh from simp_worker.mem
+    ) -> None:
+        self._world_size = world_size
+        self._workers = workers
+        self._mesh = mesh
+
+    def shutdown(self) -> None:
+        """Shutdown the local memory driver and its mesh."""
+        self._mesh.shutdown()
+
+    @property
+    def world_size(self) -> int:
+        return self._world_size
+
+    @property
+    def workers(self) -> list[Interpreter]:
+        """Worker interpreters (for backward compatibility)."""
+        return self._workers
+
+    def submit(
+        self, rank: int, graph: Graph, inputs: list[Any], job_id: str | None = None
+    ) -> Future[Any]:
+        """Submit execution to local worker thread."""
+        return cast(
+            "Future[Any]",
+            self._mesh.executor.submit(
+                self._run_worker, rank, graph, inputs, job_id=job_id
+            ),
+        )
+
+    def collect(self, futures: list[Future[Any]]) -> list[Any]:
+        """Wait for threads and collect results."""
+        done, _ = concurrent.futures.wait(
+            futures, return_when=concurrent.futures.FIRST_EXCEPTION
+        )
+        for f in done:
+            exc = f.exception()
+            if exc:
+                for nf in futures:
+                    nf.cancel()
+                self._mesh.shutdown(wait=False)
+                raise exc
+        return [f.result() for f in futures]
+
+    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]
+
+    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
+        ]
+
+
+def make_simulator(
+    world_size: int,
+    *,
+    cluster_spec: Any = None,
+    enable_tracing: bool = False,
+) -> Interpreter:
+    """Create an Interpreter configured for local SIMP simulation.
+
+    This factory creates a MemCluster with workers and returns an
+    Interpreter with the simp dialect state attached.
+
+    Args:
+        world_size: Number of simulated parties.
+        cluster_spec: Optional ClusterSpec for metadata.
+        enable_tracing: If True, enable execution tracing.
+
+    Returns:
+        Configured Interpreter with simp state attached.
+
+    Example:
+        >>> interp = make_simulator(2)
+        >>> with interp:
+        ...     result = my_func()
+    """
+    from mplang.v2.backends.simp_driver.ops import DRIVER_HANDLERS
+
+    if cluster_spec is None:
+        from mplang.v2.libs.device import ClusterSpec
+
+        cluster_spec = ClusterSpec.simple(world_size)
+
+    cluster = MemCluster(
+        world_size=world_size,
+        cluster_spec=cluster_spec,
+        enable_tracing=enable_tracing,
+    )
+    state = cluster.create_state()
+
+    handlers: dict[str, Callable[..., Any]] = {**DRIVER_HANDLERS}  # type: ignore[dict-item]
+    interp = Interpreter(
+        name="HostInterpreter",
+        root_dir=cluster.host_root,
+        handlers=handlers,
+        tracer=cluster.tracer,
+        store=ObjectStore(fs_root=str(cluster.host_root)),
+    )
+    interp.set_dialect_state("simp", state)
+
+    # Keep cluster alive (prevent GC)
+    interp._simp_cluster = cluster  # type: ignore[attr-defined]
+    interp._cluster_spec = cluster_spec  # type: ignore[attr-defined]
+
+    return interp
+
+
+# Backward compatibility alias
+LocalCluster = MemCluster