run-llama/workflows-py
1
0
Fork 0
mirror of https://github.com/run-llama/workflows-py.git synced 2026-07-16 05:31:56 -04:00
Code Issues 18 Packages Projects Releases 10 Wiki Activity

Changes to facilitate pluggable server runtime (#348)

Browse Source
This commit is contained in:
Adrian Lyjak
2026-02-09 15:44:08 -05:00
parent 92b0d02343
commit 3c22216cfd
15 changed files with 1011 additions and 228 deletions
Download Patch File Download Diff File Expand all files Collapse all files
.changeset/neat-chicken-bake.md
+5
View File
@@ -0,0 +1,5 @@
---
"llama-index-workflows": minor
---
Make WorkflowTick serializable, and support switching workflow name and runtime before launch
packages/llama-index-workflows/src/workflows/context/state_store.py
+146 -77
View File
@@ -192,6 +192,121 @@ def assign_path_step(obj: Any, segment: str, value: Any) -> None:
setattr(obj, segment, value)
def get_by_path(state: Any, path: str, default: Any = Ellipsis) -> Any:
"""Get a nested value from state using a dot-separated path.
Args:
state: The root state object.
path: Dot-separated path, e.g. "user.profile.name".
default: If provided, return this when the path does not exist;
otherwise, raise ValueError.
Returns:
The resolved value.
Raises:
ValueError: If the path is invalid and no default is provided,
or if path depth exceeds MAX_DEPTH.
"""
segments = path.split(".") if path else []
if len(segments) > MAX_DEPTH:
raise ValueError(f"Path length exceeds {MAX_DEPTH} segments")
try:
value: Any = state
for segment in segments:
value = traverse_path_step(value, segment)
except Exception:
if default is not Ellipsis:
return default
raise ValueError(f"Path '{path}' not found in state")
return value
def set_by_path(state: Any, path: str, value: Any) -> None:
"""Set a nested value on state using a dot-separated path.
Intermediate dicts are created as needed.
Args:
state: The root state object (mutated in place).
path: Dot-separated path to write.
value: Value to assign.
Raises:
ValueError: If the path is empty or exceeds MAX_DEPTH.
"""
if not path:
raise ValueError("Path cannot be empty")
segments = path.split(".")
if len(segments) > MAX_DEPTH:
raise ValueError(f"Path length exceeds {MAX_DEPTH} segments")
current = state
for segment in segments[:-1]:
try:
current = traverse_path_step(current, segment)
except (KeyError, AttributeError, IndexError, TypeError):
intermediate: Any = {}
assign_path_step(current, segment, intermediate)
current = intermediate
assign_path_step(current, segments[-1], value)
def merge_state(current_state: MODEL_T, incoming: BaseModel) -> MODEL_T:
"""Replace or merge incoming state onto current state.
If incoming is the same type (or subclass) of current, it replaces directly.
If current's type is a subclass of incoming's type (parent provided),
fields are merged preserving child-specific fields.
Args:
current_state: The existing state.
incoming: The new state to apply.
Returns:
The resulting state after merge/replace.
Raises:
ValueError: If the types are not compatible.
"""
current_type = type(current_state)
new_type = type(incoming)
if isinstance(incoming, current_type):
return incoming # type: ignore[return-value]
elif issubclass(current_type, new_type):
parent_data = incoming.model_dump()
return current_type.model_validate(
{**current_state.model_dump(), **parent_data}
)
else:
raise ValueError(
f"State must be of type {current_type.__name__} or a parent type, "
f"got {new_type.__name__}"
)
def create_cleared_state(state_type: Type[MODEL_T]) -> MODEL_T:
"""Create a default instance of the state type, wrapping ValidationError.
Args:
state_type: The state model class to instantiate.
Returns:
A new default instance.
Raises:
ValueError: If the model cannot be instantiated from defaults.
"""
try:
return state_type()
except ValidationError:
raise ValueError("State must have defaults for all fields")
# Only warn once about unserializable keys
class UnserializableKeyWarning(Warning):
pass
@@ -377,27 +492,8 @@ class InMemoryStateStore(Generic[MODEL_T]):
Raises:
ValueError: If the types are not compatible (neither same nor parent).
"""
current_type = type(self._state)
new_type = type(state)
if isinstance(state, current_type):
# Exact match or subclass - direct replacement
async with self._lock:
self._state = state
elif issubclass(current_type, new_type):
# Parent type provided - merge fields onto current state
# This preserves child-specific fields while updating parent fields
async with self._lock:
# Get the fields from the parent type and update them on the current state
parent_data = state.model_dump()
self._state = current_type.model_validate(
{**self._state.model_dump(), **parent_data}
)
else:
raise ValueError(
f"State must be of type {current_type.__name__} or a parent type, "
f"got {new_type.__name__}"
)
async with self._lock:
self._state = merge_state(self._state, state)
def to_dict(self, serializer: "BaseSerializer") -> dict[str, Any]:
"""Serialize the state and model metadata for persistence.
@@ -465,9 +561,6 @@ class InMemoryStateStore(Generic[MODEL_T]):
async def get(self, path: str, default: Any = Ellipsis) -> Any:
"""Get a nested value using dot-separated paths.
Supports dict keys, list indices, and attribute access transparently at
each segment.
Args:
path (str): Dot-separated path, e.g. "user.profile.name".
default (Any): If provided, return this when the path does not
@@ -480,30 +573,12 @@ class InMemoryStateStore(Generic[MODEL_T]):
ValueError: If the path is invalid and no default is provided or if
the path depth exceeds limits.
"""
segments = path.split(".") if path else []
if len(segments) > MAX_DEPTH:
raise ValueError(f"Path length exceeds {MAX_DEPTH} segments")
async with self._lock:
try:
value: Any = self._state
for segment in segments:
value = traverse_path_step(value, segment)
except Exception:
if default is not Ellipsis:
return default
msg = f"Path '{path}' not found in state"
raise ValueError(msg)
return value
return get_by_path(self._state, path, default)
async def set(self, path: str, value: Any) -> None:
"""Set a nested value using dot-separated paths.
Intermediate containers are created as needed. Dicts, lists, tuples, and
Pydantic models are supported where appropriate.
Args:
path (str): Dot-separated path to write.
value (Any): Value to assign.
@@ -511,28 +586,8 @@ class InMemoryStateStore(Generic[MODEL_T]):
Raises:
ValueError: If the path is empty or exceeds the maximum depth.
"""
if not path:
raise ValueError("Path cannot be empty")
segments = path.split(".")
if len(segments) > MAX_DEPTH:
raise ValueError(f"Path length exceeds {MAX_DEPTH} segments")
async with self._lock:
current = self._state
# Navigate/create intermediate segments
for segment in segments[:-1]:
try:
current = traverse_path_step(current, segment)
except (KeyError, AttributeError, IndexError, TypeError):
# Create intermediate object and assign it
intermediate: Any = {}
assign_path_step(current, segment, intermediate)
current = intermediate
# Assign the final value
assign_path_step(current, segments[-1], value)
set_by_path(self._state, path, value)
async def clear(self) -> None:
"""Reset the state to its type defaults.
@@ -541,10 +596,33 @@ class InMemoryStateStore(Generic[MODEL_T]):
ValueError: If the model type cannot be instantiated from defaults
(i.e., fields missing default values).
"""
await self.set_state(create_cleared_state(self._state.__class__))
def deserialize_dict_state_data(
data: dict[str, Any],
serializer: BaseSerializer,
) -> DictState:
"""Deserialize DictState from {"_data": {...}} format.
Args:
data: Dict with {"_data": {...}} structure containing serialized values.
serializer: Strategy for decoding values.
Returns:
DictState with deserialized values.
Raises:
ValueError: If deserialization fails for any key.
"""
_data_serialized = data.get("_data", {})
deserialized_data = {}
for key, value in _data_serialized.items():
try:
await self.set_state(self._state.__class__())
except ValidationError:
raise ValueError("State must have defaults for all fields")
deserialized_data[key] = serializer.deserialize(value)
except Exception as e:
raise ValueError(f"Failed to deserialize state value for key {key}: {e}")
return DictState(_data=deserialized_data)
def deserialize_state_from_dict(
@@ -570,16 +648,7 @@ def deserialize_state_from_dict(
state_type_name = serialized_state.get("state_type", "DictState")
if state_type_name == "DictState":
_data_serialized = state_data.get("_data", {})
deserialized_data = {}
for key, value in _data_serialized.items():
try:
deserialized_data[key] = serializer.deserialize(value)
except Exception as e:
raise ValueError(
f"Failed to deserialize state value for key {key}: {e}"
)
return DictState(_data=deserialized_data)
return deserialize_dict_state_data(state_data, serializer)
else:
return serializer.deserialize(state_data)
packages/llama-index-workflows/src/workflows/plugins/basic.py
+6 -6
View File
@@ -128,7 +128,7 @@ class InternalAsyncioAdapter(InternalRunAdapter, SnapshottableAdapter):
except asyncio.TimeoutError:
return WaitResultTimeout()
def on_tick(self, tick: WorkflowTick) -> None:
async def on_tick(self, tick: WorkflowTick) -> None:
self._queues.ticks.append(tick)
def replay(self) -> list[WorkflowTick]:
@@ -148,7 +148,8 @@ class ExternalAsyncioAdapter(
and stream events published by the workflow.
"""
def __init__(self, queues: AsyncioAdapterQueues) -> None:
def __init__(self, outer: BasicRuntime, queues: AsyncioAdapterQueues) -> None:
self._outer = outer
self._queues = queues
@property
@@ -171,9 +172,6 @@ class ExternalAsyncioAdapter(
if isinstance(item, StopEvent):
break
def on_tick(self, tick: WorkflowTick) -> None:
self._queues.ticks.append(tick)
def replay(self) -> list[WorkflowTick]:
return self._queues.ticks
@@ -196,6 +194,7 @@ class ExternalAsyncioAdapter(
"""Abort by cancelling the control loop task."""
if not self._queues.complete.done():
self._queues.complete.cancel()
self._outer._queues.pop(self.run_id, None)
@property
def init_state(self) -> BrokerState:
@@ -206,6 +205,7 @@ class BasicRuntime(Runtime):
"""Default asyncio-based runtime with no durability."""
def __init__(self) -> None:
super().__init__()
# WeakValueDictionary allows queues to be GC'd when no adapters reference them.
# The task closure in run_workflow() captures a strong reference, keeping
# queues alive for fire-and-forget workflows even if the external adapter is dropped.
@@ -323,7 +323,7 @@ class BasicRuntime(Runtime):
def get_external_adapter(self, run_id: str) -> ExternalRunAdapter:
if run_id not in self._queues:
raise RuntimeError(f"No active workflow with run_id '{run_id}'. ")
return ExternalAsyncioAdapter(self._queues[run_id])
return ExternalAsyncioAdapter(self, self._queues[run_id])
_current_run_id: ContextVar[str | None] = ContextVar("current_run_id", default=None)
packages/llama-index-workflows/src/workflows/runtime/control_loop.py
+2 -6
View File
@@ -11,7 +11,6 @@ import traceback
from dataclasses import replace
from typing import TYPE_CHECKING
from workflows.decorators import R
from workflows.errors import (
WorkflowCancelledByUser,
WorkflowRuntimeError,
@@ -50,7 +49,6 @@ from workflows.runtime.types.named_task import NamedTask
from workflows.runtime.types.plugin import (
InternalRunAdapter,
WaitResultTick,
as_snapshottable_adapter,
get_current_run,
)
from workflows.runtime.types.results import (
@@ -128,7 +126,6 @@ class _ControlLoopRunner:
# avoiding TypeError from comparing WorkflowTick objects that don't implement __lt__
self.scheduled_wakeups: list[tuple[float, int, WorkflowTick]] = []
self._wakeup_sequence = 0
self.snapshot_adapter = as_snapshottable_adapter(adapter)
# Pull task sequence counter for deterministic journaling
self._pull_sequence = 0
# Map from worker task to (step_name, worker_id) key
@@ -431,8 +428,7 @@ class _ControlLoopRunner:
)
raise
if self.snapshot_adapter is not None:
self.snapshot_adapter.on_tick(tick)
await self.adapter.on_tick(tick)
for command in commands:
try:
@@ -557,7 +553,7 @@ def _check_idle_state(state: BrokerState) -> bool:
def _process_step_result_tick(
tick: TickStepResult[R], init: BrokerState, now_seconds: float
tick: TickStepResult, init: BrokerState, now_seconds: float
) -> tuple[BrokerState, list[WorkflowCommand]]:
"""
processes the results from a step function execution
packages/llama-index-workflows/src/workflows/runtime/types/plugin.py
+78 -31
View File
@@ -7,6 +7,7 @@ A runtime interface to switch out a broker runtime (external library or service
from __future__ import annotations
import asyncio
import weakref
from abc import ABC, abstractmethod
from contextlib import contextmanager
from contextvars import ContextVar, Token
@@ -63,7 +64,7 @@ WaitResult = Union[WaitResultTick, WaitResultTimeout]
class RegisteredWorkflow:
workflow: Workflow
workflow_run_fn: WorkflowRunFunction
steps: dict[str, StepWorkerFunction[Any]]
steps: dict[str, StepWorkerFunction]
class InternalRunAdapter(ABC):
@@ -180,6 +181,17 @@ class InternalRunAdapter(ABC):
"""
pass
async def on_tick(self, tick: WorkflowTick) -> None:
"""
Called whenever a tick event is processed by the control loop.
This method is invoked for both external ticks (sent via send_event)
and internal ticks (generated by step completions, timeouts, etc.).
Adapters can override to record ticks, persist them, etc.
Default is no-op.
"""
pass
async def wait_for_next_task(
self,
task_set: list[NamedTask],
@@ -310,7 +322,7 @@ class RunContext:
workflow: Workflow
run_adapter: InternalRunAdapter
context: Context
steps: dict[str, StepWorkerFunction[Any]]
steps: dict[str, StepWorkerFunction]
_current_run: ContextVar[RunContext | None] = ContextVar("current_run", default=None)
@@ -334,6 +346,48 @@ def get_current_run() -> RunContext:
return ctx
class WorkflowSet:
"""Identity-based weak set for tracking Workflow instances.
Uses id() as the key and weakref.ref with cleanup callbacks to
avoid hashability requirements and memory leaks.
"""
def __init__(self) -> None:
self._refs: dict[int, weakref.ref[Workflow]] = {}
def add(self, workflow: Workflow) -> None:
obj_id = id(workflow)
if obj_id in self._refs:
return
def _cleanup(ref: weakref.ref[Workflow], _id: int = obj_id) -> None:
self._refs.pop(_id, None)
self._refs[obj_id] = weakref.ref(workflow, _cleanup)
def discard(self, workflow: Workflow) -> None:
self._refs.pop(id(workflow), None)
def __contains__(self, workflow: Workflow) -> bool:
ref = self._refs.get(id(workflow))
if ref is None:
return False
return ref() is not None
def __iter__(self) -> Generator[Workflow, None, None]:
for ref in list(self._refs.values()):
obj = ref()
if obj is not None:
yield obj
def __len__(self) -> int:
return sum(1 for _ in self)
def __bool__(self) -> bool:
return any(ref() is not None for ref in self._refs.values())
class Runtime(ABC):
"""
Abstract base class for workflow execution runtimes.
@@ -351,9 +405,13 @@ class Runtime(ABC):
Use registering() context manager for implicit workflow registration.
"""
def __init__(self) -> None:
self._pending: WorkflowSet = WorkflowSet()
self._launched: bool = False
_token: Token[Runtime | None]
def get_or_register(self, workflow: "Workflow") -> RegisteredWorkflow:
def get_or_register(self, workflow: Workflow) -> RegisteredWorkflow:
"""Get the registered workflow if available, otherwise register it."""
registered = self.get_registered(workflow)
if registered is None:
@@ -361,7 +419,7 @@ class Runtime(ABC):
return registered
@abstractmethod
def register(self, workflow: "Workflow") -> RegisteredWorkflow:
def register(self, workflow: Workflow) -> RegisteredWorkflow:
"""
Register a workflow with the runtime.
@@ -399,7 +457,7 @@ class Runtime(ABC):
...
@abstractmethod
def get_internal_adapter(self, workflow: "Workflow") -> InternalRunAdapter:
def get_internal_adapter(self, workflow: Workflow) -> InternalRunAdapter:
"""
Get the internal adapter for a workflow run.
@@ -429,7 +487,11 @@ class Runtime(ABC):
For many runtime's, this must be called before running workflows.
"""
pass
self._launched = True
for wf in self._pending:
self.register(wf)
wf._runtime_locked = True
self._pending = WorkflowSet()
def destroy(self) -> None:
"""
@@ -439,17 +501,19 @@ class Runtime(ABC):
"""
pass
def track_workflow(self, workflow: "Workflow") -> None:
def track_workflow(self, workflow: Workflow) -> None:
"""
Track a workflow instance for registration at launch time.
Called by Workflow.__init__ to register with the runtime.
Override in runtimes that need to track workflows.
Default implementation is a no-op.
"""
pass
self._pending.add(workflow)
def get_registered(self, workflow: "Workflow") -> RegisteredWorkflow | None:
def untrack_workflow(self, workflow: Workflow) -> None:
"""Remove a workflow from this runtime's tracking set."""
self._pending.discard(workflow)
def get_registered(self, workflow: Workflow) -> RegisteredWorkflow | None:
"""
Get the registered workflow if available.
@@ -462,8 +526,7 @@ class Runtime(ABC):
"""
Context manager for implicit workflow registration.
Workflows created inside this block will automatically register
with this runtime. Does NOT call launch() on exit.
Workflows created inside this block will automatically set this runtime as their runtime.
"""
token = _current_runtime.set(self)
try:
@@ -478,35 +541,19 @@ class SnapshottableAdapter(ABC):
This is a standalone mixin (not inheriting from InternalRunAdapter or
ExternalRunAdapter) that can be combined with adapter implementations
to add tick recording for debugging or replay purposes.
Adapters that implement this interface can record ticks as they occur
and replay them later. `on_tick` is called whenever a tick event is
received externally OR as a result from an internal command (e.g., a
step function completing, a timeout occurring, etc.)
to add init_state and replay capabilities for state reconstruction.
Use `as_snapshottable_adapter()` to check if an adapter supports snapshotting.
"""
@property
@abstractmethod
def init_state(self) -> "BrokerState":
def init_state(self) -> BrokerState:
"""
Get the initial state of the adapter.
"""
...
@abstractmethod
def on_tick(self, tick: WorkflowTick) -> None:
"""
Called whenever a tick event is received.
This method is invoked for both external ticks (sent via send_event)
and internal ticks (generated by step completions, timeouts, etc.).
Implementations should record the tick for later replay.
"""
...
@abstractmethod
def replay(self) -> list[WorkflowTick]:
"""
packages/llama-index-workflows/src/workflows/runtime/types/results.py
+63 -49
View File
@@ -7,19 +7,21 @@ import dataclasses
from contextvars import ContextVar
from dataclasses import dataclass
from typing import (
TYPE_CHECKING,
Any,
Generic,
Literal,
TypeVar,
Union,
)
from workflows.decorators import R
from pydantic import BaseModel, ConfigDict, model_serializer, model_validator
from workflows.events import Event
if TYPE_CHECKING:
pass
from workflows.runtime.types.serialization_helpers import (
SerializableEvent,
SerializableEventType,
SerializableException,
SerializableOptionalEvent,
)
EventType = TypeVar("EventType", bound=Event)
@@ -29,7 +31,7 @@ EventType = TypeVar("EventType", bound=Event)
@dataclass(frozen=True)
class StepWorkerContext(Generic[R]):
class StepWorkerContext:
"""
Base state passed to step functions and returned by step functions.
"""
@@ -37,7 +39,7 @@ class StepWorkerContext(Generic[R]):
# immutable state of the step events at start of the step function execution
state: StepWorkerState
# add commands here to mutate the internal worker state after step execution
returns: Returns[R]
returns: Returns
@dataclass(frozen=True)
@@ -79,13 +81,13 @@ class StepWorkerWaiter(Generic[EventType]):
@dataclass()
class Returns(Generic[R]):
class Returns:
"""
Mutate to add return values to the step function. These are only executed after the
step function has completed (including errors!)
"""
return_values: list[StepFunctionResult[R]]
return_values: list[StepFunctionResult]
class WaitingForEvent(Exception, Generic[EventType]):
@@ -110,71 +112,83 @@ StepWorkerStateContextVar = ContextVar[StepWorkerContext]("step_worker")
###################################
@dataclass
class StepWorkerResult(Generic[R]):
"""
Returned after a step function has been successfully executed.
"""
class StepWorkerResult(BaseModel):
"""Returned after a step function has been successfully executed."""
result: R
model_config = ConfigDict(frozen=True, arbitrary_types_allowed=True)
type: Literal["result"] = "result"
result: SerializableOptionalEvent = None
@dataclass
class StepWorkerFailed(Generic[R]):
"""
Returned after a step function has failed
"""
class StepWorkerFailed(BaseModel):
"""Returned after a step function has failed."""
exception: Exception
model_config = ConfigDict(frozen=True, arbitrary_types_allowed=True)
type: Literal["failed"] = "failed"
exception: SerializableException
failed_at: float
@dataclass
class DeleteWaiter:
"""
Returned after a waiter condition has been successfully resolved.
"""
class DeleteWaiter(BaseModel):
"""Returned after a waiter condition has been successfully resolved."""
model_config = ConfigDict(frozen=True)
type: Literal["delete_waiter"] = "delete_waiter"
waiter_id: str
@dataclass
class DeleteCollectedEvent:
"""
Returned after a collected event has been successfully resolved.
"""
class DeleteCollectedEvent(BaseModel):
"""Returned after a collected event has been successfully resolved."""
model_config = ConfigDict(frozen=True)
type: Literal["delete_collected"] = "delete_collected"
event_id: str
@dataclass
class AddCollectedEvent:
"""
Returned after a collected event has been added, and is not yet resolved.
"""
class AddCollectedEvent(BaseModel):
"""Returned after a collected event has been added, and is not yet resolved."""
model_config = ConfigDict(frozen=True, arbitrary_types_allowed=True)
type: Literal["add_collected"] = "add_collected"
event_id: str
event: Event
event: SerializableEvent
@dataclass
class AddWaiter(Generic[EventType]):
"""
Returned after a waiter has been added, and is not yet resolved.
"""
class AddWaiter(BaseModel, Generic[EventType]):
"""Returned after a waiter has been added, and is not yet resolved."""
model_config = ConfigDict(frozen=True, arbitrary_types_allowed=True)
type: Literal["add_waiter"] = "add_waiter"
waiter_id: str
waiter_event: Event | None
requirements: dict[str, Any]
timeout: float | None
event_type: type[EventType]
waiter_event: SerializableOptionalEvent = None
requirements: dict[str, Any] = {}
timeout: float | None = None
event_type: SerializableEventType
has_requirements: bool = False
@model_serializer(mode="wrap")
def _serialize(self, handler: Any) -> dict[str, Any]:
data = handler(self)
# Always serialize requirements as {} and record whether they existed
data["has_requirements"] = bool(self.requirements)
data["requirements"] = {}
return data
@model_validator(mode="wrap")
@classmethod
def _validate(cls, data: Any, handler: Any) -> AddWaiter:
if isinstance(data, dict):
# Strip has_requirements before validation (it's computed)
data = dict(data)
data.pop("has_requirements", None)
return handler(data)
# A step function result "command" communicates back to the workflow how the step function was resolved
# e.g. are we collecting events, waiting for an event, or just returning a result?
StepFunctionResult = Union[
StepWorkerResult[R],
StepWorkerFailed[R],
StepWorkerResult,
StepWorkerFailed,
AddCollectedEvent,
DeleteCollectedEvent,
AddWaiter[Event],
packages/llama-index-workflows/src/workflows/runtime/types/serialization_helpers.py
+100
View File
@@ -0,0 +1,100 @@
# SPDX-License-Identifier: MIT
# Copyright (c) 2026 LlamaIndex Inc.
"""Annotated types for Pydantic serialization of tricky tick/result fields.
Provides custom serializers/validators for:
- Events (polymorphic Pydantic models)
- Exceptions (not natively serializable)
- Event types (type[Event] as qualified name strings)
"""
from __future__ import annotations
from typing import Annotated, Any, Optional
from pydantic import PlainSerializer, PlainValidator
from workflows.context.serializers import JsonSerializer
from workflows.context.utils import (
import_module_from_qualified_name,
)
from workflows.events import Event
_json_serializer = JsonSerializer()
def _serialize_event(event: Event) -> Any:
return _json_serializer.serialize_value(event)
def _deserialize_event(data: Any) -> Event:
return _json_serializer.deserialize_value(data)
SerializableEvent = Annotated[
Event,
PlainSerializer(_serialize_event, return_type=Any),
PlainValidator(_deserialize_event),
]
def _serialize_optional_event(event: Event | None) -> Any:
if event is None:
return None
return _json_serializer.serialize_value(event)
def _deserialize_optional_event(data: Any) -> Event | None:
if data is None:
return None
return _json_serializer.deserialize_value(data)
SerializableOptionalEvent = Annotated[
Optional[Event],
PlainSerializer(_serialize_optional_event, return_type=Any),
PlainValidator(_deserialize_optional_event),
]
def _serialize_exception(exc: Exception) -> dict[str, Any]:
exc_type = type(exc)
qualified_name = f"{exc_type.__module__}.{exc_type.__qualname__}"
return {
"exception_type": qualified_name,
"exception_message": str(exc),
}
def _deserialize_exception(data: Any) -> Exception:
if isinstance(data, Exception):
return data
exc_message = data["exception_message"]
try:
exc_cls = import_module_from_qualified_name(data["exception_type"])
return exc_cls(exc_message)
except (ImportError, AttributeError, ValueError):
return Exception(exc_message)
SerializableException = Annotated[
Exception,
PlainSerializer(_serialize_exception, return_type=dict[str, Any]),
PlainValidator(_deserialize_exception),
]
def _serialize_event_type(event_type: type[Event]) -> str:
return f"{event_type.__module__}.{event_type.__qualname__}"
def _deserialize_event_type(data: Any) -> type[Event]:
if isinstance(data, type):
return data
return import_module_from_qualified_name(data)
SerializableEventType = Annotated[
type[Event],
PlainSerializer(_serialize_event_type, return_type=str),
PlainValidator(_deserialize_event_type),
]
packages/llama-index-workflows/src/workflows/runtime/types/step_function.py
+23 -17
View File
@@ -7,11 +7,11 @@ import asyncio
import functools
import time
from contextvars import copy_context
from typing import TYPE_CHECKING, Any, Awaitable, Callable, Generic, Protocol
from typing import TYPE_CHECKING, Any, Awaitable, Callable, Optional, Protocol, TypeVar
from llama_index_instrumentation.dispatcher import instrument_tags
from llama_index_instrumentation.span import active_span_id
from workflows.decorators import P, R, StepConfig
from workflows.decorators import P, StepConfig
from workflows.errors import WorkflowRuntimeError
from workflows.events import (
Event,
@@ -41,24 +41,26 @@ from workflows.workflow import Workflow
if TYPE_CHECKING:
from workflows.context.context import Context
StepReturnT = TypeVar("StepReturnT", bound=Optional[Event])
class StepWorkerFunction(Protocol, Generic[R]):
class StepWorkerFunction(Protocol):
def __call__(
self,
state: StepWorkerState,
step_name: str,
event: Event,
workflow: Workflow,
) -> Awaitable[list[StepFunctionResult[R]]]: ...
) -> Awaitable[list[StepFunctionResult]]: ...
async def partial(
func: Callable[..., R],
func: Callable[..., Any],
step_config: StepConfig,
event: Event,
context: Context,
workflow: Workflow,
) -> Callable[[], R]:
) -> Callable[[], Any]:
kwargs: dict[str, Any] = {}
kwargs[step_config.event_name] = event
if step_config.context_parameter:
@@ -76,14 +78,16 @@ async def partial(
def as_step_worker_functions(workflow: Workflow) -> dict[str, StepWorkerFunction]:
step_funcs = workflow._get_steps()
step_workers: dict[str, StepWorkerFunction[Any]] = {
step_workers: dict[str, StepWorkerFunction] = {
name: as_step_worker_function(getattr(func, "__func__", func))
for name, func in step_funcs.items()
}
return step_workers
def as_step_worker_function(func: Callable[P, Awaitable[R]]) -> StepWorkerFunction[R]:
def as_step_worker_function(
func: Callable[P, Awaitable[StepReturnT]],
) -> StepWorkerFunction:
"""
Wrap a step function, setting context variables and handling exceptions to instead
return the appropriate StepFunctionResult.
@@ -91,7 +95,7 @@ def as_step_worker_function(func: Callable[P, Awaitable[R]]) -> StepWorkerFuncti
# Keep original function reference for free-function steps; for methods we
# will resolve the currently-bound method from the provided workflow at call time.
original_func: Callable[..., Awaitable[R]] = func
original_func: Callable[..., Awaitable[StepReturnT]] = func
# Avoid functools.wraps here because it would set __wrapped__ to the bound
# method (when present), which would strongly reference the workflow
@@ -101,11 +105,11 @@ def as_step_worker_function(func: Callable[P, Awaitable[R]]) -> StepWorkerFuncti
step_name: str,
event: Event,
workflow: Workflow,
) -> list[StepFunctionResult[R]]:
) -> list[StepFunctionResult]:
from workflows.context.context import Context
internal_context = Context._create_internal(workflow=workflow)
returns = Returns[R](return_values=[])
returns = Returns(return_values=[])
token = StepWorkerStateContextVar.set(
StepWorkerContext(state=state, returns=returns)
@@ -135,15 +139,17 @@ def as_step_worker_function(func: Callable[P, Awaitable[R]]) -> StepWorkerFuncti
# run_in_executor doesn't accept **kwargs, so we need to use partial
copy = copy_context()
result: R = await asyncio.get_event_loop().run_in_executor(
None,
lambda: copy.run(partial_func), # type: ignore
result: StepReturnT = (
await asyncio.get_event_loop().run_in_executor(
None,
lambda: copy.run(partial_func), # type: ignore
)
)
else:
result = await partial_func()
if result is not None and not isinstance(result, Event):
msg = f"Step function {step_name} returned {type(result).__name__} instead of an Event instance."
raise WorkflowRuntimeError(msg)
if result is not None and not isinstance(result, Event):
msg = f"Step function {step_name} returned {type(result).__name__} instead of an Event instance."
raise WorkflowRuntimeError(msg)
returns.return_values.append(StepWorkerResult(result=result))
except WaitingForEvent as e:
await asyncio.sleep(0)
packages/llama-index-workflows/src/workflows/runtime/types/ticks.py
+27 -22
View File
@@ -1,6 +1,5 @@
# SPDX-License-Identifier: MIT
# Copyright (c) 2026 LlamaIndex Inc.
"""
Ticks (events) that drive the control loop.
@@ -16,55 +15,61 @@ events that can occur during workflow execution:
from __future__ import annotations
from dataclasses import dataclass
from typing import Generic, Union
from typing import Annotated, Literal, Union
from workflows.decorators import R
from workflows.events import Event
from pydantic import BaseModel, ConfigDict, Discriminator, TypeAdapter
from workflows.runtime.types.results import StepFunctionResult
from workflows.runtime.types.serialization_helpers import SerializableEvent
@dataclass(frozen=True)
class TickStepResult(Generic[R]):
class TickStepResult(BaseModel):
"""When processed, executes a step function and publishes the result"""
model_config = ConfigDict(frozen=True, arbitrary_types_allowed=True)
type: Literal["step_result"] = "step_result"
step_name: str
worker_id: int
event: Event
result: list[StepFunctionResult[R]]
event: SerializableEvent
result: list[Annotated[StepFunctionResult, Discriminator("type")]]
@dataclass(frozen=True)
class TickAddEvent:
class TickAddEvent(BaseModel):
"""When sent, adds an event to the workflow's event queue"""
event: Event
model_config = ConfigDict(frozen=True, arbitrary_types_allowed=True)
type: Literal["add_event"] = "add_event"
event: SerializableEvent
step_name: str | None = None
attempts: int | None = None
first_attempt_at: float | None = None
@dataclass(frozen=True)
class TickCancelRun:
class TickCancelRun(BaseModel):
"""When processed, cancels the workflow run"""
pass
model_config = ConfigDict(frozen=True)
type: Literal["cancel_run"] = "cancel_run"
@dataclass(frozen=True)
class TickPublishEvent:
class TickPublishEvent(BaseModel):
"""When sent, publishes an event to workflow consumers, e.g. a UI or a callback"""
event: Event
model_config = ConfigDict(frozen=True, arbitrary_types_allowed=True)
type: Literal["publish_event"] = "publish_event"
event: SerializableEvent
@dataclass(frozen=True)
class TickTimeout:
class TickTimeout(BaseModel):
"""When processed, times the workflow out, cancelling it"""
model_config = ConfigDict(frozen=True)
type: Literal["timeout"] = "timeout"
timeout: float
WorkflowTick = Union[
TickStepResult[R], TickAddEvent, TickCancelRun, TickPublishEvent, TickTimeout
WorkflowTick = Annotated[
Union[TickStepResult, TickAddEvent, TickCancelRun, TickPublishEvent, TickTimeout],
Discriminator("type"),
]
WorkflowTickAdapter: TypeAdapter[WorkflowTick] = TypeAdapter(WorkflowTick)
packages/llama-index-workflows/src/workflows/workflow.py
+24
View File
@@ -153,6 +153,7 @@ class Workflow(metaclass=WorkflowMeta):
self._resource_manager = resource_manager or ResourceManager()
# Instrumentation
self._dispatcher = dispatcher
self._runtime_locked = False
# Runtime registration: explicit > context-scoped > basic_runtime
from workflows.plugins._context import get_current_runtime
@@ -183,6 +184,18 @@ class Workflow(metaclass=WorkflowMeta):
"""The runtime this workflow is registered with."""
return self._runtime
def _switch_runtime(self, new_runtime: Runtime) -> None:
if new_runtime is self._runtime:
return
if self._runtime_locked:
raise RuntimeError(
"Cannot reassign runtime after workflow has been launched"
)
old = self._runtime
old.untrack_workflow(self)
self._runtime = new_runtime
new_runtime.track_workflow(self)
@property
def workflow_name(self) -> str:
"""
@@ -203,6 +216,13 @@ class Workflow(metaclass=WorkflowMeta):
cls = self.__class__
return f"{cls.__module__}.{cls.__qualname__}"
def _switch_workflow_name(self, name: str) -> None:
if self._runtime_locked and name != self._workflow_name:
raise RuntimeError(
"Cannot change workflow_name after workflow has been launched"
)
self._workflow_name = name
def _ensure_start_event_class(self) -> type[StartEvent]:
"""
Returns the StartEvent type used in this workflow.
@@ -414,6 +434,10 @@ class Workflow(metaclass=WorkflowMeta):
"""
from workflows.context import Context
if not self._runtime_locked:
# don't allow switching runtime after a workflow has been launched
self._runtime_locked = True
# Manually manage span to keep it open until workflow completes
# llama-index-instrumentation currently does not manage Awaitable's well (i.e. the workflow handler)
# this pattern is unusual enough to special case it here
packages/llama-index-workflows/tests/conftest.py
+7 -6
View File
@@ -6,10 +6,12 @@ from typing import Any, AsyncGenerator
import pytest
from pydantic import Field
from workflows.context import Context
from workflows.context.state_store import DictState, InMemoryStateStore
from workflows.context.external_context import ExternalContext
from workflows.decorators import step
from workflows.events import Event, StartEvent, StopEvent
from workflows.plugins.basic import AsyncioAdapterQueues, ExternalAsyncioAdapter
from workflows.plugins.basic import (
BasicRuntime,
)
from workflows.runtime.types.internal_state import BrokerState
from workflows.workflow import Workflow
@@ -52,16 +54,15 @@ def events() -> list:
@pytest.fixture()
async def ctx(workflow: Workflow) -> AsyncGenerator[Context[Any], None]:
from workflows.context.external_context import ExternalContext
runtime = BasicRuntime()
queues = AsyncioAdapterQueues(
_ = runtime._get_or_create_queues(
run_id="test-run",
init_state=BrokerState.from_workflow(workflow),
state_store=InMemoryStateStore(DictState()),
)
ctx = Context._create_external(
workflow=workflow,
external_adapter=ExternalAsyncioAdapter(queues=queues),
external_adapter=runtime.get_external_adapter("test-run"),
)
assert isinstance(ctx._face, ExternalContext)
try:
packages/llama-index-workflows/tests/runtime/conftest.py
+1 -1
View File
@@ -124,7 +124,7 @@ class MockRunAdapter(
workers={},
)
def on_tick(self, tick: WorkflowTick) -> None:
async def on_tick(self, tick: WorkflowTick) -> None:
"""Record a tick for replay."""
self._ticks.append(tick)
packages/llama-index-workflows/tests/runtime/test_control_loop_transformations.py
+12 -13
View File
@@ -1,6 +1,5 @@
# SPDX-License-Identifier: MIT
# Copyright (c) 2026 LlamaIndex Inc.
"""
Unit tests for control loop transformation functions.
@@ -10,7 +9,7 @@ testing them in isolation without running the full async control loop.
from __future__ import annotations
from typing import Any, cast
from typing import cast
import pytest
from workflows.decorators import StepConfig
@@ -253,7 +252,7 @@ def test_step_worker_failed_with_retry(base_state: BrokerState) -> None:
event = MyTestEvent(value=42)
add_worker(base_state, event)
tick: TickStepResult[Any] = TickStepResult(
tick: TickStepResult = TickStepResult(
step_name="test_step",
worker_id=0,
event=event,
@@ -278,7 +277,7 @@ def test_step_worker_failed_without_retry(base_state: BrokerState) -> None:
event = MyTestEvent(value=42)
add_worker(base_state, event)
tick: TickStepResult[Any] = TickStepResult(
tick: TickStepResult = TickStepResult(
step_name="test_step",
worker_id=0,
event=event,
@@ -297,7 +296,7 @@ def test_collected_events(base_state: BrokerState) -> None:
add_worker(base_state, event)
# Add event
tick: TickStepResult[Any] = TickStepResult(
tick: TickStepResult = TickStepResult(
step_name="test_step",
worker_id=0,
event=event,
@@ -334,12 +333,12 @@ def test_waiters(base_state: BrokerState) -> None:
event_type=OtherEvent,
)
# Add waiter
tick: TickStepResult[Any] = TickStepResult(
tick: TickStepResult = TickStepResult(
step_name="test_step",
worker_id=0,
event=event,
result=[
cast(StepFunctionResult[Any], result),
cast(StepFunctionResult, result),
],
)
new_state, _ = _process_step_result_tick(tick, base_state, now_seconds=110.0)
@@ -634,7 +633,7 @@ def test_step_worker_failed_retry_preserves_delay(base_state: BrokerState) -> No
event = MyTestEvent(value=42)
add_worker(base_state, event)
tick: TickStepResult[Any] = TickStepResult(
tick: TickStepResult = TickStepResult(
step_name="test_step",
worker_id=0,
event=event,
@@ -676,7 +675,7 @@ def test_step_worker_failed_retry_preserves_first_attempt_at(
)
)
tick: TickStepResult[Any] = TickStepResult(
tick: TickStepResult = TickStepResult(
step_name="test_step",
worker_id=0,
event=event,
@@ -761,11 +760,11 @@ def test_idle_event_emitted_on_transition_to_idle(base_state: BrokerState) -> No
event_type=OtherEvent,
)
tick: TickStepResult[Any] = TickStepResult[Any](
tick: TickStepResult = TickStepResult(
step_name="test_step",
worker_id=0,
event=event,
result=[cast(StepFunctionResult[Any], result)],
result=[cast(StepFunctionResult, result)],
)
new_state, commands = _process_step_result_tick(tick, base_state, now_seconds=110.0)
@@ -853,7 +852,7 @@ def test_no_idle_event_when_work_remains(base_state: BrokerState) -> None:
EventAttempt(event=MyTestEvent(value=99))
)
tick: TickStepResult[Any] = TickStepResult(
tick: TickStepResult = TickStepResult(
step_name="test_step",
worker_id=0,
event=event,
@@ -887,7 +886,7 @@ def test_no_idle_event_when_workflow_completes(base_state: BrokerState) -> None:
base_state.workers["test_step"].collected_waiters.append(waiter)
# Complete the workflow with StopEvent
tick: TickStepResult[Any] = TickStepResult(
tick: TickStepResult = TickStepResult(
step_name="test_step",
worker_id=0,
event=event,
packages/llama-index-workflows/tests/runtime/test_tick_serialization.py
+269
View File
@@ -0,0 +1,269 @@
# SPDX-License-Identifier: MIT
# Copyright (c) 2026 LlamaIndex Inc.
from __future__ import annotations
import json
import time
import pytest
from pydantic import TypeAdapter
from workflows.events import Event, StartEvent, StopEvent
from workflows.runtime.types.results import (
AddCollectedEvent,
AddWaiter,
DeleteCollectedEvent,
DeleteWaiter,
StepWorkerFailed,
StepWorkerResult,
)
from workflows.runtime.types.serialization_helpers import (
_deserialize_event,
_deserialize_event_type,
_deserialize_exception,
_serialize_event,
_serialize_event_type,
_serialize_exception,
)
from workflows.runtime.types.ticks import (
TickAddEvent,
TickCancelRun,
TickPublishEvent,
TickStepResult,
TickTimeout,
WorkflowTick,
)
class MyEvent(Event):
value: str = "hello"
# -- Serialization helper roundtrip tests --
def test_event_roundtrip() -> None:
event = MyEvent(value="world")
serialized = _serialize_event(event)
result = _deserialize_event(serialized)
assert isinstance(result, MyEvent)
assert result.value == "world"
def test_exception_roundtrip() -> None:
exc = ValueError("something went wrong")
serialized = _serialize_exception(exc)
result = _deserialize_exception(serialized)
assert isinstance(result, ValueError)
assert str(result) == "something went wrong"
def test_exception_roundtrip_unimportable() -> None:
CustomError = type("CustomError", (Exception,), {})
exc = CustomError("oops")
serialized = _serialize_exception(exc)
result = _deserialize_exception(serialized)
assert type(result) is Exception
assert str(result) == "oops"
def test_event_type_roundtrip() -> None:
serialized = _serialize_event_type(MyEvent)
result = _deserialize_event_type(serialized)
assert result is MyEvent
# -- Tick roundtrip tests --
@pytest.mark.parametrize(
"tick",
[
pytest.param(
TickAddEvent(
event=StartEvent(),
step_name="my_step",
attempts=3,
first_attempt_at=1234567890.0,
),
id="add_event",
),
pytest.param(
TickPublishEvent(event=MyEvent(value="world")),
id="publish_event",
),
pytest.param(
TickCancelRun(),
id="cancel_run",
),
pytest.param(
TickTimeout(timeout=30.5),
id="timeout",
),
pytest.param(
TickStepResult(
step_name="process",
worker_id=42,
event=MyEvent(value="trigger"),
result=[StepWorkerResult(result=StopEvent(result="done"))],
),
id="step_result_with_event",
),
pytest.param(
TickStepResult(
step_name="process",
worker_id=1,
event=StartEvent(),
result=[StepWorkerResult(result=None)],
),
id="step_result_with_none",
),
pytest.param(
TickStepResult(
step_name="collector",
worker_id=2,
event=StartEvent(),
result=[
AddCollectedEvent(
event_id="evt-1", event=MyEvent(value="collected")
)
],
),
id="step_result_add_collected_event",
),
pytest.param(
TickStepResult(
step_name="collector",
worker_id=3,
event=StartEvent(),
result=[DeleteCollectedEvent(event_id="evt-2")],
),
id="step_result_delete_collected_event",
),
pytest.param(
TickStepResult(
step_name="cleanup",
worker_id=5,
event=StartEvent(),
result=[DeleteWaiter(waiter_id="w-2")],
),
id="step_result_delete_waiter",
),
],
)
def test_tick_roundtrip(tick: WorkflowTick) -> None:
serialized = tick.model_dump(mode="json")
roundtripped = json.loads(json.dumps(serialized))
result = type(tick).model_validate(roundtripped)
assert result == tick
# -- Tick roundtrip tests with lossy serialization --
def test_tick_step_result_with_failed_value_error() -> None:
failed_at = time.time()
tick = TickStepResult(
step_name="broken_step",
worker_id=7,
event=StartEvent(),
result=[
StepWorkerFailed(
exception=ValueError("something went wrong"), failed_at=failed_at
)
],
)
serialized = tick.model_dump(mode="json")
roundtripped = json.loads(json.dumps(serialized))
result = TickStepResult.model_validate(roundtripped)
assert isinstance(result, TickStepResult)
r = result.result[0]
assert isinstance(r, StepWorkerFailed)
assert isinstance(r.exception, ValueError)
assert str(r.exception) == "something went wrong"
assert r.failed_at == failed_at
def test_tick_step_result_with_failed_unimportable_exception() -> None:
CustomError = type("CustomError", (Exception,), {})
failed_at = time.time()
tick = TickStepResult(
step_name="broken_step",
worker_id=8,
event=StartEvent(),
result=[StepWorkerFailed(exception=CustomError("oops"), failed_at=failed_at)],
)
serialized = tick.model_dump(mode="json")
roundtripped = json.loads(json.dumps(serialized))
result = TickStepResult.model_validate(roundtripped)
assert isinstance(result, TickStepResult)
r = result.result[0]
assert isinstance(r, StepWorkerFailed)
assert type(r.exception) is Exception
assert str(r.exception) == "oops"
assert r.failed_at == failed_at
def test_tick_step_result_with_add_waiter() -> None:
tick = TickStepResult(
step_name="waiter_step",
worker_id=4,
event=StartEvent(),
result=[
AddWaiter(
waiter_id="w-1",
waiter_event=MyEvent(value="waiting"),
requirements={"key": "value"},
timeout=60.0,
event_type=MyEvent,
)
],
)
serialized = tick.model_dump(mode="json")
# Verify the serialized form captures has_requirements correctly
waiter_data = serialized["result"][0]
assert waiter_data["has_requirements"] is True
assert waiter_data["requirements"] == {}
roundtripped = json.loads(json.dumps(serialized))
result = TickStepResult.model_validate(roundtripped)
assert isinstance(result, TickStepResult)
r = result.result[0]
assert isinstance(r, AddWaiter)
assert r.waiter_id == "w-1"
assert isinstance(r.waiter_event, MyEvent)
assert r.waiter_event.value == "waiting"
# Requirements are always {} after deserialization
assert r.requirements == {}
assert r.timeout == 60.0
assert r.event_type is MyEvent
# -- WorkflowTick discriminated union tests --
def test_workflow_tick_discriminated_union_roundtrip() -> None:
"""Verify that WorkflowTick TypeAdapter can roundtrip all tick types."""
adapter = TypeAdapter(WorkflowTick)
ticks = [
TickAddEvent(event=StartEvent(), step_name="s"),
TickPublishEvent(event=MyEvent(value="x")),
TickCancelRun(),
TickTimeout(timeout=10.0),
TickStepResult(
step_name="s",
worker_id=0,
event=StartEvent(),
result=[StepWorkerResult(result=None)],
),
]
for tick in ticks:
dumped = adapter.dump_python(tick, mode="json")
roundtripped = json.loads(json.dumps(dumped))
restored = adapter.validate_python(roundtripped)
assert type(restored) is type(tick)
packages/llama-index-workflows/tests/runtime/test_workflow_set_and_tracking.py
+248
View File
@@ -0,0 +1,248 @@
# SPDX-License-Identifier: MIT
# Copyright (c) 2026 LlamaIndex Inc.
"""Tests for WorkflowSet, runtime tracking, and workflow mutation."""
from __future__ import annotations
import gc
from typing import Any
import pytest
from workflows import Workflow, step
from workflows.events import StartEvent, StopEvent
from workflows.plugins import BasicRuntime
from workflows.runtime.types.plugin import WorkflowSet
class SimpleWorkflow(Workflow):
@step
async def start(self, ev: StartEvent) -> StopEvent:
return StopEvent(result="done")
class UnhashableWorkflow(Workflow):
def __init__(self, **kwargs: Any) -> None:
super().__init__(**kwargs)
self.data = [1, 2, 3] # Makes it unhashable
@step
async def start(self, ev: StartEvent) -> StopEvent:
return StopEvent(result="done")
@pytest.fixture
def basic_runtime() -> BasicRuntime:
return BasicRuntime()
@pytest.fixture
def workflow_set() -> WorkflowSet:
return WorkflowSet()
# ---------------------------------------------------------------------------
# WorkflowSet tests
# ---------------------------------------------------------------------------
def test_workflow_set_add_and_contains(
workflow_set: WorkflowSet, basic_runtime: BasicRuntime
) -> None:
wf = SimpleWorkflow(runtime=basic_runtime)
workflow_set.add(wf)
assert wf in workflow_set
def test_workflow_set_add_unhashable_workflow(
workflow_set: WorkflowSet, basic_runtime: BasicRuntime
) -> None:
wf = UnhashableWorkflow(runtime=basic_runtime)
workflow_set.add(wf)
assert wf in workflow_set
items = list(workflow_set)
assert wf in items
def test_workflow_set_discard(
workflow_set: WorkflowSet, basic_runtime: BasicRuntime
) -> None:
wf = SimpleWorkflow(runtime=basic_runtime)
workflow_set.add(wf)
assert wf in workflow_set
workflow_set.discard(wf)
assert wf not in workflow_set
def test_workflow_set_len_and_bool(
workflow_set: WorkflowSet, basic_runtime: BasicRuntime
) -> None:
assert not workflow_set
assert len(workflow_set) == 0
wf = SimpleWorkflow(runtime=basic_runtime)
workflow_set.add(wf)
assert workflow_set
assert len(workflow_set) == 1
def test_workflow_set_iter(
workflow_set: WorkflowSet, basic_runtime: BasicRuntime
) -> None:
wf1 = SimpleWorkflow(runtime=basic_runtime)
wf2 = SimpleWorkflow(runtime=basic_runtime)
wf3 = SimpleWorkflow(runtime=basic_runtime)
workflow_set.add(wf1)
workflow_set.add(wf2)
workflow_set.add(wf3)
items = set(id(w) for w in workflow_set)
assert items == {id(wf1), id(wf2), id(wf3)}
def test_workflow_set_gc_cleanup(
workflow_set: WorkflowSet, basic_runtime: BasicRuntime
) -> None:
wf = SimpleWorkflow(runtime=basic_runtime)
workflow_set.add(wf)
assert len(workflow_set) == 1
del wf
gc.collect()
assert len(workflow_set) == 0
def test_workflow_set_add_idempotent(
workflow_set: WorkflowSet, basic_runtime: BasicRuntime
) -> None:
wf = SimpleWorkflow(runtime=basic_runtime)
workflow_set.add(wf)
workflow_set.add(wf)
assert len(workflow_set) == 1
# ---------------------------------------------------------------------------
# Runtime tracking tests
# ---------------------------------------------------------------------------
def test_track_workflow_adds_to_set(basic_runtime: BasicRuntime) -> None:
wf = SimpleWorkflow(runtime=basic_runtime)
assert wf in basic_runtime._pending
def test_untrack_workflow_removes_from_set(basic_runtime: BasicRuntime) -> None:
wf = SimpleWorkflow(runtime=basic_runtime)
assert wf in basic_runtime._pending
basic_runtime.untrack_workflow(wf)
assert wf not in basic_runtime._pending
def test_launch_locks_tracked_workflows(basic_runtime: BasicRuntime) -> None:
wf1 = SimpleWorkflow(runtime=basic_runtime)
wf2 = SimpleWorkflow(runtime=basic_runtime)
basic_runtime.launch()
assert wf1._runtime_locked is True
assert wf2._runtime_locked is True
def test_relaunch_locks_new_workflows(basic_runtime: BasicRuntime) -> None:
wf1 = SimpleWorkflow(runtime=basic_runtime)
basic_runtime.launch()
assert wf1._runtime_locked is True
wf2 = SimpleWorkflow(runtime=basic_runtime)
assert wf2._runtime_locked is False
basic_runtime.launch()
assert wf1._runtime_locked is True
assert wf2._runtime_locked is True
def test_weak_reference_cleanup(basic_runtime: BasicRuntime) -> None:
wf = SimpleWorkflow(runtime=basic_runtime)
assert len(basic_runtime._pending) == 1
del wf
gc.collect()
assert len(basic_runtime._pending) == 0
def test_basic_runtime_launch_sets_launched_flag(basic_runtime: BasicRuntime) -> None:
assert basic_runtime._launched is False
basic_runtime.launch()
assert basic_runtime._launched is True
# ---------------------------------------------------------------------------
# Workflow mutation tests
# ---------------------------------------------------------------------------
def test_workflow_name_setter(basic_runtime: BasicRuntime) -> None:
wf = SimpleWorkflow(runtime=basic_runtime)
wf._switch_workflow_name("custom-name")
assert wf.workflow_name == "custom-name"
def test_workflow_name_setter_raises_after_launch() -> None:
rt = BasicRuntime()
wf = SimpleWorkflow(runtime=rt)
wf._switch_workflow_name("before-launch")
assert wf.workflow_name == "before-launch"
rt.launch()
with pytest.raises(RuntimeError, match="Cannot change workflow_name"):
wf._switch_workflow_name("after-launch")
def test_runtime_setter_swaps_tracking() -> None:
rt1 = BasicRuntime()
rt2 = BasicRuntime()
wf = SimpleWorkflow(runtime=rt1)
assert wf in rt1._pending
assert wf not in rt2._pending
wf._switch_runtime(rt2)
assert wf not in rt1._pending
assert wf in rt2._pending
def test_runtime_setter_post_launch_raises() -> None:
rt1 = BasicRuntime()
rt2 = BasicRuntime()
wf = SimpleWorkflow(runtime=rt1)
rt1.launch()
with pytest.raises(RuntimeError, match="Cannot reassign runtime"):
wf._switch_runtime(rt2)
def test_runtime_setter_same_runtime_after_launch_is_noop() -> None:
rt = BasicRuntime()
wf = SimpleWorkflow(runtime=rt)
rt.launch()
# Assigning the same runtime should not raise
wf._switch_runtime(rt)
assert wf.runtime is rt
@pytest.mark.asyncio
async def test_run_locks_runtime() -> None:
rt1 = BasicRuntime()
rt2 = BasicRuntime()
wf = SimpleWorkflow(runtime=rt1)
assert wf._runtime_locked is False
handler = wf.run()
assert wf._runtime_locked is True
with pytest.raises(RuntimeError, match="Cannot reassign runtime"):
wf._switch_runtime(rt2)
await handler
def test_runtime_setter_before_launch_then_launch_locks() -> None:
rt1 = BasicRuntime()
rt2 = BasicRuntime()
wf = SimpleWorkflow(runtime=rt1)
wf._switch_runtime(rt2)
assert wf._runtime_locked is False
rt2.launch()
assert wf._runtime_locked is True