Coverage for src/arraybridge/oom_recovery.py: 93%

1""" 

2GPU Out of Memory (OOM) recovery utilities. 

3 

4Provides comprehensive OOM detection and cache clearing for all supported 

5GPU frameworks in OpenHCS. 

6 

7OOM classification and cache cleanup delegate to each ``MemoryType`` declaration. 

8""" 

9 

10import gc 

11import logging 

12 

13from arraybridge.types import MemoryType 

14 

15logger = logging.getLogger(__name__) 

16 

17 

18def _is_oom_error(e: Exception, memory_type: str) -> bool: 

19 """ 

20 Detect Out of Memory errors for all GPU frameworks. 

21 

22 Args: 

23 e: Exception to check 

24 memory_type: Memory type string (e.g., 'torch', 'cupy') 

25 

26 Returns: 

27 True if exception is an OOM error for the given framework 

28 """ 

29 try: 

30 mem_type_enum = MemoryType(memory_type) 

31 except ValueError: 

32 return False 

33 

34 return mem_type_enum.is_oom_error(e) 

35 

36 

37def _clear_cache_for_memory_type(memory_type: str, device_id: int | None = None): 

38 """ 

39 Clear GPU cache for specific memory type. 

40 

41 Args: 

42 memory_type: Memory type string (e.g., 'torch', 'cupy') 

43 device_id: Optional framework-local GPU device ID. ``None`` cleans all. 

44 """ 

45 try: 

46 mem_type_enum = MemoryType(memory_type) 

47 except ValueError: 

48 logger.warning(f"Unknown memory type for cache clearing: {memory_type}") 

49 gc.collect() 

50 return 

51 

52 try: 

53 mem_type_enum.cleanup_loaded(device_id) 

54 except Exception as e: 

55 logger.warning(f"Failed to clear cache for {memory_type}: {e}") 

56 

57 # Always trigger Python garbage collection 

58 gc.collect() 

59 

60 

61def _execute_with_oom_recovery( 

62 func_callable, 

63 memory_type: str, 

64 max_retries: int = 2, 

65 device_id: int | None = None, 

66): 

67 """ 

68 Execute function with automatic OOM recovery. 

69 

70 Args: 

71 func_callable: Function to execute 

72 memory_type: Memory type from MemoryType enum 

73 max_retries: Maximum number of retry attempts 

74 device_id: Optional framework-local device whose cache should be cleared 

75 

76 Returns: 

77 Function result 

78 

79 Raises: 

80 Original exception if not OOM or retries exhausted 

81 """ 

82 for attempt in range(max_retries + 1): 

83 try: 

84 return func_callable() 

85 except Exception as e: 

86 if not _is_oom_error(e, memory_type) or attempt == max_retries: 

87 raise 

88 

89 # Clear cache and retry 

90 _clear_cache_for_memory_type(memory_type, device_id)