vellum-workflow-server 1.9.7.post1__tar.gz → 1.11.0.post1__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (36) hide show
  1. {vellum_workflow_server-1.9.7.post1 → vellum_workflow_server-1.11.0.post1}/PKG-INFO +4 -3
  2. {vellum_workflow_server-1.9.7.post1 → vellum_workflow_server-1.11.0.post1}/pyproject.toml +4 -3
  3. vellum_workflow_server-1.11.0.post1/src/workflow_server/api/tests/test_workflow_view_async_exec.py +410 -0
  4. {vellum_workflow_server-1.9.7.post1 → vellum_workflow_server-1.11.0.post1}/src/workflow_server/api/tests/test_workflow_view_stream_workflow_route.py +69 -0
  5. {vellum_workflow_server-1.9.7.post1 → vellum_workflow_server-1.11.0.post1}/src/workflow_server/api/workflow_view.py +4 -3
  6. {vellum_workflow_server-1.9.7.post1 → vellum_workflow_server-1.11.0.post1}/src/workflow_server/code_exec_runner.py +4 -3
  7. {vellum_workflow_server-1.9.7.post1 → vellum_workflow_server-1.11.0.post1}/src/workflow_server/core/executor.py +4 -25
  8. {vellum_workflow_server-1.9.7.post1 → vellum_workflow_server-1.11.0.post1}/src/workflow_server/core/utils.py +4 -0
  9. {vellum_workflow_server-1.9.7.post1 → vellum_workflow_server-1.11.0.post1}/src/workflow_server/core/workflow_executor_context.py +13 -1
  10. {vellum_workflow_server-1.9.7.post1 → vellum_workflow_server-1.11.0.post1}/src/workflow_server/start.py +2 -2
  11. {vellum_workflow_server-1.9.7.post1 → vellum_workflow_server-1.11.0.post1}/README.md +0 -0
  12. {vellum_workflow_server-1.9.7.post1 → vellum_workflow_server-1.11.0.post1}/src/workflow_server/__init__.py +0 -0
  13. {vellum_workflow_server-1.9.7.post1 → vellum_workflow_server-1.11.0.post1}/src/workflow_server/api/__init__.py +0 -0
  14. {vellum_workflow_server-1.9.7.post1 → vellum_workflow_server-1.11.0.post1}/src/workflow_server/api/auth_middleware.py +0 -0
  15. {vellum_workflow_server-1.9.7.post1 → vellum_workflow_server-1.11.0.post1}/src/workflow_server/api/healthz_view.py +0 -0
  16. {vellum_workflow_server-1.9.7.post1 → vellum_workflow_server-1.11.0.post1}/src/workflow_server/api/status_view.py +0 -0
  17. {vellum_workflow_server-1.9.7.post1 → vellum_workflow_server-1.11.0.post1}/src/workflow_server/api/tests/__init__.py +0 -0
  18. {vellum_workflow_server-1.9.7.post1 → vellum_workflow_server-1.11.0.post1}/src/workflow_server/api/tests/test_input_display_mapping.py +0 -0
  19. {vellum_workflow_server-1.9.7.post1 → vellum_workflow_server-1.11.0.post1}/src/workflow_server/api/tests/test_workflow_view.py +0 -0
  20. {vellum_workflow_server-1.9.7.post1 → vellum_workflow_server-1.11.0.post1}/src/workflow_server/config.py +0 -0
  21. {vellum_workflow_server-1.9.7.post1 → vellum_workflow_server-1.11.0.post1}/src/workflow_server/core/__init__.py +0 -0
  22. {vellum_workflow_server-1.9.7.post1 → vellum_workflow_server-1.11.0.post1}/src/workflow_server/core/cancel_workflow.py +0 -0
  23. {vellum_workflow_server-1.9.7.post1 → vellum_workflow_server-1.11.0.post1}/src/workflow_server/core/events.py +0 -0
  24. {vellum_workflow_server-1.9.7.post1 → vellum_workflow_server-1.11.0.post1}/src/workflow_server/logging_config.py +0 -0
  25. {vellum_workflow_server-1.9.7.post1 → vellum_workflow_server-1.11.0.post1}/src/workflow_server/server.py +0 -0
  26. {vellum_workflow_server-1.9.7.post1 → vellum_workflow_server-1.11.0.post1}/src/workflow_server/utils/__init__.py +0 -0
  27. {vellum_workflow_server-1.9.7.post1 → vellum_workflow_server-1.11.0.post1}/src/workflow_server/utils/exit_handler.py +0 -0
  28. {vellum_workflow_server-1.9.7.post1 → vellum_workflow_server-1.11.0.post1}/src/workflow_server/utils/log_proxy.py +0 -0
  29. {vellum_workflow_server-1.9.7.post1 → vellum_workflow_server-1.11.0.post1}/src/workflow_server/utils/oom_killer.py +0 -0
  30. {vellum_workflow_server-1.9.7.post1 → vellum_workflow_server-1.11.0.post1}/src/workflow_server/utils/sentry.py +0 -0
  31. {vellum_workflow_server-1.9.7.post1 → vellum_workflow_server-1.11.0.post1}/src/workflow_server/utils/system_utils.py +0 -0
  32. {vellum_workflow_server-1.9.7.post1 → vellum_workflow_server-1.11.0.post1}/src/workflow_server/utils/tests/__init__.py +0 -0
  33. {vellum_workflow_server-1.9.7.post1 → vellum_workflow_server-1.11.0.post1}/src/workflow_server/utils/tests/test_sentry_integration.py +0 -0
  34. {vellum_workflow_server-1.9.7.post1 → vellum_workflow_server-1.11.0.post1}/src/workflow_server/utils/tests/test_system_utils.py +0 -0
  35. {vellum_workflow_server-1.9.7.post1 → vellum_workflow_server-1.11.0.post1}/src/workflow_server/utils/tests/test_utils.py +0 -0
  36. {vellum_workflow_server-1.9.7.post1 → vellum_workflow_server-1.11.0.post1}/src/workflow_server/utils/utils.py +0 -0
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.1
2
2
  Name: vellum-workflow-server
