Утилита для мониторинга потребления CPU/RAM процессом и его дочерними процессами

procmon.py

#!/usr/bin/env python
import argparse
import subprocess
import sys
import time
import typing

import matplotlib.pyplot as plt
import psutil

# Константы для цветов
COLOR_CPU = "#F44336"  # Красный
COLOR_RAM = "#9C27B0"  # Фиолетовый
COLOR_TEXT = "#424242"  # Темно-серый
COLOR_BACKGROUND = "#FAFAFA"  # Светло-серый фон
COLOR_AXES_BACKGROUND = "#FFFFFF"  # Белый фон для осей
COLOR_SPINES = "#BDBDBD"  # Серый цвет для рамок
COLOR_OVERLAY_TEXT = "#212121"  # Темно-серый для текста поверх графиков


def monitor_process(
    proc: subprocess.Popen, interval: float
) -> tuple[list[float], list[float], list[float]]:
    cpu_usage: list[float] = []
    ram_usage: list[float] = []
    timestamps: list[float] = []
    process = psutil.Process(proc.pid)
    st = time.monotonic()
    try:
        while proc.poll() is None:
            cpu: float = (
                process.cpu_percent(interval=interval) / psutil.cpu_count()
            )
            ram: float = process.memory_info().rss / (1024 * 1024)

            cpu_usage.append(cpu)
            ram_usage.append(ram)
            timestamps.append(time.monotonic() - st)

            for child in process.children(recursive=True):
                try:
                    cpu += child.cpu_percent(interval=0) / psutil.cpu_count()
                    ram += child.memory_info().rss / (1024 * 1024)
                except psutil.NoSuchProcess:
                    continue

            time.sleep(interval)

    except psutil.NoSuchProcess:
        pass

    return cpu_usage, ram_usage, timestamps


def plot_usage(
    cpu_usage: list[float],
    ram_usage: list[float],
    timestamps: list[float],
    width: int,
    height: int,
) -> None:
    fig, ax1 = plt.subplots(figsize=(width / 100, height / 100))

    ax2 = ax1.twinx()
    ax1.set_xlabel("Time (s)", color=COLOR_TEXT)
    ax1.set_ylabel("CPU Usage (%)", color=COLOR_CPU)
    ax2.set_ylabel("RAM Usage (MB)", color=COLOR_RAM)

    fig.patch.set_facecolor(COLOR_BACKGROUND)
    ax1.set_facecolor(COLOR_AXES_BACKGROUND)
    ax2.set_facecolor(COLOR_AXES_BACKGROUND)

    for spine in ax1.spines.values():
        spine.set_edgecolor(COLOR_SPINES)

    for spine in ax2.spines.values():
        spine.set_edgecolor(COLOR_SPINES)

    ax1.plot(
        timestamps,
        cpu_usage,
        color=COLOR_CPU,
    )

    ax2.plot(
        timestamps,
        ram_usage,
        color=COLOR_RAM,
    )

    cpu_text = f"CPU: min: {min(cpu_usage):.2f}%, avg: {sum(cpu_usage)/len(cpu_usage):.2f}%, max: {max(cpu_usage):.2f}%"
    ram_text = f"RAM: min: {min(ram_usage):.2f}MB, avg: {sum(ram_usage)/len(ram_usage):.2f}MB, max: {max(ram_usage):.2f}MB"

    plt.figtext(0.5, 0.95, cpu_text, ha="center", fontsize=10, color=COLOR_CPU)
    plt.figtext(0.5, 0.05, ram_text, ha="center", fontsize=10, color=COLOR_RAM)

    plt.tight_layout(rect=[0, 0.05, 1, 0.95])
    plt.show()


def run_command(args: argparse.Namespace) -> None:
    proc: subprocess.Popen = subprocess.Popen(args.command)
    cpu_usage, ram_usage, timestamps = monitor_process(proc, args.interval)
    proc.wait()

    plot_usage(cpu_usage, ram_usage, timestamps, args.width, args.height)


def main(argv: typing.Sequence[str] | None = None) -> int | None:
    parser: argparse.ArgumentParser = argparse.ArgumentParser(
        description="Monitor CPU and RAM usage of a process and its subprocesses in cyberpunk style."
    )
    parser.add_argument(
        "--width", type=int, default=960, help="Width of the graph in pixels"
    )
    parser.add_argument(
        "--height", type=int, default=640, help="Height of the graph in pixels"
    )
    parser.add_argument(
        "--interval",
        type=float,
        default=0.2,
        help="Interval for checking CPU and RAM usage in seconds",
    )

    subparsers = parser.add_subparsers(dest="command")

    run_parser = subparsers.add_parser(
        "run", help="Run a command and monitor its resource usage"
    )
    run_parser.add_argument(
        "command", nargs=argparse.REMAINDER, help="Command to run"
    )
    run_parser.set_defaults(func=run_command)
    args = parser.parse_args(argv)

    if hasattr(args, "func"):
        return args.func(args)

    parser.print_help()
    return 1


if __name__ == "__main__":
    sys.exit(main(sys.argv[1:]))

./procmon.py run vim procmon.py