Task Consumer

The task consumer is a TaskManager which is also a Workers that consumes tasks from the task queue and executes them. It can be imported from fluid.scheduler:

from fluid.scheduler import TaskConsumer

fluid.scheduler.TaskConsumer

TaskConsumer(
    *,
    deps=None,
    config=None,
    name="",
    stopping_grace_period=None,
    **kwargs
)

Bases: TaskManager, Workers

The Task Consumer is a Task Manager responsible for consuming tasks from a task queue

Create the task consumer and the workers it runs.

Workers added via add_workers are started by startup, therefore they only run in a process which consumes the task queue.

The async dispatcher is added as an async context manager instead, so it is started by __aenter__ and runs in every process which uses the task manager. A cpu bound task runs in a process of its own, which executes a single task and never starts the consumer workers, and it still needs to dispatch the lifecycle events its plugins and handlers subscribe to.

PARAMETER	DESCRIPTION
deps	Application dependencies available to every task run. See the Task Dependencies tutorial. TYPE: Any DEFAULT: None
config	Task manager configuration. Built from the extra keyword arguments when not provided. TYPE: TaskManagerConfig | None DEFAULT: None
name	Worker's name, if not provided it is evaluated from the class name TYPE: str DEFAULT: ''
stopping_grace_period	Grace period in seconds to wait for workers to stop running when this worker is shutdown. It defaults to the FLUID_STOPPING_GRACE_PERIOD environment variable or 10 seconds. TYPE: float | None DEFAULT: None
**kwargs	Configuration fields, used when config is not provided. TYPE: Any DEFAULT: {}

Source code in fluid/scheduler/consumer.py

def __init__(
    self,
    *,
    deps: Annotated[
        Any,
        Doc("""
            Application dependencies available to every task run.

            See the [Task Dependencies](../tutorials/task_deps.md)
            tutorial.
            """),
    ] = None,
    config: Annotated[
        TaskManagerConfig | None,
        Doc("""
            Task manager configuration.

            Built from the extra keyword arguments when not provided.
            """),
    ] = None,
    name: Annotated[
        str,
        Doc("Worker's name, if not provided it is evaluated from the class name"),
    ] = "",
    stopping_grace_period: Annotated[
        float | None,
        Doc(
            "Grace period in seconds to wait for workers to stop running "
            "when this worker is shutdown. "
            "It defaults to the `FLUID_STOPPING_GRACE_PERIOD` "
            "environment variable or 10 seconds."
        ),
    ] = None,
    **kwargs: Annotated[
        Any,
        Doc("Configuration fields, used when `config` is not provided."),
    ],
) -> None:
    """Create the task consumer and the workers it runs.

    Workers added via `add_workers` are started by `startup`, therefore
    they only run in a process which consumes the task queue.

    The async dispatcher is added as an async context manager instead, so
    it is started by `__aenter__` and runs in every process which uses the
    task manager. A cpu bound task runs in a process of its own, which
    executes a single task and never starts the consumer workers, and it
    still needs to dispatch the lifecycle events its plugins and handlers
    subscribe to.
    """
    super().__init__(deps=deps, config=config, **kwargs)
    if stopping_grace_period is None:
        stopping_grace_period = settings.STOPPING_GRACE_PERIOD
    Workers.__init__(
        self, name=name, stopping_grace_period=2 * stopping_grace_period
    )
    self._async_dispatcher_worker = AsyncConsumer(
        AsyncTaskDispatcher(), stopping_grace_period=stopping_grace_period
    )
    self._in_process_queue = InProcessTaskQueue(
        self, stopping_grace_period=stopping_grace_period
    )
    self.add_async_context_manager(self._async_dispatcher_worker)
    self.add_workers(self._in_process_queue)
    for i in range(self.config.max_concurrent_tasks):
        worker_name = f"task-worker-{i+1}"
        self.add_workers(
            WorkerFunction(
                partial(self._consume_tasks, worker_name),
                name=worker_name,
                stopping_grace_period=stopping_grace_period,
            )
        )
    self.add_workers(
        WorkerFunction(
            self._ping_status,
            heartbeat=1.0,
            name="manager-status",
            stopping_grace_period=stopping_grace_period,
        )
    )

worker_state property

worker_state

The running state of the worker

worker_name property

worker_name

The name of the worker

num_workers property

num_workers

deps instance-attribute

deps = deps if deps is not None else State()

Dependencies for the task manager.