3
- Version: 1.9.7.post1
3
+ Version: 1.11.0.post1
4
4
  Summary:
5
5
  License: AGPL
6
6
  Requires-Python: >=3.9.0,<4
@@ -24,12 +24,13 @@ Requires-Dist: cryptography (==43.0.3)
24
24
  Requires-Dist: flask (==2.3.3)
25
25
  Requires-Dist: gunicorn (==23.0.0)
26
26
  Requires-Dist: orderly-set (==5.2.2)
27
+ Requires-Dist: orjson (==3.11.4)
27
28
  Requires-Dist: pebble (==5.0.7)
28
29
  Requires-Dist: pyjwt (==2.10.0)
29
- Requires-Dist: python-dotenv (==1.0.1)
30
+ Requires-Dist: python-dotenv (==1.2.1)
30
31
  Requires-Dist: retrying (==1.3.4)
31
32
  Requires-Dist: sentry-sdk[flask] (==2.20.0)
32
- Requires-Dist: vellum-ai (==1.9.7)
33
+ Requires-Dist: vellum-ai (==1.11.0)
33
34
  Description-Content-Type: text/markdown
34
35
 
35
36
  # Vellum Workflow Runner Server
@@ -3,7 +3,7 @@ name = "vellum-workflow-server"
3
3
 
4
4
  [tool.poetry]
5
5
  name = "vellum-workflow-server"
6
- version = "1.9.7.post1"
6
+ version = "1.11.0.post1"
7
7
  description = ""
8
8
  readme = "README.md"
9
9
  authors = []
@@ -45,8 +45,9 @@ flask = "2.3.3"
45
45
  orderly-set = "5.2.2"
46
46
  pebble = "5.0.7"
47
47
  gunicorn = "23.0.0"
48
- vellum-ai = "1.9.7"
49
- python-dotenv = "1.0.1"
48
+ orjson = "3.11.4"
49
+ vellum-ai = "1.11.0"
50
+ python-dotenv = "1.2.1"
50
51
  retrying = "1.3.4"
51
52
  sentry-sdk = {extras = ["flask"], version = "2.20.0"}
52
53
 
