mplang-nightly 0.1.dev155__py3-none-any.whl → 0.1.dev156__py3-none-any.whl

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/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/kernels/crypto.py

@@ -27,15 +27,19 @@ __all__: list[str] = []  # flat kernels only
 
 
 def _get_rng() -> np.random.Generator:
-    """Get (and lazily create) per-rank RNG for crypto kernels."""
+    """Get (and lazily create) per-rank RNG for crypto kernels.
+
+    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/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/kernels/spu.py

@@ -63,14 +63,10 @@ class SpuValue:
         return f"SpuValue({self.shape},{self.dtype},{self.vtype})"
 
 
-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)
@@ -84,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 = (
@@ -102,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")
@@ -197,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/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/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.dev155.dist-info → mplang_nightly-0.1.dev156.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: mplang-nightly
-Version: 0.1.dev155
+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.dev155.dist-info → mplang_nightly-0.1.dev156.dist-info}/RECORD

@@ -27,15 +27,15 @@ mplang/core/expr/utils.py,sha256=VDTJ_-CsdHtVy9wDaGa7XdFxQ7o5lYYaeqcgsAhkbNI,262
 mplang/core/expr/visitor.py,sha256=2Ge-I5N-wH8VVXy8d2WyNaEv8x6seiRx9peyH9S2BYU,2044
 mplang/core/expr/walk.py,sha256=lXkGJEEuvKGDqQihbxXPxfz2RfR1Q1zYUlt11iooQW0,11889
 mplang/kernels/__init__.py,sha256=2WE4cmW96Xkzyq2yRRYNww4kZ5o6u6NbPV0BxqZG698,581
-mplang/kernels/base.py,sha256=eizxj16sWkUvBvXWS0Zl-S9uuqalJmNUzB1xLhBg8S8,4920
+mplang/kernels/base.py,sha256=-YV4Aj5fs6GT4ehS6Tyi8WQ-amxn5edHTFJRQzyjHXY,3826
 mplang/kernels/builtin.py,sha256=nSuM79cn7M6M27A6Y8ycilXT_qAlB1ktkwkRX6dv_VQ,7052
-mplang/kernels/context.py,sha256=EiqsYq9Cw4Z_2lLvRNc8xNsu5EHkG_738QKFJ3Wf0zE,11329
-mplang/kernels/crypto.py,sha256=pENGeF869VmPszhCOhvyLSFhAeSnpj6gAtp4ofO08N0,3737
-mplang/kernels/mock_tee.py,sha256=ifVqb03FgQ5EDK8r14PMpZKMMaSeqMDorcpGLiy00mM,2233
+mplang/kernels/context.py,sha256=yJjQUyQmBBl6btLb7KlmA-Ejf4-cgK3KmtC-m0sBbb8,13594
+mplang/kernels/crypto.py,sha256=s7R0yd4Fk5cI2Qd3LpLc-kmbVuk8fFsbKbfKi43R0aE,3892
+mplang/kernels/mock_tee.py,sha256=173QSzPgkrLo0zn0jsx6nNmq1WvfUhIM67FM4Dn30aA,2297
 mplang/kernels/phe.py,sha256=8-_1IFPOaGECGj9mbYja8XoqbMYnYqfpDNVyMJd8J1Y,65247
-mplang/kernels/spu.py,sha256=-xJ1GJ5SmHSBg3tUsTzUujpkeVFXKDtaD1mTK-5MewM,9252
+mplang/kernels/spu.py,sha256=LkM8tNzhwTa8lufNgClHfnI4LNu25cdWLQZdJsMDEO8,9301
 mplang/kernels/sql_duckdb.py,sha256=UN1Ev6-MxF_-65zMExUsLScC9PlmEIEcN8YziIoX_rY,1724
-mplang/kernels/stablehlo.py,sha256=jDsu-lIHRAm94FcUcxZgK02c6BhFJpbO8cf2hP2DFgk,2937
+mplang/kernels/stablehlo.py,sha256=M25_1G332uFCUcIIjpcTMwIdSDW1TMjfhFWxbz5OBZ0,2966
 mplang/ops/__init__.py,sha256=dpe7WWiYapOFzJeGoKFYBr5mnd6P5SdOyvdYaM2Nhm0,1408
 mplang/ops/base.py,sha256=h67_SHWNZGUuTCuMll-9kDgGvlPhlFov7WAQCHTmUvw,18258
 mplang/ops/builtin.py,sha256=D7T8rRF9g05VIw9T72lsncF5cDQqaT37eapBieRKvRI,9363
@@ -53,7 +53,7 @@ mplang/runtime/__init__.py,sha256=IRPP3TtpFC4iSt7_uaq-S4dL7CwrXL0XBMeaBoEYLlg,94
 mplang/runtime/cli.py,sha256=WehDodeVB4AukSWx1LJxxtKUqGmLPY4qjayrPlOg3bE,14438
 mplang/runtime/client.py,sha256=vkJUFSDcKIdbKiGUM5AosCKTZygl9g8uZFEjw2xwKig,15249
 mplang/runtime/communicator.py,sha256=Lek6_h_Wmr_W-_JpT-vMxL3CHxcVZdtf7jdaLGuxPgQ,3199
-mplang/runtime/data_providers.py,sha256=hH2butEOYNGq2rRZjVBDfXLxe3YUin2ftAF6htbTfLA,8226
+mplang/runtime/data_providers.py,sha256=GX10_nch8PmEyok32mSC4p5rDowvmXrJ-4J5-LvY6ig,8206
 mplang/runtime/driver.py,sha256=pq2EQFZK9tH90Idops_yeF6fj0cfFVD_5mFcmy4Hzco,11089
 mplang/runtime/exceptions.py,sha256=c18U0xK20dRmgZo0ogTf5vXlkix9y3VAFuzkHxaXPEk,981
 mplang/runtime/http_api.md,sha256=-re1DhEqMplAkv_wnqEU-PSs8tTzf4-Ml0Gq0f3Go6s,4883
@@ -70,8 +70,8 @@ mplang/utils/crypto.py,sha256=rvPomBFtznRHc3RPi6Aip9lsU8zW2oxBqGv1K3vn7Rs,1052
 mplang/utils/func_utils.py,sha256=vCJcZmu0bEbqhOQKdpttV2_MBllIcPSN0b8U4WjNGGo,5164
 mplang/utils/spu_utils.py,sha256=S3L9RBkBe2AvSuMSQQ12cBY5Y1NPthubvErSX_7nj1A,4158
 mplang/utils/table_utils.py,sha256=aC-IZOKkSmFkpr3NZchLM0Wt0GOn-rg_xHBHREWBwAU,2202
-mplang_nightly-0.1.dev155.dist-info/METADATA,sha256=26YKbEAbRU0dK9hZ8xFftRx5UDw_Ij5th2CySJgVlFo,16547
-mplang_nightly-0.1.dev155.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-mplang_nightly-0.1.dev155.dist-info/entry_points.txt,sha256=mG1oJT-GAjQR834a62_QIWb7litzWPPyVnwFqm-rWuY,55
-mplang_nightly-0.1.dev155.dist-info/licenses/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
-mplang_nightly-0.1.dev155.dist-info/RECORD,,
+mplang_nightly-0.1.dev156.dist-info/METADATA,sha256=iaTA5SM0ALRdvM5h40B44CV8WMxsMQXP6DqsPrqeCY4,16547
+mplang_nightly-0.1.dev156.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+mplang_nightly-0.1.dev156.dist-info/entry_points.txt,sha256=mG1oJT-GAjQR834a62_QIWb7litzWPPyVnwFqm-rWuY,55
+mplang_nightly-0.1.dev156.dist-info/licenses/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
+mplang_nightly-0.1.dev156.dist-info/RECORD,,