Production applications requires global dependencies to be available to all tasks. This can be achieved by setting the deps attribute of the task manager to an object with the required dependencies.

Each task can cast the dependencies to the required type.

state instance-attribute

state = State()

State for the task manager. This can be used by plugins to store state in the task manager.

config instance-attribute

config = config or TaskManagerConfig(**kwargs)

Task manager configuration

dispatcher instance-attribute

dispatcher = TaskDispatcher()

A dispatcher of TaskRun events.

Application can register handlers to listen for events happening during the lifecycle of a task run.

broker instance-attribute

broker = TaskBroker.from_url(self.config.broker_url)

manager_id instance-attribute

manager_id = self.broker.new_uuid()

registry property

registry

The task registry

type property

type

The type of the task manager

has_started

has_started()

Source code in fluid/utils/worker.py

def has_started(self) -> bool:
    return self._worker_state != WorkerState.INIT

is_running

is_running()

Source code in fluid/utils/worker.py

def is_running(self) -> bool:
    return self._worker_state == WorkerState.RUNNING

is_stopping

is_stopping()

Source code in fluid/utils/worker.py

def is_stopping(self) -> bool:
    return self._worker_state == WorkerState.STOPPING

is_stopped

is_stopped()

Source code in fluid/utils/worker.py

def is_stopped(self) -> bool:
    return self._worker_state in (WorkerState.STOPPED, WorkerState.FORCE_STOPPED)

gracefully_stop

gracefully_stop()

Try to gracefully stop the workers and this worker

Source code in fluid/utils/worker.py

def gracefully_stop(self) -> None:
    """Try to gracefully stop the workers and this worker"""
    super().gracefully_stop()
    for worker in self._workers:
        worker.gracefully_stop()

after_shutdown

after_shutdown(reason, code)

Called after shutdown of worker

By default it does nothing, but can be overriden to do something such as exit the process.

Source code in fluid/utils/worker.py

def after_shutdown(self, reason: str, code: int) -> None:  # noqa: B027
    """Called after shutdown of worker

    By default it does nothing, but can be overriden to do something such as
    exit the process.
    """

status async

status()

Source code in fluid/utils/worker.py

async def status(self) -> dict:
    status_workers = await asyncio.gather(
        *[worker.status() for worker in self._workers],
    )
    return {
        worker.worker_name: status
        for worker, status in zip(self._workers, status_workers, strict=False)
    }

on_startup async

on_startup()

Source code in fluid/scheduler/consumer.py

async def on_startup(self) -> None:
    await self.__aenter__()

on_shutdown async

on_shutdown()

Source code in fluid/scheduler/consumer.py

async def on_shutdown(self) -> None:
    await self.__aexit__(None, None, None)

shutdown async

shutdown()

Shutdown a running worker and wait for it to stop

This method will try to gracefully stop the worker and wait for it to stop. If the worker does not stop in the grace period, it will force shutdown by cancelling the task.

Source code in fluid/utils/worker.py

async def shutdown(self) -> None:
    """Shutdown a running worker and wait for it to stop

    This method will try to gracefully stop the worker and wait for it to stop.
    If the worker does not stop in the grace period, it will force shutdown
    by cancelling the task.
    """
    if self._worker_task_runner is not None:
        await self._worker_task_runner.shutdown()

wait_for_shutdown async

wait_for_shutdown()

Wait for the worker to stop

This method will wait for the worker to stop running, but doesn't try to gracefully stop it nor force shutdown.

Source code in fluid/utils/worker.py

async def wait_for_shutdown(self) -> None:
    """Wait for the worker to stop

    This method will wait for the worker to stop running, but doesn't
    try to gracefully stop it nor force shutdown.
    """
    if self._worker_task_runner is not None:
        await self._worker_task_runner.wait_for_shutdown()

workers

workers()

Source code in fluid/utils/worker.py

def workers(self) -> Iterator[Worker]:
    return iter(self._workers)

run async

run()

Source code in fluid/utils/worker.py

async def run(self) -> None:
    while self.is_running():
        for worker in self._workers:
            if not worker.has_started():
                await worker.startup()
            if not worker.is_running():
                self.gracefully_stop()
                break
        await asyncio.sleep(self._heartbeat)
    await self._wait_for_workers()

add_workers

add_workers(*workers)

add workers to the workers

They can be added while the worker is running.

Source code in fluid/utils/worker.py

def add_workers(self, *workers: Worker) -> None:
    """add workers to the workers

    They can be added while the worker is running.
    """
    for worker in workers:
        if worker not in self._workers:
            self._workers.append(worker)

