Add system monitor

Dockerfile

@@ -44,6 +44,8 @@ RUN pyenv install ${PYTHON_VERSION} && \
 RUN pip install --no-cache-dir -U torch==1.13.1 torchvision==0.14.1
 COPY --chown=1000 requirements.txt /tmp/requirements.txt
 RUN pip install --no-cache-dir -U -r /tmp/requirements.txt
+COPY --chown=1000 requirements-monitor.txt /tmp/requirements-monitor.txt
+RUN pip install --no-cache-dir -U -r /tmp/requirements-monitor.txt
 
 COPY --chown=1000 . ${HOME}/app
 RUN cd Tune-A-Video && patch -p1 < ../patch

app_system_monitor.py

@@ -0,0 +1,87 @@
+#!/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()

app_training.py

@@ -6,6 +6,7 @@ import os
 
 import gradio as gr
 
+from app_system_monitor import create_monitor_demo
 from constants import UploadTarget
 from inference import InferencePipeline
 from trainer import Trainer
@@ -13,6 +14,11 @@ from trainer import Trainer
 
 def create_training_demo(trainer: Trainer,
                          pipe: InferencePipeline | None = None) -> gr.Blocks:
+    def read_log() -> str:
+        with open(trainer.log_file) as f:
+            lines = f.readlines()
+        return ''.join(lines[-10:])
+
     hf_token = os.getenv('HF_TOKEN')
     with gr.Blocks() as demo:
         with gr.Row():
@@ -108,8 +114,14 @@ def create_training_demo(trainer: Trainer,
         run_button = gr.Button('Start Training')
 
         with gr.Box():
-            gr.Markdown('Output message')
-            output_message = gr.Markdown()
+            gr.Text(label='Log',
+                    value=read_log,
+                    lines=10,
+                    max_lines=10,
+                    every=1)
+            if not os.getenv('DISABLE_SYSTEM_MONITOR'):
+                with gr.Accordion(label='System info', open=False):
+                    create_monitor_demo()
 
         if pipe is not None:
             run_button.click(fn=pipe.clear)
@@ -136,8 +148,7 @@ def create_training_demo(trainer: Trainer,
                              upload_to,
                              remove_gpu_after_training,
                              input_token,
-                         ],
-                         outputs=output_message)
+                         ])
     return demo
 
 

requirements-monitor.txt

@@ -0,0 +1,4 @@
+nvitop==1.1.1
+pandas==2.0.0
+plotly==5.14.1
+psutil==5.9.4

trainer.py

@@ -32,6 +32,9 @@ class Trainer:
         self.checkpoint_dir = pathlib.Path('checkpoints')
         self.checkpoint_dir.mkdir(exist_ok=True)
 
+        self.log_file = pathlib.Path('log.txt')
+        self.log_file.touch(exist_ok=True)
+
     def download_base_model(self, base_model_id: str) -> str:
         model_dir = self.checkpoint_dir / base_model_id
         if not model_dir.exists():
@@ -72,7 +75,7 @@ class Trainer:
         upload_to: str,
         remove_gpu_after_training: bool,
         input_token: str,
-    ) -> str:
+    ) -> None:
         if SPACE_ID == ORIGINAL_SPACE_ID:
             raise gr.Error(
                 'This Space does not work on this Shared UI. Duplicate the Space and attribute a GPU'
@@ -134,15 +137,19 @@ class Trainer:
             OmegaConf.save(config, f)
 
         command = f'accelerate launch Tune-A-Video/train_tuneavideo.py --config {config_path}'
-        subprocess.run(shlex.split(command))
+        with open(self.log_file, 'w') as f:
+            subprocess.run(shlex.split(command),
+                           stdout=f,
+                           stderr=subprocess.STDOUT,
+                           text=True)
         save_model_card(save_dir=output_dir,
                         base_model=base_model,
                         training_prompt=training_prompt,
                         test_prompt=validation_prompt,
                         test_image_dir='samples')
 
-        message = 'Training completed!'
-        print(message)
+        with open(self.log_file, 'a') as f:
+            f.write('Training completed!\n')
 
         if upload_to_hub:
             upload_message = self.model_uploader.upload_model(
@@ -152,8 +159,8 @@ class Trainer:
                 private=use_private_repo,
                 delete_existing_repo=delete_existing_repo,
                 input_token=input_token)
-            print(upload_message)
-            message = message + '\n' + upload_message
+            with open(self.log_file, 'a') as f:
+                f.write(upload_message)
 
         if remove_gpu_after_training:
             space_id = os.getenv('SPACE_ID')
@@ -162,5 +169,3 @@ class Trainer:
                     token=self.hf_token if self.hf_token else input_token)
                 api.request_space_hardware(repo_id=space_id,
                                            hardware='cpu-basic')
-
-        return message