import os
import torch

def configure_gpu(target_gpu: int = 0,
                  max_mem_fraction: float = 0.40,
                  preallocate_megabytes: int | None = None) -> None:
    """
    Configure GPU visibility and memory use for a PyTorch process.

    Parameters
    ----------
    target_gpu : int
        Index of the GPU that this process should see (0-based).
        If you want to let PyTorch pick any free GPU, skip this argument.
    max_mem_fraction : float
        Fraction (0–1) of that GPU's total memory PyTorch may allocate.
        Uses torch.cuda.set_per_process_memory_fraction (PyTorch ≥ 1.7).
    preallocate_megabytes : int | None
        If set, immediately reserve this many MB with an empty tensor so
        PyTorch never exceeds the requested size.
    """
    # 1 — Hide all other GPUs
    os.environ["CUDA_VISIBLE_DEVICES"] = str(target_gpu)

    if not torch.cuda.is_available():
        print("CUDA not available, using CPU only.")
        return

    device_index = 0   # because after CUDA_VISIBLE_DEVICES we see only one GPU
    torch.cuda.set_device(device_index)

    # 2 — Limit allocator to a fraction of free memory
    torch.cuda.set_per_process_memory_fraction(max_mem_fraction, device_index)

    # 3 — Optionally pre-allocate a fixed block
    if preallocate_megabytes:
        dummy = torch.empty(
            (preallocate_megabytes * 256),   # 4 bytes × 256 = 1 MB
            dtype=torch.float32,
            device=f"cuda:{device_index}"
        )
        del dummy  # immediately free it; allocator keeps the block reserved

    # Sanity print
    total_mem = torch.cuda.get_device_properties(device_index).total_memory / 1024**3
    print(f"GPU {target_gpu} visible, capped at {max_mem_fraction*100:.0f}% "
          f"≈ {total_mem*max_mem_fraction:.1f} GiB usable by PyTorch.")

# Example usage ---------------------------------------------------------------
if __name__ == "__main__":
    configure_gpu(target_gpu=0, max_mem_fraction=0.2, preallocate_megabytes=4096)