add_async_context_manager

add_async_context_manager(cm)

Add an async context manager to the task manager

These context managers are entered when the task manager starts

Source code in fluid/scheduler/consumer.py

def add_async_context_manager(self, cm: Any) -> None:
    """Add an async context manager to the task manager

    These context managers are entered when the task manager starts
    """
    self._async_contexts.append(cm)

register_task

register_task(task, tags=None)

Register a task with the task manager

PARAMETER	DESCRIPTION
task	Task to register TYPE: Task
tags	Extra tags to add to the task before registering it TYPE: Sequence[str] | None DEFAULT: None

Source code in fluid/scheduler/consumer.py

def register_task(
    self,
    task: Annotated[Task, Doc("Task to register")],
    tags: Annotated[
        Sequence[str] | None,
        Doc("Extra tags to add to the task before registering it"),
    ] = None,
) -> None:
    """Register a task with the task manager"""
    if tags:
        task = task._replace(tags=task.tags | frozenset(tags))
    self.broker.register_task(task)

execute async

execute(task, *, run_id='', priority=None, **params)

Execute a task and wait for it to finish

This method is an async method that should be used in an asynchronous context when one need to wait for the task to finish execution.

PARAMETER	DESCRIPTION
task	The task or task name, if a task name it must be registered with the task manager. TYPE: str | Task
run_id	Unique ID for the task run. If not provided a new UUID is generated. TYPE: str DEFAULT: ''
priority	Override the default task priority if provided TYPE: TaskPriority | None DEFAULT: None
**params	The optional parameters for the task run. They must match the task params model TYPE: Any DEFAULT: {}

Source code in fluid/scheduler/consumer.py

async def execute(
    self,
    task: Annotated[
        str | Task,
        Doc(
            "The task or task name,"
            " if a task name it must be registered with the task manager."
        ),
    ],
    *,
    run_id: Annotated[
        str,
        Doc("Unique ID for the task run. If not provided a new UUID is generated."),
    ] = "",
    priority: Annotated[
        TaskPriority | None, Doc("Override the default task priority if provided")
    ] = None,
    **params: Annotated[
        Any,
        Doc(
            "The optional parameters for the task run. "
            "They must match the task params model"
        ),
    ],
) -> TaskRun:
    """Execute a task and wait for it to finish

    This method is an async method that should be used in an asynchronous
    context when one need to wait for the task to finish execution.
    """
    task_run = self.create_task_run(
        task,
        run_id=run_id,
        priority=priority,
        **params,
    )
    try:
        await task_run._execute()
    except TaskAbortedError as exc:
        await self.broker.set_task_aborted(task_run.id, str(exc))
    return task_run

execute_sync

execute_sync(task, *, run_id='', priority=None, **params)

Execute a task synchronously

This method is a blocking method that should be used in a synchronous context.

PARAMETER	DESCRIPTION
task	The task or task name, if a task name it must be registered with the task manager. TYPE: str | Task
run_id	Unique ID for the task run. If not provided a new UUID is generated. TYPE: str DEFAULT: ''
priority	Override the default task priority if provided TYPE: TaskPriority | None DEFAULT: None
**params	The optional parameters for the task run. They must match the task params model TYPE: Any DEFAULT: {}

Source code in fluid/scheduler/consumer.py

def execute_sync(
    self,
    task: Annotated[
        str | Task,
        Doc(
            "The task or task name,"
            " if a task name it must be registered with the task manager."
        ),
    ],
    *,
    run_id: Annotated[
        str,
        Doc("Unique ID for the task run. If not provided a new UUID is generated."),
    ] = "",
    priority: Annotated[
        TaskPriority | None, Doc("Override the default task priority if provided")
    ] = None,
    **params: Annotated[
        Any,
        Doc(
            "The optional parameters for the task run. "
            "They must match the task params model"
        ),
    ],
) -> TaskRun:
    """Execute a task synchronously

    This method is a blocking method that should be used in a synchronous
    context.
    """
    return asyncio.run(
        self._execute_and_exit(
            task,
            run_id=run_id,
            priority=priority,
            **params,
        )
    )

queue async

queue(
    task,
    *,
    run_id="",
    priority=None,
    from_task_run=None,
    **params
)

Queue a task for execution

This methods fires two events:

• init: when the task run is created
• queued: after the task is queued

It returns the TaskRun object

