dense-evolution 8.3.0__py3-none-win_amd64.whl

dense_evolution/backends/chunk/geometry.py

@@ -0,0 +1,95 @@
+from typing import Tuple
+
+import numpy as np
+import jax.numpy as jnp
+
+from .guard import _device_memory_budget_bytes, HAS_JAX
+
+__all__ = ["MemoryChunker", "get_dynamic_chunk"]
+
+
+def get_dynamic_chunk(dtype_target) -> int:
+    available_bytes, _source = _device_memory_budget_bytes()
+    safe_bytes = available_bytes * 0.85
+    if HAS_JAX and dtype_target is jnp.complex128:
+        bpe = 16
+    elif dtype_target is np.complex128:
+        bpe = 16
+    else:
+        bpe = 8
+    max_elements = safe_bytes / bpe
+    max_bits = int(np.floor(np.log2(max(max_elements, 2.0))))
+    # Upper bound raised 27 -> 30 (prog.txt / user request): the real
+    # constraint is already _device_memory_budget_bytes's own VRAM/RAM
+    # reading (a 27-bit chunk is only 2.1GB regardless of dtype -- any
+    # GPU with more than ~2.5GB free was being capped well below what it
+    # could actually hold). 30 keeps a ceiling at all (avoids ever
+    # returning a chunk_size_bits so large a single chunk's own indexing
+    # arithmetic elsewhere risks int32 overflow) without arbitrarily
+    # discarding real headroom on 16GB+ GPUs (a Kaggle T4/P100 at
+    # ~0.85*16GB/16B per complex128 element already computes max_bits
+    # near 30 on its own -- this line was the thing throwing that away).
+    return max(16, min(max_bits, 30))
+
+
+def _dtype_for_qubits(n_qubits: int):
+    xp = jnp if HAS_JAX else np
+    return xp.complex64 if n_qubits > 26 else xp.complex128
+
+
+# ─────────────────────────────────────────────────────────────────────────────────
+# MemoryChunker  (chunk1)
+# ─────────────────────────────────────────────────────────────────────────────────
+
+class MemoryChunker:
+    """
+    Geometry calculator for chunked simulation.
+
+    Attributes
+    ----------
+    n_qubits        int   — requested logical qubit count
+    dtype                 — numpy/jax dtype for the statevector
+    chunk_size_bits int   — safe qubit-width that fits in RAM
+    num_chunks      int   — number of statevector chunks required
+    chunk_dim       int   — 2 ** chunk_size_bits
+    """
+
+    def __init__(self, n_qubits: int):
+        # Looked up through the package's own top-level attribute, not
+        # this module's local `get_dynamic_chunk` name directly -- tests
+        # (and any caller) override chunk sizing via
+        # `dense_evolution.chunk.get_dynamic_chunk = ...`, the package's
+        # public attribute; a deferred `from . import` re-reads that
+        # attribute at call time instead of capturing this module's own
+        # function object once at import time, so an override actually
+        # takes effect on the next MemoryChunker() constructed.
+        from . import get_dynamic_chunk as _get_dynamic_chunk
+        self.n_qubits        = n_qubits
+        self.dtype           = _dtype_for_qubits(n_qubits)
+        self.chunk_size_bits = _get_dynamic_chunk(self.dtype)
+
+        if self.n_qubits <= self.chunk_size_bits:
+            self.num_chunks = 1
+            self.chunk_dim  = 2 ** self.n_qubits
+        else:
+            self.num_chunks = 2 ** (self.n_qubits - self.chunk_size_bits)
+            self.chunk_dim  = 2 ** self.chunk_size_bits
+
+    def geometry(self) -> Tuple[int, int, int]:
+        """(num_chunks, chunk_dim, chunk_size_bits)"""
+        return self.num_chunks, self.chunk_dim, self.chunk_size_bits
+
+    def memory_mb(self) -> float:
+        """Estimated RAM per chunk in MB."""
+        bpe = np.dtype(self.dtype).itemsize
+        return (self.chunk_dim * bpe) / (1024 * 1024)
+
+    def __repr__(self) -> str:
+        return (
+            f"MemoryChunker(n_qubits={self.n_qubits}, "
+            f"num_chunks={self.num_chunks}, "
+            f"chunk_dim={self.chunk_dim}, "
+            f"chunk_size_bits={self.chunk_size_bits}, "
+            f"dtype={self.dtype}, "
+            f"mem_per_chunk={self.memory_mb():.2f} MB)"
+        )

dense_evolution/backends/chunk/guard.py

