nvflare.fuel.utils.memory_utils module

Memory management utilities for federated learning.

This module provides memory cleanup utilities to help manage RSS (Resident Set Size) in long-running FL jobs using Python + PyTorch + glibc/jemalloc.

Allocator Support: - glibc: Uses malloc_trim() to return freed pages to OS - jemalloc: Relies on auto-decay (MALLOC_CONF), no manual action needed

Best Practices: - Client: Set MALLOC_ARENA_MAX=2, cleanup every round - Server: Set MALLOC_ARENA_MAX=4, cleanup every 5 rounds - jemalloc: Set MALLOC_CONF=”dirty_decay_ms:5000,muzzy_decay_ms:5000”

Usage:

from nvflare.fuel.utils.memory_utils import cleanup_memory, get_allocator_type

# Check which allocator is in use allocator = get_allocator_type() # “glibc”, “jemalloc”, or “unknown”

# At end of each round (client) or every N rounds (server) cleanup_memory(cuda_empty_cache=True) # True for PyTorch GPU clients

cleanup_memory(cuda_empty_cache: bool = False) → None[source]

Perform allocator-aware memory cleanup to reduce RSS.

This function: 1. Runs Python garbage collection (gc.collect) 2. For glibc: Releases free heap pages to OS (malloc_trim)

For jemalloc: Relies on auto-decay (no manual action needed)

  1. Optionally clears PyTorch CUDA cache

Parameters:

cuda_empty_cache – If True, also call torch.cuda.empty_cache(). Only applicable to PyTorch GPU clients.

Note

Call this at the end of each FL round (client) or every N rounds (server). The function automatically detects the allocator type and applies the appropriate cleanup strategy.

get_allocator_type() → str[source]

Detect which memory allocator is in use at runtime.

Returns:

jemalloc is loaded (recommended for PyTorch) “glibc”: Standard glibc malloc is in use “unknown”: Could not detect allocator type

Return type:

“jemalloc”

Note

  • jemalloc is typically loaded via LD_PRELOAD

  • Detection is cached after first call

  • Safe to call frequently (no overhead after first call)

try_malloc_trim() → int | None[source]

Attempt to release free memory back to the OS (glibc only).

This calls glibc’s malloc_trim(0) to return free heap pages to the OS, which helps reduce RSS (Resident Set Size) after memory-intensive operations.

Returns:

1 if memory was released, 0 if not, None if malloc_trim is not available.

Note

  • Only works on Linux with glibc (not musl/Alpine)

  • Safe no-op on other platforms

  • Very low overhead, safe to call frequently