vLLM मेमोरी लीक डिबगिंग: Heap के पार UCX और mmap का रहस्य

सारांश:

• समस्या की स्थिति: vLLM के Prefill/Decode अलग-अलग (disaggregated) serving environment में प्रति मिनट 400MB system memory (RSS) लीक हो रही थी, लेकिन सामान्य Python profiler इसे पकड़ नहीं पा रहे थे।
• कारण विश्लेषण: Heaptrack और pmap से पता चला कि लीक heap में नहीं बल्कि anonymous memory mappings (mmap) में हो रही थी, और BPFtrace व automated GDB script के जरिए कारण का पता लगाया गया।
• असली वजह: high-performance communication library UCX optimization के लिए mmap/munmap calls को intercept कर रही थी, और freed memory को तुरंत लौटाने के बजाय अनिश्चित काल तक queue में जमा कर रही थी।
• समाधान: environment variable UCX_MEM_MMAP_HOOK_MODE=none सेट करके UCX की memory hooking feature को disable कर समस्या हल की गई।

विस्तृत सारांश:

1. रहस्यमय मेमोरी लीक

Mistral AI टीम ने vLLM का उपयोग करने वाले Prefill/Decode disaggregated serving (NIXL-आधारित) environment में प्रति मिनट 400MB system memory के linear रूप से बढ़ने की समस्या पाई।

• लक्षण: Python heap memory स्थिर थी, लेकिन operating system स्तर का RSS (Resident Set Size) लगातार बढ़ता रहा और अंततः OOM (Out of Memory) तक पहुंच गया।
• शुरुआती कोशिशें विफल: Memray, Guppy 3 जैसे Python tools सब सामान्य दिखा रहे थे, standard GDB process को crash कर दे रहा था, और Valgrind इतना धीमा था कि उसका उपयोग संभव नहीं था।

2. kernel level पर गहन विश्लेषण

टीम ने महसूस किया कि समस्या का कारण application level (Python/C++) पर नहीं, बल्कि उससे नीचे के स्तर पर है, इसलिए system tools का उपयोग किया गया।

• Heaptrack: visual रूप से पुष्टि हुई कि heap allocations (malloc/free) स्थिर थे, लेकिन RSS बढ़ रहा था। इससे संकेत मिला कि लीक glibc की heap management के बाहर anonymous memory mappings में हो रही थी।
• pmap: /proc/<pid>/maps की निगरानी करके पाया गया कि एक खास anonymous mapping region लगातार बड़ा हो रहा था और उसका address बदल रहा था। इसका मतलब था कि mremap या munmap के बाद mmap cycle बार-बार दोहराया जा रहा था।
• BPFtrace: ऐसे system calls को trace करने के लिए, जिन्हें LD_PRELOAD भी नहीं पकड़ पा रहा था (यानी जो glibc को bypass कर रहे थे), BPFtrace का इस्तेमाल किया गया। इससे पता चला कि mmap calls सीधे syscall के जरिए हो रही थीं।

3. दोषी की पहचान: automated GDB scripting

BPFtrace से समस्या वाले system call address की पहचान करने के बाद, GDB से ऐसा script लिखा गया जो सिर्फ उसी address (SYS_mmap) पर रुकता था।

इस्तेमाल किया गया GDB script उदाहरण:

# mmap system call (number 9) पर conditional breakpoint सेट करें  
break syscall if $rdi == 9  
commands  
  silent  
  # system call return point पर temporary breakpoint सेट करें  
  tbreak *0x00007ffff7d9525d  
  commands  
    silent  
    # stack trace और returned address प्रिंट करें  
    bt  
    printf "Syscall returned: rax = 0x%012lx\n", $rax  
    continue  
  end  
  continue  
end  
  

इस stack trace से निर्णायक सुराग मिला कि UCX (Unified Communication X) library Python के mmap/munmap calls को बीच में intercept कर रही थी।

4. कारण: UCX का अत्यधिक optimization

UCX InfiniBand transfer performance बढ़ाने के लिए memory allocation/deallocation को hook करता है।

• मैकेनिज़्म: जब munmap call होती है, तो memory को तुरंत operating system को वापस नहीं किया जाता, बल्कि बाद में reuse या cleanup के लिए 'invalidation queue' में डाल दिया जाता है।
• बग: default setting (UCX_RCACHE_MAX_UNRELEASED=inf) की वजह से यह queue अनंत तक बढ़ सकती थी, और vLLM के खास usage pattern में cleanup logic (ucp_worker_progress) सही से काम नहीं कर रही थी, जिससे memory लगातार जमा होती गई।

5. समाधान

vLLM के मामले में सिर्फ एक बड़े KVCache memory region को register करना था, इसलिए UCX की जटिल memory hooking feature की वास्तव में जरूरत नहीं थी।

• तत्काल समाधान: environment variable UCX_MEM_MMAP_HOOK_MODE=none सेट करके UCX की memory hooking को पूरी तरह बंद कर दिया गया, जिससे leak रुक गया।
• विकल्प: UCX_RCACHE_MAX_UNRELEASED=1024 की तरह queue size limit करके भी forced cleanup कराया जा सकता है।
• कार्रवाई: vLLM community के लिए यह fix merge कर दी गई है, और भविष्य के NIXL releases में default behavior भी बेहतर किया जाएगा।