PARAMETER	DESCRIPTION
task	The task or task name, if a task name it must be registered with the task manager. TYPE: str | Task
run_id	Unique ID for the task run. If not provided a new UUID is generated. TYPE: str DEFAULT: ''
priority	Override the default task priority if provided TYPE: TaskPriority | None DEFAULT: None
from_task_run	The task run queueing this one, if any. Prefer TaskRun.queue, which passes it for you. TYPE: TaskRun | None DEFAULT: None
**params	The optional parameters for the task run. They must match the task params model TYPE: Any DEFAULT: {}

Source code in fluid/scheduler/consumer.py

async def queue(
    self,
    task: Annotated[
        str | Task,
        Doc(
            "The task or task name,"
            " if a task name it must be registered with the task manager."
        ),
    ],
    *,
    run_id: Annotated[
        str,
        Doc("Unique ID for the task run. If not provided a new UUID is generated."),
    ] = "",
    priority: Annotated[
        TaskPriority | None, Doc("Override the default task priority if provided")
    ] = None,
    from_task_run: Annotated[
        TaskRun | None,
        Doc(
            "The task run queueing this one, if any. "
            "Prefer [TaskRun.queue][fluid.scheduler.TaskRun.queue], "
            "which passes it for you."
        ),
    ] = None,
    **params: Annotated[
        Any,
        Doc(
            "The optional parameters for the task run. "
            "They must match the task params model"
        ),
    ],
) -> TaskRun:
    """Queue a task for execution

    This methods fires two events:

    - `init`: when the task run is created
    - `queued`: after the task is queued

    It returns the [TaskRun][fluid.scheduler.TaskRun] object
    """
    task_run = self.create_task_run(
        task,
        run_id=run_id,
        priority=priority,
        from_task_run=from_task_run,
        **params,
    )
    return await self._queue_task_run(task_run)

create_task_run

create_task_run(
    task,
    *,
    run_id="",
    priority=None,
    from_task_run=None,
    **params
)

Create a TaskRun in init state

PARAMETER	DESCRIPTION
task	The task or task name, if a task name it must be registered with the task manager. TYPE: str | Task
run_id	Unique ID for the task run. If not provided a new UUID is generated. TYPE: str DEFAULT: ''
priority	Override the default task priority if provided TYPE: TaskPriority | None DEFAULT: None
from_task_run	The task run creating this one, if any. It records the chain. TYPE: TaskRun | None DEFAULT: None
**params	The optional parameters for the task run. They must match the task params model TYPE: Any DEFAULT: {}

Source code in fluid/scheduler/consumer.py

def create_task_run(
    self,
    task: Annotated[
        str | Task,
        Doc(
            "The task or task name,"
            " if a task name it must be registered with the task manager."
        ),
    ],
    *,
    run_id: Annotated[
        str,
        Doc("Unique ID for the task run. If not provided a new UUID is generated."),
    ] = "",
    priority: Annotated[
        TaskPriority | None, Doc("Override the default task priority if provided")
    ] = None,
    from_task_run: Annotated[
        TaskRun | None,
        Doc("The task run creating this one, if any. It records the chain."),
    ] = None,
    **params: Annotated[
        Any,
        Doc(
            "The optional parameters for the task run. "
            "They must match the task params model"
        ),
    ],
) -> TaskRun:
    """Create a [TaskRun][fluid.scheduler.TaskRun] in `init` state"""
    task = self.broker.task_from_registry(task)
    run_id = run_id or self.broker.new_uuid()
    return TaskRun(
        id=run_id,
        task=task,
        priority=priority or task.priority,
        params=task.params_model(**params),
        task_manager=self,
        from_run_id=from_task_run.id if from_task_run else "",
        # the first run in a chain has no root of its own, it is the root
        root_run_id=(
            (from_task_run.root_run_id or from_task_run.id) if from_task_run else ""
        ),
    )

register_from_module

register_from_module(module, tags=None)

Register tasks from a python module

PARAMETER	DESCRIPTION
module	Python module with tasks implementations - can contain any object, only instances of Task are registered TYPE: ModuleType
tags	Extra tags to add to every registered task TYPE: Sequence[str] | None DEFAULT: None

Source code in fluid/scheduler/consumer.py

def register_from_module(
    self,
    module: Annotated[
        ModuleType,
        Doc(
            "Python module with tasks implementations "
            "- can contain any object, only instances of Task are registered"
        ),
    ],
    tags: Annotated[
        Sequence[str] | None,
        Doc("Extra tags to add to every registered task"),
    ] = None,
) -> None:
    """Register tasks from a python module"""
    for name in dir(module):
        if name.startswith("_"):
            continue
        if isinstance(obj := getattr(module, name), Task):
            self.register_task(obj, tags=tags)

