How to use streaq in dynamic environments (not in import-time)?

[espdev]

Hello,

streaq is a very interesting library. I'm thinking about using it in our project. Unfortunately, I don't quite understand how you can use the library without creating a module-level worker.

The documentation says:

from contextlib import asynccontextmanager
from dataclasses import dataclass
from typing import AsyncIterator
from httpx import AsyncClient
from streaq import Worker

@dataclass
class WorkerContext:
    """
    Type safe way of defining the dependencies of your tasks.
    e.g. HTTP client, database connection, settings.
    """
    http_client: AsyncClient

@asynccontextmanager
async def lifespan(worker: Worker[WorkerContext]) -> AsyncIterator[WorkerContext]:
    """
    Here, we initialize the worker's dependencies.
    You can also do any startup/shutdown work here
    """
    async with AsyncClient() as http_client:
        yield Context(http_client)

worker = Worker(redis_url="redis://localhost:6379", lifespan=lifespan)

...
@worker.task(timeout=5)
async def fetch(url: str) -> int:
    # worker.context here is of type WorkerContext
    res = await worker.context.http_client.get(url)
    return len(res.text)

...
async with worker:
    await fetch.enqueue("https://tastyware.dev/")
    # enqueue returns a task object that can be used to get results/info
    task = await fetch.enqueue("https://github.com/tastyware/streaq").start(delay=3)
    print(await task.info())
    print(await task.result(timeout=5))

And running the worker:

streaq script.worker

Let's imagine, I would like to initialize a worker not in import-time but in run-time (using a worker_factory func), for example, something like this:

from my_app.tasks import foo, bar

@dataclass
class AppContext:
    config: Config = field(default_factory=Config)
    # more context data here

    worker: Worker = field(init=False)
    
    @asynccontextmanager
    async def lifespan(self, worker) -> Self:
        yield self
    
    def __post_init__(self):
        self.worker = Worker(redis_url=self.config.redis_url, lifespan=self.lifespan)
        
        # for example, register tasks dynamically here or anywhere else
        self.worker.register_task(foo)
        self.worker.register_task(bar)

def worker_factory() -> Worker:
    app_context = AppContext()
    return app_context.worker

And run the worker:

streaq --factory app.worker_factory

And somewhere in my code I can access worker (by dependency injection, for example) and enqueue tasks.
I need to get rid of any logic in import-time for more control, configurability and flexibility.

The problem with import-time is that I need to have an initialized configured Worker already at the function-task definition stage. But in real applications, I can only configure the Worker after some initial loading and configuration of my application. I can't manage the application configuration at the time of importing modules.

For example, see taskiq dynamic envs (I can't use taskiq because there is no any way to abort tasks).

Any thoughts on this.

[Graeme22]

Hello! I would start by looking at the code in the CLI used to run a worker:
https://github.com/tastyware/streaq/blob/master/streaq%2Fcli.py#L111

Really it just sets up the log config and calls Worker.run_sync(). If you wanted to do something similar to the taskiq example, you can call Worker.run_async(), taking into account that you're running everything in the same event loop of course. (Take a look at the tests for some examples of this.)

If it's important to separate the work out from the worker code, you could try using threads to spawn workers (but not processes since you'd want the worker to access the same task definitions in memory).

Anyways my instinct is that this is possible, give it a try and let me know how it goes.

[espdev]

Yes, I guess I'll have to make a serializer/deserializer for Pydantic models. We have already done this when working with arq. It should be trivial.

[wwarne]

Hi @espdev

I really like the idea of TaskRegistry.

However, since we need to use worker.context inside tasks (and either worker.context or worker.task_context() inside middlewares), I'm currently experimenting with a shared_worker approach.

So I have

# shared_worker.py
import time
from typing import Any

from streaq import Worker
from streaq.types import ReturnCoroutine

from dataclasses import dataclass

