diff --git a/api/core/app/apps/advanced_chat/app_runner.py b/api/core/app/apps/advanced_chat/app_runner.py index cf3943a6ccc..4b303663e08 100644 --- a/api/core/app/apps/advanced_chat/app_runner.py +++ b/api/core/app/apps/advanced_chat/app_runner.py @@ -11,7 +11,9 @@ from core.app.apps.base_app_queue_manager import AppQueueManager from core.app.apps.workflow.command_channels import ( CelerySignalCommandChannel, CombinedCommandChannel, + StopFlagCommandChannel, ) +from core.app.apps.workflow.stop_aware_ready_queue import attach_stop_aware_ready_queue from core.app.apps.workflow_app_runner import WorkflowBasedAppRunner from core.app.entities.app_invoke_entities import ( AdvancedChatAppGenerateEntity, @@ -229,9 +231,11 @@ class AdvancedChatAppRunner(WorkflowBasedAppRunner): shutdown_state_getter=celery_warm_shutdown_started, abort_reason=WORKFLOW_WARM_SHUTDOWN_ABORT_REASON, ) + attach_stop_aware_ready_queue(graph_runtime_state, task_id=task_id) command_channel = CombinedCommandChannel( ( RedisChannel(redis_client, channel_key), + StopFlagCommandChannel(task_id=task_id), celery_signal_channel, ) ) diff --git a/api/core/app/apps/execution_coordinator.py b/api/core/app/apps/execution_coordinator.py index 1372d55155e..e0fb9d36c65 100644 --- a/api/core/app/apps/execution_coordinator.py +++ b/api/core/app/apps/execution_coordinator.py @@ -26,11 +26,24 @@ def app_task_command_channel_key(task_id: str) -> str: return f"workflow:{task_id}:commands" +def app_task_stop_flag_key(task_id: str) -> str: + """Redis key of the legacy generate-task stop flag.""" + return f"generate_task_stopped:{task_id}" + + def set_app_task_stop_flag(task_id: str) -> None: if not task_id: return - redis_client.setex(f"generate_task_stopped:{task_id}", 600, 1) + redis_client.setex(app_task_stop_flag_key(task_id), 600, 1) + + +def is_app_task_stop_flag_set(task_id: str) -> bool: + """Return whether the legacy Redis stop flag is currently armed.""" + if not task_id: + return False + + return redis_client.get(app_task_stop_flag_key(task_id)) is not None def clear_app_task_cancellation_signals(task_id: str) -> None: @@ -47,7 +60,7 @@ def clear_app_task_cancellation_signals(task_id: str) -> None: return try: - redis_client.delete(f"generate_task_stopped:{task_id}") + redis_client.delete(app_task_stop_flag_key(task_id)) except Exception: logger.exception("Failed to clear stop flag for app task %s", task_id) diff --git a/api/core/app/apps/workflow/app_runner.py b/api/core/app/apps/workflow/app_runner.py index 1ea7b26da29..b2a3b9dd77a 100644 --- a/api/core/app/apps/workflow/app_runner.py +++ b/api/core/app/apps/workflow/app_runner.py @@ -9,7 +9,9 @@ from core.app.apps.workflow.app_config_manager import WorkflowAppConfig from core.app.apps.workflow.command_channels import ( CelerySignalCommandChannel, CombinedCommandChannel, + StopFlagCommandChannel, ) +from core.app.apps.workflow.stop_aware_ready_queue import attach_stop_aware_ready_queue from core.app.apps.workflow_app_runner import WorkflowBasedAppRunner from core.app.entities.app_invoke_entities import DifyRunContext, InvokeFrom, WorkflowAppGenerateEntity from core.app.workflow.layers.persistence import PersistenceWorkflowInfo, WorkflowPersistenceLayer @@ -159,9 +161,11 @@ class WorkflowAppRunner(WorkflowBasedAppRunner): shutdown_state_getter=celery_warm_shutdown_started, abort_reason=WORKFLOW_WARM_SHUTDOWN_ABORT_REASON, ) + attach_stop_aware_ready_queue(graph_runtime_state, task_id=task_id) command_channel = CombinedCommandChannel( ( RedisChannel(redis_client, channel_key), + StopFlagCommandChannel(task_id=task_id), celery_signal_channel, ) ) diff --git a/api/core/app/apps/workflow/command_channels.py b/api/core/app/apps/workflow/command_channels.py index 526476f8c24..6dab7bebbf4 100644 --- a/api/core/app/apps/workflow/command_channels.py +++ b/api/core/app/apps/workflow/command_channels.py @@ -4,6 +4,7 @@ import logging from collections.abc import Callable, Sequence from typing import final, override +from core.app.apps.execution_coordinator import is_app_task_stop_flag_set from graphon.graph_engine.command_channels import CommandChannel from graphon.graph_engine.entities.commands import AbortCommand, GraphEngineCommand @@ -65,3 +66,34 @@ class CelerySignalCommandChannel(CommandChannel): @override def send_command(self, command: GraphEngineCommand) -> None: _ = command + + +@final +class StopFlagCommandChannel(CommandChannel): + """Translate the legacy Redis stop flag into one GraphEngine abort command.""" + + _task_id: str + _abort_reason: str + _abort_emitted: bool + + def __init__( + self, + *, + task_id: str, + abort_reason: str = "User requested stop", + ) -> None: + self._task_id = task_id + self._abort_reason = abort_reason + self._abort_emitted = False + + @override + def fetch_commands(self) -> list[GraphEngineCommand]: + if self._abort_emitted or not is_app_task_stop_flag_set(self._task_id): + return [] + + self._abort_emitted = True + return [AbortCommand(reason=self._abort_reason)] + + @override + def send_command(self, command: GraphEngineCommand) -> None: + _ = command diff --git a/api/core/app/apps/workflow/stop_aware_ready_queue.py b/api/core/app/apps/workflow/stop_aware_ready_queue.py new file mode 100644 index 00000000000..dadc402da1e --- /dev/null +++ b/api/core/app/apps/workflow/stop_aware_ready_queue.py @@ -0,0 +1,67 @@ +"""Ready-queue wrapper that honors stop before the next node is scheduled.""" + +from __future__ import annotations + +from typing import final + +from core.app.apps.execution_coordinator import is_app_task_stop_flag_set +from graphon.graph_engine.ready_queue import ReadyQueue, ReadyTask +from graphon.runtime.graph_runtime_state import GraphExecutionProtocol, GraphRuntimeState + + +@final +class StopAwareReadyQueue: + """Reject newly ready nodes once the run has been stopped. + + GraphEngine drain still enqueues successors after abort. Drop those puts so + later nodes do not start; the in-flight node can finish. + """ + + def __init__( + self, + inner: ReadyQueue, + *, + task_id: str, + graph_execution: GraphExecutionProtocol, + ) -> None: + self._inner = inner + self._task_id = task_id + self._graph_execution = graph_execution + + def _should_reject(self) -> bool: + return self._graph_execution.aborted or is_app_task_stop_flag_set(self._task_id) + + def put(self, item: ReadyTask) -> None: + if self._should_reject(): + return + self._inner.put(item) + + def get(self, timeout: float | None = None) -> ReadyTask: + return self._inner.get(timeout=timeout) + + def task_done(self) -> None: + self._inner.task_done() + + def qsize(self) -> int: + return self._inner.qsize() + + def drain(self) -> list[ReadyTask]: + return self._inner.drain() + + def dumps(self) -> str: + return self._inner.dumps() + + def loads(self, data: str) -> None: + self._inner.loads(data) + + +def attach_stop_aware_ready_queue(graph_runtime_state: GraphRuntimeState, *, task_id: str) -> None: + """Install stop-aware enqueue policy on an existing runtime state.""" + current = graph_runtime_state.ready_queue + if isinstance(current, StopAwareReadyQueue): + return + graph_runtime_state._ready_queue = StopAwareReadyQueue( + current, + task_id=task_id, + graph_execution=graph_runtime_state.graph_execution, + ) diff --git a/api/tests/unit_tests/core/app/apps/test_execution_coordinator.py b/api/tests/unit_tests/core/app/apps/test_execution_coordinator.py index b926e694202..5ab95b961cb 100644 --- a/api/tests/unit_tests/core/app/apps/test_execution_coordinator.py +++ b/api/tests/unit_tests/core/app/apps/test_execution_coordinator.py @@ -6,7 +6,9 @@ from core.app.apps.execution_coordinator import ( AppExecutionCoordinator, AppExecutionState, app_task_command_channel_key, + app_task_stop_flag_key, clear_app_task_cancellation_signals, + is_app_task_stop_flag_set, ) @@ -143,3 +145,22 @@ def test_stop_flag_failure_does_not_block_graph_stop(caplog: pytest.LogCaptureFi reason="test abort", ) assert "Failed to set stop flag for app execution task=task" in caplog.text + + +def test_stop_flag_key_matches_legacy_redis_flag() -> None: + assert app_task_stop_flag_key("task") == "generate_task_stopped:task" + + +def test_is_app_task_stop_flag_set_reads_redis() -> None: + with patch("core.app.apps.execution_coordinator.redis_client") as redis_client: + redis_client.get.return_value = b"1" + assert is_app_task_stop_flag_set("task") is True + redis_client.get.assert_called_once_with("generate_task_stopped:task") + + +def test_is_app_task_stop_flag_set_is_false_when_missing_or_empty() -> None: + with patch("core.app.apps.execution_coordinator.redis_client") as redis_client: + redis_client.get.return_value = None + assert is_app_task_stop_flag_set("task") is False + assert is_app_task_stop_flag_set("") is False + redis_client.get.assert_called_once_with("generate_task_stopped:task") diff --git a/api/tests/unit_tests/core/app/apps/workflow/test_command_channels.py b/api/tests/unit_tests/core/app/apps/workflow/test_command_channels.py index 1a7e7700358..30d0afe6cfa 100644 --- a/api/tests/unit_tests/core/app/apps/workflow/test_command_channels.py +++ b/api/tests/unit_tests/core/app/apps/workflow/test_command_channels.py @@ -1,12 +1,14 @@ from __future__ import annotations from types import SimpleNamespace +from unittest.mock import patch import pytest from core.app.apps.workflow.command_channels import ( CelerySignalCommandChannel, CombinedCommandChannel, + StopFlagCommandChannel, ) from graphon.graph_engine.entities.commands import AbortCommand, PauseCommand @@ -105,3 +107,35 @@ def test_celery_signal_command_channel_send_command_is_noop() -> None: channel.send_command(command) assert channel.fetch_commands() == [] + + +def test_stop_flag_command_channel_emits_abort_when_flag_is_set() -> None: + channel = StopFlagCommandChannel(task_id="task-1", abort_reason="User requested stop") + + with patch("core.app.apps.workflow.command_channels.is_app_task_stop_flag_set", return_value=False): + assert channel.fetch_commands() == [] + + with patch("core.app.apps.workflow.command_channels.is_app_task_stop_flag_set", return_value=True): + commands = channel.fetch_commands() + + assert len(commands) == 1 + assert isinstance(commands[0], AbortCommand) + assert commands[0].reason == "User requested stop" + + +def test_stop_flag_command_channel_emits_abort_once_per_instance() -> None: + channel = StopFlagCommandChannel(task_id="task-1") + + with patch("core.app.apps.workflow.command_channels.is_app_task_stop_flag_set", return_value=True): + assert len(channel.fetch_commands()) == 1 + assert channel.fetch_commands() == [] + + +def test_stop_flag_command_channel_send_command_is_noop() -> None: + channel = StopFlagCommandChannel(task_id="task-1") + command = PauseCommand(reason="pause") + + channel.send_command(command) + + with patch("core.app.apps.workflow.command_channels.is_app_task_stop_flag_set", return_value=False): + assert channel.fetch_commands() == [] diff --git a/api/tests/unit_tests/core/app/apps/workflow/test_stop_aware_ready_queue.py b/api/tests/unit_tests/core/app/apps/workflow/test_stop_aware_ready_queue.py new file mode 100644 index 00000000000..90737cd25ff --- /dev/null +++ b/api/tests/unit_tests/core/app/apps/workflow/test_stop_aware_ready_queue.py @@ -0,0 +1,98 @@ +from __future__ import annotations + +from unittest.mock import Mock, patch + +from core.app.apps.workflow.stop_aware_ready_queue import ( + StopAwareReadyQueue, + attach_stop_aware_ready_queue, +) +from graphon.graph_engine.domain.graph_execution import GraphExecution +from graphon.graph_engine.ready_queue import StartTask +from graphon.runtime.graph_runtime_state import GraphRuntimeState +from graphon.runtime.variable_pool import VariablePool + + +def _start_task(node_id: str = "next-node") -> StartTask: + return StartTask(frame_id="root", node_id=node_id) + + +def _graph_execution(*, aborted: bool = False) -> GraphExecution: + return GraphExecution(workflow_id="workflow-1", aborted=aborted) + + +def test_ready_queue_accepts_work_while_run_is_active() -> None: + inner = Mock() + queue = StopAwareReadyQueue(inner, task_id="task-1", graph_execution=_graph_execution()) + task = _start_task() + + with patch( + "core.app.apps.workflow.stop_aware_ready_queue.is_app_task_stop_flag_set", + return_value=False, + ): + queue.put(task) + + inner.put.assert_called_once_with(task) + + +def test_ready_queue_rejects_next_node_after_graph_abort() -> None: + inner = Mock() + queue = StopAwareReadyQueue(inner, task_id="task-1", graph_execution=_graph_execution(aborted=True)) + + with patch( + "core.app.apps.workflow.stop_aware_ready_queue.is_app_task_stop_flag_set", + return_value=False, + ): + queue.put(_start_task()) + + inner.put.assert_not_called() + + +def test_ready_queue_rejects_next_node_when_stop_flag_is_set() -> None: + inner = Mock() + queue = StopAwareReadyQueue(inner, task_id="task-1", graph_execution=_graph_execution()) + + with patch( + "core.app.apps.workflow.stop_aware_ready_queue.is_app_task_stop_flag_set", + return_value=True, + ) as stop_flag: + queue.put(_start_task("code-node")) + + stop_flag.assert_called_once_with("task-1") + inner.put.assert_not_called() + + +def test_attach_stop_aware_ready_queue_wraps_once() -> None: + inner = Mock() + runtime_state = GraphRuntimeState( + variable_pool=VariablePool(), + start_at=0, + ready_queue=inner, + graph_execution=_graph_execution(), + ) + + attach_stop_aware_ready_queue(runtime_state, task_id="task-1") + first = runtime_state.ready_queue + attach_stop_aware_ready_queue(runtime_state, task_id="task-1") + + assert isinstance(first, StopAwareReadyQueue) + assert runtime_state.ready_queue is first + + +def test_stop_aware_queue_delegates_reads() -> None: + inner = Mock() + inner.get.return_value = _start_task("queued") + inner.qsize.return_value = 1 + inner.drain.return_value = [_start_task("queued")] + inner.dumps.return_value = "{}" + queue = StopAwareReadyQueue(inner, task_id="task-1", graph_execution=_graph_execution()) + + assert queue.get(timeout=0.1) == _start_task("queued") + queue.task_done() + assert queue.qsize() == 1 + assert queue.drain() == [_start_task("queued")] + assert queue.dumps() == "{}" + queue.loads("{}") + + inner.get.assert_called_once_with(timeout=0.1) + inner.task_done.assert_called_once() + inner.loads.assert_called_once_with("{}")