Fix stream mode not respected in subgraphs

- The default applied for subgraphs should be applied only when stream mode arg not passed in
This commit is contained in:
Nuno Campos
2025-05-23 16:40:51 -07:00
parent 8fb91569b9
commit 36478eb745
+6 -4
View File
@@ -2214,12 +2214,14 @@ class Pregel(PregelProtocol):
validate_keys(output_keys, self.channels)
interrupt_before = interrupt_before or self.interrupt_before_nodes
interrupt_after = interrupt_after or self.interrupt_after_nodes
stream_mode = stream_mode if stream_mode is not None else self.stream_mode
if stream_mode is None and CONFIG_KEY_TASK_ID in config.get(CONF, {}):
# if being called as a node in another graph, default to values mode
# but don't overwrite stream_mode arg if provided
stream_mode = ["values"]
elif stream_mode is None:
stream_mode = self.stream_mode
if not isinstance(stream_mode, list):
stream_mode = [stream_mode]
if CONFIG_KEY_TASK_ID in config.get(CONF, {}):
# if being called as a node in another graph, always use values mode
stream_mode = ["values"]
if self.checkpointer is False:
checkpointer: BaseCheckpointSaver | None = None
elif CONFIG_KEY_CHECKPOINTER in config.get(CONF, {}):