missing upload

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