[Airflow] RuntimeError: Cannot re-initialize CUDA in forked subprocess 에러 해결

실행 환경

• OS: MacOS
• apache-airflow: 2.2.3 버전

 

문제 상황

Airflow로 개발중인 기능은 다음과 같았다.

1. HuggingFace에 업로드 된 크롤링 데이터를 서버로 가져와 학습시킬 수 있는 형태로 전처리한다.
2. 전처리된 데이터를 Lora 활용하여 kullm(구름)의 fine-tuning을 진행한다.

2번 기능을 테스트하던 중 아래와 같은 에러가 발생했다.

RuntimeError: Cannot re-initialize CUDA in forked subprocess. To use CUDA with multiprocessing, you must use the 'spawn' start method

 

airflow 없이 독립적으로 실행할 때는 정상적으로 동작하던 코드였는데, Airflow로 스케줄링하니 위와 같은 문제가 발생한 것이다.

 

문제 분석

에러 메세지를 읽어보면, CUDA 환경에서 프로세스를 fork하여 발생한 에러임을 확인할 수 있다. 부모 프로세스를 fork하여 자식 프로세스를 생성할 때, 부모 프로세스의 메모리는 read-only로 자식에게 공유되는데, CUDA 환경에서는 이러한 방식으로 메모리를 공유할 수 없고, 메모리를 복사하여 넘겨주어야 한다.

Airflow 아키텍쳐

Airflow 환경에서 Executor가 worker를 생성하여 DAG의 작업을 실행시킬때, 프로세스가 fork되면서 위 에러가 발생하는 것으로 추측된다.

이 문제를 해결하기 위한 가장 간단한 방법은 multiprocessing 방식을 fork에서 spawn으로 수정하는 것이다. 

torch.multiprocessing.set_start_method('spawn')

 

그러나, 나는 위 코드를 추가해도 multiprocessing 방식이 수정되지 않았다. 그 이유는 Airflow가 내부적으로 worker를 생성할 때 os.fork를 호출하여 multiprocessing 방식을 수정하더라도 프로세스가 fork되는 것으로 추측하고 있다. (Gunicorn에서 동일한 이슈를 찾을 수 있다.)

 

문제 해결

여러 다른 해결 방법들을 시도한 끝에 내가 찾은 해결책은 PythonOperator 대신 PythonVirtualenvOperator를 사용하는 것이다.

PythonVirtualenvOperator를 사용하면 가상 환경을 생성하여 가상 환경 안에서 작업을 실행시키는데, 이로 인해 프로세스가 fork되지 않기 때문에 문제를 해결할 수 있었다.