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

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

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

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

@@ -6,7 +6,7 @@ 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=4XaBbhzQte4bRbbUPhrqAXDBXTXoc67Twzr1Pp5bPpw,21800
 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
@@ -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.0.post1.dist-info/METADATA,sha256=fY2l-73g7b-QbDRuJzHYfA40CIhE3t5o6CaXCdidCKI,2273
+vellum_workflow_server-1.4.0.post1.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
+vellum_workflow_server-1.4.0.post1.dist-info/entry_points.txt,sha256=uB_0yPkr7YV6RhEXzvFReUM8P4OQBlVXD6TN6eb9-oc,277
+vellum_workflow_server-1.4.0.post1.dist-info/RECORD,,

workflow_server/api/workflow_view.py

@@ -49,8 +49,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 +184,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 +193,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 +230,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 +258,7 @@ def stream_workflow_route() -> Response:
                     process.kill()
                 if process:
                     increment_process_count(-1)
+                    remove_active_span_id(span_id)
 
                 break
 
@@ -332,10 +337,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()),

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: