mplang-nightly 0.1.dev277__py3-none-any.whl → 0.1.dev279__py3-none-any.whl

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/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/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/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/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)
 

mplang/tool/__init__.py

@@ -0,0 +1,46 @@
+# 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.
+
+"""Tool-layer APIs for MPLang.
+
+This package contains utilities that are intentionally *not* part of the core
+EDSL execution surface. In particular, compile/execute decoupling lives here:
+- build a portable `CompiledProgram`
+- pack/unpack to a container format
+
+These helpers must not depend on user source code being available at execution.
+"""
+
+from __future__ import annotations
+
+from mplang.edsl.program import CompiledProgram, FlatIOSignature
+from mplang.tool.program import (
+    compile_program,
+    inspect_artifact,
+    pack,
+    pack_to_path,
+    unpack,
+    unpack_path,
+)
+
+__all__ = [
+    "CompiledProgram",
+    "FlatIOSignature",
+    "compile_program",
+    "inspect_artifact",
+    "pack",
+    "pack_to_path",
+    "unpack",
+    "unpack_path",
+]

mplang/tool/program.py

@@ -0,0 +1,335 @@
+# 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 io
+import json
+import tarfile
+from datetime import UTC, datetime
+from pathlib import Path
+from typing import Any, Literal
+
+import mplang
+from mplang.edsl import serde
+from mplang.edsl.graph import Graph
+from mplang.edsl.program import (
+    CompiledProgram,
+    FlatIOSignature,
+    compute_graph_digest,
+)
+from mplang.edsl.tracer import TracedFunction, trace
+
+DEFAULT_MAX_ARTIFACT_JSON_BYTES = 512 * 1024 * 1024  # 512 MiB
+
+
+def _utc_now_iso() -> str:
+    return datetime.now(UTC).isoformat()
+
+
+def _iter_graphs(root: Graph) -> list[Graph]:
+    # Use an explicit stack to avoid Python recursion limits.
+    # Also guard against potential region graph cycles.
+    out: list[Graph] = []
+    stack: list[Graph] = [root]
+    visited: set[int] = set()
+    while stack:
+        graph = stack.pop()
+        graph_id = id(graph)
+        if graph_id in visited:
+            continue
+        visited.add(graph_id)
+        out.append(graph)
+
+        for op in graph.operations:
+            if op.regions:
+                stack.extend(op.regions)
+
+    return out
+
+
+def _collect_opcodes(graph: Graph) -> set[str]:
+    opcodes: set[str] = set()
+    for g in _iter_graphs(graph):
+        for op in g.operations:
+            opcodes.add(op.opcode)
+    return opcodes
+
+
+def _collect_parties(graph: Graph) -> set[int]:
+    parties: set[int] = set()
+    for g in _iter_graphs(graph):
+        for op in g.operations:
+            raw = op.attrs.get("parties")
+            if raw is None:
+                continue
+
+            if not isinstance(raw, (list, tuple, set)):
+                raise TypeError(
+                    "Invalid 'parties' attribute: expected list/tuple/set of ints, "
+                    f"got {type(raw).__name__}"
+                )
+            for p in raw:
+                p_int = int(p)
+                if p_int < 0:
+                    raise ValueError("Invalid 'parties' attribute: negative party id")
+                parties.add(p_int)
+    return parties
+
+
+def _compute_required_world_size(graph: Graph) -> int:
+    parties = _collect_parties(graph)
+    if not parties:
+        return 0
+    return max(parties) + 1
+
+
+def _validate_traced_for_artifact(traced: TracedFunction) -> None:
+    # Restriction: no closure captures
+    if traced.captured:
+        raise ValueError(
+            "CompiledProgram does not support closure captures; "
+            "please refactor to pass all values explicitly."
+        )
+
+    # Restriction: no constant outputs (out_imms)
+    if traced.out_imms:
+        raise ValueError(
+            "CompiledProgram does not support constant outputs (out_imms); "
+            "return only traced values (graph outputs)."
+        )
+
+    # Restriction: signature is flat positional list I/O.
+    # We do not preserve (args, kwargs) pytree metadata.
+    # We therefore require all runtime-provided inputs correspond exactly to graph.inputs.
+    if len(traced.graph.inputs) != len(traced.in_var_pos):
+        raise ValueError(
+            "CompiledProgram requires flat positional inputs that map 1:1 to graph.inputs; "
+            f"got graph.inputs={len(traced.graph.inputs)} but in_var_pos={len(traced.in_var_pos)}."
+        )
+
+
+def _validate_program(program: CompiledProgram) -> None:
+    if program.signature.kind != FlatIOSignature.kind:
+        raise ValueError(f"Unsupported signature kind: {program.signature.kind}")
+
+    if program.signature.input_arity != len(program.graph.inputs):
+        raise ValueError(
+            "Signature input_arity does not match graph.inputs: "
+            f"input_arity={program.signature.input_arity}, inputs={len(program.graph.inputs)}"
+        )
+    if program.signature.output_arity != len(program.graph.outputs):
+        raise ValueError(
+            "Signature output_arity does not match graph.outputs: "
+            f"output_arity={program.signature.output_arity}, outputs={len(program.graph.outputs)}"
+        )
+
+    expected_opcodes = sorted(_collect_opcodes(program.graph))
+    if sorted(program.required_opcodes) != expected_opcodes:
+        raise ValueError(
+            "required_opcodes mismatch with graph content; "
+            "artifact may be corrupted or constructed inconsistently."
+        )
+
+    actual_digest = compute_graph_digest(program.graph)
+    if program.graph_digest and program.graph_digest != actual_digest:
+        raise ValueError(
+            "Graph digest mismatch: "
+            f"expected={program.graph_digest}, actual={actual_digest}"
+        )
+
+    expected_world_size = _compute_required_world_size(program.graph)
+    if (
+        program.required_world_size is not None
+        and program.required_world_size != expected_world_size
+    ):
+        raise ValueError(
+            "required_world_size mismatch with graph content; "
+            f"expected={expected_world_size}, got={program.required_world_size}."
+        )
+
+    # Ensure JSON serialization works (fail fast for non-serde attrs).
+    serde.to_json(program)
+
+
+def compile_program(
+    fn_or_traced: Any,
+    *args: Any,
+    context: Any | None = None,
+    name: str | None = None,
+    **kwargs: Any,
+) -> CompiledProgram:
+    """Compile (trace) into a source-free executable `CompiledProgram`.
+
+    Restrictions (enforced):
+    - no closure captures
+    - no constant outputs (`out_imms` must be empty)
+    - signature is flat list (positional) I/O
+
+    Note: `in_imms` (compile-time constants) are allowed: they are baked into the graph.
+    """
+
+    traced: TracedFunction
+    if isinstance(fn_or_traced, TracedFunction):
+        traced = fn_or_traced
+    else:
+        if context is not None:
+            with context:
+                traced = trace(fn_or_traced, *args, **kwargs)
+        else:
+            traced = trace(fn_or_traced, *args, **kwargs)
+
+    _validate_traced_for_artifact(traced)
+
+    signature = FlatIOSignature(
+        input_arity=len(traced.graph.inputs),
+        output_arity=len(traced.graph.outputs),
+    )
+
+    required_opcodes = sorted(_collect_opcodes(traced.graph))
+    graph_digest = compute_graph_digest(traced.graph)
+    required_world_size = _compute_required_world_size(traced.graph)
+
+    program = CompiledProgram(
+        graph=traced.graph,
+        signature=signature,
+        required_opcodes=required_opcodes,
+        graph_digest=graph_digest,
+        required_world_size=required_world_size,
+        created_at=_utc_now_iso(),
+        mplang_version=getattr(mplang, "__version__", None),
+        name=name or traced.name,
+    )
+    return program
+
+
+def pack(program: CompiledProgram, *, compress: bool = True) -> bytes:
+    """Pack a `CompiledProgram` into portable bytes.
+
+    Container format (recommended): a `tar.gz` archive containing a single
+    human-readable JSON file `artifact.json`.
+
+    This allows users to inspect artifacts via:
+        `tar -xzf program.tar.gz && cat artifact.json`
+
+    If `compress=False`, returns an uncompressed tar archive (still extractable
+    via `tar -xf`).
+    """
+
+    artifact_json = json.dumps(
+        serde.to_json(program),
+        ensure_ascii=False,
+        indent=2,
+        sort_keys=True,
+    ).encode("utf-8")
+
+    buf = io.BytesIO()
+    mode: Literal["w:gz", "w"] = "w:gz" if compress else "w"
+    with tarfile.open(fileobj=buf, mode=mode) as tf:
+        info = tarfile.TarInfo(name="artifact.json")
+        info.size = len(artifact_json)
+        tf.addfile(info, io.BytesIO(artifact_json))
+
+    return buf.getvalue()
+
+
+def pack_to_path(
+    program: CompiledProgram, path: str | Path, *, compress: bool = True
+) -> Path:
+    """Pack and write artifact to disk.
+
+    Args:
+        program: Program to pack.
+        path: Output path (typically ends with `.tar.gz`).
+        compress: Whether to gzip the tar archive.
+
+    Returns:
+        The resolved output path.
+    """
+
+    out_path = Path(path).expanduser().resolve()
+    out_path.write_bytes(pack(program, compress=compress))
+    return out_path
+
+
+def unpack(
+    data: bytes, *, max_artifact_json_bytes: int = DEFAULT_MAX_ARTIFACT_JSON_BYTES
+) -> CompiledProgram:
+    """Unpack bytes into a `CompiledProgram`.
+
+    Supported container format: tar(.gz) containing `artifact.json`.
+    """
+
+    try:
+        with tarfile.open(fileobj=io.BytesIO(data), mode="r:*") as tf:
+            member = tf.getmember("artifact.json")
+
+            if not member.isfile():
+                raise ValueError("artifact.json is not a regular file")
+
+            if member.size < 0:
+                raise ValueError("Invalid artifact.json size in tar header")
+
+            if member.size > max_artifact_json_bytes:
+                raise ValueError(
+                    "artifact.json is too large to unpack safely: "
+                    f"size={member.size} bytes, limit={max_artifact_json_bytes} bytes"
+                )
+
+            f = tf.extractfile(member)
+            if f is None:
+                raise ValueError("artifact.json not found in tar archive")
+            payload = json.loads(f.read().decode("utf-8"))
+    except (tarfile.ReadError, KeyError, OSError, json.JSONDecodeError) as exc:
+        raise ValueError(
+            "Invalid artifact container: expected tar(.gz) with artifact.json"
+        ) from exc
+
+    program = serde.from_json(payload)
+    if not isinstance(program, CompiledProgram):
+        raise TypeError(
+            f"Expected artifact.json to deserialize to CompiledProgram, got {type(program).__name__}"
+        )
+
+    _validate_program(program)
+    return program
+
+
+def unpack_path(path: str | Path) -> CompiledProgram:
+    """Read an artifact from disk and unpack it."""
+
+    in_path = Path(path).expanduser().resolve()
+    return unpack(in_path.read_bytes())
+
+
+def inspect_artifact(data: bytes) -> dict[str, Any]:
+    """Return a JSON-friendly inspection report without executing."""
+
+    program = unpack(data)
+    return {
+        "schema_version": program.schema_version,
+        "name": program.name,
+        "mplang_version": program.mplang_version,
+        "created_at": program.created_at,
+        "graph_digest": program.graph_digest,
+        "required_world_size": program.required_world_size,
+        "signature": program.signature.to_json(),
+        "required_opcodes": program.required_opcodes,
+        "graph": {
+            "inputs": len(program.graph.inputs),
+            "ops": len(program.graph.operations),
+            "outputs": len(program.graph.outputs),
+            "region_count": sum(len(op.regions) for op in program.graph.operations),
+        },
+    }

