Coverage for gco / services / central_queue_worker.py: 100.00%
247 statements
« prev ^ index » next coverage.py v7.13.5, created at 2026-09-14 22:07 +0000
1"""Lease-fenced regional worker for the global DynamoDB Job queue."""
3from __future__ import annotations
5import asyncio
6import contextlib
7import logging
8import os
9import uuid
10from dataclasses import dataclass
11from datetime import UTC, datetime
12from typing import Any
14from kubernetes.client.rest import ApiException
16from gco.services.manifest_processor import (
17 ManifestProcessor,
18 QueuedJobNotCreatedError,
19 RetryableQueuedJobApplyError,
20)
21from gco.services.spot_price_gate import SpotPriceGate, should_persist_observation
22from gco.services.structured_logging import sanitize_log_value
23from gco.services.template_store import JobStatus, JobStore
25# <pyflowchart-code-diagram> BEGIN - auto-inserted, do not edit
26# Generated at (UTC): 2026-08-30T12:00:00Z
27# Generated from Git commit: affbf6eccf3773dc3cfeac202e2cc6cbf92d4fc7
28# Flowchart(s) generated from this file:
29# * ``process_queued_jobs_once`` -> ``diagrams/code_diagrams/gco/services/central_queue_worker.process_queued_jobs_once.html``
30# (PNG: ``diagrams/code_diagrams/gco/services/central_queue_worker.process_queued_jobs_once.png``)
31# * ``reconcile_active_jobs_once`` -> ``diagrams/code_diagrams/gco/services/central_queue_worker.reconcile_active_jobs_once.html``
32# (PNG: ``diagrams/code_diagrams/gco/services/central_queue_worker.reconcile_active_jobs_once.png``)
33# Regenerate with ``SOURCE_DATE_EPOCH=<unix-seconds> GCO_DIAGRAM_SOURCE_COMMIT=<40-char-sha> python diagrams/generate.py --code-only``.
34# <pyflowchart-code-diagram> END
37logger = logging.getLogger(__name__)
39_TERMINAL_STATUSES = frozenset({JobStatus.SUCCEEDED.value, JobStatus.FAILED.value})
40_MAX_ERROR_LENGTH = 2_000
43def _utc_now() -> str:
44 return datetime.now(UTC).isoformat()
47def _bounded_error(value: object) -> str:
48 """Bound user/runtime error text before persisting it in DynamoDB."""
49 text = str(value)
50 return text if len(text) <= _MAX_ERROR_LENGTH else f"{text[:_MAX_ERROR_LENGTH]}...[truncated]"
53def _worker_identity(region: str) -> str:
54 """Return a process-unique owner including region and Kubernetes pod identity."""
55 pod_identity = os.getenv("POD_UID") or os.getenv("HOSTNAME") or "local"
56 return f"{region}/{pod_identity}/{uuid.uuid4().hex}"
59async def _lease_heartbeat(
60 store: JobStore,
61 *,
62 job_id: str,
63 target_region: str,
64 claimed_by: str,
65 claim_token: str,
66 claim_generation: int,
67 interval_seconds: float,
68 done: asyncio.Event,
69 lost: asyncio.Event,
70) -> None:
71 """Renew one active apply lease until completion or fencing."""
72 while not done.is_set():
73 try:
74 await asyncio.wait_for(done.wait(), timeout=interval_seconds)
75 return
76 except TimeoutError:
77 pass
79 try:
80 renewed = await asyncio.to_thread(
81 store.renew_claim,
82 job_id,
83 target_region,
84 claimed_by,
85 claim_token,
86 claim_generation,
87 )
88 except asyncio.CancelledError:
89 raise
90 except Exception: # noqa: BLE001 - loss must fence the in-flight result
91 logger.exception(
92 "Lease renewal failed for central queue job %s",
93 sanitize_log_value(job_id),
94 )
95 lost.set()
96 return
97 if not renewed:
98 logger.warning(
99 "Central queue job %s lost claim generation %d",
100 sanitize_log_value(job_id),
101 claim_generation,
102 )
103 lost.set()
104 return
107async def _stop_heartbeat(task: asyncio.Task[None], done: asyncio.Event) -> None:
108 done.set()
109 with contextlib.suppress(asyncio.CancelledError):
110 await task
113async def _defer_price_gated_job(
114 store: JobStore,
115 gate: SpotPriceGate,
116 queued_job: dict[str, Any],
117 job_id: str,
118) -> dict[str, Any] | None:
119 """Return a deferral record when the job's spot price gate is closed.
121 ``None`` means the job carries no gate or its gate is open — dispatch
122 proceeds. Closed gates optionally persist a throttled observation so
123 ``gco queue get`` can show why the job is waiting.
124 """
125 decision = await asyncio.to_thread(gate.evaluate, queued_job)
126 if decision is None or not decision.gated:
127 return None
128 if should_persist_observation(queued_job):
129 observed = (
130 f"{decision.observed_price:.6f}" if decision.observed_price is not None else "unknown"
131 )
132 try:
133 await asyncio.to_thread(
134 store.record_spot_gate_observation,
135 job_id,
136 observed_price=observed,
137 )
138 except Exception: # noqa: BLE001 - observations are advisory only
139 logger.exception(
140 "Failed to persist spot gate observation for %s",
141 sanitize_log_value(job_id),
142 )
143 logger.info(
144 "Deferring price-gated central queue job %s: %s",
145 sanitize_log_value(job_id),
146 decision.reason,
147 )
148 return {
149 "job_id": job_id,
150 "status": "price_gated",
151 "instance_type": decision.instance_type,
152 "max_spot_price": decision.max_price,
153 "observed_spot_price": decision.observed_price,
154 "reason": decision.reason,
155 }
158async def process_queued_jobs_once(
159 processor: ManifestProcessor,
160 store: JobStore,
161 *,
162 limit: int,
163 owner_id: str | None = None,
164 lease_renewal_seconds: float | None = None,
165 stop_event: asyncio.Event | None = None,
166 spot_gate: SpotPriceGate | None = None,
167) -> tuple[int, list[dict[str, Any]]]:
168 """Claim and apply one bounded batch for the processor's region.
170 Price-gated jobs whose spot cap is not currently met are deferred (left
171 queued) without consuming the apply budget. To keep a run of gated
172 high-priority jobs from starving dispatchable work behind them, the
173 candidate fetch is wider than ``limit``; at most ``limit`` jobs are
174 claimed/applied per pass.
175 """
176 owner = owner_id or _worker_identity(processor.region)
177 renewal_seconds = lease_renewal_seconds or max(
178 5.0,
179 min(float(store.claim_lease_seconds) / 3.0, 60.0),
180 )
181 gate = spot_gate or SpotPriceGate(processor.region)
182 migration = await asyncio.to_thread(
183 store.migrate_legacy_records_for_region,
184 processor.region,
185 max(limit * 20, 100),
186 )
187 migrated = int(migration.get("migrated", 0))
188 migration_failed = int(migration.get("failed", 0))
189 if migrated or migration_failed:
190 logger.warning(
191 "Central queue migration for %s: backfilled=%d safely_failed=%d complete=%s",
192 processor.region,
193 migrated,
194 migration_failed,
195 migration.get("complete", False),
196 )
197 fetch_limit = min(max(limit * 4, 20), 100)
198 queued_jobs = await asyncio.to_thread(
199 store.get_queued_jobs_for_region,
200 processor.region,
201 fetch_limit,
202 )
203 processed: list[dict[str, Any]] = []
204 attempted = 0
206 for queued_job in queued_jobs:
207 if stop_event is not None and stop_event.is_set():
208 break
209 if attempted >= limit:
210 break
211 job_id = str(queued_job.get("job_id", ""))
212 if not job_id:
213 logger.error("Ignoring central queue record without a job_id")
214 continue
216 deferral = await _defer_price_gated_job(store, gate, queued_job, job_id)
217 if deferral is not None:
218 processed.append(deferral)
219 continue
221 claimed: dict[str, Any] | None = None
222 try:
223 try:
224 claimed = await asyncio.to_thread(
225 store.claim_job,
226 job_id,
227 processor.region,
228 owner,
229 )
230 except Exception:
231 # Unlike an explicit None (a known conditional-write loss), a
232 # transport/service exception can be post-commit. Count that
233 # ambiguous outcome so this pass cannot exceed its claim bound.
234 attempted += 1
235 raise
236 if not claimed:
237 continue
238 # Only a claim this worker actually acquired consumes the apply
239 # budget. A conditional-write race lost to another worker must not
240 # prevent later candidates from being considered in this pass.
241 attempted += 1
242 claim_token = str(claimed.get("claim_token") or "")
243 claim_generation = int(claimed.get("claim_generation", 0))
244 if not claim_token or claim_generation <= 0:
245 raise RuntimeError("JobStore returned an incomplete fenced claim")
247 applying = await asyncio.to_thread(
248 store.transition_job,
249 job_id,
250 target_region=processor.region,
251 expected_status=JobStatus.CLAIMED,
252 status=JobStatus.APPLYING,
253 message="Applying deterministic Kubernetes Job",
254 claimed_by=owner,
255 claim_token=claim_token,
256 claim_generation=claim_generation,
257 )
258 if not applying:
259 processed.append({"job_id": job_id, "status": "fenced"})
260 continue
262 manifest = claimed.get("manifest")
263 namespace = claimed.get("namespace")
264 if not isinstance(manifest, dict) or not isinstance(namespace, str):
265 raise QueuedJobNotCreatedError(
266 "Queued job contains an invalid manifest or namespace"
267 )
269 heartbeat_done = asyncio.Event()
270 claim_lost = asyncio.Event()
271 heartbeat = asyncio.create_task(
272 _lease_heartbeat(
273 store,
274 job_id=job_id,
275 target_region=processor.region,
276 claimed_by=owner,
277 claim_token=claim_token,
278 claim_generation=claim_generation,
279 interval_seconds=renewal_seconds,
280 done=heartbeat_done,
281 lost=claim_lost,
282 ),
283 name=f"central-queue-lease-{job_id}",
284 )
285 try:
286 resource = await asyncio.to_thread(
287 processor.apply_queued_job,
288 manifest,
289 namespace,
290 job_id,
291 )
292 finally:
293 await _stop_heartbeat(heartbeat, heartbeat_done)
295 if claim_lost.is_set():
296 processed.append({"job_id": job_id, "status": "fenced"})
297 continue
299 pending = await asyncio.to_thread(
300 store.transition_job,
301 job_id,
302 target_region=processor.region,
303 expected_status=JobStatus.APPLYING,
304 status=JobStatus.PENDING,
305 message="Applied to Kubernetes, waiting for scheduling",
306 k8s_job_name=resource.name,
307 k8s_job_namespace=resource.namespace,
308 k8s_job_uid=resource.uid,
309 claimed_by=owner,
310 claim_token=claim_token,
311 claim_generation=claim_generation,
312 )
313 if pending:
314 processed.append(
315 {
316 "job_id": job_id,
317 "status": "applied",
318 "k8s_job_name": resource.name,
319 "k8s_job_uid": resource.uid,
320 }
321 )
322 else:
323 processed.append({"job_id": job_id, "status": "fenced"})
324 except RetryableQueuedJobApplyError as exc:
325 # The API may have accepted a deterministic create before the
326 # response was lost. Keep APPLYING fenced and let lease recovery
327 # return it to QUEUED, where the next attempt adopts by full queue ID.
328 error = _bounded_error(exc)
329 logger.warning(
330 "Deferring central queue job %s after an inconclusive Kubernetes result",
331 sanitize_log_value(job_id),
332 )
333 processed.append({"job_id": job_id, "status": "retryable", "error": error})
334 except Exception as exc: # noqa: BLE001 - isolate malformed/permanent records
335 error = _bounded_error(exc)
336 logger.exception(
337 "Failed to process central queue job %s",
338 sanitize_log_value(job_id),
339 )
340 claim = claimed
341 token = claim.get("claim_token") if isinstance(claim, dict) else None
342 generation = claim.get("claim_generation") if isinstance(claim, dict) else None
343 if token and generation:
344 transition_options: dict[str, Any] = {}
345 if isinstance(exc, QueuedJobNotCreatedError):
346 transition_options["workload_not_created"] = True
347 try:
348 failed = await asyncio.to_thread(
349 store.transition_job,
350 job_id,
351 target_region=processor.region,
352 expected_status=JobStatus.APPLYING,
353 status=JobStatus.FAILED,
354 message="Failed to apply deterministic Kubernetes Job",
355 error=error,
356 claimed_by=owner,
357 claim_token=str(token),
358 claim_generation=int(generation),
359 **transition_options,
360 )
361 except Exception: # noqa: BLE001 - lease recovery remains the fallback
362 logger.exception(
363 "Unable to persist failure for central queue job %s",
364 sanitize_log_value(job_id),
365 )
366 failed = None
367 processed.append(
368 {"job_id": job_id, "status": "failed" if failed else "fenced", "error": error}
369 )
370 else:
371 processed.append({"job_id": job_id, "status": "fenced", "error": error})
373 return len(queued_jobs), processed
376def _observed_job_state(job: Any) -> tuple[str, str | None]:
377 """Map a Kubernetes Job object to the central queue lifecycle."""
378 status = job.status
379 for condition in status.conditions or []:
380 if condition.type == "Complete" and condition.status == "True":
381 return JobStatus.SUCCEEDED.value, None
382 if condition.type == "Failed" and condition.status == "True":
383 detail = condition.message or condition.reason or "Kubernetes Job failed"
384 return JobStatus.FAILED.value, _bounded_error(detail)
385 if (status.active or 0) > 0:
386 return JobStatus.RUNNING.value, None
387 return JobStatus.PENDING.value, None
390async def reconcile_active_jobs_once(
391 processor: ManifestProcessor,
392 store: JobStore,
393 *,
394 limit: int,
395 stop_event: asyncio.Event | None = None,
396) -> int:
397 """Reconcile pending/running queue records with their exact Kubernetes Jobs."""
398 jobs = await asyncio.to_thread(store.get_active_jobs_for_region, processor.region, limit)
399 transitions = 0
401 for queued_job in jobs:
402 if stop_event is not None and stop_event.is_set():
403 break
404 job_id = str(queued_job.get("job_id", ""))
405 job_name = queued_job.get("k8s_job_name")
406 namespace = queued_job.get("k8s_job_namespace")
407 expected_uid = str(queued_job.get("k8s_job_uid") or "")
408 current = str(queued_job.get("status") or "")
409 if (
410 not job_id
411 or not isinstance(job_name, str)
412 or not isinstance(namespace, str)
413 or not expected_uid
414 or current not in {JobStatus.PENDING.value, JobStatus.RUNNING.value}
415 ):
416 logger.error(
417 "Ignoring active queue record %s with incomplete Kubernetes identity", job_id
418 )
419 continue
421 observed: str
422 error: str | None
423 try:
424 k8s_job = await asyncio.to_thread(
425 processor.read_queued_job,
426 job_name,
427 namespace,
428 )
429 except ApiException as exc:
430 if exc.status == 404:
431 observed, error = JobStatus.FAILED.value, "Kubernetes Job no longer exists"
432 else:
433 logger.warning(
434 "Unable to reconcile central queue job %s: %s",
435 sanitize_log_value(job_id),
436 exc,
437 )
438 continue
439 except Exception as exc: # noqa: BLE001 - reconciliation is best-effort per record
440 logger.warning(
441 "Unable to reconcile central queue job %s: %s",
442 sanitize_log_value(job_id),
443 exc,
444 )
445 continue
446 else:
447 actual_uid = str(getattr(k8s_job.metadata, "uid", "") or "")
448 if actual_uid != expected_uid:
449 observed, error = JobStatus.FAILED.value, "Kubernetes Job identity changed"
450 else:
451 observed, error = _observed_job_state(k8s_job)
453 if observed == current:
454 continue
455 if observed not in {
456 status.value
457 for status in (
458 JobStatus.RUNNING,
459 JobStatus.SUCCEEDED,
460 JobStatus.FAILED,
461 )
462 }:
463 # Running Jobs never regress to pending.
464 continue
466 message = {
467 JobStatus.RUNNING.value: "Kubernetes Job is running",
468 JobStatus.SUCCEEDED.value: "Kubernetes Job completed successfully",
469 JobStatus.FAILED.value: "Kubernetes Job failed",
470 }[observed]
471 updated = await asyncio.to_thread(
472 store.transition_job,
473 job_id,
474 target_region=processor.region,
475 expected_status=current,
476 status=observed,
477 message=message,
478 error=error,
479 expected_k8s_uid=expected_uid,
480 )
481 if updated:
482 transitions += 1
484 return transitions
487@dataclass
488class CentralQueueWorker:
489 """Continuously activate and reconcile Jobs for one regional cluster."""
491 processor: ManifestProcessor
492 store: JobStore
493 poll_interval_seconds: float = 10.0
494 batch_size: int = 5
495 reconcile_limit: int = 100
496 lease_renewal_seconds: float | None = None
497 owner_id: str | None = None
499 def __post_init__(self) -> None:
500 self.poll_interval_seconds = min(max(float(self.poll_interval_seconds), 1.0), 300.0)
501 self.batch_size = min(max(int(self.batch_size), 1), 20)
502 self.reconcile_limit = min(max(int(self.reconcile_limit), 1), 500)
503 default_renewal = max(5.0, min(float(self.store.claim_lease_seconds) / 3.0, 60.0))
504 requested_renewal = (
505 default_renewal
506 if self.lease_renewal_seconds is None
507 else float(self.lease_renewal_seconds)
508 )
509 self.lease_renewal_seconds = min(
510 max(requested_renewal, 1.0),
511 max(float(self.store.claim_lease_seconds) / 2.0, 1.0),
512 )
513 self.owner_id = self.owner_id or _worker_identity(self.processor.region)
514 self._stop_event = asyncio.Event()
515 # One gate per worker so its price cache spans polling passes.
516 self._spot_gate = SpotPriceGate(self.processor.region)
517 self.running = False
518 self.stopping = False
519 self.last_pass_started_at: str | None = None
520 self.last_successful_pass_at: str | None = None
521 self.last_error: str | None = None
523 def stop(self) -> None:
524 """Stop accepting new records and finish the current bounded SDK call."""
525 self.stopping = True
526 self._stop_event.set()
528 def health(self) -> dict[str, Any]:
529 """Return operator-safe worker health without exposing claim tokens."""
530 return {
531 "enabled": True,
532 "running": self.running,
533 "stopping": self.stopping,
534 "owner": self.owner_id,
535 "last_pass_started_at": self.last_pass_started_at,
536 "last_successful_pass_at": self.last_successful_pass_at,
537 "last_error": self.last_error,
538 }
540 async def run(self) -> None:
541 """Run until stopped; isolate pass failures and keep polling."""
542 self.running = True
543 logger.info(
544 "Central queue worker started for %s (interval=%ss, batch=%d)",
545 self.processor.region,
546 self.poll_interval_seconds,
547 self.batch_size,
548 )
549 try:
550 while not self._stop_event.is_set():
551 self.last_pass_started_at = _utc_now()
552 try:
553 recovered = await asyncio.to_thread(
554 self.store.requeue_expired_jobs,
555 self.processor.region,
556 self.reconcile_limit,
557 )
558 if self._stop_event.is_set():
559 break
560 polled, processed = await process_queued_jobs_once(
561 self.processor,
562 self.store,
563 limit=self.batch_size,
564 owner_id=self.owner_id,
565 lease_renewal_seconds=self.lease_renewal_seconds,
566 stop_event=self._stop_event,
567 spot_gate=self._spot_gate,
568 )
569 transitions = 0
570 if not self._stop_event.is_set():
571 transitions = await reconcile_active_jobs_once(
572 self.processor,
573 self.store,
574 limit=self.reconcile_limit,
575 stop_event=self._stop_event,
576 )
577 self.last_successful_pass_at = _utc_now()
578 self.last_error = None
579 if recovered or polled or transitions:
580 logger.info(
581 "Central queue pass: recovered=%d polled=%d processed=%d transitions=%d",
582 recovered,
583 polled,
584 len(processed),
585 transitions,
586 )
587 except asyncio.CancelledError:
588 raise
589 except Exception as exc: # noqa: BLE001 - transient pass failures are retried
590 self.last_error = _bounded_error(exc)
591 logger.exception("Central queue worker pass failed")
593 with contextlib.suppress(TimeoutError):
594 await asyncio.wait_for(
595 self._stop_event.wait(),
596 timeout=self.poll_interval_seconds,
597 )
598 finally:
599 self.running = False
600 logger.info("Central queue worker stopped for %s", self.processor.region)