@@ -0,0 +1,410 @@
1
+ import pytest
2
+ import logging
3
+ import time
4
+ from uuid import uuid4
5
+
6
+ from workflow_server.server import create_app
7
+ from workflow_server.utils.system_utils import get_active_process_count
8
+
9
+
10
+ @pytest.fixture(autouse=True)
11
+ def drain_background_threads():
12
+ """
13
+ Ensures background threads from previous tests complete before starting the next test.
14
+ This prevents cross-test interference in process count assertions.
15
+ """
16
+ baseline = get_active_process_count()
17
+ yield
18
+
19
+ deadline = time.time() + 15
20
+ while time.time() < deadline:
21
+ current_count = get_active_process_count()
22
+ if current_count == baseline:
23
+ break
24
+ time.sleep(0.1)
25
+
26
+
27
+ def test_async_exec_route__happy_path():
28
+ """
29
+ Tests that the async-exec route successfully accepts a valid workflow and returns immediately.
30
+ """
31
+ # GIVEN a Flask application
32
+ flask_app = create_app()
33
+
34
+ # AND a valid workflow request
35
+ span_id = uuid4()
36
+ request_body = {
37
+ "execution_id": str(span_id),
38
+ "inputs": [],
39
+ "environment_api_key": "test",
40
+ "module": "workflow",
41
+ "timeout": 360,
42
+ "files": {
43
+ "__init__.py": "",
44
+ "workflow.py": """\
45
+ from vellum.workflows import BaseWorkflow
46
+
47
+ class Workflow(BaseWorkflow):
48
+ class Outputs(BaseWorkflow.Outputs):
49
+ foo = "hello"
50
+ """,
51
+ },
52
+ }
53
+
54
+ # WHEN we make a request to the async-exec route
55
+ with flask_app.test_client() as test_client:
56
+ response = test_client.post("/workflow/async-exec", json=request_body)
57
+
58
+ # THEN we should get a 200 response
59
+ assert response.status_code == 200
60
+
61
+ # AND the response should indicate success
62
+ assert response.json == {"success": True}
63
+
64
+
65
+ def test_async_exec_route__with_inputs():
66
+ """
67
+ Tests that the async-exec route handles workflows with inputs correctly.
68
+ """
69
+ # GIVEN a Flask application
70
+ flask_app = create_app()
71
+
72
+ # AND a valid workflow request with inputs
73
+ span_id = uuid4()
74
+ request_body = {
75
+ "execution_id": str(span_id),
76
+ "inputs": [
77
+ {"name": "foo", "type": "STRING", "value": "hello"},
78
+ ],
79
+ "environment_api_key": "test",
80
+ "module": "workflow",
81
+ "timeout": 360,
82
+ "files": {
83
+ "__init__.py": "",
84
+ "workflow.py": """\
85
+ from vellum.workflows import BaseWorkflow
86
+ from vellum.workflows.state import BaseState
87
+ from .inputs import Inputs
88
+
89
+ class Workflow(BaseWorkflow[Inputs, BaseState]):
90
+ class Outputs(BaseWorkflow.Outputs):
91
+ foo = "hello"
92
+ """,
93
+ "inputs.py": """\
94
+ from vellum.workflows.inputs import BaseInputs
95
+
96
+ class Inputs(BaseInputs):
97
+ foo: str
98
+ """,
99
+ },
100
+ }
101
+
102
+ # WHEN we make a request to the async-exec route
103
+ with flask_app.test_client() as test_client:
104
+ response = test_client.post("/workflow/async-exec", json=request_body)
105
+
106
+ # THEN we should get a 200 response
107
+ assert response.status_code == 200
108
+
109
+ # AND the response should indicate success
110
+ assert response.json == {"success": True}
111
+
112
+
113
+ def test_async_exec_route__with_state():
114
+ """
115
+ Tests that the async-exec route handles workflows with state correctly.
116
+ """
117
+ # GIVEN a Flask application
118
+ flask_app = create_app()
119
+
120
+ # AND a valid workflow request with state
121
+ span_id = uuid4()
122
+ request_body = {
123
+ "execution_id": str(span_id),
124
+ "state": {"foo": "bar"},
125
+ "environment_api_key": "test",
126
+ "module": "workflow",
127
+ "timeout": 360,
128
+ "files": {
129
+ "__init__.py": "",
130
+ "workflow.py": """\
131
+ from vellum.workflows import BaseWorkflow
132
+ from vellum.workflows.inputs import BaseInputs
133
+ from .state import State
134
+
135
+ class Workflow(BaseWorkflow[BaseInputs, State]):
136
+ class Outputs(BaseWorkflow.Outputs):
137
+ foo = State.foo
138
+ """,
139
+ "state.py": """\
140
+ from vellum.workflows.state import BaseState
141
+
142
+ class State(BaseState):
143
+ foo: str
144
+ """,
145
+ },
146
+ }
147
+
148
+ # WHEN we make a request to the async-exec route
149
+ with flask_app.test_client() as test_client:
150
+ response = test_client.post("/workflow/async-exec", json=request_body)
151
+
152
+ # THEN we should get a 200 response
153
+ assert response.status_code == 200
154
+
155
+ # AND the response should indicate success
156
+ assert response.json == {"success": True}
157
+
158
+
159
+ def test_async_exec_route__invalid_context():
160
+ """
161
+ Tests that the async-exec route returns 400 for invalid request context.
162
+ """
163
+ # GIVEN a Flask application
164
+ flask_app = create_app()
165
+
166
+ # AND an invalid request missing required fields
167
+ request_body = {
168
+ "inputs": [],
169
+ }
170
+
171
+ # WHEN we make a request to the async-exec route
172
+ with flask_app.test_client() as test_client:
173
+ response = test_client.post("/workflow/async-exec", json=request_body)
174
+
175
+ # THEN we should get a 400 response
176
+ assert response.status_code == 400
177
+
178
+ # AND the response should contain error details
179
+ assert "detail" in response.json
180
+ assert "Invalid context" in response.json["detail"]
181
+
182
+
183
+ def test_async_exec_route__missing_files():
184
+ """
185
+ Tests that the async-exec route returns 400 when files are missing.
186
+ """
187
+ # GIVEN a Flask application
188
+ flask_app = create_app()
189
+
190
+ span_id = uuid4()
191
+ request_body = {
192
+ "execution_id": str(span_id),
193
+ "inputs": [],
194
+ "environment_api_key": "test",
195
+ "module": "workflow",
196
+ "timeout": 360,
197
+ }
198
+
199
+ # WHEN we make a request to the async-exec route
200
+ with flask_app.test_client() as test_client:
201
+ response = test_client.post("/workflow/async-exec", json=request_body)
202
+
203
+ # THEN we should get a 400 response
204
+ assert response.status_code == 400
205
+
206
+ # AND the response should contain error details
207
+ assert "detail" in response.json
208
+ assert "Invalid context" in response.json["detail"]
209
+
210
+
211
+ def test_async_exec_route__with_syntax_error_in_workflow():
212
+ """
213
+ Tests that the async-exec route handles workflows with syntax errors gracefully.
214
+ """
215
+ # GIVEN a Flask application
216
+ flask_app = create_app()
217
+
218
+ span_id = uuid4()
219
+ request_body = {
220
+ "execution_id": str(span_id),
221
+ "inputs": [],
222
+ "environment_api_key": "test",
223
+ "module": "workflow",
224
+ "timeout": 360,
225
+ "files": {
226
+ "__init__.py": "",
227
+ "workflow.py": """\
228
+ from vellum.workflows import BaseWorkflow
229
+
230
+ class Workflow(BaseWorkflow)
231
+ class Outputs(BaseWorkflow.Outputs):
232
+ foo = "hello"
233
+ """,
234
+ },
235
+ }
236
+
237
+ # WHEN we make a request to the async-exec route
238
+ with flask_app.test_client() as test_client:
239
+ response = test_client.post("/workflow/async-exec", json=request_body)
240
+
241
+ # THEN we should get a 200 response (async execution is accepted)
242
+ assert response.status_code == 200
243
+
244
+ # AND the response should indicate success
245
+ assert response.json == {"success": True}
246
+
247
+
248
+ def test_async_exec_route__with_invalid_inputs():
249
+ """
250
+ Tests that the async-exec route handles workflows with invalid inputs gracefully.
251
+ """
252
+ # GIVEN a Flask application
253
+ flask_app = create_app()
254
+
255
+ span_id = uuid4()
256
+ request_body = {
257
+ "execution_id": str(span_id),
258
+ "inputs": [],
259
+ "environment_api_key": "test",
260
+ "module": "workflow",
261
+ "timeout": 360,
262
+ "files": {
263
+ "__init__.py": "",
264
+ "workflow.py": """\
265
+ from vellum.workflows import BaseWorkflow
266
+ from vellum.workflows.state import BaseState
267
+ from .inputs import Inputs
268
+
269
+ class Workflow(BaseWorkflow[Inputs, BaseState]):
270
+ class Outputs(BaseWorkflow.Outputs):
271
+ foo = "hello"
272
+ """,
273
+ "inputs.py": """\
274
+ from vellum.workflows.inputs import BaseInputs
275
+
276
+ class Inputs(BaseInputs):
277
+ foo: str
278
+ """,
279
+ },
280
+ }
281
+
282
+ # WHEN we make a request to the async-exec route
283
+ with flask_app.test_client() as test_client:
284
+ response = test_client.post("/workflow/async-exec", json=request_body)
285
+
286
+ # THEN we should get a 200 response (async execution is accepted)
287
+ assert response.status_code == 200
288
+
289
+ # AND the response should indicate success
290
+ assert response.json == {"success": True}
291
+
292
+
293
+ def test_async_exec_route__background_thread_completes(caplog):
294
+ """
295
+ Verifies that the async background worker thread runs to completion.
296
+ """
297
+ # GIVEN a Flask application with log capture enabled
298
+ caplog.set_level(logging.INFO, logger="workflow_server.api.workflow_view")
299
+ flask_app = create_app()
300
+
301
+ baseline = get_active_process_count()
302
+
303
+ # AND a valid workflow request
304
+ span_id = uuid4()
305
+ request_body = {
306
+ "execution_id": str(span_id),
307
+ "inputs": [],
308
+ "environment_api_key": "test",
309
+ "module": "workflow",
310
+ "timeout": 360,
311
+ "files": {
312
+ "__init__.py": "",
313
+ "workflow.py": """\
314
+ from vellum.workflows import BaseWorkflow
315
+
316
+ class Workflow(BaseWorkflow):
317
+ class Outputs(BaseWorkflow.Outputs):
318
+ foo = "hello"
319
+ """,
320
+ },
321
+ }
322
+
323
+ # WHEN we call the async-exec route
324
+ with flask_app.test_client() as test_client:
325
+ response = test_client.post("/workflow/async-exec", json=request_body)
326
+
327
+ # THEN we get immediate acceptance
328
+ assert response.status_code == 200
329
+ assert response.json == {"success": True}
330
+
331
+ # AND the background thread should complete
332
+ completion_deadline = time.time() + 15
333
+ saw_completion_log = False
334
+ while time.time() < completion_deadline:
335
+ if any("Workflow async exec completed" in rec.message for rec in caplog.records):
336
+ saw_completion_log = True
337
+ break
338
+ time.sleep(0.1)
339
+
340
+ # THEN we should observe the completion log
341
+ assert saw_completion_log, "Did not observe background completion log within 15 seconds"
342
+
343
+ cleanup_deadline = time.time() + 15
344
+ process_count_returned = False
345
+ while time.time() < cleanup_deadline:
346
+ current_count = get_active_process_count()
347
+ if current_count == baseline:
348
+ process_count_returned = True
349
+ break
350
+ time.sleep(0.1)
351
+
352
+ current_count = get_active_process_count()
353
+ assert process_count_returned, (
354
+ f"Process count did not return to baseline within 15 seconds after completion log. "
355
+ f"Expected: {baseline}, Current: {current_count}"
356
+ )
357
+
358
+
359
+ def test_async_exec_route__background_thread_completes_on_error(caplog):
360
+ """
361
+ Verifies that the background worker completes even when the workflow fails early.
362
+ """
363
+ # GIVEN a Flask application with log capture enabled
364
+ caplog.set_level(logging.INFO, logger="workflow_server.api.workflow_view")
365
+ flask_app = create_app()
366
+
367
+ baseline = get_active_process_count()
368
+
369
+ span_id = uuid4()
370
+ request_body = {
371
+ "execution_id": str(span_id),
372
+ "inputs": [],
373
+ "environment_api_key": "test",
374
+ "module": "workflow",
375
+ "timeout": 360,
376
+ "files": {
377
+ "__init__.py": "",
378
+ "workflow.py": """\
379
+ from vellum.workflows import BaseWorkflow
380
+
381
+ class Workflow(BaseWorkflow)
382
+ class Outputs(BaseWorkflow.Outputs):
383
+ foo = "hello"
384
+ """,
385
+ },
386
+ }
387
+
388
+ # WHEN we call the async-exec route
389
+ with flask_app.test_client() as test_client:
390
+ response = test_client.post("/workflow/async-exec", json=request_body)
391
+
392
+ # THEN we get immediate acceptance
393
+ assert response.status_code == 200
394
+ assert response.json == {"success": True}
395
+
396
+ # AND the background thread should complete and clean up resources
397
+ deadline = time.time() + 15
398
+ process_count_returned = False
399
+ while time.time() < deadline:
400
+ current_count = get_active_process_count()
401
+ if current_count == baseline:
402
+ process_count_returned = True
403
+ break
404
+ time.sleep(0.1)
405
+
406
+ current_count = get_active_process_count()
407
+ assert process_count_returned, (
408
+ f"Process count did not return to baseline on error within 15 seconds. "
409
+ f"Expected: {baseline}, Current: {current_count}"
410
+ )
@@ -1241,3 +1241,72 @@ class InvalidWorkflow(BaseWorkflow):
1241
1241
  assert events[3]["name"] == "vembda.execution.fulfilled"