{mplang_nightly-0.1.dev277.dist-info → mplang_nightly-0.1.dev279.dist-info}/METADATA

@@ -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.dist-info → mplang_nightly-0.1.dev279.dist-info}/RECORD

@@ -1,4 +1,4 @@
-mplang/__init__.py,sha256=QzbdmzTomqQwTZmLKayovb2cl2qcFVzvM_qxlfFfRQ0,13607
+mplang/__init__.py,sha256=PsUcGqKGQym3N_CU1Rav304YQwGVj8CLVx9S_6UTD9c,14519
 mplang/cli.py,sha256=NW0GmxZeRC4rrYg8RVBlZiDjkihBXGcHmrll-JFOqWM,20317
 mplang/cli_guide.md,sha256=hKC6AKgJn-lM_wZ0CzZIP2QUBxGPnT0Op_1YyeUhCfI,3581
 mplang/logging_config.py,sha256=6Zm1Z_EBnzxAfeKr94xiLpWYDD8fa4ZEs2g_kMoH8eI,7579
@@ -26,7 +26,7 @@ mplang/backends/simp_driver/values.py,sha256=Lz1utNSIzH-dCzZAEjU6JRcxPsfKGfUJrYl
 mplang/backends/simp_worker/__init__.py,sha256=gdrSY1-MDkupCoJ8xwwH7em7fgVWv3J4gBJ45uHdzgg,961
 mplang/backends/simp_worker/http.py,sha256=90nJnNLSM9TUVRxhAFq9pyNk0LwmSmvgnv3Tb8KFWSE,12660
 mplang/backends/simp_worker/mem.py,sha256=tMGiRppeca0TnY8WdqYQMQvsx5UVswCqdeOhiDlLQBs,3574
