时间戳

首页 > 解决方案 > Python Redis Queue (rq) - 如何避免为每个作业预加载 ML 模型？

python - Python Redis Queue (rq) - 如何避免为每个作业预加载 ML 模型？

问题描述

我想使用 rq 对我的 ml 预测进行排队。示例代码（伪ish）：

predict.py：

import tensorflow as tf

def predict_stuff(foo):
    model = tf.load_model()
    result = model.predict(foo)
    return result

app.py：

from rq import Queue
from redis import Redis
from predict import predict_stuff

q = Queue(connection=Redis())
for foo in baz:
    job = q.enqueue(predict_stuff, foo)

worker.py：

import sys
from rq import Connection, Worker

# Preload libraries
import tensorflow as tf

with Connection():
    qs = sys.argv[1:] or ['default']

    w = Worker(qs)
    w.work()

我已经阅读了 rq 文档，解释说您可以预加载库以避免每次运行作业时都导入它们（因此在示例代码中，我在工作代码中导入 tensorflow）。但是，我还想移动模型加载，predict_stuff以避免每次工作人员运行作业时加载模型。我该怎么办？

标签： pythonredispython-rq

解决方案

我不确定这是否有帮助，但是按照此处的示例：

https://github.com/rq/rq/issues/720

您可以共享模型，而不是共享连接池。

伪代码：

import tensorflow as tf

from rq import Worker as _Worker
from rq.local import LocalStack

_model_stack = LocalStack()

def get_model():
    """Get Model."""
    m = _model_stack.top
    try:
        assert m
    except AssertionError:
        raise('Run outside of worker context')
    return m

class Worker(_Worker):
    """Worker Class."""

    def work(self, burst=False, logging_level='WARN'):
        """Work."""
        _model_stack.push(tf.load_model())
        return super().work(burst, logging_level)

def predict_stuff_job(foo):
    model = get_model()
    result = model.predict(foo)
    return result

我将类似的东西用于我编写的“全局”文件阅读器。将实例加载到 LocalStack 并让工作人员读取堆栈。

推荐阅读