1242
1242
  assert events[3]["span_id"] == str(span_id)
1243
1243
  assert events[3]["body"]["exit_code"] == 0
1244
+
1245
+
1246
+ @mock.patch("workflow_server.api.workflow_view.get_is_oom_killed")
1247
+ def test_stream_workflow_route__oom_does_not_set_timed_out_flag(mock_get_is_oom_killed):
1248
+ """
1249
+ Tests that when an OOM error occurs, we don't set the timed_out flag in the vembda fulfilled event.
1250
+ """
1251
+ # GIVEN a workflow that takes some time to execute
1252
+ span_id = uuid4()
1253
+ request_body = {
1254
+ "timeout": 10,
1255
+ "execution_id": str(span_id),
1256
+ "inputs": [],
1257
+ "environment_api_key": "test",
1258
+ "module": "workflow",
1259
+ "files": {
1260
+ "__init__.py": "",
1261
+ "workflow.py": """\
1262
+ import time
1263
+
1264
+ from vellum.workflows.nodes.bases.base import BaseNode
1265
+ from vellum.workflows.workflows.base import BaseWorkflow
1266
+
1267
+
1268
+ class SlowNode(BaseNode):
1269
+ class Outputs(BaseNode.Outputs):
1270
+ value: str
1271
+
1272
+ def run(self) -> Outputs:
1273
+ time.sleep(2)
1274
+ return self.Outputs(value="hello world")
1275
+
1276
+
1277
+ class OOMWorkflow(BaseWorkflow):
1278
+ graph = SlowNode
1279
+ class Outputs(BaseWorkflow.Outputs):
1280
+ final_value = SlowNode.Outputs.value
1281
+
1282
+ """,
1283
+ },
1284
+ }
1285
+
1286
+ # WHEN we mock the OOM killer to trigger after a few checks
1287
+ call_count = [0]
1288
+
1289
+ def mock_oom_side_effect():
1290
+ call_count[0] += 1
1291
+ if call_count[0] > 3:
1292
+ return True
1293
+ return False
1294
+
1295
+ mock_get_is_oom_killed.side_effect = mock_oom_side_effect
1296
+
1297
+ # AND we call the stream route
1298
+ status_code, events = flask_stream(request_body)
1299
+
1300
+ # THEN we get a 200 response
1301
+ assert status_code == 200
1302
+
1303
+ # AND we get the expected events
1304
+ event_names = [e["name"] for e in events]
1305
+
1306
+ assert "vembda.execution.initiated" in event_names
1307
+
1308
+ # THEN the key assertion: if there's a vembda.execution.fulfilled event, it should NOT have timed_out=True
1309
+ vembda_fulfilled_event = next(e for e in events if e["name"] == "vembda.execution.fulfilled")
1310
+ assert (
1311
+ vembda_fulfilled_event["body"].get("timed_out") is not True
1312
+ ), "timed_out flag should not be set when OOM occurs"
@@ -15,6 +15,7 @@ from uuid import uuid4
15
15
  from typing import Any, Dict, Generator, Iterator, Optional, Union, cast
