Skip to content

Commit

Permalink
Add disk space status to gui
Browse files Browse the repository at this point in the history
  • Loading branch information
larsevj authored and eivindjahren committed Jan 15, 2025
1 parent 1119c50 commit 77b2851
Show file tree
Hide file tree
Showing 8 changed files with 142 additions and 16 deletions.
16 changes: 3 additions & 13 deletions src/ert/config/model_config.py
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -94,15 +93,15 @@ 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 (
percentage_used > FULL_DISK_PERCENTAGE_THRESHOLD
and free_space < MINIMUM_BYTES_LEFT_ON_DISK_THRESHOLD
):
msg = (
f"Low disk space: {byte_with_unit(free_space)} free on {mount_dir !s}."
f"Low disk space: {byte_with_unit(free_space)} free on {mount_dir!s}."
" Consider freeing up some space to ensure successful simulation runs."
)
ConfigWarning.warn(msg)
Expand Down Expand Up @@ -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
10 changes: 9 additions & 1 deletion src/ert/gui/simulation/run_dialog.py
Original file line number Diff line number Diff line change
Expand Up @@ -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}"

Expand Down Expand Up @@ -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")
Expand All @@ -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)
Expand Down Expand Up @@ -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)}"
Expand Down
3 changes: 2 additions & 1 deletion src/ert/gui/simulation/view/__init__.py
Original file line number Diff line number Diff line change
@@ -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"]
78 changes: 78 additions & 0 deletions src/ert/gui/simulation/view/disk_space_widget.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,78 @@
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:

Check failure on line 40 in src/ert/gui/simulation/view/disk_space_widget.py

View workflow job for this annotation

GitHub Actions / type-checking (3.12)

Missing return statement
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)

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)
10 changes: 10 additions & 0 deletions src/ert/shared/status/utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,7 @@
import os
import resource
import sys
from pathlib import Path


def byte_with_unit(byte_count: float) -> str:
Expand Down Expand Up @@ -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
2 changes: 1 addition & 1 deletion tests/ert/unit_tests/config/test_model_config.py
Original file line number Diff line number Diff line change
Expand Up @@ -103,7 +103,7 @@ def test_warning_when_full_disk(
tmp_path, recwarn, total_space, used_space, to_warn, expected_warning
):
Path(tmp_path / "simulations").mkdir()
runpath = f"{tmp_path !s}/simulations/realization-%d/iter-%d"
runpath = f"{tmp_path!s}/simulations/realization-%d/iter-%d"
with patch(
"ert.config.model_config.shutil.disk_usage",
return_value=(total_space, used_space, total_space - used_space),
Expand Down
3 changes: 3 additions & 0 deletions tests/ert/unit_tests/gui/simulation/test_run_dialog.py
Original file line number Diff line number Diff line change
Expand Up @@ -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


Expand Down
36 changes: 36 additions & 0 deletions tests/ert/unit_tests/gui/simulation/view/test_disk_space_widget.py
Original file line number Diff line number Diff line change
@@ -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()

0 comments on commit 77b2851

Please sign in to comment.