store.py file-level

at main · View file ↗ · Intel ↗

History
1 files
1 commits
0 hotspots
0 🧊 dead
0 💥 blast risk
sha256:c Add T4 job cancellation retry and validation · · Jun 11, 2026
1 """Thread-safe job and artifact store for the in-process fake worker."""
2
3 from __future__ import annotations
4
5 import hashlib
6 import json
7 import threading
8 from dataclasses import dataclass, field
9 from datetime import UTC, datetime
10 from pathlib import Path
11 from typing import Iterable
12
13 from scooling_lab.contracts import (
14 TrainingJobRequest,
15 TrainingJobStatus,
16 require_artifact_id,
17 require_job_id,
18 )
19 from scooling_lab.errors import ApiError, ErrorCode
20 from scooling_lab.provenance import ProvenanceRecord, validate_provenance_record
21 from scooling_lab.retention import (
22 RetentionPolicy,
23 expires_at,
24 is_expired,
25 retention_policy_from_mapping,
26 )
27 from scooling_lab.state_machine import cancel, transition
28
29
30 def utc_now_iso() -> str:
31 """Return a UTC timestamp formatted for public metadata."""
32
33 return datetime.now(UTC).replace(microsecond=0).isoformat().replace("+00:00", "Z")
34
35
36 @dataclass(frozen=True)
37 class ArtifactMetadata:
38 """Placeholder artifact metadata registered by the fake worker."""
39
40 id: str
41 job_id: str
42 dataset_hash: str
43 artifact_hash: str
44 created_at: str
45 provenance_id: str
46 retention_policy: RetentionPolicy
47
48 def to_dict(self) -> dict[str, object]:
49 """Serialize the artifact metadata for API responses and persistence."""
50
51 return {
52 "id": self.id,
53 "jobId": self.job_id,
54 "datasetHash": self.dataset_hash,
55 "artifactHash": self.artifact_hash,
56 "createdAt": self.created_at,
57 "expiresAt": expires_at(self.created_at, self.retention_policy),
58 "provenanceRecordId": self.provenance_id,
59 "retentionPolicy": self.retention_policy.to_public_dict(),
60 }
61
62 @classmethod
63 def from_dict(cls, payload: dict[str, object]) -> "ArtifactMetadata":
64 """Rehydrate persisted artifact metadata."""
65
66 return cls(
67 id=str(payload["id"]),
68 job_id=str(payload["jobId"]),
69 dataset_hash=str(payload["datasetHash"]),
70 artifact_hash=str(payload["artifactHash"]),
71 created_at=str(payload["createdAt"]),
72 provenance_id=str(payload["provenanceRecordId"]),
73 retention_policy=retention_policy_from_mapping(payload["retentionPolicy"]),
74 )
75
76
77 @dataclass(frozen=True)
78 class DeletionReceipt:
79 """Content-free result for idempotent artifact deletion."""
80
81 job_id: str
82 artifact_id: str
83 deleted: bool
84 already_deleted: bool
85 verified: bool
86
87 def to_dict(self) -> dict[str, str | bool]:
88 """Serialize deletion status without hashes, paths, or request content."""
89
90 return {
91 "alreadyDeleted": self.already_deleted,
92 "artifactId": self.artifact_id,
93 "deleted": self.deleted,
94 "jobId": self.job_id,
95 "verified": self.verified,
96 }
97
98
99 @dataclass
100 class TrainingJobRecord:
101 """Internal record for one fixture training job."""
102
103 id: str
104 request: TrainingJobRequest | None
105 status: TrainingJobStatus = TrainingJobStatus.QUEUED
106 created_at: str = field(default_factory=utc_now_iso)
107 updated_at: str = field(default_factory=utc_now_iso)
108 artifacts: list[ArtifactMetadata] = field(default_factory=list)
109 provenance: ProvenanceRecord | None = None
110 deleted_at: str | None = None
111 retry_of_job_id: str | None = None
112
113 def to_public_dict(self) -> dict[str, object]:
114 """Serialize the safe job status shape for API responses."""
115
116 if self.status == TrainingJobStatus.DELETED:
117 payload: dict[str, object] = {
118 "id": self.id,
119 "status": self.status.value,
120 "createdAt": self.created_at,
121 "updatedAt": self.updated_at,
122 }
123 if self.deleted_at is not None:
124 payload["deletedAt"] = self.deleted_at
125 if self.retry_of_job_id is not None:
126 payload["retryOfJobId"] = self.retry_of_job_id
127 return payload
128 if self.request is None:
129 raise ApiError(ErrorCode.INTERNAL_ERROR, 500)
130 payload = {
131 "id": self.id,
132 "status": self.status.value,
133 "createdAt": self.created_at,
134 "updatedAt": self.updated_at,
135 "request": self.request.to_public_dict(),
136 }
137 if self.retry_of_job_id is not None:
138 payload["retryOfJobId"] = self.retry_of_job_id
139 return payload
140
141 def to_persisted_dict(self) -> dict[str, object]:
142 """Serialize the full non-secret record to the server-controlled store."""
143
144 if self.status == TrainingJobStatus.DELETED:
145 payload: dict[str, object] = {
146 "id": self.id,
147 "status": self.status.value,
148 "createdAt": self.created_at,
149 "updatedAt": self.updated_at,
150 }
151 if self.deleted_at is not None:
152 payload["deletedAt"] = self.deleted_at
153 if self.provenance is not None:
154 payload["provenance"] = self.provenance.to_dict()
155 if self.retry_of_job_id is not None:
156 payload["retryOfJobId"] = self.retry_of_job_id
157 return payload
158 if self.request is None:
159 raise ApiError(ErrorCode.INTERNAL_ERROR, 500)
160 payload = {
161 "id": self.id,
162 "status": self.status.value,
163 "createdAt": self.created_at,
164 "updatedAt": self.updated_at,
165 "request": {
166 "idempotencyKey": self.request.idempotency_key,
167 "datasetId": self.request.dataset_id,
168 "modelId": self.request.model_id,
169 "requestedBy": self.request.requested_by,
170 "retentionPolicy": self.request.retention_policy.to_public_dict(),
171 "trainingParameters": dict(self.request.training_parameters),
172 },
173 "artifacts": [artifact.to_dict() for artifact in self.artifacts],
174 "provenance": self.provenance.to_dict()
175 if self.provenance is not None
176 else None,
177 }
178 if self.retry_of_job_id is not None:
179 payload["retryOfJobId"] = self.retry_of_job_id
180 return payload
181
182 @classmethod
183 def from_persisted_dict(cls, payload: dict[str, object]) -> "TrainingJobRecord":
184 """Rehydrate a job record from server-controlled JSON."""
185
186 status = TrainingJobStatus(str(payload["status"]))
187 if status == TrainingJobStatus.DELETED:
188 deleted_at = payload.get("deletedAt")
189 provenance_payload = payload.get("provenance")
190 provenance: ProvenanceRecord | None = None
191 if provenance_payload is not None:
192 if not isinstance(provenance_payload, dict):
193 raise ApiError(ErrorCode.INTERNAL_ERROR, 500)
194 provenance = ProvenanceRecord.from_mapping(provenance_payload)
195 return cls(
196 id=str(payload["id"]),
197 request=None,
198 status=status,
199 created_at=str(payload["createdAt"]),
200 updated_at=str(payload["updatedAt"]),
201 deleted_at=str(deleted_at) if deleted_at is not None else None,
202 provenance=provenance,
203 retry_of_job_id=str(payload["retryOfJobId"])
204 if payload.get("retryOfJobId") is not None
205 else None,
206 )
207 request_payload = payload["request"]
208 if not isinstance(request_payload, dict):
209 raise ApiError(ErrorCode.INTERNAL_ERROR, 500)
210 artifacts_payload = payload.get("artifacts", [])
211 if not isinstance(artifacts_payload, list):
212 raise ApiError(ErrorCode.INTERNAL_ERROR, 500)
213 provenance_payload = payload.get("provenance")
214 provenance = None
215 if provenance_payload is not None:
216 if not isinstance(provenance_payload, dict):
217 raise ApiError(ErrorCode.INTERNAL_ERROR, 500)
218 provenance = ProvenanceRecord.from_mapping(provenance_payload)
219 return cls(
220 id=str(payload["id"]),
221 request=TrainingJobRequest.from_mapping(request_payload),
222 status=status,
223 created_at=str(payload["createdAt"]),
224 updated_at=str(payload["updatedAt"]),
225 artifacts=[
226 ArtifactMetadata.from_dict(artifact)
227 for artifact in artifacts_payload
228 if isinstance(artifact, dict)
229 ],
230 provenance=provenance,
231 retry_of_job_id=str(payload["retryOfJobId"])
232 if payload.get("retryOfJobId") is not None
233 else None,
234 )
235
236
237 class TrainingJobStore:
238 """In-memory job store with optional server-controlled JSON persistence."""
239
240 def __init__(
241 self,
242 persistence_path: Path | None = None,
243 queue_limit: int = 5,
244 max_concurrent_running: int = 1,
245 ) -> None:
246 """Create a store with bounded queued/running capacity and a concurrency limit.
247
248 ``queue_limit`` caps the total number of queued+running jobs.
249 ``max_concurrent_running`` caps how many jobs may be in the running
250 state simultaneously (default 1 — FIFO serial execution).
251 """
252
253 self._lock = threading.RLock()
254 self._jobs: dict[str, TrainingJobRecord] = {}
255 self._persistence_path = persistence_path
256 self._queue_limit = queue_limit
257 self._max_concurrent_running = max_concurrent_running
258 if persistence_path is not None and persistence_path.exists():
259 self._load()
260
261 def create(self, request: TrainingJobRequest) -> TrainingJobRecord:
262 """Create or return the deterministic duplicate job for a request."""
263
264 with self._lock:
265 job_id = request.stable_job_id()
266 existing = self._jobs.get(job_id)
267 if existing is not None:
268 return existing
269 if self.active_count() >= self._queue_limit:
270 raise ApiError(ErrorCode.QUEUE_LIMIT_EXCEEDED, 429)
271 record = TrainingJobRecord(id=job_id, request=request)
272 self._jobs[job_id] = record
273 self._save()
274 return record
275
276 def cancel(self, job_id: str) -> TrainingJobRecord:
277 """Cancel queued/running jobs and promote queued work after running cancel."""
278
279 with self._lock:
280 record = self.get(job_id)
281 previous_status = record.status
282 next_status = cancel(record.status)
283 if next_status != record.status:
284 record.status = next_status
285 record.updated_at = utc_now_iso()
286 if previous_status == TrainingJobStatus.RUNNING:
287 self._promote_queued_locked()
288 self._save()
289 return record
290
291 def retry(self, job_id: str) -> TrainingJobRecord:
292 """Create a new queued job linked to a failed or cancelled original."""
293
294 with self._lock:
295 original = self.get(job_id)
296 if original.status not in {
297 TrainingJobStatus.FAILED,
298 TrainingJobStatus.CANCELLED,
299 }:
300 raise ApiError(ErrorCode.INVALID_TRANSITION, 409)
301 if original.request is None:
302 raise ApiError(ErrorCode.INVALID_TRANSITION, 409)
303 if self.active_count() >= self._queue_limit:
304 raise ApiError(ErrorCode.QUEUE_LIMIT_EXCEEDED, 429)
305
306 retry_id = self._next_retry_id_locked(original.id)
307 record = TrainingJobRecord(
308 id=retry_id,
309 request=original.request,
310 retry_of_job_id=original.id,
311 )
312 self._jobs[retry_id] = record
313 self._save()
314 return record
315
316 def get(self, job_id: str) -> TrainingJobRecord:
317 """Return one job or raise a safe not-found error."""
318
319 with self._lock:
320 record = self._jobs.get(job_id)
321 if record is None:
322 raise ApiError(ErrorCode.NOT_FOUND, 404)
323 return record
324
325 def update_status(
326 self, job_id: str, target_status: TrainingJobStatus
327 ) -> TrainingJobRecord:
328 """Apply the state machine and persist the updated job."""
329
330 with self._lock:
331 record = self.get(job_id)
332 next_status = transition(record.status, target_status)
333 if next_status != record.status:
334 previous_status = record.status
335 record.status = next_status
336 record.updated_at = utc_now_iso()
337 if (
338 previous_status == TrainingJobStatus.RUNNING
339 and record.status != TrainingJobStatus.RUNNING
340 ):
341 self._promote_queued_locked()
342 self._save()
343 return record
344
345 def register_artifact(
346 self, job_id: str, artifact: ArtifactMetadata, provenance: ProvenanceRecord
347 ) -> TrainingJobRecord:
348 """Attach a placeholder artifact once, preserving retry stability."""
349
350 require_job_id(job_id)
351 require_artifact_id(artifact.id)
352 validate_provenance_record(provenance)
353 with self._lock:
354 record = self.get(job_id)
355 if all(existing.id != artifact.id for existing in record.artifacts):
356 record.artifacts.append(artifact)
357 record.provenance = provenance
358 record.updated_at = utc_now_iso()
359 self._save()
360 return record
361
362 def list_artifacts(self, job_id: str) -> list[ArtifactMetadata]:
363 """Return artifacts for one job with no cross-job scan exposure."""
364
365 require_job_id(job_id)
366 with self._lock:
367 record = self.get(job_id)
368 if record.status == TrainingJobStatus.DELETED:
369 return []
370 return list(record.artifacts)
371
372 def get_provenance(self, job_id: str) -> ProvenanceRecord:
373 """Return the validated provenance record for one completed job.
374
375 Provenance is readable even for deleted tombstones when the deletion
376 was triggered by retention expiry (the record retains its provenance).
377 Explicitly deleted artifacts have provenance wiped and return NOT_FOUND.
378 """
379
380 require_job_id(job_id)
381 with self._lock:
382 record = self.get(job_id)
383 if record.provenance is None:
384 raise ApiError(ErrorCode.NOT_FOUND, 404)
385 return record.provenance
386
387 def evaluate_expiry(
388 self, job_id: str, now: datetime | None = None
389 ) -> TrainingJobRecord:
390 """Evaluate a job's artifact expiry and delete it when TTL has elapsed."""
391
392 require_job_id(job_id)
393 with self._lock:
394 record = self.get(job_id)
395 self._delete_expired_artifact_locked(record, now or datetime.now(UTC))
396 return record
397
398 def sweep_expired(self, now: datetime | None = None) -> dict[str, object]:
399 """Delete every expired artifact and return a content-free sweep summary."""
400
401 deleted_job_ids: list[str] = []
402 sweep_time = now or datetime.now(UTC)
403 with self._lock:
404 for record in sorted(self._jobs.values(), key=lambda item: item.id):
405 receipt = self._delete_expired_artifact_locked(record, sweep_time)
406 if receipt is not None and receipt.deleted:
407 deleted_job_ids.append(record.id)
408 return {
409 "deletedJobIds": deleted_job_ids,
410 "deletedCount": len(deleted_job_ids),
411 }
412
413 def delete_artifact(self, job_id: str, artifact_id: str) -> DeletionReceipt:
414 """Delete an artifact, provenance, and job content as an idempotent cascade."""
415
416 require_job_id(job_id)
417 require_artifact_id(artifact_id)
418 with self._lock:
419 record = self.get(job_id)
420 if record.status == TrainingJobStatus.DELETED:
421 return DeletionReceipt(
422 job_id=job_id,
423 artifact_id=artifact_id,
424 deleted=True,
425 already_deleted=True,
426 verified=True,
427 )
428 artifact = self._find_artifact(record, artifact_id)
429 if artifact is None:
430 raise ApiError(ErrorCode.NOT_FOUND, 404)
431 return self._delete_artifact_locked(record, artifact)
432
433 def verify_hash_absence(self, hash_values: Iterable[str]) -> bool:
434 """Verify supplied hashes are absent from every store serialization."""
435
436 hashes = tuple(value for value in hash_values if value)
437 if not hashes:
438 return False
439 with self._lock:
440 output = json.dumps(
441 {
442 "artifacts": {
443 record.id: [artifact.to_dict() for artifact in record.artifacts]
444 for record in self._jobs.values()
445 },
446 "jobs": [
447 record.to_public_dict()
448 for record in sorted(self._jobs.values(), key=lambda item: item.id)
449 ],
450 "persisted": [
451 record.to_persisted_dict()
452 for record in sorted(self._jobs.values(), key=lambda item: item.id)
453 ],
454 "provenance": {
455 record.id: record.provenance.to_dict()
456 for record in self._jobs.values()
457 if record.provenance is not None
458 },
459 },
460 sort_keys=True,
461 )
462 return all(hash_value not in output for hash_value in hashes)
463
464 def active_count(self) -> int:
465 """Count queued and running jobs for quota enforcement."""
466
467 return sum(
468 1
469 for record in self._jobs.values()
470 if record.status
471 in {TrainingJobStatus.QUEUED, TrainingJobStatus.RUNNING}
472 )
473
474 def running_count(self) -> int:
475 """Count jobs currently in the running state."""
476
477 return sum(
478 1
479 for record in self._jobs.values()
480 if record.status == TrainingJobStatus.RUNNING
481 )
482
483 def can_run_now(self) -> bool:
484 """Return True when the running-job slot is available."""
485
486 return self.running_count() < self._max_concurrent_running
487
488 def _promote_queued_locked(self) -> None:
489 while self.running_count() < self._max_concurrent_running:
490 queued = next(
491 (
492 record
493 for record in self._jobs.values()
494 if record.status == TrainingJobStatus.QUEUED
495 ),
496 None,
497 )
498 if queued is None:
499 return
500 queued.status = transition(queued.status, TrainingJobStatus.RUNNING)
501 queued.updated_at = utc_now_iso()
502
503 def _next_retry_id_locked(self, original_job_id: str) -> str:
504 attempt = 1
505 while True:
506 digest = hashlib.sha256(
507 f"{original_job_id}:retry:{attempt}".encode("utf-8")
508 ).hexdigest()[:24]
509 retry_id = f"job_{digest}"
510 if retry_id not in self._jobs:
511 return retry_id
512 attempt += 1
513
514 def queue_state(self) -> dict[str, object]:
515 """Return a content-free snapshot of the job queue state."""
516
517 with self._lock:
518 queued = sum(
519 1
520 for r in self._jobs.values()
521 if r.status == TrainingJobStatus.QUEUED
522 )
523 running = self.running_count()
524 return {
525 "activeCount": queued + running,
526 "maxConcurrentRunning": self._max_concurrent_running,
527 "queueLimit": self._queue_limit,
528 "queuedCount": queued,
529 "runningCount": running,
530 }
531
532 def _delete_expired_artifact_locked(
533 self, record: TrainingJobRecord, now: datetime
534 ) -> DeletionReceipt | None:
535 if record.status != TrainingJobStatus.SUCCEEDED:
536 return None
537 for artifact in record.artifacts:
538 if is_expired(artifact.created_at, artifact.retention_policy, now):
539 return self._delete_artifact_locked(
540 record, artifact, verify=False, retain_provenance=True
541 )
542 return None
543
544 def _delete_artifact_locked(
545 self,
546 record: TrainingJobRecord,
547 artifact: ArtifactMetadata,
548 verify: bool = True,
549 retain_provenance: bool = False,
550 ) -> DeletionReceipt:
551 hashes = self._hashes_for_artifact(record, artifact)
552 record.status = transition(record.status, TrainingJobStatus.DELETED)
553 record.request = None
554 record.artifacts = []
555 if not retain_provenance:
556 record.provenance = None
557 record.deleted_at = utc_now_iso()
558 record.updated_at = record.deleted_at
559 self._save()
560 return DeletionReceipt(
561 job_id=record.id,
562 artifact_id=artifact.id,
563 deleted=True,
564 already_deleted=False,
565 verified=self.verify_hash_absence(hashes) if verify else True,
566 )
567
568 def _find_artifact(
569 self, record: TrainingJobRecord, artifact_id: str
570 ) -> ArtifactMetadata | None:
571 for artifact in record.artifacts:
572 if artifact.id == artifact_id:
573 return artifact
574 return None
575
576 def _hashes_for_artifact(
577 self, record: TrainingJobRecord, artifact: ArtifactMetadata
578 ) -> tuple[str, ...]:
579 hashes = [artifact.dataset_hash, artifact.artifact_hash]
580 if record.provenance is not None:
581 hashes.extend(
582 [
583 record.provenance.dataset_hash,
584 record.provenance.artifact_hash,
585 record.provenance.training_config_hash,
586 ]
587 )
588 return tuple(hashes)
589
590 def _save(self) -> None:
591 if self._persistence_path is None:
592 return
593 payload = {
594 "jobs": [
595 record.to_persisted_dict()
596 for record in sorted(self._jobs.values(), key=lambda item: item.id)
597 ]
598 }
599 self._persistence_path.parent.mkdir(parents=True, exist_ok=True)
600 self._persistence_path.write_text(
601 json.dumps(payload, indent=2, sort_keys=True) + "\n",
602 encoding="utf-8",
603 )
604
605 def _load(self) -> None:
606 if self._persistence_path is None:
607 return
608 payload = json.loads(self._persistence_path.read_text(encoding="utf-8"))
609 jobs = payload.get("jobs", [])
610 if not isinstance(jobs, list):
611 raise ApiError(ErrorCode.INTERNAL_ERROR, 500)
612 self._jobs = {
613 record.id: record
614 for record in (
615 TrainingJobRecord.from_persisted_dict(job)
616 for job in jobs
617 if isinstance(job, dict)
618 )
619 }