I have a celery task that uses torch library, which internally uses CUDA. When I run the task, it fails saying "Cannot re-initialize CUDA in forked subprocess. To use CUDA with multiprocessing, you must use the 'spawn' start method"
When I browsed on this a little, I got this - https://github.com/celery/celery/issues/6036 This issue says Celery supports only fork and not spawn.
Is there any workaround/alternative to this?
you should load pytorch model and use it as global variable. The instance of model and task will run in the same process
from celery.signals import worker_process_init
# for more information about worker_process_init, read here:
# https://docs.celeryproject.org/en/stable/userguide/signals.html#worker-process-init
pytorch_model = None
@worker_process_init.connect()
def init_worker_process(**kwargs):
"""
load model before running tasks
:param kwargs:
:return:
"""
global pytorch_model
pytorch_model = load_model()
@app.task
def predict_task(image: np.ndarray):
return pytorch_model.predict(image)