from mscclpp import Transport, TcpBootstrap, Communicator
from mscclpp._mscclpp import Context, RawGpuBuffer
import cupy as cp
cp.cuda.Device(0).use()
bootstrap = TcpBootstrap.create(0, 1)
bootstrap.initialize(bootstrap.create_unique_id(), 60)
comm = Communicator(bootstrap)
for i in range(100):
    if i % 10 == 0:
        print(f"{i=}", flush=True)
    mem = RawGpuBuffer(2 ** 30)
    reg = comm.register_memory(mem.data(), mem.bytes(), Transport.CudaIpc)
    del reg, mem