Coverage for src/arraybridge/gpu_cleanup.py: 100%

1""" 

2GPU memory cleanup utilities for different frameworks. 

3 

4This module provides unified GPU memory cleanup functions for PyTorch, CuPy, 

5TensorFlow, JAX, and pyclesperanto. The cleanup functions are designed to be called 

6after processing steps to free up GPU memory that's no longer needed. 

7 

8Framework-specific cleanup behavior belongs to each ``MemoryType`` declaration. 

9""" 

10 

11import logging 

12from types import MappingProxyType 

13 

14from arraybridge.types import MemoryType 

15 

16logger = logging.getLogger(__name__) 

17 

18 

19def _cleanup_declared(mem_type: MemoryType, device_id: int | None = None) -> None: 

20 try: 

21 mem_type.cleanup_loaded(device_id) 

22 except Exception as error: 

23 logger.warning( 

24 "Failed to cleanup %s GPU memory: %s", 

25 mem_type.display_name, 

26 error, 

27 ) 

28 

29 

30def _create_cleanup_function(mem_type: MemoryType): 

31 """Create one compatibility cleanup function from its enum declaration.""" 

32 

33 def cleanup(device_id: int | None = None) -> None: 

34 """Clean an already-loaded framework without importing it.""" 

35 _cleanup_declared(mem_type, device_id) 

36 

37 cleanup.__name__ = f"cleanup_{mem_type.import_name}_gpu" 

38 cleanup.__doc__ = f"Clean already-loaded {mem_type.display_name} GPU resources." 

39 

40 return cleanup 

41 

42 

43# Auto-generate all cleanup functions 

44for mem_type in MemoryType: 

45 cleanup_func = _create_cleanup_function(mem_type) 

46 globals()[cleanup_func.__name__] = cleanup_func 

47 

48 

49# Auto-generate cleanup registry 

50MEMORY_TYPE_CLEANUP_REGISTRY = MappingProxyType( 

51 {mem_type.value: globals()[f"cleanup_{mem_type.import_name}_gpu"] for mem_type in MemoryType} 

52) 

53 

54 

55def cleanup_all_gpu_frameworks(device_id: int | None = None) -> None: 

56 """ 

57 Clean up GPU memory for all available frameworks. 

58 

59 This function calls cleanup for all GPU frameworks that are currently loaded. 

60 It's safe to call even if some frameworks aren't available. 

61 

62 Args: 

63 device_id: Optional GPU device ID. If None, cleans all devices. 

64 """ 

65 logger.debug(f"🔥 GPU CLEANUP: Starting cleanup for all GPU frameworks (device_id={device_id})") 

66 

67 for mem_type in MemoryType: 

68 if mem_type.is_gpu: 

69 _cleanup_declared(mem_type, device_id) 

70 

71 logger.debug("🔥 GPU CLEANUP: Completed cleanup for all GPU frameworks") 

72 

73 

74# Export all cleanup functions and utilities 

75__all__ = [ 

76 "cleanup_all_gpu_frameworks", 

77 "MEMORY_TYPE_CLEANUP_REGISTRY", 

78] + [f"cleanup_{mem_type.import_name}_gpu" for mem_type in MemoryType]