Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
10 changes: 8 additions & 2 deletions api/core/workflow/graph_engine/graph_engine.py
Original file line number Diff line number Diff line change
Expand Up @@ -8,6 +8,7 @@
import contextvars
import logging
import queue
import threading
from collections.abc import Generator
from typing import TYPE_CHECKING, cast, final

Expand Down Expand Up @@ -70,10 +71,13 @@ def __init__(
scale_down_idle_time: float | None = None,
) -> None:
"""Initialize the graph engine with all subsystems and dependencies."""
# stop event
self._stop_event = threading.Event()

# Bind runtime state to current workflow context
self._graph = graph
self._graph_runtime_state = graph_runtime_state
self._graph_runtime_state.stop_event = self._stop_event
self._graph_runtime_state.configure(graph=cast("GraphProtocol", graph))
self._command_channel = command_channel

Expand Down Expand Up @@ -169,6 +173,7 @@ def __init__(
max_workers=self._max_workers,
scale_up_threshold=self._scale_up_threshold,
scale_down_idle_time=self._scale_down_idle_time,
stop_event=self._stop_event,
)

# === Orchestration ===
Expand Down Expand Up @@ -199,6 +204,7 @@ def __init__(
event_handler=self._event_handler_registry,
execution_coordinator=self._execution_coordinator,
event_emitter=self._event_manager,
stop_event=self._stop_event,
)

# === Validation ===
Expand Down Expand Up @@ -319,6 +325,7 @@ def _start_execution(self, *, resume: bool = False) -> None:
paused_nodes: list[str] = []
if resume:
paused_nodes = self._graph_runtime_state.consume_paused_nodes()
self._stop_event.clear()

# Start worker pool (it calculates initial workers internally)
self._worker_pool.start()
Expand All @@ -343,13 +350,12 @@ def _start_execution(self, *, resume: bool = False) -> None:

def _stop_execution(self) -> None:
"""Stop execution subsystems."""
self._stop_event.set()
self._dispatcher.stop()
self._worker_pool.stop()
# Don't mark complete here as the dispatcher already does it

# Notify layers
logger = logging.getLogger(__name__)

for layer in self._layers:
try:
layer.on_graph_end(self._graph_execution.error)
Expand Down
7 changes: 3 additions & 4 deletions api/core/workflow/graph_engine/orchestration/dispatcher.py
Original file line number Diff line number Diff line change
Expand Up @@ -44,6 +44,7 @@ def __init__(
event_queue: queue.Queue[GraphNodeEventBase],
event_handler: "EventHandler",
execution_coordinator: ExecutionCoordinator,
stop_event: threading.Event,
event_emitter: EventManager | None = None,
) -> None:
"""
Expand All @@ -61,24 +62,22 @@ def __init__(
self._event_emitter = event_emitter

self._thread: threading.Thread | None = None
self._stop_event = threading.Event()
self._stop_event = stop_event
self._start_time: float | None = None

def start(self) -> None:
"""Start the dispatcher thread."""
if self._thread and self._thread.is_alive():
return

self._stop_event.clear()
self._start_time = time.time()
self._thread = threading.Thread(target=self._dispatcher_loop, name="GraphDispatcher", daemon=True)
self._thread.start()

def stop(self) -> None:

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Since Dispatcher.stop() no longer sets _stop_event, callers that invoke stop() without first setting the shared event may leave the dispatcher thread running (the join() just times out).

Fix This in Augment

🤖 Was this useful? React with 👍 or 👎

"""Stop the dispatcher thread."""
self._stop_event.set()
if self._thread and self._thread.is_alive():
self._thread.join(timeout=10.0)
self._thread.join(timeout=2.0)

def _dispatcher_loop(self) -> None:
"""Main dispatcher loop."""
Expand Down
10 changes: 7 additions & 3 deletions api/core/workflow/graph_engine/worker.py
Original file line number Diff line number Diff line change
Expand Up @@ -42,6 +42,7 @@ def __init__(
event_queue: queue.Queue[GraphNodeEventBase],
graph: Graph,
layers: Sequence[GraphEngineLayer],
stop_event: threading.Event,
worker_id: int = 0,
flask_app: Flask | None = None,
context_vars: contextvars.Context | None = None,
Expand All @@ -65,13 +66,16 @@ def __init__(
self._worker_id = worker_id
self._flask_app = flask_app
self._context_vars = context_vars
self._stop_event = threading.Event()
self._last_task_time = time.time()
self._stop_event = stop_event
self._layers = layers if layers is not None else []

def stop(self) -> None:

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

With Worker.stop() now a no-op, WorkerPool._remove_worker() (used for scale-down) can no longer actually terminate an individual worker, so “removed” workers may keep running and processing tasks outside the pool’s accounting.

Fix This in Augment

🤖 Was this useful? React with 👍 or 👎

"""Signal the worker to stop processing."""
self._stop_event.set()
"""Worker is controlled via shared stop_event from GraphEngine.

