Skip to content

Commit 48bbc46

Browse files
feat(app): handle preparation errors as node errors
We were not handling node preparation errors as node errors before. Here's the explanation, copied from a comment that is no longer required: --- TODO(psyche): Sessions only support errors on nodes, not on the session itself. When an error occurs outside node execution, it bubbles up to the processor where it is treated as a queue item error. Nodes are pydantic models. When we prepare a node in `session.next()`, we set its inputs. This can cause a pydantic validation error. For example, consider a resize image node which has a constraint on its `width` input field - it must be greater than zero. During preparation, if the width is set to zero, pydantic will raise a validation error. When this happens, it breaks the flow before `invocation` is set. We can't set an error on the invocation because we didn't get far enough to get it - we don't know its id. Hence, we just set it as a queue item error. --- This change wraps the node preparation step with exception handling. A new `NodeInputError` exception is raised when there is a validation error. This error has the node (in the state it was in just prior to the error) and an identifier of the input that failed. This allows us to mark the node that failed preparation as errored, correctly making such errors _node_ errors and not _processor_ errors. It's much easier to diagnose these situations. The error messages look like this: > Node b5ac87c6-0678-4b8c-96b9-d215aee12175 has invalid incoming input for height Some of the exception handling logic is cleaned up.
1 parent 434dbb0 commit 48bbc46

File tree

2 files changed

+48
-12
lines changed

2 files changed

+48
-12
lines changed

invokeai/app/services/session_processor/session_processor_default.py

Lines changed: 10 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,7 @@
2121
)
2222
from invokeai.app.services.session_processor.session_processor_common import CanceledException
2323
from invokeai.app.services.session_queue.session_queue_common import SessionQueueItem
24+
from invokeai.app.services.shared.graph import NodeInputError
2425
from invokeai.app.services.shared.invocation_context import InvocationContextData, build_invocation_context
2526
from invokeai.app.util.profiler import Profiler
2627

@@ -66,18 +67,16 @@ def run(self, queue_item: SessionQueueItem):
6667

6768
# Loop over invocations until the session is complete or canceled
6869
while True:
69-
# TODO(psyche): Sessions only support errors on nodes, not on the session itself. When an error occurs outside
70-
# node execution, it bubbles up to the processor where it is treated as a queue item error.
71-
#
72-
# Nodes are pydantic models. When we prepare a node in `session.next()`, we set its inputs. This can cause a
73-
# pydantic validation error. For example, consider a resize image node which has a constraint on its `width`
74-
# input field - it must be greater than zero. During preparation, if the width is set to zero, pydantic will
75-
# raise a validation error.
76-
#
77-
# When this happens, it breaks the flow before `invocation` is set. We can't set an error on the invocation
78-
# because we didn't get far enough to get it - we don't know its id. Hence, we just set it as a queue item error.
70+
try:
71+
invocation = queue_item.session.next()
72+
# Anything other than a `NodeInputError` is handled as a processor error
73+
except NodeInputError as e:
74+
# Must extract the exception traceback here to not lose its stacktrace when we change scope
75+
traceback = e.__traceback__
76+
assert traceback is not None
77+
self._on_node_error(e.node, queue_item, type(e), e, traceback)
78+
break
7979

80-
invocation = queue_item.session.next()
8180
if invocation is None or self._cancel_event.is_set():
8281
break
8382
self.run_node(invocation, queue_item)

invokeai/app/services/shared/graph.py

Lines changed: 38 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,7 @@
88
from pydantic import (
99
BaseModel,
1010
GetJsonSchemaHandler,
11+
ValidationError,
1112
field_validator,
1213
)
1314
from pydantic.fields import Field
@@ -190,6 +191,39 @@ class UnknownGraphValidationError(ValueError):
190191
pass
191192

192193

194+
class NodeInputError(ValueError):
195+
"""Raised when a node fails preparation. This occurs when a node's inputs are being set from its incomers, but an
196+
input fails validation.
197+
198+
Attributes:
199+
node: The node that failed preparation. Note: only successfully set fields will be accurate. Review the error to
200+
determine which field caused the failure.
201+
"""
202+
203+
def __init__(self, node: BaseInvocation, e: ValidationError):
204+
self.original_error = e
205+
self.node = node
206+
# When preparing a node, we set each input one-at-a-time. We may thus safely assume that the first error
207+
# represents the first input that failed.
208+
self.failed_input = loc_to_dot_sep(e.errors()[0]["loc"])
209+
super().__init__(f"Node {node.id} has invalid incoming input for {self.failed_input}")
210+
211+
212+
def loc_to_dot_sep(loc: tuple[Union[str, int], ...]) -> str:
213+
"""Helper to pretty-print pydantic error locations as dot-separated strings.
214+
Taken from https://docs.pydantic.dev/latest/errors/errors/#customize-error-messages
215+
"""
216+
path = ""
217+
for i, x in enumerate(loc):
218+
if isinstance(x, str):
219+
if i > 0:
220+
path += "."
221+
path += x
222+
else:
223+
path += f"[{x}]"
224+
return path
225+
226+
193227
@invocation_output("iterate_output")
194228
class IterateInvocationOutput(BaseInvocationOutput):
195229
"""Used to connect iteration outputs. Will be expanded to a specific output."""
@@ -821,7 +855,10 @@ def next(self) -> Optional[BaseInvocation]:
821855

822856
# Get values from edges
823857
if next_node is not None:
824-
self._prepare_inputs(next_node)
858+
try:
859+
self._prepare_inputs(next_node)
860+
except ValidationError as e:
861+
raise NodeInputError(next_node, e)
825862

826863
# If next is still none, there's no next node, return None
827864
return next_node

0 commit comments

Comments
 (0)