超簡便Python任務隊列:huey

項目地址:coleifer/huey
文檔地址:huey, a little task queue

Why

其實這種需求很常見 。在跑仿真的時候,爲了得到多組數據以供統計分析比對,我們需要將一些程序按照不同參數配置運行多次。爲了利用現代計算機的多核性能,我們常常以多進程的方式來運行這些仿真程序以最大化利用硬件資源。

這是我開始使用huey的時候的需求,但是事實上huey的應用遠不止於此。和celery類似,你可以把huey用來管理和調度異步任務,或者是用來運行週期性任務以及定時服務。

我最開始的方法是寫了一個工具腳本,在這個腳本內部開了一個線程池,然後在將計算任務提交給這個線程池。但是這樣有一個問題:如果我已經開始了一組計算任務,而我又希望再增加一組任務,我只能等前一組任務完全結束之後才能再開始下一組。不然前後兩組任務會同時開始運行,爭搶計算資源,而不是按照我設想的任務的隊列的模式。

設想一個具體的例子。我的計算機是8核的,而我的一組計算任務需要以不同的參數運行16次。通過在進程內創建一個8線程容量的線程池,我可以讓這16個任務進入隊列序次執行,同時執行的現成數不超過8。倘若前一次的16個任務還剩餘4個未運行,而我又要再增加16個計算任務,此時以這種模式,我無法將新增的任務放到前一個進程的任務隊列裏。如果直接強行運行,則會導致前一組任務的4個線程在運行,而後一個任務會調度出8個線程,總計12個線程,這並不能最優地使用多核。

那麼這個時候我就需要一個任務隊列系統來管理我的計算任務了:每次我要運行一組任務的時候,我將計算任務提交給任務隊列系統,由這個系統來調度worker(worker的數量常常設定爲CPU的線程數)執行。簡而言之,任務隊列系統提供了跨進程的隊列。

其實之前我用過celery這個任務隊列框架,可是celery是用起來“太重”,配置較爲複雜。因此我後來選擇了huey這個框架。這個框架的好處是:

  1. 配置起來非常簡單
  2. 支持自動重試失敗的任務

下面我來簡要介紹一下這個框架的配置方法。當然更多的功能你可以前往其官網研究。

How

安裝

pip install huey

配置

# worker.py
from huey import RedisHuey

huey = RedisHuey()

@huey.task()
def add_numbers(a, b):
    return a + b

使用RedisHuey需要安裝redis,這個請自行Google安裝方法。如果覺得Redis還不夠輕量級。可以選擇Sqlite作爲broker,具體方法參見huey的文檔,我這裏不做贅述。

沒錯,設置就是怎麼簡單。huey中任務的定義圍繞着huey.task這個裝飾器來進行。如果相對任務做更多的功能設置,可以通過爲huey.task傳入更多的參數實現。例如設置自動重試:

@huey.task(retries=3)
def add_numbers(a, b):
    return a + b

此時任務如果遭遇失敗,會嘗試重試3次。如果重試之後仍然失敗則會拋出異常。

啓動Huey Worker

huey提供了huey_consumer.py這個腳本來快捷啓動worker。命令的基本使用方式如下:

huey_consumer.py worker.huey -k process -w 4

注意其中的worker.huey,你需要在worker.py的目錄下執行這個腳本,而worker.huey是你import這個對象需要的路徑形式。其餘常用的flag定義如下:

-w n                 worker的數量
-k process/thread/greenlet    worker使用線程還是進程還是greenlet
-v           log輸出詳細的debug信息
-0          

使用Worker

在做好上述配置之後,將計算任務發送給worker執行的過程同普通的函數調用並沒有什麼分別。

from worker import huey, add_numbers

add_numbers(1, 2)

儘管看起來和普通的函數調用無異,但其實huey在背後將函數及其參數序列化後通過Redis這個broker,提交給了huey_consumer進程。這個進程在反向解析出函數和參數後,調度worker來執行這個函數。值得注意的是,這裏的函數調用就變成了異步的,而在調用者側函數調用的返回結果,則成爲一個特殊打包對象,我們可以通過這個對象來獲取任務執行的結果。

from worker import huey, add_numbers

res = add_numbers(1, 2)
add_result = res(blocking=True)  # get 3 

小結

以上是對huey這個任務隊列工具的簡單介紹,適合之前沒有接觸過任務隊列工具的朋友入門的。事實上huey還爲我們提供了更多的功能,包括週期性任務調度,延時調用,以及計算任務的暫停,取消和重新調度,consumer事件管理等等。關於這些內容我會後續更新文章介紹。

這裏我給大家提供一個較爲通用的用來執行一個Bash命令行命令的worker模板:

import subprocess
import os

from huey import RedisHuey

huey = RedisHuey()
@huey.task(retries=3)
def run_bash_task(cmd, std_out, working_dir):
    print("Entring %s" % working_dir)
    print("Running command %s, stdout to %s" % (cmd, std_out))
    os.chdir(working_dir)
    if std_out is None:
        subprocess.check_call(cmd)
    else:
        with open(std_out, "w") as f:
            subprocess.check_call(cmd, stdout=f)
    print("finish task %s\n" % cmd)

發表評論
所有評論
還沒有人評論,想成為第一個評論的人麼? 請在上方評論欄輸入並且點擊發布.
相關文章