-mplang/backends/simp_worker/ops.py,sha256=X_mAhrSZkiN9GmgTJgJs1c50CFTNB_lo9zoPI4rIYU4,5510
+mplang/backends/simp_worker/ops.py,sha256=ntxfkD4e6Il4w7FshK1ODcUCUPMlipt33pDY_x5iC0U,5661
 mplang/backends/simp_worker/state.py,sha256=nIu0ybvdYqRqp0TkoSneUF2u31evDHucCRduVBaDals,1445
 mplang/dialects/__init__.py,sha256=CYMmkeQVU0Znr9n3_5clZKb16u7acJ5jl5Zjbx4Tn1U,1478
 mplang/dialects/bfv.py,sha256=m5YfobFCBqn0lg2zBM9RNs2AC7i4PUQH2qXjHLHwSy4,22332
@@ -42,13 +42,14 @@ mplang/dialects/table.py,sha256=i9ruyh91_tSWu9rsLomrBUfqRdbHiZMMMJzNKfMrAUc,1353
 mplang/dialects/tee.py,sha256=BMFSbeK-Ck2jQP4qY9bZeNYTxEa7uEtUWLZLC4BPQxk,10111
 mplang/dialects/tensor.py,sha256=7aAYKaMaFjJ8N25yPFnmVhUuUdKJYy-M-a4NsZGE7kY,39893
 mplang/edsl/README.md,sha256=viflvdRojOa6Xk_UMRPqpuPGXcPGmdlv2-XR6LO7B58,7592
