vellum-workflow-server 1.4.0__tar.gz → 1.4.0.post1__tar.gz

{vellum_workflow_server-1.4.0 → vellum_workflow_server-1.4.0.post1}/PKG-INFO

@@ -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 → vellum_workflow_server-1.4.0.post1}/pyproject.toml

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

{vellum_workflow_server-1.4.0 → vellum_workflow_server-1.4.0.post1}/src/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()),

{vellum_workflow_server-1.4.0 → vellum_workflow_server-1.4.0.post1}/src/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)

{vellum_workflow_server-1.4.0 → vellum_workflow_server-1.4.0.post1}/src/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: