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