Spaces:
Paused
Paused
File size: 3,522 Bytes
d1ac27a | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 | """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
|