dramatiq 学习Dramatiq是一个简单好用的python任务队列。任务队列工作方式类似于生产者-消费者模型里的

Dramatiq是一个简单好用的python任务队列。今天来学习一下。

前言

任务队列，是做什么的呢？

举个例子，在一个web服务中，有时会做一些耗时的操作。但web服务的接口，需要在几秒之内返回响应。这种情况下，可以把耗时操作放在任务队列中执行。任务完成时把结果放在数据库中，web服务可以去数据库看结果。

任务队列工作方式类似于生产者-消费者模型里的消费者。在以上例子中，web服务就像是一个生产者，生产一些消息放到消息队列中。任务队列(下文中的worker)就像消费者，从消息队列中拿出消息，执行消息对应的任务(下文中的actor)。

安装

pip install -U 'dramatiq[redis, watch]'

即可安装。它支持rabbitmq/redis做broker。这里我选的是redis，因为我熟悉redis，可以在redis中看它的数据是如何储存的。另外本地有个redis，端口是6379。

注意如果你本地没运行redis/rabbitmq，或者不是默认端口，你应该在代码里进行相关设置后，再使用dramatiq。参考下文的消息中间人章节。

基本的使用方式

一个不使用任务队列的例子

我们先用同步的方式，来执行一个耗时的操作先创建一个模块count_words.py, 内容如下

import requests
url = "https://dramatiq.io/"
def count_words(url):
    response = requests.get(url)
    count = len(response.text.split())
    print(f"There are {count} words at {url}.")

再运行count_words函数。

In [2]: from count_words import count_words, url

In [3]: count_words(url)
There are 1008 words at https://dramatiq.io/.

count_words运行了大概一秒的时间。

Actors和Workers

下面我们借助dramatiq，让此函数在任务队列中运行。这样我们当前的线程不会阻塞住。修改count_words.py

import requests
import dramatiq
url = "https://dramatiq.io/"
@dramatiq.actor
def count_words(url):
    response = requests.get(url)
    count = len(response.text.split())
    print(f"There are {count} words at {url}.")

这里我们用@dramatiq.actor装饰器来作用在count_words函数上，创建了一个actor(任务)。它可以在任务队列中执行。启动dramatiq任务队列

(learn) ➜  learn_python dramatiq count_words   
[2023-03-18 22:40:49,360] [PID 239667] [MainThread] [dramatiq.MainProcess] [INFO] Dramatiq '1.14.1' is booting up.
[2023-03-18 22:40:49,303] [PID 239668] [MainThread] [dramatiq.WorkerProcess(0)] [INFO] Worker process is ready for action.
[2023-03-18 22:40:49,317] [PID 239669] [MainThread] [dramatiq.WorkerProcess(1)] [INFO] Worker process is ready for action.
...

然后让count_words在任务队列中执行

In [1]: from count_words import count_words, url

In [2]: count_words.send(url)
Out[2]: Message(queue_name='default', actor_name='count_words', args=('https://dramatiq.io/',), kwargs={}, options={'redis_message_id': '950441b8-b989-4bd6-aa49-b29b2827ae8d'}, message_id='a63556cf-3a89-469f-81d8-e0f7319788ad', message_timestamp=1679150532924)

可以看到没有阻塞，也没有打印出相关的消息。因为count_words函数没有在当前的线程、进程中运行，只是把消息发给了任务队列。而在dramatiq任务队列的log，可以找到count_words函数的输出。

...
[2023-03-18 22:40:49,360] [PID 239861] [MainThread] [dramatiq.ForkProcess(0)] [INFO] Fork process 'dramatiq.middleware.prometheus:_run_exposition_server' is ready for action.
There are 1008 words at https://dramatiq.io/.

在我们启动dramatiq命令时，dramatiq会启动多个worker进程。这些进程从消息队列中把消息拿出来，发给相关的actor函数，执行。

错误处理

我们现在给actor函数一个错误的输入，试试它会怎样处理错误。