@dataclass
class WorkerContext:
    """
    Type safe way of defining the dependencies of your tasks.
    e.g. HTTP client, database connection, settings.
    """
    answer: str
    current_env: str

# Explicitly specify WorkerContext for better IDE and mypy support
shared_worker: Worker[WorkerContext] = Worker(redis_url="NOT SET YET")

# I also put some middlewares here like
@shared_worker.middleware
def timer(task: ReturnCoroutine) -> ReturnCoroutine:
    async def wrapper(*args, **kwargs) -> Any:
        start_time = time.perf_counter()
        result = await task(*args, **kwargs)
        ctx = shared_worker.task_context()
        print(f"Executed task {ctx.task_id} in {time.perf_counter() - start_time:.3f}s")
        return result

    return wrapper

Then I use it in tasks like

# tasks.py
import asyncio
import datetime as dt

from streaq_experiment.shared_worker import shared_worker

@shared_worker.task(timeout=5)
async def fetch(url: str) -> str:
    await asyncio.sleep(3)
    return f'Answer is: {shared_worker.context.answer} - {url}'


@shared_worker.cron('* * * * *')
async def task_every_minute() -> str:
    now = dt.datetime.now(dt.timezone.utc)
    print(f'[{shared_worker.context.current_env}] Now is {now.isoformat()}')
    return f'success at {now.isoformat()}'
    

And in my fastapi endpoints like this

    async with shared_worker:
        await fetch.enqueue("https://tastyware.dev/")
        task = await fetch.enqueue("https://github.com/tastyware/streaq").start(delay=3)
        print(await task.info())
        print(await task.result(timeout=10))

To actually run everything, we need:

A real worker, which will execute tasks. It should have a lifespan function to initialize dependencies. This will typically run in its own process, container, or K8s pod.

An initialized shared_worker for enqueueing tasks. Here we don't want to activate the lifespan function—especially if we're loading large data models or other heavyweight dependencies.

As noted in the docs:

"worker should have the same Redis URL, serializer/deserializer, and queue name as the worker defined elsewhere"

Note – We can simply use create_worker() every time, as long as the lifespan function isn't doing anything heavy. I just wanted to explore whether I could separate them.

So I wrote simple functions

# worker.py
from collections.abc import AsyncIterator
from contextlib import asynccontextmanager

from streaq import Worker
from .worker_context import WorkerContext
from .shared_worker import shared_worker

@asynccontextmanager
async def lifespan(worker: Worker[WorkerContext]) -> AsyncIterator[WorkerContext]:
    """
    Here, we initialize the worker's dependencies.
    You can also do any startup/shutdown work here
    """
    print('HELLO I AM LIFESPAN START')
    yield WorkerContext(answer='42', current_env='dev')
    print('HELLO I AM LIFESPAN END')

def create_worker() -> Worker:
    # here we load dynamic settings from env variables, etc...
    real_worker = Worker(
        redis_url='redis://localhost:6379',  # hardcoded it for demostration
        lifespan=lifespan,
    )
    # copy everything to shared worker
    shared_worker.redis = real_worker.redis
    shared_worker.queue_name = real_worker.queue_name
    shared_worker.serializer = real_worker.serializer
    shared_worker.deserializer = real_worker.deserializer
    shared_worker.lifespan = real_worker.lifespan
    shared_worker.loop = real_worker.loop
    
    return shared_worker

def init_shared_worker() -> None:
    # Same as above but without passing lifespan 
    real_worker = Worker(
        redis_url='redis://localhost:6379',
    )
    # copy everything to shared worker
    shared_worker.redis = real_worker.redis
    shared_worker.queue_name = real_worker.queue_name
    shared_worker.serializer = real_worker.serializer
    shared_worker.deserializer = real_worker.deserializer
    shared_worker.loop = real_worker.loop

Finally I have worker entrypoint. This file does not get imported anywhere I run it as python -m streaq_experiment.worker_entrypoint

# worker_entrypoint.py
import logging

