[python] gpu(cuda) 프로세스 정리하기

CUDA를 사용하는 프로세스 종료하기

저희 flask 서버에서는 gpu를 사용하는 모델을 2개 돌립니다.
- spleeter (mr제거)
- cnn 분류기
 
spleeter 작업 후 cnn분류기가 연이어 있습니다.
문제는 spleeter가 끝나더라도 gpu 메모리가 곧장 내려가지는 않습니다.
정리해줘야 합니다. Out Of Memory가 잘 터지기 때문이죠!
 
그래서 gpu 정리 기능을 추가했습니다.
gpu 점유 프로세스를 모두 정리하는 코드입니다.
(의도치않은 gpu 프로세스가 종료되지 않는지 검토하고 사용해야 합니다.)
(cpu 환경에서는 CudaAPIError가 터지기 때문에, 아무 것도 안하도록 예외처리했습니다.)

from numba import cuda
from numba.cuda.cudadrv.devices import CudaAPIError

# gpu 메모리 완전 정리
def clean_gpu():
    try:
        device = cuda.get_current_device()
        device.reset()
    except CudaAPIError:
        # CPU 모드에서는 정리할 것이 없다
        pass

현재 작업중인 cuda device에 대한 정보가 나옵니다. 간단하게는 device를 print해보면 gpu명이 나옵니다. 
예를 들면 <CUDA device 0 'b'NVIDIA GeForce MX450''> 이렇게.. 
 
 
참고) Cpu환경에서의 CudaAPIError

---------------------------------------------------------------------------
CudaAPIError                              Traceback (most recent call last)
c:\Users\kye33\Desktop\mr_remover_test\mr_remover_test.ipynb Cell 8 line 9
      6     print(device)
      7     device.reset()
----> 9 clean_gpu()

c:\Users\kye33\Desktop\mr_remover_test\mr_remover_test.ipynb Cell 8 line 5
      4 def clean_gpu():
----> 5     device = cuda.get_current_device()
      6     print(device)
      7     device.reset()

File c:\Users\kye33\miniconda3\envs\gpuenv\lib\site-packages\numba\cuda\api.py:443, in get_current_device()
    441 def get_current_device():
    442     "Get current device associated with the current thread"
--> 443     return current_context().device

File c:\Users\kye33\miniconda3\envs\gpuenv\lib\site-packages\numba\cuda\cudadrv\devices.py:220, in get_context(devnum)
    216 def get_context(devnum=None):
    217     """Get the current device or use a device by device number, and
    218     return the CUDA context.
    219     """
--> 220     return _runtime.get_or_create_context(devnum)
...
    397 if retcode == enums.CUDA_ERROR_NOT_INITIALIZED:
    398     self._detect_fork()
--> 399 raise CudaAPIError(retcode, msg)

CudaAPIError: [2] Call to cuDevicePrimaryCtxRetain results in CUDA_ERROR_OUT_OF_MEMORY

 
 
사용 예시!
gpu 정리가 필요한 곳에 배치했습니다.

 
 
효과가 있는 모습입니다.
OOM 문제를 해결했습니다. 끝!