81 lines
3.1 KiB
Python
81 lines
3.1 KiB
Python
from __future__ import annotations
|
|
from contextlib import ContextDecorator, contextmanager, nullcontext
|
|
from typing import Any
|
|
|
|
from .formatter import _format_node_description, get_node_input_names
|
|
from .proxy import NodeErrorReportingProxy
|
|
|
|
|
|
def attach_node_failure_info(node: Any, exc: Exception, fallback_inputs: list[str] | None = None) -> Exception:
|
|
"""Enriches an exception in-place with the failing node and input socket names, preventing duplicate wrapping."""
|
|
if getattr(exc, "_nodeslicer_node", None) is not None:
|
|
return exc
|
|
try:
|
|
setattr(exc, "_nodeslicer_node", node)
|
|
input_names = get_node_input_names(node, fallback_inputs=fallback_inputs)
|
|
setattr(exc, "_nodeslicer_inputs", input_names)
|
|
|
|
node_str = _format_node_description(node) or "Unknown Node"
|
|
inputs_str = ", ".join(input_names) if input_names else "(None)"
|
|
prefix = f"Execution failed at node {node_str} [Inputs: {inputs_str}]: "
|
|
|
|
if exc.args:
|
|
orig_msg = str(exc.args[0])
|
|
if not orig_msg.startswith("Execution failed at node "):
|
|
exc.args = (prefix + orig_msg, *exc.args[1:])
|
|
else:
|
|
exc.args = (prefix + type(exc).__name__,)
|
|
|
|
print(f"[NodeSlicer ERROR] {str(exc)}")
|
|
except Exception:
|
|
pass
|
|
return exc
|
|
|
|
|
|
class NodeErrorHandler(ContextDecorator):
|
|
"""Context manager and decorator that intercepts exceptions in node execution
|
|
and enriches them with node failure info (node name, inputs) without intruding
|
|
into the node's internal logic."""
|
|
|
|
def __init__(self, node: Any, fallback_inputs: list[str] | None = None):
|
|
self.node = node
|
|
self.fallback_inputs = fallback_inputs
|
|
|
|
def __enter__(self):
|
|
return self
|
|
|
|
def __exit__(self, exc_type, exc_val, exc_tb):
|
|
if exc_val is not None and isinstance(exc_val, Exception):
|
|
attach_node_failure_info(self.node, exc_val, fallback_inputs=self.fallback_inputs)
|
|
return False # Do not suppress exception
|
|
|
|
|
|
def handle_node_errors(node: Any, fallback_inputs: list[str] | None = None) -> NodeErrorHandler:
|
|
"""Creates a modular context decorator to handle and enrich node execution errors."""
|
|
return NodeErrorHandler(node, fallback_inputs=fallback_inputs)
|
|
|
|
|
|
@contextmanager
|
|
def node_execution_scope(
|
|
node: Any,
|
|
operator: Any = None,
|
|
fallback_inputs: list[str] | None = None,
|
|
):
|
|
"""Modular, non-intrusive execution scope that combines profiling (if active),
|
|
operator error reporting proxy (if operator provided), and exception enrichment
|
|
without tightly coupling any of them into node classes."""
|
|
try:
|
|
from ..profiling import profile_node_execution
|
|
prof_ctx = profile_node_execution(node)
|
|
except Exception:
|
|
prof_ctx = nullcontext()
|
|
|
|
wrapped_op = NodeErrorReportingProxy(operator, node) if operator is not None else None
|
|
|
|
class ExecutionScope:
|
|
def __init__(self, op: Any):
|
|
self.operator = op
|
|
|
|
with prof_ctx:
|
|
with handle_node_errors(node, fallback_inputs=fallback_inputs):
|
|
yield ExecutionScope(wrapped_op)
|