from streaq_experiment.worker import create_worker
# we need to import all tasks here so they are registered
from streaq_experiment import tasks

logging.basicConfig(level=logging.INFO)

if __name__ == '__main__':
    worker = create_worker()
    worker.run_sync()

And I use init_shared_worker() in my FastAPI lifespan so I can use async with shared_worker: later on.

OR here's example of simple script which enqueues some tasks

import asyncio

from streaq_experiment.worker import init_shared_worker
from streaq_experiment.tasks import fetch
from streaq_experiment.shared_worker import shared_worker

async def main() -> None:
    init_shared_worker()

    async with shared_worker:
        await fetch.enqueue("https://tastyware.dev/")
        task = await fetch.enqueue("https://github.com/tastyware/streaq").start(delay=3)
        print(await task.info())
        print(await task.result(timeout=10))


if __name__ == '__main__':
    asyncio.run(main())

It's just one possible approach to this problem. It works, but I haven’t tested it in a real-life project yet. I plan to try it out in a pet project when I have more free time.

Maybe something here will inspire your own solution.

[espdev]

@wwarne

I decided to go the simplest possible route with minimal code addition. I defined the tasks in Pydantic models with minimal wrappers to register and enqueue tasks. It looks something like this:

Task defenitions:

from pydantic import (
    BaseModel,
    ConfigDict,
    Field,
    ImportString,
    PrivateAttr,
    TypeAdapter,
    ValidationInfo,
    field_validator,
)

class QueueTask(BaseModel):
    _name: str | None = PrivateAttr(None)
    
    fn_path: str
    
    max_tries: int = 1
    timeout: timedelta | int | None = None
    ttl: timedelta | int | None = timedelta(minutes=5)
    unique: bool = False

    @cached_property
    def fn(self) -> Callable[[Any], Any]:
        return TypeAdapter(ImportString[Callable[[Any], Any]]).validate_python(self.fn_path)
    
    def enqueue(self, worker: Worker[Any], *args, **kwargs) -> Task[Any]:
        return worker.enqueue_unsafe(self._name, *args, **kwargs)

class BaseQueueTasks(BaseModel):
    model_config = ConfigDict(validate_default=True)

    @field_validator('*', mode='after')
    def _set_name(cls, task: QueueTask, info: ValidationInfo) -> QueueTask:
        task._name = info.field_name
        return task

    def register(self, worker: Worker[Any]):
        for name in type(self).model_fields:
            task: QueueTask = getattr(self, name)

            worker.task(
                max_tries=task.max_tries,
                timeout=task.timeout,
                ttl=task.ttl,
                unique=task.unique,
                name=name,
            )(task.fn)

class QueueTasks(BaseQueueTasks):
    """Queue tasks to register in job queue"""

    foo: QueueTask = QueueTask(fn_path='my_app.tasks.foo')
    bar: QueueTask = QueueTask(fn_path='my_app.tasks.bar', ttl=timedelta(minutes=10))

queue_tasks = QueueTasks()

For an application, I create an AppContext class in which I setup the application context, create FastAPI app and Worker, create and initialize all the necessary resources, etc. The AppContext instance is stored to a module-level variable and then we can access all the application resources in API handlers through a dependency or middleware.

App context:

@dataclass
class AppContext:
    config: Config = field(default_factory=Config)
    worker_mode: bool

    app: FastAPI | None = field(default=None, init=False)
    worker: Worker['AppContext'] = field(init=False)
        
    def __post_init__(self):
        self._create_worker()
        
        if self.worker_mode:
            queue_tasks.register(self.worker)
        else:
            self._create_app()

    @asynccontextmanager
    async def _app_lifespan(self, _app: FastAPI):
        await self._on_startup()
        try:
            async with self.worker:
                yield
        finally:
            await self._on_shutdown()

    @asynccontextmanager
    async def _worker_lifespan(self, _worker: Worker):
        if self.worker_mode:
            await self._on_startup()
        try:
            yield self
        finally:
            if self.worker_mode:
                await self._on_shutdown()