This method is a no-op retained for backward compatibility.
"""
pass

@property
def is_idle(self) -> bool:
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -41,6 +41,7 @@ def __init__(
event_queue: queue.Queue[GraphNodeEventBase],
graph: Graph,
layers: list[GraphEngineLayer],
stop_event: threading.Event,
flask_app: "Flask | None" = None,
context_vars: "Context | None" = None,
min_workers: int | None = None,
Expand Down Expand Up @@ -81,6 +82,7 @@ def __init__(
self._worker_counter = 0
self._lock = threading.RLock()
self._running = False
self._stop_event = stop_event

# No longer tracking worker states with callbacks to avoid lock contention

Expand Down Expand Up @@ -128,14 +130,10 @@ def stop(self) -> None:
if worker_count > 0:
logger.debug("Stopping worker pool: %d workers", worker_count)

# Stop all workers
for worker in self._workers:
worker.stop()

# Wait for workers to finish
for worker in self._workers:
if worker.is_alive():
worker.join(timeout=10.0)
worker.join(timeout=2.0)

self._workers.clear()

Expand All @@ -152,6 +150,7 @@ def _create_worker(self) -> None:
worker_id=worker_id,
flask_app=self._flask_app,
context_vars=self._context_vars,
stop_event=self._stop_event,
)

worker.start()
Expand Down
19 changes: 19 additions & 0 deletions api/core/workflow/nodes/base/node.py
Original file line number Diff line number Diff line change
Expand Up @@ -264,6 +264,10 @@ def _run(self) -> NodeRunResult | Generator[NodeEventBase, None, None]:
"""
raise NotImplementedError

def _should_stop(self) -> bool:
"""Check if execution should be stopped."""
return self.graph_runtime_state.stop_event and self.graph_runtime_state.stop_event.is_set()

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

_should_stop() can return None when graph_runtime_state.stop_event is unset, despite being annotated -> bool; that can leak non-bool values into callers/tests.

Fix This in Augment

🤖 Was this useful? React with 👍 or 👎


def run(self) -> Generator[GraphNodeEventBase, None, None]:
execution_id = self.ensure_execution_id()
self._start_at = naive_utc_now()
Expand Down Expand Up @@ -332,6 +336,21 @@ def run(self) -> Generator[GraphNodeEventBase, None, None]:
yield event
else:
yield event

if self._should_stop():
error_message = "Execution cancelled"
yield NodeRunFailedEvent(
id=self.execution_id,
node_id=self._node_id,
node_type=self.node_type,
start_at=self._start_at,
node_run_result=NodeRunResult(
status=WorkflowNodeExecutionStatus.FAILED,

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Cancellation emits a NodeRunFailedEvent with WorkflowNodeExecutionStatus.FAILED; since there is also a STOPPED status, marking cancelled nodes as failed may misclassify user stops in UI/metrics.

Fix This in Augment

🤖 Was this useful? React with 👍 or 👎

error=error_message,
),
error=error_message,
)
return
except Exception as e:
logger.exception("Node %s failed to run", self._node_id)
result = NodeRunResult(
Expand Down
2 changes: 2 additions & 0 deletions api/core/workflow/runtime/graph_runtime_state.py
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,7 @@

import importlib
import json
import threading
from collections.abc import Mapping, Sequence
from copy import deepcopy
from dataclasses import dataclass
Expand Down Expand Up @@ -168,6 +169,7 @@ def __init__(
self._pending_response_coordinator_dump: str | None = None
self._pending_graph_execution_workflow_id: str | None = None
self._paused_nodes: set[str] = set()
self.stop_event: threading.Event | None = None

if graph is not None:
self.attach_graph(graph)
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -3,6 +3,7 @@
from __future__ import annotations

import queue
import threading
from unittest import mock

from core.workflow.entities.pause_reason import SchedulingPause
Expand Down Expand Up @@ -36,6 +37,7 @@ def test_dispatcher_should_consume_remains_events_after_pause():
event_queue=event_queue,
event_handler=event_handler,
execution_coordinator=execution_coordinator,
stop_event=threading.Event(),
)
dispatcher._dispatcher_loop()
assert event_queue.empty()
Expand Down Expand Up @@ -96,6 +98,7 @@ def _run_dispatcher_for_event(event) -> int:
event_queue=event_queue,
event_handler=event_handler,
execution_coordinator=coordinator,
stop_event=threading.Event(),
)

dispatcher._dispatcher_loop()
Expand Down Expand Up @@ -181,6 +184,7 @@ def test_dispatcher_drain_event_queue():
event_queue=event_queue,
event_handler=event_handler,
execution_coordinator=coordinator,
stop_event=threading.Event(),
)

dispatcher._dispatcher_loop()
Expand Down
Loading