Python Redis Queue (rq) - how to avoid preloading ML model for each job?

Question

I want to queue my ml predictions using rq. Example code (pesudo-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()

I've read rq docs explaining that you can preload libraries to avoid importing them every time a job is run (so in example code I import tensorflow in the worker code). However, I also want to move model loading from predict_stuff to avoid loading the model every time the worker runs a job. How can I go about that?

Answer 1

I'm not sure if this is something that can help but, following the example here:

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

Instead of sharing a connection pool, you can share the model.

pseudo code:

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

I use something similar to this for a "global" file reader I wrote. Load up the instance into the LocalStack and have the workers read off the stack.