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

mplang/v2/backends/simp_driver/ops.py

@@ -0,0 +1,135 @@
+# 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 ops (DRIVER_HANDLERS).
+
+Unified SPMD dispatch pattern for all SIMP operations.
+All ops: wrap → dispatch to ALL workers → collect DriverVar(s).
+Op-specific logic lives in Worker handlers (simp_worker/ops.py).
+"""
+
+from __future__ import annotations
+
+from typing import Any
+
+from mplang.v2.backends.simp_driver.values import DriverVar
+from mplang.v2.dialects import simp
+from mplang.v2.edsl.graph import Graph, Operation
+from mplang.v2.edsl.typing import CustomType
+
+
+def _get_driver_context(interpreter: Any) -> Any:
+    """Get the simp driver state from interpreter."""
+    state = interpreter.get_dialect_state("simp")
+    if state is None:
+        raise RuntimeError("Interpreter must have simp dialect state attached")
+    return state
+
+
+def _wrap_op_as_graph(op: Operation) -> Graph:
+    """Wrap an Operation into a single-op Graph for worker submission."""
+    g = Graph()
+    any_type = CustomType("Any")
+
+    # Create graph inputs
+    graph_inputs = [g.add_input(f"in_{i}", any_type) for i in range(len(op.inputs))]
+
+    # Determine output types
+    output_types = [out.type for out in op.outputs] if op.outputs else [any_type]
+
+    # Add the operation (this handles outputs and value registration)
+    g.add_op(
+        opcode=op.opcode,
+        inputs=graph_inputs,
+        output_types=output_types,
+        attrs=op.attrs.copy(),
+        regions=op.regions,
+    )
+
+    # Mark outputs
+    for v in g.operations[-1].outputs:
+        g.add_output(v)
+
+    return g
+
+
+def _collect_to_hostvars(results: list[Any], num_outputs: int, world_size: int) -> Any:
+    """Collect worker results into DriverVar(s).
+
+    Args:
+        results: List of results from each worker (length = world_size).
+                 Each result is a list of URIs (one per output).
+        num_outputs: Number of outputs per worker
+        world_size: Total number of workers
+
+    Returns:
+        Single DriverVar if num_outputs == 1, else list of DriverVars
+    """
+    if num_outputs == 0:
+        return None
+
+    # Transpose [worker][output] -> [output][worker]
+    # results[worker_idx] is a list of URIs for that worker's outputs
+    transposed = []
+    for out_idx in range(num_outputs):
+        transposed.append(
+            DriverVar([res[out_idx] if res is not None else None for res in results])
+        )
+
+    if num_outputs == 1:
+        return transposed[0]
+    return transposed
+
+
+def _generic_simp_dispatch(interpreter: Any, op: Operation, *args: Any) -> Any:
+    """Unified SIMP dispatch: wrap op, SPMD submit, collect DriverVar(s).
+
+    This is the ONLY driver handler needed for all SIMP ops.
+    Worker handlers implement the actual op-specific logic.
+    """
+    driver = _get_driver_context(interpreter)
+    world_size = driver.world_size
+
+    # 1. Wrap operation into a Graph
+    wrapper_graph = _wrap_op_as_graph(op)
+
+    # 2. SPMD dispatch to ALL workers
+    futures = []
+    for rank in range(world_size):
+        # Extract per-party inputs from DriverVars
+        party_inputs = [
+            arg[rank] if isinstance(arg, DriverVar) else arg for arg in args
+        ]
+        futures.append(driver.submit(rank, wrapper_graph, party_inputs))
+
+    # 3. Collect results
+    results = driver.collect(futures)
+
+    # 4. Assemble into DriverVar(s)
+    num_outputs = len(op.outputs) if op.outputs else 1
+    return _collect_to_hostvars(results, num_outputs, world_size)
+
+
+# =============================================================================
+# All SIMP ops use unified dispatch
+# =============================================================================
+
+DRIVER_HANDLERS = {
+    simp.pcall_static_p.name: _generic_simp_dispatch,
+    simp.pcall_dynamic_p.name: _generic_simp_dispatch,
+    simp.shuffle_static_p.name: _generic_simp_dispatch,
+    simp.converge_p.name: _generic_simp_dispatch,
+    simp.uniform_cond_p.name: _generic_simp_dispatch,
+    simp.while_loop_p.name: _generic_simp_dispatch,
+}

mplang/v2/backends/simp_driver/state.py

@@ -0,0 +1,60 @@
+# 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.
+
+"""SimpDriver abstract base class."""
+
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+from typing import TYPE_CHECKING, Any
+
+from mplang.v2.runtime.dialect_state import DialectState
+
+if TYPE_CHECKING:
+    from concurrent.futures import Future
+
+    from mplang.v2.edsl.graph import Graph
+
+
+class SimpDriver(DialectState, ABC):
+    """Abstract base class for Simp Host drivers.
+
+    All simp drivers must implement submit/fetch/collect interface
+    for dispatching work to workers.
+    """
+
+    dialect_name: str = "simp"
+
+    @property
+    @abstractmethod
+    def world_size(self) -> int:
+        """Number of workers."""
+        ...
+
+    @abstractmethod
+    def submit(
+        self, rank: int, graph: Graph, inputs: list[Any], job_id: str | None = None
+    ) -> Future[Any]:
+        """Submit graph execution to a worker."""
+        ...
+
+    @abstractmethod
+    def fetch(self, rank: int, uri: str) -> Future[Any]:
+        """Fetch data from a worker."""
+        ...
+
+    @abstractmethod
+    def collect(self, futures: list[Future[Any]]) -> list[Any]:
+        """Collect results from futures."""
+        ...

mplang/v2/backends/simp_driver/values.py

@@ -0,0 +1,52 @@
+# 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 values (DriverVar)."""
+
+from __future__ import annotations
+
+from typing import Any, ClassVar, Self
+
+from mplang.v2.edsl import serde
+from mplang.v2.runtime.value import Value
+
+
+@serde.register_class
+class DriverVar(Value):
+    """A value replicated (or sharded) on the Driver.
+
+    A DriverVar holds a list of values, one for each party in the computation.
+    """
+
+    _serde_kind: ClassVar[str] = "simp.DriverVar"
+
+    def __init__(self, values: list[Any]):
+        self.values = values
+
+    @property
+    def world_size(self) -> int:
+        return len(self.values)
+
+    def __repr__(self) -> str:
+        return f"DriverVar({self.values})"
+
+    def __getitem__(self, idx: int) -> Any:
+        return self.values[idx]
+
+    def to_json(self) -> dict[str, Any]:
+        return {"values": self.values}
+
+    @classmethod
+    def from_json(cls, data: dict[str, Any]) -> Self:
+        return cls(values=data["values"])

