diff --git a/src/ert/config/model_config.py b/src/ert/config/model_config.py index 968bfee8758..11c9809eb76 100644 --- a/src/ert/config/model_config.py +++ b/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", "", 1) format_string = format_string.replace("%d", "", 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 diff --git a/src/ert/gui/simulation/run_dialog.py b/src/ert/gui/simulation/run_dialog.py index dee4e941a0b..c0a06ff404b 100644 --- a/src/ert/gui/simulation/run_dialog.py +++ b/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)}" diff --git a/src/ert/gui/simulation/view/__init__.py b/src/ert/gui/simulation/view/__init__.py index f905ab18f24..200d686ce21 100644 --- a/src/ert/gui/simulation/view/__init__.py +++ b/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"] diff --git a/src/ert/gui/simulation/view/disk_space_widget.py b/src/ert/gui/simulation/view/disk_space_widget.py new file mode 100644 index 00000000000..e1324d709b6 --- /dev/null +++ b/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) diff --git a/src/ert/shared/status/utils.py b/src/ert/shared/status/utils.py index 99c905f542e..1b74108551e 100644 --- a/src/ert/shared/status/utils.py +++ b/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 diff --git a/tests/ert/unit_tests/gui/simulation/test_run_dialog.py b/tests/ert/unit_tests/gui/simulation/test_run_dialog.py index 686939f032a..39f8d629e6d 100644 --- a/tests/ert/unit_tests/gui/simulation/test_run_dialog.py +++ b/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 diff --git a/tests/ert/unit_tests/gui/simulation/view/test_disk_space_widget.py b/tests/ert/unit_tests/gui/simulation/view/test_disk_space_widget.py new file mode 100644 index 00000000000..b127114a2ef --- /dev/null +++ b/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()