simulate_gpu_load.py

"""
GPU stress script for the Scenario 3 demo.
Fills the GPU with matrix multiplications to push utilization above the
85% threshold so the EMOS postprocessor triggers the model hot-swap live.

Usage:
    python3 simulate_gpu_load.py          # CUDA (NVIDIA)
    python3 simulate_gpu_load.py --cpu    # CPU fallback (raises load, not GPU%)

Press Ctrl+C to stop — the EMOS postprocessor will detect the drop below
60% and swap back to the full model automatically.
"""

import sys
import time

USE_CPU = "--cpu" in sys.argv

if USE_CPU:
    print("Running CPU stress (GPU% will not rise — use only as fallback).")
    import threading

    def cpu_burn():
        while True:
            _ = sum(i * i for i in range(100_000))

    for _ in range(4):
        threading.Thread(target=cpu_burn, daemon=True).start()

    try:
        while True:
            time.sleep(1)
    except KeyboardInterrupt:
        print("Stopped.")

else:
    try:
        import cupy as cp
    except ImportError:
        sys.exit("CuPy not found. Install with: pip install cupy-cuda12x  (or use --cpu)")

    # Test if GPU is actually accessible by CuPy
    try:
        cp.cuda.Device(0).compute_capability
    except cp.cuda.runtime.CUDARuntimeError:
        sys.exit("No CUDA GPU detected. Use --cpu flag or check your CUDA install.")

    SIZE = 8192
    print(f"Stressing GPU with {SIZE}×{SIZE} matrix multiplications. Press Ctrl+C to stop.")

    # Create random matrices directly on the GPU
    a = cp.random.randn(SIZE, SIZE, dtype=cp.float32)
    b = cp.random.randn(SIZE, SIZE, dtype=cp.float32)

    try:
        while True:
            # Matrix multiplication
            _ = cp.matmul(a, b)
            # Force synchronization so the GPU doesn't queue asynchronously and drop load
            cp.cuda.Stream.null.synchronize()
    except KeyboardInterrupt:
        print("\nStopped — GPU load released.")