Coverage for src/finbot/ml/artifacts.py: 87%
171 statements
« prev ^ index » next coverage.py v7.16.0, created at 2026-09-04 12:05 +0000
1import csv
2import hashlib
3import io
4import json
5from collections.abc import Iterable, Mapping
6from dataclasses import dataclass
7from datetime import UTC, datetime
8from enum import StrEnum
9from pathlib import Path
10from uuid import UUID, uuid4
12from finbot.backtest.artifacts import ExperimentStore
13from finbot.backtest.models import BacktestResult
14from finbot.market.cache import HistoricalDataCache
15from finbot.ml.experiment import MLExperimentConfig, MLExperimentResult, MLExperimentRunner
16from finbot.models import Model, load_model
18ML_EXPERIMENT_SCHEMA_VERSION = 1
19REQUIRED_ARTIFACTS = (
20 "model.json",
21 "feature_schema.json",
22 "predictions.csv",
23 "trades.csv",
24 "metrics.json",
25 "comparison.json",
26)
29class MLRunStatus(StrEnum):
30 RUNNING = "running"
31 FAILED = "failed"
32 COMPLETED = "completed"
35class MLArtifactError(RuntimeError):
36 pass
39@dataclass(frozen=True)
40class StoredMLExperiment:
41 run_id: UUID
42 metadata: dict[str, object]
43 model: Model
44 predictions: tuple[dict[str, str], ...]
45 metrics: dict[str, object]
46 comparison: dict[str, object]
49class MLExperimentStore(ExperimentStore):
50 def __init__(self, root: str | Path) -> None:
51 super().__init__(Path(root))
53 def execute(
54 self,
55 config: MLExperimentConfig,
56 runner: MLExperimentRunner,
57 *,
58 git_commit: str,
59 run_id: UUID | None = None,
60 ) -> StoredMLExperiment:
61 resolved_run_id = run_id or uuid4()
62 run_dir = self._begin(resolved_run_id, config, git_commit)
63 try:
64 result = runner.run(config)
65 self._complete(run_dir, resolved_run_id, config, git_commit, result)
66 except Exception as exc:
67 self._write_metadata(
68 run_dir,
69 self._metadata(
70 resolved_run_id,
71 config,
72 git_commit,
73 MLRunStatus.FAILED,
74 error={"type": type(exc).__name__, "message": str(exc)},
75 ),
76 )
77 raise
78 return self.load(resolved_run_id)
80 def load(self, run_id: UUID | str) -> StoredMLExperiment:
81 run_dir = self.root / str(run_id)
82 metadata = self._read_object(run_dir / "metadata.json")
83 status = metadata.get("status")
84 if status != MLRunStatus.COMPLETED:
85 raise MLArtifactError(f"ML experiment {run_id} is {status or 'incomplete'}.")
86 if ( 86 ↛ 91line 86 didn't jump to line 91 because the condition on line 86 was never true
87 metadata.get("schema_version") != ML_EXPERIMENT_SCHEMA_VERSION
88 or metadata.get("run_type") != "ml_experiment"
89 or metadata.get("run_id") != str(run_id)
90 ):
91 raise MLArtifactError("ML experiment metadata is inconsistent.")
93 manifest = self._read_object(run_dir / "manifest.json")
94 if manifest.get("schema_version") != ML_EXPERIMENT_SCHEMA_VERSION: 94 ↛ 95line 94 didn't jump to line 95 because the condition on line 94 was never true
95 raise MLArtifactError("ML experiment manifest version is unsupported.")
96 files = manifest.get("files")
97 if not isinstance(files, dict) or set(files) != set(REQUIRED_ARTIFACTS):
98 raise MLArtifactError("ML experiment manifest is incomplete or inconsistent.")
99 for name, expected_hash in files.items():
100 path = run_dir / name
101 if not path.is_file() or not isinstance(expected_hash, str): 101 ↛ 102line 101 didn't jump to line 102 because the condition on line 101 was never true
102 raise MLArtifactError(f"ML experiment is missing required artifact {name}.")
103 if _sha256(path.read_bytes()) != expected_hash:
104 raise MLArtifactError(f"ML experiment artifact {name} failed integrity validation.")
106 try:
107 model = load_model(run_dir / "model.json")
108 predictions = tuple(csv.DictReader(io.StringIO((run_dir / "predictions.csv").read_text(encoding="utf-8"))))
109 metrics = self._read_object(run_dir / "metrics.json")
110 comparison = self._read_object(run_dir / "comparison.json")
111 feature_schema = self._read_object(run_dir / "feature_schema.json")
112 except (OSError, ValueError) as exc:
113 raise MLArtifactError(f"ML experiment artifacts are invalid: {exc}") from exc
115 if metadata.get("model_id") != model.metadata.model_id: 115 ↛ 116line 115 didn't jump to line 116 because the condition on line 115 was never true
116 raise MLArtifactError("Stored model identity does not match experiment metadata.")
117 if feature_schema.get("feature_names") != list(model.metadata.feature_schema): 117 ↛ 118line 117 didn't jump to line 118 because the condition on line 117 was never true
118 raise MLArtifactError("Stored feature schema does not match the model artifact.")
119 if metadata.get("feature_dataset_id") != feature_schema.get("feature_dataset_id"): 119 ↛ 120line 119 didn't jump to line 120 because the condition on line 119 was never true
120 raise MLArtifactError("Stored feature dataset identity is inconsistent.")
121 if any(row.get("model_id") != model.metadata.model_id for row in predictions): 121 ↛ 122line 121 didn't jump to line 122 because the condition on line 121 was never true
122 raise MLArtifactError("Stored predictions do not match the model artifact.")
123 return StoredMLExperiment(
124 run_id=UUID(str(run_id)),
125 metadata=metadata,
126 model=model,
127 predictions=predictions,
128 metrics=metrics,
129 comparison=comparison,
130 )
132 def _begin(
133 self,
134 run_id: UUID,
135 config: MLExperimentConfig,
136 git_commit: str,
137 ) -> Path:
138 if not git_commit.strip():
139 raise ValueError("git_commit must not be empty")
140 self.root.mkdir(parents=True, exist_ok=True)
141 run_dir = self.root / str(run_id)
142 run_dir.mkdir(exist_ok=False)
143 self._write_metadata(
144 run_dir,
145 self._metadata(run_id, config, git_commit, MLRunStatus.RUNNING),
146 )
147 return run_dir
149 def _complete(
150 self,
151 run_dir: Path,
152 run_id: UUID,
153 config: MLExperimentConfig,
154 git_commit: str,
155 result: MLExperimentResult,
156 ) -> None:
157 result.model.save(run_dir / "model.json")
158 self._write_json(
159 run_dir / "feature_schema.json",
160 {
161 "feature_set": config.feature_set,
162 "feature_dataset_id": result.feature_dataset.dataset_id,
163 "feature_names": list(result.feature_dataset.feature_names),
164 "schema": [entry.to_dict() for entry in result.feature_dataset.metadata.feature_schema],
165 },
166 )
167 self._write_predictions(run_dir / "predictions.csv", result)
168 self._write_trades(run_dir / "trades.csv", result)
169 self._write_json(
170 run_dir / "metrics.json",
171 {
172 "validation": result.validation_metrics.to_dict(),
173 "test": result.test_metrics.to_dict(),
174 "backtest": result.comparison["logistic_regression"].to_dict(),
175 },
176 )
177 self._write_json(
178 run_dir / "comparison.json",
179 {
180 "assumptions": {
181 "dataset_id": config.dataset_id,
182 "symbols": list(config.symbols),
183 "evaluation_period": {
184 "start": config.test_period.start.isoformat(),
185 "end": config.test_period.end.isoformat(),
186 },
187 "initial_cash": config.initial_cash,
188 "fee_rate": config.fee_rate,
189 "slippage_rate": config.slippage_rate,
190 "execution_policy": config.execution_policy.value,
191 },
192 "parameters": {
193 "logistic_regression": {
194 "model_id": result.model.metadata.model_id,
195 "signal_policy": config.to_dict()["signal_policy"],
196 },
197 "sma_crossover": config.to_dict()["baselines"]["sma_crossover"],
198 "buy_and_hold": {},
199 },
200 "results": {
201 name: metrics.to_dict() for name, metrics in result.comparison.items()
202 },
203 },
204 )
205 manifest = {
206 "schema_version": ML_EXPERIMENT_SCHEMA_VERSION,
207 "files": {
208 name: _sha256((run_dir / name).read_bytes())
209 for name in REQUIRED_ARTIFACTS
210 },
211 }
212 self._write_json(run_dir / "manifest.json", manifest)
213 self._write_metadata(
214 run_dir,
215 self._metadata(
216 run_id,
217 config,
218 git_commit,
219 MLRunStatus.COMPLETED,
220 result=result,
221 ),
222 )
224 @staticmethod
225 def _metadata(
226 run_id: UUID,
227 config: MLExperimentConfig,
228 git_commit: str,
229 status: MLRunStatus,
230 *,
231 result: MLExperimentResult | None = None,
232 error: dict[str, str] | None = None,
233 ) -> dict[str, object]:
234 metadata: dict[str, object] = {
235 "schema_version": ML_EXPERIMENT_SCHEMA_VERSION,
236 "run_type": "ml_experiment",
237 "run_id": str(run_id),
238 "status": status.value,
239 "updated_at": datetime.now(UTC).isoformat(),
240 "git_commit": git_commit,
241 "config": config.to_dict(),
242 }
243 if result is not None:
244 metadata.update(
245 {
246 "source_dataset_id": config.dataset_id,
247 "feature_dataset_id": result.feature_dataset.dataset_id,
248 "supervised_dataset_id": result.supervised_dataset.dataset_id,
249 "training_dataset_id": result.model.metadata.training_configuration.dataset_id,
250 "target_definition": {"name": result.supervised_dataset.target_name, "horizon": config.target_horizon},
251 "model_id": result.model.metadata.model_id,
252 "model_type": result.model.metadata.model_type,
253 "model_artifact": "model.json",
254 },
255 )
256 if error is not None:
257 metadata["error"] = error
258 return metadata
260 @classmethod
261 def _write_metadata(cls, run_dir: Path, payload: dict[str, object]) -> None:
262 cls._write_json(run_dir / "metadata.json", payload)
264 @staticmethod
265 def _write_json(path: Path, payload: object) -> None:
266 temporary = path.with_suffix(f"{path.suffix}.tmp")
267 temporary.write_text(
268 json.dumps(payload, indent=2, sort_keys=True, allow_nan=False) + "\n",
269 encoding="utf-8",
270 )
271 temporary.replace(path)
273 @staticmethod
274 def _write_predictions(path: Path, result: MLExperimentResult) -> None:
275 output = io.StringIO(newline="")
276 fields = ["split", "symbol", "timestamp", "probability_positive", "model_id", "model_type"]
277 writer = csv.DictWriter(output, fieldnames=fields, lineterminator="\n")
278 writer.writeheader()
279 for split, predictions in (
280 ("validation", result.validation_predictions),
281 ("test", result.test_predictions),
282 ):
283 for prediction in predictions:
284 writer.writerow(
285 {
286 "split": split,
287 "symbol": prediction.symbol,
288 "timestamp": prediction.timestamp.isoformat(),
289 "probability_positive": repr(prediction.probability_positive),
290 "model_id": prediction.model_id,
291 "model_type": prediction.model_type,
292 },
293 )
294 _write_text_atomic(path, output.getvalue())
296 @staticmethod
297 def _write_trades(path: Path, result: MLExperimentResult) -> None:
298 write_backtest_trades(path, result.backtests)
300 @staticmethod
301 def _read_object(path: Path) -> dict[str, object]:
302 try:
303 payload = json.loads(path.read_text(encoding="utf-8"))
304 except (FileNotFoundError, OSError, json.JSONDecodeError) as exc:
305 raise MLArtifactError(f"Missing or invalid ML artifact {path.name}.") from exc
306 if not isinstance(payload, dict): 306 ↛ 307line 306 didn't jump to line 307 because the condition on line 306 was never true
307 raise MLArtifactError(f"ML artifact {path.name} must contain an object.")
308 return payload
311def write_backtest_trades(
312 path: Path,
313 backtests: Mapping[str, BacktestResult],
314) -> None:
315 fields = [
316 "strategy", "symbol", "entry_time", "exit_time", "quantity",
317 "entry_price", "exit_price", "entry_fee", "exit_fee", "net_pnl", "return_pct",
318 ]
319 output = io.StringIO(newline="")
320 writer = csv.DictWriter(output, fieldnames=fields, lineterminator="\n")
321 writer.writeheader()
322 for strategy, backtest in backtests.items():
323 for trade in backtest.trades: 323 ↛ 324line 323 didn't jump to line 324 because the loop on line 323 never started
324 writer.writerow(
325 {
326 "strategy": strategy,
327 "symbol": trade.symbol,
328 "entry_time": trade.entry_time.isoformat(),
329 "exit_time": trade.exit_time.isoformat(),
330 "quantity": repr(trade.quantity),
331 "entry_price": repr(trade.entry_price),
332 "exit_price": repr(trade.exit_price),
333 "entry_fee": repr(trade.entry_fee),
334 "exit_fee": repr(trade.exit_fee),
335 "net_pnl": repr(trade.net_pnl),
336 "return_pct": repr(trade.return_pct),
337 },
338 )
339 _write_text_atomic(path, output.getvalue())
342def run_and_store(
343 config: MLExperimentConfig,
344 *,
345 datasets: HistoricalDataCache,
346 output_dir: Path,
347 git_commit: str,
348) -> StoredMLExperiment:
349 return MLExperimentStore(output_dir).execute(
350 config,
351 MLExperimentRunner(datasets),
352 git_commit=git_commit,
353 )
356def _write_text_atomic(path: Path, value: str) -> None:
357 temporary = path.with_suffix(f"{path.suffix}.tmp")
358 temporary.write_text(value, encoding="utf-8", newline="")
359 temporary.replace(path)
362def _sha256(payload: bytes) -> str:
363 return hashlib.sha256(payload).hexdigest()
366def compare_ml_experiments(
367 experiments: Iterable["StoredMLExperiment"],
368) -> list[dict[str, object]]:
369 """Return persisted result sets only when their evaluation assumptions match."""
370 items = list(experiments)
371 if not items: 371 ↛ 372line 371 didn't jump to line 372 because the condition on line 371 was never true
372 return []
373 reference = items[0].comparison.get("assumptions")
374 if not isinstance(reference, dict): 374 ↛ 375line 374 didn't jump to line 375 because the condition on line 374 was never true
375 raise ValueError("ML comparison is missing its evaluation assumptions.")
376 reports: list[dict[str, object]] = []
377 for experiment in items:
378 assumptions = experiment.comparison.get("assumptions")
379 results = experiment.comparison.get("results")
380 if assumptions != reference:
381 raise ValueError("ML experiments use incompatible evaluation assumptions.")
382 if not isinstance(results, dict): 382 ↛ 383line 382 didn't jump to line 383 because the condition on line 382 was never true
383 raise ValueError("ML comparison is missing strategy results.")
384 reports.append(
385 {
386 "run_id": str(experiment.run_id),
387 "parameters": experiment.comparison.get("parameters", {}),
388 "results": results,
389 },
390 )
391 return reports