compressed_tensors.offload.cache.disk_utils

@contextlib.contextmanager
def disk_load_context() -> Iterator[None]:
    """
    Hold safetensors file handles open for the duration of the block.

    `safe_open` parses a header describing every tensor in the file, so the cost
    of opening is proportional to the number of tensors the file holds. Opening
    once per tensor therefore makes a group of N reads from one shard O(N**2).
    Reusing the handle makes the same group linear.

    This is opt in. Outside the context, `DiskCache.onload` opens and closes per
    read exactly as before, so nothing changes for callers that do not use it.

    Nesting is allowed; handles are closed when the outermost context exits.

    Example:
        with disk_load_context():
            for subgraph in subgraphs:
                subgraph(batch)
    """
    if _open_files.depth == 0:
        _open_files.cache = OrderedDict()
    _open_files.depth += 1
    try:
        yield
    finally:
        _open_files.depth -= 1
        if _open_files.depth == 0:
            cache, _open_files.cache = _open_files.cache, None
            for handle in cache.values():
                handle.__exit__(None, None, None)