16
16
 
17
17
  from flask import Blueprint, Response, current_app as app, request, stream_with_context
18
+ import orjson
18
19
  from pydantic import ValidationError
19
20
  from vellum_ee.workflows.display.nodes.get_node_display_class import get_node_display_class
20
21
  from vellum_ee.workflows.display.types import WorkflowDisplayContext
@@ -115,7 +116,7 @@ def stream_workflow_route() -> Response:
115
116
  for row in workflow_events:
116
117
  yield "\n"
117
118
  if isinstance(row, dict):
118
- dump = json.dumps(row)
119
+ dump = orjson.dumps(row).decode("utf-8")
119
120
  yield dump
120
121
  else:
121
122
  yield row
@@ -500,11 +501,11 @@ def stream_node_route() -> Response:
500
501
  break
501
502
 
502
503
  def generator() -> Generator[str, None, None]:
503
- yield json.dumps(vembda_initiated_event.model_dump(mode="json"))
504
+ yield orjson.dumps(vembda_initiated_event.model_dump(mode="json")).decode("utf-8")
504
505
 
505
506
  for row in node_events():
506
507
  yield "\n"
507
- yield json.dumps(row)
508
+ yield orjson.dumps(row).decode("utf-8")
508
509
 
509
510
  headers = {
510
511
  "X-Vellum-SDK-Version": vembda_initiated_event.body.sdk_version,
@@ -1,11 +1,12 @@
1
1
  from datetime import datetime
2
- import json
3
2
  import logging
4
3
  import os
5
4
  from threading import Event as ThreadingEvent
6
5
  from uuid import uuid4
7
6
  from typing import Optional
8
7
 
8
+ import orjson
9
+
9
10
  from workflow_server.core.events import VembdaExecutionInitiatedBody, VembdaExecutionInitiatedEvent
10
11
  from workflow_server.core.executor import stream_workflow
11
12
  from workflow_server.core.utils import serialize_vembda_rejected_event
@@ -29,7 +30,7 @@ def run_code_exec_stream() -> None:
29
30
  split_input = input_raw.split("\n--vellum-input-stop--\n")
30
31
  input_json = split_input[0]
31
32
 
32
- input_data = json.loads(input_json)
33
+ input_data = orjson.loads(input_json)
33
34
  context = WorkflowExecutorContext.model_validate(input_data)
34
35
 
35
36
  print("--vellum-output-start--") # noqa: T201
@@ -53,7 +54,7 @@ def run_code_exec_stream() -> None:
53
54
  cancel_signal=ThreadingEvent(),
54
55
  )
