mplang-nightly 0.1.dev252__py3-none-any.whl → 0.1.dev254__py3-none-any.whl

mplang/v2/__init__.py

@@ -52,15 +52,19 @@ from mplang.v2.edsl import (
     TracedFunction,
     Tracer,
     Value,
+    find_context,
+    find_context_with_state,
+    find_interpreter,
     format_graph,
     get_current_context,
     get_default_context,
-    get_root_context,
+    is_tracing,
     jit,
     pop_context,
     primitive,
     push_context,
     register_default_context_factory,
+    set_root_context,
     trace,
 )
 from mplang.v2.edsl.registry import get_profiler
@@ -97,34 +101,6 @@ from mplang.v2.runtime.interpreter import Interpreter
 # =============================================================================
 
 
-def set_root_context(context: Interpreter, force: bool = False) -> None:
-    """Set the global/root execution context.
-
-    This explicitly sets the provided interpreter as the Root Context.
-    All subsequent operations (compile, evaluate, device resolution) will
-    use this context as the default environment.
-
-    Args:
-        context: Interpreter to use as the root context.
-        force: If True, clears the existing context stack before setting.
-               If False (default), pushes onto the stack.
-    """
-    from mplang.v2.edsl.context import _context_stack, get_current_context
-
-    if force:
-        _context_stack.clear()
-        _context_stack.append(context)
-        return
-
-    if get_current_context() is not None:
-        raise RuntimeError(
-            "Cannot set root context: Context stack is not empty. "
-            "Use force=True to overwrite the existing root context."
-        )
-
-    push_context(context)
-
-
 def _get_context(context: Interpreter | None) -> Interpreter:
     """Get context from parameter or context stack."""
     if context is not None:
@@ -204,6 +180,30 @@ def fetch(
 ) -> Any:
     """Fetch results from interpreter context to Python.
 
+    This is a meta-function that operates at execution boundaries, not a traced
+    dialect operation. It brings data from the distributed/MPC runtime back to
+    the Python host.
+
+    Behavior in different contexts:
+    - **Tracing (compile)**: Returns the input unchanged (identity). The graph
+      outputs are determined by the function's return statement, not fetch calls.
+    - **Execution (evaluate)**: Actually fetches data from workers/parties.
+
+    Design Note (A vs B tradeoff):
+        Two designs were considered for fetch behavior during tracing:
+
+        - **Design A (chosen)**: fetch = identity during tracing. Graph outputs
+          are determined solely by the return statement. This is simpler and
+          avoids ambiguity when fetch and return reference different values.
+
+        - **Design B (alternative)**: fetch marks output points in the graph.
+          This would allow fetch(a), fetch(b), return b to output both a and b.
+          However, it complicates the semantics and requires tracking fetch
+          points separately from return values.
+
+        Design A was chosen for simplicity. If a value needs to be an output,
+        it should be returned. fetch's role is purely for execution-time I/O.
+
     Args:
         result: Object(s) to fetch. Can be a single InterpObject, DriverVar,
             or nested structure containing them.
@@ -216,14 +216,22 @@ def fetch(
         Fetched Python values. For device objects with follow_device=True,
         returns single value from the device's rank(s). Otherwise returns
         list of values (one per party) or single value for world_size=1.
+        During tracing, returns the input unchanged.
     """
-
     from jax.tree_util import tree_map
 
     from mplang.v2.backends.simp_driver.values import DriverVar
+    from mplang.v2.edsl.context import is_tracing
     from mplang.v2.runtime.interpreter import InterpObject
     from mplang.v2.runtime.value import WrapValue
 
+    # Check if we are in tracing context - if so, return identity
+    if is_tracing():
+        # Design A: fetch = identity during tracing
+        # Graph outputs are determined by return statement, not fetch calls
+        return result
+
+    # Execution context - actually fetch data
     interp = _get_context(context)
 
     def _fetch_single(var: Any) -> Any:
