aaronrene / scooling-lab public
service.py python
187 lines 7.1 KB
Raw
sha256:ca1d0e687bff8686b37126eb8e4fb6d38b40e189352a4920d66ae89c7274e340 Add T4 job cancellation retry and validation Human minor ⚠ breaking 41 days ago
1 """Application service for Scooling Lab training API operations."""
2
3 from __future__ import annotations
4
5 import threading
6 from datetime import datetime
7 from typing import Iterable
8
9 from scooling_lab.contracts import (
10 TrainingJobRequest,
11 TrainingJobStatus,
12 require_artifact_id,
13 require_job_id,
14 )
15 from scooling_lab.dataset_review import (
16 DatasetStore,
17 dataset_shape_from_registration,
18 validate_review_request,
19 )
20 from scooling_lab.errors import ApiError, ErrorCode
21 from scooling_lab.fake_worker import FakeTrainingWorker
22 from scooling_lab.store import TrainingJobRecord, TrainingJobStore
23
24
25 class TrainingApiService:
26 """Implements the T2/T3 API contract over a store, fake worker, and dataset store.
27
28 Dataset approval is checked before any job is created. A ``threading.Semaphore``
29 enforces the ``max_concurrent_running`` bound from the store so that concurrent
30 callers see at most that many jobs in the running state simultaneously.
31 """
32
33 def __init__(
34 self,
35 store: TrainingJobStore,
36 auto_run_worker: bool = True,
37 dataset_store: DatasetStore | None = None,
38 ) -> None:
39 """Create a service with optional synchronous fake-worker completion.
40
41 If ``dataset_store`` is omitted a default store pre-approving the
42 synthetic fixture dataset is used, so existing callers are unaffected.
43 """
44
45 self._store = store
46 self._worker = FakeTrainingWorker(store)
47 self._auto_run_worker = auto_run_worker
48 self._dataset_store = dataset_store if dataset_store is not None else DatasetStore()
49 # Semaphore mirrors the store's max_concurrent_running for thread safety.
50 self._run_semaphore = threading.Semaphore(store._max_concurrent_running)
51
52 # ------------------------------------------------------------------ dataset
53
54 def register_dataset(self, payload: dict[str, object]) -> dict[str, object]:
55 """Register a dataset for the review lifecycle."""
56
57 dataset_id, synthetic_shape = dataset_shape_from_registration(payload)
58 record = self._dataset_store.register_shape(dataset_id, synthetic_shape)
59 return record.to_public_dict()
60
61 def submit_dataset_for_review(self, dataset_id: str) -> dict[str, object]:
62 """Advance a registered dataset to pending_review."""
63
64 record = self._dataset_store.submit_for_review(dataset_id)
65 return record.to_public_dict()
66
67 def review_dataset(
68 self, dataset_id: str, payload: dict[str, object]
69 ) -> dict[str, object]:
70 """Apply an approve or reject decision to a pending-review dataset.
71
72 The ``payload`` must contain ``action: "approve" | "reject"`` and, for
73 rejections, a bounded ``reasonCode`` enum value. No free text is
74 accepted or echoed.
75 """
76
77 action, reason_code = validate_review_request(payload)
78 if action == "approve":
79 record = self._dataset_store.approve(dataset_id)
80 else:
81 assert reason_code is not None
82 record = self._dataset_store.reject(dataset_id, reason_code)
83 return record.to_public_dict()
84
85 def get_dataset(self, dataset_id: str) -> dict[str, object]:
86 """Return the current review status for one dataset."""
87
88 return self._dataset_store.get(dataset_id).to_public_dict()
89
90 # ------------------------------------------------------------------- queue
91
92 def get_queue_state(self) -> dict[str, object]:
93 """Return a content-free snapshot of the job queue state."""
94
95 return self._store.queue_state()
96
97 # -------------------------------------------------------------------- jobs
98
99 def create_training_job(self, payload: dict[str, object]) -> dict[str, object]:
100 """Validate, create, and optionally complete a fixture training job.
101
102 Raises ``DATASET_NOT_APPROVED`` (403) if the requested dataset has not
103 passed the review lifecycle.
104 """
105
106 request = TrainingJobRequest.from_mapping(payload)
107 if not self._dataset_store.is_approved(request.dataset_id):
108 raise ApiError(ErrorCode.DATASET_NOT_APPROVED, 403)
109 job = self._store.create(request)
110 job = self._run_if_configured(job.id)
111 return job.to_public_dict()
112
113 def get_training_job(self, job_id: str) -> dict[str, object]:
114 """Return the public status for one training job."""
115
116 require_job_id(job_id)
117 return self._store.evaluate_expiry(job_id).to_public_dict()
118
119 def cancel_training_job(self, job_id: str) -> dict[str, object]:
120 """Cancel a queued or running training job."""
121
122 require_job_id(job_id)
123 return self._store.cancel(job_id).to_public_dict()
124
125 def retry_training_job(self, job_id: str) -> dict[str, object]:
126 """Retry a failed or cancelled job with a fresh lineage-linked job id."""
127
128 require_job_id(job_id)
129 retry = self._store.retry(job_id)
130 retry = self._run_if_configured(retry.id)
131 return retry.to_public_dict()
132
133 def list_artifacts(self, job_id: str) -> dict[str, object]:
134 """Return placeholder artifacts registered for one job."""
135
136 require_job_id(job_id)
137 self._store.evaluate_expiry(job_id)
138 artifacts = [artifact.to_dict() for artifact in self._store.list_artifacts(job_id)]
139 return {"jobId": job_id, "artifacts": artifacts}
140
141 def get_provenance(self, job_id: str) -> dict[str, object]:
142 """Return the validated provenance record for one completed fixture job.
143
144 Provenance is also returned for expired (tombstone) jobs when the
145 deletion was triggered by retention TTL, not an explicit delete call.
146 """
147
148 require_job_id(job_id)
149 self._store.evaluate_expiry(job_id)
150 return self._store.get_provenance(job_id).to_dict()
151
152 def delete_artifact(self, job_id: str, artifact_id: str) -> dict[str, object]:
153 """Delete an artifact and all derived content-bearing metadata."""
154
155 require_job_id(job_id)
156 require_artifact_id(artifact_id)
157 return self._store.delete_artifact(job_id, artifact_id).to_dict()
158
159 def sweep_expired_artifacts(self, now: datetime | None = None) -> dict[str, object]:
160 """Evaluate all retention policies and delete expired artifacts."""
161
162 return self._store.sweep_expired(now)
163
164 def verify_deleted_artifact_absence(self, hash_values: Iterable[str]) -> bool:
165 """Verify deleted artifact hashes are absent from all store outputs."""
166
167 return self._store.verify_hash_absence(hash_values)
168
169 def _run_if_configured(self, job_id: str) -> TrainingJobRecord:
170 """Run a queued/running job through the fake worker when configured."""
171
172 if not self._auto_run_worker:
173 return self._store.get(job_id)
174 job = self._store.get(job_id)
175 if job.status not in {
176 TrainingJobStatus.QUEUED,
177 TrainingJobStatus.RUNNING,
178 }:
179 return job
180 acquired = self._run_semaphore.acquire(blocking=True, timeout=10)
181 try:
182 if acquired:
183 return self._worker.run_job(job_id)
184 return self._store.get(job_id)
185 finally:
186 if acquired:
187 self._run_semaphore.release()
File History 1 commit
sha256:ca1d0e687bff8686b37126eb8e4fb6d38b40e189352a4920d66ae89c7274e340 Add T4 job cancellation retry and validation Human minor ⚠ 41 days ago