Coverage for src / cufile_patcher / plugins / system.py: 100%

1from __future__ import annotations 

2 

3import ctypes 

4 

5from ..cufile_types import CUfileDescr, CUfileError, CUfileHandle_t, DescrUnion 

6from .base import CuFileBackend 

7 

8 

9class SystemCuFileBackend(CuFileBackend): 

10 """Backend that calls the system-installed libcufile via ctypes.""" 

11 

12 def __init__(self, library_name: str = "libcufile.so") -> None: 

13 try: 

14 self._lib = ctypes.CDLL(library_name) 

15 except OSError as exc: 

16 raise RuntimeError( 

17 "unable to load libcufile.so; ensure cuFile is installed" 

18 ) from exc 

19 self._configure_signatures() 

20 

21 def _configure_signatures(self) -> None: 

22 self._lib.cuFileDriverOpen.restype = CUfileError 

23 self._lib.cuFileDriverClose.restype = CUfileError 

24 self._lib.cuFileHandleRegister.restype = CUfileError 

25 self._lib.cuFileHandleDeregister.restype = CUfileError 

26 self._lib.cuFileRead.restype = ctypes.c_size_t 

27 self._lib.cuFileWrite.restype = ctypes.c_size_t 

28 

29 self._lib.cuFileHandleRegister.argtypes = [ 

30 ctypes.POINTER(CUfileHandle_t), 

31 ctypes.POINTER(CUfileDescr), 

32 ] 

33 self._lib.cuFileHandleDeregister.argtypes = [CUfileHandle_t] 

34 self._lib.cuFileRead.argtypes = [ 

35 CUfileHandle_t, 

36 ctypes.c_void_p, 

37 ctypes.c_size_t, 

38 ctypes.c_longlong, 

39 ctypes.c_longlong, 

40 ] 

41 self._lib.cuFileWrite.argtypes = [ 

42 CUfileHandle_t, 

43 ctypes.c_void_p, 

44 ctypes.c_size_t, 

45 ctypes.c_longlong, 

46 ctypes.c_longlong, 

47 ] 

48 

49 @staticmethod 

50 def _check_status(status: CUfileError, name: str) -> None: 

51 if status.err != 0: 

52 raise RuntimeError( 

53 f"{name} failed (cuFile err={status.err}, cuda_err={status.cu_err})" 

54 ) 

55 

56 def driver_open(self) -> None: 

57 self._check_status(self._lib.cuFileDriverOpen(), "cuFileDriverOpen") 

58 

59 def driver_close(self) -> None: 

60 self._check_status(self._lib.cuFileDriverClose(), "cuFileDriverClose") 

61 

62 def handle_register(self, fd: int) -> CUfileHandle_t: 

63 descr = CUfileDescr(type=1, handle=DescrUnion(fd=fd), fs_ops=None) 

64 handle = CUfileHandle_t() 

65 status = self._lib.cuFileHandleRegister(ctypes.byref(handle), ctypes.byref(descr)) 

66 self._check_status(status, "cuFileHandleRegister") 

67 return handle 

68 

69 def handle_deregister(self, handle: CUfileHandle_t) -> None: 

70 self._check_status( 

71 self._lib.cuFileHandleDeregister(handle), "cuFileHandleDeregister" 

72 ) 

73 

74 def read( 

75 self, 

76 handle: CUfileHandle_t, 

77 buf: ctypes.c_void_p, 

78 size: int, 

79 file_offset: int, 

80 dev_offset: int, 

81 ) -> int: 

82 return int(self._lib.cuFileRead(handle, buf, size, file_offset, dev_offset)) 

83 

84 def write( 

85 self, 

86 handle: CUfileHandle_t, 

87 buf: ctypes.c_void_p, 

88 size: int, 

89 file_offset: int, 

90 dev_offset: int, 

91 ) -> int: 

92 return int(self._lib.cuFileWrite(handle, buf, size, file_offset, dev_offset))