aaronrene / scooling-lab public
dataset_review.py python
404 lines 13.8 KB
Raw
sha256:ca1d0e687bff8686b37126eb8e4fb6d38b40e189352a4920d66ae89c7274e340 Add T4 job cancellation retry and validation Human minor ⚠ breaking 41 days ago
1 """Dataset registration and review lifecycle for the Scooling Lab training API.
2
3 Datasets move through a bounded state machine before any job can reference them:
4 registered → pending_review → approved | rejected
5
6 Only approved datasets are eligible for job submission. Rejection is communicated
7 as a machine-readable ``RejectionReasonCode`` enum; no caller-supplied text is ever
8 echoed back in API responses or logs.
9 """
10
11 from __future__ import annotations
12
13 import re
14 import threading
15 from dataclasses import dataclass, field
16 from enum import Enum
17 from types import MappingProxyType
18 from typing import Mapping
19
20 from scooling_lab.errors import ApiError, ErrorCode
21
22
23 DATASET_ID_RE = re.compile(r"^[A-Za-z0-9._:-]{3,96}$")
24 FORBIDDEN_STRING_RE = re.compile(
25 r"(?i)(https?://|file://|ssh://|[;&|`$<>]|\.\./|/\w|[A-Za-z]:\\)"
26 )
27 FORBIDDEN_FIELD_TERMS: tuple[str, ...] = (
28 "content",
29 "credential",
30 "document",
31 "email",
32 "file",
33 "instruction",
34 "message",
35 "password",
36 "path",
37 "prompt",
38 "secret",
39 "text",
40 "url",
41 )
42 SYNTHETIC_ROW_COUNT_MIN = 1
43 SYNTHETIC_ROW_COUNT_MAX = 10_000
44 SYNTHETIC_DATASET_SCHEMA: MappingProxyType[str, str] = MappingProxyType(
45 {
46 "exampleId": "string",
47 "inputTokenCount": "integer",
48 "outputTokenCount": "integer",
49 "split": "string",
50 }
51 )
52 ALLOWED_DATASET_REGISTRATION_KEYS: frozenset[str] = frozenset(
53 {"datasetId", "rowCount", "declaredSchema"}
54 )
55
56 # The fixture dataset is pre-approved so existing tests require no changes.
57 FIXTURE_APPROVED_DATASET_IDS: frozenset[str] = frozenset(
58 {"fixture:synthetic-tiny-v1"}
59 )
60
61
62 class DatasetStatus(str, Enum):
63 """States in the dataset review lifecycle."""
64
65 REGISTERED = "registered"
66 PENDING_REVIEW = "pending_review"
67 APPROVED = "approved"
68 REJECTED = "rejected"
69
70
71 class RejectionReasonCode(str, Enum):
72 """Bounded machine-readable reason codes for dataset rejection.
73
74 Free text is never accepted or echoed. Only one of these values may
75 appear in a review decision.
76 """
77
78 SYNTHETIC_LIMIT = "SYNTHETIC_LIMIT"
79 FORMAT_INVALID = "FORMAT_INVALID"
80 POLICY_VIOLATION = "POLICY_VIOLATION"
81 SCHEMA_MISMATCH = "SCHEMA_MISMATCH"
82 DUPLICATE_SUBMISSION = "DUPLICATE_SUBMISSION"
83
84
85 @dataclass(frozen=True)
86 class SyntheticDatasetShape:
87 """Content-free metadata used to validate a synthetic dataset submission."""
88
89 row_count: int
90 declared_schema: Mapping[str, str]
91 parse_reason: RejectionReasonCode | None = None
92
93
94 _DATASET_TRANSITIONS: MappingProxyType[
95 DatasetStatus, frozenset[DatasetStatus]
96 ] = MappingProxyType(
97 {
98 DatasetStatus.REGISTERED: frozenset({DatasetStatus.PENDING_REVIEW}),
99 DatasetStatus.PENDING_REVIEW: frozenset(
100 {DatasetStatus.APPROVED, DatasetStatus.REJECTED}
101 ),
102 DatasetStatus.APPROVED: frozenset(),
103 DatasetStatus.REJECTED: frozenset(),
104 }
105 )
106
107
108 def dataset_transition(
109 current: DatasetStatus, target: DatasetStatus
110 ) -> DatasetStatus:
111 """Apply the dataset review state machine or raise for invalid moves.
112
113 Idempotent same-state replay is allowed for every state.
114 """
115
116 if current == target:
117 return current
118 if target in _DATASET_TRANSITIONS[current]:
119 return target
120 raise ApiError(ErrorCode.INVALID_TRANSITION, 409)
121
122
123 def require_dataset_id(value: object) -> str:
124 """Validate a dataset id against the safe-identifier contract."""
125
126 if not isinstance(value, str) or not DATASET_ID_RE.fullmatch(value):
127 raise ApiError(ErrorCode.VALIDATION_ERROR, 400)
128 if FORBIDDEN_STRING_RE.search(value):
129 raise ApiError(ErrorCode.VALIDATION_ERROR, 400)
130 return value
131
132
133 def require_rejection_reason(value: object) -> RejectionReasonCode:
134 """Validate a rejection reason code; reject unknown or free-text values."""
135
136 if not isinstance(value, str):
137 raise ApiError(ErrorCode.VALIDATION_ERROR, 400)
138 try:
139 return RejectionReasonCode(value)
140 except ValueError as exc:
141 raise ApiError(ErrorCode.VALIDATION_ERROR, 400) from exc
142
143
144 def dataset_shape_from_registration(
145 payload: Mapping[str, object],
146 ) -> tuple[str, SyntheticDatasetShape]:
147 """Validate dataset registration metadata without storing content.
148
149 The dataset id is still schema-validated as a request contract. The optional
150 synthetic shape metadata is stored only as bounded counts and schema labels;
151 submit-time validation converts malformed or policy-unsafe shapes into
152 ``RejectionReasonCode`` values without echoing caller-provided text.
153 """
154
155 dataset_id = require_dataset_id(payload.get("datasetId"))
156 unknown_keys = set(payload).difference(ALLOWED_DATASET_REGISTRATION_KEYS)
157 if unknown_keys:
158 return dataset_id, default_dataset_shape(RejectionReasonCode.FORMAT_INVALID)
159
160 has_shape_metadata = "rowCount" in payload or "declaredSchema" in payload
161 if not has_shape_metadata:
162 return dataset_id, default_dataset_shape()
163
164 row_count = payload.get("rowCount")
165 declared_schema = payload.get("declaredSchema")
166 if (
167 not isinstance(row_count, int)
168 or isinstance(row_count, bool)
169 or not isinstance(declared_schema, Mapping)
170 ):
171 return dataset_id, default_dataset_shape(RejectionReasonCode.FORMAT_INVALID)
172
173 parsed_schema: dict[str, str] = {}
174 for field_name, field_type in declared_schema.items():
175 if not isinstance(field_name, str) or not isinstance(field_type, str):
176 return dataset_id, default_dataset_shape(RejectionReasonCode.FORMAT_INVALID)
177 parsed_schema[field_name] = field_type
178
179 return dataset_id, SyntheticDatasetShape(
180 row_count=row_count,
181 declared_schema=MappingProxyType(parsed_schema),
182 )
183
184
185 def default_dataset_shape(
186 parse_reason: RejectionReasonCode | None = None,
187 ) -> SyntheticDatasetShape:
188 """Return the default valid synthetic shape used by compatibility fixtures."""
189
190 return SyntheticDatasetShape(
191 row_count=SYNTHETIC_ROW_COUNT_MIN,
192 declared_schema=SYNTHETIC_DATASET_SCHEMA,
193 parse_reason=parse_reason,
194 )
195
196
197 def validate_synthetic_dataset_shape(
198 shape: SyntheticDatasetShape,
199 ) -> RejectionReasonCode | None:
200 """Return a bounded rejection code for invalid synthetic dataset metadata."""
201
202 if shape.parse_reason is not None:
203 return shape.parse_reason
204 if not SYNTHETIC_ROW_COUNT_MIN <= shape.row_count <= SYNTHETIC_ROW_COUNT_MAX:
205 return RejectionReasonCode.SYNTHETIC_LIMIT
206 for field_name, field_type in shape.declared_schema.items():
207 if has_forbidden_field_term(field_name) or has_forbidden_field_term(field_type):
208 return RejectionReasonCode.POLICY_VIOLATION
209 if dict(shape.declared_schema) != dict(SYNTHETIC_DATASET_SCHEMA):
210 return RejectionReasonCode.SCHEMA_MISMATCH
211 return None
212
213
214 def has_forbidden_field_term(value: str) -> bool:
215 """Return true when a schema label carries unsafe or content-bearing terms."""
216
217 lowered = value.lower()
218 return FORBIDDEN_STRING_RE.search(value) is not None or any(
219 term in lowered for term in FORBIDDEN_FIELD_TERMS
220 )
221
222
223 def _utc_now_iso() -> str:
224 """Return a UTC timestamp formatted for public metadata."""
225
226 from datetime import UTC, datetime
227
228 return datetime.now(UTC).replace(microsecond=0).isoformat().replace("+00:00", "Z")
229
230
231 @dataclass
232 class DatasetRecord:
233 """Internal record for one registered fixture dataset."""
234
235 dataset_id: str
236 status: DatasetStatus = DatasetStatus.REGISTERED
237 registered_at: str = field(default_factory=_utc_now_iso)
238 updated_at: str = field(default_factory=_utc_now_iso)
239 rejection_reason: RejectionReasonCode | None = None
240 synthetic_shape: SyntheticDatasetShape = field(default_factory=default_dataset_shape)
241
242 def to_public_dict(self) -> dict[str, object]:
243 """Serialize the safe dataset status shape for API responses."""
244
245 payload: dict[str, object] = {
246 "datasetId": self.dataset_id,
247 "status": self.status.value,
248 "registeredAt": self.registered_at,
249 "updatedAt": self.updated_at,
250 }
251 if self.rejection_reason is not None:
252 payload["rejectionReasonCode"] = self.rejection_reason.value
253 return payload
254
255
256 class DatasetStore:
257 """Thread-safe, in-memory store for dataset review lifecycle records.
258
259 The store is pre-populated with the fixture dataset IDs from
260 ``FIXTURE_APPROVED_DATASET_IDS`` so that existing job submission flows
261 continue to work without an explicit registration step.
262 """
263
264 def __init__(self) -> None:
265 """Create a store pre-approving all fixture dataset ids."""
266
267 self._lock = threading.RLock()
268 self._datasets: dict[str, DatasetRecord] = {}
269 ts = _utc_now_iso()
270 for dataset_id in FIXTURE_APPROVED_DATASET_IDS:
271 self._datasets[dataset_id] = DatasetRecord(
272 dataset_id=dataset_id,
273 status=DatasetStatus.APPROVED,
274 registered_at=ts,
275 updated_at=ts,
276 )
277
278 def register(self, dataset_id: str) -> DatasetRecord:
279 """Register a dataset; idempotent if already registered.
280
281 Raises ``CONFLICT`` (409) if the same id has already been approved or
282 rejected so that callers cannot silently re-enter a terminal-path
283 dataset into review.
284 """
285
286 require_dataset_id(dataset_id)
287 return self.register_shape(dataset_id, default_dataset_shape())
288
289 def register_shape(
290 self, dataset_id: str, synthetic_shape: SyntheticDatasetShape
291 ) -> DatasetRecord:
292 """Register a dataset with bounded synthetic validation metadata."""
293
294 with self._lock:
295 existing = self._datasets.get(dataset_id)
296 if existing is not None:
297 if existing.status in {DatasetStatus.APPROVED, DatasetStatus.REJECTED}:
298 raise ApiError(ErrorCode.CONFLICT, 409)
299 return existing
300 record = DatasetRecord(
301 dataset_id=dataset_id,
302 synthetic_shape=synthetic_shape,
303 )
304 self._datasets[dataset_id] = record
305 return record
306
307 def submit_for_review(self, dataset_id: str) -> DatasetRecord:
308 """Validate a registered dataset and record the deterministic decision."""
309
310 require_dataset_id(dataset_id)
311 with self._lock:
312 record = self._get_locked(dataset_id)
313 record.status = dataset_transition(
314 record.status, DatasetStatus.PENDING_REVIEW
315 )
316 rejection_reason = validate_synthetic_dataset_shape(record.synthetic_shape)
317 if rejection_reason is None:
318 record.status = dataset_transition(record.status, DatasetStatus.APPROVED)
319 record.rejection_reason = None
320 else:
321 record.status = dataset_transition(record.status, DatasetStatus.REJECTED)
322 record.rejection_reason = rejection_reason
323 record.updated_at = _utc_now_iso()
324 return record
325
326 def approve(self, dataset_id: str) -> DatasetRecord:
327 """Approve a dataset that is pending review."""
328
329 require_dataset_id(dataset_id)
330 with self._lock:
331 record = self._get_locked(dataset_id)
332 record.status = dataset_transition(record.status, DatasetStatus.APPROVED)
333 record.updated_at = _utc_now_iso()
334 return record
335
336 def reject(
337 self, dataset_id: str, reason_code: RejectionReasonCode
338 ) -> DatasetRecord:
339 """Reject a dataset that is pending review with a bounded reason code."""
340
341 require_dataset_id(dataset_id)
342 with self._lock:
343 record = self._get_locked(dataset_id)
344 if record.status == DatasetStatus.REJECTED:
345 return record
346 record.status = dataset_transition(record.status, DatasetStatus.REJECTED)
347 record.rejection_reason = reason_code
348 record.updated_at = _utc_now_iso()
349 return record
350
351 def get(self, dataset_id: str) -> DatasetRecord:
352 """Return one dataset record or raise a safe not-found error."""
353
354 require_dataset_id(dataset_id)
355 with self._lock:
356 return self._get_locked(dataset_id)
357
358 def is_approved(self, dataset_id: str) -> bool:
359 """Return True only if the dataset is in the approved state."""
360
361 try:
362 require_dataset_id(dataset_id)
363 except ApiError:
364 return False
365 with self._lock:
366 record = self._datasets.get(dataset_id)
367 return record is not None and record.status == DatasetStatus.APPROVED
368
369 def _get_locked(self, dataset_id: str) -> DatasetRecord:
370 record = self._datasets.get(dataset_id)
371 if record is None:
372 raise ApiError(ErrorCode.NOT_FOUND, 404)
373 return record
374
375
376 def validate_review_request(payload: Mapping[str, object]) -> tuple[str, RejectionReasonCode | None]:
377 """Parse and validate a dataset review action payload.
378
379 Returns ``(action, reason_code)`` where ``reason_code`` is non-None only
380 when ``action == "reject"``. No caller-supplied text is echoed.
381 """
382
383 allowed_keys = {"action", "reasonCode"}
384 if not isinstance(payload, Mapping):
385 raise ApiError(ErrorCode.VALIDATION_ERROR, 400)
386 unknown = set(payload) - allowed_keys
387 if unknown:
388 raise ApiError(ErrorCode.VALIDATION_ERROR, 400)
389
390 action = payload.get("action")
391 if action not in {"approve", "reject"}:
392 raise ApiError(ErrorCode.VALIDATION_ERROR, 400)
393
394 reason_code: RejectionReasonCode | None = None
395 if action == "reject":
396 raw_reason = payload.get("reasonCode")
397 if raw_reason is None:
398 raise ApiError(ErrorCode.VALIDATION_ERROR, 400)
399 reason_code = require_rejection_reason(raw_reason)
400 else:
401 if "reasonCode" in payload:
402 raise ApiError(ErrorCode.VALIDATION_ERROR, 400)
403
404 return str(action), reason_code
File History 1 commit
sha256:ca1d0e687bff8686b37126eb8e4fb6d38b40e189352a4920d66ae89c7274e340 Add T4 job cancellation retry and validation Human minor ⚠ 41 days ago