55
56
  for line in stream_iterator:
56
- print(f"{_EVENT_LINE}{json.dumps(line)}") # noqa: T201
57
+ print(f"{_EVENT_LINE}{orjson.dumps(line).decode('utf-8')}") # noqa: T201
57
58
  except Exception as e:
58
59
  logger.exception(e)
59
60
 
@@ -1,6 +1,5 @@
1
1
  from datetime import datetime, timezone
2
2
  from io import StringIO
3
- import json
4
3
  import logging
5
4
  from multiprocessing import Process, Queue
6
5
  import os
@@ -11,8 +10,9 @@ from threading import Event as ThreadingEvent
11
10
  import time
12
11
  from traceback import format_exc
13
12
  from uuid import UUID, uuid4
14
- from typing import Any, Callable, Generator, Iterator, Optional, Tuple, Type
13
+ from typing import Any, Callable, Generator, Iterator, Optional, Tuple
15
14
 
15
+ import orjson
16
16
  from vellum_ee.workflows.display.utils.events import event_enricher
17
17
  from vellum_ee.workflows.server.virtual_file_loader import VirtualFileFinder
18
18
 
@@ -104,7 +104,7 @@ def _stream_workflow_wrapper(
104
104
  span_id_emitted = True
105
105
 
106
106
  for event in stream_iterator:
107
- queue.put(json.dumps(event))
107
+ queue.put(orjson.dumps(event).decode("utf-8"))
108
108
 
109
109
  except Exception as e:
110
110
  if not span_id_emitted:
@@ -273,32 +273,11 @@ def stream_node(
273
273
  disable_redirect: bool = True,
274
274
  ) -> Iterator[dict]:
275
275
  workflow, namespace = _create_workflow(executor_context)
276
- Node: Optional[Type[BaseNode]] = None
277
-
278
- for workflow_node in workflow.get_nodes():
279
- if executor_context.node_id and workflow_node.__id__ == executor_context.node_id:
280
- Node = workflow_node
281
- break
282
- elif (
283
- executor_context.node_module
284
- and executor_context.node_name
285
- and workflow_node.__name__ == executor_context.node_name
286
- and workflow_node.__module__ == f"{namespace}.{executor_context.node_module}"
287
- ):
288
- Node = workflow_node
289
- break
290
-
291
- if not Node:
292
- identifier = executor_context.node_id or f"{executor_context.node_module}.{executor_context.node_name}"
293
- raise WorkflowInitializationException(
294
- message=f"Node '{identifier}' not found in workflow",
295
- workflow_definition=workflow.__class__,
296
- )
297
276
 
298
277
  def call_node() -> Generator[dict[str, Any], Any, None]:
299
278
  executor_context.stream_start_time = time.time_ns()
300
279
 
301
- for event in workflow.run_node(Node, inputs=executor_context.inputs): # type: ignore[arg-type]
280
+ for event in workflow.run_node(executor_context.node_ref, inputs=executor_context.inputs):
302
281
  yield event.model_dump(mode="json")
303
282
 
304
283
  return _call_stream(
@@ -2,6 +2,7 @@ from datetime import datetime
2
2
  from uuid import uuid4
3
3
  from typing import Optional
4
4
 
5
+ from workflow_server.config import IS_ASYNC_MODE
5
6
  from workflow_server.core.events import VembdaExecutionFulfilledBody, VembdaExecutionFulfilledEvent
6
7
  from workflow_server.core.workflow_executor_context import BaseExecutorContext
7
8
 
@@ -46,6 +47,9 @@ def serialize_vembda_rejected_event(
46
47
 
47
48
 
48
49
  def is_events_emitting_enabled(executor_context: Optional[BaseExecutorContext]) -> bool:
50
+ if IS_ASYNC_MODE:
51
+ return True
52
+
49
53
  if not executor_context:
50
54
  return False
51
55
 
@@ -3,7 +3,7 @@ from functools import cached_property
3
3
  import os
4
4
  import time
5
5
  from uuid import UUID
6
- from typing import Any, Optional
6
+ from typing import Any, Optional, Union
7
7
  from typing_extensions import Self
8
8
 
9
9
  from flask import has_request_context, request
@@ -91,6 +91,18 @@ class NodeExecutorContext(BaseExecutorContext):
91
91
  node_module: Optional[str] = None
92
92
  node_name: Optional[str] = None
93
93
 
94
+ @property
95
+ def node_ref(self) -> Union[UUID, str]:
96
+ """
97
+ Returns the node reference for use with workflow.run_node().
98
+
99
+ Returns node_id if it exists, otherwise returns the combination
100
+ of node_module and node_name as a fully qualified string.
101
+ """
102
+ if self.node_id:
103
+ return self.node_id
104
+ return f"{self.node_module}.{self.node_name}"
105
+
94
106
  @model_validator(mode="after")
95
107
  def validate_node_identification(self) -> Self:
96
108
  if not self.node_id and not (self.node_module and self.node_name):
@@ -64,8 +64,8 @@ def start() -> None:
64
64
  "workers": int(os.getenv("GUNICORN_WORKERS", 2)),
65
65
  "threads": int(os.getenv("GUNICORN_THREADS", 9 if ENABLE_PROCESS_WRAPPER else 6)),
66
66
  # Aggressively try to avoid memory leaks when using non process mode
67
- "max_requests": 120 if ENABLE_PROCESS_WRAPPER else 20,
68
- "max_requests_jitter": 30 if ENABLE_PROCESS_WRAPPER else 10,
67
+ "max_requests": int(os.getenv("GUNICORN_MAX_REQUESTS", 120 if ENABLE_PROCESS_WRAPPER else 20)),
68
+ "max_requests_jitter": int(os.getenv("GUNICORN_MAX_REQUESTS_JITTER", 30 if ENABLE_PROCESS_WRAPPER else 10)),
69
69
  "worker_class": "gthread",
70
70
  "timeout": max_workflow_runtime_seconds,
71
71
  "logger_class": CustomGunicornLogger,