mplang-nightly 0.1.dev269__py3-none-any.whl → 0.1.dev271__py3-none-any.whl

mplang/v1/runtime/session.py

@@ -1,270 +0,0 @@
-# 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.
-
-"""Core Session model (pure, no global registries).
-
-Contents:
-    * SessionState dataclass
-    * LinkCommFactory (SPU link reuse cache)
-    * Session (topology derivation, runtime init, SPU env seeding, local symbol/computation storage)
-
-Process-wide registries (sessions, global symbols) live in the server layer
-(`server.py`) so this module remains portable and easy to unit test.
-"""
-
-from __future__ import annotations
-
-import time
-from dataclasses import dataclass, field
-from functools import cached_property
-from typing import TYPE_CHECKING, Any, cast
-
-import spu.libspu as libspu
-
-from mplang.v1.core.cluster import ClusterSpec
-from mplang.v1.core.comm import ICommunicator
-from mplang.v1.core.expr.ast import Expr
-from mplang.v1.core.expr.evaluator import IEvaluator, create_evaluator
-from mplang.v1.core.mask import Mask
-from mplang.v1.kernels.context import RuntimeContext
-from mplang.v1.kernels.spu import PFunction  # type: ignore
-from mplang.v1.kernels.value import Value
-from mplang.v1.runtime.communicator import HttpCommunicator
-from mplang.v1.runtime.exceptions import ResourceNotFound
-from mplang.v1.runtime.link_comm import LinkCommunicator
-from mplang.v1.utils.spu_utils import parse_field, parse_protocol
-
-if TYPE_CHECKING:  # pragma: no cover - import only for type checking
-    from mplang.v1.core.cluster import ClusterSpec, Node, RuntimeInfo
-
-
-@dataclass
-class Symbol:
-    name: str
-    mptype: Any
-    data: Any
-
-
-@dataclass
-class Computation:
-    name: str
-    expr: Expr
-
-
-@dataclass
-class SessionState:
-    runtime: RuntimeContext | None = None
-    computations: dict[str, Computation] = field(default_factory=dict)
-    symbols: dict[str, Symbol] = field(default_factory=dict)
-    spu_seeded: bool = False
-    created_ts: float = field(default_factory=time.time)
-    last_access_ts: float = field(default_factory=time.time)
-
-
-class Session:
-    """Represents the per-rank execution context.
-
-    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,
-        communicator: ICommunicator,
-    ):
-        self.name = name
-        self.rank = rank
-        self.cluster_spec = cluster_spec
-        self.state = SessionState()
-        self.communicator = communicator
-
-    # --- Derived topology ---
-    @cached_property
-    def node(self) -> Node:
-        return self.cluster_spec.get_node_by_rank(self.rank)
-
-    @property
-    def runtime_info(self) -> RuntimeInfo:
-        return self.node.runtime_info
-
-    @property
-    def endpoints(self) -> list[str]:
-        return self.cluster_spec.endpoints
-
-    @cached_property
-    def spu_device(self):  # type: ignore
-        devs = self.cluster_spec.get_devices_by_kind("SPU")
-        if len(devs) != 1:
-            raise RuntimeError(
-                f"Expected exactly one SPU device, got {len(devs)} (session={self.name})"
-            )
-        return devs[0]
-
-    @cached_property
-    def spu_mask(self) -> Mask:
-        return Mask.from_ranks([m.rank for m in self.spu_device.members])
-
-    @property
-    def spu_protocol(self) -> str:
-        return cast(str, self.spu_device.config.get("protocol", "SEMI2K"))
-
-    @property
-    def spu_field(self) -> str:
-        return cast(str, self.spu_device.config.get("field", "FM64"))
-
-    @property
-    def is_spu_party(self) -> bool:
-        return self.rank in self.spu_mask
-
-    # --- Runtime helpers ---
-    def ensure_runtime(self) -> RuntimeContext:
-        if self.state.runtime is None:
-            self.state.runtime = RuntimeContext(
-                rank=self.rank,
-                world_size=len(self.cluster_spec.nodes),  # type: ignore[attr-defined]
-                initial_bindings=(
-                    self.runtime_info.op_bindings if self.runtime_info else {}
-                ),
-            )
-        return self.state.runtime
-
-    def ensure_spu_env(self) -> None:
-        """Ensure SPU kernel env (config/world[/link]) registered on this runtime.
-
-        Previous logic only seeded SPU parties; non-participating ranks then raised
-        a hard error when the evaluator encountered SPU ops in the global program,
-        because the kernel pocket lacked config/world. For now we register the
-        config/world on ALL parties (idempotent) and only attach a link context for
-        participating SPU ranks. Non-parties will still error later if they try to
-        execute a link-dependent SPU kernel (which should be guarded by masks in the
-        IR), but they will no longer fail early with a misleading
-        "SPU kernel state not initialized" message.
-        """
-        if self.state.spu_seeded:
-            return
-
-        link_ctx = None
-
-        if self.is_spu_party:
-            # Use Channels mode to reuse existing HttpCommunicator
-            # This eliminates the need for separate BRPC ports (SPU_PORT_OFFSET)
-            from mplang.v1.core.comm import CommunicatorBase
-
-            # Type assertion: ICommunicator is actually CommunicatorBase
-            comm = cast(CommunicatorBase, self.communicator)
-            link_ctx = LinkCommunicator(
-                rank=self.rank,
-                comm=comm,
-                spu_mask=self.spu_mask,
-            )
-
-        spu_config = libspu.RuntimeConfig(
-            protocol=parse_protocol(self.spu_protocol),
-            field=parse_field(self.spu_field),
-            fxp_fraction_bits=18,
-        )
-        seed_pfunc = PFunction(
-            fn_type="spu.seed_env",
-            ins_info=(),
-            outs_info=(),
-            config=spu_config,
-            world=self.spu_mask.num_parties(),
-            link=link_ctx,
-        )
-        self.ensure_runtime().run_kernel(seed_pfunc, [])
-        self.state.spu_seeded = True
-
-    # --- Computations & Symbols (instance-local) ---
-    def add_computation(self, computation: Computation) -> None:
-        self.state.computations[computation.name] = computation
-
-    def get_computation(self, name: str) -> Computation | None:
-        return self.state.computations.get(name)
-
-    def add_symbol(self, symbol: Symbol) -> None:
-        self.state.symbols[symbol.name] = symbol
-
-    def get_symbol(self, name: str) -> Symbol | None:
-        return self.state.symbols.get(name)
-
-    def list_symbols(self) -> list[str]:  # pragma: no cover - trivial
-        return list(self.state.symbols.keys())
-
-    def delete_symbol(self, name: str) -> bool:
-        if name in self.state.symbols:
-            del self.state.symbols[name]
-            return True
-        return False
-
-    def list_computations(self) -> list[str]:  # pragma: no cover - trivial
-        return list(self.state.computations.keys())
-
-    def delete_computation(self, name: str) -> bool:
-        if name in self.state.computations:
-            del self.state.computations[name]
-            return True
-        return False
-
-    # --- Execution ---
-    def execute(
-        self, computation: Computation, input_names: list[str], output_names: list[str]
-    ) -> None:
-        env: dict[str, Any] = {}
-        for in_name in input_names:
-            sym = self.get_symbol(in_name)
-            if sym is None:
-                raise ResourceNotFound(
-                    f"Input symbol '{in_name}' not found in session '{self.name}'"
-                )
-            env[in_name] = sym.data
-        rt = self.ensure_runtime()
-        self.ensure_spu_env()
-        evaluator: IEvaluator = create_evaluator(
-            rank=self.rank, env=env, comm=self.communicator, runtime=rt
-        )
-        results = evaluator.evaluate(computation.expr)
-        if results and len(results) != len(output_names):
-            raise RuntimeError(
-                f"Expected {len(output_names)} results, got {len(results)}"
-            )
-        for name, val in zip(output_names, results, strict=True):
-            # In pure SIMP model, all nodes should have the same symbol table.
-            # Non-participating nodes get None values.
-            if val is not None and not isinstance(val, Value):
-                raise TypeError(
-                    "Session executions must produce kernel Value outputs; "
-                    f"got {type(val).__name__} for symbol '{name}'"
-                )
-            self.add_symbol(Symbol(name=name, mptype={}, data=val))
-
-
-# --- 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/v1/runtime/simulation.py

@@ -1,324 +0,0 @@
-# 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.
-
-from __future__ import annotations
-
-import concurrent.futures
-import faulthandler
-import logging
-import sys
-import threading
-import traceback
-from collections.abc import Sequence
-from typing import Any, cast
-
-import spu.libspu as libspu
-
-from mplang.v1.core import (
-    ClusterSpec,
-    CollectiveMixin,
-    CommunicatorBase,
-    InterpContext,
-    InterpVar,
-    IrReader,
-    IrWriter,
-    Mask,
-    MPObject,
-    MPType,
-    PFunction,  # for spu.seed_env kernel seeding
-    TensorLike,
-)
-from mplang.v1.core.expr.ast import Expr
-from mplang.v1.core.expr.evaluator import IEvaluator, create_evaluator
-from mplang.v1.kernels.context import RuntimeContext
-from mplang.v1.runtime.link_comm import LinkCommunicator
-from mplang.v1.utils.spu_utils import parse_field, parse_protocol
-
-
-class ThreadCommunicator(CommunicatorBase, CollectiveMixin):
-    """Thread-based communicator for in-memory communication between threads"""
-
-    def __init__(self, rank: int, world_size: int):
-        super().__init__(rank, world_size)
-        self.peers: list[ThreadCommunicator] = []
-        logging.debug(
-            f"ThreadCommunicator initialized with rank={self.rank}, world_size={self.world_size}"
-        )
-
-    def set_peers(self, peers: list[ThreadCommunicator]) -> None:
-        assert self.world_size == len(peers)
-        self.peers = peers
-
-    def send(self, to: int, key: str, data: Any) -> None:
-        assert 0 <= to < self.world_size
-        # print(f"send {key}: {self.rank} -> {to_rank}")
-        self.peers[to].onSent(self.rank, key, data)
-
-
-class SimVar(InterpVar):
-    """A variable that references a value in an interpreter.
-
-    SimVar represents a value that has been computed and exists
-    in the interpreter's variable store.
-    """
-
-    def __init__(self, ctx: Simulator, mptype: MPType, values: list[Any]):
-        # Initialize the parent InterpVar with a generated name
-        super().__init__(ctx, mptype)
-        self._values = values
-
-    @property
-    def values(self) -> list[Any]:
-        """Converted values across all ranks for user inspection."""
-        return [v.to_numpy() if hasattr(v, "to_numpy") else v for v in self._values]
-
-    def __repr__(self) -> str:
-        return f"SimVar({self.mptype})"
-
-
-class Simulator(InterpContext):
-    def __init__(
-        self,
-        cluster_spec: ClusterSpec,
-        *,
-        trace_ranks: list[int] | None = None,
-    ) -> None:
-        """Initialize a simulator with the given cluster specification.
-
-        Args:
-            cluster_spec: The cluster specification defining the simulation environment.
-            trace_ranks: List of ranks to trace execution for debugging.
-                Per-node op binding overrides should now be provided via
-                each node's `runtime_info.op_bindings` in the supplied
-                `cluster_spec`.
-        """
-        super().__init__(cluster_spec)
-        self._trace_ranks = trace_ranks or []
-
-        spu_devices = cluster_spec.get_devices_by_kind("SPU")
-        if not spu_devices:
-            raise ValueError("No SPU device found in the cluster specification")
-        if len(spu_devices) > 1:
-            raise ValueError("Multiple SPU devices found in the cluster specification")
-        spu_device = spu_devices[0]
-
-        # compute spu_mask from spu_device members
-        spu_mask = Mask.from_ranks([member.rank for member in spu_device.members])
-
-        # Convert protocol and field from config using utility functions
-        spu_protocol = parse_protocol(spu_device.config["protocol"])
-        spu_field = parse_field(spu_device.config["field"])
-
-        world_size = self.world_size()
-
-        # Setup communicators
-        self._comms = [
-            ThreadCommunicator(rank, world_size) for rank in range(world_size)
-        ]
-        for comm in self._comms:
-            comm.set_peers(self._comms)
-
-        # Prepare link contexts for SPU parties (store for evaluator-time initialization)
-        # Use Channels mode to reuse ThreadCommunicator instead of separate mem_link
-        self._spu_link_ctxs: list[LinkCommunicator | None] = [None] * world_size
-
-        # Create LinkCommunicators in parallel to avoid deadlock
-        # (create_with_channels does handshake via TestSend/TestRecv)
-        exceptions: dict[int, Exception] = {}
-
-        def create_link(g_rank: int) -> None:
-            try:
-                self._spu_link_ctxs[g_rank] = LinkCommunicator(
-                    rank=g_rank,
-                    comm=self._comms[g_rank],
-                    spu_mask=spu_mask,
-                )
-            except Exception as e:
-                exceptions[g_rank] = e
-
-        threads = [
-            threading.Thread(target=create_link, args=(g_rank,)) for g_rank in spu_mask
-        ]
-        for t in threads:
-            t.start()
-        for t in threads:
-            t.join()
-
-        # Check for exceptions during link creation
-        if exceptions:
-            first_exc = next(iter(exceptions.values()))
-            raise RuntimeError(
-                f"Failed to create SPU link contexts for ranks {list(exceptions.keys())}"
-            ) from first_exc
-
-        self._spu_runtime_cfg = libspu.RuntimeConfig(
-            protocol=spu_protocol, field=spu_field
-        )
-        self._spu_world = spu_mask.num_parties()
-        self._spu_mask = spu_mask
-
-        # Persistent per-rank RuntimeContext instances (reused across evaluates).
-        # We no longer pre-create evaluators since each evaluate has different env bindings.
-        # Build per-rank runtime contexts.
-        self._runtimes: list[RuntimeContext] = []
-        for rank in range(self.world_size()):
-            node = self.cluster_spec.get_node_by_rank(rank)
-            rt = RuntimeContext(
-                rank=rank,
-                world_size=self.world_size(),
-                initial_bindings=node.runtime_info.op_bindings,
-            )
-            self._runtimes.append(rt)
-
-    @classmethod
-    def simple(
-        cls,
-        world_size: int,
-        op_bindings: dict[str, str] | None = None,
-        **kwargs: Any,
-    ) -> Simulator:
-        """Create a simple simulator with the given number of parties.
-
-        This is a convenience method that creates a ClusterSpec.simple()
-        configuration for quick testing and prototyping.
-
-        Args:
-            world_size: Number of simulated parties.
-            **kwargs: Additional arguments passed to the Simulator constructor.
-
-        Returns:
-            A Simulator instance with a simple cluster configuration.
-        """
-        cluster_spec = ClusterSpec.simple(world_size)
-        if op_bindings:
-            # Apply the same op_bindings to every node's runtime_info for convenience
-            for node in cluster_spec.nodes.values():
-                node.runtime_info.op_bindings.update(op_bindings)
-        return cls(cluster_spec, **kwargs)
-
-    def _do_evaluate(self, expr: Expr, evaluator_engine: IEvaluator) -> Any:
-        """
-        Helper function to simulate real-world MPIR serialization/deserialization
-        process instead of direct expr.accept execution.
-
-        This exposes potential MPIR serialization bugs by forcing expressions
-        to go through the full serialize->deserialize cycle.
-        """
-        writer = IrWriter()
-        graph_proto = writer.dumps(expr)
-
-        reader = IrReader()
-        deserialized_expr = reader.loads(graph_proto)
-
-        if deserialized_expr is None:
-            raise ValueError("Failed to deserialize expression")
-
-        return evaluator_engine.evaluate(deserialized_expr)
-
-    # override
-    def fetch(self, obj: MPObject) -> list[TensorLike]:
-        if not isinstance(obj, SimVar):
-            raise ValueError(f"Expected SimVar, got {type(obj)}")
-        return [v.to_numpy() if hasattr(v, "to_numpy") else v for v in obj._values]
-
-    # override
-    def evaluate(self, expr: Expr, bindings: dict[str, MPObject]) -> Sequence[MPObject]:
-        # sanity check for bindings.
-        for name, var in bindings.items():
-            if var.ctx is not self:
-                raise ValueError(f"Variable {name} not in this context, got {var.ctx}.")
-
-        pts_env = [
-            {name: cast(SimVar, var)._values[rank] for name, var in bindings.items()}
-            for rank in range(self.world_size())
-        ]
-
-        # Build per-rank evaluators with the per-party environment (runtime reused)
-        pts_evaluators: list[IEvaluator] = []
-        for rank in range(self.world_size()):
-            runtime = self._runtimes[rank]
-            ev = create_evaluator(
-                rank,
-                pts_env[rank],
-                self._comms[rank],
-                runtime,
-                None,
-            )
-            # Seed SPU once per runtime (idempotent logical requirement)
-            # Use setdefault to both retrieve and create metadata dict in one step.
-            spu_meta = runtime.state.setdefault("_spu", {})
-            if not spu_meta.get("inited", False):
-                link_ctx = self._spu_link_ctxs[rank]
-                seed_fn = PFunction(
-                    fn_type="spu.seed_env",
-                    ins_info=(),
-                    outs_info=(),
-                    config=self._spu_runtime_cfg,
-                    world=self._spu_world,
-                    link=link_ctx,
-                )
-                ev.runtime.run_kernel(seed_fn, [])  # type: ignore[arg-type]
-                spu_meta["inited"] = True
-            pts_evaluators.append(ev)
-
-        # Collect evaluation results from all parties
-        pts_results: list[Any] = []
-
-        with concurrent.futures.ThreadPoolExecutor() as executor:
-            futures = [
-                executor.submit(self._do_evaluate, expr, evaluator)
-                for evaluator in pts_evaluators
-            ]
-
-            # Collect results with proper exception handling
-            for i, future in enumerate(futures):
-                try:
-                    result = future.result(100)  # 100 second timeout
-                    pts_results.append(result)
-                except concurrent.futures.TimeoutError:
-                    faulthandler.dump_traceback(file=sys.stderr, all_threads=True)
-                    raise
-                except Exception as e:
-                    print(
-                        f"Exception in party {i}: {type(e).__name__}: {e}",
-                        file=sys.stderr,
-                    )
-                    traceback.print_exc(file=sys.stderr)
-                    executor.shutdown(wait=False, cancel_futures=True)
-                    raise
-
-        # Convert results to SimVar objects
-        # pts_results is a list of party results, where each party result is a list of values
-        # We need to transpose this to get (n_outputs, n_parties) structure
-        assert len(pts_results) == self.world_size()
-
-        # Ensure all parties returned the same number of outputs (matrix validation)
-        if pts_results and not all(
-            len(row) == len(pts_results[0]) for row in pts_results
-        ):
-            raise ValueError("Inconsistent number of outputs across parties")
-
-        # Transpose: (n_parties, n_outputs) -> (n_outputs, n_parties)
-        output_values = list(zip(*pts_results, strict=False))
-
-        # Get the output types from the expression
-        output_types = expr.mptypes
-
-        # Create SimVar objects for each output
-        sim_vars = []
-        for values, mptype in zip(output_values, output_types, strict=False):
-            sim_var = SimVar(self, mptype, list(values))
-            sim_vars.append(sim_var)
-
-        return sim_vars

mplang/v1/simp/__init__.py

@@ -1,13 +0,0 @@
-# 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.