216 lines
7.7 KiB
Python
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),
|
|
)
|