Enforce workflow concurrency through DBOS worker queues (#732)

`DBOSRuntime` starts every control loop directly, so
`Workflow(num_concurrent_runs=...)` only limits runs under the basic
runtime. DBOS workers cannot apply the limit.

Now `register()` declares one DBOS queue per workflow, named
`_llamaindex_workflow_queue:<workflow_name>`, with `worker_concurrency`
set from `num_concurrent_runs`. Workflows with a limit submit through
the queue, and runs beyond the limit wait as `ENQUEUED`, admitting
within about `polling_interval_sec`. Workflows without a limit keep
starting directly, so the default path has no added latency. The
constructor rejects zero, negative, boolean, and non-integer limits.

The queue is declared even for unlimited workflows, so removing a limit
still leaves a listener for rows that were `ENQUEUED` under the old one.
Declarations live in a process-level map because DBOS's registry
survives `DBOS.destroy()` and rejects redeclaring a name. Recreating the
runtime reuses the queue object and updates its limit in place, which
DBOS's poller reads live.

The limit is per worker, so deployment capacity is the limit times the
number of live workers. Applications that restrict `DBOS.listen_queues`
collect `runtime.workflow_queues` after registering workflows and before
launch.
This commit is contained in:
Adrian Lyjak
2026-08-14 01:19:25 +00:00
committed by GitHub
parent 2eba0b0b16
commit 386f19043f
13 changed files with 677 additions and 15 deletions
+5
View File
@@ -0,0 +1,5 @@
---
"llama-agents-dbos": minor
---
Enforce `num_concurrent_runs` on `DBOSRuntime` as per-worker DBOS queue concurrency. Workflows without a limit keep starting directly.
+5
View File
@@ -0,0 +1,5 @@
---
"llama-index-workflows": patch
---
Reject invalid workflow concurrency limits during construction.
@@ -8,6 +8,50 @@ Workflows can run steps at the same time. When several steps are independent and
The usual pattern is fan-out and fan-in. You split the work into pieces, run them at the same time, then join the results back together. You write this directly in the step signatures. Return a `list` from a step and it fans out, with one event per element. Take a `list` parameter and it fans in, firing once on the whole batch. The `@step` decorator reads those types. The validator and the [visualization](/python/llamaagents/workflows/drawing) then connect each producer step to the steps that consume its events, with no extra work from you. When you need to emit events that do not follow from the signature, you can send them yourself with `ctx.send_event`. The [dynamic API](#the-dynamic-api) at the end of this page covers that.
## Limit simultaneous workflow runs
Whole-run concurrency and step concurrency are separate. Set
`num_concurrent_runs` on a workflow to limit how many calls to `run()` may be
active at once:
```python
workflow = ParallelFlow(num_concurrent_runs=4)
```
The value must be a positive integer or `None`. The default is `None`, which
allows unlimited runs. The limit restricts concurrency within a process:
additional calls wait until an active run finishes.
[DBOS-backed workflows](/python/llamaagents/workflows/dbos#run-concurrency-limits)
also support it, applied per worker.
Use `@step(num_workers=...)` for concurrency inside each run. It controls how
many copies of that step can process events at once and does not limit the
number of workflow runs.
## Limit active steps across runs
To limit how many copies of a step run at once across all runs, for example to
cap concurrent calls to an external API, share an `asyncio.Semaphore` between
them. Steps are plain async functions, so this needs no library support:
```python
import asyncio
API_SLOTS = asyncio.Semaphore(3)
class Fetcher(Workflow):
@step
async def fetch(self, ev: FetchRequest) -> FetchResult:
async with API_SLOTS:
data = await call_external_api(ev.url)
return FetchResult(data=data)
```
Every run shares the one semaphore, so at most three `fetch` steps are in the
API call at any moment, whichever runs they belong to. The semaphore lives in
one Python process. If you deploy multiple processes, each has its own.
## Fan-out: return a list
Return a `list` from a step and each element fires as its own event. Here five `Task`s run concurrently under `work`:
@@ -231,6 +231,22 @@ For production multi-replica deployments, [DBOS Conductor](https://docs.dbos.dev
Understanding the DBOS execution model helps you write workflows that behave correctly across restarts and replicas.
### Workflow identity
A workflow's name is its durable identity. Journal entries, step registrations, and the admission queue are all keyed by it, so everything DBOS has recorded for a workflow is filed under that name. The default is the module-qualified class name (e.g. `my_app.CounterWorkflow`), which means moving or renaming the class silently changes the identity. For anything long-lived, set the name explicitly and treat it as permanent:
```python
wf = CounterWorkflow(runtime=runtime, workflow_name="counter-v1")
```
A worker only looks for recorded work under the names it registers. After a rename, in-flight and queued runs filed under the old name are invisible to the new deployment — keep workers registering the old name running until that work finishes.
When using a server, the name passed to `add_workflow` is the HTTP route name, independent of the workflow's durable name:
```python
server.add_workflow("counter", CounterWorkflow(runtime=runtime, workflow_name="counter-v1"))
```
### Replica ownership
Each replica is identified by its `executor_id` and **owns** every workflow it starts. A workflow and all of its steps run in the same process — there is no distribution of individual steps across replicas. This means your steps can safely rely on local state like in-memory caches, local files, or process-level singletons. The trade-off is that a single workflow's workload cannot be spread across multiple replicas.
@@ -251,19 +267,37 @@ Replica IDs and replica counts must be stable. If you scale down and remove a re
Since resumption is based on journal replay, changing a workflow's code while historical runs are still in progress can cause non-determinism — for example, a step that now accepts a different set of events than when the run was originally started. To avoid this:
- **Drain in-flight workflows** before deploying code changes, or
- **Register the updated workflow under a new name** so that old runs continue against the original code and new runs use the updated version
- **Register the updated workflow under a new name** (e.g. `workflow_name="counter-v2"`) so that old runs continue against the original code and new runs use the updated version. This is a deliberate identity change — see [Workflow identity](#workflow-identity) for what the name keys.
A workflow's name defaults to its module-qualified class name (e.g. `my_app.CounterWorkflow`). You can set it explicitly with the `workflow_name` parameter:
### Run concurrency limits
`Workflow(num_concurrent_runs=N)` limits active runs of that workflow to N per
replica, so deployment capacity is roughly N times the number of replicas.
Runs beyond the limit wait in a DBOS queue and start within about a second of
a slot opening. The queue is shared across replicas: a waiting run has no
affinity to the replica that submitted it, and any replica with a free slot
can pick it up. Leaving the value unset keeps runs starting directly, with no
queue in the path.
```python
wf = CounterWorkflow(runtime=runtime, workflow_name="counter-v2")
wf = CounterWorkflow(runtime=runtime, num_concurrent_runs=4)
```
When using a server, the name passed to `add_workflow` is the HTTP route name, independent of the workflow's internal name:
The queue is keyed by the workflow's durable name (see
[Workflow identity](#workflow-identity)), so runs waiting under an old name
are invisible after a rename. Changing or removing the limit itself is safe:
waiting runs stay on the same queue and keep admitting. A new limit does not
count runs that started before it, so a replica can briefly exceed it while
those finish.
```python
server.add_workflow("counter-v2", CounterWorkflow(runtime=runtime, workflow_name="counter-v2"))
```
A waiting run cannot be cancelled until it starts, because cancellation is a
message delivered to the running workflow. The request is saved, and the run
stops itself as soon as it is admitted.
DBOS normally watches every queue automatically. An application that instead
passes an explicit list to `DBOS.listen_queues` must add this runtime's
queues to it (`runtime.workflow_queues`), collected after registering
workflows and before launch.
### Event streaming behavior
@@ -16,6 +16,36 @@ Each DBOS replica is configured with a unique `executor_id` (e.g. `"replica-8001
The `executor_id` model means horizontal scaling works by adding replicas that each own a slice of the workload, not by distributing individual workflow steps across nodes.
## Workflow Admission Queues
`DBOSRuntime.register()` declares one DBOS queue per workflow, named
`_llamaindex_workflow_queue:<workflow_name>`, with `worker_concurrency` taken
from the workflow's `num_concurrent_runs`. Workflows with a limit submit runs
through the queue. Workflows without one start directly and never touch it.
The queue is declared either way, so removing a limit still leaves a listener
for rows that were `ENQUEUED` under the old one.
DBOS remembers every queue ever declared in the process, even after
`DBOS.destroy()`, and declaring the same name twice raises. So the adapter
keeps its own module-level dict of queue objects. A recreated runtime reuses
the existing object and just changes `worker_concurrency` on it, which DBOS
reads fresh on every poll.
DBOS tags every run with an application version, a fingerprint of the
registered workflow functions, and a worker only picks up runs whose
fingerprint matches its own. This protects a run from being resumed by code
that changed under it. Under this adapter every workflow registers the same
control-loop wrapper, so the fingerprint does not change when users edit
their step code. It does change when a workflow is added or removed, or when
the `dbos` or `llama-agents-dbos` package changes. After such a deployment,
runs tagged with the old fingerprint can only be finished by workers running
the old code, so keep those workers up until that work drains.
Cancellation uses `TickCancelRun` instead of DBOS hard cancellation. A run
still waiting in the queue cannot process it yet; the request is saved, and
the run publishes `WorkflowCancelledEvent` and stops as soon as it is
admitted.
## Process Layout
```
+38
View File
@@ -42,8 +42,46 @@ async def main():
asyncio.run(main())
```
## Workflow concurrency
Set `num_concurrent_runs` to limit how many runs of a workflow may be active
at once on each DBOS worker:
```python
workflow = MyWorkflow(runtime=runtime, num_concurrent_runs=8)
```
The default is `None`, which is unlimited. Unlimited workflows start directly,
with no queue in the path. Limited workflows submit through a DBOS queue named
`_llamaindex_workflow_queue:<workflow_name>`, and runs beyond the limit wait as
`ENQUEUED`. Admission takes about the configured
`DBOSRuntime(queue_polling_interval_sec=...)`, one second by default. Capacity across
a deployment is the limit times the number of workers. The queue is shared, so
an enqueued run has no affinity to the replica that submitted it. Any worker
with a free slot can pick it up.
The runtime declares the queue for every workflow, limited or not, so turning a
limit on or off never strands queued work. A new limit does not count runs that
started before it, so a worker can briefly exceed the limit while those finish.
Waiting runs are rows in the database, filed under the workflow's name
(`workflow_name`, defaulting to the Python module and class name). A worker
only looks for waiting work under the names it knows, so if you rename a
workflow, rows filed under the old name are invisible to the new deployment.
Keep old workers running until they finish that work.
A run that is still waiting in the queue cannot be cancelled yet, because
cancellation is a message delivered to the running workflow. The request is
saved, and the run stops itself as soon as it starts.
DBOS normally watches every queue automatically. An application that instead
passes an explicit list to `DBOS.listen_queues` must add this runtime's
queues to it (`runtime.workflow_queues`), or waiting runs are never picked
up. Build the list after registering workflows and before launch.
## Features
- Durable workflow execution backed by DBOS
- Automatic step recording and replay
- Distributed workers and recovery support
- Per-worker workflow concurrency with an unlimited default
@@ -15,6 +15,7 @@ import sqlite3
import threading
import time
from collections.abc import AsyncIterator, Awaitable, Callable
from dataclasses import dataclass
from typing import Any, AsyncGenerator, TypedDict, cast
import asyncpg
@@ -73,13 +74,15 @@ from workflows.runtime.types.named_task import (
from workflows.runtime.types.plugin import (
ExternalRunAdapter,
InternalRunAdapter,
RegisteredWorkflow,
Runtime,
WaitForNextTaskResult,
WaitResult,
WaitResultTick,
WaitResultTimeout,
)
from workflows.runtime.types.plugin import (
RegisteredWorkflow as BaseRegisteredWorkflow,
)
from workflows.runtime.types.step_function import (
StepWorkerFunction,
as_step_worker_functions,
@@ -88,7 +91,7 @@ from workflows.runtime.types.step_function import (
from workflows.runtime.types.ticks import WorkflowTick
from workflows.workflow import Workflow
from dbos import DBOS, SetWorkflowID, WorkflowHandleAsync
from dbos import DBOS, Queue, SetWorkflowID, WorkflowHandleAsync
from dbos._context import get_local_dbos_context
from dbos._dbos import _get_dbos_instance
from dbos._error import DBOSNonExistentWorkflowError
@@ -113,6 +116,35 @@ STATE_TABLE_NAME = "workflow_state"
logger = logging.getLogger(__name__)
@dataclass
class RegisteredWorkflow(BaseRegisteredWorkflow):
queue: Queue
_workflow_queues: dict[str, Queue] = {}
_workflow_queues_lock = threading.Lock()
def _declare_workflow_queue(
name: str,
worker_concurrency: int | None,
polling_interval_sec: float,
) -> Queue:
with _workflow_queues_lock:
queue = _workflow_queues.get(name)
if queue is None:
queue = Queue(
name,
worker_concurrency=worker_concurrency,
polling_interval_sec=polling_interval_sec,
)
_workflow_queues[name] = queue
else:
queue.worker_concurrency = worker_concurrency
queue.polling_interval_sec = polling_interval_sec
return queue
class DBOSWorkflowStore(AbstractWorkflowStore):
"""Lazy proxy that defers dialect resolution until first use.
@@ -217,6 +249,7 @@ class DBOSRuntimeConfig(TypedDict, total=False):
"""
polling_interval_sec: float
queue_polling_interval_sec: float | None
run_migrations_on_launch: bool
schema: str | None
state_table_name: str
@@ -292,6 +325,10 @@ class DBOSRuntime(Runtime):
Args:
**kwargs: Configuration options. See DBOSRuntimeConfig for details.
polling_interval_sec: Interval for polling workflow results. Default 1.0.
queue_polling_interval_sec: Interval for polling workflow
admission queues. Bounds how long a run limited by
``num_concurrent_runs`` waits for a free slot. Defaults to
``polling_interval_sec``.
run_migrations_on_launch: Auto-run migrations on launch(). Default True.
schema: Database schema name. Default: auto-detected at launch
("dbos" for PostgreSQL, None for SQLite). Pass None explicitly
@@ -379,6 +416,14 @@ class DBOSRuntime(Runtime):
self._tracked_workflows.append(workflow)
self._tracked_workflow_ids.add(wf_id)
@property
def workflow_queues(self) -> tuple[Queue, ...]:
"""Queues declared by this runtime, in registration order."""
queues = dict.fromkeys(
registered.queue for registered in self._registered.values()
)
return tuple(queues)
def get_registered(self, workflow: Workflow) -> RegisteredWorkflow | None:
"""Get the registered workflow if available."""
return self._registered.get(id(workflow))
@@ -398,6 +443,14 @@ class DBOSRuntime(Runtime):
# Use workflow's name directly
name = workflow.workflow_name
queue_polling_interval = self.config.get("queue_polling_interval_sec")
if queue_polling_interval is None:
queue_polling_interval = self.config.get("polling_interval_sec", 1.0)
queue = _declare_workflow_queue(
f"_llamaindex_workflow_queue:{name}",
workflow._num_concurrent_runs,
queue_polling_interval,
)
# Create DBOS-wrapped control loop with stable name
wf_kwargs: dict[str, Any] = {"name": f"{name}.control_loop"}
@@ -426,7 +479,10 @@ class DBOSRuntime(Runtime):
}
registered = RegisteredWorkflow(
workflow=workflow, workflow_run_fn=_dbos_control_loop, steps=wrapped_steps
workflow=workflow,
workflow_run_fn=_dbos_control_loop,
steps=wrapped_steps,
queue=queue,
)
self._registered[id(workflow)] = registered
return registered
@@ -612,6 +668,13 @@ class DBOSRuntime(Runtime):
await store.ensure_seeded()
try:
if workflow._num_concurrent_runs is not None:
return await registered.queue.enqueue_async(
registered.workflow_run_fn,
init_state,
start_event,
get_dispatcher().capture_propagation_context(),
)
return await DBOS.start_workflow_async(
registered.workflow_run_fn,
init_state,
@@ -43,7 +43,9 @@ async def main() -> None:
"executor_id": args.executor_id or f"test-replica-{args.port}",
}
DBOS(config=config)
dbos_runtime = DBOSRuntime(polling_interval_sec=0.01)
dbos_runtime = DBOSRuntime(
polling_interval_sec=0.01, queue_polling_interval_sec=1.0
)
wf = workflow_class(runtime=dbos_runtime)
await dbos_runtime.launch()
+4 -1
View File
@@ -64,4 +64,7 @@ def setup_dbos(db_url: str, app_name: str = "test-workflow") -> DBOSRuntime:
"notification_listener_polling_interval_sec": 0.01,
}
DBOS(config=config)
return DBOSRuntime(polling_interval_sec=0.01)
# Keep queue polling at its slow default: these workflows are unlimited, and
# a 10ms queue poll writes-locks the SQLite file hard enough to starve the
# workflow's own journal writes on slow CI runners.
return DBOSRuntime(polling_interval_sec=0.01, queue_polling_interval_sec=1.0)
@@ -17,6 +17,7 @@ from unittest.mock import AsyncMock, patch
import asyncpg
import pytest
from dbos import DBOS, DBOSConfig
from dbos._context import get_local_dbos_context
from llama_agents.dbos import DBOSRuntime
from llama_agents.dbos.journal.crud import SqliteJournalCrud
from llama_agents.dbos.journal.task_journal import TaskJournal
@@ -341,12 +342,24 @@ async def test_run_workflow_seeds_state_store_from_durable_handle() -> None:
return StopEvent(result="done")
runtime = DBOSRuntime(polling_interval_sec=0.01)
workflow = SimpleWf(runtime=runtime)
runtime._dbos_launched = True
workflow = SimpleWf()
workflow_store = RecordingWorkflowStore()
serialized_state = {"store_type": "sqlite", "run_id": "old-run"}
serializer = JsonSerializer()
fake_handle = AsyncMock()
start_observations: dict[str, Any] = {}
async def start_workflow_async(*args: Any, **kwargs: Any) -> Any:
context = get_local_dbos_context()
assert context is not None
start_observations["run_id"] = context.id_assigned_for_next_workflow
start_observations["state_seeded"] = (
workflow_store.state_store.ensure_seeded_called
)
return fake_handle
start_mock = AsyncMock(side_effect=start_workflow_async)
with (
patch.object(runtime, "create_workflow_store", return_value=workflow_store),
@@ -359,7 +372,7 @@ async def test_run_workflow_seeds_state_store_from_durable_handle() -> None:
),
patch(
"llama_agents.dbos.runtime.DBOS.start_workflow_async",
new=AsyncMock(return_value=fake_handle),
new=start_mock,
),
):
adapter = runtime.run_workflow(
@@ -376,6 +389,11 @@ async def test_run_workflow_seeds_state_store_from_durable_handle() -> None:
assert workflow_store.create_state_store_calls == [
("run-1", DictState, serialized_state, serializer)
]
start_mock.assert_awaited_once()
start_args = start_mock.await_args
assert start_args is not None
assert start_args.args[0] is workflow_run_fn
assert start_observations == {"run_id": "run-1", "state_seeded": True}
@pytest.mark.asyncio
@@ -0,0 +1,384 @@
# SPDX-License-Identifier: MIT
# Copyright (c) 2026 LlamaIndex Inc.
from __future__ import annotations
import asyncio
import threading
from pathlib import Path
from typing import Any
import pytest
from dbos import DBOS, DBOSConfig
from llama_agents.dbos import DBOSRuntime
from workflows.decorators import step
from workflows.errors import WorkflowCancelledByUser
from workflows.events import StartEvent, StopEvent, WorkflowCancelledEvent
from workflows.workflow import Workflow
class _AdmissionGate:
def __init__(self) -> None:
self._condition = threading.Condition()
self._release = threading.Event()
self.started: list[str] = []
async def enter(self, run: str) -> None:
with self._condition:
self.started.append(run)
self._condition.notify_all()
await asyncio.to_thread(self._release.wait)
def wait_for_started(self, count: int, timeout: float = 5.0) -> bool:
with self._condition:
return self._condition.wait_for(
lambda: len(self.started) >= count,
timeout=timeout,
)
def release(self) -> None:
self._release.set()
def _dbos_config(db_path: Path, name: str) -> DBOSConfig:
return {
"name": name,
"system_database_url": (
f"sqlite+pysqlite:///{db_path}?check_same_thread=false"
),
"run_admin_server": False,
} # type: ignore[return-value]
def _blocking_workflow_type(
gate: _AdmissionGate,
*,
class_name: str,
) -> type[Workflow]:
class BlockingWorkflow(Workflow):
@step
async def block(self, ev: StartEvent) -> StopEvent:
run = ev.get("run")
assert isinstance(run, str)
await gate.enter(run)
return StopEvent(result=run)
BlockingWorkflow.__name__ = class_name
BlockingWorkflow.__qualname__ = class_name
return BlockingWorkflow
async def _run_limited_admission_case(tmp_path: Path) -> None:
DBOS(
config=_dbos_config(
tmp_path / "limited-admission.sqlite3",
"dbos-limited-admission-test",
)
)
runtime = DBOSRuntime(polling_interval_sec=0.01, queue_polling_interval_sec=0.01)
gate = _AdmissionGate()
workflow = _blocking_workflow_type(
gate,
class_name="LimitedAdmissionWorkflow",
)(
runtime=runtime,
workflow_name="tests.dbos.limited-admission",
num_concurrent_runs=1,
timeout=10,
)
handlers: dict[str, Any] = {}
try:
await runtime.launch()
handlers = {
run_id: workflow.run(run=run_id, run_id=run_id)
for run_id in ("limited-first", "limited-second")
}
admitted = await asyncio.to_thread(gate.wait_for_started, 1)
assert admitted
assert len(gate.started) == 1
queued_run = (
"limited-second" if gate.started == ["limited-first"] else "limited-first"
)
status = await DBOS.get_workflow_status_async(queued_run)
assert status is not None
assert status.status == "ENQUEUED"
gate.release()
results = await asyncio.wait_for(
asyncio.gather(*handlers.values()),
timeout=10,
)
assert set(results) == {"limited-first", "limited-second"}
assert set(gate.started) == {"limited-first", "limited-second"}
finally:
gate.release()
if handlers:
await asyncio.gather(*handlers.values(), return_exceptions=True)
await runtime.destroy()
def test_limited_workflow_queues_excess_runs(tmp_path: Path) -> None:
asyncio.run(_run_limited_admission_case(tmp_path))
async def _run_unlimited_admission_case(tmp_path: Path) -> None:
DBOS(
config=_dbos_config(
tmp_path / "unlimited-admission.sqlite3",
"dbos-unlimited-admission-test",
)
)
runtime = DBOSRuntime(polling_interval_sec=0.01, queue_polling_interval_sec=0.01)
gate = _AdmissionGate()
workflow = _blocking_workflow_type(
gate,
class_name="UnlimitedAdmissionWorkflow",
)(
runtime=runtime,
workflow_name="tests.dbos.unlimited-admission",
timeout=10,
)
handlers: dict[str, Any] = {}
try:
await runtime.launch()
handlers = {
run_id: workflow.run(run=run_id, run_id=run_id)
for run_id in ("unlimited-first", "unlimited-second")
}
admitted = await asyncio.to_thread(gate.wait_for_started, 2)
assert admitted
assert set(gate.started) == {"unlimited-first", "unlimited-second"}
statuses = await asyncio.gather(
*(DBOS.get_workflow_status_async(run_id) for run_id in handlers)
)
assert all(status is not None for status in statuses)
assert all(status.status != "ENQUEUED" for status in statuses if status)
gate.release()
results = await asyncio.wait_for(
asyncio.gather(*handlers.values()),
timeout=10,
)
assert set(results) == {"unlimited-first", "unlimited-second"}
finally:
gate.release()
if handlers:
await asyncio.gather(*handlers.values(), return_exceptions=True)
await runtime.destroy()
def test_unlimited_workflow_starts_runs_directly(tmp_path: Path) -> None:
asyncio.run(_run_unlimited_admission_case(tmp_path))
async def _run_queue_declaration_case(tmp_path: Path) -> None:
DBOS(
config=_dbos_config(
tmp_path / "queue-declaration.sqlite3",
"dbos-queue-declaration-test",
)
)
runtime = DBOSRuntime(queue_polling_interval_sec=0.125)
_blocking_workflow_type(
_AdmissionGate(),
class_name="LimitedQueueWorkflow",
)(
runtime=runtime,
workflow_name="tests.dbos.limited-queue",
num_concurrent_runs=3,
)
_blocking_workflow_type(
_AdmissionGate(),
class_name="UnlimitedQueueWorkflow",
)(
runtime=runtime,
workflow_name="tests.dbos.unlimited-queue",
)
try:
await runtime.launch()
queues = {queue.name: queue for queue in runtime.workflow_queues}
assert list(queues) == [
"_llamaindex_workflow_queue:tests.dbos.limited-queue",
"_llamaindex_workflow_queue:tests.dbos.unlimited-queue",
]
assert (
queues[
"_llamaindex_workflow_queue:tests.dbos.limited-queue"
].worker_concurrency
== 3
)
assert (
queues[
"_llamaindex_workflow_queue:tests.dbos.unlimited-queue"
].worker_concurrency
is None
)
finally:
await runtime.destroy()
def test_runtime_declares_queue_for_every_workflow(tmp_path: Path) -> None:
asyncio.run(_run_queue_declaration_case(tmp_path))
async def _run_destroy_relaunch_case(tmp_path: Path) -> None:
config = _dbos_config(
tmp_path / "reused-queue.sqlite3",
"dbos-reused-queue-test",
)
workflow_name = "tests.dbos.reused-queue"
DBOS(config=config)
first_runtime = DBOSRuntime(
polling_interval_sec=0.01, queue_polling_interval_sec=0.01
)
_blocking_workflow_type(
_AdmissionGate(),
class_name="FirstReusedQueueWorkflow",
)(
runtime=first_runtime,
workflow_name=workflow_name,
num_concurrent_runs=1,
)
await first_runtime.launch()
queue = first_runtime.workflow_queues[0]
await first_runtime.destroy(destroy_dbos=True)
DBOS(config=config)
second_runtime = DBOSRuntime(
polling_interval_sec=0.01, queue_polling_interval_sec=0.02
)
_blocking_workflow_type(
_AdmissionGate(),
class_name="SecondReusedQueueWorkflow",
)(
runtime=second_runtime,
workflow_name=workflow_name,
num_concurrent_runs=2,
)
try:
await second_runtime.launch()
assert second_runtime.workflow_queues == (queue,)
assert queue.worker_concurrency == 2
assert queue.polling_interval_sec == 0.02
finally:
await second_runtime.destroy(destroy_dbos=True)
def test_destroy_relaunch_reuses_queue_with_updated_limit(tmp_path: Path) -> None:
asyncio.run(_run_destroy_relaunch_case(tmp_path))
async def _run_restricted_listener_case(tmp_path: Path) -> None:
DBOS(
config=_dbos_config(
tmp_path / "restricted-listener.sqlite3",
"dbos-restricted-listener-test",
)
)
runtime = DBOSRuntime(polling_interval_sec=0.01, queue_polling_interval_sec=0.01)
gate = _AdmissionGate()
workflow = _blocking_workflow_type(
gate,
class_name="RestrictedListenerWorkflow",
)(
runtime=runtime,
workflow_name="tests.dbos.restricted-listener",
num_concurrent_runs=1,
timeout=10,
)
runtime.register(workflow)
DBOS.listen_queues(list(runtime.workflow_queues))
handlers: dict[str, Any] = {}
try:
await runtime.launch()
handlers = {
run_id: workflow.run(run=run_id, run_id=run_id)
for run_id in ("restricted-first", "restricted-second")
}
admitted = await asyncio.to_thread(gate.wait_for_started, 1)
assert admitted
gate.release()
results = await asyncio.wait_for(
asyncio.gather(*handlers.values()),
timeout=10,
)
assert set(results) == {"restricted-first", "restricted-second"}
finally:
gate.release()
if handlers:
await asyncio.gather(*handlers.values(), return_exceptions=True)
await runtime.destroy()
def test_restricted_listener_admits_runtime_queues(tmp_path: Path) -> None:
asyncio.run(_run_restricted_listener_case(tmp_path))
async def _run_queued_cancellation_case(tmp_path: Path) -> None:
DBOS(
config=_dbos_config(
tmp_path / "queued-cancellation.sqlite3",
"dbos-queued-cancellation-test",
)
)
runtime = DBOSRuntime(polling_interval_sec=0.01, queue_polling_interval_sec=0.01)
gate = _AdmissionGate()
class QueuedCancellationWorkflow(Workflow):
@step
async def block(self, ev: StartEvent) -> StopEvent:
run = ev.get("run")
assert isinstance(run, str)
await gate.enter(run)
await asyncio.sleep(0.2)
return StopEvent(result=run)
workflow = QueuedCancellationWorkflow(
runtime=runtime,
workflow_name="tests.dbos.queued-cancellation",
num_concurrent_runs=1,
timeout=10,
)
handlers: dict[str, Any] = {}
try:
await runtime.launch()
handlers = {
run_id: workflow.run(run=run_id, run_id=run_id)
for run_id in ("cancel-first", "cancel-second")
}
admitted = await asyncio.to_thread(gate.wait_for_started, 1)
assert admitted
queued_run = (
"cancel-second" if gate.started == ["cancel-first"] else "cancel-first"
)
queued_handler = handlers[queued_run]
await queued_handler._external_adapter.cancel()
status = await DBOS.get_workflow_status_async(queued_run)
assert status is not None
assert status.status == "ENQUEUED"
gate.release()
await handlers[gate.started[0]]
with pytest.raises(WorkflowCancelledByUser):
await queued_handler
assert isinstance(queued_handler.get_stop_event(), WorkflowCancelledEvent)
finally:
gate.release()
if handlers:
await asyncio.gather(*handlers.values(), return_exceptions=True)
await runtime.destroy()
def test_enqueued_cancellation_waits_for_admission(tmp_path: Path) -> None:
asyncio.run(_run_queued_cancellation_case(tmp_path))
@@ -116,7 +116,11 @@ class Workflow(metaclass=WorkflowMeta):
verbose (bool): If True, print step activity.
resource_manager (ResourceManager | None): Custom resource manager
for dependency injection.
num_concurrent_runs (int | None): Limit on concurrent `run()` calls.
num_concurrent_runs (int | None): Maximum number of active runs for
this workflow. Must be a positive integer or `None`. The
default, `None`, allows unlimited runs. How the limit is
scoped is up to the runtime; the basic runtime applies it to
this workflow instance within the process.
runtime (Runtime | None): Optional runtime to use for this workflow.
If not provided, uses the current context-scoped runtime or
falls back to basic_runtime.
@@ -139,6 +143,14 @@ class Workflow(metaclass=WorkflowMeta):
)
# Configuration
if num_concurrent_runs is not None and (
isinstance(num_concurrent_runs, bool)
or not isinstance(num_concurrent_runs, int)
or num_concurrent_runs <= 0
):
raise WorkflowValidationError(
"num_concurrent_runs must be an integer greater than 0 or None"
)
self._timeout = timeout
self._verbose = verbose
self._disable_validation = disable_validation
@@ -395,6 +395,30 @@ class DummyWorkflowForConcurrentRunsTest(Workflow):
return self.num_active_runs_history
@pytest.mark.parametrize("num_concurrent_runs", [0, -1, True, 1.5, "1"])
def test_workflow_rejects_invalid_num_concurrent_runs(
num_concurrent_runs: Any,
) -> None:
with pytest.raises(
WorkflowValidationError,
match="num_concurrent_runs must be an integer greater than 0",
):
DummyWorkflowForConcurrentRunsTest(
num_concurrent_runs=num_concurrent_runs,
)
@pytest.mark.parametrize("num_concurrent_runs", [None, 1, 8])
def test_workflow_accepts_valid_num_concurrent_runs(
num_concurrent_runs: int | None,
) -> None:
workflow = DummyWorkflowForConcurrentRunsTest(
num_concurrent_runs=num_concurrent_runs,
)
assert workflow._num_concurrent_runs == num_concurrent_runs
@pytest.mark.asyncio
@pytest.mark.parametrize(
(