app_system_monitor.py

#!/usr/bin/env python

from __future__ import annotations

import collections

import gradio as gr
import nvitop
import pandas as pd
import plotly.express as px
import psutil


class SystemMonitor:
    MAX_SIZE = 61

    def __init__(self):
        self.devices = nvitop.Device.all()
        self.cpu_memory_usage = collections.deque(
            [0 for _ in range(self.MAX_SIZE)], maxlen=self.MAX_SIZE)
        self.cpu_memory_usage_str = ''
        self.gpu_memory_usage = collections.deque(
            [0 for _ in range(self.MAX_SIZE)], maxlen=self.MAX_SIZE)
        self.gpu_util = collections.deque([0 for _ in range(self.MAX_SIZE)],
                                          maxlen=self.MAX_SIZE)
        self.gpu_memory_usage_str = ''
        self.gpu_util_str = ''

    def update(self) -> None:
        self.update_cpu()
        self.update_gpu()

    def update_cpu(self) -> None:
        memory = psutil.virtual_memory()
        self.cpu_memory_usage.append(memory.percent)
        self.cpu_memory_usage_str = f'{memory.used / 1024**3:0.2f}GiB / {memory.total / 1024**3:0.2f}GiB ({memory.percent}%)'

    def update_gpu(self) -> None:
        if not self.devices:
            return
        device = self.devices[0]
        self.gpu_memory_usage.append(device.memory_percent())
        self.gpu_util.append(device.gpu_utilization())
        self.gpu_memory_usage_str = f'{device.memory_usage()} ({device.memory_percent()}%)'
        self.gpu_util_str = f'{device.gpu_utilization()}%'

    def get_json(self) -> dict[str, str]:
        return {
            'CPU memory usage': self.cpu_memory_usage_str,
            'GPU memory usage': self.gpu_memory_usage_str,
            'GPU Util': self.gpu_util_str,
        }

    def get_graph_data(self) -> dict[str, list[int | float]]:
        return {
            'index': list(range(-self.MAX_SIZE + 1, 1)),
            'CPU memory usage': self.cpu_memory_usage,
            'GPU memory usage': self.gpu_memory_usage,
            'GPU Util': self.gpu_util,
        }

    def get_graph(self):
        df = pd.DataFrame(self.get_graph_data())
        return px.line(df,
                       x='index',
                       y=[
                           'CPU memory usage',
                           'GPU memory usage',
                           'GPU Util',
                       ],
                       range_y=[-5,
                                105]).update_layout(xaxis_title='Time',
                                                    yaxis_title='Percentage')


def create_monitor_demo() -> gr.Blocks:
    monitor = SystemMonitor()
    with gr.Blocks() as demo:
        gr.JSON(value=monitor.update, every=1, visible=False)
        gr.JSON(value=monitor.get_json, show_label=False, every=1)
        gr.Plot(value=monitor.get_graph, show_label=False, every=1)
    return demo


if __name__ == '__main__':
    demo = create_monitor_demo()
    demo.queue(api_open=False).launch()