DevLab / infra / memory-leak-detector-for-python-processes
INFRA PYTHON INTERMEDIATE

Memory Leak Detector for Python Processes

Long-running Python services—such as web workers, queue consumers, and background schedulers—frequently suffer from memory leaks. In Python, these leaks ty

Python3.10+
OSLinux
Lines~219
Packages2 deps
Install Dependencies
$ pip install psutil tracemalloc
Usage
$ python3 mem_leak.py --pid 1234 --interval 10 --alert-mb 500
#infra#python#intermediate#psutil#tracemalloc

Problem Statement

Long-running Python services—such as web workers, queue consumers, and background schedulers—frequently suffer from memory leaks. In Python, these leaks typically stem from unbounded in-memory caches, circular references involving custom __del__ methods, retained global references, or unmanaged native memory in C-extensions (e.g., XML parsers, image processing libraries). Left undetected, a leaking process will eventually exhaust system resources, leading to host instability or abrupt Out-Of-Memory (OOM) kernel kills.

Diagnosing these issues in production without restarting the process or introducing heavy performance overhead is challenging. This script solves the problem by providing a non-invasive, continuous memory monitoring tool. It tracks external process memory footprints via psutil and calculates memory growth velocity over time. Additionally, it integrates Python's native tracemalloc engine for profiling local executions, allowing engineers to identify leaks before they cause outages.


How It Works

The detector operates by establishing a baseline memory metric for a target process using its Process ID (PID). At user-defined intervals, the script samples the target process's Resident Set Size (RSS), Virtual Memory Size (VMS), thread count, and open file descriptors using psutil. It maintains a sliding window of these snapshots to distinguish between harmless, temporary memory spikes (which garbage collection eventually reclaims) and sustained linear memory growth (true leaks).

If the monitored process breaches a predefined total memory ceiling (--alert-mb) or exhibits a continuous upward drift over a configurable sampling window (--growth-threshold-mb), the script triggers structured alerts. When run locally or in profiling mode, the script leverages Python's built-in tracemalloc module to capture diffs between memory allocations, pin-pointing the exact file names and line numbers responsible for object retention.


Full Script

#!/usr/bin/env python3
"""
Memory Leak Detector for Python Processes

Monitors a running process for memory leaks using psutil for system-level statistics
and tracemalloc for detailed in-process memory snapshot analysis.
"""

from __future__ import annotations

import argparse
import logging
import os
import sys
import time
import tracemalloc
from dataclasses import dataclass
from datetime import datetime, timezone
from typing import List, Optional, Tuple

import psutil


# ============================================================================
# Configuration & Defaults
# ============================================================================
DEFAULT_INTERVAL = float(os.environ.get("MEM_LEAK_INTERVAL", "5.0"))
DEFAULT_ALERT_MB = float(os.environ.get("MEM_LEAK_ALERT_MB", "500.0"))
DEFAULT_GROWTH_THRESHOLD_MB = float(os.environ.get("MEM_LEAK_GROWTH_THRESHOLD", "50.0"))
DEFAULT_WINDOW_SIZE = int(os.environ.get("MEM_LEAK_WINDOW_SIZE", "12"))
LOG_LEVEL = os.environ.get("MEM_LEAK_LOG_LEVEL", "INFO").upper()

logging.basicConfig(
    level=getattr(logging, LOG_LEVEL, logging.INFO),
    format="%(asctime)s [%(levelname)s] %(message)s",
    handlers=[logging.StreamHandler(sys.stdout)],
)
logger = logging.getLogger("mem_leak_detector")


# ============================================================================
# Data Structures
# ============================================================================
@dataclass(frozen=True)
class MemorySnapshot:
    """Stores a point-in-time memory measurement for a monitored process."""
    timestamp: datetime
    rss_mb: float
    vms_mb: float
    num_threads: int
    open_files_count: int


@dataclass
class LeakReport:
    """Summary report detailing memory growth and potential leaks."""
    pid: int
    duration_seconds: float
    start_rss_mb: float
    current_rss_mb: float
    peak_rss_mb: float
    net_growth_mb: float
    is_leaking: bool
    warning_reason: str


