Coverage for src/arraybridge/utils.py: 90%

1""" 

2Memory conversion utility functions for arraybridge. 

3 

4This module provides utility functions for memory conversion operations, 

5supporting Clause 251 (Declarative Memory Conversion Interface) and 

6Clause 65 (Fail Loudly). 

7""" 

8 

9import importlib 

10import logging 

11from typing import Any 

12 

13from arraybridge.types import MemoryType 

14 

15from .exceptions import MemoryConversionError 

16 

17logger = logging.getLogger(__name__) 

18 

19 

20def optional_import(module_name: str) -> Any | None: 

21 """Import an optional module, returning ``None`` when it is unavailable. 

22 

23 Args: 

24 module_name: Name of the module to import 

25 

26 Returns: 

27 The imported module if available, otherwise ``None``. 

28 

29 Example: 

30 ```python 

31 # Import torch if available 

32 torch = optional_import("torch") 

33 

34 if torch is not None: 

35 # Use torch 

36 tensor = torch.tensor([1, 2, 3]) 

37 else: 

38 # Handle the case where torch is not available 

39 raise ImportError("PyTorch is required for this function") 

40 ``` 

41 """ 

42 try: 

43 return importlib.import_module(module_name) 

44 except ModuleNotFoundError as error: 

45 missing_name = error.name or "" 

46 if module_name == missing_name or module_name.startswith(f"{missing_name}."): 

47 return None 

48 raise 

49 

50 

51def _ensure_module(module_name: str) -> Any: 

52 """ 

53 Ensure a module is imported and meets version requirements. 

54 

55 Args: 

56 module_name: The name of the module to import 

57 

58 Returns: 

59 The imported module 

60 

61 Raises: 

62 ImportError: If the module cannot be imported or does not meet version requirements 

63 RuntimeError: If the module has known issues with specific versions 

64 """ 

65 try: 

66 module = importlib.import_module(module_name) 

67 except ImportError: 

68 raise ImportError( 

69 f"Module {module_name} is required for this operation " f"but is not installed" 

70 ) 

71 

72 return module 

73 

74 

75def _supports_cuda_array_interface(obj: Any) -> bool: 

76 """ 

77 Check if an object supports the CUDA Array Interface. 

78 

79 Args: 

80 obj: The object to check 

81 

82 Returns: 

83 True if the object supports the CUDA Array Interface, False otherwise 

84 """ 

85 return hasattr(obj, "__cuda_array_interface__") 

86 

87 

88def _supports_dlpack(obj: Any) -> bool: 

89 """Return whether an object exposes a standard DLPack export protocol.""" 

90 

91 return any( 

92 callable(getattr(obj, attribute, None)) 

93 for attribute in ("__dlpack__", "toDlpack", "to_dlpack") 

94 ) 

95 

96 

97# Compatibility adapters over declaration-owned device operations. 

98 

99 

100def _get_device_id(data: Any, memory_type: str) -> int | None: 

101 """ 

102 Get the declaration-owned GPU device ID from a data object. 

103 

104 Args: 

105 data: The data object 

106 memory_type: The memory type 

107 

108 Returns: 

109 The GPU device ID or None if not applicable 

110 

111 Raises: 

112 MemoryConversionError: If the device ID cannot be determined for a GPU memory type 

113 """ 

114 mem_type = MemoryType(memory_type) 

115 try: 

116 return mem_type.device_id_of(data) 

117 except Exception as e: 

118 raise MemoryConversionError( 

119 source_type=memory_type, 

120 target_type=memory_type, 

121 method="device_identification", 

122 reason=f"Failed to identify the {mem_type.value} device: {e}", 

123 ) from e 

124 

125 

126def _set_device(memory_type: str, device_id: int) -> None: 

127 """ 

128 Set the current device through its memory-type declaration. 

129 

130 Args: 

131 memory_type: The memory type 

132 device_id: The GPU device ID 

133 

134 Raises: 

135 MemoryConversionError: If the device cannot be set 

136 """ 

137 mem_type = MemoryType(memory_type) 

138 try: 

139 mem_type.activate_device(device_id) 

140 except Exception as e: 

141 raise MemoryConversionError( 

142 source_type=memory_type, 

143 target_type=memory_type, 

144 method="device_selection", 

145 reason=f"Failed to set {mem_type.value} device to {device_id}: {e}", 

146 ) from e 

147 

148 

149def _move_to_device(data: Any, memory_type: str, device_id: int) -> Any: 

150 """ 

151 Move data through its memory-type declaration. 

152 

153 Args: 

154 data: The data to move 

155 memory_type: The memory type 

156 device_id: The target GPU device ID 

157 

158 Returns: 

159 The data on the target device 

160 

161 Raises: 

162 MemoryConversionError: If the data cannot be moved to the specified device 

163 """ 

164 mem_type = MemoryType(memory_type) 

165 try: 

166 return mem_type.move_to_device(data, device_id) 

167 except Exception as e: 

168 raise MemoryConversionError( 

169 source_type=memory_type, 

170 target_type=memory_type, 

171 method="device_movement", 

172 reason=f"Failed to move {mem_type.value} array to device {device_id}: {e}", 

173 ) from e