aaronrene / scooling-lab public
store.py python
229 lines 7.8 KB
Raw
sha256:2313b6a6353294a0ca08ff579624198da250b99163a86215aed31a905912e360 Land Scooling Lab T0/T2 training contract and seven-tier tests Human minor ⚠ breaking 43 days ago
1 """Thread-safe job and artifact store for the in-process fake worker."""
2
3 from __future__ import annotations
4
5 import json
6 import threading
7 from dataclasses import dataclass, field
8 from datetime import UTC, datetime
9 from pathlib import Path
10
11 from scooling_lab.contracts import TrainingJobRequest, TrainingJobStatus
12 from scooling_lab.errors import ApiError, ErrorCode
13 from scooling_lab.state_machine import transition
14
15
16 def utc_now_iso() -> str:
17 """Return a UTC timestamp formatted for public metadata."""
18
19 return datetime.now(UTC).replace(microsecond=0).isoformat().replace("+00:00", "Z")
20
21
22 @dataclass(frozen=True)
23 class ArtifactMetadata:
24 """Placeholder artifact metadata registered by the fake worker."""
25
26 id: str
27 job_id: str
28 dataset_hash: str
29 artifact_hash: str
30 created_at: str
31
32 def to_dict(self) -> dict[str, str]:
33 """Serialize the artifact metadata for API responses and persistence."""
34
35 return {
36 "id": self.id,
37 "jobId": self.job_id,
38 "datasetHash": self.dataset_hash,
39 "artifactHash": self.artifact_hash,
40 "createdAt": self.created_at,
41 }
42
43 @classmethod
44 def from_dict(cls, payload: dict[str, str]) -> "ArtifactMetadata":
45 """Rehydrate persisted artifact metadata."""
46
47 return cls(
48 id=payload["id"],
49 job_id=payload["jobId"],
50 dataset_hash=payload["datasetHash"],
51 artifact_hash=payload["artifactHash"],
52 created_at=payload["createdAt"],
53 )
54
55
56 @dataclass
57 class TrainingJobRecord:
58 """Internal record for one fixture training job."""
59
60 id: str
61 request: TrainingJobRequest
62 status: TrainingJobStatus = TrainingJobStatus.QUEUED
63 created_at: str = field(default_factory=utc_now_iso)
64 updated_at: str = field(default_factory=utc_now_iso)
65 artifacts: list[ArtifactMetadata] = field(default_factory=list)
66
67 def to_public_dict(self) -> dict[str, object]:
68 """Serialize the safe job status shape for API responses."""
69
70 return {
71 "id": self.id,
72 "status": self.status.value,
73 "createdAt": self.created_at,
74 "updatedAt": self.updated_at,
75 "request": self.request.to_public_dict(),
76 }
77
78 def to_persisted_dict(self) -> dict[str, object]:
79 """Serialize the full non-secret record to the server-controlled store."""
80
81 return {
82 "id": self.id,
83 "status": self.status.value,
84 "createdAt": self.created_at,
85 "updatedAt": self.updated_at,
86 "request": {
87 "idempotencyKey": self.request.idempotency_key,
88 "datasetId": self.request.dataset_id,
89 "modelId": self.request.model_id,
90 "requestedBy": self.request.requested_by,
91 "trainingParameters": dict(self.request.training_parameters),
92 },
93 "artifacts": [artifact.to_dict() for artifact in self.artifacts],
94 }
95
96 @classmethod
97 def from_persisted_dict(cls, payload: dict[str, object]) -> "TrainingJobRecord":
98 """Rehydrate a job record from server-controlled JSON."""
99
100 request_payload = payload["request"]
101 if not isinstance(request_payload, dict):
102 raise ApiError(ErrorCode.INTERNAL_ERROR, 500)
103 artifacts_payload = payload.get("artifacts", [])
104 if not isinstance(artifacts_payload, list):
105 raise ApiError(ErrorCode.INTERNAL_ERROR, 500)
106 return cls(
107 id=str(payload["id"]),
108 request=TrainingJobRequest.from_mapping(request_payload),
109 status=TrainingJobStatus(str(payload["status"])),
110 created_at=str(payload["createdAt"]),
111 updated_at=str(payload["updatedAt"]),
112 artifacts=[
113 ArtifactMetadata.from_dict(artifact)
114 for artifact in artifacts_payload
115 if isinstance(artifact, dict)
116 ],
117 )
118
119
120 class TrainingJobStore:
121 """In-memory job store with optional server-controlled JSON persistence."""
122
123 def __init__(self, persistence_path: Path | None = None, queue_limit: int = 5) -> None:
124 """Create a store with a bounded queued/running fixture capacity."""
125
126 self._lock = threading.RLock()
127 self._jobs: dict[str, TrainingJobRecord] = {}
128 self._persistence_path = persistence_path
129 self._queue_limit = queue_limit
130 if persistence_path is not None and persistence_path.exists():
131 self._load()
132
133 def create(self, request: TrainingJobRequest) -> TrainingJobRecord:
134 """Create or return the deterministic duplicate job for a request."""
135
136 with self._lock:
137 job_id = request.stable_job_id()
138 existing = self._jobs.get(job_id)
139 if existing is not None:
140 return existing
141 if self.active_count() >= self._queue_limit:
142 raise ApiError(ErrorCode.QUEUE_LIMIT_EXCEEDED, 429)
143 record = TrainingJobRecord(id=job_id, request=request)
144 self._jobs[job_id] = record
145 self._save()
146 return record
147
148 def get(self, job_id: str) -> TrainingJobRecord:
149 """Return one job or raise a safe not-found error."""
150
151 with self._lock:
152 record = self._jobs.get(job_id)
153 if record is None:
154 raise ApiError(ErrorCode.NOT_FOUND, 404)
155 return record
156
157 def update_status(
158 self, job_id: str, target_status: TrainingJobStatus
159 ) -> TrainingJobRecord:
160 """Apply the state machine and persist the updated job."""
161
162 with self._lock:
163 record = self.get(job_id)
164 next_status = transition(record.status, target_status)
165 if next_status != record.status:
166 record.status = next_status
167 record.updated_at = utc_now_iso()
168 self._save()
169 return record
170
171 def register_artifact(
172 self, job_id: str, artifact: ArtifactMetadata
173 ) -> TrainingJobRecord:
174 """Attach a placeholder artifact once, preserving retry stability."""
175
176 with self._lock:
177 record = self.get(job_id)
178 if all(existing.id != artifact.id for existing in record.artifacts):
179 record.artifacts.append(artifact)
180 record.updated_at = utc_now_iso()
181 self._save()
182 return record
183
184 def list_artifacts(self, job_id: str) -> list[ArtifactMetadata]:
185 """Return artifacts for one job with no cross-job scan exposure."""
186
187 with self._lock:
188 return list(self.get(job_id).artifacts)
189
190 def active_count(self) -> int:
191 """Count queued and running jobs for quota enforcement."""
192
193 return sum(
194 1
195 for record in self._jobs.values()
196 if record.status
197 in {TrainingJobStatus.QUEUED, TrainingJobStatus.RUNNING}
198 )
199
200 def _save(self) -> None:
201 if self._persistence_path is None:
202 return
203 payload = {
204 "jobs": [
205 record.to_persisted_dict()
206 for record in sorted(self._jobs.values(), key=lambda item: item.id)
207 ]
208 }
209 self._persistence_path.parent.mkdir(parents=True, exist_ok=True)
210 self._persistence_path.write_text(
211 json.dumps(payload, indent=2, sort_keys=True) + "\n",
212 encoding="utf-8",
213 )
214
215 def _load(self) -> None:
216 if self._persistence_path is None:
217 return
218 payload = json.loads(self._persistence_path.read_text(encoding="utf-8"))
219 jobs = payload.get("jobs", [])
220 if not isinstance(jobs, list):
221 raise ApiError(ErrorCode.INTERNAL_ERROR, 500)
222 self._jobs = {
223 record.id: record
224 for record in (
225 TrainingJobRecord.from_persisted_dict(job)
226 for job in jobs
227 if isinstance(job, dict)
228 )
229 }
File History 1 commit
sha256:2313b6a6353294a0ca08ff579624198da250b99163a86215aed31a905912e360 Land Scooling Lab T0/T2 training contract and seven-tier tests Human minor ⚠ 43 days ago