# ============================================================================
# Core Detector Logic
# ============================================================================
class MemoryLeakDetector:
    """Monitors process memory usage and determines sustained growth patterns."""

    def __init__(
        self,
        pid: int,
        alert_mb: float = DEFAULT_ALERT_MB,
        growth_threshold_mb: float = DEFAULT_GROWTH_THRESHOLD_MB,
        window_size: int = DEFAULT_WINDOW_SIZE,
    ) -> None:
        self.pid = pid
        self.alert_mb = alert_mb
        self.growth_threshold_mb = growth_threshold_mb
        self.window_size = window_size
        self.snapshots: List[MemorySnapshot] = []
        
        try:
            self.process = psutil.Process(self.pid)
        except psutil.NoSuchProcess:
            logger.error("Process with PID %d does not exist.", self.pid)
            sys.exit(1)
        except psutil.AccessDenied:
            logger.error("Access denied when accessing PID %d. Run with higher privileges.", self.pid)
            sys.exit(1)

    def capture_snapshot(self) -> MemorySnapshot:
        """Collects current memory metrics from the target process."""
        mem_info = self.process.memory_info()
        
        try:
            open_files = len(self.process.open_files())
        except (psutil.AccessDenied, psutil.NoSuchProcess):
            open_files = -1

        snapshot = MemorySnapshot(
            timestamp=datetime.now(timezone.utc),
            rss_mb=mem_info.rss / (1024 * 1024),
            vms_mb=mem_info.vms / (1024 * 1024),
            num_threads=self.process.num_threads(),
            open_files_count=open_files,
        )
        
        self.snapshots.append(snapshot)
        if len(self.snapshots) > self.window_size:
            self.snapshots.pop(0)

        return snapshot

    def analyze(self) -> LeakReport:
        """Analyzes historical snapshots to detect memory leak behavior."""
        if not self.snapshots:
            raise ValueError("No snapshots available to analyze.")

        first = self.snapshots[0]
        latest = self.snapshots[-1]
        
        peak_rss = max(s.rss_mb for s in self.snapshots)
        net_growth = latest.rss_mb - first.rss_mb
        duration = (latest.timestamp - first.timestamp).total_seconds()

        is_leaking = False
        reasons = []

        # Threshold Rule 1: Exceeding absolute limit
        if latest.rss_mb >= self.alert_mb:
            is_leaking = True
            reasons.append(f"Memory ({latest.rss_mb:.2f} MB) exceeded ceiling ({self.alert_mb:.2f} MB)")

        # Threshold Rule 2: Sustained growth over the snapshot window
        if len(self.snapshots) >= self.window_size and net_growth >= self.growth_threshold_mb:
            is_leaking = True
            reasons.append(
                f"Sustained growth detected: +{net_growth:.2f} MB over last {self.window_size} samples"
            )

        reason_str = " | ".join(reasons) if reasons else "Memory levels nominal"

        return LeakReport(
            pid=self.pid,
            duration_seconds=duration,
            start_rss_mb=first.rss_mb,
            current_rss_mb=latest.rss_mb,
            peak_rss_mb=peak_rss,
            net_growth_mb=net_growth,
            is_leaking=is_leaking,
            warning_reason=reason_str,
        )


def profile_internal_allocations(top_n: int = 5) -> None:
    """Uses tracemalloc to print top internal allocations for self-monitoring."""
    if not tracemalloc.is_tracing():
        logger.warning("tracemalloc is not active. Starting tracemalloc...")
        tracemalloc.start()
        return

    snapshot = tracemalloc.take_snapshot()
    top_stats = snapshot.statistics("lineno")

    logger.info("--- Top %d Memory Allocations (tracemalloc) ---", top_n)
    for index, stat in enumerate(top_stats[:top_n], 1):
        frame = stat.traceback[0]
        logger.info(
            "#%d: %s:%s: %.2f KB",
            index,
            frame.filename,
            frame.lineno,
            stat.size / 1024,
        )


# ============================================================================
# Main Loop & CLI
# ============================================================================
def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description="Monitor a Python or system process for memory leaks."
    )
    parser.add_argument(
        "--pid",
        type=int,
        default=os.getpid(),
        help="PID of the process to monitor (defaults to self PID).",
    )
    parser.add_argument(
        "--interval",
        type=float,
        default=DEFAULT_INTERVAL,
        help=f"Polling interval in seconds (default: {DEFAULT_INTERVAL}).",
    )
    parser.add_argument(
        "--alert-mb",
        type=float,
        default=DEFAULT_ALERT_MB,
        help=f"Absolute memory threshold in MB to trigger alert (default: {DEFAULT_ALERT_MB}).",
    )
    parser.add_argument(
        "--growth-threshold-mb",
        type=float,
        default=DEFAULT_GROWTH_THRESHOLD_MB,
        help=f"Windowed growth threshold in MB to trigger leak alert (default: {DEFAULT_GROWTH_THRESHOLD_MB}).",
    )
    parser.add_argument(
        "--enable-tracemalloc",
        action="store_true",
        help="Enable internal tracemalloc profiling (only useful when monitoring self).",
    )
    return parser.parse_args()


