new: disk usage profiling

Author: evilsocket

dyana/cli.py

@@ -12,6 +12,7 @@
 from dyana.tracer.tracee import Tracer
 from dyana.view import (
     view_disk_events,
+    view_disk_usage,
     view_extra,
     view_gpus,
     view_header,
@@ -106,6 +107,7 @@ def summary(trace_path: pathlib.Path = typer.Option(help="Path to the trace file
     view_header(trace)
     view_ram(trace["run"])
     view_gpus(trace["run"])
+    view_disk_usage(trace["run"])
 
     view_process_executions(trace)
     view_network_events(trace)

dyana/loaders/base/dyana.py

@@ -1,4 +1,5 @@
 import resource
+import shutil
 import sys
 import typing as t
 from contextlib import contextmanager
@@ -8,6 +9,7 @@
 class Profiler:
     def __init__(self, gpu: bool = False):
         self._errors: dict[str, str] = {}
+        self._disk: dict[str, int] = {"start": get_disk_usage()}
         self._ram: dict[str, int] = {"start": get_peak_rss()}
         self._gpu: dict[str, list[dict[str, t.Any]]] = {"start": get_gpu_usage()} if gpu else {}
         self._imports_at_start = get_current_imports()
@@ -18,6 +20,9 @@ def track_memory(self, event: str) -> None:
         if self._gpu:
             self._gpu[event] = get_gpu_usage()
 
+    def track_disk(self, event: str) -> None:
+        self._disk[event] = get_disk_usage()
+
     def track_error(self, event: str, error: str) -> None:
         self._errors[event] = error
 
@@ -30,6 +35,7 @@ def as_dict(self) -> dict[str, t.Any]:
 
         as_dict: dict[str, t.Any] = {
             "ram": self._ram,
+            "disk": self._disk,
             "errors": self._errors,
             "extra": {"imports": imported},
         } | self._additionals
@@ -63,6 +69,14 @@ def capture_output() -> t.Generator[tuple[StringIO, StringIO], None, None]:
         sys.stderr = old_stderr
 
 
+def get_disk_usage() -> int:
+    """
+    Get the disk usage.
+    """
+    _, used, _ = shutil.disk_usage("/")
+    return used
+
+
 def get_peak_rss() -> int:
     """
     Get the peak RSS memory usage of the current process.

dyana/loaders/loader.py

@@ -32,6 +32,7 @@ class Run(BaseModel):
     errors: dict[str, str] | None = None
     ram: dict[str, int] | None = None
     gpu: dict[str, list[GpuDeviceUsage]] | None = None
+    disk: dict[str, int] | None = None
     stdout: str | None = None
     stderr: str | None = None
     exit_code: int | None = None
@@ -103,8 +104,8 @@ def __init__(
                 print(
                     f":whale: [bold]loader[/]: using image [green]{self.image.tags[0]}[/] [dim]({self.image.id})[/] ({self.platform})"
                 )
-            else:
-                print(f":whale: [bold]loader[/]: using image [green]{self.image.tags[0]}[/] [dim]({self.image.id})[/]")
+            # else:
+                # print(f":whale: [bold]loader[/]: using image [green]{self.image.tags[0]}[/] [dim]({self.image.id})[/]")
 
     def _reader_thread(self) -> None:
         if not self.container:
@@ -161,13 +162,15 @@ def run(self, allow_network: bool = False, allow_gpus: bool = True, allow_volume
                 else:
                     arguments.append(arg.value)
 
-        if allow_network:
-            print(
-                ":popcorn: [bold]loader[/]: [yellow]warning: allowing bridged network access to the model container[/]"
-            )
+        if self.settings and self.settings.network:
+            allow_network = True
+            print(":popcorn: [bold]loader[/]: [yellow]required bridged network access[/]")
+
+        elif allow_network:
+            print(":popcorn: [bold]loader[/]: [yellow]warning: allowing bridged network access to the container[/]")
 
         if allow_volume_write:
-            print(":popcorn: [bold]loader[/]: [yellow]warning: allowing volume write to the model container[/]")
+            print(":popcorn: [bold]loader[/]: [yellow]warning: allowing volume write to the container[/]")
 
         if arguments:
             print(f":popcorn: [bold]loader[/]: executing with arguments [dim]{arguments}[/] ...")

dyana/loaders/settings.py

@@ -19,6 +19,7 @@ class LoaderSettings(BaseModel):
     description: str
     build_args: dict[str, str] | None = None
     args: list[LoaderArgument] | None = None
+    network: bool | None = False
 
     def _parse_arg_name_from(self, name: str, args: list[str]) -> str | None:
         found_pre = False

dyana/view.py

@@ -131,6 +131,22 @@ def view_gpus(run: dict[str, t.Any]) -> None:
                     print()
 
 
+def view_disk_usage(run: dict[str, t.Any]) -> None:
+    disk = run["disk"]
+    if disk:
+        print("[bold yellow]Disk Usage:[/]")
+        disk_stages = list(disk.keys())
+        prev_stage = None
+        for stage in disk_stages:
+            if prev_stage is None:
+                print(f"  * {stage} : {sizeof_fmt(disk[stage])}")
+            else:
+                print(f"  * {stage} : {delta_fmt(disk[prev_stage], disk[stage])}")
+            prev_stage = stage
+
+        print()
+
+
 def view_process_executions(trace: dict[str, t.Any]) -> None:
     proc_execs = [event for event in trace["events"] if event["eventName"] == "sched_process_exec"]
     visualized = []