DJL 모델 메모리 최적화와 JVM 튜닝 해결기

궁금하면 500원·2025년 7월 25일

AI모델 DJL JVM Java OutofMemoryError PyTorch Spring 개발자 기술블로그 딥러닝 메모리최적화 백엔드 성능개선 트러블슈팅

미생의 개발 이야기

목록 보기

54/58

DJL 모델 메모리 부족 문제 해결 사례

문제 분석

대용량 배치 처리를 위해 **DJL(Deep Java Library)**을 사용하면서 java.lang.OutOfMemoryError: Java heap space 에러가 발생했습니다.

문제의 원인은 크게 세 가지로 분석했습니다.

대용량 모델 로딩: PyTorch 모델 하나가 메모리에 로드될 때마다 2GB 이상의 힙 메모리를 사용했습니다.
비효율적인 동시 처리: 배치 작업 중 여러 스레드가 동시에 모델 인스턴스를 생성하면서 메모리 사용량이 급증했습니다.
부족한 JVM 힙 설정: 기본 JVM 힙 메모리(1GB)가 대용량 모델을 감당하기에 턱없이 부족했습니다.

해결 방안 및 최적화 과정

효율적인 자원 관리와 성능 향상을 위해 다음과 같은 단계로 트러블 슈팅을 진행했습니다.

1단계: 메모리 사용량 모니터링

정확한 문제 진단을 위해 **MemoryMXBean**을 활용한 모니터링 로직을 추가했습니다. 이를 통해 모델 로딩 전후의 힙 메모리 사용량을 실시간으로 확인하고, 최적화 효과를 명확하게 측정할 수 있었습니다.

@Component
public class ModelMemoryMonitor {
    private final MemoryMXBean memoryBean = ManagementFactory.getMemoryMXBean();
    
    public void logMemoryUsage(String context) {
        MemoryUsage heapUsage = memoryBean.getHeapMemoryUsage();
        long usedMB = heapUsage.getUsed() / (1024 * 1024);
        long maxMB = heapUsage.getMax() / (1024 * 1024);
        
        log.info("Memory Usage [{}]: {}MB / {}MB ({}%)", 
            context, usedMB, maxMB, (usedMB * 100 / maxMB));
    }
}

2단계: 모델 캐싱 및 동시성 제어

매번 새로운 모델 인스턴스를 생성하는 비효율을 개선하기 위해 모델 캐싱 전략을 도입했습니다. 또한, **Semaphore**를 사용하여 동시에 로드되는 모델의 수를 제한함으로써, 메모리 과부하를 방지했습니다.

@Service
public class OptimizedDjlAnalysisTasklet {
    private static final int MAX_CONCURRENT_MODELS = 2;
    private final Semaphore modelSemaphore = new Semaphore(MAX_CONCURRENT_MODELS);
    private final ConcurrentHashMap<String, Model> modelCache = new ConcurrentHashMap<>();
    
    public DjlSentimentResult analyzeSentiment(String text) {
        modelSemaphore.acquire();
        try {
            Model model = getOrCreateModel();
            // 분석 로직
            return performAnalysis(model, text);
        } finally {
            modelSemaphore.release();
        }
    }
    
    private Model getOrCreateModel() {
        return modelCache.computeIfAbsent("sentiment-model", key -> {
            try {
                return Model.newInstance("sentiment");
            } catch (Exception e) {
                throw new RuntimeException("Failed to create model", e);
            }
        });
    }
}

3단계: JVM 힙 메모리 튜닝

대용량 모델을 안정적으로 운영하기 위해 JVM 힙 설정을 확장했습니다. **-Xms2g -Xmx4g** 옵션을 적용하여 초기 힙 크기를 2GB로, 최대 크기를 4GB로 설정했습니다. 또한, G1GC(Garbage-First Garbage Collector)를 사용하여 효율적인 가비지 컬렉션을 유도했습니다.

# JVM 옵션 예시
java -Xms2g -Xmx4g -XX:+UseG1GC -XX:MaxGCPauseMillis=200 -jar your-application.jar

4단계: DJL 모델 로딩 최적화

DJL 자체의 **ModelManager**를 활용하여 모델 캐싱을 더욱 효율적으로 관리했습니다. 설정 파일을 통해 캐시 크기와 만료 시간을 지정하여, 메모리 관리를 자동화하고 애플리케이션의 유연성을 높였습니다.

@Configuration
public class DjlOptimizationConfig {
    
    @Bean
    @ConditionalOnProperty(name = "ai.djl.performance.enable-caching", havingValue = "true")
    public ModelManager modelManager() {
        return ModelManager.builder()
            .setMaxCacheSize(500)
            .setCacheExpiration(Duration.ofHours(1))
            .build();
    }
}