diff --git a/gptqmodel/utils/disk.py b/gptqmodel/utils/disk.py new file mode 100644 index 000000000..2739f710e --- /dev/null +++ b/gptqmodel/utils/disk.py @@ -0,0 +1,52 @@ +"""Disk IO utilities.""" + +from __future__ import annotations + +import os +import tempfile +import time +from typing import Final + +from .logger import setup_logger + + +log = setup_logger() + +DEFAULT_WRITE_SIZE: Final[int] = 100 * 1024 * 1024 +BUFFER_SIZE: Final[int] = 32 * 1024 * 1024 + + +def estimate_disk_io_speed(write_size_bytes: int = DEFAULT_WRITE_SIZE) -> float: + """Return sequential write speed in bytes/s measured via a temporary file. + + The helper writes ``write_size_bytes`` to a temporary file using a 32 MB CPU + buffer and measures wall-clock time. The temporary file is fsync'd before + computing the rate. Returns 0.0 on failure. + """ + + if write_size_bytes <= 0: + raise ValueError("write_size_bytes must be positive") + + chunk_size = min(BUFFER_SIZE, write_size_bytes) + chunk = b"\0" * chunk_size + total_written = 0 + + try: + start_time = time.perf_counter() + with tempfile.NamedTemporaryFile(delete=True) as tmp_file: + with open(tmp_file.name, "wb", buffering=BUFFER_SIZE) as handle: + while total_written < write_size_bytes: + to_write = min(chunk_size, write_size_bytes - total_written) + handle.write(chunk if to_write == chunk_size else chunk[:to_write]) + total_written += to_write + handle.flush() + os.fsync(handle.fileno()) + elapsed = time.perf_counter() - start_time + except OSError as exc: # pragma: no cover - depends on runtime environment + log.warn(f"Disk IO speed estimation failed: {exc}") + return 0.0 + + if elapsed <= 0: + return float("inf") + + return total_written / elapsed