Coverage for collectra/pipelines/handlers.py: 0.00%

1"""Artefact handling and persistence classes for Collectra workflows. 

2 

3This module provides specialized handlers for saving workflow configurations 

4and artifacts to different storage formats. It supports both directory-based 

5and zip-based storage options with proper model artifact management. 

6 

7The handlers manage: 

8 - Workflow configuration serialization 

9 - Machine learning model artifact storage 

10 - File system operations for different output formats 

11 - Model versioning and cleanup operations 

12 

13Classes: 

14 DataHandler: Abstract base class for data handling operations 

15 DirectoryHandler: Handler for directory-based workflow storage 

16 ZipHandler: Handler for zip-based workflow storage 

17""" 

18 

19import os 

20import shutil 

21import tempfile 

22import zipfile 

23from dataclasses import dataclass 

24from datetime import datetime 

25from pathlib import Path 

26 

27import yaml 

28 

29from collectra.pipelines.base import Collectra 

30from collectra.tasks.base import Task 

31from collectra.utils import error_msg, processing_msg, success_msg 

32 

33 

34@dataclass(kw_only=True) 

35class DataHandler: 

36 """Abstract base class for workflow data handling operations. 

37 

38 Provides common functionality for saving workflow configurations and 

39 managing machine learning model artifacts across different storage formats. 

40 

41 Attributes: 

42 pipeline (Collectra): The workflow pipeline to save. 

43 """ 

44 

45 pipeline: Collectra 

46 

47 def save(self): 

48 """Save the workflow configuration and artifacts. 

49 

50 Raises: 

51 NotImplementedError: This method must be implemented by subclasses. 

52 """ 

53 raise NotImplementedError( 

54 "If this is being called, the incorrect handler is not assigned." 

55 ) 

56 

57 def get_models(self, task: Task) -> tuple[str, str]: 

58 """Extract current and previous model paths from a machine learning task. 

59 

60 Args: 

61 task: A task exposing a model attribute. 

62 

63 Returns: 

64 tuple[str, str]: Current model path and old model path. 

65 """ 

66 task_config = self.pipeline.data.get(task.name, {}) 

67 return str(getattr(task, "model", "") or ""), str( 

68 task_config.get("old_model", "") 

69 ) 

70 

71 def has_model(self, task: Task) -> bool: 

72 """Return whether a task exposes a model artifact. 

73 

74 Args: 

75 task: Task to inspect. 

76 

77 Returns: 

78 bool: True if the task is a machine learning task, False otherwise. 

79 """ 

80 if not hasattr(task, "model"): 

81 error_message = f"Task {task.name} does not expose a model. Skipping..." 

82 print(error_msg(error_message)) 

83 return False 

84 return True 

85 

86 def _save_config(self, tmp_dir: Path) -> None: 

87 """Save the workflow configuration to a YAML file. 

88 

89 Args: 

90 tmp_dir (Path): Temporary directory to save the configuration file. 

91 """ 

92 pipeline_yaml = tmp_dir / "pipeline.yaml" 

93 data = self.pipeline.get_config() 

94 with open(pipeline_yaml, "w") as f: 

95 for key in data: 

96 f.write( 

97 yaml.dump( 

98 {key: data[key]}, default_flow_style=False, sort_keys=False 

99 ) 

100 ) 

101 f.write("\n") 

102 

103 

104@dataclass(kw_only=True) 

105class DirectoryHandler(DataHandler): 

106 """Handler for saving workflows as directory structures. 

107 

108 Saves workflow configurations and artifacts to a directory structure, 

109 making them easy to inspect and modify manually. 

110 """ 

111 

112 def save(self): 

113 """Save the workflow and its artifacts to a directory structure. 

114 

115 Creates a temporary directory for staging files, then moves them to 

116 the final output directory location. 

117 """ 

118 with tempfile.TemporaryDirectory() as tmpdirname: 

119 out_dir = self.pipeline.get_out_dir() 

120 tmp_dir = Path(tmpdirname) 

121 self._save_artifacts(tmp_dir) 

122 self._save_config(tmp_dir) 

123 out_dir.mkdir( 

124 parents=True, exist_ok=True 

125 ) # Ensure the output directory exists 

126 for item in tmp_dir.iterdir(): 

127 shutil.move(item, out_dir / item.name) 

128 success_message = ( 

129 f"Workflow {self.pipeline.name} was saved successfully to {out_dir}" 

130 ) 

131 print(success_msg(success_message)) 

132 

133 def _save_artifacts(self, tmp_dir: Path) -> None: 

134 """Save machine learning model artifacts to the temporary directory. 

135 

136 Processes all machine learning tasks in the pipeline to save their model 

137 artifacts. Handles model versioning, file movement, and cleanup of old models. 

138 

139 Args: 

140 tmp_dir (Path): Temporary directory to save artifacts to. 

141 """ 

142 old_models: dict[str, Path] = dict() 

143 for task_item in self.pipeline.tasks: 

144 if not self.has_model(task_item): 

145 continue 

146 task = task_item 

147 model, old_model = self.get_models(task) 

148 if not model: 

