26 lines
660 B
Python
26 lines
660 B
Python
import gc
|
|
import logging
|
|
|
|
import psutil
|
|
|
|
|
|
class MemoryMonitor:
|
|
"""Monitor and manage memory usage during batch processing"""
|
|
|
|
@staticmethod
|
|
def get_memory_usage_mb() -> float:
|
|
"""Get current memory usage in MB"""
|
|
process = psutil.Process()
|
|
return process.memory_info().rss / 1024 / 1024
|
|
|
|
@staticmethod
|
|
def cleanup_memory():
|
|
"""Force garbage collection"""
|
|
gc.collect()
|
|
|
|
@staticmethod
|
|
def log_memory_usage(step_name: str):
|
|
"""Log current memory usage"""
|
|
memory_mb = MemoryMonitor.get_memory_usage_mb()
|
|
logging.info(f"Memory usage after {step_name}: {memory_mb:.1f} MB")
|