vellum-workflow-server 1.7.2__tar.gz → 1.7.2.post2__tar.gz

{vellum_workflow_server-1.7.2 → vellum_workflow_server-1.7.2.post2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: vellum-workflow-server
-Version: 1.7.2
+Version: 1.7.2.post2
 Summary: 
 License: AGPL
 Requires-Python: >=3.9.0,<4

{vellum_workflow_server-1.7.2 → vellum_workflow_server-1.7.2.post2}/pyproject.toml

@@ -3,7 +3,7 @@ name = "vellum-workflow-server"
 
 [tool.poetry]
 name = "vellum-workflow-server"
-version = "1.7.2"
+version = "1.7.2.post2"
 description = ""
 readme = "README.md"
 authors = []

{vellum_workflow_server-1.7.2 → vellum_workflow_server-1.7.2.post2}/src/workflow_server/api/tests/test_workflow_view_stream_workflow_route.py

@@ -571,6 +571,65 @@ class BasicCancellableWorkflow(BaseWorkflow):
         )
 
 
+def test_stream_workflow_route__timeout_emits_rejection_events():
+    """
+    Tests that when a workflow times out, we emit node and workflow rejection events.
+    """
+
+    span_id = uuid4()
+    request_body = {
+        "timeout": 1,
+        "execution_id": str(span_id),
+        "inputs": [],
+        "environment_api_key": "test",
+        "module": "workflow",
+        "files": {
+            "__init__.py": "",
+            "workflow.py": """\
+import time
+
+from vellum.workflows.nodes.bases.base import BaseNode
+from vellum.workflows.workflows.base import BaseWorkflow
+
+
+class LongRunningNode(BaseNode):
+    class Outputs(BaseNode.Outputs):
+        value: str
+
+    def run(self) -> Outputs:
+        time.sleep(30)
+        return self.Outputs(value="hello world")
+
+
+class TimeoutWorkflow(BaseWorkflow):
+    graph = LongRunningNode
+    class Outputs(BaseWorkflow.Outputs):
+        final_value = LongRunningNode.Outputs.value
+
+""",
+        },
+    }
+
+    status_code, events = flask_stream(request_body)
+
+    assert status_code == 200
+
+    event_names = [e["name"] for e in events]
+
+    assert "vembda.execution.initiated" in event_names
+    assert "workflow.execution.initiated" in event_names
+    assert "node.execution.initiated" in event_names
+
+    # TODO: Re-enable once solved SDK-side
+    # assert "node.execution.rejected" in event_names, "Should emit node.execution.rejected on timeout"
+
+    assert "workflow.execution.rejected" in event_names, "Should emit workflow.execution.rejected on timeout"
+
+    assert "vembda.execution.fulfilled" in event_names
+    vembda_fulfilled = next(e for e in events if e["name"] == "vembda.execution.fulfilled")
+    assert vembda_fulfilled["body"]["timed_out"] is True
+
+
 def test_stream_workflow_route__very_large_events(both_stream_types):
     # GIVEN a valid request body
     span_id = uuid4()

{vellum_workflow_server-1.7.2 → vellum_workflow_server-1.7.2.post2}/src/workflow_server/api/workflow_view.py

@@ -3,12 +3,11 @@ import importlib
 import inspect
 import json
 import logging
-from multiprocessing import Process, Queue, set_start_method
+from multiprocessing import Event as MultiprocessingEvent, Process, Queue, set_start_method
 import os
 import pkgutil
 from queue import Empty
 import sys
-from threading import Event as ThreadingEvent
 import time
 import traceback
 from uuid import uuid4
@@ -123,8 +122,8 @@ def stream_workflow_route() -> Response:
             headers=headers,
         )
 
-    cancel_signal = ThreadingEvent()
-    timeout_signal = ThreadingEvent()
+    cancel_signal = MultiprocessingEvent()
+    timeout_signal = MultiprocessingEvent()
 
     process: Optional[Process] = None
     if ENABLE_PROCESS_WRAPPER:

{vellum_workflow_server-1.7.2 → vellum_workflow_server-1.7.2.post2}/src/workflow_server/code_exec_runner.py

@@ -50,6 +50,7 @@ def run_code_exec_stream() -> None:
             disable_redirect=True,
             # Timeouts are handled at the code exec level right now so just passing in an unused threading event
             timeout_signal=ThreadingEvent(),
+            cancel_signal=ThreadingEvent(),
         )
         for line in stream_iterator:
             print(f"{_EVENT_LINE}{json.dumps(line)}")  # noqa: T201

{vellum_workflow_server-1.7.2 → vellum_workflow_server-1.7.2.post2}/src/workflow_server/core/cancel_workflow.py

@@ -1,11 +1,13 @@
 import logging
-from threading import Event, Event as ThreadingEvent, Thread
+from threading import Thread
 import time
 from uuid import UUID
 from typing import Optional
 
 import requests
 
+from workflow_server.core.utils import CancelSignal
+
 _TIMER_INTERVAL = 5
 
 logger = logging.getLogger(__name__)
