vellum-workflow-server 1.4.0__py3-none-any.whl → 1.4.1.post1__py3-none-any.whl

{vellum_workflow_server-1.4.0.dist-info → vellum_workflow_server-1.4.1.post1.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: vellum-workflow-server
-Version: 1.4.0
+Version: 1.4.1.post1
 Summary: 
 License: AGPL
 Requires-Python: >=3.9.0,<4
@@ -29,7 +29,7 @@ Requires-Dist: pyjwt (==2.10.0)
 Requires-Dist: python-dotenv (==1.0.1)
 Requires-Dist: retrying (==1.3.4)
 Requires-Dist: sentry-sdk[flask] (==2.20.0)
-Requires-Dist: vellum-ai (==1.4.0)
+Requires-Dist: vellum-ai (==1.4.1)
 Description-Content-Type: text/markdown
 
 # Vellum Workflow Runner Server

{vellum_workflow_server-1.4.0.dist-info → vellum_workflow_server-1.4.1.post1.dist-info}/RECORD

@@ -6,13 +6,13 @@ workflow_server/api/tests/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJW
 workflow_server/api/tests/test_input_display_mapping.py,sha256=drBZqMudFyB5wgiUOcMgRXz7E7ge-Qgxbstw4E4f0zE,2211
 workflow_server/api/tests/test_workflow_view.py,sha256=RlAw1tHeIlnOXGrFQN-w3EOLPZkhp6Dfy6d1r7kU5oc,22573
 workflow_server/api/tests/test_workflow_view_stream_workflow_route.py,sha256=Qo8u6mPyRCmE2jamY1yIh8l44hgo4-Nwlq03z61ND5g,27031
-workflow_server/api/workflow_view.py,sha256=Ufc49WJ4LON3FFxzY2MeP4q7uQyorMNMu9zLRuJvdp4,21523
+workflow_server/api/workflow_view.py,sha256=_WhjNgimTPoS10C-npRWDfJixzg4eHTJ5xIKACStZf4,21943
 workflow_server/code_exec_runner.py,sha256=lBnMIorPZL8zZBye6TjeCIs06WTJM7P2HR07B1fjJJI,2533
 workflow_server/config.py,sha256=qmmTr6ty3ZN5LDOFs3TfUxYshYe6Mmn_LanplHHeE9Q,1796
 workflow_server/core/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 workflow_server/core/cancel_workflow.py,sha256=Ffkc3mzmrdMEUcD-sHfEhX4IwVrka-E--SxKA1dUfIU,2185
 workflow_server/core/events.py,sha256=24MA66DVQuaLJJcZrS8IL1Zq4Ohi9CoouKZ5VgoH3Cs,1402
-workflow_server/core/executor.py,sha256=srpxx2xkjAE0QNBo1sfQV6Ovwyq3tWQdjicTySxSrpc,18092
+workflow_server/core/executor.py,sha256=bNryNvXBbg7IHW3BFzzFaGoCO1rGufeEE4FrY9TDA90,18294
 workflow_server/core/utils.py,sha256=aIpSINstLGslP2PIoDLM82_1GlJ1uC_0AIrP-V7Yobo,3230
 workflow_server/core/workflow_executor_context.py,sha256=w3OhV_AXpgh7AxpjEsc0vo-IJypgJcgr5DXJCqGptOU,1587
 workflow_server/server.py,sha256=QBU12AaAfAgLqfCDBd24qIJl_mbheiq0-hfcWV7rZM4,1234
@@ -20,15 +20,15 @@ workflow_server/start.py,sha256=pkwRcms6I4tkVHP06LdrZY6rG_DFHfBx4ioY5X91W5k,2264
 workflow_server/utils/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 workflow_server/utils/exit_handler.py,sha256=_FacDVi4zc3bfTA3D2mJsISePlJ8jpLrnGVo5-xZQFs,743
 workflow_server/utils/log_proxy.py,sha256=nugi6fOgAYKX2X9DIc39TG366rsmmDUPoEtG3gzma_Y,3088
-workflow_server/utils/oom_killer.py,sha256=4Sag_iRQWqbp62iIBn6nKP-pxUHguOF93DdVXZTtJDk,2809
+workflow_server/utils/oom_killer.py,sha256=AprKFXC_wT3lQZcKDxU5O6dtJwi6meRxjo7nhQtQ8T0,2955
 workflow_server/utils/sentry.py,sha256=pqx3X_4W3yOzmz8QMJYUEi39skIKWtrTN5nyFhaPkbk,1597
-workflow_server/utils/system_utils.py,sha256=SVO7ofsoEXl7XduBtpiwHKMx7DTPmdC-uH6PiwRcN6M,2203
+workflow_server/utils/system_utils.py,sha256=3jNv113zRkKJ0928i2Vm6TqFHrDulteQu1kjseP2B0Y,3271
 workflow_server/utils/tests/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 workflow_server/utils/tests/test_sentry_integration.py,sha256=LGmWiaLhFrx-jslrRjRq9JY6Z5ShLZyx_N_L0-FU6OI,2100
 workflow_server/utils/tests/test_system_utils.py,sha256=_4GwXvVvU5BrATxUEWwQIPg0bzQXMWBtiBmjP8MTxJM,4314
 workflow_server/utils/tests/test_utils.py,sha256=0Nq6du8o-iBtTrip9_wgHES53JSiJbVdSXaBnPobw3s,6930
 workflow_server/utils/utils.py,sha256=ZPoM1Suhid22dpB8oEFLux8wx-9iyzmSfWuYxSCrgWk,4774
-vellum_workflow_server-1.4.0.dist-info/METADATA,sha256=Xj4ZNLbf9KnNGWKA5Hm9EFX8dCdhxc5BS5TthF5z3Us,2267
-vellum_workflow_server-1.4.0.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
-vellum_workflow_server-1.4.0.dist-info/entry_points.txt,sha256=uB_0yPkr7YV6RhEXzvFReUM8P4OQBlVXD6TN6eb9-oc,277
-vellum_workflow_server-1.4.0.dist-info/RECORD,,
+vellum_workflow_server-1.4.1.post1.dist-info/METADATA,sha256=sXR0B1bNZZvRQswy2hK4IqACrHG67fw6XLErZfqRoRE,2273
+vellum_workflow_server-1.4.1.post1.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
+vellum_workflow_server-1.4.1.post1.dist-info/entry_points.txt,sha256=uB_0yPkr7YV6RhEXzvFReUM8P4OQBlVXD6TN6eb9-oc,277
+vellum_workflow_server-1.4.1.post1.dist-info/RECORD,,

workflow_server/api/workflow_view.py

@@ -23,6 +23,7 @@ from vellum_ee.workflows.display.types import WorkflowDisplayContext
 from vellum_ee.workflows.display.workflows import BaseWorkflowDisplay
 from vellum_ee.workflows.server.virtual_file_loader import VirtualFileFinder
 
+from vellum.workflows import BaseWorkflow
 from vellum.workflows.exceptions import WorkflowInitializationException
 from vellum.workflows.nodes import BaseNode
 from workflow_server.config import ENABLE_PROCESS_WRAPPER, MEMORY_LIMIT_MB
@@ -49,8 +50,10 @@ from workflow_server.core.workflow_executor_context import (
 )
 from workflow_server.utils.oom_killer import get_is_oom_killed
 from workflow_server.utils.system_utils import (
+    add_active_span_id,
     get_active_process_count,
     increment_process_count,
+    remove_active_span_id,
     wait_for_available_process,
 )
 from workflow_server.utils.utils import convert_json_inputs_to_vellum, get_version
@@ -182,10 +185,7 @@ def stream_workflow_route() -> Response:
         first_item = next(stream_iterator)
         increment_process_count(1)
 
-    if isinstance(first_item, str) and first_item.startswith(SPAN_ID_EVENT):
-        span_id = first_item.split(":")[1]
-        headers["X-Vellum-Workflow-Span-Id"] = span_id
-    else:
+    if not isinstance(first_item, str) or not first_item.startswith(SPAN_ID_EVENT):
         logger.error("Workflow stream did not start with span id event")
         return Response(
             json.dumps({"detail": "Internal Server Error"}),
@@ -194,6 +194,10 @@ def stream_workflow_route() -> Response:
             headers=headers,
         )
 
+    span_id = first_item.split(":")[1]
+    headers["X-Vellum-Workflow-Span-Id"] = span_id
+    add_active_span_id(span_id)
+
     logger.info(f"Starting Workflow Stream, execution ID: {span_id}, ")
 
     def process_events(queue: Queue) -> Iterator[Union[str, dict]]:
@@ -227,6 +231,7 @@ def stream_workflow_route() -> Response:
 
                 if not ENABLE_PROCESS_WRAPPER or process:
                     increment_process_count(-1)
+                    remove_active_span_id(span_id)
 
                 yield VembdaExecutionFulfilledEvent(
                     id=uuid4(),
@@ -254,6 +259,7 @@ def stream_workflow_route() -> Response:
                     process.kill()
                 if process:
                     increment_process_count(-1)
+                    remove_active_span_id(span_id)
 
                 break
 
@@ -332,10 +338,12 @@ def stream_workflow_route() -> Response:
                         process.kill()
                     if process:
                         increment_process_count(-1)
+                        remove_active_span_id(span_id)
                 except Exception as e:
                     logger.error("Failed to kill process", e)
             else:
                 increment_process_count(-1)
+                remove_active_span_id(span_id)
 
     resp = Response(
         stream_with_context(generator()),
@@ -451,7 +459,10 @@ def serialize_route() -> Response:
             try:
                 result = BaseWorkflowDisplay.serialize_module(namespace, client=client)
             except Exception as e:
-                raise WorkflowInitializationException(str(e)) from e
+                raise WorkflowInitializationException(
+                    message=str(e),
+                    workflow_definition=BaseWorkflow,
+                ) from e
 
             return Response(
                 json.dumps(result.model_dump()),

workflow_server/core/executor.py

@@ -174,7 +174,7 @@ def stream_workflow(
     cancel_signal: Optional[ThreadingEvent] = None,
 ) -> tuple[Iterator[dict], UUID]:
     workflow, namespace = _gather_workflow(executor_context)
-    workflow_inputs = _get_workflow_inputs(executor_context)
+    workflow_inputs = _get_workflow_inputs(executor_context, workflow.__class__)
     display_context = _gather_display_context(workflow, namespace)
     workflow_state = (
         workflow.deserialize_state(
@@ -258,21 +258,18 @@ def stream_node(
     executor_context: NodeExecutorContext,
     disable_redirect: bool = True,
 ) -> Iterator[dict]:
-    namespace = _get_file_namespace(executor_context)
-
     def call_node() -> Generator[dict[str, Any], Any, None]:
-        sys.meta_path.append(VirtualFileFinder(executor_context.files, namespace))
-        workflow_context = _create_workflow_context(executor_context)
+        workflow, namespace = _gather_workflow(executor_context)
         node_module = importlib.import_module(f"{namespace}.{executor_context.node_module}")
 
         Node = getattr(node_module, executor_context.node_name)
 
-        workflow_inputs = _get_workflow_inputs(executor_context)
+        workflow_inputs = _get_workflow_inputs(executor_context, workflow.__class__)
         workflow_state = _get_workflow_state(executor_context, workflow_inputs=workflow_inputs)
 
         node = Node(
             state=workflow_state,
-            context=workflow_context,
+            context=workflow._context,
         )
 
         executor_context.stream_start_time = time.time_ns()
@@ -339,7 +336,7 @@ def _call_stream(
         yield vembda_fulfilled_event.model_dump(mode="json")
 
 
-def _create_workflow(executor_context: WorkflowExecutorContext, namespace: str) -> BaseWorkflow:
+def _create_workflow(executor_context: BaseExecutorContext, namespace: str) -> BaseWorkflow:
     workflow_context = _create_workflow_context(executor_context)
     Workflow = BaseWorkflow.load_from_module(namespace)
     VembdaExecutionFulfilledEvent.model_rebuild(
@@ -431,7 +428,9 @@ def _dump_event(event: BaseEvent, executor_context: BaseExecutorContext, client:
     return dump
 
 
-def _get_workflow_inputs(executor_context: BaseExecutorContext) -> Optional[BaseInputs]:
+def _get_workflow_inputs(
+    executor_context: BaseExecutorContext, workflow_class: Type[BaseWorkflow]
+) -> Optional[BaseInputs]:
     if not executor_context.inputs:
         return None
 
@@ -443,17 +442,22 @@ def _get_workflow_inputs(executor_context: BaseExecutorContext) -> Optional[Base
     try:
         inputs_module = importlib.import_module(inputs_module_path)
     except Exception as e:
-        raise WorkflowInitializationException(f"Failed to initialize workflow inputs: {e}") from e
+        raise WorkflowInitializationException(
+            message=f"Failed to initialize workflow inputs: {e}",
+            workflow_definition=workflow_class,
+        ) from e
 
     if not hasattr(inputs_module, "Inputs"):
         raise WorkflowInitializationException(
-            f"Inputs module {inputs_module_path} does not have a required Inputs class"
+            message=f"Inputs module {inputs_module_path} does not have a required Inputs class",
+            workflow_definition=workflow_class,
         )
 
     if not issubclass(inputs_module.Inputs, BaseInputs):
         raise WorkflowInitializationException(
-            f"""The class {inputs_module_path}.Inputs was expected to be a subclass of BaseInputs, \
-but found {inputs_module.Inputs.__class__.__name__}"""
+            message=f"""The class {inputs_module_path}.Inputs was expected to be a subclass of BaseInputs, \
+but found {inputs_module.Inputs.__class__.__name__}""",
+            workflow_definition=workflow_class,
         )
 
     return inputs_module.Inputs(**executor_context.inputs)
@@ -493,7 +497,7 @@ def _get_run_from_node(executor_context: WorkflowExecutorContext, workflow: Base
     return None
 
 
-def _gather_workflow(context: WorkflowExecutorContext) -> Tuple[BaseWorkflow, str]:
+def _gather_workflow(context: BaseExecutorContext) -> Tuple[BaseWorkflow, str]:
     try:
         namespace = _get_file_namespace(context)
         if namespace != LOCAL_WORKFLOW_MODULE:
@@ -505,7 +509,10 @@ def _gather_workflow(context: WorkflowExecutorContext) -> Tuple[BaseWorkflow, st
         return workflow, namespace
     except Exception as e:
         logger.exception("Failed to initialize Workflow")
-        raise WorkflowInitializationException(f"Failed to initialize workflow: {e}") from e
+        raise WorkflowInitializationException(
+            message=f"Failed to initialize workflow: {e}",
+            workflow_definition=BaseWorkflow,
+        ) from e
 
 
 def _gather_display_context(workflow: BaseWorkflow, namespace: str) -> Optional["WorkflowEventDisplayContext"]:

workflow_server/utils/oom_killer.py

@@ -15,6 +15,7 @@ from workflow_server.utils.system_utils import (
     FORCE_GC_MEMORY_PERCENT,
     WARN_MEMORY_PERCENT,
     get_active_process_count,
+    get_active_span_ids,
     get_memory_in_use_mb,
 )
 
@@ -71,8 +72,10 @@ class OomKillerThread(Thread):
 
             if memory_mb > (MEMORY_LIMIT_MB * _MAX_MEMORY_PERCENT):
                 self._kill_switch.set()
+                active_span_ids = get_active_span_ids()
                 logger.error(
-                    f"Workflow server OOM killed, memory: {memory_mb}MB, Process Count: {get_active_process_count()}"
+                    f"Workflow server OOM killed, memory: {memory_mb}MB, Process Count: {get_active_process_count()}",
+                    extra={"active_span_ids": active_span_ids},
                 )
                 # Give time for the threads to get our kill switch
                 sleep(_KILL_GRACE_PERIOD)

workflow_server/utils/system_utils.py

@@ -17,6 +17,7 @@ _MEMORY_CHECK_INTERVAL_SECONDS = 2
 _MAX_MEMORY_CHECK_ATTEMPTS = 3
 _ACTIVE_PROCESS_COUNT = multiprocessing.Value("i", 0)
 _ACTIVE_PROCESS_LOCK = multiprocessing.Lock()
+_ACTIVE_SPAN_IDS = multiprocessing.Manager().list()
 
 
 def increment_process_count(change: int) -> None:
@@ -36,6 +37,39 @@ def get_active_process_count() -> int:
     return _ACTIVE_PROCESS_COUNT.value  # type: ignore
 
 
+def get_active_span_ids() -> list[str]:
+    """Get a copy of currently active span IDs"""
+    with _ACTIVE_PROCESS_LOCK:
+        return list(_ACTIVE_SPAN_IDS)
+
+
+def add_active_span_id(span_id: str) -> None:
+    """Add a span ID to the active tracking list"""
+    result = _ACTIVE_PROCESS_LOCK.acquire(timeout=5)
+    try:
+        if result:
+            _ACTIVE_SPAN_IDS.append(span_id)
+        else:
+            logger.error("Failed to lock workflow server span ID tracking.")
+    finally:
+        if result:
+            _ACTIVE_PROCESS_LOCK.release()
+
+
+def remove_active_span_id(span_id: str) -> None:
+    """Remove a span ID from the active tracking list"""
+    result = _ACTIVE_PROCESS_LOCK.acquire(timeout=5)
+    try:
+        if result and span_id in _ACTIVE_SPAN_IDS:
+            _ACTIVE_SPAN_IDS.remove(span_id)
+        else:
+            if not result:
+                logger.error("Failed to lock workflow server span ID tracking.")
+    finally:
+        if result:
+            _ACTIVE_PROCESS_LOCK.release()
+
+
 def get_memory_in_use_mb() -> Optional[float]:
     try:
         with open("/sys/fs/cgroup/memory/memory.usage_in_bytes", "r") as file: