fix(auth): reuse a process-wide wasmtime engine

    Create the wasmtime Engine lazily once per process and reuse it
    across OPAPolicy instances. This avoids repeatedly initializing JIT
    and trap-handling state as auth-enabled test clients come and go,
    which caused native xdist worker crashes without Python tracebacks.

    Remove the obsolete scoped-access-key worker pin while retaining
    coverage that revoked keys are reported as REVOKED and rejected by
    authentication and workspace endpoints.

Signed-off-by: Ryan S <267728323+ironcommit@users.noreply.github.com>

Author: ironcommit

.github/workflows/ci.yaml

@@ -1073,12 +1073,12 @@ jobs:
           NMP_FILES_HF_RETRY_ATTEMPTS: "7"
           NMP_FILES_HF_RETRY_INITIAL_DELAY_SECONDS: "1"
           NMP_FILES_HF_RETRY_MAX_DELAY_SECONDS: "30"
+          PYTEST_EXTRA: "--capture=no -p no:faulthandler"
           PYTEST_WORKERS: "4"
           _TYPER_FORCE_DISABLE_TERMINAL: "1"
           # Per-worker crash dumps; pytest's own faulthandler would point them at stderr, which
           # xdist never forwards. See conftest.py.
           PYTEST_CRASH_DUMP_DIR: crash-dumps
-          PYTEST_EXTRA: -p no:faulthandler
       - name: Upload test artifacts
         if: always()
         uses: actions/upload-artifact@043fb46d1a93c77aae656e7c1c64a875d1fc6a0a # v7.0.1

Makefile

@@ -382,12 +382,11 @@ test: test-unit  ## Run all Python unit tests (fast tests without infrastructure
 PYTEST_VERBOSITY := $(if $(filter true,$(CI)),-q,-v)
 PYTEST_WORKERS ?= auto
 PYTEST_MAX_WORKERS ?= 16
-PYTEST_DIST ?= loadscope
 PYTEST_MAX_WORKER_RESTART ?= 2
-PYTEST_CMD = env PYTHONWARNINGS="ignore::UserWarning:pytest_only.version" $(UV) run --frozen \
+PYTEST_DIST ?= loadscope
+PYTEST_CMD = env PYTHONFAULTHANDLER=1 PYTHONWARNINGS="ignore::UserWarning:pytest_only.version" $(UV) run --frozen \
 	pytest \
-	-n $(PYTEST_WORKERS) --maxprocesses=$(PYTEST_MAX_WORKERS) \
-	--max-worker-restart=$(PYTEST_MAX_WORKER_RESTART) \
+	-n $(PYTEST_WORKERS) --maxprocesses=$(PYTEST_MAX_WORKERS) --max-worker-restart=$(PYTEST_MAX_WORKER_RESTART) \
 	--dist $(PYTEST_DIST) --timeout=120 $(PYTEST_VERBOSITY) $(PYTEST_EXTRA)
 
 PYTEST_CI_OPTS := --cov=src --cov=packages \
@@ -411,9 +410,10 @@ PYTEST_CI_CMD = timeout --kill-after=60s $(PYTEST_CI_TIMEOUT)s $(PYTEST_CMD)
 # ``loadscope``.
 test-integration test-integration-ci: PYTEST_DIST := loadgroup
 
-# A killed integration worker never tears down its containers and ports, so replacing it re-runs
-# the group against resources the dead worker still holds. Unit workers own nothing external.
-test-integration test-integration-ci: PYTEST_MAX_WORKER_RESTART := 0
+# In CI, a crashed xdist worker should fail the job immediately. Restarting the
+# worker can hide the crashing test and leave pytest waiting until the wrapper
+# timeout kills the process.
+test-unit-ci test-integration-ci test-gpu-integration-ci: PYTEST_MAX_WORKER_RESTART := 0
 
 .PHONY: test-unit
 test-unit: ## Run Python unit tests across all packages and services

services/core/auth/src/nmp/core/auth/app/embedded_pdp/engine.py

@@ -6,6 +6,7 @@
 import json
 import logging
 import threading
+from functools import cache
 from typing import Any, Dict, List, Optional, cast
 
 from nmp.core.auth.app.embedded_pdp.policy_wasm import ensure_embedded_policy_wasm
@@ -23,13 +24,27 @@ class PolicyEngineError(Exception):
     """Error during policy evaluation."""
 
 
+@cache
+def _get_engine() -> Engine:
+    """Return the process-wide wasmtime Engine.
+
+    Engine setup is the heavyweight step in wasmtime — JIT code generation and process-wide
+    trap/signal-handling registration — and wasmtime's own guidance is one Engine per process with
+    many cheap Stores created from it. Creating a fresh Engine per thread-local OPAPolicy would
+    churn through many Engines over an xdist worker's lifetime; that churn, not any single test, is
+    what was crashing workers with no traceback ("node down: Not properly terminated").
+    """
+    config = Config()
+    config.consume_fuel = True
+    return Engine(config)
+
+
 class OPAPolicy:
     """Wrapper for OPA WASM policy evaluation."""
 
     def __init__(self, wasm_path: str, *, fuel_limit: int = 200_000_000, memory_limit_mb: int = 32):
