EC2 Django에서 GPU 할당 제한하기

2020년 12월 07일 기록

문제 상황

views.py 에서 ML 모델 Predict 하는 함수를 구현했는데, tensorflow는 gpu를 사용한다고 한다. 그런데 기본적으로 gpu는 남은 메모리의 95%를 모두 사용하도록 되어 있기 때문에 너무 많은 메모리를 잡아먹어서 함께 다른 작업을 할 수 없는 상황이 발생했다.

점유할 gpu 메모리를 제한하는 코드가 있었는데, views.py에서 아무리 호출해줘도 적용되지 않았다. 지금 와서 생각해보니 처음에는 됐는데 매번 요청할 때마다 메모리를 할당해서? 결국 Out Of Memory 에러가 나는 것 같다.

해결 방법

처음 Django 서버를 실행할 때 딱 한 번만 gpu 메모리를 할당하도록 __init__.py 파일에 구현해주면 잘 되었다.

### __init__.py

import tensorflow as tf
import os

gpus = tf.config.experimental.list_physical_devices('GPU')

if gpus:
    # 특정 GPU에 1GB 메모리만 할당하도록 제한
    try:
        tf.config.experimental.set_visible_devices(gpus[0], 'GPU')
        tf.config.experimental.set_virtual_device_configuration(
            gpus[0],
            [tf.config.experimental.VirtualDeviceConfiguration(memory_limit=4000)])
    except RuntimeError as e:
        # 프로그램 시작시에 가상 장치가 설정되어야만 합니다
        print(e)