register_from_dict

register_from_dict(data, tags=None)

Register tasks from a python dictionary

PARAMETER	DESCRIPTION
data	Python dictionary with tasks implementations - can contain any object, only instances of Task are registered TYPE: dict[str, Any]
tags	Extra tags to add to every registered task TYPE: Sequence[str] | None DEFAULT: None

Source code in fluid/scheduler/consumer.py

def register_from_dict(
    self,
    data: Annotated[
        dict[str, Any],
        Doc(
            "Python dictionary with tasks implementations "
            "- can contain any object, only instances of Task are registered"
        ),
    ],
    tags: Annotated[
        Sequence[str] | None,
        Doc("Extra tags to add to every registered task"),
    ] = None,
) -> None:
    """Register tasks from a python dictionary"""
    for name, obj in data.items():
        if name.startswith("_"):
            continue
        if isinstance(obj, Task):
            self.register_task(obj, tags=tags)

with_plugin

with_plugin(plugin)

Register a plugin with the task manager

PARAMETER	DESCRIPTION
plugin	The plugin to register TYPE: TaskManagerPlugin

Source code in fluid/scheduler/consumer.py

def with_plugin(
    self,
    plugin: Annotated[TaskManagerPlugin, Doc("The plugin to register")],
) -> Self:
    """Register a plugin with the task manager"""
    self._plugins.append(plugin)
    plugin.register(self)
    return self

startup async

startup()

Start the task consumer workers.

A cpu bound process executes a single task and exits, it never consumes the queue. Reaching this point means the entry point ignored the exec command, so it cannot run cpu bound tasks.

Source code in fluid/scheduler/consumer.py

async def startup(self) -> None:
    """Start the task consumer workers.

    A cpu bound process executes a single task and exits, it never consumes
    the queue. Reaching this point means the entry point ignored the `exec`
    command, so it cannot run cpu bound tasks.
    """
    if is_in_cpu_process():
        raise CpuBoundEntryPointError(
            "a task consumer cannot start in a cpu bound process: "
            "running cpu bound tasks requires the application entry point "
            "to be a TaskManagerCLI"
        )
    await super().startup()

sync_queue

sync_queue(task, delay=0)

Queue a task synchronously

Source code in fluid/scheduler/consumer.py

def sync_queue(self, task: str | Task | TaskRun, delay: float = 0) -> None:
    """Queue a task synchronously"""
    self._in_process_queue.queue(task, delay=delay)

queue_and_wait async

queue_and_wait(task, *, timeout=None, **params)

Queue a task and wait for it to finish

PARAMETER	DESCRIPTION
task	The task or task name, if a task name it must be registered with the task manager. TYPE: str | Task
timeout	Timeout for waiting the task to finish TYPE: int | None DEFAULT: None
**params	The optional parameters for the task run. They must match the task params model TYPE: Any DEFAULT: {}

Source code in fluid/scheduler/consumer.py

async def queue_and_wait(
    self,
    task: Annotated[
        str | Task,
        Doc(
            "The task or task name,"
            " if a task name it must be registered with the task manager."
        ),
    ],
    *,
    timeout: Annotated[
        int | None, Doc("Timeout for waiting the task to finish")
    ] = None,
    **params: Annotated[
        Any,
        Doc(
            "The optional parameters for the task run. "
            "They must match the task params model"
        ),
    ],
) -> TaskRun:
    """Queue a task and wait for it to finish"""
    with TaskRunWaiter(self) as waiter:
        task_run = await self.queue(task, **params)
        return await waiter.wait(task_run, timeout=timeout)

register_async_handler

register_async_handler(event, handler)

Source code in fluid/scheduler/consumer.py

def register_async_handler(self, event: Event | str, handler: AsyncHandler) -> None:
    event = Event.from_string_or_event(event)
    self.dispatcher.register_handler(
        f"{event.type}.async_dispatch",
        self._async_dispatcher_worker.send,
    )
    self._async_dispatcher_worker.dispatcher.register_handler(event, handler)

unregister_async_handler

unregister_async_handler(event)

Source code in fluid/scheduler/consumer.py

def unregister_async_handler(self, event: Event | str) -> AsyncHandler | None:
    return self._async_dispatcher_worker.dispatcher.unregister_handler(event)