Add disk space status to gui

src/ert/config/model_config.py

@@ -5,13 +5,12 @@
 import os.path
 import shutil
 from datetime import datetime
-from pathlib import Path
 from typing import no_type_check
 
 from pydantic import field_validator
 from pydantic.dataclasses import dataclass
 
-from ert.shared.status.utils import byte_with_unit
+from ert.shared.status.utils import byte_with_unit, get_mount_directory
 
 from .parsing import (
     ConfigDict,
@@ -94,7 +93,7 @@ def validate_runpath(cls, runpath_format_string: str) -> str:
             ConfigWarning.warn(msg)
             logger.warning(msg)
         with contextlib.suppress(Exception):
-            mount_dir = _get_mount_directory(runpath_format_string)
+            mount_dir = get_mount_directory(runpath_format_string)
             total_space, used_space, free_space = shutil.disk_usage(mount_dir)
             percentage_used = used_space / total_space
             if (
@@ -166,12 +165,3 @@ def _replace_runpath_format(format_string: str) -> str:
     format_string = format_string.replace("%d", "<IENS>", 1)
     format_string = format_string.replace("%d", "<ITER>", 1)
     return format_string
-
-
-def _get_mount_directory(runpath: str) -> Path:
-    path = Path(runpath).absolute()
-
-    while not path.is_mount():
-        path = path.parent
-
-    return path

src/ert/gui/simulation/run_dialog.py

@@ -64,11 +64,12 @@
     byte_with_unit,
     file_has_content,
     format_running_time,
+    get_mount_directory,
 )
 
 from ..find_ert_info import find_ert_info
 from .queue_emitter import QueueEmitter
-from .view import ProgressWidget, RealizationWidget, UpdateWidget
+from .view import DiskSpaceWidget, ProgressWidget, RealizationWidget, UpdateWidget
 
 _TOTAL_PROGRESS_TEMPLATE = "Total progress {total_progress}% — {iteration_label}"
 
@@ -220,6 +221,9 @@ def __init__(
 
         self.running_time = QLabel("")
         self.memory_usage = QLabel("")
+        self.disk_space = DiskSpaceWidget(
+            get_mount_directory(run_model.run_paths._runpath_format)
+        )
 
         self.kill_button = QPushButton("Terminate experiment")
         self.restart_button = QPushButton("Rerun failed")
@@ -245,6 +249,8 @@ def __init__(
         button_layout.addStretch()
         button_layout.addWidget(self.memory_usage)
         button_layout.addStretch()
+        button_layout.addWidget(self.disk_space)
+        button_layout.addStretch()
         button_layout.addWidget(self.copy_debug_info_button)
         button_layout.addWidget(self.kill_button)
         button_layout.addWidget(self.restart_button)
@@ -425,6 +431,8 @@ def _on_ticker(self) -> None:
 
         maximum_memory_usage = self._snapshot_model.root.max_memory_usage
 
+        self.disk_space.update_status()
+
         if maximum_memory_usage:
             self.memory_usage.setText(
                 f"Maximal realization memory usage: {byte_with_unit(maximum_memory_usage)}"

src/ert/gui/simulation/view/__init__.py

@@ -1,5 +1,6 @@
+from .disk_space_widget import DiskSpaceWidget
 from .progress_widget import ProgressWidget
 from .realization import RealizationWidget
 from .update import UpdateWidget
 
-__all__ = ["ProgressWidget", "RealizationWidget", "UpdateWidget"]
+__all__ = ["DiskSpaceWidget", "ProgressWidget", "RealizationWidget", "UpdateWidget"]

src/ert/gui/simulation/view/disk_space_widget.py

@@ -0,0 +1,79 @@
+import contextlib
+import shutil
+from pathlib import Path
+
+from qtpy.QtCore import Qt
+from qtpy.QtWidgets import QHBoxLayout, QLabel, QProgressBar, QWidget
+
+from ert.shared.status.utils import byte_with_unit
+
+CRITICAL_RED = "#e74c3c"
+WARNING_YELLOW = "#f1c40f"
+NORMAL_GREEN = "#2ecc71"
+
+
+class DiskSpaceWidget(QWidget):
+    def __init__(self, mount_path: Path, parent: QWidget | None = None) -> None:
+        super().__init__(parent)
+
+        self.mount_path = mount_path
+
+        layout = QHBoxLayout(self)
+        layout.setContentsMargins(0, 0, 0, 0)
+        layout.setSpacing(10)
+
+        # Text label
+        self.usage_label = QLabel(self)
+        self.space_left_label = QLabel(self)
+
+        # Progress bar
+        self.progress_bar = QProgressBar(self)
+        self.progress_bar.setRange(0, 100)
+        self.progress_bar.setTextVisible(True)
+        self.progress_bar.setFixedWidth(100)
+        self.progress_bar.setAlignment(Qt.AlignCenter)  # type: ignore
+
+        layout.addWidget(self.usage_label)
+        layout.addWidget(self.progress_bar)
+        layout.addWidget(self.space_left_label)
+
+    def _get_status(self) -> tuple[float, str] | None:
+        with contextlib.suppress(Exception):
+            disk_info = shutil.disk_usage(self.mount_path)
+            percentage_used = (disk_info.used / disk_info.total) * 100
+            return percentage_used, byte_with_unit(disk_info.free)
+        return None
+
+    def update_status(self) -> None:
+        """Update both the label and progress bar with current disk usage"""
+        if (disk_info := self._get_status()) is not None:
+            usage, space_left = disk_info
+            self.usage_label.setText("Disk space runpath:")
+            self.progress_bar.setValue(int(usage))
+            self.progress_bar.setFormat(f"{usage:.1f}%")
+
+            # Set color based on usage threshold
+            if usage >= 90:
+                color = CRITICAL_RED
+            elif usage >= 70:
+                color = WARNING_YELLOW
+            else:
+                color = NORMAL_GREEN
+
+            self.progress_bar.setStyleSheet(f"""
+                QProgressBar {{
+                    border: 1px solid #ccc;
+                    border-radius: 2px;
+                    text-align: center;
+                }}
+
+                QProgressBar::chunk {{
+                    background-color: {color};
+                }}
+            """)
+
+            self.space_left_label.setText(f"{space_left} free")
+
+            self.setVisible(True)
+        else:
+            self.setVisible(False)

src/ert/shared/status/utils.py

@@ -2,6 +2,7 @@
 import os
 import resource
 import sys
+from pathlib import Path
 
 
 def byte_with_unit(byte_count: float) -> str:
@@ -82,3 +83,12 @@ def get_ert_memory_usage() -> int:
         rss_scale = 1000
 
     return usage.ru_maxrss // rss_scale
+
+
+def get_mount_directory(runpath: str) -> Path:
+    path = Path(runpath).absolute()
+
+    while not path.is_mount():
+        path = path.parent
+
+    return path

tests/ert/unit_tests/gui/simulation/test_run_dialog.py

@@ -48,6 +48,9 @@ def run_model():
     run_model.format_error.return_value = ""
     run_model.get_runtime.return_value = 1
     run_model.support_restart = True
+    run_paths_mock = MagicMock()
+    run_paths_mock._runpath_format = "/"
+    run_model.run_paths = run_paths_mock
     return run_model
 
 

tests/ert/unit_tests/gui/simulation/view/test_disk_space_widget.py

@@ -0,0 +1,36 @@
+from unittest.mock import MagicMock
+
+from ert.gui.simulation.view import DiskSpaceWidget
+from ert.gui.simulation.view.disk_space_widget import (
+    CRITICAL_RED,
+    NORMAL_GREEN,
+    WARNING_YELLOW,
+)
+
+
+def test_disk_space_widget(qtbot):
+    disk_space_widget = DiskSpaceWidget("/tmp")
+    qtbot.addWidget(disk_space_widget)
+    disk_space_widget._get_status = MagicMock()
+
+    disk_space_widget._get_status.return_value = (20.0, "2 jiggabytes")
+    disk_space_widget.update_status()
+    assert disk_space_widget.space_left_label.text() == "2 jiggabytes free"
+    assert disk_space_widget.progress_bar.value() == 20
+    assert NORMAL_GREEN in disk_space_widget.progress_bar.styleSheet()
+
+    disk_space_widget._get_status.return_value = (88.0, "2mb")
+    disk_space_widget.update_status()
+    assert disk_space_widget.space_left_label.text() == "2mb free"
+    assert disk_space_widget.progress_bar.value() == 88
+    assert WARNING_YELLOW in disk_space_widget.progress_bar.styleSheet()
+
+    disk_space_widget._get_status.return_value = (99.9, "2 bytes")
+    disk_space_widget.update_status()
+    assert disk_space_widget.space_left_label.text() == "2 bytes free"
+    assert disk_space_widget.progress_bar.value() == 99
+    assert CRITICAL_RED in disk_space_widget.progress_bar.styleSheet()
+
+    disk_space_widget._get_status.return_value = None
+    disk_space_widget.update_status()
+    assert not disk_space_widget.isVisible()