mplang-nightly 0.1.dev154__tar.gz → 0.1.dev156__tar.gz

{mplang_nightly-0.1.dev154 → mplang_nightly-0.1.dev156}/PKG-INFO

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

{mplang_nightly-0.1.dev154 → mplang_nightly-0.1.dev156}/mplang/core/cluster.py

@@ -255,7 +255,7 @@ class ClusterSpec:
             Optional explicit endpoint list of length ``world_size``. Each element may
             include scheme (``http://``) or not; stored verbatim. If not provided we
             synthesize ``localhost:{5000 + i}`` (5000 is a fixed default; pass explicit
-            endpoints for control). Deprecated ``base_port`` legacy kwarg can adjust it.
+            endpoints for control).
         spu_protocol / spu_field:
             SPU device config values.
         runtime_version / runtime_platform:

{mplang_nightly-0.1.dev154 → mplang_nightly-0.1.dev156}/mplang/core/expr/evaluator.py

@@ -66,7 +66,6 @@ class EvalSemantic:
     """Shared evaluation semantics and utilities for evaluators.
 
     Minimal dataclass carrying runtime execution context (rank/env/comm/runtime).
-    Legacy handler-based execution (pfunc_handles) has been fully removed.
     """
 
     rank: int

{mplang_nightly-0.1.dev154 → mplang_nightly-0.1.dev156}/mplang/core/mptype.py

@@ -400,39 +400,5 @@ class MPType:
                 "Table object detection for non-pandas objects not fully implemented yet"
             )
 
-        # Check if it's a table-like object (legacy check for backward compatibility)
-        if hasattr(obj, "dtypes") and hasattr(obj, "columns"):
-            # Basic pandas DataFrame support
-            try:
-                import pandas as pd
-
-                if isinstance(obj, pd.DataFrame):
-                    from mplang.core.dtype import DType
-
-                    schema_dict = {}
-                    for col_name in obj.columns:
-                        pandas_dtype = obj[col_name].dtype
-                        # Convert pandas dtype to DType
-                        if pandas_dtype.kind in (
-                            "O",
-                            "U",
-                            "S",
-                        ):  # object, unicode, string
-                            schema_dict[col_name] = (
-                                DType.from_numpy(pandas_dtype)
-                                if pandas_dtype.kind != "O"
-                                else STRING
-                            )
-                        else:
-                            schema_dict[col_name] = DType.from_numpy(pandas_dtype)
-                    schema = TableType.from_dict(schema_dict)
-                    return cls(schema, pmask, attrs)
-            except ImportError:
-                pass
-            # For other table-like objects without pandas
-            raise NotImplementedError(
-                "Table object detection not fully implemented yet"
-            )
-
         # Otherwise treat as tensor-like
         return cls.from_tensor(obj, pmask, **attrs)

{mplang_nightly-0.1.dev154 → mplang_nightly-0.1.dev156}/mplang/kernels/base.py

@@ -34,7 +34,10 @@ from __future__ import annotations
 import contextvars
 from collections.abc import Callable
 from dataclasses import dataclass