The application can be run in both "web app" and "worker" mode. In "worker" mode all tasks will be registered. In "web app" the FastAPI app and worker are created, and I can call queue_tasks.foo.enqueue(worker) without having to import the task functions.

[kuhnroyal]

Having some built-in form of modularization similar to FastAPIs Router that can lazily be added to a Worker would really help structuring code in larger applications. The type safety of streaq is really nice but it kind of gets defeated when I have to redeclare my task definitions, this could be avoided with some modularization.

[vikigenius]

@Graeme22 I know you said it's a niche use case. But I think it comes up all the time when you have a worker system as a part of a larger application/server like FastAPI etc.

So anything that would make this more ergonomic would really help

[Graeme22]

There's basically two different ways this could be done. We could do something similar to FastAPI routers:

group_a = TaskGroup("a")
group_b = TaskGroup("b")
group_a.add_subgroup(group_b)

@group_a.task()
async def barfoo() -> bool: ...

@group_b.task(timeout=1)
async def sleeper() -> None: ...

worker = Worker()
worker.add_group(group_a)

@worker.task()
async def foobar() -> int: ...

This works great. However, it wouldn't support enqueuing jobs unless you've already entered the worker's async context manager. The other option would be something like this:

worker_a = Worker()
worker_b = Worker()
worker_a.include(worker_b)

@worker_a.task()
async def barfoo() -> bool: ...

@worker_b.task(timeout=1)
async def sleeper() -> None: ...

worker = Worker()
worker.include(worker_a)

@worker.task()
async def foobar() -> int: ...

This would allow for enqueuing tasks in a sub-worker's context manager without needing to worry about the root worker, plus you could run sub-workers individually, so it's more flexible and powerful. I'll have to think about how I'd implement this though. Thoughts @espdev @wwarne @kuhnroyal @vikigenius ?

[Graeme22]

Yes they'd need to have the same URLs/queue name

[vikigenius]

Also why would we want to enqueue tasks outside of the worker context manager? API 2 seems a bit unintuitive to me with the subworkers concept especially with how they have to share the same URLs/queue names.

TaskGroups seem a bit more intuitive as an abstraction. It allows to separate worker definitions completely. A worker can handle multiple TaskGroups.

I can have multiple workers with different queue names that care of different sets of TaskGroups and start them in separate processes/containers.

The only question in API 1 is type safety and how do I know the type of dependencies from WorkerContext in a Task

[kuhnroyal]

For me option 1 looks more intuitive. I assume with option 1 we can enqueue only on the worker and with option 2 we can enqueue on the (sub)workers as well?

[Graeme22]

For me option 1 looks more intuitive. I assume with option 1 we can enqueue only on the worker and with option 2 we can enqueue on the (sub)workers as well?

That's right, which is why option 2 is a lot more powerful. Presumably if you need to modularize the task definitions you'll also often need to modularize enqueuing.

[Graeme22]

Also why would we want to enqueue tasks outside of the worker context manager? API 2 seems a bit unintuitive to me with the subworkers concept especially with how they have to share the same URLs/queue names.

See above, it's likely to be important in many scenarios.

I can have multiple workers with different queue names that care of different sets of TaskGroups and start them in separate processes/containers.

This could be done with option 2 as well, you just don't combine the workers you want to keep separate.

The only question in API 1 is type safety and how do I know the type of dependencies from WorkerContext in a Task

Type safety would be lost because the dependency type wouldn't be known until runtime.

[Graeme22]

I have added a simple way to modularize workers to the master branch: https://streaq.readthedocs.io/en/latest/worker.html#modularizing-workers
Please give it a try and let me know how it works for everyone, if feedback is good I'll include it in the upcoming v6 release.

[Graeme22]

Y'all might find #113 interesting

[Graeme22]

Please see #125, I think this resolves all modularity concerns in a satisfactory way