aaronrene / scooling-lab public
test_integration.py python
183 lines 7.7 KB
Raw
sha256:ca1d0e687bff8686b37126eb8e4fb6d38b40e189352a4920d66ae89c7274e340 Add T4 job cancellation retry and validation Human minor ⚠ breaking 41 days ago
1 """Integration tier tests for Scooling Lab API and BOM generation."""
2
3 from __future__ import annotations
4
5 from datetime import UTC, datetime, timedelta
6 import unittest
7
8 from scooling_lab_helpers import PROJECT_ROOT, valid_payload
9
10 from scooling_lab.bom import collect_entries, render_markdown
11 from scooling_lab.contracts import TrainingJobRequest, TrainingJobStatus
12 from scooling_lab.errors import ApiError
13 from scooling_lab.provenance import validate_provenance_record
14 from scooling_lab.service import TrainingApiService
15 from scooling_lab.store import TrainingJobStore
16
17
18 class ScoolingLabIntegrationTests(unittest.TestCase):
19 """Integration tests spanning service, fake worker, store, and BOM."""
20
21 def test_integration_api_to_fake_worker_lifecycle(self) -> None:
22 """createTrainingJob completes through the fake worker and lists metadata."""
23
24 service = TrainingApiService(TrainingJobStore())
25 created = service.create_training_job(valid_payload("integration"))
26 self.assertEqual(created["status"], "succeeded")
27 job_id = str(created["id"])
28
29 fetched = service.get_training_job(job_id)
30 self.assertEqual(fetched["id"], job_id)
31 self.assertEqual(fetched["status"], "succeeded")
32
33 artifacts = service.list_artifacts(job_id)
34 self.assertEqual(len(artifacts["artifacts"]), 1)
35 artifact = artifacts["artifacts"][0]
36 self.assertEqual(artifact["jobId"], job_id)
37 self.assertTrue(str(artifact["datasetHash"]))
38 self.assertTrue(str(artifact["artifactHash"]))
39 self.assertTrue(str(artifact["provenanceRecordId"]))
40
41 def test_integration_fake_worker_emits_valid_provenance(self) -> None:
42 """Worker completion writes a schema-valid provenance record."""
43
44 service = TrainingApiService(TrainingJobStore())
45 created = service.create_training_job(valid_payload("provenance-integration"))
46 job_id = str(created["id"])
47
48 provenance = service.get_provenance(job_id)
49
50 validate_provenance_record(provenance)
51 self.assertEqual(provenance["jobId"], job_id)
52 self.assertEqual(provenance["baseModelId"], "fixture-tiny-llm")
53
54 def test_integration_t4_submit_validates_and_decides_dataset_status(self) -> None:
55 """Dataset submit deterministically approves valid shapes and rejects invalid ones."""
56
57 service = TrainingApiService(TrainingJobStore())
58 valid = service.register_dataset(
59 {
60 "datasetId": "integration-valid-dataset",
61 "rowCount": 4,
62 "declaredSchema": {
63 "exampleId": "string",
64 "inputTokenCount": "integer",
65 "outputTokenCount": "integer",
66 "split": "string",
67 },
68 }
69 )
70 rejected = service.register_dataset(
71 {
72 "datasetId": "integration-reject-dataset",
73 "rowCount": 0,
74 "declaredSchema": {
75 "exampleId": "string",
76 "inputTokenCount": "integer",
77 "outputTokenCount": "integer",
78 "split": "string",
79 },
80 }
81 )
82
83 valid_decision = service.submit_dataset_for_review(str(valid["datasetId"]))
84 rejected_decision = service.submit_dataset_for_review(str(rejected["datasetId"]))
85
86 self.assertEqual(valid_decision["status"], "approved")
87 self.assertEqual(rejected_decision["status"], "rejected")
88 self.assertEqual(rejected_decision["rejectionReasonCode"], "SYNTHETIC_LIMIT")
89
90 def test_integration_t4_cancel_running_job_frees_slot_for_next_queued(self) -> None:
91 """Cancelling a running job starts the next queued job without oversubscription."""
92
93 store = TrainingJobStore(queue_limit=4, max_concurrent_running=1)
94 service = TrainingApiService(store, auto_run_worker=False)
95 first = service.create_training_job(valid_payload("cancel-slot-a"))
96 second = service.create_training_job(valid_payload("cancel-slot-b"))
97 store.update_status(str(first["id"]), TrainingJobStatus.RUNNING)
98
99 cancelled = service.cancel_training_job(str(first["id"]))
100 queue_state = service.get_queue_state()
101
102 self.assertEqual(cancelled["status"], "cancelled")
103 self.assertEqual(service.get_training_job(str(second["id"]))["status"], "running")
104 self.assertEqual(queue_state["runningCount"], 1)
105 self.assertEqual(queue_state["queuedCount"], 0)
106
107 def test_integration_t4_failed_job_retries_to_new_successful_provenance(self) -> None:
108 """Retrying a failed job creates a new id and independently valid provenance."""
109
110 store = TrainingJobStore()
111 request = TrainingJobRequest.from_mapping(valid_payload("retry-success"))
112 original = store.create(request)
113 store.update_status(original.id, TrainingJobStatus.FAILED)
114 service = TrainingApiService(store)
115
116 retried = service.retry_training_job(original.id)
117 provenance = service.get_provenance(str(retried["id"]))
118
119 self.assertEqual(retried["status"], "succeeded")
120 self.assertNotEqual(retried["id"], original.id)
121 self.assertEqual(retried["retryOfJobId"], original.id)
122 self.assertEqual(service.get_training_job(original.id)["status"], "failed")
123 validate_provenance_record(provenance)
124 self.assertEqual(provenance["jobId"], retried["id"])
125
126 def test_integration_deletion_cascades_store_and_provenance(self) -> None:
127 """deleteArtifact removes artifact metadata and provenance together."""
128
129 service = TrainingApiService(TrainingJobStore())
130 created = service.create_training_job(valid_payload("delete-integration"))
131 job_id = str(created["id"])
132 artifact = service.list_artifacts(job_id)["artifacts"][0]
133 artifact_id = str(artifact["id"])
134 deleted_hashes = (
135 str(artifact["datasetHash"]),
136 str(artifact["artifactHash"]),
137 str(service.get_provenance(job_id)["trainingConfigHash"]),
138 )
139
140 receipt = service.delete_artifact(job_id, artifact_id)
141
142 self.assertTrue(receipt["deleted"])
143 self.assertTrue(service.verify_deleted_artifact_absence(deleted_hashes))
144 self.assertEqual(service.list_artifacts(job_id)["artifacts"], [])
145 with self.assertRaises(ApiError):
146 service.get_provenance(job_id)
147
148 def test_integration_sweep_over_mixed_policy_fixtures(self) -> None:
149 """Explicit sweep deletes expired artifacts and keeps unexpired artifacts."""
150
151 service = TrainingApiService(TrainingJobStore())
152 expired = service.create_training_job(
153 valid_payload(
154 "sweep-expired",
155 {"policyClass": "ephemeral", "ttlSeconds": 60},
156 )
157 )
158 retained = service.create_training_job(
159 valid_payload(
160 "sweep-retained",
161 {"policyClass": "extended", "ttlSeconds": 86_400},
162 )
163 )
164 sweep_at = datetime.now(UTC) + timedelta(seconds=120)
165
166 summary = service.sweep_expired_artifacts(sweep_at)
167
168 self.assertEqual(summary["deletedCount"], 1)
169 self.assertEqual(service.get_training_job(str(expired["id"]))["status"], "deleted")
170 self.assertEqual(service.get_training_job(str(retained["id"]))["status"], "succeeded")
171
172 def test_integration_bom_generation_over_real_project_files(self) -> None:
173 """The real project pyproject and lockfile generate an allowlisted BOM."""
174
175 entries = collect_entries(PROJECT_ROOT)
176 markdown = render_markdown(entries)
177 self.assertIn("scooling-lab", markdown)
178 self.assertIn("Apache-2.0", markdown)
179 self.assertNotIn("AGPL-3.0 |", markdown)
180
181
182 if __name__ == "__main__":
183 unittest.main()
File History 1 commit
sha256:ca1d0e687bff8686b37126eb8e4fb6d38b40e189352a4920d66ae89c7274e340 Add T4 job cancellation retry and validation Human minor ⚠ 41 days ago