Tune-A-Video-Training-UI / app_system_monitor.py
hysts's picture
hysts HF staff
Add system monitor
400839c
raw history blame
No virus
2.95 kB
#!/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()