Search code examples
pytorchcelery

Unable to use Pytorch with CUDA in Celery task

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?

Solution

  • You may start celery worker with "solo" pool type

    celery -A tasks worker --pool=solo --loglevel=info

    Solo creates only one thread and runs celery tasks using that thread. Concurrency number cannot be provided here.