feat(adapters): caching interceptor

fern/versions/latest/pages/model-server/adapters-caching.mdx

@@ -0,0 +1,26 @@
+---
+title: "Caching Interceptor"
+description: "Disk-backed cache that short-circuits the chain on cache hit."
+position: 5
+---
+
+Disk-backed cache keyed by a SHA-256 hash of the canonicalized request body (+ optional session prefix). On cache hit, the interceptor short-circuits the chain and returns the stored response without invoking the upstream.
+
+| Name | Stage | Purpose |
+|------|-------|---------|
+| `caching` | request → response | Disk-backed cache (sqlite). Hits short-circuit upstream. |
+
+### `caching`
+
+```yaml
+- name: caching
+  config:
+    cache_dir: /var/cache/gym-adapter
+    bypass: false
+```
+
+Cache keys include the session prefix when present (via `ctx.extra["session_id"]`), so the same body in different sessions does not collide.
+
+<Note>
+Per-replica proxy or server instance means per-replica disk cache: if multiple replicas share the same `cache_dir`, sqlite write contention is real. Configure a per-replica path (e.g. `cache_dir: /tmp/adapter_cache_${pod_id}`) or accept the race.
+</Note>

nemo_gym/adapters/cache/__init__.py

@@ -0,0 +1,14 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# 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.

nemo_gym/adapters/cache/disk_cache.py

@@ -0,0 +1,112 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# 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.
+from __future__ import annotations
+
+import asyncio
+import hashlib
+import json
+import logging
+import os
+import sqlite3
+import threading
+import time
+from typing import Any
+
+
+logger = logging.getLogger(__name__)
+
+_RELEVANT_KEYS = ("model", "messages", "tools", "temperature", "max_tokens", "top_p", "seed")
+
+_SCHEMA = """
+CREATE TABLE IF NOT EXISTS cache_v1 (
+    key        TEXT PRIMARY KEY,
+    value      TEXT NOT NULL,
+    created_at REAL NOT NULL
+);
+"""
+
+_UPSERT = """
+INSERT INTO cache_v1 (key, value, created_at) VALUES (?, ?, ?)
+ON CONFLICT(key) DO UPDATE SET value=excluded.value, created_at=excluded.created_at;
+"""
+
+
+class DiskCache:
+    def __init__(self, cache_dir: str) -> None:
+        os.makedirs(cache_dir, exist_ok=True)
+        self._db_path = os.path.join(cache_dir, "adapter_cache.db")
+        self._lock = threading.Lock()
+        self._init_db()
+
+    def _init_db(self) -> None:
+        try:
+            with self._lock:
+                conn = sqlite3.connect(self._db_path)
+                try:
+                    conn.execute("PRAGMA journal_mode=WAL")
+                    conn.execute("PRAGMA busy_timeout=5000")
+                    conn.execute(_SCHEMA)
+                    conn.commit()
+                finally:
+                    conn.close()
+        except sqlite3.Error:
+            logger.warning("Failed to initialize cache DB at %s", self._db_path, exc_info=True)
+
+    @staticmethod
+    def cache_key(body: dict[str, Any], *, session_prefix: str = "") -> str:
+        canonical: dict[str, Any] = {}
+        for k in _RELEVANT_KEYS:
+            if k in body:
+                canonical[k] = body[k]
+        if "extra_body" in body and isinstance(body["extra_body"], dict):
+            canonical["extra_body"] = body["extra_body"]
+        raw = json.dumps(canonical, sort_keys=True, ensure_ascii=False)
+        if session_prefix:
+            raw = session_prefix + "|" + raw
+        return hashlib.sha256(raw.encode()).hexdigest()
+
+    def _get_sync(self, key: str) -> dict | None:
+        try:
+            with self._lock:
+                conn = sqlite3.connect(self._db_path)
+                try:
+                    row = conn.execute("SELECT value FROM cache_v1 WHERE key = ?", (key,)).fetchone()
+                finally:
+                    conn.close()
+            if row is None:
+                return None
+            return json.loads(row[0])
+        except (sqlite3.Error, json.JSONDecodeError):
+            logger.warning("Cache get failed for key=%s", key[:16], exc_info=True)
+            return None
+
+    def _set_sync(self, key: str, value: dict) -> None:
+        try:
+            serialized = json.dumps(value, ensure_ascii=False)
+            with self._lock:
+                conn = sqlite3.connect(self._db_path)
+                try:
+                    conn.execute(_UPSERT, (key, serialized, time.time()))
+                    conn.commit()
+                finally:
+                    conn.close()
+        except sqlite3.Error:
+            logger.warning("Cache set failed for key=%s", key[:16], exc_info=True)
+
+    async def get(self, key: str) -> dict | None:
+        return await asyncio.to_thread(self._get_sync, key)
+
+    async def set(self, key: str, value: dict) -> None:
+        await asyncio.to_thread(self._set_sync, key, value)

nemo_gym/adapters/interceptors/caching.py

@@ -0,0 +1,67 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# 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.
+from __future__ import annotations
+
+import logging
+
+from nemo_gym.adapters.cache.disk_cache import DiskCache
+from nemo_gym.adapters.types import (
+    AdapterRequest,
+    AdapterResponse,
+    RequestToResponseInterceptor,
+    ResponseInterceptor,
+)
+
+
+logger = logging.getLogger(__name__)
+
+
+class Interceptor(RequestToResponseInterceptor, ResponseInterceptor):
+    stream_safe = False
+
+    def __init__(self, cache_dir: str, *, bypass: bool = False) -> None:
+        self._bypass = bypass
+        self._cache = DiskCache(cache_dir)
+
+    async def intercept_request(
+        self,
+        req: AdapterRequest,
+    ) -> AdapterRequest | AdapterResponse:
+        if self._bypass:
+            return req
+        session_prefix = req.ctx.extra.get("session_id", "")
+        key = DiskCache.cache_key(req.body, session_prefix=session_prefix)
+        hit = await self._cache.get(key)
+        if hit is not None:
+            logger.debug("cache hit key=%s", key[:16])
+            req.ctx.extra["cache_hit"] = True
+            return AdapterResponse(
+                status_code=200,
+                headers={},
+                body=hit,
+                latency_ms=0.0,
+                ctx=req.ctx,
+            )
+        req.ctx.extra["cache_key"] = key
+        return req
+
+    async def intercept_response(self, resp: AdapterResponse) -> AdapterResponse:
+        key = resp.ctx.extra.get("cache_key")
+        if key is None or not resp.ok:
+            return resp
+        if isinstance(resp.body, dict):
+            await self._cache.set(key, resp.body)
+        resp.ctx.extra.pop("cache_key", None)
+        return resp

nemo_gym/adapters/registry.py

@@ -52,6 +52,9 @@
 # Each family adds entries under its own family-named comment marker so
 # different families don't fight for the same diff context.
 
+# Caching family.
+_BUILTIN["caching"] = "nemo_gym.adapters.interceptors.caching"
+
 # External / plugin registrations at runtime.
 _EXTRA: dict[str, str] = {}
 

tests/unit_tests/test_adapter_cache_keys.py

@@ -0,0 +1,68 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# 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.
+"""Cache-key derivation golden tests — ported from NEL.
+
+The golden SHA-256 strings are byte-identical to NEL's: the cache_key
+algorithm is the same code re-rooted at ``nemo_gym.adapters.cache``, so
+both implementations must produce the same digests for the same inputs.
+This file protects against accidental drift.
+"""
+
+from nemo_gym.adapters.cache.disk_cache import DiskCache
+
+
+def _compute_key(body):
+    return DiskCache.cache_key(body)
+
+
+def test_golden_key_simple():
+    body = {"model": "gpt-4", "messages": [{"role": "user", "content": "hello"}], "temperature": 0.7}
+    key = _compute_key(body)
+    assert key == _compute_key(body)
+    assert key == "92232807c9b5c0bab68fd1abb38bc83c884fb04b0eb9b8b453c19daead7d682f"  # pragma: allowlist secret
+
+
+def test_golden_key_with_tools():
+    body = {
+        "model": "gpt-4",
+        "messages": [{"role": "user", "content": "hello"}],
+        "tools": [{"type": "function", "function": {"name": "get_weather"}}],
+    }
+    key = _compute_key(body)
+    assert key == _compute_key(body)
+    assert key == "ce43d524c627c96eaadbb9abdac38ad5eb60a93d15b1d26880fa1fe0f86d4c68"  # pragma: allowlist secret
+
+
+def test_key_ignores_irrelevant_fields():
+    base = {"model": "gpt-4", "messages": [{"role": "user", "content": "hello"}]}
+    with_stream = {**base, "stream": True}
+    assert _compute_key(base) == _compute_key(with_stream)
+
+
+def test_key_changes_with_temperature():
+    body_a = {"model": "gpt-4", "messages": [{"role": "user", "content": "hello"}], "temperature": 0.7}
+    body_b = {"model": "gpt-4", "messages": [{"role": "user", "content": "hello"}], "temperature": 0.9}
+    assert _compute_key(body_a) != _compute_key(body_b)
+
+
+def test_key_differs_with_session_prefix():
+    """Session prefix ensures repeats of the same problem never share cache entries."""
+    body = {"model": "gpt-4", "messages": [{"role": "user", "content": "hello"}]}
+    k_none = DiskCache.cache_key(body)
+    k_a = DiskCache.cache_key(body, session_prefix="repeat-0")
+    k_b = DiskCache.cache_key(body, session_prefix="repeat-1")
+    assert k_none != k_a
+    assert k_a != k_b
+    assert k_none == DiskCache.cache_key(body, session_prefix="")

tests/unit_tests/test_adapter_disk_cache.py

@@ -0,0 +1,68 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# 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.
+"""DiskCache behavior tests — ported from NEL."""
+
+import asyncio
+
+import pytest
+
+from nemo_gym.adapters.cache.disk_cache import DiskCache
+
+
+@pytest.fixture
+def cache(tmp_path):
+    return DiskCache(str(tmp_path / "cache"))
+
+
+def test_cache_key_deterministic():
+    body = {"model": "gpt-4", "messages": [{"role": "user", "content": "hello"}]}
+    assert DiskCache.cache_key(body) == DiskCache.cache_key(body)
+
+
+def test_cache_key_varies_with_model():
+    base = {"model": "gpt-4", "messages": [{"role": "user", "content": "hello"}]}
+    other = {**base, "model": "gpt-3.5-turbo"}
+    assert DiskCache.cache_key(base) != DiskCache.cache_key(other)
+
+
+async def test_get_set_roundtrip(cache):
+    key = DiskCache.cache_key({"model": "x", "messages": []})
+    value = {"choices": [{"message": {"content": "hi"}}]}
+    await cache.set(key, value)
+    result = await cache.get(key)
+    assert result == value
+
+
+async def test_get_missing_returns_none(cache):
+    result = await cache.get("nonexistent_key")
+    assert result is None
+
+
+async def test_concurrent_writes(cache):
+    pairs = [(f"key_{i}", {"i": i}) for i in range(50)]
+
+    await asyncio.gather(*(cache.set(k, v) for k, v in pairs))
+
+    for k, v in pairs:
+        assert await cache.get(k) == v
+
+
+async def test_degrade_on_bad_path():
+    bad = DiskCache.__new__(DiskCache)
+    bad._db_path = "/dev/null/bad/adapter_cache.db"
+    bad._lock = __import__("threading").Lock()
+
+    assert await bad.get("anything") is None
+    await bad.set("anything", {"x": 1})