@@ -23,18 +25,18 @@ def get_is_workflow_cancelled(execution_id: UUID, vembda_public_url: Optional[st
 
 
 class CancelWorkflowWatcherThread(Thread):
-    _kill_switch: Event
+    _kill_switch: CancelSignal
     _execution_id: UUID
-    _cancel_signal: ThreadingEvent
+    _cancel_signal: CancelSignal
     _vembda_public_url: Optional[str]
     _start = time.time()
 
     def __init__(
         self,
-        kill_switch: Event,
+        kill_switch: CancelSignal,
         execution_id: UUID,
         vembda_public_url: Optional[str],
-        cancel_signal: ThreadingEvent,
+        cancel_signal: CancelSignal,
         timeout_seconds: int,
     ) -> None:
         Thread.__init__(self)

{vellum_workflow_server-1.7.2 → vellum_workflow_server-1.7.2.post2}/src/workflow_server/core/executor.py

@@ -8,7 +8,6 @@ import os
 import random
 import string
 import sys
-import threading
 from threading import Event as ThreadingEvent
 import time
 from traceback import format_exc
@@ -43,6 +42,7 @@ from workflow_server.core.events import (
     VembdaExecutionFulfilledEvent,
 )
 from workflow_server.core.utils import (
+    CancelSignal,
     create_vembda_rejected_event,
     is_events_emitting_enabled,
     serialize_vembda_rejected_event,
@@ -87,8 +87,8 @@ def _stream_node_wrapper(executor_context: NodeExecutorContext, queue: Queue) ->
 def _stream_workflow_wrapper(
     executor_context: WorkflowExecutorContext,
     queue: Queue,
-    cancel_signal: Optional[ThreadingEvent],
-    timeout_signal: ThreadingEvent,
+    cancel_signal: CancelSignal,
+    timeout_signal: CancelSignal,
 ) -> None:
     span_id_emitted = False
     try:
@@ -119,8 +119,8 @@ def _stream_workflow_wrapper(
 def stream_workflow_process_timeout(
     executor_context: WorkflowExecutorContext,
     queue: Queue,
-    cancel_signal: ThreadingEvent,
-    timeout_signal: ThreadingEvent,
+    cancel_signal: CancelSignal,
+    timeout_signal: CancelSignal,
 ) -> Process:
     workflow_process = Process(
         target=_stream_workflow_wrapper,
@@ -141,9 +141,9 @@ def stream_workflow_process_timeout(
 
 def stream_workflow(
     executor_context: WorkflowExecutorContext,
-    timeout_signal: ThreadingEvent,
+    timeout_signal: CancelSignal,
+    cancel_signal: CancelSignal,
     disable_redirect: bool = True,
-    cancel_signal: Optional[ThreadingEvent] = None,
 ) -> tuple[Iterator[dict], UUID]:
     cancel_watcher_kill_switch = ThreadingEvent()
     try:
@@ -172,7 +172,8 @@ def stream_workflow(
             state=workflow_state,
             node_output_mocks=node_output_mocks,
             event_filter=all_workflow_event_filter,
-            cancel_signal=cancel_signal,
+            # TODO: Fix type in SDK
+            cancel_signal=cancel_signal,  # type: ignore[arg-type]
             entrypoint_nodes=[run_from_node] if run_from_node else None,
             previous_execution_id=executor_context.previous_execution_id,
         )
@@ -235,12 +236,6 @@ def stream_workflow(
         finally:
             cancel_watcher_kill_switch.set()
 
-        emitter_thread = next(
-            (t for t in threading.enumerate() if t.name.endswith(".background_thread") and t.is_alive()), None
-        )
-        if emitter_thread:
-            emitter_thread.join()
-
         workflow.join()
 
     return (
@@ -298,7 +293,7 @@ def stream_node(
 def _call_stream(
     executor_context: BaseExecutorContext,
     stream_generator: Callable[[], Generator[dict[str, Any], Any, None]],
-    timeout_signal: ThreadingEvent,
+    timeout_signal: CancelSignal,
     disable_redirect: bool = True,
 ) -> Iterator[dict]:
     log_redirect: Optional[StringIO] = None

{vellum_workflow_server-1.7.2 → vellum_workflow_server-1.7.2.post2}/src/workflow_server/core/utils.py

@@ -1,10 +1,14 @@
 from datetime import datetime
+from multiprocessing.synchronize import Event as MultiprocessingEvent
+from threading import Event as ThreadingEvent
 from uuid import uuid4
-from typing import Optional
+from typing import Optional, Union
 
 from workflow_server.core.events import VembdaExecutionFulfilledBody, VembdaExecutionFulfilledEvent
 from workflow_server.core.workflow_executor_context import BaseExecutorContext
 
+CancelSignal = Union[MultiprocessingEvent, ThreadingEvent]
+
 
 def _create_vembda_rejected_event_base(
     executor_context: Optional[BaseExecutorContext], error_message: str, timed_out: bool