|
| 1 | +"""A Python script to monitor system resources for a given PID and optionally create |
| 2 | +a py-spy profiler report.""" |
| 3 | + |
| 4 | +import psutil |
| 5 | +import pynvml |
| 6 | +import time |
| 7 | +import subprocess |
| 8 | +import click |
| 9 | +import threading |
| 10 | + |
| 11 | + |
| 12 | +@click.command() |
| 13 | +@click.option("--pid", type=int, required=True, help="Process ID of the Python script") |
| 14 | +@click.option("--interval", type=int, default=2, help="Monitoring interval (seconds)") |
| 15 | +@click.option( |
| 16 | + "--duration", type=int, default=30, help="Total monitoring duration (seconds)" |
| 17 | +) |
| 18 | +@click.option( |
| 19 | + "--output", |
| 20 | + type=str, |
| 21 | + default=None, |
| 22 | + help="File to save system resource logs (optional)", |
| 23 | +) |
| 24 | +@click.option("--spy", is_flag=True, help="Enable py-spy profiling") |
| 25 | +@click.option( |
| 26 | + "--spy-output", type=str, default="pyspy_profile.svg", help="Py-Spy output file" |
| 27 | +) |
| 28 | +def monitor_resources( |
| 29 | + pid: int, interval: int, duration: int, output: str, spy: bool, spy_output: str |
| 30 | +): |
| 31 | + """Monitor system resources for a given PID and optionally create a py-spy profiler |
| 32 | + report. |
| 33 | +
|
| 34 | + Args: |
| 35 | + pid (int): Process ID of the Python script. |
| 36 | + interval (int): Monitoring interval in seconds. |
| 37 | + duration (int): Total monitoring duration in seconds. |
| 38 | + output (str): File to save logs (optional). |
| 39 | + spy (bool): Enable py-spy profiling. |
| 40 | + spy_output (str): Py-Spy output file. |
| 41 | + """ |
| 42 | + if not psutil.pid_exists(pid): |
| 43 | + click.echo(click.style(f"Error: Process with PID {pid} not found.", fg="red")) |
| 44 | + return |
| 45 | + |
| 46 | + click.echo( |
| 47 | + click.style( |
| 48 | + f"Monitoring system resources for PID {pid} for {duration} seconds...", |
| 49 | + fg="green", |
| 50 | + ) |
| 51 | + ) |
| 52 | + |
| 53 | + def run_py_spy(): |
| 54 | + """Run py-spy profiler for deep profiling.""" |
| 55 | + click.echo(click.style("Running py-spy for deep profiling...", fg="green")) |
| 56 | + spy_cmd = f"py-spy record -o {spy_output} --pid {pid} --duration {duration}" |
| 57 | + try: |
| 58 | + subprocess.run( |
| 59 | + spy_cmd, shell=True, check=True, capture_output=True, text=True |
| 60 | + ) |
| 61 | + click.echo( |
| 62 | + click.style(f"Py-Spy flame graph saved to {spy_output}", fg="green") |
| 63 | + ) |
| 64 | + except subprocess.CalledProcessError as e: |
| 65 | + click.echo(click.style(f"Error running py-spy: {e.stderr}", fg="red")) |
| 66 | + |
| 67 | + # Start py-spy profiling in a separate thread if enabled. |
| 68 | + if spy: |
| 69 | + spy_thread = threading.Thread(target=run_py_spy) |
| 70 | + spy_thread.start() |
| 71 | + |
| 72 | + # Start main resources monitoring loop. |
| 73 | + pynvml.nvmlInit() |
| 74 | + end_time = time.time() + duration |
| 75 | + logs, cpu_usages, ram_usages, gpu_usages, vram_usages = [], [], [], [], [] |
| 76 | + while time.time() < end_time: |
| 77 | + try: |
| 78 | + # General system usage. |
| 79 | + process = psutil.Process(pid) |
| 80 | + cpu_usage = process.cpu_percent(interval=interval) |
| 81 | + ram_usage = process.memory_info().rss / (1024 * 1024) # MB |
| 82 | + |
| 83 | + # GPU usage. |
| 84 | + process_gpu_usage = 0 |
| 85 | + process_vram_usage = 0 |
| 86 | + device_count = pynvml.nvmlDeviceGetCount() |
| 87 | + for i in range(device_count): |
| 88 | + handle = pynvml.nvmlDeviceGetHandleByIndex(i) |
| 89 | + processes = pynvml.nvmlDeviceGetComputeRunningProcesses(handle) |
| 90 | + for proc_info in processes: |
| 91 | + if proc_info.pid == pid: |
| 92 | + process_gpu_usage = pynvml.nvmlDeviceGetUtilizationRates( |
| 93 | + handle |
| 94 | + ).gpu |
| 95 | + process_vram_usage = proc_info.usedGpuMemory / ( |
| 96 | + 1024 * 1024 |
| 97 | + ) # MB |
| 98 | + break |
| 99 | + |
| 100 | + # Collect and log resource usage. |
| 101 | + log_entry = f"CPU: {cpu_usage:.2f}%, RAM: {ram_usage:.2f}MB, GPU: {process_gpu_usage:.2f}%, VRAM: {process_vram_usage:.2f}MB" |
| 102 | + click.echo(log_entry) |
| 103 | + logs.append(log_entry) |
| 104 | + cpu_usages.append(cpu_usage) |
| 105 | + ram_usages.append(ram_usage) |
| 106 | + gpu_usages.append(process_gpu_usage) |
| 107 | + vram_usages.append(process_vram_usage) |
| 108 | + except psutil.NoSuchProcess: |
| 109 | + click.echo(click.style("Error: Process terminated!")) |
| 110 | + break |
| 111 | + |
| 112 | + pynvml.nvmlShutdown() |
| 113 | + |
| 114 | + # Calculate and log average resource usage. |
| 115 | + avg_cpu_usage = sum(cpu_usages) / len(cpu_usages) if cpu_usages else 0 |
| 116 | + avg_ram_usage = sum(ram_usages) / len(ram_usages) if ram_usages else 0 |
| 117 | + avg_gpu_usage = sum(gpu_usages) / len(gpu_usages) if gpu_usages else 0 |
| 118 | + avg_vram_usage = sum(vram_usages) / len(vram_usages) if vram_usages else 0 |
| 119 | + avg_log_entry = f"AVERAGE - CPU: {avg_cpu_usage:.2f}%, RAM: {avg_ram_usage:.2f}MB, GPU: {avg_gpu_usage:.2f}%, VRAM: {avg_vram_usage:.2f}MB" |
| 120 | + click.echo(avg_log_entry) |
| 121 | + logs.append(avg_log_entry) |
| 122 | + |
| 123 | + # Save logs if output file is provided. |
| 124 | + if output: |
| 125 | + with open(output, "w") as f: |
| 126 | + f.write("\n".join(logs)) |
| 127 | + click.echo(click.style(f"Logs saved to {output}", fg="green")) |
| 128 | + |
| 129 | + # Wait for py-spy thread to finish if it was started. |
| 130 | + if spy: |
| 131 | + spy_thread.join() |
| 132 | + |
| 133 | + |
| 134 | +if __name__ == "__main__": |
| 135 | + monitor_resources() |
0 commit comments