Spaces:
Paused
Paused
File size: 4,028 Bytes
5e0b58b | 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 | from __future__ import annotations
import csv
import hashlib
import json
import platform
import sys
from dataclasses import asdict, is_dataclass
from datetime import datetime
from pathlib import Path
from typing import Any, cast
import yaml
class RunArtifacts:
def __init__(
self,
config: dict[str, Any],
run_id: str | None = None,
root: str | Path = "artifacts/runs",
):
if run_id is None:
run_id = datetime.now().strftime("%Y%m%d_%H%M%S") + f"_seed{config['seed']}"
self.run_dir = Path(root) / run_id
self.run_dir.mkdir(parents=True, exist_ok=True)
self.config = config
self.run_id = run_id
canonical = json.dumps(config, sort_keys=True, separators=(",", ":"))
self.config_hash = hashlib.sha256(canonical.encode("utf-8")).hexdigest()[:16]
self._metrics_handle = self._spike_handle = self._stimulus_handle = None
self._metrics_writer = None
self.summary: dict[str, Any] = {}
self._topology_data = None
with (self.run_dir / "environment.json").open("w", encoding="utf-8") as f:
json.dump(
{
"python_version": sys.version,
"platform": platform.platform(),
"machine": platform.machine(),
"processor": platform.processor(),
"run_id": run_id,
"seed": config["seed"],
"timestamp": datetime.now().isoformat(),
"config_hash": self.config_hash,
},
f,
indent=2,
)
with (self.run_dir / "effective_config.yaml").open("w", encoding="utf-8") as f:
yaml.safe_dump(config, f, sort_keys=False)
(self.run_dir / "config_hash.txt").write_text(
self.config_hash, encoding="utf-8"
)
def __enter__(self):
self._metrics_handle = (self.run_dir / "metrics.csv").open(
"w", newline="", encoding="utf-8"
)
self._spike_handle = (self.run_dir / "spikes.jsonl").open("w", encoding="utf-8")
self._stimulus_handle = (self.run_dir / "stimulus.jsonl").open(
"w", encoding="utf-8"
)
return self
def __exit__(
self,
*_args: object,
) -> None:
for handle in (self._metrics_handle, self._spike_handle, self._stimulus_handle):
if handle:
handle.close()
if self._topology_data is not None:
(self.run_dir / "topology.json").write_text(
json.dumps(self._topology_data, indent=2), encoding="utf-8"
)
(self.run_dir / "run_summary.json").write_text(
json.dumps(self.summary, indent=2), encoding="utf-8"
)
def log_metrics(self, metric_dict: dict[str, Any]) -> None:
if self._metrics_writer is None:
self._metrics_writer = csv.DictWriter(
self._metrics_handle, fieldnames=list(metric_dict.keys()) # type: ignore[reportArgumentType]
)
self._metrics_writer.writeheader()
self._metrics_writer.writerow(metric_dict)
def log_spikes(self, tick: int, spike_ids: tuple[int, ...]) -> None:
self._spike_handle.write( # type: ignore[reportOptionalMemberAccess]
json.dumps({"tick": tick, "spikes": list(spike_ids)}) + "\n"
)
def log_stimulus(self, stimulus_result: Any) -> None:
if is_dataclass(stimulus_result):
data: dict[str, Any] = asdict(stimulus_result) # type: ignore[reportArgumentType]
else:
data = cast("dict[str, Any]", vars(stimulus_result))
self._stimulus_handle.write(json.dumps(data) + "\n") # type: ignore[reportOptionalMemberAccess]
def save_topology(self, topology_data: dict[str, Any]) -> None:
self._topology_data = topology_data
def save_summary(self, summary: dict[str, Any]) -> None:
self.summary.update(summary)
|