In [3]: count_words.send("foo")
Out[3]: Message(queue_name='default', actor_name='count_words', args=('foo',), kwargs={}, options={'redis_message_id': '2491df6f-8008-47a9-a3db-75dc101cd2fc'}, message_id='f1942f1e-2b65-45c3-9632-fdf9bd978625', message_timestamp=1679207648677)

  File "/home/xyc/.virtualenvs/learn/lib/python3.11/site-packages/requests/models.py", line 439, in prepare_url
    raise MissingSchema(
requests.exceptions.MissingSchema: Invalid URL 'foo': No scheme supplied. Perhaps you meant https://foo?
[2023-03-19 14:34:08,760] [PID 239674] [Thread-9] [dramatiq.middleware.retries.Retries] [INFO] Retrying message 'f1942f1e-2b65-45c3-9632-fdf9bd978625' in 9840 milliseconds.
...
  File "/home/xyc/.virtualenvs/learn/lib/python3.11/site-packages/requests/models.py", line 439, in prepare_url
    raise MissingSchema(
requests.exceptions.MissingSchema: Invalid URL 'foo': No scheme supplied. Perhaps you meant https://foo?
[2023-03-19 14:34:19,301] [PID 239677] [Thread-11] [dramatiq.middleware.retries.Retries] [INFO] Retrying message 'f1942f1e-2b65-45c3-9632-fdf9bd978625' in 23824 milliseconds.
...

可以看出它在进行错误重试我们改一下代码，在actor函数里把这个错误给处理掉。

@dramatiq.actor
def count_words(url):
    try:
        response = requests.get(url)
        count = len(response.text.split())
        print(f"There are {count} words at {url}.")
    except requests.exceptions.MissingSchema:
        print(f"{url} doesn't look like a valid URL.")

然后重启dramatiq队列，然后重试count_words.send("foo")

[2023-03-19 14:39:23,922] [PID 248106] [MainThread] [dramatiq.ForkProcess(0)] [INFO] Fork process 'dramatiq.middleware.prometheus:_run_exposition_server' is ready for action.
foo doesn't look like a valid URL.

这次果然它能正确在函数里正确处理异常了。

代码重载

为了使新的代码生效，除了以上方法之外还有以下两种：

向主进程发送 HUP 信号

我把代码修改成以下形式

@dramatiq.actor
def count_words(url):
    print("test code reload by sending signal HUP")
    try:
        response = requests.get(url)

从正在运行的dramatiq的log中找到这一行[2023-03-19 14:39:23,921] [PID 247910] [MainThread] [dramatiq.MainProcess] [INFO] Dramatiq '1.14.1' is booting up.可知MainProcess的PID为247910。发送信号kill -s HUP 247910。果然重载了代码

[2023-03-19 14:51:54,430] [PID 247910] [MainThread] [dramatiq.MainProcess] [INFO] Sending signal 1 to subprocesses...
[2023-03-19 14:51:54,431] [PID 247911] [MainThread] [dramatiq.WorkerProcess(0)] [INFO] Stopping worker process...
...
[2023-03-19 14:51:57,323] [PID 247910] [MainThread] [dramatiq.MainProcess] [INFO] Dramatiq '1.14.1' is booting up.
[2023-03-19 14:51:57,264] [PID 251573] [MainThread] [dramatiq.WorkerProcess(0)] [INFO] Worker process is ready for action.

使用`--watch`参数启动dramatiq。

我们使用dramatiq --watch . count_words启动任务队列。修改代码后果然重载了。这里--watch后面的参数是个directory，其中的任意源文件变更都会引起代码重载。

细粒度地控制任务运行

前面我们学习了怎样让任务运行。这里我们学习如何细粒度地控制任务的运行。

延迟启动任务

使用send_with_options, 可以指定任务的一些参数，例如delay可以指定延迟毫秒数。

In [7]: count_words.send_with_options(args=(url,), delay=1500)
Out[7]: Message(queue_name='default.DQ', actor_name='count_words', args=('https://dramatiq.io/',), kwargs={}, options={'redis_message_id': '906f1242-238b-453b-bf5a-7a5c0c721b23', 'eta': 1679215416822}, message_id='87a1be27-1fcf-4778-a172-7778baf7ee13', message_timestamp=1679215415321)

以上代码让任务1.5s后执行。 send和send_with_options有什么不一样？

    def send(self, *args: P.args, **kwargs: P.kwargs) -> Message[R]:
        # ...
        return self.send_with_options(args=args, kwargs=kwargs)

    def send_with_options(
        self, *,
        args: tuple = (),
        kwargs: Optional[Dict[str, Any]] = None,
        delay: Optional[int] = None,
        **options,
    ) -> Message[R]:
        # ...
        message = self.message_with_options(args=args, kwargs=kwargs, **options)
        return self.broker.enqueue(message, delay=delay)

查看源码可知，send也是send_with_options实现的。send的入参的签名和原函数相同。send_with_options通过args和kwargs把入参传入原函数，所以可以支持传入options来传入设置。

消息中间人(Message Brokers)

使用send/send_with_options创建任务时，会把消息（包括要传给原函数的参数和其它的一些信息），保存到数据库中。这个数据库就是Message Broker

在本次使用的代码中，没有指定broker。但通过日志可以看出，dramatiq自动地设置了redis作为broker。

我们使用count_words.send_with_options(args=(url,), delay=10000)来创建一个延迟高达10s的任务，所以任务信息可以在redis里保存较长时间。来看一下redis

127.0.0.1:6379> keys *
1) "dramatiq:default.DQ.msgs"
2) "dramatiq:__heartbeats__"
3) "dramatiq:__acks__.06da3859-a327-4238-aeb9-af15e3262074.default.DQ"
127.0.0.1:6379>

这证明dramatiq自动使用了这个redis作为它的broker 当然也可以设置rabbitmq为它的broker。我本地有个docker里运行着的rabbitmq，端口是5672 首先进行相关的安装pip install 'dramatiq[rabbitmq, watch]' 然后在代码这样设置，即可使用rabbitmq作为broker

import dramatiq
from dramatiq.brokers.rabbitmq import RabbitmqBroker
dramatiq.set_broker(RabbitmqBroker(url="amqp://localhost:5672"))
...

总结

dramatiq比celery用起来简单一点。使用过程中，我发现它的类型提示做的挺好。现在只学习了基础的使用方式，高级的使用方式以后用到了再学吧。

dramatiq 学习

前言

安装