Coverage for scripts / live_release_validation / checks / central_queue.py: 100.00%
209 statements
« prev ^ index » next coverage.py v7.13.5, created at 2026-09-14 22:07 +0000
1"""DynamoDB central-queue identity, reconciliation, and polling helpers."""
3from __future__ import annotations
5import copy
6import hashlib
7import re
8import time
9import uuid
10from typing import Any, cast
11from urllib.parse import quote
13from boto3.dynamodb.types import TypeDeserializer
15from ..checks.jobs import (
16 _central_workload_identity,
17 _job_appearance_timeout,
18 _load_manifest,
19 _response_json,
20 _run_token,
21)
22from ..constants import (
23 _CENTRAL_QUEUE_IDEMPOTENCY_NAMESPACE,
24 _PATH_JOB_LABEL,
25 _TERMINAL_QUEUE_STATUSES,
26)
27from ..models import RunContext
30def _central_manifest(ctx: RunContext) -> tuple[dict[str, Any], str, str, str]:
31 manifests, _name, namespace = _load_manifest(ctx, "api-smoke-job.yaml")
32 manifest = copy.deepcopy(manifests[0])
33 token = _run_token(ctx.settings.run_id)
34 name = f"gco-live-ddb-{token}"[:63].rstrip("-")
35 marker = f"GCO_LIVE_DDB_{token}"
36 manifest["metadata"]["name"] = name
37 manifest["metadata"]["labels"][_PATH_JOB_LABEL] = "dynamodb"
38 manifest["spec"]["template"]["metadata"]["labels"][_PATH_JOB_LABEL] = "dynamodb"
39 manifest["spec"]["template"]["spec"]["containers"][0]["command"] = [
40 "sh",
41 "-c",
42 f"echo {marker}",
43 ]
44 return manifest, name, namespace, marker
47def _deserialize_item(item: dict[str, Any]) -> dict[str, Any]:
48 deserializer = TypeDeserializer()
49 return {key: deserializer.deserialize(value) for key, value in item.items()}
52def _read_central_job_item(ctx: RunContext, job_id: str) -> dict[str, Any]:
53 table_name = f"{ctx.config.project_name}-jobs"
54 response = ctx.session.client("dynamodb", region_name=ctx.config.global_region).get_item(
55 TableName=table_name,
56 Key={"job_id": {"S": job_id}},
57 ConsistentRead=True,
58 )
59 item = response.get("Item")
60 if not item:
61 raise RuntimeError(f"DynamoDB item {job_id} was not found in {table_name}")
62 return _deserialize_item(item)
65def _central_queue_job_id(idempotency_key: str) -> str:
66 return str(uuid.uuid5(_CENTRAL_QUEUE_IDEMPOTENCY_NAMESPACE, idempotency_key))
69def _central_queue_kubernetes_job_name(original_name: str, job_id: str) -> str:
70 suffix = hashlib.sha256(job_id.encode("utf-8")).hexdigest()[:16]
71 prefix = re.sub(r"[^a-z0-9-]+", "-", original_name.lower()).strip("-")
72 prefix = prefix[: 63 - len(suffix) - 1].rstrip("-") or "gco-job"
73 return f"{prefix}-{suffix}"
76def _central_persisted_kubernetes_identity(
77 job: dict[str, Any],
78 *,
79 required: bool,
80) -> tuple[str, str, str] | None:
81 raw = (
82 job.get("k8s_job_name"),
83 job.get("k8s_job_namespace"),
84 job.get("k8s_job_uid"),
85 )
86 populated = [value is not None for value in raw]
87 if not any(populated):
88 if required:
89 raise RuntimeError("Central DynamoDB record omitted worker Kubernetes identity")
90 return None
91 if not all(populated):
92 raise RuntimeError("Central DynamoDB record contains a partial Kubernetes identity")
93 identity = tuple(str(value or "") for value in raw)
94 if not all(identity):
95 raise RuntimeError("Central DynamoDB record contains an empty Kubernetes identity field")
96 return cast(tuple[str, str, str], identity)
99def _validate_central_checkpoint_kubernetes_identity(
100 central_record: dict[str, Any],
101 identity: dict[str, str],
102) -> None:
103 """Reject partial or conflicting central identity before mutating either record."""
104 previous = {key: central_record.get(key) for key in identity}
105 populated = [value is not None for value in previous.values()]
106 if any(populated) and not all(populated):
107 raise RuntimeError("Checkpoint central record contains a partial Kubernetes identity")
108 source = central_record.get("k8s_identity_source")
109 if source is not None and source != "dynamodb":
110 raise RuntimeError("Central checkpoint Kubernetes identity has an unexpected source")
111 if source is not None and not any(populated):
112 raise RuntimeError("Central checkpoint identity source has no Kubernetes identity")
113 for key, value in identity.items():
114 if previous[key] is not None and previous[key] != value:
115 raise RuntimeError(f"Central checkpoint Kubernetes identity changed: {key}")
118def _reconcile_central_workload_identity(
119 ctx: RunContext,
120 central_record: dict[str, Any],
121 persisted_job: dict[str, Any],
122 *,
123 workload_record: dict[str, Any] | None = None,
124 require_identity: bool = True,
125) -> dict[str, Any]:
126 """Bind exact worker evidence without mutating requested replay identity."""
127 _validate_central_job_identity(central_record, persisted_job)
128 identity = _central_persisted_kubernetes_identity(
129 persisted_job,
130 required=require_identity,
131 )
132 record = workload_record or _central_workload_record(ctx, central_record)
133 if identity is None:
134 return record
136 actual_name, actual_namespace, actual_uid = identity
137 expected_name = _central_queue_kubernetes_job_name(
138 str(central_record["job_name"]),
139 str(central_record["job_id"]),
140 )
141 if actual_name != expected_name:
142 raise RuntimeError(
143 f"Central worker persisted unexpected Kubernetes Job name {actual_name!r}; "
144 f"expected {expected_name!r}"
145 )
146 if actual_namespace != central_record["namespace"]:
147 raise RuntimeError("Central worker persisted a different Kubernetes namespace")
149 central_identity = {
150 "k8s_job_name": actual_name,
151 "k8s_job_namespace": actual_namespace,
152 "k8s_job_uid": actual_uid,
153 }
154 _validate_central_checkpoint_kubernetes_identity(central_record, central_identity)
155 ctx.bind_central_job_identity(
156 record,
157 job_id=str(central_record["job_id"]),
158 name=actual_name,
159 namespace=actual_namespace,
160 uid=actual_uid,
161 appearance_timeout_seconds=_job_appearance_timeout(ctx),
162 )
163 with ctx.state_lock:
164 _validate_central_checkpoint_kubernetes_identity(central_record, central_identity)
165 central_record.update(central_identity)
166 central_record["k8s_identity_source"] = "dynamodb"
167 central_record["workload_appearance_deadline"] = record.get("appearance_deadline")
168 ctx.persist_callback(ctx.checkpoint)
169 return record
172def _register_central_job(
173 ctx: RunContext,
174 *,
175 job_id: str,
176 idempotency_key: str,
177 record: dict[str, Any],
178 marker: str,
179 body: dict[str, Any],
180) -> dict[str, Any]:
181 candidate = {
182 "job_id": job_id,
183 "idempotency_key": idempotency_key,
184 "job_name": record["name"],
185 "namespace": record["namespace"],
186 "target_region": record["region"],
187 "transport_region": record.get("transport_region"),
188 "marker": marker,
189 "body": copy.deepcopy(body),
190 }
191 with ctx.state_lock:
192 raw_central_jobs = ctx.checkpoint.state.setdefault("central_jobs", [])
193 if not isinstance(raw_central_jobs, list) or any(
194 not isinstance(item, dict) for item in raw_central_jobs
195 ):
196 raise RuntimeError("Checkpoint central_jobs must be a list of objects")
197 central_jobs = cast(list[dict[str, Any]], raw_central_jobs)
198 matches = [
199 item
200 for item in central_jobs
201 if item.get("job_id") == job_id or item.get("idempotency_key") == idempotency_key
202 ]
203 if len(matches) > 1:
204 raise RuntimeError(f"Checkpoint contains duplicate central Job records for {job_id}")
205 if matches:
206 central_record = matches[0]
207 for key, value in candidate.items():
208 if central_record.get(key) != value:
209 raise RuntimeError(f"Central Job identity changed for {job_id}: {key}")
210 else:
211 central_record = {
212 **candidate,
213 "submission_state": str(record.get("submission_state") or "prepared"),
214 "appearance_deadline": record.get("appearance_deadline"),
215 "cleanup_complete": False,
216 }
217 central_jobs.append(central_record)
218 if central_record.get("appearance_deadline") is None:
219 central_record["appearance_deadline"] = record.get("appearance_deadline")
220 ctx.persist_callback(ctx.checkpoint)
221 return central_record
224def _central_workload_record(
225 ctx: RunContext,
226 central_record: dict[str, Any],
227) -> dict[str, Any]:
228 raw_jobs = ctx.checkpoint.state.get("jobs", [])
229 if not isinstance(raw_jobs, list) or any(not isinstance(item, dict) for item in raw_jobs):
230 raise RuntimeError("Checkpoint jobs must be a list of objects")
231 jobs = cast(list[dict[str, Any]], raw_jobs)
232 matches = [
233 record
234 for record in jobs
235 if record.get("path") == "dynamodb"
236 and record.get("name") == central_record.get("job_name")
237 and record.get("namespace") == central_record.get("namespace")
238 and record.get("region") == central_record.get("target_region")
239 and record.get("transport_region") == central_record.get("transport_region")
240 ]
241 if len(matches) != 1:
242 raise RuntimeError(
243 "Central queue record does not resolve to exactly one checkpointed workload: "
244 f"{central_record.get('job_id')}"
245 )
246 return matches[0]
249def _validate_central_job_identity(central_record: dict[str, Any], job: dict[str, Any]) -> None:
250 expected = {
251 "job_id": central_record["job_id"],
252 "job_name": central_record["job_name"],
253 "namespace": central_record["namespace"],
254 "target_region": central_record["target_region"],
255 }
256 for key, value in expected.items():
257 if str(job.get(key) or "") != str(value):
258 raise RuntimeError(f"Central queue returned a different {key} for {value!r}")
259 observed_key = job.get("idempotency_key")
260 if observed_key is not None and observed_key != central_record["idempotency_key"]:
261 raise RuntimeError("Central queue idempotency key changed")
264def _reconcile_central_cleanup_workload(
265 ctx: RunContext,
266 central_record: dict[str, Any],
267 persisted_job: dict[str, Any],
268) -> tuple[dict[str, Any], bool]:
269 """Reconcile workload authority from terminal DynamoDB cleanup evidence."""
270 _validate_central_job_identity(central_record, persisted_job)
271 terminal_status = str(persisted_job.get("status") or "unknown")
272 if terminal_status not in _TERMINAL_QUEUE_STATUSES:
273 raise RuntimeError("Central cleanup evidence is not terminal")
275 worker_proved_not_created = persisted_job.get("workload_not_created") is True
276 if worker_proved_not_created and terminal_status != "failed":
277 raise RuntimeError(
278 "Central worker no-workload proof is valid only for a failed queue record"
279 )
280 persisted_identity = _central_persisted_kubernetes_identity(
281 persisted_job,
282 required=terminal_status == "succeeded"
283 or (terminal_status == "failed" and not worker_proved_not_created),
284 )
285 workload_record = _central_workload_record(ctx, central_record)
287 if worker_proved_not_created:
288 if persisted_identity is not None:
289 raise RuntimeError(
290 "Failed central Job has both no-workload proof and Kubernetes identity"
291 )
292 if _central_workload_identity(workload_record) is not None:
293 raise RuntimeError(
294 "Worker-proven uncreated central Job already has checkpointed workload identity"
295 )
296 if _central_workload_identity(central_record) is not None:
297 raise RuntimeError(
298 "Worker-proven uncreated central Job already has central checkpoint identity"
299 )
300 if workload_record.get("uid"):
301 raise RuntimeError(
302 "Worker-proven uncreated central Job already has Kubernetes UID authority"
303 )
304 job_id = str(central_record["job_id"])
305 state = str(workload_record.get("submission_state") or "registered")
306 prior_proof = workload_record.get("central_worker_not_created_job_id")
307 if state == "deleted":
308 if prior_proof != job_id:
309 raise RuntimeError(
310 "Deleted central workload lacks matching worker no-workload proof"
311 )
312 else:
313 ctx.mark_central_job_not_created_by_worker(workload_record, job_id=job_id)
314 return workload_record, True
316 if terminal_status != "cancelled":
317 return (
318 _reconcile_central_workload_identity(
319 ctx,
320 central_record,
321 persisted_job,
322 workload_record=workload_record,
323 ),
324 False,
325 )
327 if persisted_identity is not None or _central_workload_identity(workload_record) is not None:
328 raise RuntimeError("Cancelled-before-claim central Job unexpectedly has workload identity")
329 if workload_record.get("uid"):
330 raise RuntimeError(
331 "Cancelled-before-claim central Job already has Kubernetes UID authority"
332 )
333 job_id = str(central_record["job_id"])
334 state = str(workload_record.get("submission_state") or "registered")
335 prior_proof = workload_record.get("central_cancelled_before_claim_job_id")
336 if state == "deleted":
337 if prior_proof != job_id:
338 raise RuntimeError("Deleted central workload lacks matching cancellation proof")
339 else:
340 ctx.mark_central_job_cancelled_before_claim(workload_record, job_id=job_id)
341 return workload_record, True
344def _get_central_queue_job(
345 ctx: RunContext, central_record: dict[str, Any]
346) -> dict[str, Any] | None:
347 response = ctx.aws_client.make_authenticated_request(
348 method="GET",
349 path=f"/api/v1/queue/jobs/{quote(str(central_record['job_id']), safe='')}",
350 target_region=central_record.get("transport_region"),
351 )
352 if response.status_code == 404:
353 return None
354 if not response.ok:
355 raise RuntimeError(f"Central queue lookup failed: {response.status_code} {response.text}")
356 data = _response_json(response, "Central queue lookup")
357 job = data.get("job")
358 if not isinstance(job, dict):
359 raise RuntimeError("Central queue lookup omitted job")
360 _validate_central_job_identity(central_record, job)
361 return job
364def _wait_for_central_queue_appearance(
365 ctx: RunContext,
366 central_record: dict[str, Any],
367 *,
368 raise_on_timeout: bool = True,
369) -> dict[str, Any] | None:
370 raw_deadline = central_record.get("appearance_deadline")
371 if raw_deadline is None:
372 deadline = time.time() + _job_appearance_timeout(ctx)
373 central_record["appearance_deadline"] = deadline
374 ctx.persist()
375 else:
376 deadline = float(raw_deadline)
377 while True:
378 job = _get_central_queue_job(ctx, central_record)
379 if job is not None:
380 return job
381 if time.time() >= deadline:
382 if raise_on_timeout:
383 raise TimeoutError(
384 f"Central queue job {central_record['job_id']} did not appear before "
385 "the bounded submission deadline"
386 )
387 return None
388 time.sleep(ctx.settings.poll_interval_seconds)
391def _wait_for_central_queue_terminal(
392 ctx: RunContext,
393 central_record: dict[str, Any],
394) -> tuple[dict[str, Any], list[dict[str, Any]]]:
395 deadline = time.monotonic() + ctx.settings.job_timeout_seconds
396 history: list[dict[str, Any]] = []
397 while True:
398 job = _get_central_queue_job(ctx, central_record)
399 if job is None:
400 raise RuntimeError(
401 f"Central queue job {central_record['job_id']} disappeared after observation"
402 )
403 status = str(job.get("status") or "unknown")
404 history.append({"status": status, "at": time.time()})
405 if status in _TERMINAL_QUEUE_STATUSES:
406 return job, history
407 if time.monotonic() >= deadline:
408 raise TimeoutError(
409 f"Central queue job {central_record['job_id']} did not reach a terminal status"
410 )
411 time.sleep(ctx.settings.poll_interval_seconds)