@@ -0,0 +1,190 @@
+import gc
+from typing import Tuple
+
+import psutil
+
+# JAX is now a mandatory dependency (see registry.py) -- kept as an
+# import-time flag for the dual-path code below, which stays as-is.
+import jax
+HAS_JAX = True
+
+__all__ = ["MemoryPressureError", "SafeMemoryGuard"]
+
+
+def _device_memory_budget_bytes(device=None) -> Tuple[float, str]:
+    """Real, dynamic memory budget -- reads the ACTIVE compute device's
+    own memory via JAX's device.memory_stats() (bytes_limit -
+    bytes_in_use) when available, falling back to host RAM via
+    psutil.virtual_memory() otherwise (some backends, e.g. plain CPU,
+    return None from memory_stats() -- confirmed directly, so on CPU
+    this is identical to the old psutil-only behavior). get_dynamic_chunk
+    and SafeMemoryGuard used to read psutil.virtual_memory()
+    UNCONDITIONALLY, which is only correct when the compute device IS the
+    host (true on CPU, false on GPU/TPU) -- sizing chunks off host RAM
+    while the data actually lives in smaller GPU VRAM is exactly why
+    Chunk's own README benchmark ("~2GB constant RAM regardless of qubit
+    count") holds on CPU but silently mis-sized chunks on GPU, causing
+    MemoryPressureError/real OOM well before the GPU's own VRAM was
+    actually exhausted. Verified on a real Colab T4 GPU: before this fix,
+    chunk_size_bits was computed from host RAM (often >11GB free on
+    Colab) while the data lived in the T4's 11-15GB VRAM pool -- after
+    this fix, device.memory_stats() correctly reports VRAM instead.
+
+    Returns (available_bytes, source_description) -- the raw number, not
+    a safety-margined one; callers apply their own margin."""
+    if device is None:
+        device = jax.devices()[0] if HAS_JAX else None
+    stats = device.memory_stats() if device is not None and hasattr(device, "memory_stats") else None
+    if stats and stats.get("bytes_limit"):
+        available = stats["bytes_limit"] - stats.get("bytes_in_use", 0)
+        return float(available), f"device.memory_stats() on {device}"
+    vm = psutil.virtual_memory()
+    return float(vm.available), "psutil (host RAM) -- device.memory_stats() unavailable"
+
+
+def _device_total_bytes(device=None) -> float:
+    """Companion to _device_memory_budget_bytes for the TOTAL (not
+    available) capacity -- used by SafeMemoryGuard to compute free_pct
+    against the right pool (device VRAM on GPU/TPU, host RAM on CPU)."""
+    if device is None:
+        device = jax.devices()[0] if HAS_JAX else None
+    stats = device.memory_stats() if device is not None and hasattr(device, "memory_stats") else None
+    if stats and stats.get("bytes_limit"):
+        return float(stats["bytes_limit"])
+    return float(psutil.virtual_memory().total)
+
+
+# ─────────────────────────────────────────────────────────────────────────────────
+# SafeMemoryGuard  — Anti-OOM block
+# ─────────────────────────────────────────────────────────────────────────────────
+
+class MemoryPressureError(RuntimeError):
+    """
+    Raised when available system RAM drops below the configured safety threshold.
+    Catches the condition *before* the allocator attempts and crashes with
+    jaxlib.xla_extension.XlaRuntimeError: RESOURCE_EXHAUSTED.
+    """
+    pass
+
+
+class SafeMemoryGuard:
+    """
+    Monitors system RAM before every high-memory operation and blocks execution
+    if free RAM falls below ``threshold_pct`` of total physical memory.
+    """
+
+    _WARN_MULTIPLIER = 2.0
+    # Below this, check_allocation's own margin-after-allocation math
+    # reduces to "is ambient free RAM currently below threshold_pct?",
+    # independent of required_mb -- for a request this small that's not
+    # protecting against anything, since real machines don't run out of
+    # memory over a request this size regardless of ambient load. Sized
+    # to comfortably cover a dense statevector up to 24 qubits at
+    # complex128 (2**24 * 16 bytes = 268.4 MB) with a real margin -- the
+    # same 24-qubit figure dashboard_core.system_limits already treats as
+    # MPS's own RAM-independent ceiling, and well inside the 16-27 qubit
+    # range dense_evolution.chunk.get_dynamic_chunk considers safe on any
+    # machine (27 qubits / 2 GB is that function's own practical ceiling
+    # for a single dense block). An 8 GB minimum-spec machine can always
+    # absorb an allocation this size without real risk, even while
+    # running low on RAM from unrelated processes. A 2-qubit statevector
+    # (64 bytes) used to be refused on a machine sitting at 8% free RAM
+    # from unrelated processes, which was not a real anti-OOM save, just
+    # a false positive from a threshold designed for actual multi-GB
+    # allocations. Only bypasses when available_mb is itself comfortably
+    # above this floor too (see check_allocation below) -- a machine
+    # truly down to its last few hundred MB still gets the real check.
+    _NEGLIGIBLE_MB = 512.0
+
+    def __init__(self, threshold_pct: float = 0.15, gc_before_check: bool = True):
+        if not 0.0 < threshold_pct < 1.0:
+            raise ValueError(f"threshold_pct must be in (0, 1), got {threshold_pct}")
+        self.threshold_pct   = threshold_pct
+        self.gc_before_check = gc_before_check
+        # Device VRAM total on GPU/TPU, host RAM total on CPU (identical
+        # to the old psutil-only value there) -- see
+        # _device_memory_budget_bytes's docstring for why this matters.
+        self._total_mb       = _device_total_bytes() / (1024 * 1024)
+
+    def status(self) -> dict:
+        available_bytes, _source = _device_memory_budget_bytes()
+        available_mb = available_bytes / (1024 * 1024)
+        free_pct     = available_mb / self._total_mb if self._total_mb > 0 else 0.0
+        return {
+            "total_mb"    : self._total_mb,
+            "available_mb": available_mb,
+            "used_pct"    : (1.0 - free_pct) * 100.0,
+            "free_pct"    : free_pct * 100.0,
+            "safe"        : free_pct >= self.threshold_pct,
+        }
+
+    def check(self, context: str = "") -> None:
+        if self.gc_before_check:
+            gc.collect()
+
+        s   = self.status()
+        tag = f"[{context}] " if context else ""
+        free_frac = s["free_pct"] / 100.0
+
+        if not s["safe"]:
+            raise MemoryPressureError(
+                f"\n{'─'*60}\n"
+                f"  {tag}MEMORIA CRITICA — simulazione bloccata\n"
+                f"  Disponibile : {s['available_mb']:.0f} MB  "
+                f"({s['free_pct']:.1f}% libera)\n"
+                f"  Soglia      : {self.threshold_pct * 100:.0f}%  "
+                f"({self._total_mb * self.threshold_pct:.0f} MB)\n"
+                f"  Azione      : liberare RAM o ridurre n_qubits / chunk_size.\n"
+                f"{'─'*60}"
+            )
+
+        warn_threshold = self.threshold_pct * self._WARN_MULTIPLIER
+        if free_frac < warn_threshold:
+            print(
+                f"  [WARN] {tag}RAM bassa: {s['available_mb']:.0f} MB liberi "
+                f"({s['free_pct']:.1f}%) — soglia critica al "
+                f"{self.threshold_pct * 100:.0f}%."
+            )
+
+    def check_allocation(self, required_mb: float, context: str = "") -> None:
+        """Like check(), but sized: verifies that *required_mb* can be
+        allocated while still leaving threshold_pct free afterwards —
+        check() alone only looks at RAM free *right now*, independent of
+        what's about to be allocated (needed for Chunk's multi-chunk path,
+        where several chunk-sized simulators are held in RAM at once)."""
+        if self.gc_before_check:
+            gc.collect()
+
+        s = self.status()
+        if required_mb < self._NEGLIGIBLE_MB and s["available_mb"] > self._NEGLIGIBLE_MB:
+            return
+
+        tag = f"[{context}] " if context else ""
+        available_after_mb = s["available_mb"] - required_mb
+        free_frac_after = available_after_mb / self._total_mb if self._total_mb > 0 else 0.0
+
+        if available_after_mb < 0 or free_frac_after < self.threshold_pct:
+            import socket
+            raise MemoryPressureError(
+                f"\n{'─'*60}\n"
+                f"  {tag}MEMORIA INSUFFICIENTE per l'allocazione richiesta\n"
+                f"  Macchina     : {socket.gethostname()} ({self._total_mb / 1024:.1f} GB RAM totali)\n"
+                f"  Richiesti    : {required_mb:.0f} MB\n"
+                f"  Disponibile  : {s['available_mb']:.0f} MB ({s['free_pct']:.1f}% libera)\n"
+                f"  Dopo alloc.  : {available_after_mb:.0f} MB ({free_frac_after * 100:.1f}% libera)\n"
+                f"  Soglia       : {self.threshold_pct * 100:.0f}% libera dopo l'allocazione\n"
+                f"  Azione       : ridurre n_qubits, liberare RAM, oppure passare\n"
+                f"                 allow_disk_overflow=True a Chunk() per continuare\n"
+                f"                 oltre la RAM disponibile usando lo storage su disco\n"
+                f"                 (vedi disk_overflow.py).\n"
+                f"{'─'*60}"
+            )
+
+    def __repr__(self) -> str:
+        s = self.status()
+        return (
+            f"SafeMemoryGuard("
+            f"threshold={self.threshold_pct*100:.0f}%, "
+            f"available={s['available_mb']:.0f} MB / {s['free_pct']:.1f}% free, "
+            f"safe={s['safe']})"
+        )