fleetmesh-node 0.2.0__py3-none-any.whl

node/planes/services/compute/executors/wasm.py

@@ -0,0 +1,244 @@
+"""fleetmesh-wasi-v1: the WASI preview1 runner on wasmtime-py, and WasmExecutor.
+
+wasmtime is optional. Everything the deterministic profile fixes (clock, random,
+NaN canonicalisation, no filesystem, no network) is set up in run_wasi_module.
+"""
+from __future__ import annotations
+
+import hashlib
+import os
+from typing import Any, Mapping, Sequence
+
+from fleetmesh_core.crypto import crypto
+from ..common import ComputeError, resolve_blob
+from ..sandbox import ComputeLimits, NODE_MAX_MEMORY_BYTES
+from .base import ComputeExecutionResult, ComputeExecutor, ComputeQuote
+
+# -------------------------------------------------------------------------
+# fleetmesh-wasi-v1 runtime (wasmtime-py)
+# -------------------------------------------------------------------------
+#
+# Profile: WASI preview1, stdin = job input, stdout = result, no preopened
+# directories (no filesystem), no environment, no network (preview1 has no
+# sockets), clock_time_get / clock_res_get fixed at WASM_VIRTUAL_EPOCH_NS,
+# random_get a SHA-256 stream seeded from the job nonce, NaNs canonicalised,
+# threads/SIMD-relaxed off. Two nodes running the same module on the same
+# input and nonce must produce byte-identical output and fuel.
+#
+# wasmtime is optional: a node without it refuses WASM jobs with a clear
+# error instead of returning a fake result.
+
+WASM_DEFAULT_MEMORY_BYTES: int = 256 * 1024 * 1024
+WASM_MAX_OUTPUT_BYTES: int = 16 * 1024 * 1024
+WASM_VIRTUAL_EPOCH_NS: int = 1_700_000_000 * 1_000_000_000
+
+_WASM_ENGINE: Any = None
+
+
+def wasmtime_available() -> bool:
+    # importlib.util.find_spec rather than a bare import: the module object
+    # is not needed here, and pyflakes cannot see through a noqa.
+    import importlib.util
+    return importlib.util.find_spec('wasmtime') is not None
+
+
+def _wasm_engine() -> Any:
+    global _WASM_ENGINE
+    if _WASM_ENGINE is None:
+        import wasmtime
+        cfg = wasmtime.Config()
+        cfg.consume_fuel = True
+        cfg.cranelift_nan_canonicalization = True
+        cfg.wasm_threads = False
+        cfg.wasm_relaxed_simd = False
+        cfg.parallel_compilation = False
+        # Default wasmtime reserves 4 GiB+guard of address space per memory,
+        # which cannot be mapped under the sandbox's RLIMIT_AS. Reserve
+        # nothing up front and let memories move on growth instead.
+        cfg.memory_reservation = 0
+        cfg.memory_guard_size = 64 * 1024
+        cfg.memory_reservation_for_growth = 1024 * 1024
+        cfg.memory_may_move = True
+        if hasattr(type(cfg), 'macos_use_mach_ports'):
+            try:
+                cfg.macos_use_mach_ports = False
+            except Exception:
+                pass
+        _WASM_ENGINE = wasmtime.Engine(cfg)
+    return _WASM_ENGINE
+
+
+def run_wasi_module(
+    module_bytes: bytes,
+    input_bytes: bytes,
+    fuel_cap: int,
+    memory_bytes: int = WASM_DEFAULT_MEMORY_BYTES,
+    seed: str = '',
+    args: Sequence[str] = (),
+) -> tuple[bytes, int]:
+    """Run a WASI command module under fleetmesh-wasi-v1. Returns (stdout, fuel)."""
+    try:
+        import wasmtime
+    except ImportError as e:
+        raise ComputeError('WASM executor unavailable: install wasmtime (pip install wasmtime) [%s]' % e)
+    import tempfile
+
+    engine = _wasm_engine()
+    try:
+        module = wasmtime.Module(engine, module_bytes)
+    except wasmtime.WasmtimeError as e:
+        raise ComputeError('invalid WASM module: %s' % str(e).splitlines()[0])
+
+    store = wasmtime.Store(engine)
+    store.set_fuel(int(fuel_cap))
+    store.set_limits(memory_size=int(memory_bytes), instances=1, tables=4, memories=1)
+
+    out = bytearray()
+    overflow = [False]
+
+    def _stdout(chunk: bytes) -> None:
+        if len(out) + len(chunk) > WASM_MAX_OUTPUT_BYTES:
+            overflow[0] = True
+            return None
+        out.extend(chunk)
+        return None
+
+    with tempfile.TemporaryDirectory(prefix='fleetmesh-wasi-') as tmp:
+        in_path = os.path.join(tmp, 'stdin')
+        with open(in_path, 'wb') as f:
+            f.write(input_bytes)
+        wc = wasmtime.WasiConfig()
+        wc.argv = ['job.wasm', *args]
+        wc.stdin_file = in_path
+        wc.stdout_custom = _stdout
+        store.set_wasi(wc)
+
+        linker = wasmtime.Linker(engine)
+        linker.define_wasi()
+        linker.allow_shadowing = True
+        i32, i64 = wasmtime.ValType.i32(), wasmtime.ValType.i64()
+
+        def _clock_get(caller: Any, _cid: int, _prec: int, ptr: int) -> int:
+            caller['memory'].write(caller, WASM_VIRTUAL_EPOCH_NS.to_bytes(8, 'little'), ptr)
+            return 0
+
+        def _clock_res(caller: Any, _cid: int, ptr: int) -> int:
+            caller['memory'].write(caller, (1).to_bytes(8, 'little'), ptr)
+            return 0
+
+        rng_state = {'ctr': 0}
+        rng_key = hashlib.sha256(b'fleetmesh-wasi-v1/random:' + seed.encode('utf-8')).digest()
+
+        def _random(caller: Any, ptr: int, n: int) -> int:
+            buf = bytearray()
+            while len(buf) < n:
+                buf += hashlib.sha256(rng_key + rng_state['ctr'].to_bytes(8, 'big')).digest()
+                rng_state['ctr'] += 1
+            caller['memory'].write(caller, bytes(buf[:n]), ptr)
+            return 0
+
+        ns = 'wasi_snapshot_preview1'
+        linker.define_func(ns, 'clock_time_get', wasmtime.FuncType([i32, i64, i32], [i32]), _clock_get, access_caller=True)
+        linker.define_func(ns, 'clock_res_get', wasmtime.FuncType([i32, i32], [i32]), _clock_res, access_caller=True)
+        linker.define_func(ns, 'random_get', wasmtime.FuncType([i32, i32], [i32]), _random, access_caller=True)
+
+        try:
+            instance = linker.instantiate(store, module)
+            start = instance.exports(store).get('_start')
+            # A module may export '_start' as a memory/table/global to confuse
+            # a runner; only a real function is a WASI command entry point.
+            if not isinstance(start, wasmtime.Func):
+                raise ComputeError('WASM module has no _start export (not a WASI command)')
+            start(store)
+        except wasmtime.ExitTrap as e:
+            if e.code != 0:
+                raise ComputeError('WASM module exited with status %d' % e.code)
+        except wasmtime.Trap as e:
+            if e.trap_code == wasmtime.TrapCode.OUT_OF_FUEL:
+                raise ComputeError('job execution exceeded fuel ceiling (%d)' % fuel_cap)
+            raise ComputeError('WASM trap: %s' % str(e).splitlines()[0])
+        except wasmtime.WasmtimeError as e:
+            raise ComputeError('WASM error: %s' % str(e).splitlines()[0])
+
+    if overflow[0]:
+        raise ComputeError('WASM output exceeded %d bytes' % WASM_MAX_OUTPUT_BYTES)
+    return bytes(out), int(fuel_cap) - int(store.get_fuel())
+
+
+class WasmExecutor(ComputeExecutor):
+    """Deterministic WASI runtime driver (fleetmesh-wasi-v1 profile)."""
+
+    engine_id: str = 'wasmtime'
+    wasi_profile_digest: str = hashlib.sha256(
+        b'fleetmesh-wasi-v1:network=none:fs=memory:clock=virtual:prng=seeded'
+    ).hexdigest()
+
+    def is_live_runnable(self) -> bool:
+        """Check whether wasmtime engine is installed and available on this host."""
+        return wasmtime_available()
+
+    def quote(
+        self,
+        module_spec: Mapping[str, Any],
+        input_spec: Mapping[str, Any],
+        limits: ComputeLimits,
+    ) -> ComputeQuote:
+        mod_len = len(str(module_spec.get('sha256', '')))
+        inp_len = len(str(input_spec.get('sha256', '')))
+        fuel_est = min(25000 + (mod_len + inp_len) * 250, limits.max_fuel)
+        return ComputeQuote(
+            fuel_estimate=fuel_est,
+            rate_unit='msat_per_mfuel',
+            rate_amount=45,
+            engine_id=self.engine_id,
+        )
+
+    def execute(
+        self,
+        job_payload: Mapping[str, Any],
+        fuel_cap: int,
+    ) -> ComputeExecutionResult:
+        module_spec = job_payload.get('module', {})
+        input_spec = job_payload.get('input', {})
+
+        blob_store = job_payload.get('blob_store')
+        mod_bytes = resolve_blob(module_spec, blob_store=blob_store, blossom_servers=job_payload.get('blossom_servers'))
+        inp_bytes = resolve_blob(input_spec, blob_store=blob_store, blossom_servers=job_payload.get('blossom_servers'))
+        if mod_bytes is None:
+            raise ComputeError('WASM module not resolvable (sha256=%s)' % module_spec.get('sha256', '?'))
+        want = str(module_spec.get('sha256') or '').lower()
+        if want.startswith('baf'):
+            # A CID commits to the same digest; decode it rather than skip
+            # the check. A CID that does not decode is a refusal, not a pass.
+            from node.protocols.nostr import blobstore
+            try:
+                want = blobstore.cid_to_sha256(want)
+            except Exception as e:
+                raise ComputeError('WASM module CID does not decode: %s' % e)
+        if want and hashlib.sha256(mod_bytes).hexdigest() != want:
+            raise ComputeError('WASM module bytes do not match declared sha256')
+
+        # The requester may ask for less linear memory than the operator
+        # ceiling, never more: run_job passes the clamped ComputeLimits
+        # ceiling as '_memory_ceiling_bytes'; without it (direct execute)
+        # NODE_MAX_MEMORY_BYTES applies. A wasm cap above the rlimit would let
+        # the module get killed by the sandbox instead of trapping cleanly.
+        ceiling = int(job_payload.get('_memory_ceiling_bytes') or NODE_MAX_MEMORY_BYTES)
+        ceiling = max(1, min(ceiling, NODE_MAX_MEMORY_BYTES))
+        requested = int(job_payload.get('wasm_max_memory_bytes') or WASM_DEFAULT_MEMORY_BYTES)
+        memory_bytes = max(1, min(requested, ceiling))
+
+        raw_output, fuel_used = run_wasi_module(
+            mod_bytes, inp_bytes or b'', fuel_cap,
+            memory_bytes=memory_bytes,
+            seed=str(job_payload.get('nonce') or job_payload.get('ask') or ''),
+            args=[str(a) for a in (job_payload.get('args') or [])],
+        )
+        result_hash = crypto.sha256(raw_output)
+        return ComputeExecutionResult(
+            result_hash=result_hash,
+            raw_output=raw_output,
+            fuel_used=fuel_used,
+            engine_id=self.engine_id,
+            wasi_profile_digest=self.wasi_profile_digest,
+        )