@@ -375,10 +383,14 @@ __all__ = [  # noqa: RUF022
     "compile",
     "evaluate",
     "fetch",
+    "find_context",
+    "find_context_with_state",
+    "find_interpreter",
     "format_graph",
     "function",
     "get_current_context",
     "get_default_context",
+    "is_tracing",
     "jit",
     "mplang",
     "pop_context",
@@ -404,7 +416,6 @@ __all__ = [  # noqa: RUF022
     # Dialects
     "dialects",
     "register_default_context_factory",
-    "get_root_context",
     "get_profiler",
 ]
 

mplang/v2/backends/simp_driver/mem.py

@@ -154,6 +154,10 @@ class SimpMemDriver(SimpDriver):
         self._workers = workers
         self._mesh = mesh
 
+    def shutdown(self) -> None:
+        """Shutdown the local memory driver and its mesh."""
+        self._mesh.shutdown()
+
     @property
     def world_size(self) -> int:
         return self._world_size

mplang/v2/cli.py

@@ -62,7 +62,7 @@ def run_worker(
     signal.signal(signal.SIGINT, signal.SIG_DFL)
     signal.signal(signal.SIGTERM, signal.SIG_DFL)
 
-    from mplang.v2.backends.simp_http_worker import create_worker_app
+    from mplang.v2.backends.simp_worker.http import create_worker_app
 
     app = create_worker_app(rank, world_size, endpoints, spu_endpoints)
 

mplang/v2/edsl/__init__.py

@@ -33,12 +33,16 @@ from . import typing as typing
 # Context management
 from .context import (
     Context,
+    find_context,
+    find_context_with_state,
+    find_interpreter,
     get_current_context,
     get_default_context,
-    get_root_context,
+    is_tracing,
     pop_context,
     push_context,
     register_default_context_factory,
+    set_root_context,
 )
 
 # Graph IR
@@ -77,15 +81,19 @@ __all__ = [
     "Tracer",
     "Value",
     "VectorObject",
+    "find_context",
+    "find_context_with_state",
+    "find_interpreter",
     "format_graph",
     "get_current_context",
     "get_default_context",
-    "get_root_context",
+    "is_tracing",
     "jit",
     "pop_context",
     "primitive",
     "push_context",
     "register_default_context_factory",
+    "set_root_context",
     "trace",
     "typing",
 ]

mplang/v2/edsl/context.py

@@ -27,6 +27,16 @@ Contexts can be used directly with Python's 'with' statement:
     with tracer:
         # Operations run under tracer context
         result = primitive.bind(x, y)
+
+State Management:
+    Contexts can carry arbitrary named state via set_state/get_state.
+    This allows different layers (device, ml, analytics) to attach their
+    own state without the EDSL layer knowing about specific state types.
+
+    State key conventions:
+    - "dialect.{name}": Dialect runtime state (e.g., "dialect.simp")
+    - "device.cluster": Device/cluster configuration
+    - "ml.{component}": ML pipeline components
 """
 
 from __future__ import annotations
@@ -42,14 +52,65 @@ if TYPE_CHECKING:
 
 
 class Context(ABC):
-    """Base class for EDSL execution contexts.
+    """Base class for EDSL execution contexts with extensible state slots.
 
     A Context represents an environment where primitives are executed.
     There are two types of contexts:
     - Tracer: Records operations to Graph IR (compile-time)
     - Interpreter: Execution context (executes operations immediately)
+
+    State Management:
+        Contexts can carry arbitrary named state. Different layers can attach
+        their own state without the EDSL layer knowing specifics:
+
+        >>> ctx.set_state("device.cluster", cluster_spec)
+        >>> ctx.set_state("dialect.simp", simp_driver)
+        >>> cluster = ctx.get_state("device.cluster")
     """
 
+    def __init__(self) -> None:
+        self._states: dict[str, Any] = {}
+
+    # =========================================================================
+    # State Management
+    # =========================================================================
+
+    def set_state(self, key: str, value: Any) -> None:
+        """Attach state to this context.
+
+        Args:
+            key: State key (e.g., "dialect.simp", "device.cluster")
+            value: State value
+        """
+        self._states[key] = value
+
+    def get_state(self, key: str, default: Any = None) -> Any:
+        """Get attached state by key.
+
+        Args:
+            key: State key
+            default: Default value if key not found
+
+        Returns:
+            State value or default
+        """
+        return self._states.get(key, default)
+
+    def has_state(self, key: str) -> bool:
+        """Check if state exists.
+
+        Args:
+            key: State key
+
+        Returns:
+            True if state exists
+        """
+        return key in self._states
+
+    # =========================================================================
+    # Abstract Methods
+    # =========================================================================
+
     @abstractmethod
     def bind_primitive(
         self, primitive: Primitive, args: tuple[Any, ...], kwargs: dict[str, Any]
@@ -80,6 +141,10 @@ class Context(ABC):
             Object in the context's native type (TraceObject or InterpObject)
         """
 
+    # =========================================================================
+    # Context Manager
+    # =========================================================================
+
     def __enter__(self) -> Self:
         """Enter context manager (push context onto stack)."""
         push_context(self)
@@ -90,9 +155,31 @@ class Context(ABC):
         pop_context()
 
 
-# ============================================================================
+# =============================================================================
+# Abstract Interpreter Interface
+# =============================================================================
+
+
+class AbstractInterpreter(Context):
+    """Abstract interface for Interpreters.
+
+    This allows EDSL components (like JIT) to depend on the Interpreter interface
+    without depending on the concrete Runtime implementation (which may depend on
+    ObjectStore, Backends, etc.).
+    """
+
+    @abstractmethod
+    def evaluate_graph(self, graph: Graph, inputs: list[Any]) -> Any:
+        """Execute a Graph IR with given inputs."""
+
+    @abstractmethod
+    def lift(self, obj: Any) -> Any:
+        """Lift a python object to an interpreter object."""
+
+
+# =============================================================================
 # Global Context Stack Management
-# ============================================================================
+# =============================================================================
 
 _context_stack: list[Context] = []
 _default_context: Context | None = None
@@ -100,20 +187,25 @@ _default_context_factory: Callable[[], Context] | None = None
 
 
 def get_current_context() -> Context | None:
-    """Get the current active context.
+    """Get the current active context (top of stack).
 
-    Returns None if no context is active (will use default context).
+    Returns None if no context is active.
     """
-
     return _context_stack[-1] if _context_stack else None
 
 
-def get_root_context() -> Context | None:
-    """Get the root context (bottom of the stack).
+def push_context(context: Context) -> None:
+    """Push a context onto the stack (enter context)."""
+    _context_stack.append(context)
+
+
+def pop_context() -> Context | None:
+    """Pop a context from the stack (exit context).
 
-    This context typically holds the global environment state (e.g. ClusterSpec).
+    Returns:
+        The popped context, or None if stack was empty.
     """
-    return _context_stack[0] if _context_stack else None
+    return _context_stack.pop() if _context_stack else None
 
 
 def find_context(predicate: Callable[[Context], bool]) -> Context | None:
@@ -122,9 +214,6 @@ def find_context(predicate: Callable[[Context], bool]) -> Context | None:
     Traverses from top (most recent) to bottom of the context stack,
     returning the first context for which predicate(ctx) returns True.
 
-    This is a general-purpose utility for finding contexts with specific
-    attributes or capabilities without hardcoding business logic here.
-
     Args:
         predicate: A callable that takes a Context and returns True if it matches.
 
@@ -133,13 +222,7 @@ def find_context(predicate: Callable[[Context], bool]) -> Context | None:
 
     Example:
         >>> # Find context with simp dialect state
-        >>> ctx = find_context(
-        ...     lambda c: hasattr(c, "get_dialect_state")
-        ...     and c.get_dialect_state("simp") is not None
-        ... )
-        >>>
-        >>> # Find context with cluster spec
-        >>> ctx = find_context(lambda c: getattr(c, "_cluster_spec", None) is not None)
+        >>> ctx = find_context(lambda c: c.has_state("dialect.simp"))
     """
     for ctx in reversed(_context_stack):
         if predicate(ctx):
@@ -147,15 +230,41 @@ def find_context(predicate: Callable[[Context], bool]) -> Context | None:
     return None
 
 
-def push_context(context: Context) -> None:
-    """Push a context onto the stack (enter context)."""
-    _context_stack.append(context)
+def find_context_with_state(key: str) -> Context | None:
+    """Find first context that has the specified state.
+
+    Args:
+        key: State key to look for
+
+    Returns:
+        First context with the state, or None
+    """
+    return find_context(lambda c: c.has_state(key))
+
+
+def find_interpreter() -> Context | None:
+    """Find first Interpreter in the context stack.
+
+    Returns:
+        First Interpreter context, or None if not found.
+    """
+    return find_context(lambda c: isinstance(c, AbstractInterpreter))
 
 
-def pop_context() -> None:
-    """Pop a context from the stack (exit context)."""
-    if _context_stack:
-        _context_stack.pop()
+def is_tracing() -> bool:
+    """Check if current context is a Tracer.
+
+    Returns:
+        True if the top of the context stack is a Tracer.
+    """
+    from mplang.v2.edsl.tracer import Tracer
+
+    return isinstance(get_current_context(), Tracer)
+
+
+# =============================================================================
+# Default Context Management
+# =============================================================================
 
 
 def register_default_context_factory(factory: Callable[[], Context]) -> None:
@@ -177,18 +286,26 @@ def get_default_context() -> Context:
     return _default_context
 
 
-class AbstractInterpreter(Context):
-    """Abstract interface for Interpreters.
+def set_root_context(context: Context, force: bool = False) -> None:
+    """Set the root/default execution context.
 
-    This allows EDSL components (like JIT) to depend on the Interpreter interface
-    without depending on the concrete Runtime implementation (which may depend on
-    ObjectStore, Backends, etc.).
-    """
-
-    @abstractmethod
-    def evaluate_graph(self, graph: Graph, inputs: list[Any]) -> Any:
-        """Execute a Graph IR with given inputs."""
+    This sets the provided context as the base of the context stack.
+    All subsequent operations will use this context as the default environment.
 
-    @abstractmethod
-    def lift(self, obj: Any) -> Any:
-        """Lift a python object to an interpreter object."""
+    Args:
+        context: Context to set as root.
+        force: If True, clears the existing context stack before setting.
+               If False (default), raises error if stack is not empty.
+    """
+    if force:
+        _context_stack.clear()
+        _context_stack.append(context)
+        return
+
+    if get_current_context() is not None:
+        raise RuntimeError(
+            "Cannot set root context: Context stack is not empty. "
+            "Use force=True to overwrite the existing context."
+        )
+
+    push_context(context)

mplang/v2/libs/device/api.py

@@ -732,25 +732,23 @@ def fetch(obj: Object) -> Any:
     """
     from mplang.v2.backends.simp_driver.state import SimpDriver
     from mplang.v2.backends.simp_driver.values import DriverVar
-    from mplang.v2.backends.table_impl import TableValue
-    from mplang.v2.backends.tensor_impl import TensorValue
     from mplang.v2.edsl.context import get_current_context
     from mplang.v2.runtime.interpreter import InterpObject, Interpreter
+    from mplang.v2.runtime.value import WrapValue
 
     def _unwrap_value(val: Any) -> Any:
-        """Unwrap Value types to get the underlying data."""
-        if isinstance(val, TensorValue):
-            return val.data
-        elif isinstance(val, TableValue):
+        """Unwrap WrapValue to get the underlying data."""
+        if isinstance(val, WrapValue):
             return val.data
         return val
 
-    # Get device info
+    # 1. Ensure is object and is device obj
     if not is_device_obj(obj):
         raise DeviceError(
             "Object does not have device attribute. Use mp.fetch() directly."
         )
 
+    # 2. Get device information according to device id
     dev_id = get_dev_attr(obj)
     cluster = _resolve_cluster()
     dev_info = cluster.devices[dev_id]
@@ -760,9 +758,10 @@ def fetch(obj: Object) -> Any:
     if not isinstance(ctx, Interpreter):
         raise RuntimeError("No interpreter context available for fetch")
 
-    simp_state = cast(SimpDriver | None, ctx.get_dialect_state("simp"))
+    simp_state = ctx.get_dialect_state("simp")
+    assert isinstance(simp_state, SimpDriver), "DriverVar requires simp state"
 
-    # Unwrap InterpObject
+    # Unwrap InterpObject to get runtime value
     assert isinstance(obj, InterpObject), f"Expected InterpObject, got {type(obj)}"
     runtime_obj = obj.runtime_obj
 
@@ -770,19 +769,26 @@ def fetch(obj: Object) -> Any:
         """Fetch value from a rank (DriverVar values are always URIs)."""
         uri = runtime_obj.values[rank]
         assert isinstance(uri, str) and "://" in uri, f"Expected URI, got {uri}"
-        assert simp_state is not None, "No simp state for fetch"
-        return _unwrap_value(simp_state.fetch(rank, uri).result())
+        return simp_state.fetch(rank, uri).result()
 
-    # Handle DriverVar
+    # 3. Match device type and do corresponding fetch action
     if isinstance(runtime_obj, DriverVar):
-        # For PPU/TEE: single member
+        # 3.1 PPU/TEE: single member, fetch directly
         if dev_info.kind.upper() in ("PPU", "TEE"):
             assert len(dev_info.members) == 1
-            return _fetch_from_rank(dev_info.members[0].rank)
+            result = _fetch_from_rank(dev_info.members[0].rank)
+            # 4. Unwrap if WrapValue
+            return _unwrap_value(result)
 
-        # For SPU: fetch from first member (should be revealed first)
+        # 3.2 SPU: fetch from all ranks and reconstruct
         elif dev_info.kind.upper() == "SPU":
-            return _fetch_from_rank(dev_info.members[0].rank)
-
-    # Direct value
+            # Fetch shares from all SPU members
+            shares = [_fetch_from_rank(m.rank) for m in dev_info.members]
+            # For now, just return the first share (TODO: implement spu.reconstruct)
+            # In practice, SPU values should be revealed to a PPU first
+            result = shares[0] if shares else None
+            # 4. Unwrap if WrapValue
+            return _unwrap_value(result)
+
+    # Direct value (not DriverVar)
     return _unwrap_value(runtime_obj)

mplang/v2/libs/ml/__init__.py

@@ -0,0 +1,23 @@
+# 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.
+
+"""Machine Learning algorithms for secure multi-party computation."""
+
+from mplang.v2.libs.ml.sgb import SecureBoost, Tree, TreeEnsemble
+
+__all__ = [
+    "SecureBoost",
+    "Tree",
+    "TreeEnsemble",
+]