Coverage for /home/admin/Documents/AI/applications/lexigram-dev/lexigram/experimental/ai/lexigram-ai-evaluation/src/lexigram/ai/evaluation/checkpoints.py: 38%

1"""Digest-verified file checkpoints for the AI evaluation framework. 

2 

3The :class:`FileCheckpointStore` persists run state as JSON documents 

4under ``<root>/runs/<run_id>/checkpoints/<slug>.json``. Every payload 

5is written with a SHA-256 digest of its canonicalized JSON, and loads 

6re-verify that digest so tampered or corrupted checkpoints are never 

7returned as valid state. 

8""" 

9 

10from __future__ import annotations 

11 

12from datetime import UTC, datetime 

13import hashlib 

14from pathlib import Path 

15from typing import Any 

16 

17from lexigram.ai.evaluation.exceptions import CheckpointError 

18from lexigram.ai.evaluation.tracking import canonical_json 

19from lexigram.contracts.ai.experiment import Checkpoint, CheckpointStoreProtocol 

20from lexigram.logging import get_logger 

21from lexigram.serialization import dumps_str, loads_str 

22 

23logger = get_logger(__name__) 

24 

25 

26def _payload_digest(payload: dict[str, Any]) -> str: 

27 """Return the SHA-256 digest of a canonicalized payload.""" 

28 return hashlib.sha256(canonical_json(payload).encode()).hexdigest() 

29 

30 

31class FileCheckpointStore(CheckpointStoreProtocol): 

32 """Filesystem checkpoint store with digest verification. 

33 

34 Args: 

35 root: Base directory for run artifacts. Defaults to ``runs``. 

36 """ 

37 

38 def __init__(self, root: str | Path = "runs") -> None: 

39 self._root = Path(root) 

40 

41 def _checkpoint_file(self, run_id: str, slug: str) -> Path: 

42 return self._root / "runs" / run_id / "checkpoints" / f"{slug}.json" 

43 

44 async def save(self, run_id: str, slug: str, payload: dict[str, Any]) -> Checkpoint: 

45 """Persist a checkpoint for a run. 

46 

47 Args: 

48 run_id: Run identifier. 

49 slug: Stable name of the checkpoint within the run. 

50 payload: State to checkpoint. 

51 

52 Returns: 

53 The stored checkpoint with its content digest. 

54 

55 Raises: 

56 CheckpointError: If the checkpoint cannot be written. 

57 """ 

58 checkpoint = Checkpoint( 

59 run_id=run_id, 

60 slug=slug, 

61 digest=_payload_digest(payload), 

62 payload=payload, 

63 created_at=datetime.now(UTC).isoformat(), 

64 ) 

65 path = self._checkpoint_file(run_id, slug) 

66 try: 

67 path.parent.mkdir(parents=True, exist_ok=True) 

68 path.write_text( 

69 dumps_str( 

70 { 

71 "run_id": run_id, 

72 "slug": slug, 

73 "digest": checkpoint.digest, 

74 "payload": payload, 

75 "created_at": checkpoint.created_at, 

76 }, 

77 sort_keys=True, 

78 ) 

79 ) 

80 except OSError as exc: 

81 raise CheckpointError( 

82 f"cannot write checkpoint {slug!r} for run {run_id!r}: {exc}" 

83 ) from exc 

84 return checkpoint 

85 

86 async def load(self, run_id: str, slug: str) -> Checkpoint | None: 

87 """Load a checkpoint, verifying its content digest. 

88 

89 Args: 

90 run_id: Run identifier. 

91 slug: Checkpoint name. 

92 

93 Returns: 

94 The verified checkpoint, or ``None`` when absent or tampered. 

95 """ 

96 path = self._checkpoint_file(run_id, slug) 

97 if not path.exists(): 

98 return None 

99 try: 

100 data = loads_str(path.read_text()) 

101 except (OSError, ValueError) as exc: 

102 raise CheckpointError( 

103 f"cannot read checkpoint {slug!r} for run {run_id!r}: {exc}" 

104 ) from exc 

105 payload = data["payload"] 

106 expected = data.get("digest", "") 

107 actual = _payload_digest(payload) 

108 if actual != expected: 

109 logger.warning("checkpoint_digest_mismatch", run_id=run_id, slug=slug) 

110 return None 

111 return Checkpoint( 

112 run_id=run_id, 

113 slug=slug, 

114 digest=actual, 

115 payload=payload, 

116 created_at=data.get("created_at", ""), 

117 ) 

118 

119 async def list(self, run_id: str) -> list[Checkpoint]: 

120 """List all checkpoints for a run. 

121 

122 Args: 

123 run_id: Run identifier. 

124 

125 Returns: 

126 Checkpoints in creation order. 

127 """ 

128 directory = self._root / "runs" / run_id / "checkpoints" 

129 if not directory.is_dir(): 

130 return [] 

131 checkpoints: list[Checkpoint] = [] 

132 for path in sorted(directory.glob("*.json")): 

133 loaded = await self.load(run_id, path.stem) 

134 if loaded is not None: 

135 checkpoints.append(loaded) 

136 return checkpoints 

137 

138 

139__all__ = ["FileCheckpointStore"]