-        config = Config()
-        config.consume_fuel = True
-        engine = Engine(config)
+        self._owner_thread_id = threading.get_ident()
+        engine = _get_engine()
 
         self.fuel_limit = fuel_limit
         self.store = Store(engine)
@@ -73,20 +88,29 @@ def __init__(self, wasm_path: str, *, fuel_limit: int = 200_000_000, memory_limi
         self._base_heap = self._export_func("opa_heap_ptr_get")(self.store)
         self._data_heap = self._base_heap
         self._data_addr: Optional[int] = None
-        self._lock = threading.Lock()
+
+    def _assert_owner_thread(self) -> None:
+        current_thread_id = threading.get_ident()
+        if current_thread_id != self._owner_thread_id:
+            raise RuntimeError(
+                f"OPAPolicy used from a different thread (owner={self._owner_thread_id}, current={current_thread_id})"
+            )
 
     def _export_func(self, name: str) -> Func:
+        self._assert_owner_thread()
         return cast(Func, self.exports[name])
 
     def _write_json(self, data: Any) -> int:
         """Write JSON to WASM memory, return OPA value address."""
+        self._assert_owner_thread()
         json_bytes = json.dumps(data).encode("utf-8")
         addr = self._export_func("opa_malloc")(self.store, len(json_bytes))
         self.memory.write(self.store, json_bytes, addr)
         return self._export_func("opa_json_parse")(self.store, addr, len(json_bytes))
 
     def _read_json(self, addr: int) -> Any:
         """Read OPA value as JSON from WASM memory."""
+        self._assert_owner_thread()
         json_addr = self._export_func("opa_json_dump")(self.store, addr)
         mem = self.memory.data_ptr(self.store)
         end = json_addr
@@ -96,74 +120,151 @@ def _read_json(self, addr: int) -> Any:
 
     def set_data(self, data: Dict[str, Any]) -> None:
         """Set the base data document."""
-        with self._lock:
-            self.store.set_fuel(DATA_LOADING_FUEL)
-            self._export_func("opa_heap_ptr_set")(self.store, self._base_heap)
-            self._data_addr = self._write_json(data)
-            self._data_heap = self._export_func("opa_heap_ptr_get")(self.store)
+        self._assert_owner_thread()
+        self.store.set_fuel(DATA_LOADING_FUEL)
+        self._export_func("opa_heap_ptr_set")(self.store, self._base_heap)
+        self._data_addr = self._write_json(data)
+        self._data_heap = self._export_func("opa_heap_ptr_get")(self.store)
 
     def evaluate(self, input_data: Dict[str, Any], entrypoint: int = 0) -> Any:
         """Evaluate policy with given input."""
+        self._assert_owner_thread()
         if self._data_addr is None:
             raise PolicyEngineError("Policy data not loaded — refusing to evaluate (fail-closed)")
 
+        self.store.set_fuel(self.fuel_limit)
+
+        heap_base = getattr(self, "_data_heap", self._base_heap)
+        self._export_func("opa_heap_ptr_set")(self.store, heap_base)
+
+        ctx = self._export_func("opa_eval_ctx_new")(self.store)
+        self._export_func("opa_eval_ctx_set_input")(self.store, ctx, self._write_json(input_data))
+        self._export_func("opa_eval_ctx_set_data")(self.store, ctx, self._data_addr)
+        self._export_func("opa_eval_ctx_set_entrypoint")(self.store, ctx, entrypoint)
+
+        self._export_func("eval")(self.store, ctx)
+        return self._read_json(self._export_func("opa_eval_ctx_get_result")(self.store, ctx))
+
+
+class _PolicyRuntimeManager:
+    """Owns policy data snapshots and thread-local WASM policy runtimes."""
+
+    def __init__(self) -> None:
+        self._local = threading.local()
+        self._lock = threading.Lock()
+        self._data: Dict[str, Any] = {}
+        self._data_loaded = False
+        self._data_version = 0
+        self._generation = 0
+
+    def _clear_thread_policy(self) -> None:
+        for attr in ("policy", "policy_generation", "policy_data_version"):
+            if hasattr(self._local, attr):
+                delattr(self._local, attr)
+
+    def _create_policy(self) -> OPAPolicy:
+        from nmp.common.config import get_service_config
+        from nmp.core.auth.config import AuthServiceConfig
+
+        cfg = get_service_config(AuthServiceConfig)
+        path = ensure_embedded_policy_wasm(auto_build=cfg.embedded_pdp_auto_build_wasm)
+        return OPAPolicy(
+            str(path),
+            fuel_limit=cfg.embedded_pdp_cpu_limit * 1_000_000,
+            memory_limit_mb=cfg.embedded_pdp_memory_limit_mb,
+        )
+
+    def _data_snapshot(self) -> tuple[Dict[str, Any], int, bool]:
         with self._lock:
-            self.store.set_fuel(self.fuel_limit)
+            return self._data, self._data_version, self._data_loaded
 
-            heap_base = getattr(self, "_data_heap", self._base_heap)
-            self._export_func("opa_heap_ptr_set")(self.store, heap_base)
+    def _generation_snapshot(self) -> int:
+        with self._lock:
+            return self._generation
+
+    def _get_thread_policy(self) -> Optional[OPAPolicy]:
+        return cast(Optional[OPAPolicy], getattr(self._local, "policy", None))
+
+    def _sync_data_if_needed(self, policy: OPAPolicy) -> None:
+        while True:
+            local_version = cast(int, getattr(self._local, "policy_data_version", -1))
+            data, data_version, data_loaded = self._data_snapshot()
+            if local_version == data_version:
+                return
+
+            if data_loaded:
+                policy.set_data(data)
+
+            _, current_data_version, _ = self._data_snapshot()
+            if data_version == current_data_version:
+                self._local.policy_data_version = data_version
+                return
+
+    def get_policy(self) -> OPAPolicy:
+        """Get or create the current thread's policy runtime."""
+        while True:
+            generation = self._generation_snapshot()
+            policy = self._get_thread_policy()
+            policy_generation = cast(int, getattr(self._local, "policy_generation", -1))
+            if policy is None or policy_generation != generation:
+                policy = self._create_policy()
+                self._local.policy = policy
+                self._local.policy_generation = generation
+                self._local.policy_data_version = -1
+
+            self._sync_data_if_needed(policy)
+            if generation == self._generation_snapshot():
+                return policy
 
-            ctx = self._export_func("opa_eval_ctx_new")(self.store)
-            self._export_func("opa_eval_ctx_set_input")(self.store, ctx, self._write_json(input_data))
-            self._export_func("opa_eval_ctx_set_data")(self.store, ctx, self._data_addr)
-            self._export_func("opa_eval_ctx_set_entrypoint")(self.store, ctx, entrypoint)
+    def set_data(self, data: Dict[str, Any]) -> None:
+        """Set policy data (principals, roles, etc.)."""
+        with self._lock:
+            self._data = data
+            self._data_loaded = True
+            self._data_version += 1
 
-            self._export_func("eval")(self.store, ctx)
-            return self._read_json(self._export_func("opa_eval_ctx_get_result")(self.store, ctx))
+        policy = self._get_thread_policy()
+        if policy is not None and getattr(self._local, "policy_generation", -1) == self._generation_snapshot():
+            self._sync_data_if_needed(policy)
 
+    def reload(self) -> None:
+        """Force each thread to rebuild its policy runtime on next access."""
+        with self._lock:
+            self._generation += 1
+        self._clear_thread_policy()
+        self.get_policy()
 
-# Module-level singleton
-_policy: Optional[OPAPolicy] = None
-_policy_lock = threading.Lock()
-_policy_data: Dict[str, Any] = {}
+    def reset_for_testing(self) -> None:
+        """Reset policy runtime state between tests."""
+        with self._lock:
+            self._data = {}
+            self._data_loaded = False
+            self._data_version += 1
+            self._generation += 1
+        self._clear_thread_policy()
+
+
+_policy_runtime = _PolicyRuntimeManager()
 
 
 def get_policy() -> OPAPolicy:
-    """Get or create the singleton policy instance (thread-safe, double-checked locking)."""
-    global _policy
-    if _policy is None:
-        with _policy_lock:
-            if _policy is None:
-                from nmp.common.config import get_service_config
-                from nmp.core.auth.config import AuthServiceConfig
-
-                cfg = get_service_config(AuthServiceConfig)
-                path = ensure_embedded_policy_wasm(auto_build=cfg.embedded_pdp_auto_build_wasm)
-                _policy = OPAPolicy(
-                    str(path),
-                    fuel_limit=cfg.embedded_pdp_cpu_limit * 1_000_000,
-                    memory_limit_mb=cfg.embedded_pdp_memory_limit_mb,
-                )
-                if _policy_data:
-                    _policy.set_data(_policy_data)
-    return _policy
+    """Get or create the current thread's policy runtime."""
+    return _policy_runtime.get_policy()
 
 
 def set_policy_data(data: Dict[str, Any]) -> None:
     """Set policy data (principals, roles, etc.)."""
-    global _policy_data
-    with _policy_lock:
-        _policy_data = data
-        if _policy is not None:
-            _policy.set_data(data)
+    _policy_runtime.set_data(data)
 
 
 def reload_policy() -> None:
-    """Force reload the policy."""
-    global _policy
-    with _policy_lock:
-        _policy = None
-    get_policy()
+    """Force each thread to rebuild its policy runtime on next access."""
+    _policy_runtime.reload()
+
+
+def _reset_policy_state_for_testing() -> None:
+    """Reset module policy state between tests."""
+    _policy_runtime.reset_for_testing()
 
 
 def evaluate(entrypoint: str, input_data: Dict[str, Any]) -> Dict[str, Any]: