feat: implement R141 snapshot inventory indexing

This commit is contained in:
rookiestar28
2026-03-13 23:42:15 +08:00
parent d641c04bde
commit 9904ecea49
5 changed files with 508 additions and 40 deletions
+12 -5
View File
@@ -20,8 +20,8 @@ if __package__ and "." in __package__:
from ..models.schemas import MAX_BODY_SIZE
from ..services.access_control import is_loopback, require_admin_token
from ..services.preflight import (
_get_model_inventory,
_get_node_class_mappings,
get_model_inventory_snapshot,
run_preflight_check,
)
from ..services.rate_limit import check_rate_limit
@@ -30,8 +30,8 @@ else: # pragma: no cover (test-only import mode)
from models.schemas import MAX_BODY_SIZE # type: ignore
from services.access_control import is_loopback, require_admin_token # type: ignore
from services.preflight import ( # type: ignore
_get_model_inventory,
_get_node_class_mappings,
get_model_inventory_snapshot,
run_preflight_check,
)
from services.rate_limit import check_rate_limit # type: ignore
@@ -195,11 +195,18 @@ async def inventory_handler(request: web.Request) -> web.Response:
nodes_map = _get_node_class_mappings()
node_classes = sorted(list(nodes_map.keys()))
# Models
models_map = _get_model_inventory()
inventory_snapshot = get_model_inventory_snapshot()
return web.json_response(
{"ok": True, "nodes": node_classes, "models": models_map}
{
"ok": True,
"nodes": node_classes,
"models": inventory_snapshot["models"],
"snapshot_ts": inventory_snapshot["snapshot_ts"],
"scan_state": inventory_snapshot["scan_state"],
"stale": inventory_snapshot["stale"],
"last_error": inventory_snapshot["last_error"],
}
)
except Exception as e:
logger.exception("Inventory fetch failed")
+30 -3
View File
@@ -4,6 +4,10 @@ const { defineConfig } = require('@playwright/test');
const e2ePort = process.env.OPENCLAW_E2E_PORT || '3000';
const e2eBase = `http://127.0.0.1:${e2ePort}`;
function isWslDrvFs() {
return process.platform === 'linux' && !!process.env.WSL_DISTRO_NAME && process.cwd().startsWith('/mnt/');
}
function resolveWorkers() {
const raw = process.env.OPENCLAW_PLAYWRIGHT_WORKERS;
if (raw) {
@@ -18,16 +22,39 @@ function resolveWorkers() {
// CRITICAL: WSL on /mnt/* has repeatable E2E harness instability under
// high parallelism; cap workers to keep the repo acceptance gate deterministic.
if (process.platform === 'linux' && process.env.WSL_DISTRO_NAME && process.cwd().startsWith('/mnt/')) {
if (isWslDrvFs()) {
return 1;
}
return undefined;
}
function resolveTimeoutMs() {
const raw = process.env.OPENCLAW_PLAYWRIGHT_TIMEOUT_MS;
if (raw) {
const parsed = Number.parseInt(raw, 10);
if (!Number.isInteger(parsed) || parsed < 1) {
throw new Error(
`OPENCLAW_PLAYWRIGHT_TIMEOUT_MS must be a positive integer, got '${raw}'`,
);
}
return parsed;
}
// IMPORTANT: WSL on /mnt/* can push single-test wall time well past 30s even
// with one worker, especially when a spec reloads the harness inside one test.
if (isWslDrvFs()) {
return 60_000;
}
return 30_000;
}
const timeoutMs = resolveTimeoutMs();
module.exports = defineConfig({
testDir: 'tests/e2e/specs',
timeout: 30_000,
timeout: timeoutMs,
retries: 0,
workers: resolveWorkers(),
use: {
@@ -39,6 +66,6 @@ module.exports = defineConfig({
command: `${process.env.PYTHON || (process.platform === 'win32' ? 'python' : 'python3')} -m http.server ${e2ePort}`,
url: `${e2eBase}/tests/e2e/test-harness.html`,
reuseExistingServer: true,
timeout: 30_000,
timeout: timeoutMs,
},
});
+175 -32
View File
@@ -6,6 +6,7 @@ checking for missing node classes and models.
"""
import logging
import threading
import time
from typing import Any, Dict, List, Set, Tuple
@@ -28,6 +29,19 @@ except Exception: # pragma: no cover
_CACHE = {}
_CACHE_TTL = 60 # seconds
_INVENTORY_SCAN_STATE_IDLE = "idle"
_INVENTORY_SCAN_STATE_REFRESHING = "refreshing"
_INVENTORY_SCAN_STATE_ERROR = "error"
_INVENTORY_SNAPSHOT_KEY = "inventory_snapshot"
_INVENTORY_SNAPSHOT_TS_KEY = "inventory_snapshot_ts"
_INVENTORY_LAST_ERROR_KEY = "inventory_last_error"
_INVENTORY_SCAN_STATE_KEY = "inventory_scan_state"
_INVENTORY_CHECKPOINT_KEY = "inventory_scan_checkpoint"
_INVENTORY_LAST_ATTEMPT_TS_KEY = "inventory_last_attempt_ts"
_LEGACY_INVENTORY_CACHE_KEY = "inventory"
_INVENTORY_LOCK = threading.RLock()
_INVENTORY_SCAN_THREAD: threading.Thread | None = None
_INVENTORY_ERROR_RETRY_SEC = 5
# Heuristic mapping: input_key -> folder_paths type
_INPUT_KEY_MAP = {
@@ -51,28 +65,7 @@ def _get_node_class_mappings() -> Dict[str, Any]:
return {}
def _get_model_inventory() -> Dict[str, List[str]]:
"""
Retrieve snapshot of available models using folder_paths.
Returns a dict mapping folder name (e.g., 'checkpoints') to list of filenames.
Cached for 60s to prevent IO spam.
"""
global _CACHE
now = time.time()
cached = _CACHE.get("inventory")
if cached:
timestamp, data = cached
if now - timestamp < _CACHE_TTL:
return data
inventory = {}
if not folder_paths:
return inventory
# Common model types to check
# We use the keys from folder_paths.folder_names_and_paths if available,
# or a hardcoded list of common ones.
def _resolve_inventory_model_types() -> List[str]:
model_types = [
"checkpoints",
"loras",
@@ -88,23 +81,173 @@ def _get_model_inventory() -> Dict[str, List[str]]:
"vae_approx",
"photomaker",
]
# Add any dynamic ones
if hasattr(folder_paths, "folder_names_and_paths"):
for k in folder_paths.folder_names_and_paths.keys():
if k not in model_types:
model_types.append(k)
for key in folder_paths.folder_names_and_paths.keys():
if key not in model_types:
model_types.append(key)
return model_types
for mtype in model_types:
def _scan_model_inventory(checkpoint: List[str] | None = None) -> Dict[str, List[str]]:
"""
Build a complete model inventory snapshot synchronously.
The caller decides whether this runs on-request or in a background worker.
"""
inventory: Dict[str, List[str]] = {}
if not folder_paths:
return inventory
model_types = _resolve_inventory_model_types()
for index, model_type in enumerate(model_types):
if checkpoint is not None:
checkpoint[:] = [str(index), model_type]
try:
files = folder_paths.get_filename_list(mtype)
files = folder_paths.get_filename_list(model_type)
if files:
inventory[mtype] = list(files)
inventory[model_type] = list(files)
except Exception:
# Some folders might not exist or raise error
# Some folders might not exist or raise error.
continue
if checkpoint is not None:
checkpoint[:] = []
return inventory
_CACHE["inventory"] = (now, inventory)
def _copy_inventory_snapshot(models: Dict[str, List[str]]) -> Dict[str, List[str]]:
return {key: list(value) for key, value in (models or {}).items()}
def _inventory_snapshot_stale_locked(now: float | None = None) -> bool:
snapshot_ts = _CACHE.get(_INVENTORY_SNAPSHOT_TS_KEY)
if not snapshot_ts:
return True
current = time.time() if now is None else now
return current - float(snapshot_ts) >= _CACHE_TTL
def _inventory_scan_running_locked() -> bool:
global _INVENTORY_SCAN_THREAD
if _INVENTORY_SCAN_THREAD is not None and not _INVENTORY_SCAN_THREAD.is_alive():
_INVENTORY_SCAN_THREAD = None
return _INVENTORY_SCAN_THREAD is not None
def _inventory_should_schedule_refresh_locked(now: float) -> bool:
if not folder_paths or not _inventory_snapshot_stale_locked(now):
return False
if _inventory_scan_running_locked():
return False
if _CACHE.get(_INVENTORY_SCAN_STATE_KEY) != _INVENTORY_SCAN_STATE_ERROR:
return True
last_attempt = float(_CACHE.get(_INVENTORY_LAST_ATTEMPT_TS_KEY) or 0.0)
return now - last_attempt >= _INVENTORY_ERROR_RETRY_SEC
def _inventory_refresh_worker() -> None:
checkpoint: List[str] = []
try:
snapshot = _scan_model_inventory(checkpoint)
with _INVENTORY_LOCK:
_CACHE[_INVENTORY_SNAPSHOT_KEY] = snapshot
_CACHE[_INVENTORY_SNAPSHOT_TS_KEY] = time.time()
_CACHE[_INVENTORY_LAST_ERROR_KEY] = None
_CACHE[_INVENTORY_SCAN_STATE_KEY] = _INVENTORY_SCAN_STATE_IDLE
_CACHE[_INVENTORY_CHECKPOINT_KEY] = None
except Exception as exc: # pragma: no cover - defensive outer guard
with _INVENTORY_LOCK:
_CACHE[_INVENTORY_LAST_ERROR_KEY] = str(exc)
_CACHE[_INVENTORY_SCAN_STATE_KEY] = _INVENTORY_SCAN_STATE_ERROR
_CACHE[_INVENTORY_CHECKPOINT_KEY] = (
checkpoint[1] if len(checkpoint) >= 2 else None
)
logger.exception("Inventory deep scan failed")
finally:
global _INVENTORY_SCAN_THREAD
with _INVENTORY_LOCK:
_INVENTORY_SCAN_THREAD = None
def _schedule_inventory_refresh_locked() -> None:
global _INVENTORY_SCAN_THREAD
if _inventory_scan_running_locked() or not folder_paths:
return
_CACHE[_INVENTORY_SCAN_STATE_KEY] = _INVENTORY_SCAN_STATE_REFRESHING
_CACHE.setdefault(_INVENTORY_LAST_ERROR_KEY, None)
_CACHE[_INVENTORY_LAST_ATTEMPT_TS_KEY] = time.time()
worker = threading.Thread(
target=_inventory_refresh_worker,
name="openclaw-inventory-refresh",
daemon=True,
)
_INVENTORY_SCAN_THREAD = worker
worker.start()
def get_model_inventory_snapshot(*, trigger_refresh: bool = True) -> Dict[str, Any]:
"""
Return the latest served inventory snapshot plus scan metadata.
This powers `/openclaw/preflight/inventory` so requests can return quickly
while a background deep scan refreshes stale or missing snapshots.
"""
now = time.time()
with _INVENTORY_LOCK:
if trigger_refresh and _inventory_should_schedule_refresh_locked(now):
_schedule_inventory_refresh_locked()
models = _copy_inventory_snapshot(_CACHE.get(_INVENTORY_SNAPSHOT_KEY, {}))
snapshot_ts = _CACHE.get(_INVENTORY_SNAPSHOT_TS_KEY)
scan_state = _CACHE.get(_INVENTORY_SCAN_STATE_KEY, _INVENTORY_SCAN_STATE_IDLE)
last_error = _CACHE.get(_INVENTORY_LAST_ERROR_KEY)
if not folder_paths:
scan_state = _INVENTORY_SCAN_STATE_IDLE
last_error = None
stale = bool(folder_paths) and _inventory_snapshot_stale_locked(now)
return {
"models": models,
"snapshot_ts": snapshot_ts,
"scan_state": scan_state,
"stale": stale,
"last_error": last_error,
}
def _reset_inventory_state_for_tests() -> None:
global _INVENTORY_SCAN_THREAD
thread = _INVENTORY_SCAN_THREAD
if thread is not None and thread.is_alive():
thread.join(timeout=2.0)
with _INVENTORY_LOCK:
_INVENTORY_SCAN_THREAD = None
for key in (
_INVENTORY_SNAPSHOT_KEY,
_INVENTORY_SNAPSHOT_TS_KEY,
_INVENTORY_LAST_ERROR_KEY,
_INVENTORY_SCAN_STATE_KEY,
_INVENTORY_CHECKPOINT_KEY,
_INVENTORY_LAST_ATTEMPT_TS_KEY,
_LEGACY_INVENTORY_CACHE_KEY,
):
_CACHE.pop(key, None)
def _get_model_inventory() -> Dict[str, List[str]]:
"""
Retrieve snapshot of available models using folder_paths.
Returns a dict mapping folder name (e.g., 'checkpoints') to list of filenames.
Cached for 60s to prevent IO spam.
"""
global _CACHE
now = time.time()
cached = _CACHE.get(_LEGACY_INVENTORY_CACHE_KEY)
if cached:
timestamp, data = cached
if now - timestamp < _CACHE_TTL:
return data
inventory = _scan_model_inventory()
_CACHE[_LEGACY_INVENTORY_CACHE_KEY] = (now, inventory)
return inventory
+246
View File
@@ -0,0 +1,246 @@
import asyncio
import threading
import time
import unittest
from unittest.mock import MagicMock, patch
import services.preflight
try:
from aiohttp import web
from aiohttp.test_utils import AioHTTPTestCase, unittest_run_loop
except Exception: # pragma: no cover
web = None # type: ignore
AioHTTPTestCase = unittest.TestCase # type: ignore
def unittest_run_loop(fn): # type: ignore
return fn
def _wait_for(predicate, timeout=1.0):
deadline = time.time() + timeout
while time.time() < deadline:
if predicate():
return True
time.sleep(0.01)
return predicate()
class TestR141InventoryIndexing(unittest.TestCase):
def setUp(self):
services.preflight._reset_inventory_state_for_tests()
def tearDown(self):
services.preflight._reset_inventory_state_for_tests()
def test_snapshot_response_schedules_background_refresh(self):
started = threading.Event()
release = threading.Event()
def slow_scan(_checkpoint=None):
started.set()
release.wait(timeout=1.0)
return {"checkpoints": ["sdxl.safetensors"]}
with (
patch.object(
services.preflight, "folder_paths", MagicMock(), create=True
) as mock_folder_paths,
patch.object(services.preflight, "_scan_model_inventory", side_effect=slow_scan),
):
mock_folder_paths.folder_names_and_paths = {"checkpoints": [("/tmp", None)]}
first = services.preflight.get_model_inventory_snapshot()
self.assertEqual(first["models"], {})
self.assertEqual(first["scan_state"], "refreshing")
self.assertTrue(first["stale"])
self.assertIsNone(first["snapshot_ts"])
self.assertTrue(started.wait(timeout=0.5))
release.set()
self.assertTrue(
_wait_for(
lambda: services.preflight.get_model_inventory_snapshot(
trigger_refresh=False
)["scan_state"]
== "idle"
)
)
final = services.preflight.get_model_inventory_snapshot(trigger_refresh=False)
self.assertEqual(
final["models"], {"checkpoints": ["sdxl.safetensors"]}
)
self.assertEqual(final["scan_state"], "idle")
self.assertFalse(final["stale"])
self.assertIsNone(final["last_error"])
self.assertIsInstance(final["snapshot_ts"], float)
def test_failed_background_refresh_surfaces_error_metadata(self):
started = threading.Event()
def exploding_scan(_checkpoint=None):
started.set()
raise RuntimeError("scan boom")
with (
patch.object(
services.preflight, "folder_paths", MagicMock(), create=True
) as mock_folder_paths,
patch.object(
services.preflight, "_scan_model_inventory", side_effect=exploding_scan
),
):
mock_folder_paths.folder_names_and_paths = {"checkpoints": [("/tmp", None)]}
first = services.preflight.get_model_inventory_snapshot()
self.assertEqual(first["scan_state"], "refreshing")
self.assertTrue(started.wait(timeout=0.5))
self.assertTrue(
_wait_for(
lambda: services.preflight.get_model_inventory_snapshot(
trigger_refresh=False
)["scan_state"]
== "error"
)
)
final = services.preflight.get_model_inventory_snapshot(trigger_refresh=False)
self.assertEqual(final["models"], {})
self.assertEqual(final["scan_state"], "error")
self.assertTrue(final["stale"])
self.assertEqual(final["last_error"], "scan boom")
self.assertIsNone(final["snapshot_ts"])
def test_error_state_respects_retry_cooldown(self):
with patch.object(
services.preflight, "folder_paths", MagicMock(), create=True
) as mock_folder_paths:
mock_folder_paths.folder_names_and_paths = {"checkpoints": [("/tmp", None)]}
services.preflight._CACHE[services.preflight._INVENTORY_SCAN_STATE_KEY] = "error"
services.preflight._CACHE[services.preflight._INVENTORY_LAST_ERROR_KEY] = (
"scan boom"
)
services.preflight._CACHE[services.preflight._INVENTORY_LAST_ATTEMPT_TS_KEY] = (
time.time()
)
with patch.object(
services.preflight, "_schedule_inventory_refresh_locked"
) as mock_schedule:
snapshot = services.preflight.get_model_inventory_snapshot()
self.assertEqual(snapshot["scan_state"], "error")
self.assertTrue(snapshot["stale"])
self.assertEqual(snapshot["last_error"], "scan boom")
mock_schedule.assert_not_called()
@unittest.skipIf(web is None, "aiohttp not installed")
class TestR141InventoryApi(AioHTTPTestCase):
async def get_application(self):
from api.preflight_handler import inventory_handler
app = web.Application()
app.router.add_get("/openclaw/preflight/inventory", inventory_handler)
return app
def setUp(self):
super().setUp()
services.preflight._reset_inventory_state_for_tests()
def tearDown(self):
services.preflight._reset_inventory_state_for_tests()
super().tearDown()
@patch("api.preflight_handler.check_rate_limit")
@patch("api.preflight_handler.require_admin_token")
@unittest_run_loop
async def test_inventory_handler_returns_metadata_fields(
self, mock_require_admin, mock_rate_limit
):
mock_rate_limit.return_value = True
mock_require_admin.return_value = (True, None)
with (
patch.object(
services.preflight,
"get_model_inventory_snapshot",
return_value={
"models": {"checkpoints": ["a.safetensors"]},
"snapshot_ts": 123.0,
"scan_state": "idle",
"stale": False,
"last_error": None,
},
),
patch("api.preflight_handler.get_model_inventory_snapshot") as mock_snapshot,
patch("api.preflight_handler._get_node_class_mappings") as mock_nodes,
):
mock_snapshot.return_value = {
"models": {"checkpoints": ["a.safetensors"]},
"snapshot_ts": 123.0,
"scan_state": "idle",
"stale": False,
"last_error": None,
}
mock_nodes.return_value = {"KSampler": object}
resp = await self.client.get("/openclaw/preflight/inventory")
self.assertEqual(resp.status, 200)
payload = await resp.json()
self.assertTrue(payload["ok"])
self.assertEqual(payload["nodes"], ["KSampler"])
self.assertEqual(payload["models"], {"checkpoints": ["a.safetensors"]})
self.assertEqual(payload["snapshot_ts"], 123.0)
self.assertEqual(payload["scan_state"], "idle")
self.assertFalse(payload["stale"])
self.assertIsNone(payload["last_error"])
@patch("api.preflight_handler.check_rate_limit")
@patch("api.preflight_handler.require_admin_token")
@unittest_run_loop
async def test_inventory_handler_returns_quickly_while_refresh_runs(
self, mock_require_admin, mock_rate_limit
):
mock_rate_limit.return_value = True
mock_require_admin.return_value = (True, None)
started = threading.Event()
release = threading.Event()
def slow_scan(_checkpoint=None):
started.set()
release.wait(timeout=1.0)
return {"checkpoints": ["late.safetensors"]}
with (
patch.object(
services.preflight, "folder_paths", MagicMock(), create=True
) as mock_folder_paths,
patch.object(services.preflight, "_scan_model_inventory", side_effect=slow_scan),
patch("api.preflight_handler._get_node_class_mappings", return_value={}),
):
mock_folder_paths.folder_names_and_paths = {"checkpoints": [("/tmp", None)]}
resp = await asyncio.wait_for(
self.client.get("/openclaw/preflight/inventory"), timeout=0.2
)
payload = await resp.json()
self.assertTrue(started.wait(timeout=0.5))
self.assertEqual(resp.status, 200)
self.assertEqual(payload["models"], {})
self.assertEqual(payload["scan_state"], "refreshing")
self.assertTrue(payload["stale"])
self.assertIsNone(payload["snapshot_ts"])
release.set()
self.assertTrue(
_wait_for(
lambda: services.preflight.get_model_inventory_snapshot(
trigger_refresh=False
)["scan_state"]
== "idle"
)
)
+45
View File
@@ -86,7 +86,13 @@ export const ExplorerTab = {
invList.style.border = "1px solid var(--border-color, #444)";
invList.style.padding = "0.5rem";
const invStatus = makeEl("div", "openclaw-inventory-status openclaw-inventory-status moltbot-inventory-status");
invStatus.style.fontSize = "0.85em";
invStatus.style.opacity = "0.8";
invStatus.style.marginBottom = "8px";
invContent.appendChild(searchInput);
invContent.appendChild(invStatus);
invContent.appendChild(invList);
// --- Snapshots Content ---
@@ -139,6 +145,7 @@ export const ExplorerTab = {
// --- Logic ---
let inventoryData = null;
let inventoryRefreshTimer = null;
// Helper: Debounce
function debounce(func, wait) {
@@ -152,16 +159,54 @@ export const ExplorerTab = {
// Fetch Inventory
async function loadInventory() {
if (inventoryRefreshTimer) {
clearTimeout(inventoryRefreshTimer);
inventoryRefreshTimer = null;
}
invList.innerHTML = "Loading...";
invStatus.textContent = "";
const res = await openclawApi.getInventory();
if (res.ok) {
inventoryData = res.data;
renderInventoryStatus(res.data);
renderInventoryList(res.data, searchInput.value);
if (res.data?.scan_state === "refreshing" || (res.data?.stale && !res.data?.last_error)) {
inventoryRefreshTimer = window.setTimeout(() => {
loadInventory().catch(() => { });
}, 1500);
}
} else {
invStatus.textContent = "";
invList.innerHTML = `<div class="error">Failed to load inventory: ${res.error}</div>`;
}
}
function renderInventoryStatus(data) {
if (!data) {
invStatus.textContent = "";
return;
}
const bits = [];
if (data.scan_state === "refreshing") {
bits.push("Refreshing inventory snapshot...");
} else if (data.scan_state === "error") {
bits.push("Inventory refresh failed.");
}
if (typeof data.snapshot_ts === "number" && Number.isFinite(data.snapshot_ts)) {
const stamp = new Date(data.snapshot_ts * 1000);
bits.push(`Snapshot: ${stamp.toLocaleString()}`);
} else {
bits.push("Snapshot: pending");
}
if (data.stale) {
bits.push("State: stale");
}
if (data.last_error) {
bits.push(`Last error: ${data.last_error}`);
}
invStatus.textContent = bits.join(" ");
}
async function loadSnapshots() {
snapList.innerHTML = "Loading...";
const res = await openclawApi.listCheckpoints();