Allow open()/close() without context manager

File registration/deregistration can become a bottleneck. Opening
cuFile handles without a context manager allows us to re-use handles
and register them once.

Author: sdimitro

cufile/cufile.py

@@ -49,37 +49,59 @@ def __init__(self, path: str, mode: str = "r", use_direct_io: bool = False):
         if use_direct_io:
             self._os_mode |= os.O_DIRECT
 
-    def __enter__(self):
+        self._handle = None
+        self._cu_file_handle = None
+
+    def __del__(self):
         """
-        Context manager entry.
+        Destructor to ensure the file is closed.
         """
+        if self.is_open:
+            self.close()
+
+    @property
+    def is_open(self) -> bool:
+        return self._handle is not None
+
+    def open(self):
+        """Opens the file and registers the handle."""
+        if self.is_open:
+            return
         self._handle = os.open(self._path, self._os_mode)
         self._cu_file_handle = cuFileHandleRegister(self._handle)
-        return self
 
-    def __exit__(self, exc_type, exc_val, exc_tb):
-        """
-        Context manager exit.
-        """
+    def close(self):
+        """Deregisters the handle and closes the file."""
+        if not self.is_open:
+            return
         cuFileHandleDeregister(self._cu_file_handle)
         os.close(self._handle)
+        self._handle = None
+        self._cu_file_handle = None
+
+    def __enter__(self):
+        """Context manager entry."""
+        self.open()
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        """Context manager exit."""
+        self.close()
 
     def read(self, dest: ctypes.c_void_p, size: int, file_offset: int = 0, dev_offset: int = 0):
-        """
-        Read from the file.
-        """
+        """Read from the file."""
+        if not self.is_open:
+            raise IOError("File is not open.")
         return cuFileRead(self._cu_file_handle, dest, size, file_offset, dev_offset)
 
     def write(self, src: ctypes.c_void_p, size: int, file_offset: int = 0, dev_offset: int = 0):
-        """
-        Write to the file.
-        """
+        """Write to the file."""
+        if not self.is_open:
+            raise IOError("File is not open.")
         return cuFileWrite(self._cu_file_handle, src, size, file_offset, dev_offset)
 
     def get_handle(self):
-        """
-        Get the file handle.
-        """
+        """Get the file handle."""
         return self._handle
 
-    
+ 

requirements.txt

@@ -2,4 +2,5 @@ pytest>=7.0.0
 pytest-cov>=4.0.0
 black>=22.0.0
 flake8>=6.0.0
-mypy>=1.0.0 
+mypy>=1.0.0
+cuda-bindings>=12.0.0

tests/test_cufile.py

@@ -20,7 +20,9 @@
 assert err == 0, f"cuInit failed: {err}"
 err, device = cuda.cuDeviceGet(CUDA_DEVICE)
 assert err == 0, f"cuDeviceGet failed: {err}"
-err, context = cuda.cuCtxCreate(0, device)
+params = cuda.CUctxCreateParams()
+assert params is not None
+err, context = cuda.cuCtxCreate(params, 0, device)
 assert err == 0, f"cuCtxCreate failed: {err}"
 err, dptr_w = cuda.cuMemAlloc(BUF_SIZE)
 assert err == 0, f"cuMemAlloc failed: {err}"
@@ -41,9 +43,8 @@ def test_cufile_context_manager():
     with CuFile(file_path, "w") as cufile:
         assert isinstance(cufile, CuFile)
 
-def test_cufile_read_write():
+def test_cufile_read_write_with_context_manager():
     """Test that CuFile can read and write to a file."""
-
     begin = time.perf_counter()
     with CuFile(file_path, "w") as cufile:
         begin_write = time.perf_counter()
@@ -69,3 +70,55 @@ def test_cufile_read_write():
     host_buf = (ctypes.c_ubyte * BUF_SIZE).from_address(hptr)
     for i in range(BUF_SIZE):
         assert host_buf[i] == PATTERN_BYTE
+
+def test_cufile_read_write():
+    """Test that CuFile can read and write to a file."""
+
+    begin = time.perf_counter()
+    cufile = CuFile(file_path, "w")
+    cufile.open()
+    begin_write = time.perf_counter()
+    ret = cufile.write(ctypes.c_void_p(int(dptr_w)), BUF_SIZE)
+    write_time = time.perf_counter() - begin_write
+    assert ret == BUF_SIZE
+    cufile.close()
+    dt = time.perf_counter() - begin
+    print(f"WRITE (w/o open/register) {ret/1024/1024:.2f}MB in {write_time*1e3:.2f}ms ({ret/write_time/1024/1024/1024:.2f}GB/s)")
+    print(f"FULL WRITE {ret/1024/1024:.2f}MB in {dt*1e3:.2f}ms ({ret/dt/1024/1024/1024:.2f}GB/s)")
+
+    begin = time.perf_counter()
+    cufile = CuFile(file_path, "r")
+    cufile.open()
+    begin_read = time.perf_counter()
+    ret = cufile.read(ctypes.c_void_p(int(dptr_r)), BUF_SIZE)
+    read_time = time.perf_counter() - begin_read
+    assert ret == BUF_SIZE
+    cufile.close()
+    dt = time.perf_counter() - begin
+    print(f"READ (w/o open/register) {ret/1024/1024:.2f}MB in {read_time*1e3:.2f}ms ({ret/read_time/1024/1024/1024:.2f}GB/s)")
+    print(f"FULL READ {ret/1024/1024:.2f}MB in {dt*1e3:.2f}ms ({ret/dt/1024/1024/1024:.2f}GB/s)")
+
+    err, = cuda.cuMemcpyDtoH(hptr, dptr_r, BUF_SIZE)
+    assert err == 0, f"cuMemcpyDtoH failed: {err}"
+    host_buf = (ctypes.c_ubyte * BUF_SIZE).from_address(hptr)
+    for i in range(BUF_SIZE):
+        assert host_buf[i] == PATTERN_BYTE
+
+def test_read_write_without_open():
+    cufile = CuFile(file_path, "r")
+    try:
+        cufile.read(ctypes.c_void_p(int(dptr_r)), BUF_SIZE)
+        raise AssertionError("Expected an exception but none was raised")
+    except IOError as e:
+        assert "File is not open." in str(e)
+    finally:
+        cufile.close()
+
+    cufile = CuFile(file_path, "w")
+    try:
+        cufile.write(ctypes.c_void_p(int(dptr_w)), BUF_SIZE)
+        raise AssertionError("Expected an exception but none was raised")
+    except IOError as e:
+        assert "File is not open." in str(e)
+    finally:
+        cufile.close()