aaronrene / scooling-lab public
fake_worker.py python
121 lines 4.1 KB
Raw
sha256:91e875d4a97bb1e35f37992f803988d5713931f1782d870c371c50054574af22 Add fixture provenance retention deletion Human minor ⚠ breaking 42 days ago
1 """In-process fake worker for the T2 training API contract."""
2
3 from __future__ import annotations
4
5 import hashlib
6 import importlib.resources
7 import json
8 from functools import cache
9
10 from scooling_lab.contracts import TrainingJobStatus
11 from scooling_lab.errors import ApiError, ErrorCode
12 from scooling_lab.provenance import ProvenanceRecord, validate_provenance_record
13 from scooling_lab.store import ArtifactMetadata, TrainingJobRecord, TrainingJobStore, utc_now_iso
14
15 FIXTURE_DATASET_RESOURCE = "synthetic_training_dataset.jsonl"
16
17
18 def fixture_dataset_bytes() -> bytes:
19 """Read the committed synthetic fixture dataset from package resources."""
20
21 return (
22 importlib.resources.files("scooling_lab.fixtures")
23 .joinpath(FIXTURE_DATASET_RESOURCE)
24 .read_bytes()
25 )
26
27
28 def fixture_dataset_hash() -> str:
29 """Return the stable SHA-256 hash for the synthetic fixture dataset."""
30
31 return hashlib.sha256(fixture_dataset_bytes()).hexdigest()
32
33
34 def placeholder_artifact_hash(job: TrainingJobRecord, dataset_hash: str) -> str:
35 """Return a stable hash for the placeholder artifact metadata."""
36
37 if job.request is None:
38 raise ApiError(ErrorCode.INTERNAL_ERROR, 500)
39 payload = json.dumps(
40 {
41 "datasetHash": dataset_hash,
42 "jobId": job.id,
43 "modelId": job.request.model_id,
44 "placeholder": "scooling-lab-fake-artifact-v1",
45 "trainingParameters": dict(sorted(job.request.training_parameters.items())),
46 },
47 separators=(",", ":"),
48 sort_keys=True,
49 )
50 return hashlib.sha256(payload.encode("utf-8")).hexdigest()
51
52
53 @cache
54 def cached_fixture_dataset_hash() -> str:
55 """Return the fixture dataset hash without repeated fixture reads."""
56
57 return fixture_dataset_hash()
58
59
60 def training_config_hash(job: TrainingJobRecord) -> str:
61 """Return a content-free hash of the bounded training configuration."""
62
63 if job.request is None:
64 raise ApiError(ErrorCode.INTERNAL_ERROR, 500)
65 payload = json.dumps(
66 {
67 "modelId": job.request.model_id,
68 "retentionPolicy": job.request.retention_policy.to_public_dict(),
69 "trainingParameters": dict(sorted(job.request.training_parameters.items())),
70 },
71 separators=(",", ":"),
72 sort_keys=True,
73 )
74 return hashlib.sha256(payload.encode("utf-8")).hexdigest()
75
76
77 class FakeTrainingWorker:
78 """Synchronous fake worker that performs no real training and writes no models."""
79
80 def __init__(self, store: TrainingJobStore) -> None:
81 """Bind the worker to a server-owned job store."""
82
83 self._store = store
84
85 def run_job(self, job_id: str) -> TrainingJobRecord:
86 """Move a queued job through running to succeeded and register metadata."""
87
88 job = self._store.get(job_id)
89 if job.status == TrainingJobStatus.SUCCEEDED:
90 return job
91 self._store.update_status(job_id, TrainingJobStatus.RUNNING)
92 running_job = self._store.get(job_id)
93 if running_job.request is None:
94 raise ApiError(ErrorCode.INTERNAL_ERROR, 500)
95 dataset_hash = cached_fixture_dataset_hash()
96 artifact_hash = placeholder_artifact_hash(running_job, dataset_hash)
97 artifact_id = f"artifact_{artifact_hash[:24]}"
98 created_at = utc_now_iso()
99 provenance = ProvenanceRecord(
100 artifact_hash=artifact_hash,
101 base_model_id=running_job.request.model_id,
102 created_at=created_at,
103 dataset_hash=dataset_hash,
104 job_id=job_id,
105 training_config_hash=training_config_hash(running_job),
106 )
107 validate_provenance_record(provenance)
108 self._store.register_artifact(
109 job_id,
110 ArtifactMetadata(
111 id=artifact_id,
112 job_id=job_id,
113 dataset_hash=dataset_hash,
114 artifact_hash=artifact_hash,
115 created_at=created_at,
116 provenance_id=f"provenance_{job_id}",
117 retention_policy=running_job.request.retention_policy,
118 ),
119 provenance,
120 )
121 return self._store.update_status(job_id, TrainingJobStatus.SUCCEEDED)
File History 1 commit
sha256:91e875d4a97bb1e35f37992f803988d5713931f1782d870c371c50054574af22 Add fixture provenance retention deletion Human minor ⚠ 42 days ago