Files
Azornes-Comfyui-Model-Resolver/core/workflow/inventory.py
T

216 lines
7.7 KiB
Python

"""Shared workflow model inventory and incremental cache management."""
import threading
import time
from collections import OrderedDict
from typing import Any, Callable, Dict, List, Optional
from ..contracts import ModelReference, WorkflowModelInventory
from ..log_system import create_module_logger
from . import analysis
log = create_module_logger(__name__)
_WORKFLOW_MODEL_INVENTORY_CACHE_TTL_SECONDS = 300.0
_WORKFLOW_MODEL_INVENTORY_CACHE_MAX_ENTRIES = 8
_WORKFLOW_MODEL_INVENTORY_CACHE_LOCK = threading.Lock()
_WORKFLOW_MODEL_INVENTORY_CACHE: OrderedDict[str, Dict[str, Any]] = OrderedDict()
def _build_workflow_node_cache(
workflow_json: Dict[str, Any],
model_refs: List[ModelReference],
) -> Dict[tuple[str, str, str], Dict[str, Any]]:
fingerprints = analysis._get_workflow_node_fingerprints(workflow_json)
refs_by_node: Dict[tuple[str, str, str], List[ModelReference]] = {
key: [] for key in fingerprints
}
for ref in model_refs:
is_top_level = ref.is_top_level is not False
key = analysis._get_workflow_scope_key(
ref.node_id,
subgraph_id=ref.subgraph_id,
is_top_level=is_top_level,
)
if key in refs_by_node:
refs_by_node[key].append(ref)
return {
key: {
"fingerprint": fingerprint,
"refs": refs_by_node.get(key, []),
}
for key, fingerprint in fingerprints.items()
}
def _get_workflow_model_inventory_cache_key(
workflow_json: Dict[str, Any],
) -> str:
return analysis._hash_normalized_workflow_value(workflow_json)
def _get_analysis_log_context(
cache_key: str,
analysis_id: Optional[str],
) -> str:
normalized_analysis_id = str(analysis_id or "").strip()
if normalized_analysis_id:
return f"analysis_id={normalized_analysis_id}"
return f"workflow_signature={cache_key[:12]}"
def invalidate_workflow_model_inventory_cache() -> None:
"""Clear shared workflow model analysis snapshots."""
with _WORKFLOW_MODEL_INVENTORY_CACHE_LOCK:
_WORKFLOW_MODEL_INVENTORY_CACHE.clear()
def get_workflow_model_inventory(
workflow_json: Dict[str, Any],
*,
force_rescan: bool = False,
progress_callback: Optional[Callable[[Dict[str, Any]], None]] = None,
analysis_id: Optional[str] = None,
) -> WorkflowModelInventory:
"""
Return the shared base model inventory for a workflow.
Missing Models and Loaded Models use this snapshot so the same unchanged
workflow is not parsed separately by both endpoints. Callers must treat the
returned tuples as read-only.
"""
from ..scanner import get_model_files
def normalize_model_refs(values: List[Any]) -> List[ModelReference]:
if not isinstance(values, (list, tuple)):
raise TypeError("workflow model references must be an array")
return [
value
if isinstance(value, ModelReference)
else ModelReference.from_mapping(value)
for value in values
]
cache_key = _get_workflow_model_inventory_cache_key(workflow_json)
analysis_context = _get_analysis_log_context(cache_key, analysis_id)
promoted_context_key = analysis._get_promoted_widget_context_cache_key(workflow_json)
current_node_fingerprints = analysis._get_workflow_node_fingerprints(workflow_json)
now = time.monotonic()
cached = None
previous_inventory = None
if force_rescan:
invalidate_workflow_model_inventory_cache()
else:
with _WORKFLOW_MODEL_INVENTORY_CACHE_LOCK:
expired_keys = [
key
for key, entry in _WORKFLOW_MODEL_INVENTORY_CACHE.items()
if now - entry["created_at"]
>= _WORKFLOW_MODEL_INVENTORY_CACHE_TTL_SECONDS
]
for key in expired_keys:
_WORKFLOW_MODEL_INVENTORY_CACHE.pop(key, None)
cached = _WORKFLOW_MODEL_INVENTORY_CACHE.get(cache_key)
if cached is not None:
_WORKFLOW_MODEL_INVENTORY_CACHE.move_to_end(cache_key)
else:
best_score = 0
for _, entry in reversed(
list(_WORKFLOW_MODEL_INVENTORY_CACHE.items())
):
if entry.get("promoted_context_key") != promoted_context_key:
continue
node_cache = entry.get("node_cache", {})
score = sum(
1
for key, fingerprint in current_node_fingerprints.items()
if node_cache.get(key, {}).get("fingerprint")
== fingerprint
)
if score > best_score:
best_score = score
previous_inventory = entry
if cached is not None:
model_refs = normalize_model_refs(cached["model_refs"])
if progress_callback:
progress_callback(
{
"stage": "analyzing",
"message": "Reusing shared workflow analysis...",
"current": len(cached.get("node_cache", {})),
"total": len(cached.get("node_cache", {})),
"cached": True,
}
)
log.debug(
f"Reusing shared workflow model analysis ({analysis_context})"
)
return WorkflowModelInventory(
available_models=tuple(cached["available_models"]),
model_refs=tuple(model_refs),
)
node_cache = {}
analysis_stats = {}
if previous_inventory is not None:
available_models = list(previous_inventory["available_models"])
model_refs = analysis.analyze_workflow_models(
workflow_json,
available_models=available_models,
progress_callback=progress_callback,
previous_node_cache=previous_inventory.get("node_cache", {}),
node_cache_out=node_cache,
analysis_stats=analysis_stats,
analysis_context=analysis_context,
)
model_refs = normalize_model_refs(model_refs)
else:
available_models = get_model_files(force_rescan=force_rescan)
model_refs = analysis.analyze_workflow_models(
workflow_json,
available_models=available_models,
progress_callback=progress_callback,
analysis_context=analysis_context,
)
model_refs = normalize_model_refs(model_refs)
node_cache = _build_workflow_node_cache(workflow_json, model_refs)
analysis_stats = {
"total_nodes": len(node_cache),
"reused_nodes": 0,
"analyzed_nodes": len(node_cache),
}
with _WORKFLOW_MODEL_INVENTORY_CACHE_LOCK:
expired_keys = [
key
for key, entry in _WORKFLOW_MODEL_INVENTORY_CACHE.items()
if now - entry["created_at"]
>= _WORKFLOW_MODEL_INVENTORY_CACHE_TTL_SECONDS
]
for key in expired_keys:
_WORKFLOW_MODEL_INVENTORY_CACHE.pop(key, None)
_WORKFLOW_MODEL_INVENTORY_CACHE[cache_key] = {
"created_at": time.monotonic(),
"available_models": available_models,
"model_refs": model_refs,
"node_cache": node_cache,
"promoted_context_key": promoted_context_key,
"analysis_stats": analysis_stats,
}
_WORKFLOW_MODEL_INVENTORY_CACHE.move_to_end(cache_key)
while (
len(_WORKFLOW_MODEL_INVENTORY_CACHE)
> _WORKFLOW_MODEL_INVENTORY_CACHE_MAX_ENTRIES
):
_WORKFLOW_MODEL_INVENTORY_CACHE.popitem(last=False)
return WorkflowModelInventory(
available_models=tuple(available_models),
model_refs=tuple(model_refs),
)