Spaces:
Paused
Paused
| """Persistence for RCA investigations under runs/rca/. | |
| Layout: | |
| runs/rca/<investigation_id>/investigation.json | |
| runs/rca/<investigation_id>/offline.csv | |
| runs/rca/<investigation_id>/shadow.csv | |
| runs/rca/<investigation_id>/meta.json | |
| """ | |
| from __future__ import annotations | |
| import json | |
| import threading | |
| from pathlib import Path | |
| from typing import Dict, List, Optional | |
| import pandas as pd | |
| from sensorflow.rca.models import Investigation | |
| from sensorflow.rca.scenario import ScenarioBundle | |
| _DEFAULT_ROOT = "runs/rca" | |
| _root = Path(_DEFAULT_ROOT) | |
| _lock = threading.Lock() | |
| _cache: Dict[str, Investigation] = {} | |
| _data_cache: Dict[str, ScenarioBundle] = {} | |
| def set_rca_root(path: str) -> None: | |
| """Point the store somewhere else (tests use a tmp dir).""" | |
| global _root | |
| _root = Path(path) | |
| reset_rca_store() | |
| def reset_rca_store() -> None: | |
| _cache.clear() | |
| _data_cache.clear() | |
| def _inv_dir(inv_id: str) -> Path: | |
| return _root / inv_id | |
| def save_investigation(inv: Investigation, | |
| bundle: Optional[ScenarioBundle] = None) -> None: | |
| with _lock: | |
| d = _inv_dir(inv.id) | |
| d.mkdir(parents=True, exist_ok=True) | |
| (d / "investigation.json").write_text( | |
| json.dumps(inv.to_json_dict(), indent=1)) | |
| if bundle is not None: | |
| bundle.offline.to_csv(d / "offline.csv", index=False) | |
| bundle.shadow.to_csv(d / "shadow.csv", index=False) | |
| (d / "meta.json").write_text(json.dumps(bundle.meta, indent=1)) | |
| _data_cache[inv.id] = bundle | |
| _cache[inv.id] = inv | |
| def load_investigation(inv_id: str) -> Optional[Investigation]: | |
| if inv_id in _cache: | |
| return _cache[inv_id] | |
| f = _inv_dir(inv_id) / "investigation.json" | |
| if not f.exists(): | |
| return None | |
| inv = Investigation.from_json_dict(json.loads(f.read_text())) | |
| _cache[inv_id] = inv | |
| return inv | |
| def load_bundle(inv_id: str) -> Optional[ScenarioBundle]: | |
| if inv_id in _data_cache: | |
| return _data_cache[inv_id] | |
| d = _inv_dir(inv_id) | |
| if not (d / "meta.json").exists(): | |
| return None | |
| meta = json.loads((d / "meta.json").read_text()) | |
| offline = pd.read_csv(d / "offline.csv") | |
| shadow = pd.read_csv(d / "shadow.csv") | |
| bundle = ScenarioBundle(cause=meta["cause"], seed=meta["seed"], | |
| offline=offline, shadow=shadow, meta=meta) | |
| _data_cache[inv_id] = bundle | |
| return bundle | |
| def list_investigations() -> List[Dict]: | |
| out: List[Dict] = [] | |
| if not _root.exists(): | |
| return out | |
| for d in sorted(_root.iterdir()): | |
| f = d / "investigation.json" | |
| if not f.is_file(): | |
| continue | |
| inv = load_investigation(d.name) | |
| if inv is None: | |
| continue | |
| done = sum(1 for s in inv.stages | |
| if s.status in ("complete", "complete_with_unknowns")) | |
| summary = { | |
| "id": inv.id, | |
| "name": inv.name, | |
| "baseline_model": inv.baseline_model, | |
| "candidate_model": inv.candidate_model, | |
| "training_mode": inv.training_mode, | |
| "revealed": inv.revealed, | |
| "created_at": inv.created_at, | |
| "claims": inv.claims, | |
| "stages_complete": done, | |
| "stages_total": len(inv.stages), | |
| } | |
| if not inv.training_mode or inv.revealed: | |
| summary["scenario_cause"] = inv.scenario_cause | |
| out.append(summary) | |
| out.sort(key=lambda r: r["created_at"], reverse=True) | |
| return out | |