mplang/v2/backends/simp_worker/__init__.py

@@ -0,0 +1,29 @@
+# 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 Worker package.
+
+Provides Worker-side state and ops for the simp dialect.
+"""
+
+from mplang.v2.backends.simp_worker.mem import LocalMesh, ThreadCommunicator
+from mplang.v2.backends.simp_worker.ops import WORKER_HANDLERS
+from mplang.v2.backends.simp_worker.state import SimpWorker
+
+__all__ = [
+    "WORKER_HANDLERS",
+    "LocalMesh",
+    "SimpWorker",
+    "ThreadCommunicator",
+]

mplang/v2/backends/simp_worker/http.py

@@ -0,0 +1,354 @@
+# 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 HTTP Worker module.
+
+Provides the HTTP-based worker entry point for distributed deployment.
+This module contains:
+- HttpCommunicator: HTTP-based inter-worker communication
+- create_worker_app: Factory for FastAPI application
+
+Usage:
+    # Start a worker server
+    from mplang.v2.backends.simp_http_worker import create_worker_app
+    import uvicorn
+
+    app = create_worker_app(rank=0, world_size=3, endpoints=[...])
+    uvicorn.run(app, host="0.0.0.0", port=8000)
+
+Security:
+    This module uses secure JSON-based serialization (serde) for computation
+    graphs and data between workers. Unlike pickle, JSON deserialization
+    cannot execute arbitrary code, making it safe for network communication.
+"""
+
+from __future__ import annotations
+
+import concurrent.futures
+import logging
+import os
+import pathlib
+import threading
+import time
+from typing import Any
+
+import httpx
+from fastapi import FastAPI, HTTPException
+from pydantic import BaseModel
+
+from mplang.v2.backends import spu_impl as _spu_impl  # noqa: F401
+from mplang.v2.backends import tensor_impl as _tensor_impl  # noqa: F401
+
+# Register operation implementations (side-effect imports)
+from mplang.v2.backends.simp_worker import SimpWorker
+from mplang.v2.backends.simp_worker import ops as _simp_worker_ops  # noqa: F401
+from mplang.v2.edsl import serde
+from mplang.v2.edsl.graph import Graph
+from mplang.v2.runtime.interpreter import ExecutionTracer, Interpreter
+from mplang.v2.runtime.object_store import ObjectStore
+
+logger = logging.getLogger(__name__)
+
+
+class HttpCommunicator:
+    """Communicator using HTTP requests for inter-worker communication.
+
+    Uses a background thread pool for sending to avoid blocking the main execution.
+
+    Attributes:
+        rank: This worker's rank
+        world_size: Total number of workers
+        endpoints: HTTP endpoints for all workers
+    """
+
+    def __init__(
+        self,
+        rank: int,
+        world_size: int,
+        endpoints: list[str],
+        tracer: ExecutionTracer | None = None,
+    ):
+        self.rank = rank
+        self.world_size = world_size
+        self.endpoints = endpoints
+        self.tracer = tracer
+        self._mailbox: dict[tuple[int, str], Any] = {}
+        self._cond = threading.Condition()
+        self._send_executor = concurrent.futures.ThreadPoolExecutor(
+            max_workers=world_size, thread_name_prefix=f"comm_send_{rank}"
+        )
+        self._pending_sends: list[concurrent.futures.Future[None]] = []
+        self.client = httpx.Client(timeout=None)
+
+    def send(self, to: int, key: str, data: Any) -> None:
+        """Send data to another rank asynchronously."""
+        future = self._send_executor.submit(self._do_send, to, key, data)
+        self._pending_sends.append(future)
+
+    def _do_send(self, to: int, key: str, data: Any) -> None:
+        """Perform the HTTP send."""
+        url = f"{self.endpoints[to]}/comm/{key}"
+        logger.debug(f"Rank {self.rank} sending to {to} key={key}")
+
+        # Detect SPU channel (tag prefix "spu:") and handle bytes
+        if key.startswith("spu:") and isinstance(data, bytes):
+            # Send raw bytes for SPU channels
+            import base64
+
+            payload = base64.b64encode(data).decode("ascii")
+            is_raw_bytes = True
+        else:
+            # Use secure JSON serialization
+            payload = serde.dumps_b64(data)
+            is_raw_bytes = False
+
+        size_bytes = len(payload)
+
+        # Log to profiler
+        if self.tracer:
+            self.tracer.log_custom_event(
+                name="comm.send",
+                start_ts=time.time(),
+                end_ts=time.time(),  # Instant event for size? Or measure duration?
+                cat="comm",
+                args={"to": to, "key": key, "bytes": size_bytes},
+            )
+
+        try:
+            t0 = time.time()
+            resp = self.client.put(
+                url,
+                json={
+                    "data": payload,
+                    "from_rank": self.rank,
+                    "is_raw_bytes": is_raw_bytes,
+                },
+            )
+            resp.raise_for_status()
+            duration = time.time() - t0
+            if self.tracer:
+                self.tracer.log_custom_event(
+                    name="comm.send_req",
+                    start_ts=t0,
+                    end_ts=t0 + duration,
+                    cat="comm",
+                    args={"to": to, "key": key, "bytes": size_bytes},
+                )
+        except Exception as e:
+            logger.error(f"Rank {self.rank} failed to send to {to}: {e}")
+            raise RuntimeError(f"Failed to send to {to} ({url}): {e}") from e
+
+    def recv(self, frm: int, key: str) -> Any:
+        """Receive data from another rank (blocking)."""
+        logger.debug(f"Rank {self.rank} waiting recv from {frm} key={key}")
+        mailbox_key = (frm, key)
+        with self._cond:
+            while mailbox_key not in self._mailbox:
+                self._cond.wait(timeout=1.0)
+            return self._mailbox.pop(mailbox_key)
+
+    def on_receive(self, from_rank: int, key: str, data: Any) -> None:
+        """Called when data is received from the HTTP endpoint."""
+        mailbox_key = (from_rank, key)
+        with self._cond:
+            if mailbox_key in self._mailbox:
+                raise RuntimeError(
+                    f"Mailbox overflow: key {mailbox_key} already exists"
+                )
+            self._mailbox[mailbox_key] = data
+            self._cond.notify_all()
+
+    def wait_pending_sends(self) -> None:
+        """Wait for all pending sends to complete."""
+        for future in self._pending_sends:
+            try:
+                future.result(timeout=60.0)
+            except Exception as e:
+                logger.error(f"Rank {self.rank} send failed: {e}")
+        self._pending_sends.clear()
+
+    def shutdown(self) -> None:
+        """Shutdown the send executor."""
+        self.wait_pending_sends()
+        self._send_executor.shutdown(wait=True)
+        self.client.close()
+
+
+class ExecRequest(BaseModel):
+    """Request model for /exec endpoint."""
+
+    graph: str
+    inputs: str
+    job_id: str | None = None
+
+
+class CommRequest(BaseModel):
+    """Request model for /comm endpoint."""
+
+    data: str
+    from_rank: int
+    is_raw_bytes: bool = False  # NEW: indicates raw bytes (not serde)
+
+
+class FetchRequest(BaseModel):
+    """Request model for /fetch endpoint."""
+
+    uri: str
+
+
+def create_worker_app(
+    rank: int,
+    world_size: int,
+    endpoints: list[str],
+    spu_endpoints: dict[int, str] | None = None,
+) -> FastAPI:
+    """Create a FastAPI app for the worker.
+
+    The app uses async endpoints with run_in_executor to allow concurrent
+    handling of /exec and /comm requests. This is essential for cross-party
+    communication where one party sends while another receives.
+
+    Args:
+        rank: The global rank of this worker.
+        world_size: Total number of workers.
+        endpoints: HTTP endpoints for all workers (for shuffle communication).
+        spu_endpoints: Optional dict mapping global_rank -> BRPC endpoint for SPU.
+
+    Returns:
+        FastAPI application instance
+    """
+    import asyncio
+
+    app = FastAPI(title=f"SIMP Worker {rank}")
+
+    # Persistence root: ${MPLANG_DATA_ROOT}/<cluster_id>/node<rank>/
+    data_root = pathlib.Path(os.environ.get("MPLANG_DATA_ROOT", ".mpl"))
+    cluster_id = os.environ.get("MPLANG_CLUSTER_ID", f"__http_{world_size}")
+    root_dir = data_root / cluster_id / f"node{rank}"
+    trace_dir = root_dir / "trace"
+
+    # Enable tracing by default for now (or make it configurable via env)
+    tracer = ExecutionTracer(enabled=True, trace_dir=trace_dir)
+    tracer.start()
+
+    comm = HttpCommunicator(rank, world_size, endpoints, tracer=tracer)
+    store = ObjectStore(fs_root=str(root_dir))
+    ctx = SimpWorker(rank, world_size, comm, store, spu_endpoints)
+
+    # Register handlers
+    from collections.abc import Callable
+    from typing import cast
+
+    from mplang.v2.backends.simp_worker.ops import WORKER_HANDLERS
+
+    # func_impl is already imported at module level for side-effects
+    handlers: dict[str, Callable[..., Any]] = {**WORKER_HANDLERS}  # type: ignore[dict-item]
+
+    worker = Interpreter(
+        tracer=tracer, root_dir=root_dir, handlers=handlers, store=store
+    )
+    # 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}"
+    )
+
+    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
+        ]
+
+        result = worker.evaluate_graph(graph, resolved_inputs)
+        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]
+
+    @app.post("/exec")
+    async def execute(req: ExecRequest) -> dict[str, str]:
+        """Execute a graph on this worker."""
+        logger.debug(f"Worker {rank} received exec request")
+        try:
+            # Use secure JSON deserialization
+            graph = serde.loads_b64(req.graph)
+            inputs = serde.loads_b64(req.inputs)
+            loop = asyncio.get_event_loop()
+            result = await loop.run_in_executor(
+                exec_pool, _do_execute, graph, inputs, req.job_id
+            )
+            return {"result": serde.dumps_b64(result)}
+        except Exception as e:
+            logger.error(f"Worker {rank} exec failed: {e}")
+            raise HTTPException(status_code=500, detail=str(e)) from e
+
+    @app.put("/comm/{key}")
+    async def receive_comm(key: str, req: CommRequest) -> dict[str, str]:
+        """Receive communication data from another worker."""
+        logger.debug(f"Worker {rank} received comm key={key} from {req.from_rank}")
+        try:
+            # Handle raw bytes (SPU channels) vs serde data
+            if req.is_raw_bytes:
+                # Decode base64 to raw bytes
+                import base64
+
+                data = base64.b64decode(req.data)
+            else:
+                # Use secure JSON deserialization
+                data = serde.loads_b64(req.data)
+
+            comm.on_receive(req.from_rank, key, data)
+            return {"status": "ok"}
+        except Exception as e:
+            logger.error(f"Worker {rank} comm failed: {e}")
+            raise HTTPException(status_code=500, detail=str(e)) from e
+
+    @app.post("/fetch")
+    async def fetch(req: FetchRequest) -> dict[str, str]:
+        """Fetch data by URI."""
+        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)
+            return {"result": serde.dumps_b64(val)}
+        except Exception as e:
+            logger.error(f"Worker {rank} fetch failed: {e}")
+            raise HTTPException(status_code=500, detail=str(e)) from e
+
+    @app.get("/objects")
+    async def list_objects() -> dict[str, list[str]]:
+        """List all objects in the worker's store."""
+        try:
+            state = cast(SimpWorker, worker.get_dialect_state("simp"))
+            return {"objects": state.store.list_objects()}
+        except Exception as e:
+            logger.error(f"Worker {rank} list_objects failed: {e}")
+            raise HTTPException(status_code=500, detail=str(e)) from e
+
+    @app.get("/health")
+    async def health() -> dict[str, str]:
+        """Health check endpoint."""
+        return {"status": "ok", "rank": str(rank), "world_size": str(world_size)}
+
+    @app.on_event("shutdown")
+    def shutdown_event() -> None:
+        """Cleanup on shutdown."""
+        comm.shutdown()
+        exec_pool.shutdown(wait=True)
+
+    return app