149 error_message = f"Task {task.name} does not have a model associated with it. Skipping..." 

150 print(error_msg(error_message)) 

151 continue 

152 # Ensure the model path is a Path object 

153 model_path = Path(model) 

154 task_model_path = f"{task.name}-{datetime.now().strftime('%Y%m%d_%H%M%S')}-best.pt" # Default model path for the task 

155 new_model = old_model != model 

156 if not (model_path.exists() and model_path.is_file()): 

157 error_message = f"Model for {task.name} does not point to a valid file or doesn't exist: {model_path}. Attempting to download..." 

158 print(error_msg(error_message)) 

159 task.load(model_path.name) 

160 model_path = Path(model_path.name) 

161 new_model = True 

162 if new_model: 

163 print( 

164 f"Moving model {model_path} to temporary directory with new name {task_model_path}..." 

165 ) 

166 shutil.move(model_path, tmp_dir / task_model_path) 

167 if old_model and Path(old_model).exists(): 

168 old_models[old_model] = Path(old_model) 

169 self.pipeline.data[task.name]["model"] = task_model_path 

170 else: 

171 old_models.pop(model, None) 

172 processing_message = f"Model {model_path} already exists. Copying to temporary directory..." 

173 print(processing_msg(processing_message)) 

174 shutil.copy(model_path, tmp_dir / model_path.name) 

175 self.pipeline.data[task.name]["model"] = model_path.name 

176 

177 for old_model in old_models: 

178 if old_models[old_model].exists(): 

179 print(f"Removing old model {old_models[old_model]}...") 

180 os.remove(old_models[old_model]) 

181 

182 

183@dataclass(kw_only=True) 

184class ZipHandler(DataHandler): 

185 """Handler for saving workflows as compressed zip files. 

186 

187 Saves workflow configurations and artifacts to a single zip file, 

188 making them portable and easy to distribute. 

189 """ 

190 

191 def save(self): 

192 """Save the workflow and its artifacts to a compressed zip file. 

193 

194 Creates a temporary directory for staging files, then compresses 

195 them into a single zip file at the output location. 

196 

197 Raises: 

198 ValueError: If the output directory is not specified. 

199 """ 

200 with tempfile.TemporaryDirectory() as tmpdirname: 

201 if not self.pipeline.out_dir: 

202 raise ValueError("Output directory is not specified.") 

203 tmp_dir = Path(tmpdirname) 

204 self._save_artifacts(tmp_dir) 

205 self._save_config(tmp_dir) 

206 with zipfile.ZipFile( 

207 f"{self.pipeline.out_dir}", "w", zipfile.ZIP_DEFLATED, allowZip64=True 

208 ) as zipf: 

209 for root, _, files in os.walk(tmp_dir): 

210 for file in files: 

211 zipf.write(os.path.join(root, file), file) 

212 processing_message = ( 

213 f"Saved {file} to {Path(self.pipeline.out_dir) / file}" 

214 ) 

215 print(processing_msg(processing_message)) 

216 success_message = f"Workflow {self.pipeline.name} was saved successfully at {self.pipeline.out_dir}" 

217 print(success_msg(success_message)) 

218 

219 def _save_artifacts(self, tmp_dir: Path) -> None: 

220 """Save machine learning model artifacts to the temporary directory for zip compression. 

221 

222 Processes all machine learning tasks in the pipeline to save their model 

223 artifacts for zip archive creation. Handles model extraction from existing 

224 zip files and version management. 

225 

226 Args: 

227 tmp_dir (Path): Temporary directory to save artifacts to. 

228 

229 Raises: 

230 ValueError: If the output directory is not specified. 

231 """ 

232 if not self.pipeline.out_dir: 

233 raise ValueError("Output directory is not specified.") 

234 for task_item in self.pipeline.tasks: 

235 if not self.has_model(task_item): 

236 continue 

237 task = task_item 

238 model, old_model = self.get_models(task) 

239 if not model: 

240 error_message = f"Task {task.name} does not have a model associated with it. Skipping..." 

241 print(error_msg(error_message)) 

242 continue 

243 model_path = Path(model) 

244 task_model_path = ( 

245 f"{task.name}-{datetime.now().strftime('%Y%m%d_%H%M%S')}-best.pt" 

246 ) 

247 new_model: bool = old_model != model 

248 with zipfile.ZipFile(self.pipeline.out_dir, "r") as zipf: 

249 if not model_path.name in zipf.namelist(): 

250 task.load(model_path.name) 

251 model_path = Path(model_path.name) 

252 new_model = True 

253 if new_model: 

254 print( 

255 f"Moving model {model_path} to temporary directory with new name {task_model_path}..." 

256 ) 

257 shutil.move(model_path, tmp_dir / task_model_path) 

258 self.pipeline.data[task.name]["model"] = task_model_path 

259 else: 

260 processing_message = f"Model {model_path} already exists. Copying to temporary directory..." 

261 zipf.extract(str(model_path.name), tmp_dir) 

262 print(processing_msg(processing_message)) 

263 shutil.copy(model_path, tmp_dir / model_path.name) 

264 self.pipeline.data[task.name]["model"] = model_path.name