|
| 1 | +"""Disk space monitoring utilities for checkpoint saves.""" |
| 2 | + |
| 3 | +import logging |
| 4 | +import re |
| 5 | +import shutil |
| 6 | +import subprocess |
| 7 | +import time |
| 8 | +from enum import Enum |
| 9 | +from pathlib import Path |
| 10 | +from typing import Any, Optional |
| 11 | + |
| 12 | +logger = logging.getLogger("DiskSpaceMonitor") |
| 13 | + |
| 14 | + |
| 15 | +class DiskLowAction(str, Enum): |
| 16 | + """Action to take when disk space is below threshold.""" |
| 17 | + |
| 18 | + STOP = "stop" |
| 19 | + WAIT = "wait" |
| 20 | + SCRIPT = "script" |
| 21 | + |
| 22 | + @classmethod |
| 23 | + def from_raw(cls, raw_value: Any) -> "DiskLowAction": |
| 24 | + """Convert a raw config/CLI value to DiskLowAction enum.""" |
| 25 | + if isinstance(raw_value, cls): |
| 26 | + return raw_value |
| 27 | + if raw_value in (None, "", "None"): |
| 28 | + return cls.STOP |
| 29 | + normalized = str(raw_value).strip().lower() |
| 30 | + try: |
| 31 | + return cls(normalized) |
| 32 | + except ValueError as exc: |
| 33 | + valid_values = ", ".join(member.value for member in cls) |
| 34 | + raise ValueError(f"Unsupported disk_low_action '{raw_value}'. Expected one of: {valid_values}") from exc |
| 35 | + |
| 36 | + |
| 37 | +def parse_size_threshold(threshold_str: Optional[str]) -> Optional[int]: |
| 38 | + """ |
| 39 | + Parse a human-readable size string into bytes. |
| 40 | +
|
| 41 | + Args: |
| 42 | + threshold_str: Size string like "100G", "50M", "1T", "500K", or plain bytes. |
| 43 | + Returns None if threshold_str is None/empty (feature disabled). |
| 44 | +
|
| 45 | + Returns: |
| 46 | + Size in bytes, or None if feature is disabled. |
| 47 | +
|
| 48 | + Raises: |
| 49 | + ValueError: If the format is invalid. |
| 50 | + """ |
| 51 | + if threshold_str in (None, "", "None"): |
| 52 | + return None |
| 53 | + |
| 54 | + threshold_str = str(threshold_str).strip().upper() |
| 55 | + |
| 56 | + match = re.match(r"^(\d+(?:\.\d+)?)\s*([KMGT]?)B?$", threshold_str) |
| 57 | + if not match: |
| 58 | + raise ValueError( |
| 59 | + f"Invalid disk_low_threshold format: '{threshold_str}'. " |
| 60 | + "Expected format like '100G', '50M', '1T', '500K', or plain bytes." |
| 61 | + ) |
| 62 | + |
| 63 | + value = float(match.group(1)) |
| 64 | + unit = match.group(2) |
| 65 | + |
| 66 | + multipliers = { |
| 67 | + "": 1, |
| 68 | + "K": 1024, |
| 69 | + "M": 1024**2, |
| 70 | + "G": 1024**3, |
| 71 | + "T": 1024**4, |
| 72 | + } |
| 73 | + |
| 74 | + return int(value * multipliers[unit]) |
| 75 | + |
| 76 | + |
| 77 | +def get_available_disk_space(path: str) -> int: |
| 78 | + """ |
| 79 | + Return available disk space in bytes for the filesystem containing path. |
| 80 | +
|
| 81 | + If the path doesn't exist, traverses parent directories to find an existing one. |
| 82 | + """ |
| 83 | + resolved_path = Path(path).resolve() |
| 84 | + while not resolved_path.exists() and resolved_path.parent != resolved_path: |
| 85 | + resolved_path = resolved_path.parent |
| 86 | + |
| 87 | + usage = shutil.disk_usage(str(resolved_path)) |
| 88 | + return usage.free |
| 89 | + |
| 90 | + |
| 91 | +def _format_bytes(num_bytes: int) -> str: |
| 92 | + """Format bytes as human-readable string.""" |
| 93 | + value = float(num_bytes) |
| 94 | + for unit in ["B", "KB", "MB", "GB", "TB"]: |
| 95 | + if abs(value) < 1024.0: |
| 96 | + return f"{value:.1f}{unit}" |
| 97 | + value /= 1024.0 |
| 98 | + return f"{value:.1f}PB" |
| 99 | + |
| 100 | + |
| 101 | +def check_disk_space( |
| 102 | + output_dir: str, |
| 103 | + threshold_bytes: int, |
| 104 | + action: DiskLowAction, |
| 105 | + script_path: Optional[str] = None, |
| 106 | + check_interval: int = 30, |
| 107 | +) -> None: |
| 108 | + """ |
| 109 | + Check if available disk space is below threshold and take configured action. |
| 110 | +
|
| 111 | + Args: |
| 112 | + output_dir: Directory to check disk space for. |
| 113 | + threshold_bytes: Minimum required free space in bytes. |
| 114 | + action: Action to take when space is low. |
| 115 | + script_path: Path to cleanup script (required when action is SCRIPT). |
| 116 | + check_interval: Seconds between checks in WAIT mode. |
| 117 | +
|
| 118 | + Raises: |
| 119 | + RuntimeError: When action is STOP, or when SCRIPT fails, or when |
| 120 | + space remains low after SCRIPT execution. |
| 121 | + """ |
| 122 | + available = get_available_disk_space(output_dir) |
| 123 | + |
| 124 | + if available >= threshold_bytes: |
| 125 | + return |
| 126 | + |
| 127 | + available_human = _format_bytes(available) |
| 128 | + threshold_human = _format_bytes(threshold_bytes) |
| 129 | + |
| 130 | + if action == DiskLowAction.STOP: |
| 131 | + raise RuntimeError( |
| 132 | + f"Disk space critically low: {available_human} available, " f"threshold is {threshold_human}. Training stopped." |
| 133 | + ) |
| 134 | + |
| 135 | + elif action == DiskLowAction.WAIT: |
| 136 | + logger.warning( |
| 137 | + "Disk space low: %s available (threshold: %s). " "Waiting for space to become available...", |
| 138 | + available_human, |
| 139 | + threshold_human, |
| 140 | + ) |
| 141 | + while available < threshold_bytes: |
| 142 | + time.sleep(check_interval) |
| 143 | + available = get_available_disk_space(output_dir) |
| 144 | + logger.info( |
| 145 | + "Disk space recovered: %s available. Resuming training.", |
| 146 | + _format_bytes(available), |
| 147 | + ) |
| 148 | + |
| 149 | + elif action == DiskLowAction.SCRIPT: |
| 150 | + if not script_path: |
| 151 | + raise RuntimeError("disk_low_action is 'script' but no disk_low_script configured.") |
| 152 | + logger.warning( |
| 153 | + "Disk space low: %s available (threshold: %s). Running cleanup script: %s", |
| 154 | + available_human, |
| 155 | + threshold_human, |
| 156 | + script_path, |
| 157 | + ) |
| 158 | + try: |
| 159 | + subprocess.run([script_path], check=True) |
| 160 | + except subprocess.CalledProcessError as exc: |
| 161 | + raise RuntimeError(f"Disk cleanup script failed with exit code {exc.returncode}") from exc |
| 162 | + except FileNotFoundError as exc: |
| 163 | + raise RuntimeError(f"Disk cleanup script not found: {script_path}") from exc |
| 164 | + |
| 165 | + available = get_available_disk_space(output_dir) |
| 166 | + if available < threshold_bytes: |
| 167 | + raise RuntimeError( |
| 168 | + f"Disk space still low after cleanup script: " |
| 169 | + f"{_format_bytes(available)} available, threshold is {threshold_human}." |
| 170 | + ) |
| 171 | + logger.info("Disk cleanup script completed. %s now available.", _format_bytes(available)) |
0 commit comments