context-bug.py

# /// script
# requires-python = ">=3.13"
# dependencies = [
#     "dask-cuda",
#     "kvikio-cu12",
#     "nvidia-ml-py",
#     "psutil",
#     "rich",
# ]
# ///
import kvikio
import psutil
import pynvml
import rich.table
from dask_cuda import LocalCUDACluster


def list_contexts() -> list[tuple[int, int, str]]:
    # Get the number of available NVIDIA GPUs
    deviceCount = pynvml.nvmlDeviceGetCount()

    records = [
        (i, proc.pid, psutil.Process(proc.pid).name())
        for i in range(deviceCount)
        for proc in pynvml.nvmlDeviceGetComputeRunningProcesses(pynvml.nvmlDeviceGetHandleByIndex(i))
    ]
    return records


def display(records: list[tuple[int, int, str]]):
    t = rich.table.Table()
    t.add_column("GPU")
    t.add_column("PID")
    t.add_column("Process")
    for i, pid, name in records:
        t.add_row(str(i), str(pid), name)

    rich.print(t)


def curead():
    kvikio.CuFile("README.md")


def main():
    pynvml.nvmlInit()
    print("visible devices:", pynvml.nvmlDeviceGetCount())
    print("single")
    before = set(list_contexts())
    curead()
    after = set(list_contexts())
    display(list(after - before))

    print("multi")
    with LocalCUDACluster() as cluster, cluster.get_client() as client:
        before = set(list_contexts())
        client.run(curead)
        after = set(list_contexts())
        display(list(after - before))


if __name__ == "__main__":
    main()