-from typing import Any
+from typing import TYPE_CHECKING, Any
+
+if TYPE_CHECKING:
+    from mplang.kernels.context import RuntimeContext
 
 __all__ = [
     "KernelContext",
@@ -48,12 +51,15 @@ __all__ = [
 
 @dataclass
 class KernelContext:
-    """Ephemeral call context set via contextvar while a kernel runs."""
+    """Ephemeral per-kernel invocation context.
+
+    Cross-kernel persistent state (RNGs, compiled artifacts, environment handles)
+    should be stored in RuntimeContext.
+    """
 
     rank: int
     world_size: int
-    state: dict[str, dict[str, Any]]  # backend namespace -> pocket
-    cache: dict[str, Any]  # runtime-level shared cache (per BackendRuntime)
+    runtime: RuntimeContext
 
 
 _CTX_VAR: contextvars.ContextVar[KernelContext | None] = contextvars.ContextVar(
@@ -62,37 +68,7 @@ _CTX_VAR: contextvars.ContextVar[KernelContext | None] = contextvars.ContextVar(
 
 
 def cur_kctx() -> KernelContext:
-    """Return the current kernel execution context (only valid inside a kernel).
-
-    Two storages:
-      - state: namespaced pockets (dict[str, dict]) for backend-local mutable helpers
-      - cache: global (per runtime) shared dict; prefer state unless truly cross-backend
-
-    Examples:
-      1) Compile cache::
-            @kernel_def("mlir.stablehlo")
-            def _exec(pfunc, args):
-                ctx = cur_kctx()
-                pocket = ctx.state.setdefault("stablehlo", {})
-                cache = pocket.setdefault("compile_cache", {})
-                text = pfunc.fn_text
-                mod = cache.get(text)
-                if mod is None:
-                    mod = compile_mlir(text)
-                    cache[text] = mod
-                return run(mod, args)
-
-      2) Deterministic RNG::
-            @kernel_def("crypto.keygen")
-            def _keygen(pfunc, args):
-                ctx = cur_kctx()
-                pocket = ctx.state.setdefault("crypto", {})
-                rng = pocket.get("rng")
-                if rng is None:
-                    rng = np.random.default_rng(1234 + ctx.rank * 7919)
-                    pocket["rng"] = rng
-                return (rng.integers(0, 256, size=(32,), dtype=np.uint8),)
-    """
+    """Return current kernel execution context (only valid inside kernel)."""
     ctx = _CTX_VAR.get()
     if ctx is None:
         raise RuntimeError("cur_kctx() called outside backend kernel execution")

{mplang_nightly-0.1.dev154 → mplang_nightly-0.1.dev156}/mplang/kernels/context.py

@@ -118,12 +118,21 @@ class RuntimeContext:
         op_type -> kernel_id and form a *template* for dispatch. After
         initialization, all (re)binding must go through ``bind_op`` /
         ``rebind_op`` on this context (scoped to THIS runtime only).
-    state / cache / stats : dict, optional
-        Mutable pockets reused across kernel invocations. If omitted, new
-        dictionaries are created.
+    state : dict, optional
+        Mutable per-runtime key/value storage for kernels. Flat key space;
+        callers SHOULD use dotted prefixes (e.g. "stablehlo.compile_cache").
+        Kernels own their *state* (functional correctness data, caches,
+        handles, compiled objects, RNGs, etc.). Runtime does not interpret
+        structure—values may themselves be dicts if a kernel wants its own
+        pocket. Created empty when omitted.
+    stats : dict, optional
+        Mutable statistics/telemetry owned by the runtime (usage counters,
+        timings, profiling aids). Kernels may increment counters but should
+        avoid storing functional state here. A default "op_calls" mapping is
+        ensured. Created empty when omitted.
     """
 
-    __slots__ = ("_ibindings", "cache", "rank", "state", "stats", "world_size")
+    __slots__ = ("_ibindings", "rank", "state", "stats", "world_size")
 
     def __init__(
         self,
@@ -131,8 +140,7 @@ class RuntimeContext:
         world_size: int,
         initial_bindings: Mapping[str, str] | None = None,
         *,
-        state: dict[str, dict[str, Any]] | None = None,
-        cache: dict[str, Any] | None = None,
+        state: dict[str, Any] | None = None,
         stats: dict[str, Any] | None = None,
     ) -> None:
         _ensure_impl_imported()
@@ -144,7 +152,6 @@ class RuntimeContext:
             **(initial_bindings or {}),
         }
         self.state = state if state is not None else {}
-        self.cache = cache if cache is not None else {}
         self.stats = stats if stats is not None else {}
         self.stats.setdefault("op_calls", {})
 
@@ -168,19 +175,15 @@ class RuntimeContext:
             if isinstance(ins_spec, TensorType):
                 _validate_tensor_arg(fn_type, idx, ins_spec, val)
                 continue
+
         # install kernel context
-        kctx = KernelContext(
-            rank=self.rank,
-            world_size=self.world_size,
-            state=self.state,
-            cache=self.cache,
-        )
-        token = base._CTX_VAR.set(kctx)  # type: ignore[attr-defined]
+        kctx = KernelContext(rank=self.rank, world_size=self.world_size, runtime=self)
+        token = base._CTX_VAR.set(kctx)
         try:
             raw = fn(pfunc, *arg_list)
         finally:
-            base._CTX_VAR.reset(token)  # type: ignore[attr-defined]
-        # Stats (best effort)
+            base._CTX_VAR.reset(token)
+
         try:
             op_calls = self.stats.setdefault("op_calls", {})
             op_calls[fn_type] = op_calls.get(fn_type, 0) + 1
@@ -213,7 +216,57 @@ class RuntimeContext:
 
     def reset(self) -> None:
         self.state.clear()
-        self.cache.clear()
+
+    # ---- runtime state API (flat key space) ----
+    # Keys are treated atomically; convention encourages dotted prefixes
+    # (e.g. 'stablehlo.compile_cache.hash', 'crypto.rng'). Implementation
+    # does NOT parse or create hierarchical dicts—any grouping is purely
+    # by string prefix. Values themselves MAY be dicts if callers want a
+    # manual pocket. This keeps semantics simple and predictable.
+
+    def ensure_state(self, key: str, factory: type | Any = dict) -> Any:
+        """Return value for key; if absent create via factory and store.
+
+        Key is not parsed; dotted forms are allowed but treated as a single
+        map key. Use consistent prefixes for grouping (e.g. 'spu.config').
+        """
+        if not key:
+            raise ValueError("empty state key")
+        val = self.state.get(key)
+        if val is None:
+            val = factory()
+            self.state[key] = val
+        return val
+
+    def get_state(self, key: str, default: Any | None = None) -> Any:
+        if not key:
+            raise ValueError("empty state key")
+        return self.state.get(key, default)
+
+    def set_state(self, key: str, value: Any) -> None:
+        if not key:
+            raise ValueError("empty state key")
+        self.state[key] = value
+
+    def del_state(self, key: str) -> None:
+        if not key:
+            raise ValueError("empty state key")
+        self.state.pop(key, None)
+
+    def list_state(self, prefix: str = "") -> dict[str, Any]:
+        """Return mapping of key -> value; optional prefix filter.
+
+        Prefix match is string-based; if prefix is non-empty include keys
+        where key == prefix or key starts with prefix + '.'.
+        """
+        if not prefix:
+            return dict(self.state)
+        pref = prefix if prefix.endswith(".") else prefix + "."
+        out: dict[str, Any] = {}
+        for k, v in self.state.items():
+            if k == prefix or k.startswith(pref):
+                out[k] = v
+        return out
 
     # ---- explicit (re)binding API ----
     def bind_op(self, op_type: str, kernel_id: str, *, force: bool = False) -> None:

{mplang_nightly-0.1.dev154 → mplang_nightly-0.1.dev156}/mplang/kernels/crypto.py

@@ -29,16 +29,17 @@ __all__: list[str] = []  # flat kernels only
 def _get_rng() -> np.random.Generator:
     """Get (and lazily create) per-rank RNG for crypto kernels.
 
-    Seed rule matches legacy handler: MPLANG_CRYPTO_SEED + rank*7919
+    Runtime state is untyped, so we narrow the type explicitly for mypy.
     """
     kctx = cur_kctx()
-    pocket = kctx.state.setdefault("crypto", {})
-    rng = pocket.get("rng")
-    if rng is None:
+    rt = kctx.runtime
+    rng_obj = rt.get_state("crypto.rng")
+    if rng_obj is None:
         seed = int(os.environ.get("MPLANG_CRYPTO_SEED", "0")) + kctx.rank * 7919
-        rng = np.random.default_rng(seed)
-        pocket["rng"] = rng
-    return rng
+        rng_obj = np.random.default_rng(seed)
+        rt.set_state("crypto.rng", rng_obj)
+    assert isinstance(rng_obj, np.random.Generator)  # narrow
+    return rng_obj
 
 
 def _keystream(key: bytes, nonce: bytes, length: int) -> bytes:

{mplang_nightly-0.1.dev154 → mplang_nightly-0.1.dev156}/mplang/kernels/mock_tee.py

@@ -28,12 +28,13 @@ __all__: list[str] = []
 
 def _rng() -> np.random.Generator:
     kctx = cur_kctx()
-    pocket = kctx.state.setdefault("tee", {})
-    r = pocket.get("rng")
+    rt = kctx.runtime
+    r = rt.get_state("tee.rng")
     if r is None:
         seed = int(os.environ.get("MPLANG_TEE_SEED", "0")) + kctx.rank * 10007
         r = np.random.default_rng(seed)
-        pocket["rng"] = r
+        rt.set_state("tee.rng", r)
+    assert isinstance(r, np.random.Generator)  # type narrowing for mypy
     return r
 
 

{mplang_nightly-0.1.dev154 → mplang_nightly-0.1.dev156}/mplang/kernels/spu.py

@@ -63,17 +63,10 @@ class SpuValue:
         return f"SpuValue({self.shape},{self.dtype},{self.vtype})"
 
 
-# SpuHandler removed (legacy handler API deprecated)
-
-
-def _get_spu_pocket() -> dict[str, Any]:
-    return cur_kctx().state.setdefault("spu", {})
-
-
 def _get_spu_config_and_world() -> tuple[libspu.RuntimeConfig, int]:
-    pocket = _get_spu_pocket()
-    cfg = pocket.get("config")
-    world = pocket.get("world")
+    kctx = cur_kctx()
+    cfg = kctx.runtime.get_state("spu.config")
+    world = kctx.runtime.get_state("spu.world")
     if cfg is None or world is None:
         raise RuntimeError("SPU kernel state not initialized (config/world)")
     return cfg, int(world)
@@ -87,12 +80,12 @@ def _register_spu_env(
     Idempotent: if config/world already set, they must match; link is recorded per rank.
     This replaces previous global fallback seeding logic.
     """
-    pocket = _get_spu_pocket()
-    prev_cfg = pocket.get("config")
-    prev_world = pocket.get("world")
+    kctx = cur_kctx()
+    prev_cfg = kctx.runtime.get_state("spu.config")
+    prev_world = kctx.runtime.get_state("spu.world")
     if prev_cfg is None:
-        pocket["config"] = config
-        pocket["world"] = world_size
+        kctx.runtime.set_state("spu.config", config)
+        kctx.runtime.set_state("spu.world", world_size)
     else:
         # libspu RuntimeConfig may not implement __eq__; compare serialized repr
         same_cfg = (
@@ -105,7 +98,7 @@ def _register_spu_env(
             raise RuntimeError("Conflicting SPU env registration")
     # Store single link per runtime (one runtime per rank)
     if link_ctx is not None:
-        pocket["link"] = link_ctx
+        kctx.runtime.set_state("spu.link", link_ctx)
 
 
 @kernel_def("spu.seed_env")
@@ -200,16 +193,16 @@ def _spu_run_mlir(pfunc: PFunction, *args: Any) -> Any:
         )
 
     cfg, _ = _get_spu_config_and_world()
-    pocket = _get_spu_pocket()
-    link_ctx: LinkCommunicator | None = pocket.get("link")
+    kctx = cur_kctx()
+    link_ctx = kctx.runtime.get_state("spu.link")
     if link_ctx is None:
         raise RuntimeError("Rank not participating in SPU; no link set via seed_env")
 
-    # Lazy runtime cache
-    spu_rt = pocket.get("runtime")
+    # Lazy runtime cache under key spu.runtime
+    spu_rt = kctx.runtime.get_state("spu.runtime")
     if spu_rt is None:
         spu_rt = spu_api.Runtime(link_ctx.get_lctx(), cfg)
-        pocket["runtime"] = spu_rt
+        kctx.runtime.set_state("spu.runtime", spu_rt)
 
     # Validate that all inputs are SpuValue objects
     for i, arg in enumerate(args):

{mplang_nightly-0.1.dev154 → mplang_nightly-0.1.dev156}/mplang/kernels/stablehlo.py

@@ -36,11 +36,14 @@ def _stablehlo_exec(pfunc: PFunction, *args: Any) -> Any:
     if isinstance(mlir_text, bytes):
         mlir_text = mlir_text.decode("utf-8")
 
-    # Simple compile cache per runtime (state pocket per backend namespace)
+    # Flat-key compile cache: stablehlo.compile_cache.<hash>
     ctx = cur_kctx()
-    pocket = ctx.state.setdefault("stablehlo", {})
-    cache = pocket.setdefault("compile_cache", {})
-    compiled = cache.get(mlir_text)
+    rt = ctx.runtime
+    import hashlib
+
+    h = hashlib.sha256(mlir_text.encode("utf-8")).hexdigest()[:16]
+    key = f"stablehlo.compile_cache.{h}"
+    compiled = rt.get_state(key)
     if compiled is None:
         backend = jax.default_backend()
         client = xla_bridge.get_backend(backend)
@@ -49,7 +52,7 @@ def _stablehlo_exec(pfunc: PFunction, *args: Any) -> Any:
             compiled = client.compile(mlir_text, compile_options)
         except Exception as e:  # pragma: no cover
             raise RuntimeError(f"StableHLO compile failed: {e}") from e
-        cache[mlir_text] = compiled
+        rt.set_state(key, compiled)
 
     # Handle JAX's unused parameter elimination via arg_keep_map
     runtime_args = args

{mplang_nightly-0.1.dev154 → mplang_nightly-0.1.dev156}/mplang/runtime/data_providers.py

@@ -173,38 +173,32 @@ class FileProvider(DataProvider):
         np.save(path, np.asarray(value))
 
 
-class _KeyedPocket:
-    """Small helper to keep a dict in KernelContext.state under a namespaced key."""
-
-    def __init__(self, ns: str):
-        self.ns = ns
-
-    def get_map(self, ctx: KernelContext) -> dict[str, Any]:
-        pocket = ctx.state.setdefault("resource.providers", {})
-        store = pocket.get(self.ns)
-        if store is None:
-            store = {}
-            pocket[self.ns] = store
-        return store  # type: ignore[return-value]
-
-
 class MemProvider(DataProvider):
     """In-memory per-runtime KV provider (per rank, per session/runtime)."""
 
-    def __init__(self) -> None:
-        self._pocket = _KeyedPocket("mem")
+    STATE_KEY = "resource.providers.mem"
+
+    @staticmethod
+    def _store(ctx: KernelContext) -> dict[str, Any]:
+        # Use ensure_state so creation is atomic & centralized; enforce dict.
+        store = ctx.runtime.ensure_state(MemProvider.STATE_KEY, dict)
+        if not isinstance(store, dict):  # pragma: no cover - defensive
+            raise TypeError(
+                f"runtime state key '{MemProvider.STATE_KEY}' expected dict, got {type(store).__name__}"
+            )
+        return store  # type: ignore[return-value]
 
     def read(
         self, uri: ResolvedURI, out_spec: TensorType | TableType, *, ctx: KernelContext
     ) -> Any:
-        store = self._pocket.get_map(ctx)
+        store = self._store(ctx)
         key = uri.raw
         if key not in store:
             raise FileNotFoundError(f"mem resource not found: {key}")
         return store[key]
 
     def write(self, uri: ResolvedURI, value: Any, *, ctx: KernelContext) -> None:
-        store = self._pocket.get_map(ctx)
+        store = self._store(ctx)
         store[uri.raw] = value
 
 

{mplang_nightly-0.1.dev154 → mplang_nightly-0.1.dev156}/mplang/runtime/session.py

@@ -184,8 +184,7 @@ class Session:
             return
 
         link_ctx = None
-        # Fixed port offset for SPU runtime link services (legacy value retained).
-        # TODO: make configurable if future deployments require dynamic offset.
+        # TODO(jint): reuse same port for mplang and spu.
         SPU_PORT_OFFSET = 100
 
         if self.is_spu_party:

{mplang_nightly-0.1.dev154 → mplang_nightly-0.1.dev156}/tests/kernels/test_builtin.py

@@ -23,9 +23,7 @@ from mplang.core.tensor import TensorType
 from mplang.kernels.context import RuntimeContext
 
 
-class TestBuiltinHandler:
-    """Tests for flat builtin kernels (legacy name retained)."""
-
+class TestBuiltin:
     def setup_method(self):
         # initialize backend context for rank 0 (world_size=1) once per test
         self.runtime = RuntimeContext(rank=0, world_size=1)

{mplang_nightly-0.1.dev154 → mplang_nightly-0.1.dev156}/tests/kernels/test_stablehlo.py

@@ -155,5 +155,3 @@ class TestStablehloKernel:
         ev = create_evaluator(0, {}, comm, runtime)
         with pytest.raises(NotImplementedError):
             ev._exec_pfunc(invalid_pfunc, [])  # type: ignore[attr-defined]
-
-        # No legacy handler path anymore; only NotImplementedError is expected.