mplang-nightly 0.1.dev277__tar.gz → 0.1.dev279__tar.gz

{mplang_nightly-0.1.dev277 → mplang_nightly-0.1.dev279}/PKG-INFO

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

{mplang_nightly-0.1.dev277 → mplang_nightly-0.1.dev279}/mplang/__init__.py

@@ -34,6 +34,7 @@ except Exception:
     # Fallback for development/editable installs when package is not installed
     __version__ = "0.0.0-dev"
 
+import mplang.tool as tool
 from mplang import dialects
 from mplang.backends.simp_driver.ops import DRIVER_HANDLERS
 from mplang.backends.simp_worker import SimpWorker
@@ -41,6 +42,8 @@ from mplang.backends.simp_worker.mem import LocalMesh
 from mplang.backends.simp_worker.ops import WORKER_HANDLERS
 from mplang.dialects.simp import make_driver, make_simulator
 from mplang.edsl import (
+    CompiledProgram,
+    FlatIOSignature,
     Graph,
     GraphPrinter,
     Object,
@@ -125,7 +128,7 @@ def _get_context(context: Interpreter | None) -> Interpreter:
 
 
 def evaluate(
-    fn: Callable[..., Any] | TracedFunction,
+    fn: Callable[..., Any] | TracedFunction | CompiledProgram,
     *args: Any,
     context: Interpreter | None = None,
     **kwargs: Any,
@@ -158,15 +161,33 @@ def evaluate(
             return val.runtime_obj
         return val
 
+    def eval_graph(graph: Graph, inputs: list[Any]) -> list[InterpObject]:
+        runtime_inputs = [unwrap_if_interp(v) for v in inputs]
+        raw_result = interp.evaluate_graph(graph, runtime_inputs)
+        return [
+            InterpObject(v, graph.outputs[i].type, interp)
+            for i, v in enumerate(raw_result)
+        ]
+
     with interp:
+        if isinstance(fn, CompiledProgram):
+            if kwargs:
+                raise TypeError(
+                    "mp.evaluate(CompiledProgram, ...) does not accept keyword arguments; "
+                    "pass flat positional inputs only."
+                )
+            if len(args) != fn.signature.input_arity:
+                raise ValueError(
+                    "CompiledProgram requires flat positional inputs matching its signature; "
+                    f"expected {fn.signature.input_arity}, got {len(args)}."
+                )
+
+            return eval_graph(fn.graph, list(args))
+
         if isinstance(fn, TracedFunction):
             inputs = fn.prepare_inputs(*args, **kwargs)
             inputs = [unwrap_if_interp(v) for v in inputs]
-            raw_result = interp.evaluate_graph(fn.graph, inputs)
-            wrapped = [
-                InterpObject(v, fn.graph.outputs[i].type, interp)
-                for i, v in enumerate(raw_result)
-            ]
+            wrapped = eval_graph(fn.graph, inputs)
             return fn.reconstruct_outputs(wrapped)
 
         return fn(*args, **kwargs)
@@ -417,6 +438,9 @@ __all__ = [  # noqa: RUF022
     "WORKER_HANDLERS",
     "make_driver",
     "make_simulator",
+    "tool",
+    "CompiledProgram",
+    "FlatIOSignature",
     # Dialects
     "dialects",
     "register_default_context_factory",

{mplang_nightly-0.1.dev277 → mplang_nightly-0.1.dev279}/mplang/backends/simp_worker/ops.py

@@ -84,16 +84,20 @@ def _shuffle_static_worker_impl(
     my_rank = worker.rank
     data = args[0]
 
+    exec_id = interpreter.current_op_exec_id()
+    graph_key = interpreter.current_graph_exec_key()
+    key_prefix = f"shuffle_{graph_key}_{op.name}_{exec_id}"
+
     for tgt, src in routing.items():
         if src == my_rank and tgt != my_rank:
-            key = f"shuffle_{op.name}_{tgt}"
+            key = f"{key_prefix}_{tgt}"
             comm.send(tgt, key, data)
 
     if my_rank in routing:
         src = routing[my_rank]
         if src == my_rank:
             return data
-        key = f"shuffle_{op.name}_{my_rank}"
+        key = f"{key_prefix}_{my_rank}"
         return comm.recv(src, key)
     else:
         return None

{mplang_nightly-0.1.dev277 → mplang_nightly-0.1.dev279}/mplang/edsl/__init__.py

@@ -53,6 +53,7 @@ from .jit import jit
 from .object import Object
 from .primitive import Primitive, primitive
 from .printer import GraphPrinter, format_graph
+from .program import CompiledProgram, FlatIOSignature
 from .tracer import TracedFunction, TraceObject, Tracer, trace
 from .typing import MPType, ScalarType, SSType, TableType, TensorType, VectorType
 
@@ -65,7 +66,9 @@ TensorObject = Object[TensorType]
 VectorObject = Object[VectorType]
 
 __all__ = [
+    "CompiledProgram",
     "Context",
+    "FlatIOSignature",
     "Graph",
     "GraphPrinter",
     "MPObject",

mplang_nightly-0.1.dev279/mplang/edsl/program.py

@@ -0,0 +1,134 @@
+# 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.
+
+from __future__ import annotations
+
+import hashlib
+import json
+from dataclasses import dataclass
+from typing import Any, ClassVar
+
+from mplang.edsl import serde
+from mplang.edsl.graph import Graph
+
+
+@dataclass(frozen=True)
+class FlatIOSignature:
+    """Portable I/O signature for source-free execution.
+
+    Only supports flat positional inputs/outputs corresponding to
+    `graph.inputs` / `graph.outputs`.
+    """
+
+    kind: ClassVar[str] = "flat_list_v0"
+    input_arity: int
+    output_arity: int
+
+    def to_json(self) -> dict[str, Any]:
+        return {
+            "kind": self.kind,
+            "input_arity": self.input_arity,
+            "output_arity": self.output_arity,
+        }
+
+    @classmethod
+    def from_json(cls, data: dict[str, Any]) -> FlatIOSignature:
+        if data.get("kind") != cls.kind:
+            raise ValueError(f"Unsupported signature kind: {data.get('kind')}")
+        return cls(
+            input_arity=int(data["input_arity"]),
+            output_arity=int(data["output_arity"]),
+        )
+
+
+@serde.register_class
+@dataclass
+class CompiledProgram:
+    """Executable program decoupled from user source.
+
+    This is a *logical model*; packaging (file/zip/etc.) is handled by tool layer.
+
+    Current constraints:
+    - signature is flat positional list I/O.
+    - no closure captures.
+    - no constant outputs (out_imms) unless future signature captures them.
+    """
+
+    _serde_kind: ClassVar[str] = "mplang.CompiledProgram"
+
+    graph: Graph
+    signature: FlatIOSignature
+    required_opcodes: list[str]
+    graph_digest: str
+    required_world_size: int | None = None
+    created_at: str | None = None
+    mplang_version: str | None = None
+    schema_version: int = 1
+    name: str | None = None
+
+    def to_json(self) -> dict[str, Any]:
+        return {
+            "schema_version": self.schema_version,
+            "name": self.name,
+            "graph": serde.to_json(self.graph),
+            "signature": self.signature.to_json(),
+            "required_opcodes": list(self.required_opcodes),
+            "graph_digest": self.graph_digest,
+            "required_world_size": self.required_world_size,
+            "created_at": self.created_at,
+            "mplang_version": self.mplang_version,
+        }
+
+    @classmethod
+    def from_json(cls, data: dict[str, Any]) -> CompiledProgram:
+        if "schema_version" not in data:
+            raise KeyError("Missing required field: schema_version")
+        schema_version = int(data["schema_version"])
+        if schema_version != 1:
+            raise ValueError(
+                f"Unsupported CompiledProgram schema_version: {schema_version}"
+            )
+
+        graph = serde.from_json(data["graph"])
+        if not isinstance(graph, Graph):
+            raise TypeError(
+                f"Expected graph to deserialize to Graph, got {type(graph).__name__}"
+            )
+
+        signature = FlatIOSignature.from_json(data["signature"])
+
+        required_world_size = data.get("required_world_size")
+        if required_world_size is not None:
+            required_world_size = int(required_world_size)
+        return cls(
+            graph=graph,
+            signature=signature,
+            required_opcodes=list(data.get("required_opcodes", [])),
+            graph_digest=str(data["graph_digest"]),
+            required_world_size=required_world_size,
+            created_at=data.get("created_at"),
+            mplang_version=data.get("mplang_version"),
+            schema_version=schema_version,
+            name=data.get("name"),
+        )
+
+
+def compute_graph_digest(graph: Graph) -> str:
+    """Compute a deterministic digest for a Graph.
+
+    We intentionally avoid `serde.dumps()` because it doesn't sort keys.
+    """
+
+    canonical = json.dumps(serde.to_json(graph), sort_keys=True, separators=(",", ":"))
+    return hashlib.sha256(canonical.encode("utf-8")).hexdigest()

{mplang_nightly-0.1.dev277 → mplang_nightly-0.1.dev279}/mplang/runtime/interpreter.py

@@ -24,17 +24,19 @@ from __future__ import annotations
 
 import collections
 import concurrent.futures
+import contextlib
+import hashlib
 import json
 import os
 import pathlib
 import queue
 import threading
 import time
-from collections.abc import Callable
+from collections.abc import Callable, Iterator
 from typing import TYPE_CHECKING, Any, cast
 
 from mplang.edsl.context import AbstractInterpreter
-from mplang.edsl.graph import Graph
+from mplang.edsl.graph import Graph, Value
 from mplang.edsl.object import Object
 from mplang.edsl.registry import get_impl
 from mplang.edsl.typing import BaseType
@@ -364,12 +366,201 @@ class Interpreter(AbstractInterpreter):
         # 2. MIMO Optimization: When one output of a multi-output op is computed,
         #    all sibling outputs are cached here to avoid re-execution.
         self._execution_cache: dict[Any, InterpObject] = {}
+
+        # -----------------------------------------------------------------
+        # Graph-local op execution ids (for deterministic communication tags)
+        # -----------------------------------------------------------------
+        # We assign a monotonically increasing exec_id to each op execution
+        # within a graph namespace, and keep it deterministic across parties.
+        #
+        # IMPORTANT:
+        # - We intentionally make exec_id grow across repeated executions of the
+        #   same region graph (e.g., while_loop iterations) to avoid tag/key reuse.
+        #
+        # Implementation:
+        # - Each evaluate_graph(graph, ...) reserves a contiguous exec_id range
+        #   [base, base + len(graph.operations)).
+        # - Op exec_id = base + op_index_in_graph.
+        # - Reservation is persisted per graph_exec_key (structural hash).
+        # - We forbid concurrent execution of the same graph_hash to avoid
+        #   message tag confusion when a backend uses only per-op tags.
+        self._exec_id_lock = threading.Lock()
+        self._graph_next_exec_base: dict[str, int] = {}
+        self._active_graph_exec_keys: set[str] = set()
+        self._tls = threading.local()
         self.executor = executor
         self.async_ops: set[str] = set()
         self.name = name
         self.trace_pid = trace_pid
         self.store: ObjectStore | None = store
 
+    @contextlib.contextmanager
+    def _tls_exec_context(
+        self,
+        *,
+        graph_exec_key: str | None = None,
+        op_exec_id: int | None = None,
+    ) -> Iterator[None]:
+        """Temporarily set execution context in thread-local storage."""
+
+        prev_graph_key = getattr(self._tls, "current_graph_exec_key", None)
+        prev_exec_id = getattr(self._tls, "current_op_exec_id", None)
+
+        if graph_exec_key is not None:
+            self._tls.current_graph_exec_key = graph_exec_key
+        if op_exec_id is not None:
+            self._tls.current_op_exec_id = op_exec_id
+
+        try:
+            yield
+        finally:
+            if graph_exec_key is not None:
+                if prev_graph_key is None:
+                    delattr(self._tls, "current_graph_exec_key")
+                else:
+                    self._tls.current_graph_exec_key = prev_graph_key
+
+            if op_exec_id is not None:
+                if prev_exec_id is None:
+                    delattr(self._tls, "current_op_exec_id")
+                else:
+                    self._tls.current_op_exec_id = prev_exec_id
+
+    def _graph_exec_key(self, graph: Graph) -> str:
+        """Return a deterministic, structural hash for a graph.
+
+        Used for:
+        - Namespacing per-graph exec_id counters
+        - Communication tag disambiguation (worker ops may include this key)
+
+        Note: we cache on the Graph object assuming graphs are immutable during
+        execution (finalized graphs / regions).
+        """
+
+        cached = getattr(graph, "_exec_key", None)
+        if cached is not None:
+            return cast(str, cached)
+
+        # NOTE: We intentionally do NOT use graph.to_json() here.
+        # graph.to_json() requires all attrs to be JSON-serializable via serde,
+        # but graphs may legitimately contain runtime-only objects (e.g. JAX
+        # PyTreeDef used by func.func). For communication tag namespaces we use
+        # a simple structural fingerprint that is deterministic across parties.
+
+        def _stable_attr_value(obj: Any) -> Any | None:
+            """Return a JSON-compatible stable value or None if unsupported.
+
+            We include only values that are likely deterministic across parties.
+            Unknown runtime objects are skipped (e.g. PyTreeDef, callables, etc.).
+            """
+
+            if obj is None or isinstance(obj, (bool, int, float, str)):
+                return obj
+
+            if isinstance(obj, (bytes, bytearray, memoryview)):
+                b = bytes(obj)
+                return {
+                    "_kind": "bytes",
+                    "len": len(b),
+                    "sha256": hashlib.sha256(b).hexdigest(),
+                }
+
+            try:
+                import numpy as np  # type: ignore
+
+                if isinstance(obj, np.ndarray):
+                    b = obj.tobytes(order="C")
+                    return {
+                        "_kind": "ndarray",
+                        "dtype": str(obj.dtype),
+                        "shape": list(obj.shape),
+                        "sha256": hashlib.sha256(b).hexdigest(),
+                    }
+                if isinstance(obj, (np.integer, np.floating)):
+                    return obj.item()
+            except Exception:
+                pass
+
+            if isinstance(obj, (list, tuple)):
+                items: list[Any] = []
+                for x in obj:
+                    sx = _stable_attr_value(x)
+                    if sx is None:
+                        return None
+                    items.append(sx)
+                return items
+
+            if isinstance(obj, dict):
+                stable_items: list[tuple[Any, Any]] = []
+                for k, v in obj.items():
+                    sk = _stable_attr_value(k)
+                    sv = _stable_attr_value(v)
+                    if sk is None or sv is None:
+                        return None
+                    stable_items.append((sk, sv))
+                stable_items.sort(
+                    key=lambda kv: json.dumps(
+                        kv[0], sort_keys=True, separators=(",", ":"), ensure_ascii=False
+                    )
+                )
+                return {"_kind": "dict", "items": stable_items}
+
+            return None
+
+        def _graph_fingerprint(g: Graph) -> Any:
+            # Map SSA Values to stable indices independent of their textual names.
+            value_to_index: dict[Value, int] = {}
+
+            def _index(v: Value) -> int:
+                if v in value_to_index:
+                    return value_to_index[v]
+                value_to_index[v] = len(value_to_index)
+                return value_to_index[v]
+
+            for v in g.inputs:
+                _index(v)
+            for op in g.operations:
+                for out in op.outputs:
+                    _index(out)
+
+            ops_fp: list[dict[str, Any]] = []
+            for op in g.operations:
+                attr_keys = sorted(op.attrs.keys())
+                stable_attr_items: list[tuple[str, Any]] = []
+                for k in attr_keys:
+                    attr_val = op.attrs.get(k)
+                    sv = _stable_attr_value(attr_val)
+                    if sv is not None:
+                        stable_attr_items.append((k, sv))
+
+                ops_fp.append({
+                    "opcode": op.opcode,
+                    "inputs": [_index(v) for v in op.inputs],
+                    "outputs": [str(v.type) for v in op.outputs],
+                    "attrs": {"keys": attr_keys, "stable": stable_attr_items},
+                    "regions": [_graph_fingerprint(r) for r in op.regions],
+                })
+
+            return {
+                "inputs": [str(v.type) for v in g.inputs],
+                "ops": ops_fp,
+                "outputs": [_index(v) for v in g.outputs],
+            }
+
+        fingerprint = _graph_fingerprint(graph)
+
+        payload = json.dumps(
+            fingerprint,
+            sort_keys=True,
+            separators=(",", ":"),
+            ensure_ascii=False,
+        ).encode("utf-8")
+        key = hashlib.sha256(payload).hexdigest()
+
+        # Store on graph to avoid id(graph) reuse pitfalls.
+        graph._exec_key = key  # type: ignore[attr-defined]
+        return key
+
     def shutdown(self) -> None:
         """Shutdown the interpreter and release resources.
 
@@ -641,18 +832,70 @@ class Interpreter(AbstractInterpreter):
         Returns:
             List of runtime execution results corresponding to graph.outputs.
         """
-        logger.debug(
-            "Evaluating graph: %d inputs, %d ops, %d outputs (job_id=%s, async=%s)",
-            len(inputs),
-            len(graph.operations),
-            len(graph.outputs),
-            job_id,
-            self.executor is not None,
-        )
-        if self.executor:
-            return self._evaluate_graph_async(graph, inputs, job_id)
-        else:
-            return self._evaluate_graph_sync(graph, inputs, job_id)
+        graph_exec_key = self._graph_exec_key(graph)
+
+        # Prevent concurrent execution of the same graph hash.
+        with self._exec_id_lock:
+            if graph_exec_key in self._active_graph_exec_keys:
+                raise RuntimeError(
+                    "Concurrent execution of the same graph is not allowed. "
+                    f"graph_exec_key={graph_exec_key}"
+                )
+            self._active_graph_exec_keys.add(graph_exec_key)
+
+        try:
+            with self._tls_exec_context(graph_exec_key=graph_exec_key):
+                logger.debug(
+                    "Evaluating graph: %d inputs, %d ops, %d outputs (job_id=%s, async=%s, graph_key=%s)",
+                    len(inputs),
+                    len(graph.operations),
+                    len(graph.outputs),
+                    job_id,
+                    self.executor is not None,
+                    graph_exec_key,
+                )
+                if self.executor:
+                    return self._evaluate_graph_async(graph, inputs, job_id)
+                else:
+                    return self._evaluate_graph_sync(graph, inputs, job_id)
+        finally:
+            with self._exec_id_lock:
+                self._active_graph_exec_keys.discard(graph_exec_key)
+
+    def _reserve_op_exec_base(self, graph: Graph) -> int:
+        """Reserve a contiguous exec_id range for a single evaluate_graph call.
+
+        Counter is namespaced by the current graph_exec_key.
+        """
+        key = self.current_graph_exec_key()
+        with self._exec_id_lock:
+            base = self._graph_next_exec_base.get(key, 0)
+            self._graph_next_exec_base[key] = base + len(graph.operations)
+        return base
+
+    def current_graph_exec_key(self) -> str:
+        """Return current graph execution key during evaluate_graph execution."""
+
+        key = getattr(self._tls, "current_graph_exec_key", None)
+        if key is None:
+            raise RuntimeError(
+                "current_graph_exec_key() called outside of evaluate_graph execution"
+            )
+        return cast(str, key)
+
+    def current_op_exec_id(self) -> int:
+        """Return current op exec_id during graph execution.
+
+        Worker-side implementations can use this to build deterministic,
+        unique communication tags without coupling to any specific op.
+        """
+
+        exec_id = getattr(self._tls, "current_op_exec_id", None)
+        if exec_id is None:
+            raise RuntimeError(
+                "current_op_exec_id() called outside of evaluate_graph execution"
+            )
+        return cast(int, exec_id)
 
     def _evaluate_graph_sync(
         self, graph: Graph, inputs: list[Any], job_id: str | None = None
@@ -661,7 +904,10 @@ class Interpreter(AbstractInterpreter):
         # Local environment: Value -> Runtime Object
         env = dict(zip(graph.inputs, inputs, strict=True))
 
-        for op in graph.operations:
+        op_exec_base = self._reserve_op_exec_base(graph)
+
+        for op_index, op in enumerate(graph.operations):
+            exec_id = op_exec_base + op_index
             # Resolve inputs
             try:
                 args = [env[val] for val in op.inputs]
@@ -685,15 +931,16 @@ class Interpreter(AbstractInterpreter):
             if not handler:
                 handler = get_impl(op.opcode)
 
-            if handler:
-                # Pass interpreter to support recursive execution (HOFs)
-                # Pass op to access attributes and regions
-                # Pass args as runtime values
-                results = handler(self, op, *args)
-            else:
-                raise NotImplementedError(
-                    f"No implementation registered for opcode: {op.opcode}"
-                )
+            with self._tls_exec_context(op_exec_id=exec_id):
+                if handler:
+                    # Pass interpreter to support recursive execution (HOFs)
+                    # Pass op to access attributes and regions
+                    # Pass args as runtime values
+                    results = handler(self, op, *args)
+                else:
+                    raise NotImplementedError(
+                        f"No implementation registered for opcode: {op.opcode}"
+                    )
 
             # Update environment with outputs
             # Handler should return a single value or a tuple/list of values
@@ -719,6 +966,9 @@ class Interpreter(AbstractInterpreter):
         self, graph: Graph, inputs: list[Any], job_id: str | None = None
     ) -> list[Any]:
         """Asynchronous execution with non-blocking DAG scheduling."""
+        graph_exec_key = self.current_graph_exec_key()
+        op_exec_base = self._reserve_op_exec_base(graph)
+        op_to_index = {op: i for i, op in enumerate(graph.operations)}
         # Tracer setup (if not provided, use a disabled stub)
         tracer: ExecutionTracer | _NullTracer
         if self.tracer:
@@ -817,6 +1067,8 @@ class Interpreter(AbstractInterpreter):
             # Extract args from env (must be ready)
             args = [env[val] for val in op.inputs]
 
+            exec_id = op_exec_base + op_to_index[op]
+
             handler = self.handlers.get(op.opcode)
             if not handler:
                 handler = get_impl(op.opcode)
@@ -833,12 +1085,15 @@ class Interpreter(AbstractInterpreter):
 
                 # Submit to executor
                 def task() -> Any:
-                    start_ts = tracer.log_start(
-                        op, pid=self.trace_pid, namespace=self.trace_pid
-                    )
-                    res = handler(self, op, *args)
-                    tracer.log_end(op, start_ts, pid=self.trace_pid)
-                    return res
+                    with self._tls_exec_context(
+                        graph_exec_key=graph_exec_key, op_exec_id=exec_id
+                    ):
+                        start_ts = tracer.log_start(
+                            op, pid=self.trace_pid, namespace=self.trace_pid
+                        )
+                        res = handler(self, op, *args)
+                        tracer.log_end(op, start_ts, pid=self.trace_pid)
+                        return res
 
                 def callback(fut: Any) -> None:
                     try:
@@ -852,12 +1107,15 @@ class Interpreter(AbstractInterpreter):
             else:
                 # Sync execution (run immediately)
                 try:
-                    start_ts = tracer.log_start(
-                        op, pid=self.trace_pid, namespace=self.trace_pid
-                    )
-                    res = handler(self, op, *args)
-                    tracer.log_end(op, start_ts, pid=self.trace_pid)
-                    on_op_done(op, res)
+                    with self._tls_exec_context(
+                        graph_exec_key=graph_exec_key, op_exec_id=exec_id
+                    ):
+                        start_ts = tracer.log_start(
+                            op, pid=self.trace_pid, namespace=self.trace_pid
+                        )
+                        res = handler(self, op, *args)
+                        tracer.log_end(op, start_ts, pid=self.trace_pid)
+                        on_op_done(op, res)
                 except Exception as e:
                     on_op_done(op, None, error=e)