mirror of
https://github.com/run-llama/workflows-py.git
synced 2026-08-24 20:01:34 -04:00
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:
@@ -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.
|
||||
@@ -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
|
||||
|
||||
```
|
||||
|
||||
@@ -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()
|
||||
|
||||
@@ -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(
|
||||
(
|
||||
|
||||
Reference in New Issue
Block a user