nvidia_cufile_gpudirect_storage_write_read_with_cupy.py

1import os
2import array
3import numpy as np
4import cupy as cp
5from pynvfile import cuFileHandle
6
7# 1. Setup Data and File
8file_path = "/tmp/test_cufile.dat"
9data_size = 1024 * 1024  # 1 MB
10# Create a dummy file for writing
11with open(file_path, "wb") as f:
12    f.write(b"\0" * data_size)
13
14# 2. Initialize GPU Data
15# Create a cupy array (GPU memory) with some data
16gpu_data_out = cp.arange(data_size // 4, dtype=cp.int32)
17gpu_data_in = cp.zeros(data_size // 4, dtype=cp.int32)
18
19# 3. Use cuFile for High-Performance I/O
20# Open file with O_DIRECT for GPUDirect Storage compatibility
21flags = os.O_RDWR | os.O_DIRECT
22fd = os.open(file_path, flags)
23
24try:
25    # Initialize cuFile handle
26    cf_handle = cuFileHandle(fd)
27
28    # Write from GPU memory directly to storage
29    cf_handle.pwrite(gpu_data_out, size=data_size, file_offset=0, dev_offset=0)
30    print("Write complete.")
31
32    # Read from storage directly into GPU memory
33    cf_handle.pread(gpu_data_in, size=data_size, file_offset=0, dev_offset=0)
34    print("Read complete.")
35
36    # Validate results
37    if cp.array_equal(gpu_data_out, gpu_data_in):
38        print("Success: Data matches!")
39    else:
40        print("Failure: Data mismatch.")
41
42finally:
43    # Cleanup
44    cf_handle.close()
45    os.close(fd)
46    if os.path.exists(file_path):
47        os.remove(file_path)