-mplang/edsl/__init__.py,sha256=ARYS7FkkSXwjWCsPLtWc9kL5OaR4wd5zCQhFTkzZUp0,2598
+mplang/edsl/__init__.py,sha256=WL4efo6uY1br781_8IaCkSi7yCUldcfJfbtFsn6Fdj4,2698
 mplang/edsl/context.py,sha256=Ln8n3bDe8_ISe42TAGzUuz8fw57-tu1APuihMfAtW1Y,10075
 mplang/edsl/graph.py,sha256=nCeCN7-bxfzyv40fmxcEXOaVUx14cOCaHfFb7A9OBnE,14968
 mplang/edsl/jit.py,sha256=7eLZHoIuL5FZo9G5eF9nI4EeayLK-OvJ0NoH3VG5vLI,2393
 mplang/edsl/object.py,sha256=dBl58q-ondjpjPNBh8zZvIEj6pJw2yEoz6TCaM_oleA,1906
 mplang/edsl/primitive.py,sha256=gDrn4FH682DUOgTqcQ2-9aqDYJau9L8E1ElswyOmmdw,10859
 mplang/edsl/printer.py,sha256=drmfRkdCNqbkRfSDmejxtO-rEAaM13QyHB3AbAmKVFk,4393
+mplang/edsl/program.py,sha256=_JdEU2-nb79VlFLcgMJf4JS30TARBeUIzno0y0SFVsg,4467
 mplang/edsl/registry.py,sha256=hudXZPUrUUueEwgksDKN0cnE3iiXucuTaDdDK8uSPmk,6822
 mplang/edsl/serde.py,sha256=8K94laE8ObeGuBoF6m7g3A-xEe98EvqQ_6ZPPspddAY,11641
 mplang/edsl/tracer.py,sha256=EWN3eMVRG-CZsamTyINOnhhEUKhgd4CYwFMWeRpjycU,23129
