| """Resource guards: keep the search inside the machine's RAM and disk.
|
|
|
| Nothing here touches the GPU; the whole project is CPU and numpy only.
|
| """
|
| from __future__ import annotations
|
|
|
| import ctypes
|
| import os
|
| import shutil
|
| import sys
|
| from pathlib import Path
|
|
|
| MB = 1024 * 1024
|
| GB = 1024 * MB
|
|
|
| _JOB_HANDLE = None
|
|
|
|
|
| def free_disk_gb(path: str | Path) -> float:
|
| return shutil.disk_usage(str(path)).free / GB
|
|
|
|
|
| def dir_size_mb(path: str | Path) -> float:
|
| total = 0
|
| for root, _dirs, files in os.walk(str(path)):
|
| for name in files:
|
| try:
|
| total += os.path.getsize(os.path.join(root, name))
|
| except OSError:
|
| continue
|
| return total / MB
|
|
|
|
|
| def free_ram_mb() -> float | None:
|
| """Physical RAM currently available, or None if it cannot be read."""
|
| if sys.platform == "win32":
|
| return _free_ram_mb_windows()
|
| try:
|
| with open("/proc/meminfo", encoding="ascii") as handle:
|
| for line in handle:
|
| if line.startswith("MemAvailable:"):
|
| return int(line.split()[1]) / 1024
|
| except OSError:
|
| return None
|
| return None
|
|
|
|
|
| def _free_ram_mb_windows() -> float | None:
|
| class MemoryStatus(ctypes.Structure):
|
| _fields_ = [
|
| ("dwLength", ctypes.c_uint32),
|
| ("dwMemoryLoad", ctypes.c_uint32),
|
| ("ullTotalPhys", ctypes.c_uint64),
|
| ("ullAvailPhys", ctypes.c_uint64),
|
| ("ullTotalPageFile", ctypes.c_uint64),
|
| ("ullAvailPageFile", ctypes.c_uint64),
|
| ("ullTotalVirtual", ctypes.c_uint64),
|
| ("ullAvailVirtual", ctypes.c_uint64),
|
| ("ullAvailExtendedVirtual", ctypes.c_uint64),
|
| ]
|
|
|
| status = MemoryStatus()
|
| status.dwLength = ctypes.sizeof(MemoryStatus)
|
| if not ctypes.windll.kernel32.GlobalMemoryStatusEx(ctypes.byref(status)):
|
| return None
|
| return status.ullAvailPhys / MB
|
|
|
|
|
| def confine_own_process(limit_mb: int) -> bool:
|
| """Hard-cap this process's memory; on Windows also forbid it from starting processes.
|
|
|
| Allocations past the cap fail with MemoryError. Returns True when the cap is in force.
|
| """
|
| if sys.platform == "win32":
|
| return _cap_memory_windows(limit_mb)
|
| try:
|
| import resource
|
|
|
| limit = limit_mb * MB
|
| resource.setrlimit(resource.RLIMIT_AS, (limit, limit))
|
| return True
|
| except (ImportError, ValueError, OSError):
|
| return False
|
|
|
|
|
| def _cap_memory_windows(limit_mb: int) -> bool:
|
| global _JOB_HANDLE
|
| from ctypes import wintypes
|
|
|
| class BasicLimits(ctypes.Structure):
|
| _fields_ = [
|
| ("PerProcessUserTimeLimit", ctypes.c_int64),
|
| ("PerJobUserTimeLimit", ctypes.c_int64),
|
| ("LimitFlags", wintypes.DWORD),
|
| ("MinimumWorkingSetSize", ctypes.c_size_t),
|
| ("MaximumWorkingSetSize", ctypes.c_size_t),
|
| ("ActiveProcessLimit", wintypes.DWORD),
|
| ("Affinity", ctypes.c_size_t),
|
| ("PriorityClass", wintypes.DWORD),
|
| ("SchedulingClass", wintypes.DWORD),
|
| ]
|
|
|
| class IoCounters(ctypes.Structure):
|
| _fields_ = [(name, ctypes.c_uint64) for name in (
|
| "ReadOperationCount", "WriteOperationCount", "OtherOperationCount",
|
| "ReadTransferCount", "WriteTransferCount", "OtherTransferCount")]
|
|
|
| class ExtendedLimits(ctypes.Structure):
|
| _fields_ = [
|
| ("BasicLimitInformation", BasicLimits),
|
| ("IoInfo", IoCounters),
|
| ("ProcessMemoryLimit", ctypes.c_size_t),
|
| ("JobMemoryLimit", ctypes.c_size_t),
|
| ("PeakProcessMemoryUsed", ctypes.c_size_t),
|
| ("PeakJobMemoryUsed", ctypes.c_size_t),
|
| ]
|
|
|
| job_object_extended_limit_information = 9
|
| job_object_limit_process_memory = 0x00000100
|
| job_object_limit_active_process = 0x00000008
|
| job_object_limit_kill_on_job_close = 0x00002000
|
|
|
| k32 = ctypes.WinDLL("kernel32", use_last_error=True)
|
| k32.CreateJobObjectW.restype = wintypes.HANDLE
|
| k32.CreateJobObjectW.argtypes = [ctypes.c_void_p, wintypes.LPCWSTR]
|
| k32.GetCurrentProcess.restype = wintypes.HANDLE
|
| k32.SetInformationJobObject.restype = wintypes.BOOL
|
| k32.SetInformationJobObject.argtypes = [
|
| wintypes.HANDLE, ctypes.c_int, ctypes.c_void_p, wintypes.DWORD]
|
| k32.AssignProcessToJobObject.restype = wintypes.BOOL
|
| k32.AssignProcessToJobObject.argtypes = [wintypes.HANDLE, wintypes.HANDLE]
|
|
|
| job = k32.CreateJobObjectW(None, None)
|
| if not job:
|
| return False
|
| info = ExtendedLimits()
|
| info.BasicLimitInformation.LimitFlags = (
|
| job_object_limit_process_memory | job_object_limit_active_process
|
| | job_object_limit_kill_on_job_close)
|
| info.BasicLimitInformation.ActiveProcessLimit = 1
|
| info.ProcessMemoryLimit = limit_mb * MB
|
| if not k32.SetInformationJobObject(
|
| job, job_object_extended_limit_information,
|
| ctypes.byref(info), ctypes.sizeof(info)):
|
| return False
|
| if not k32.AssignProcessToJobObject(job, k32.GetCurrentProcess()):
|
| return False
|
| _JOB_HANDLE = job
|
| return True
|
|
|
|
|
| def check_budget(run_dir: str | Path, parallel: int, mem_mb: int,
|
| min_free_disk_gb: float, max_run_mb: float) -> str | None:
|
| """Return a reason to stop, or None when the run is inside its budget."""
|
| free_gb = free_disk_gb(run_dir)
|
| if free_gb < min_free_disk_gb:
|
| return f"only {free_gb:.1f} GB disk free (floor is {min_free_disk_gb} GB)"
|
| used_mb = dir_size_mb(run_dir)
|
| if used_mb > max_run_mb:
|
| return f"run folder is {used_mb:.0f} MB (cap is {max_run_mb:.0f} MB)"
|
| free_mb = free_ram_mb()
|
| needed = parallel * mem_mb
|
| if free_mb is not None and needed > 0.5 * free_mb:
|
| return (f"{parallel} workers x {mem_mb} MB = {needed} MB, more than half of the "
|
| f"{free_mb:.0f} MB RAM currently free; lower --parallel or --mem-mb")
|
| return None
|
|
|