def main() -> None:
    args = parse_args()

    if args.enable_tracemalloc:
        tracemalloc.start()
        logger.info("tracemalloc tracking enabled.")

    logger.info("Starting Memory Leak Detector for PID: %d", args.pid)
    logger.info("Polling every %.1fs | Alert Ceiling: %.1f MB | Window Growth Limit: %.1f MB",
                args.interval, args.alert_mb, args.growth_threshold_mb)

    detector = MemoryLeakDetector(
        pid=args.pid,
        alert_mb=args.alert_mb,
        growth_threshold_mb=args.growth_threshold_mb,
    )

    try:
        while True:
            current = detector.capture_snapshot()
            report = detector.analyze()

            log_msg = (
                f"PID: {report.pid} | RSS: {current.rss_mb:.2f} MB | "
                f"VMS: {current.vms_mb:.2f} MB | Threads: {current.num_threads} | "
                f"Net Growth: {report.net_growth_mb:+.2f} MB"
            )

            if report.is_leaking:
                logger.warning("ALERT: %s | Details: %s", log_msg, report.warning_reason)
                if args.enable_tracemalloc and args.pid == os.getpid():
                    profile_internal_allocations()
            else:
                logger.info(log_msg)

            time.sleep(args.interval)

    except KeyboardInterrupt:
        logger.info("Monitoring stopped by user.")
    except psutil.NoSuchProcess:
        logger.error("Monitored process (PID %d) terminated unexpectedly.", args.pid)
        sys.exit(1)


if __name__ == "__main__":
    main()

Configuration

The script accepts CLI arguments, which fall back to environment variables if omitted:

Environment Variable CLI Argument Default Description
MEM_LEAK_INTERVAL --interval 5.0 Sampling rate in seconds between checks.
MEM_LEAK_ALERT_MB --alert-mb 500.0 Absolute RSS memory limit (in MB) above which an alert triggers.
MEM_LEAK_GROWTH_THRESHOLD --growth-threshold-mb 50.0 Maximum allowed net RSS growth (in MB) across the sliding window.
MEM_LEAK_WINDOW_SIZE N/A 12 Number of historical samples evaluated for trend analysis.
MEM_LEAK_LOG_LEVEL N/A INFO Logging verbosity (DEBUG, INFO, WARNING, ERROR).

Usage Examples

1. Monitor an External Background Process

To monitor a running service (e.g., Gunicorn worker or Celery node with PID 41205), checking every 10 seconds and alerting if memory breaches 300 MB or grows by more than 20 MB over the evaluation window:

python3 mem_leak.py --pid 41205 --interval 10 --alert-mb 300 --growth-threshold-mb 20

Output:

2023-10-27 10:00:00,102 [INFO] Starting Memory Leak Detector for PID: 41205
2023-10-27 10:00:00,103 [INFO] Polling every 10.0s | Alert Ceiling: 300.0 MB | Window Growth Limit: 20.0 MB
2023-10-27 10:00:10,105 [INFO] PID: 41205 | RSS: 142.30 MB | VMS: 812.10 MB | Threads: 4 | Net Growth: +0.00 MB
2023-10-27 10:00:20,108 [INFO] PID: 41205 | RSS: 165.80 MB | VMS: 835.60 MB | Threads: 4 | Net Growth: +23.50 MB
2023-10-27 10:00:30,111 [WARNING] ALERT: PID: 41205 | RSS: 168.10 MB | VMS: 838.00 MB | Threads: 4 | Net Growth: +25.80 MB | Details: Sustained growth detected: +25.80 MB over last 12 samples

2. Monitoring via Environment Variables in Containerized Environments

Configure thresholds via environment variables inside a Docker container or Kubernetes sidecar:

export MEM_LEAK_INTERVAL="2.0"
export MEM_LEAK_ALERT_MB="256.0"
export MEM_LEAK_LOG_LEVEL="DEBUG"

python3 mem_leak.py --pid 1

3. Self-Monitoring Test Mode with tracemalloc Active

Test the script against its own process while enabling internal object allocation profiling:

python3 mem_leak.py --alert-mb 10 --growth-threshold-mb 1 --enable-tracemalloc

Error Handling

  • Missing Target Process (psutil.NoSuchProcess): Handled gracefully at startup and during the main sampling loop. If the target process terminates or crashes, the script logs an error and exits cleanly without throwing an unhandled exception trace.
  • Insufficient Privileges (psutil.AccessDenied): Occurs when trying to monitor processes owned by root or another user without elevated rights. The script catches access errors during initialization and when accessing process attributes (like file descriptors), falling back to safe defaults or exiting with clear guidance.
  • Process Termination Mid-Poll: Handled within the polling loop. If a target process shuts down between sampling intervals, the loop breaks gracefully.
  • Keyboard Interrupts: Captures SIGINT (Ctrl+C) cleanly to shut down monitoring loops without displaying noisy Python stack traces.

Extending Ideas

  1. Webhook and Slack Integration: Modify the analyze() block to push JSON payloads to a Slack/Teams Webhook or PagerDuty API when report.is_leaking is true.
  2. Automated Memory Dumps: Automatically trigger a heap dump (e.g., using guppy3 or pympler) or call gc.collect() via IPC when a leak condition is met.
  3. Prometheus Metrics Exporter: Expose memory growth rates, thread counts, and leak status on an HTTP endpoint (/metrics) using prometheus_client for integration with Grafana.
  4. Linear Regression Analysis: Replace simple snapshot diffing with standard polynomial or linear regression (numpy.polyfit or standard math) to calculate exact MB/hour growth slopes and eliminate false positives caused by periodic garbage collection cycles.
// RELATED_SCRIPTS
← Back to DevLab valtersit.com