@@ -91,13 +92,15 @@ mplang/libs/mpc/vole/ldpc.py,sha256=gOmIbyOjkGE5lewyatl3p6FizNNH8LZ_1oOhp_-TOck,
 mplang/libs/mpc/vole/silver.py,sha256=EIxhpFIVNBemgeIZzCu5Cz_4wysxRm9b1Xfu0xiweVQ,12218
 mplang/runtime/__init__.py,sha256=VdUwJ3kDaI46FvGw7iMGwcsjt0HTGmmRmaBwj99xKIw,620
 mplang/runtime/dialect_state.py,sha256=HxO1i4kSOujS2tQzAF9-WmI3nChSaGgupf2_07dHetY,1277
-mplang/runtime/interpreter.py,sha256=wYTglSWeF_INwM7pN6oVxsKTygiUIQOdKcspFECLaKc,33135
+mplang/runtime/interpreter.py,sha256=wcCWXpAGylqdw_HecR4suJtwmozHLrK5x6Q8xM-Pn24,43593
 mplang/runtime/object_store.py,sha256=yT6jtKG2GUEJVmpq3gnQ8mCMvUFYzgBciC5A-J5KRdk,5998
 mplang/runtime/value.py,sha256=EqlhSgxLTJi_FF3ppyKjMe4eHS6-ROx-zK1YesG1U4o,4311
+mplang/tool/__init__.py,sha256=9K-T50W_vClUlyERcVx5xGZaeyv0Ts63SaQX6AZtjIs,1341
+mplang/tool/program.py,sha256=W3H8bpPirnoJ4ZrmyPYuMCPadJis20o__n_1MKqCsWU,11058
 mplang/utils/__init__.py,sha256=toubeyISiT6WDdITdfAvdY2iXVZU3PKVNWVeC9sYxuA,947
 mplang/utils/func_utils.py,sha256=aZ-X43w8JKJgiF-IUMS0G7QqrNeoTM5ZPzRNd-tKxpw,5180
-mplang_nightly-0.1.dev277.dist-info/METADATA,sha256=RlFkHiM8w8IXi_mWFUMhH468NdRHTbWdpkZH3rqrYvU,16783
-mplang_nightly-0.1.dev277.dist-info/WHEEL,sha256=WLgqFyCfm_KASv4WHyYy0P3pM_m7J5L9k2skdKLirC8,87
-mplang_nightly-0.1.dev277.dist-info/entry_points.txt,sha256=mG1oJT-GAjQR834a62_QIWb7litzWPPyVnwFqm-rWuY,55
-mplang_nightly-0.1.dev277.dist-info/licenses/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
-mplang_nightly-0.1.dev277.dist-info/RECORD,,
+mplang_nightly-0.1.dev279.dist-info/METADATA,sha256=Q4l1RV5WC5NfRV0heHxzHH8qVf3CkjOi6Ag1kbvsX38,16783
+mplang_nightly-0.1.dev279.dist-info/WHEEL,sha256=WLgqFyCfm_KASv4WHyYy0P3pM_m7J5L9k2skdKLirC8,87
+mplang_nightly-0.1.dev279.dist-info/entry_points.txt,sha256=mG1oJT-GAjQR834a62_QIWb7litzWPPyVnwFqm-rWuY,55
+mplang_nightly-0.1.dev279.dist-info/licenses/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
+mplang_nightly-0.1.dev279.dist-info/RECORD,,