Coverage for muutils / mlutils.py: 40%

1"miscellaneous utilities for ML pipelines" 

2 

3from __future__ import annotations 

4 

5import json 

6import os 

7import random 

8import typing 

9import warnings 

10from itertools import islice 

11from pathlib import Path 

12from typing import Any, Callable, Generator, Iterable, Optional, TypeVar, Union 

13 

14ARRAY_IMPORTS: bool 

15try: 

16 import numpy as np 

17 import torch 

18 import torch.backends.mps 

19 

20 ARRAY_IMPORTS = True 

21except ImportError as e: 

22 warnings.warn( 

23 f"Numpy or torch not installed. Array operations will not be available.\n{e}" 

24 ) 

25 ARRAY_IMPORTS = False 

26 

27DEFAULT_SEED: int = 42 

28GLOBAL_SEED: int = DEFAULT_SEED 

29 

30 

31def get_device(device: "Union[str,torch.device,None]" = None) -> "torch.device": 

32 """Get the torch.device instance on which `torch.Tensor`s should be allocated.""" 

33 if not ARRAY_IMPORTS: 

34 raise ImportError( 

35 "Numpy or torch not installed. Array operations will not be available." 

36 ) 

37 try: 

38 # if device is given 

39 assert torch, "Torch is not available, cannot get device" # pyright: ignore[reportPossiblyUnboundVariable] 

40 if device is not None: 

41 device = torch.device(device) 

42 if any( 

43 [ 

44 torch.cuda.is_available() and device.type == "cuda", 

45 torch.backends.mps.is_available() and device.type == "mps", 

46 device.type == "cpu", 

47 ] 

48 ): 

49 # if device is given and available 

50 pass 

51 else: 

52 warnings.warn( 

53 f"Specified device {device} is not available, falling back to CPU" 

54 ) 

55 return torch.device("cpu") 

56 

57 # no device given, infer from availability 

58 else: 

59 if torch.cuda.is_available(): 

60 device = torch.device("cuda") 

61 elif torch.backends.mps.is_available(): 

62 device = torch.device("mps") 

63 else: 

64 device = torch.device("cpu") 

65 

66 # put a dummy tensor on the device to check if it is available 

67 _dummy = torch.zeros(1, device=device) 

68 

69 return device 

70 

71 except Exception as e: 

72 warnings.warn( 

73 f"Error while getting device, falling back to CPU. Error: {e}", 

74 RuntimeWarning, 

75 ) 

76 return torch.device("cpu") # pyright: ignore[reportPossiblyUnboundVariable] 

77 

78 

79def set_reproducibility(seed: int = DEFAULT_SEED): 

80 """ 

81 Improve model reproducibility. See https://github.com/NVIDIA/framework-determinism for more information. 

82 

83 Deterministic operations tend to have worse performance than nondeterministic operations, so this method trades 

84 off performance for reproducibility. Set use_deterministic_algorithms to True to improve performance. 

85 """ 

86 global GLOBAL_SEED 

87 

88 GLOBAL_SEED = seed 

89 

90 random.seed(seed) 

91 

92 if ARRAY_IMPORTS: 

93 try: 

94 assert np, "Numpy is not available, cannot set seed for numpy" # pyright: ignore[reportPossiblyUnboundVariable] 

95 np.random.seed(seed) 

96 except Exception as e: 

97 warnings.warn(f"Error while setting seed for numpy: {e}", RuntimeWarning) 

98 

99 try: 

100 assert torch, "Torch is not available, cannot set seed for torch" # pyright: ignore[reportPossiblyUnboundVariable] 

101 torch.manual_seed(seed) 

102 

103 torch.use_deterministic_algorithms(True) 

104 except Exception as e: 

105 warnings.warn(f"Error while setting seed for torch: {e}", RuntimeWarning) 

106 

107 # Ensure reproducibility for concurrent CUDA streams 

108 # see https://docs.nvidia.com/cuda/cublas/index.html#cublasApi_reproducibility. 

109 os.environ["CUBLAS_WORKSPACE_CONFIG"] = ":4096:8" 

110 

111 

112T = TypeVar("T") 

113 

114 

115def chunks(it: Iterable[T], chunk_size: int) -> Generator[list[T], Any, None]: 

116 """Yield successive chunks from an iterator.""" 

117 # https://stackoverflow.com/a/61435714 

118 iterator = iter(it) 

119 while chunk := list(islice(iterator, chunk_size)): 

120 yield chunk 

121 

122 

123def get_checkpoint_paths_for_run( 

124 run_path: Path, 

125 extension: typing.Literal["pt", "zanj"], 

126 checkpoints_format: str = "checkpoints/model.iter_*.{extension}", 

127) -> list[tuple[int, Path]]: 

128 """get checkpoints of the format from the run_path 

129 

130 note that `checkpoints_format` should contain a glob pattern with: 

131 - unresolved "{extension}" format term for the extension 

132 - a wildcard for the iteration number 

133 """ 

134 

135 assert run_path.is_dir(), ( 

136 f"Model path {run_path} is not a directory (expect run directory, not model files)" 

137 ) 

138 

139 return [ 

140 (int(checkpoint_path.stem.split("_")[-1].split(".")[0]), checkpoint_path) 

141 for checkpoint_path in sorted( 

142 Path(run_path).glob(checkpoints_format.format(extension=extension)) 

143 ) 

144 ] 

145 

146 

147F = TypeVar("F", bound=Callable[..., Any]) 

148 

149 

150def register_method( 

151 method_dict: dict[str, Callable[..., Any]], 

152 custom_name: Optional[str] = None, 

153) -> Callable[[F], F]: 

154 """Decorator to add a method to the method_dict""" 

155 

156 def decorator(method: F) -> F: 

157 method_name: str 

158 if custom_name is None: 

159 method_name_orig: str | None = getattr(method, "__name__", None) 

160 if method_name_orig is None: 

161 warnings.warn( 

162 f"Method {method} does not have a name, using sanitized repr" 

163 ) 

164 from muutils.misc import sanitize_identifier 

165 

166 method_name = sanitize_identifier(repr(method)) 

167 else: 

168 method_name = method_name_orig 

169 else: 

170 method_name = custom_name 

171 # TYPING: ty complains here 

172 method.__name__ = custom_name # type: ignore[unresolved-attribute] 

173 assert method_name not in method_dict, ( 

174 f"Method name already exists in method_dict: {method_name = }, {list(method_dict.keys()) = }" 

175 ) 

176 method_dict[method_name] = method 

177 return method 

178 

179 return decorator 

180 

181 

182def pprint_summary(summary: dict): 

183 print(json.dumps(summary, indent=2))