Coverage for gco / services / template_store.py: 100.00%
790 statements
« prev ^ index » next coverage.py v7.13.5, created at 2026-09-14 22:07 +0000
1"""
2DynamoDB-backed storage for job templates, webhooks, and job records.
4This module provides persistent storage for:
5- Job templates: Reusable job configurations with parameter substitution
6- Webhooks: Event notification registrations
7- Job records: Centralized job tracking with status updates
9Tables are created in the global stack and accessed from all regional services.
11Region Configuration:
12 DynamoDB tables are deployed in the global region (e.g., us-east-2) but
13 accessed from regional services (e.g., us-east-1). The region is determined
14 by checking environment variables in this order:
15 1. DYNAMODB_REGION - Explicitly set for DynamoDB access
16 2. GLOBAL_REGION - The global stack's region
17 3. AWS_REGION - Fallback to current region
19Job Queue Architecture:
20 1. Jobs are submitted to the jobs table with target_region and status="queued"
21 2. Regional manifest processors poll for jobs targeting their region
22 3. Processor claims job (status="claimed"), applies to K8s, updates status
23 4. Status updates flow back to DynamoDB for global visibility
24"""
26from __future__ import annotations
28import base64
29import binascii
30import json
31import logging
32import os
33import uuid
34from collections.abc import Collection
35from datetime import UTC, datetime, timedelta
36from enum import StrEnum
37from typing import Any
39import boto3
40from botocore.config import Config
41from botocore.exceptions import ClientError
43# <pyflowchart-code-diagram> BEGIN - auto-inserted, do not edit
44# Generated at (UTC): 2026-09-05T22:58:10Z
45# Generated from Git commit: 745b3fa3a9af9380bfe2797a5d9716fe8ce3a557
46# Flowchart(s) generated from this file:
47# * ``JobStore.claim_job`` -> ``diagrams/code_diagrams/gco/services/template_store.JobStore_claim_job.html``
48# (PNG: ``diagrams/code_diagrams/gco/services/template_store.JobStore_claim_job.png``)
49# * ``JobStore.transition_job`` -> ``diagrams/code_diagrams/gco/services/template_store.JobStore_transition_job.html``
50# (PNG: ``diagrams/code_diagrams/gco/services/template_store.JobStore_transition_job.png``)
51# Regenerate with ``SOURCE_DATE_EPOCH=<unix-seconds> GCO_DIAGRAM_SOURCE_COMMIT=<40-char-sha> python diagrams/generate.py --code-only``.
52# <pyflowchart-code-diagram> END
55logger = logging.getLogger(__name__)
57_DEFAULT_CLAIM_LEASE_SECONDS = 5 * 60
58_MIN_CLAIM_LEASE_SECONDS = 30
59_MAX_CLAIM_LEASE_SECONDS = 60 * 60
60_MAX_LIST_EVALUATED_ITEMS = 20_000
61_MAX_LEGACY_MIGRATION_EVALUATED_ITEMS = 1_000
62_LEGACY_REGION_STATUS_INDEX = "region-status-index"
63# One worker-facing GSI serves queue priority and lease recovery. Existing
64# deployments gain only this index in the compatibility release because
65# DynamoDB permits one GSI create/delete per table update.
66_REGION_STATUS_WORK_INDEX = "region-status-work-index"
67_TERMINAL_JOB_STATUSES = frozenset({"succeeded", "failed", "cancelled"})
70def _utc_now_iso() -> str:
71 """Return current UTC time in ISO format with Z suffix."""
72 return datetime.now(UTC).isoformat().replace("+00:00", "Z")
75def _claim_lease_expiry_iso(lease_seconds: int) -> str:
76 """Return a bounded lease expiry for crash-safe regional claims."""
77 return (datetime.now(UTC) + timedelta(seconds=lease_seconds)).isoformat().replace("+00:00", "Z")
80class JobSubmissionConflict(RuntimeError):
81 """A job ID or idempotency key was reused for a different submission."""
84class JobStatus(StrEnum):
85 """Job status values for the centralized job store."""
87 QUEUED = "queued" # Submitted, waiting for regional pickup
88 CLAIMED = "claimed" # Claimed by a regional processor
89 APPLYING = "applying" # Being applied to Kubernetes
90 PENDING = "pending" # Applied, waiting for pod scheduling
91 RUNNING = "running" # Pod(s) running
92 SUCCEEDED = "succeeded" # Job completed successfully
93 FAILED = "failed" # Job failed
94 CANCELLED = "cancelled" # Job was cancelled
97_ALLOWED_JOB_TRANSITIONS: dict[str, frozenset[str]] = {
98 JobStatus.QUEUED.value: frozenset({JobStatus.CLAIMED.value, JobStatus.CANCELLED.value}),
99 JobStatus.CLAIMED.value: frozenset({JobStatus.APPLYING.value, JobStatus.FAILED.value}),
100 JobStatus.APPLYING.value: frozenset({JobStatus.PENDING.value, JobStatus.FAILED.value}),
101 JobStatus.PENDING.value: frozenset(
102 {JobStatus.RUNNING.value, JobStatus.SUCCEEDED.value, JobStatus.FAILED.value}
103 ),
104 JobStatus.RUNNING.value: frozenset({JobStatus.SUCCEEDED.value, JobStatus.FAILED.value}),
105 JobStatus.SUCCEEDED.value: frozenset(),
106 JobStatus.FAILED.value: frozenset(),
107 JobStatus.CANCELLED.value: frozenset(),
108}
111class TemplateStore:
112 """DynamoDB-backed store for job templates."""
114 def __init__(self, table_name: str | None = None, region: str | None = None):
115 """Initialize the template store.
117 Args:
118 table_name: DynamoDB table name. Defaults to env var TEMPLATES_TABLE_NAME.
119 region: AWS region for DynamoDB. Defaults to env var DYNAMODB_REGION,
120 then GLOBAL_REGION, then AWS_REGION.
121 """
122 self.table_name = table_name or os.getenv("TEMPLATES_TABLE_NAME", "gco-job-templates")
123 # DynamoDB tables are in the global region, not the regional cluster region
124 self.region = (
125 region
126 or os.getenv("DYNAMODB_REGION")
127 or os.getenv("GLOBAL_REGION")
128 or os.getenv("AWS_REGION", "us-east-1")
129 )
130 self._dynamodb = boto3.resource("dynamodb", region_name=self.region)
131 self._table = self._dynamodb.Table(self.table_name)
133 def list_templates(self) -> list[dict[str, Any]]:
134 """List all templates."""
135 try:
136 response = self._table.scan(
137 ProjectionExpression="template_name, description, created_at, updated_at"
138 )
139 items = response.get("Items", [])
141 # Handle pagination
142 while "LastEvaluatedKey" in response:
143 response = self._table.scan(
144 ProjectionExpression="template_name, description, created_at, updated_at",
145 ExclusiveStartKey=response["LastEvaluatedKey"],
146 )
147 items.extend(response.get("Items", []))
149 return [
150 {
151 "name": item["template_name"],
152 "description": item.get("description"),
153 "created_at": item.get("created_at"),
154 "updated_at": item.get("updated_at"),
155 }
156 for item in items
157 ]
158 except ClientError as e:
159 logger.error(f"Failed to list templates: {e}")
160 raise
162 def get_template(self, name: str) -> dict[str, Any] | None:
163 """Get a template by name."""
164 try:
165 response = self._table.get_item(Key={"template_name": name})
166 item = response.get("Item")
167 if not item:
168 return None
170 return {
171 "name": item["template_name"],
172 "description": item.get("description"),
173 "manifest": json.loads(item["manifest"]),
174 "parameters": json.loads(item.get("parameters", "{}")),
175 "created_at": item.get("created_at"),
176 "updated_at": item.get("updated_at"),
177 }
178 except ClientError as e:
179 logger.error(f"Failed to get template {name}: {e}")
180 raise
182 def create_template(
183 self,
184 name: str,
185 manifest: dict[str, Any],
186 description: str | None = None,
187 parameters: dict[str, Any] | None = None,
188 ) -> dict[str, Any]:
189 """Create a new template."""
190 now = _utc_now_iso()
192 item = {
193 "template_name": name,
194 "manifest": json.dumps(manifest),
195 "parameters": json.dumps(parameters or {}),
196 "created_at": now,
197 "updated_at": now,
198 }
199 if description:
200 item["description"] = description
202 try:
203 self._table.put_item(
204 Item=item,
205 ConditionExpression="attribute_not_exists(template_name)",
206 )
207 return {
208 "name": name,
209 "description": description,
210 "manifest": manifest,
211 "parameters": parameters or {},
212 "created_at": now,
213 }
214 except ClientError as e:
215 if e.response["Error"]["Code"] == "ConditionalCheckFailedException":
216 raise ValueError(f"Template '{name}' already exists") from e
217 logger.error(f"Failed to create template {name}: {e}")
218 raise
220 def update_template(
221 self,
222 name: str,
223 manifest: dict[str, Any] | None = None,
224 description: str | None = None,
225 parameters: dict[str, Any] | None = None,
226 ) -> dict[str, Any] | None:
227 """Update an existing template."""
228 now = _utc_now_iso()
230 update_expr_parts = ["updated_at = :updated_at"]
231 expr_values: dict[str, Any] = {":updated_at": now}
233 if manifest is not None:
234 update_expr_parts.append("manifest = :manifest")
235 expr_values[":manifest"] = json.dumps(manifest)
237 if description is not None:
238 update_expr_parts.append("description = :description")
239 expr_values[":description"] = description
241 if parameters is not None:
242 update_expr_parts.append("parameters = :parameters")
243 expr_values[":parameters"] = json.dumps(parameters)
245 try:
246 response = self._table.update_item(
247 Key={"template_name": name},
248 UpdateExpression="SET " + ", ".join(update_expr_parts),
249 ExpressionAttributeValues=expr_values,
250 ConditionExpression="attribute_exists(template_name)",
251 ReturnValues="ALL_NEW",
252 )
253 item = response.get("Attributes", {})
254 return {
255 "name": item["template_name"],
256 "description": item.get("description"),
257 "manifest": json.loads(item["manifest"]),
258 "parameters": json.loads(item.get("parameters", "{}")),
259 "created_at": item.get("created_at"),
260 "updated_at": item.get("updated_at"),
261 }
262 except ClientError as e:
263 if e.response["Error"]["Code"] == "ConditionalCheckFailedException":
264 return None
265 logger.error(f"Failed to update template {name}: {e}")
266 raise
268 def delete_template(self, name: str) -> bool:
269 """Delete a template."""
270 try:
271 self._table.delete_item(
272 Key={"template_name": name},
273 ConditionExpression="attribute_exists(template_name)",
274 )
275 return True
276 except ClientError as e:
277 if e.response["Error"]["Code"] == "ConditionalCheckFailedException":
278 return False
279 logger.error(f"Failed to delete template {name}: {e}")
280 raise
282 def template_exists(self, name: str) -> bool:
283 """Check if a template exists."""
284 try:
285 response = self._table.get_item(
286 Key={"template_name": name},
287 ProjectionExpression="template_name",
288 )
289 return "Item" in response
290 except ClientError as e:
291 logger.error(f"Failed to check template existence {name}: {e}")
292 raise
295class WebhookStore:
296 """DynamoDB-backed store for webhooks."""
298 def __init__(self, table_name: str | None = None, region: str | None = None):
299 """Initialize the webhook store.
301 Args:
302 table_name: DynamoDB table name. Defaults to env var WEBHOOKS_TABLE_NAME.
303 region: AWS region for DynamoDB. Defaults to env var DYNAMODB_REGION,
304 then GLOBAL_REGION, then AWS_REGION.
305 """
306 self.table_name = table_name or os.getenv("WEBHOOKS_TABLE_NAME", "gco-webhooks")
307 # DynamoDB tables are in the global region, not the regional cluster region
308 self.region = (
309 region
310 or os.getenv("DYNAMODB_REGION")
311 or os.getenv("GLOBAL_REGION")
312 or os.getenv("AWS_REGION", "us-east-1")
313 )
314 self._dynamodb = boto3.resource("dynamodb", region_name=self.region)
315 self._table = self._dynamodb.Table(self.table_name)
317 def list_webhooks(
318 self,
319 namespace: str | None = None,
320 *,
321 include_secret: bool = False,
322 ) -> list[dict[str, Any]]:
323 """List webhooks, optionally filtered by namespace.
325 Secrets are redacted by default because this method also backs the
326 public list API. Internal delivery lookups explicitly opt in so HMAC
327 signing still uses the configured secret.
328 """
329 try:
330 if namespace:
331 response = self._table.query(
332 IndexName="namespace-index",
333 KeyConditionExpression="namespace = :ns",
334 ExpressionAttributeValues={":ns": namespace},
335 )
336 items = response.get("Items", [])
337 else:
338 response = self._table.scan()
339 items = response.get("Items", [])
341 while "LastEvaluatedKey" in response:
342 response = self._table.scan(
343 ExclusiveStartKey=response["LastEvaluatedKey"],
344 )
345 items.extend(response.get("Items", []))
347 webhooks: list[dict[str, Any]] = []
348 for item in items:
349 webhook = {
350 "id": item["webhook_id"],
351 "url": item["url"],
352 "events": json.loads(item.get("events", "[]")),
353 "namespace": item.get("namespace"),
354 "created_at": item.get("created_at"),
355 }
356 if include_secret and "secret" in item:
357 webhook["secret"] = item["secret"]
358 webhooks.append(webhook)
359 return webhooks
360 except ClientError as e:
361 logger.error(f"Failed to list webhooks: {e}")
362 raise
364 def get_webhook(self, webhook_id: str) -> dict[str, Any] | None:
365 """Get a webhook by ID."""
366 try:
367 response = self._table.get_item(Key={"webhook_id": webhook_id})
368 item = response.get("Item")
369 if not item:
370 return None
372 return {
373 "id": item["webhook_id"],
374 "url": item["url"],
375 "events": json.loads(item.get("events", "[]")),
376 "namespace": item.get("namespace"),
377 "secret": item.get("secret"),
378 "created_at": item.get("created_at"),
379 }
380 except ClientError as e:
381 logger.error(f"Failed to get webhook {webhook_id}: {e}")
382 raise
384 def create_webhook(
385 self,
386 webhook_id: str,
387 url: str,
388 events: list[str],
389 namespace: str | None = None,
390 secret: str | None = None,
391 ) -> dict[str, Any]:
392 """Create a new webhook."""
393 now = _utc_now_iso()
395 item: dict[str, Any] = {
396 "webhook_id": webhook_id,
397 "url": url,
398 "events": json.dumps(events),
399 "created_at": now,
400 }
401 if namespace:
402 item["namespace"] = namespace
403 if secret:
404 item["secret"] = secret
406 try:
407 self._table.put_item(Item=item)
408 return {
409 "id": webhook_id,
410 "url": url,
411 "events": events,
412 "namespace": namespace,
413 "created_at": now,
414 }
415 except ClientError as e:
416 logger.error(f"Failed to create webhook: {e}")
417 raise
419 def delete_webhook(self, webhook_id: str) -> bool:
420 """Delete a webhook."""
421 try:
422 self._table.delete_item(
423 Key={"webhook_id": webhook_id},
424 ConditionExpression="attribute_exists(webhook_id)",
425 )
426 return True
427 except ClientError as e:
428 if e.response["Error"]["Code"] == "ConditionalCheckFailedException":
429 return False
430 logger.error(f"Failed to delete webhook {webhook_id}: {e}")
431 raise
433 def get_webhooks_for_event(
434 self, event: str, namespace: str | None = None
435 ) -> list[dict[str, Any]]:
436 """Get all webhooks subscribed to a specific event."""
437 webhooks = self.list_webhooks(namespace=namespace, include_secret=True)
438 return [w for w in webhooks if event in w.get("events", [])]
441class JobStore:
442 """DynamoDB-backed store for centralized job tracking.
444 This store enables:
445 - Global job submission with region targeting
446 - Real-time status tracking across all regions
447 - Job history and audit trail
448 - Cross-region job queries without hitting K8s APIs
449 """
451 def __init__(
452 self,
453 table_name: str | None = None,
454 region: str | None = None,
455 claim_lease_seconds: int | None = None,
456 ) -> None:
457 """Initialize the store with bounded DynamoDB timeouts and claim leases."""
458 self.table_name = table_name or os.getenv("JOBS_TABLE_NAME", "gco-jobs")
459 self.region = (
460 region
461 or os.getenv("DYNAMODB_REGION")
462 or os.getenv("GLOBAL_REGION")
463 or os.getenv("AWS_REGION", "us-east-1")
464 )
465 configured_lease = claim_lease_seconds
466 if configured_lease is None:
467 try:
468 configured_lease = int(
469 os.getenv("CENTRAL_QUEUE_LEASE_SECONDS", str(_DEFAULT_CLAIM_LEASE_SECONDS))
470 )
471 except ValueError:
472 configured_lease = _DEFAULT_CLAIM_LEASE_SECONDS
473 self.claim_lease_seconds = min(
474 max(configured_lease, _MIN_CLAIM_LEASE_SECONDS),
475 _MAX_CLAIM_LEASE_SECONDS,
476 )
477 self._dynamodb = boto3.resource(
478 "dynamodb",
479 region_name=self.region,
480 config=Config(
481 connect_timeout=3,
482 read_timeout=10,
483 retries={"max_attempts": 3, "mode": "standard"},
484 ),
485 )
486 self._table = self._dynamodb.Table(self.table_name)
487 self._legacy_migration_cursors: dict[tuple[str, str], dict[str, Any]] = {}
488 self._legacy_migration_completed_in_sweep: set[tuple[str, str]] = set()
489 self._legacy_migration_next_status: dict[str, int] = {}
491 @staticmethod
492 def _is_conditional_failure(error: ClientError) -> bool:
493 return bool(
494 error.response.get("Error", {}).get("Code") == "ConditionalCheckFailedException"
495 )
497 @staticmethod
498 def _decode_json(value: Any, default: Any) -> Any:
499 if value is None:
500 return default
501 if isinstance(value, str):
502 try:
503 return json.loads(value)
504 except TypeError, ValueError:
505 return default
506 return value
508 @classmethod
509 def _history_with(
510 cls,
511 item: dict[str, Any],
512 *,
513 status: str,
514 timestamp: str,
515 message: str | None = None,
516 error: str | None = None,
517 ) -> str:
518 history = cls._decode_json(item.get("status_history"), [])
519 if not isinstance(history, list):
520 history = []
521 entry: dict[str, str] = {"status": status, "timestamp": timestamp}
522 if message:
523 entry["message"] = message
524 if error:
525 entry["error"] = error
526 history.append(entry)
527 return json.dumps(history, separators=(",", ":"))
529 def _get_raw_job(self, job_id: str) -> dict[str, Any] | None:
530 response = self._table.get_item(Key={"job_id": job_id}, ConsistentRead=True)
531 item = response.get("Item")
532 return item if isinstance(item, dict) else None
534 @staticmethod
535 def _priority_sort_key(priority: int, submitted_at: str, job_id: str) -> str:
536 """Sort higher priorities first and preserve FIFO order for ties."""
537 return f"{100 - priority:03d}#{submitted_at}#{job_id}"
539 @staticmethod
540 def _region_status(region: str, status: str) -> str:
541 return f"{region}#{status}"
543 @staticmethod
544 def _list_filter_identity(
545 target_region: str | None,
546 status: str | None,
547 namespace: str | None,
548 ) -> dict[str, str | None]:
549 return {
550 "target_region": target_region,
551 "status": status,
552 "namespace": namespace,
553 }
555 @classmethod
556 def _encode_list_cursor(
557 cls,
558 key: dict[str, Any],
559 filters: dict[str, str | None],
560 ) -> str:
561 payload = json.dumps(
562 {"version": 1, "key": key, "filters": filters},
563 separators=(",", ":"),
564 sort_keys=True,
565 ).encode("utf-8")
566 return base64.urlsafe_b64encode(payload).decode("ascii").rstrip("=")
568 @classmethod
569 def _decode_list_cursor(
570 cls,
571 cursor: str,
572 filters: dict[str, str | None],
573 ) -> dict[str, Any]:
574 if not cursor or len(cursor) > 2_048:
575 raise ValueError("Invalid queue cursor")
576 try:
577 padding = "=" * (-len(cursor) % 4)
578 payload = json.loads(base64.urlsafe_b64decode(cursor + padding))
579 except (binascii.Error, UnicodeDecodeError, json.JSONDecodeError) as error:
580 raise ValueError("Invalid queue cursor") from error
581 if not isinstance(payload, dict) or payload.get("version") != 1:
582 raise ValueError("Invalid queue cursor")
583 if payload.get("filters") != filters:
584 raise ValueError("Queue cursor does not match the requested filters")
585 key = payload.get("key")
586 if (
587 not isinstance(key, dict)
588 or set(key) != {"job_id"}
589 or not isinstance(key.get("job_id"), str)
590 or not key["job_id"]
591 ):
592 raise ValueError("Invalid queue cursor")
593 return key
595 @staticmethod
596 def _legacy_priority(item: dict[str, Any]) -> int:
597 value = item.get("priority", 0)
598 try:
599 priority = int(value) if not isinstance(value, bool) else 0
600 except TypeError, ValueError:
601 priority = 0
602 return min(max(priority, 0), 100)
604 @staticmethod
605 def _migration_snapshot_conditions(
606 item: dict[str, Any],
607 fields: Collection[str],
608 names: dict[str, str],
609 values: dict[str, Any],
610 ) -> list[str]:
611 """Build optimistic-lock predicates for fields used by migration."""
612 conditions: list[str] = []
613 for index, field_name in enumerate(fields):
614 name_token = f"#snapshot_{index}"
615 names[name_token] = field_name
616 if field_name in item:
617 value_token = f":snapshot_{index}"
618 values[value_token] = item[field_name]
619 conditions.append(f"{name_token} = {value_token}")
620 else:
621 conditions.append(f"attribute_not_exists({name_token})")
622 return conditions
624 def _migrate_legacy_record(self, item: dict[str, Any], region: str, status: str) -> str:
625 """Repair one old-writer record or fail it when adoption is unsafe.
627 The worker reads through the legacy target-region/status index during a
628 rolling upgrade, so every derived worker key may be missing *or stale*.
629 Updates carry optimistic predicates for every source field used to
630 derive those keys. A concurrent status transition, lease renewal, or
631 identity repair therefore wins instead of being overwritten by this
632 migration's older snapshot.
633 """
634 job_id = item.get("job_id")
635 if not isinstance(job_id, str) or not job_id:
636 logger.error("Ignoring legacy queue record without a job_id")
637 return "skipped"
639 priority = self._legacy_priority(item)
640 submitted_at = str(item.get("submitted_at") or item.get("updated_at") or "")
641 priority_sort = self._priority_sort_key(priority, submitted_at, job_id)
642 snapshot_fields = ["priority", "submitted_at", "updated_at"]
644 unsafe_reason: str | None = None
645 if status in {JobStatus.CLAIMED.value, JobStatus.APPLYING.value}:
646 lease_fields = (
647 "claimed_by",
648 "claim_token",
649 "claim_generation",
650 "lease_expires_at",
651 )
652 snapshot_fields.extend(lease_fields)
653 try:
654 generation = int(item.get("claim_generation", 0))
655 except TypeError, ValueError:
656 generation = 0
657 if not (
658 item.get("claimed_by")
659 and item.get("claim_token")
660 and generation > 0
661 and item.get("lease_expires_at")
662 ):
663 unsafe_reason = (
664 "Pre-upgrade transient queue record lacks complete lease fencing and "
665 "cannot be safely replayed"
666 )
667 elif status in {JobStatus.PENDING.value, JobStatus.RUNNING.value}:
668 identity_fields = ("k8s_job_name", "k8s_job_namespace", "k8s_job_uid")
669 snapshot_fields.extend(identity_fields)
670 if not all(item.get(field) for field in identity_fields):
671 unsafe_reason = (
672 "Pre-upgrade active queue record lacks deterministic Kubernetes identity and "
673 "cannot be safely adopted"
674 )
676 if unsafe_reason is None and status in {
677 JobStatus.CLAIMED.value,
678 JobStatus.APPLYING.value,
679 }:
680 work_sort = str(item["lease_expires_at"])
681 else:
682 work_sort = priority_sort
683 expected_region_status = self._region_status(region, status)
685 if unsafe_reason is None and (
686 item.get("region_status") == expected_region_status
687 and item.get("priority_sort") == priority_sort
688 and item.get("work_sort") == work_sort
689 ):
690 return "skipped"
692 values: dict[str, Any] = {
693 ":expected": status,
694 ":target_region": region,
695 ":priority_sort": priority_sort,
696 ":work_sort": work_sort,
697 }
698 names = {"#status": "status"}
699 conditions = [
700 "attribute_exists(job_id)",
701 "target_region = :target_region",
702 "#status = :expected",
703 *self._migration_snapshot_conditions(item, snapshot_fields, names, values),
704 ]
705 if unsafe_reason is None:
706 values[":region_status"] = expected_region_status
707 conditions.append(
708 "(attribute_not_exists(region_status) OR "
709 "attribute_not_exists(priority_sort) OR "
710 "attribute_not_exists(work_sort) OR "
711 "region_status <> :region_status OR "
712 "priority_sort <> :priority_sort OR work_sort <> :work_sort)"
713 )
714 update_expression = (
715 "SET region_status = :region_status, priority_sort = :priority_sort, "
716 "work_sort = :work_sort"
717 )
718 outcome = "migrated"
719 else:
720 now = _utc_now_iso()
721 update_expression = (
722 "SET #status = :failed, region_status = :region_status, "
723 "priority_sort = :priority_sort, work_sort = :work_sort, "
724 "updated_at = :now, completed_at = :now, "
725 "error_message = :error, status_history = :history "
726 "REMOVE claimed_by, claim_token, lease_expires_at"
727 )
728 values.update(
729 {
730 ":failed": JobStatus.FAILED.value,
731 ":region_status": self._region_status(region, JobStatus.FAILED.value),
732 ":now": now,
733 ":error": unsafe_reason,
734 ":history": self._history_with(
735 item,
736 status=JobStatus.FAILED.value,
737 timestamp=now,
738 message="Record fenced during queue schema migration",
739 error=unsafe_reason,
740 ),
741 }
742 )
743 outcome = "failed"
745 try:
746 self._table.update_item(
747 Key={"job_id": job_id},
748 UpdateExpression=update_expression,
749 ConditionExpression=" AND ".join(conditions),
750 ExpressionAttributeNames=names,
751 ExpressionAttributeValues=values,
752 )
753 return outcome
754 except ClientError as error:
755 if self._is_conditional_failure(error):
756 return "skipped"
757 raise
759 def migrate_legacy_records_for_region(
760 self,
761 region: str,
762 evaluation_limit: int = _MAX_LEGACY_MIGRATION_EVALUATED_ITEMS,
763 ) -> dict[str, int | bool]:
764 """Incrementally repair records written by pre-work-index workers.
766 Every bounded invocation reserves a fair share for each unfinished
767 status partition instead of allowing a large queued backlog to starve
768 lease recovery and active-job reconciliation. The starting partition
769 rotates when a budget is smaller than the number of statuses. Completed
770 sweeps reset so a mixed-version worker's later write is repaired on a
771 subsequent pass.
772 """
773 budget = min(max(int(evaluation_limit), 1), 10_000)
774 statuses = (
775 JobStatus.QUEUED.value,
776 JobStatus.CLAIMED.value,
777 JobStatus.APPLYING.value,
778 JobStatus.PENDING.value,
779 JobStatus.RUNNING.value,
780 )
781 sweep_keys = {(region, status) for status in statuses}
782 completed_in_sweep = self._legacy_migration_completed_in_sweep
783 stats: dict[str, int | bool] = {
784 "evaluated": 0,
785 "migrated": 0,
786 "failed": 0,
787 "complete": False,
788 }
790 start = self._legacy_migration_next_status.get(region, 0) % len(statuses)
791 ordered_statuses = statuses[start:] + statuses[:start]
792 attempted: list[str] = []
793 for position, status in enumerate(ordered_statuses):
794 if int(stats["evaluated"]) >= budget:
795 break
796 migration_key = (region, status)
797 if migration_key in completed_in_sweep:
798 continue
800 unfinished = sum(
801 (region, candidate) not in completed_in_sweep
802 for candidate in ordered_statuses[position:]
803 )
804 status_budget = max(
805 1,
806 (budget - int(stats["evaluated"]) + unfinished - 1) // unfinished,
807 )
808 status_evaluated = 0
809 attempted.append(status)
810 while int(stats["evaluated"]) < budget and status_evaluated < status_budget:
811 remaining = min(
812 budget - int(stats["evaluated"]),
813 status_budget - status_evaluated,
814 100,
815 )
816 kwargs: dict[str, Any] = {
817 "IndexName": _LEGACY_REGION_STATUS_INDEX,
818 "KeyConditionExpression": (
819 "target_region = :target_region AND #status = :status"
820 ),
821 "ExpressionAttributeNames": {"#status": "status"},
822 "ExpressionAttributeValues": {
823 ":target_region": region,
824 ":status": status,
825 },
826 "Limit": remaining,
827 }
828 cursor = self._legacy_migration_cursors.get(migration_key)
829 if cursor:
830 kwargs["ExclusiveStartKey"] = cursor
831 response = self._table.query(**kwargs)
832 items = response.get("Items", [])
833 scanned = int(response.get("ScannedCount", 0))
834 if scanned <= 0 and items:
835 scanned = len(items)
836 stats["evaluated"] = int(stats["evaluated"]) + scanned
837 status_evaluated += scanned
838 for item in items:
839 if not isinstance(item, dict):
840 continue
841 outcome = self._migrate_legacy_record(item, region, status)
842 if outcome in {"migrated", "failed"}:
843 stats[outcome] = int(stats[outcome]) + 1
845 next_cursor = response.get("LastEvaluatedKey")
846 if not isinstance(next_cursor, dict) or not next_cursor:
847 completed_in_sweep.add(migration_key)
848 self._legacy_migration_cursors.pop(migration_key, None)
849 break
850 self._legacy_migration_cursors[migration_key] = next_cursor
851 if scanned <= 0:
852 break
854 if attempted:
855 self._legacy_migration_next_status[region] = (statuses.index(attempted[-1]) + 1) % len(
856 statuses
857 )
859 sweep_complete = sweep_keys.issubset(completed_in_sweep)
860 if sweep_complete:
861 completed_in_sweep.difference_update(sweep_keys)
862 self._legacy_migration_next_status.pop(region, None)
863 stats["complete"] = sweep_complete
864 return stats
866 def _query_worker_index(
867 self,
868 *,
869 index_name: str,
870 region: str,
871 status: str,
872 limit: int,
873 range_attribute: str | None = None,
874 upper_bound: str | None = None,
875 ) -> list[dict[str, Any]]:
876 """Read one worker index partition with correct DynamoDB pagination."""
877 items: list[dict[str, Any]] = []
878 exclusive_start_key: dict[str, Any] | None = None
879 while len(items) < limit:
880 key_condition = "region_status = :region_status"
881 values = {":region_status": self._region_status(region, status)}
882 if range_attribute is not None:
883 assert upper_bound is not None
884 key_condition += f" AND {range_attribute} <= :upper_bound"
885 values[":upper_bound"] = upper_bound
886 kwargs: dict[str, Any] = {
887 "IndexName": index_name,
888 "KeyConditionExpression": key_condition,
889 "ExpressionAttributeValues": values,
890 "Limit": limit - len(items),
891 "ScanIndexForward": True,
892 }
893 if exclusive_start_key:
894 kwargs["ExclusiveStartKey"] = exclusive_start_key
895 response = self._table.query(**kwargs)
896 items.extend(item for item in response.get("Items", []) if isinstance(item, dict))
897 exclusive_start_key = response.get("LastEvaluatedKey")
898 if not exclusive_start_key:
899 break
900 return items[:limit]
902 def _query_region_status(
903 self,
904 region: str,
905 status: str,
906 limit: int,
907 ) -> list[dict[str, Any]]:
908 """Read the unified worker index in priority order."""
909 pages = (
910 self._query_worker_index(
911 index_name=_REGION_STATUS_WORK_INDEX,
912 region=region,
913 status=status,
914 limit=limit,
915 ),
916 )
917 items_by_job_id: dict[str, dict[str, Any]] = {}
918 for page in pages:
919 for item in page:
920 job_id = item.get("job_id")
921 if isinstance(job_id, str) and job_id:
922 items_by_job_id.setdefault(job_id, item)
924 def priority_order(item: dict[str, Any]) -> tuple[str, str]:
925 job_id = str(item.get("job_id") or "")
926 priority_sort = item.get("priority_sort")
927 if not isinstance(priority_sort, str) or not priority_sort:
928 priority_sort = self._priority_sort_key(
929 self._legacy_priority(item),
930 str(item.get("submitted_at") or item.get("updated_at") or ""),
931 job_id,
932 )
933 return priority_sort, job_id
935 return sorted(items_by_job_id.values(), key=priority_order)[:limit]
937 def _query_expired_claims(
938 self,
939 region: str,
940 status: str,
941 expires_at_or_before: str,
942 limit: int,
943 ) -> list[dict[str, Any]]:
944 """Read expired claims from the unified worker index."""
945 pages = (
946 self._query_worker_index(
947 index_name=_REGION_STATUS_WORK_INDEX,
948 region=region,
949 status=status,
950 limit=limit,
951 range_attribute="work_sort",
952 upper_bound=expires_at_or_before,
953 ),
954 )
955 items_by_job_id: dict[str, dict[str, Any]] = {}
956 for page in pages:
957 for item in page:
958 job_id = item.get("job_id")
959 if not isinstance(job_id, str) or not job_id:
960 continue
961 existing = items_by_job_id.get(job_id)
962 if existing is None or str(item.get("lease_expires_at") or "") > str(
963 existing.get("lease_expires_at") or ""
964 ):
965 # Keep the newest value if a malformed/mock page repeats a
966 # job. Real GSI query pages contain one projection per key.
967 items_by_job_id[job_id] = item
968 return sorted(
969 items_by_job_id.values(),
970 key=lambda item: (
971 str(item.get("lease_expires_at") or ""),
972 str(item.get("job_id") or ""),
973 ),
974 )[:limit]
976 def submit_job(
977 self,
978 job_id: str,
979 manifest: dict[str, Any],
980 target_region: str,
981 namespace: str = "gco-jobs",
982 priority: int = 0,
983 labels: dict[str, str] | None = None,
984 submitted_by: str | None = None,
985 *,
986 idempotency_key: str | None = None,
987 request_hash: str | None = None,
988 spot_max_price: str | None = None,
989 spot_instance_type: str | None = None,
990 ) -> dict[str, Any]:
991 """Submit a job exactly once, replaying only identical idempotent requests.
993 ``spot_max_price`` (USD/hour, serialized as a string to avoid float
994 items in DynamoDB) and ``spot_instance_type`` together form the
995 optional spot price gate: the regional queue worker will not dispatch
996 the job until the instance type's current spot price in the target
997 region drops to or below the threshold.
998 """
999 now = _utc_now_iso()
1000 job_name = manifest.get("metadata", {}).get("name", job_id)
1001 priority_sort = self._priority_sort_key(priority, now, job_id)
1002 item: dict[str, Any] = {
1003 "job_id": job_id,
1004 "job_name": job_name,
1005 "target_region": target_region,
1006 "namespace": namespace,
1007 "status": JobStatus.QUEUED.value,
1008 "region_status": self._region_status(target_region, JobStatus.QUEUED.value),
1009 "priority": priority,
1010 "priority_sort": priority_sort,
1011 "work_sort": priority_sort,
1012 "manifest": json.dumps(manifest, separators=(",", ":"), sort_keys=True),
1013 "submitted_at": now,
1014 "updated_at": now,
1015 "claim_generation": 0,
1016 "status_history": json.dumps(
1017 [{"status": JobStatus.QUEUED.value, "timestamp": now, "message": "Job submitted"}],
1018 separators=(",", ":"),
1019 ),
1020 }
1021 if labels:
1022 item["labels"] = json.dumps(labels, separators=(",", ":"), sort_keys=True)
1023 if submitted_by:
1024 item["submitted_by"] = submitted_by
1025 if idempotency_key:
1026 item["idempotency_key"] = idempotency_key
1027 item["request_hash"] = request_hash or ""
1028 if spot_max_price and spot_instance_type:
1029 item["spot_max_price"] = spot_max_price
1030 item["spot_instance_type"] = spot_instance_type
1032 try:
1033 self._table.put_item(
1034 Item=item,
1035 ConditionExpression="attribute_not_exists(job_id)",
1036 )
1037 return self._parse_job_item(item)
1038 except ClientError as error:
1039 if not self._is_conditional_failure(error):
1040 logger.error("Failed to submit job %s: %s", job_id, error)
1041 raise
1043 existing = self._get_raw_job(job_id)
1044 if (
1045 idempotency_key
1046 and existing
1047 and existing.get("idempotency_key") == idempotency_key
1048 and existing.get("request_hash") == (request_hash or "")
1049 ):
1050 replay = self._parse_job_item(existing)
1051 replay["idempotent_replay"] = True
1052 return replay
1053 raise JobSubmissionConflict("job ID or idempotency key is already in use")
1055 def record_job_failure(
1056 self,
1057 job_id: str,
1058 *,
1059 target_region: str,
1060 namespace: str,
1061 error: str,
1062 message: str | None = None,
1063 priority: int = 0,
1064 submitted_at: str | None = None,
1065 job_name: str | None = None,
1066 ) -> bool:
1067 """Create a terminal FAILED record for a job no other actor tracks.
1069 The SQS submission path (``gco jobs submit-sqs`` consumed by
1070 ``gco.services.queue_processor``) enqueues a ``job_id`` without
1071 writing a queue record, so a submission whose runs could not be
1072 applied had nothing to transition and stayed invisible outside the
1073 queue itself. This writes the record directly in
1074 ``JobStatus.FAILED`` — a terminal status the regional queue workers
1075 never claim, so the record can never be mistaken for dispatchable
1076 work.
1078 Deliberately conditional on ``attribute_not_exists(job_id)``: an
1079 existing record belongs to the centralized queue lifecycle and its
1080 fenced ``transition_job`` discipline, and this method must never
1081 stomp one. Returns ``True`` when the failure record was created and
1082 ``False`` when a record already exists (left untouched). Callers
1083 are responsible for bounding ``error`` text.
1084 """
1085 now = _utc_now_iso()
1086 submitted = submitted_at or now
1087 priority_sort = self._priority_sort_key(priority, submitted, job_id)
1088 history_entry: dict[str, str] = {
1089 "status": JobStatus.FAILED.value,
1090 "timestamp": now,
1091 }
1092 if message:
1093 history_entry["message"] = message
1094 history_entry["error"] = error
1095 item: dict[str, Any] = {
1096 "job_id": job_id,
1097 "job_name": job_name or job_id,
1098 "target_region": target_region,
1099 "namespace": namespace,
1100 "status": JobStatus.FAILED.value,
1101 "region_status": self._region_status(target_region, JobStatus.FAILED.value),
1102 "priority": priority,
1103 "priority_sort": priority_sort,
1104 "work_sort": priority_sort,
1105 "submitted_at": submitted,
1106 "updated_at": now,
1107 "completed_at": now,
1108 "claim_generation": 0,
1109 "error_message": error,
1110 "status_history": json.dumps([history_entry], separators=(",", ":")),
1111 }
1112 try:
1113 self._table.put_item(
1114 Item=item,
1115 ConditionExpression="attribute_not_exists(job_id)",
1116 )
1117 return True
1118 except ClientError as record_error:
1119 if self._is_conditional_failure(record_error):
1120 return False
1121 logger.error("Failed to record failure for job %s: %s", job_id, record_error)
1122 raise
1124 def claim_job(
1125 self,
1126 job_id: str,
1127 target_region: str,
1128 claimed_by: str,
1129 ) -> dict[str, Any] | None:
1130 """Claim a queued job with a unique token and monotonic fencing generation."""
1131 item = self._get_raw_job(job_id)
1132 if (
1133 item is None
1134 or item.get("status") != JobStatus.QUEUED.value
1135 or item.get("target_region") != target_region
1136 ):
1137 return None
1139 now = _utc_now_iso()
1140 lease_expires_at = _claim_lease_expiry_iso(self.claim_lease_seconds)
1141 claim_token = uuid.uuid4().hex
1142 generation = int(item.get("claim_generation", 0)) + 1
1143 history = self._history_with(
1144 item,
1145 status=JobStatus.CLAIMED.value,
1146 timestamp=now,
1147 message=f"Claimed by {claimed_by}",
1148 )
1149 try:
1150 response = self._table.update_item(
1151 Key={"job_id": job_id},
1152 UpdateExpression=(
1153 "SET #status = :claimed, region_status = :region_status, "
1154 "claimed_by = :claimed_by, claim_token = :claim_token, "
1155 "claim_generation = :generation, claimed_at = :now, "
1156 "updated_at = :now, lease_expires_at = :lease_expires_at, "
1157 "work_sort = :work_sort, status_history = :history"
1158 ),
1159 ConditionExpression=(
1160 "attribute_exists(job_id) AND #status = :queued AND "
1161 "target_region = :target_region AND updated_at = :expected_updated_at"
1162 ),
1163 ExpressionAttributeNames={"#status": "status"},
1164 ExpressionAttributeValues={
1165 ":claimed": JobStatus.CLAIMED.value,
1166 ":queued": JobStatus.QUEUED.value,
1167 ":region_status": self._region_status(target_region, JobStatus.CLAIMED.value),
1168 ":target_region": target_region,
1169 ":claimed_by": claimed_by,
1170 ":claim_token": claim_token,
1171 ":generation": generation,
1172 ":now": now,
1173 ":expected_updated_at": item.get("updated_at"),
1174 ":lease_expires_at": lease_expires_at,
1175 ":work_sort": lease_expires_at,
1176 ":history": history,
1177 },
1178 ReturnValues="ALL_NEW",
1179 )
1180 return self._parse_job_item(response.get("Attributes", {}), include_internal=True)
1181 except ClientError as error:
1182 if self._is_conditional_failure(error):
1183 return None
1184 logger.error("Failed to claim job %s: %s", job_id, error)
1185 raise
1187 def renew_claim(
1188 self,
1189 job_id: str,
1190 target_region: str,
1191 claimed_by: str,
1192 claim_token: str,
1193 claim_generation: int,
1194 ) -> bool:
1195 """Renew an unexpired claim; an expired or fenced owner cannot regain it."""
1196 now = _utc_now_iso()
1197 lease_expires_at = _claim_lease_expiry_iso(self.claim_lease_seconds)
1198 try:
1199 self._table.update_item(
1200 Key={"job_id": job_id},
1201 UpdateExpression=(
1202 "SET lease_expires_at = :lease_expires_at, work_sort = :work_sort, "
1203 "lease_renewed_at = :now"
1204 ),
1205 ConditionExpression=(
1206 "attribute_exists(job_id) AND target_region = :target_region AND "
1207 "#status IN (:claimed, :applying) AND claimed_by = :claimed_by AND "
1208 "claim_token = :claim_token AND claim_generation = :generation AND "
1209 "lease_expires_at > :now"
1210 ),
1211 ExpressionAttributeNames={"#status": "status"},
1212 ExpressionAttributeValues={
1213 ":target_region": target_region,
1214 ":claimed": JobStatus.CLAIMED.value,
1215 ":applying": JobStatus.APPLYING.value,
1216 ":claimed_by": claimed_by,
1217 ":claim_token": claim_token,
1218 ":generation": claim_generation,
1219 ":now": now,
1220 ":lease_expires_at": lease_expires_at,
1221 ":work_sort": lease_expires_at,
1222 },
1223 )
1224 return True
1225 except ClientError as error:
1226 if self._is_conditional_failure(error):
1227 return False
1228 logger.error("Failed to renew claim for job %s: %s", job_id, error)
1229 raise
1231 def transition_job(
1232 self,
1233 job_id: str,
1234 *,
1235 target_region: str,
1236 expected_status: JobStatus | str,
1237 status: JobStatus | str,
1238 message: str | None = None,
1239 error: str | None = None,
1240 k8s_job_name: str | None = None,
1241 k8s_job_namespace: str | None = None,
1242 k8s_job_uid: str | None = None,
1243 claimed_by: str | None = None,
1244 claim_token: str | None = None,
1245 claim_generation: int | None = None,
1246 expected_k8s_uid: str | None = None,
1247 workload_not_created: bool | None = None,
1248 ) -> dict[str, Any] | None:
1249 """Apply one fenced compare-and-set lifecycle transition.
1251 ``None`` means another actor won the race or the caller lost its lease.
1252 Terminal records are immutable because the transition matrix has no
1253 outgoing terminal edges.
1254 """
1255 expected = (
1256 expected_status.value if isinstance(expected_status, JobStatus) else expected_status
1257 )
1258 destination = status.value if isinstance(status, JobStatus) else status
1259 if destination not in _ALLOWED_JOB_TRANSITIONS.get(expected, frozenset()):
1260 raise ValueError(f"Invalid job transition: {expected} -> {destination}")
1261 if workload_not_created is not None:
1262 if workload_not_created is not True:
1263 raise ValueError("workload_not_created proof must be exactly true")
1264 if expected != JobStatus.APPLYING.value:
1265 raise ValueError("workload_not_created proof is valid only from the applying state")
1266 if destination != JobStatus.FAILED.value:
1267 raise ValueError("workload_not_created proof is valid only for failed jobs")
1268 if any((k8s_job_name, k8s_job_namespace, k8s_job_uid)):
1269 raise ValueError("workload_not_created proof cannot accompany Kubernetes identity")
1271 item = self._get_raw_job(job_id)
1272 if (
1273 item is None
1274 or item.get("status") != expected
1275 or item.get("target_region") != target_region
1276 ):
1277 return None
1278 if workload_not_created is True and any(
1279 attribute in item for attribute in ("k8s_job_name", "k8s_job_namespace", "k8s_job_uid")
1280 ):
1281 raise ValueError(
1282 "workload_not_created proof requires a record without Kubernetes identity"
1283 )
1285 claim_is_required = expected in {JobStatus.CLAIMED.value, JobStatus.APPLYING.value}
1286 if claim_is_required:
1287 if claimed_by is None or claim_token is None or claim_generation is None:
1288 raise ValueError(f"Transition from {expected} requires complete claim fencing")
1289 if (
1290 item.get("claimed_by") != claimed_by
1291 or item.get("claim_token") != claim_token
1292 or int(item.get("claim_generation", -1)) != claim_generation
1293 ):
1294 return None
1295 if expected_k8s_uid is not None and str(item.get("k8s_job_uid") or "") != str(
1296 expected_k8s_uid
1297 ):
1298 return None
1300 now = _utc_now_iso()
1301 priority_sort = str(
1302 item.get("priority_sort")
1303 or self._priority_sort_key(
1304 self._legacy_priority(item),
1305 str(item.get("submitted_at") or item.get("updated_at") or now),
1306 job_id,
1307 )
1308 )
1309 work_sort = (
1310 str(item.get("lease_expires_at") or priority_sort)
1311 if destination in {JobStatus.CLAIMED.value, JobStatus.APPLYING.value}
1312 else priority_sort
1313 )
1314 update_parts = [
1315 "#status = :destination",
1316 "region_status = :region_status",
1317 "priority_sort = :priority_sort",
1318 "work_sort = :work_sort",
1319 "updated_at = :now",
1320 "status_history = :history",
1321 ]
1322 remove_parts: list[str] = []
1323 values: dict[str, Any] = {
1324 ":destination": destination,
1325 ":expected": expected,
1326 ":region_status": self._region_status(target_region, destination),
1327 ":priority_sort": priority_sort,
1328 ":work_sort": work_sort,
1329 ":target_region": target_region,
1330 ":now": now,
1331 ":expected_updated_at": item.get("updated_at"),
1332 ":history": self._history_with(
1333 item,
1334 status=destination,
1335 timestamp=now,
1336 message=message,
1337 error=error,
1338 ),
1339 }
1340 conditions = [
1341 "attribute_exists(job_id)",
1342 "#status = :expected",
1343 "target_region = :target_region",
1344 "updated_at = :expected_updated_at",
1345 ]
1347 if claim_is_required:
1348 conditions.extend(
1349 [
1350 "claimed_by = :claimed_by",
1351 "claim_token = :claim_token",
1352 "claim_generation = :generation",
1353 "lease_expires_at > :now",
1354 ]
1355 )
1356 values.update(
1357 {
1358 ":claimed_by": claimed_by,
1359 ":claim_token": claim_token,
1360 ":generation": claim_generation,
1361 }
1362 )
1363 if expected_k8s_uid is not None:
1364 conditions.append("k8s_job_uid = :expected_k8s_uid")
1365 values[":expected_k8s_uid"] = expected_k8s_uid
1366 if workload_not_created is True:
1367 update_parts.append("workload_not_created = :workload_not_created")
1368 values[":workload_not_created"] = True
1369 conditions.extend(
1370 [
1371 "attribute_not_exists(workload_not_created)",
1372 "attribute_not_exists(k8s_job_name)",
1373 "attribute_not_exists(k8s_job_namespace)",
1374 "attribute_not_exists(k8s_job_uid)",
1375 ]
1376 )
1378 for attribute, value, placeholder in (
1379 ("k8s_job_name", k8s_job_name, ":k8s_job_name"),
1380 ("k8s_job_namespace", k8s_job_namespace, ":k8s_job_namespace"),
1381 ("k8s_job_uid", k8s_job_uid, ":k8s_job_uid"),
1382 ):
1383 if value:
1384 update_parts.append(f"{attribute} = {placeholder}")
1385 values[placeholder] = value
1387 if error:
1388 update_parts.append("error_message = :error")
1389 values[":error"] = error
1390 elif destination != JobStatus.FAILED.value:
1391 remove_parts.append("error_message")
1393 if destination in _TERMINAL_JOB_STATUSES:
1394 update_parts.append("completed_at = :now")
1395 if destination not in {JobStatus.CLAIMED.value, JobStatus.APPLYING.value}:
1396 remove_parts.extend(["claimed_by", "claim_token", "lease_expires_at"])
1398 update_expression = "SET " + ", ".join(update_parts)
1399 if remove_parts:
1400 update_expression += " REMOVE " + ", ".join(dict.fromkeys(remove_parts))
1402 try:
1403 response = self._table.update_item(
1404 Key={"job_id": job_id},
1405 UpdateExpression=update_expression,
1406 ConditionExpression=" AND ".join(conditions),
1407 ExpressionAttributeNames={"#status": "status"},
1408 ExpressionAttributeValues=values,
1409 ReturnValues="ALL_NEW",
1410 )
1411 return self._parse_job_item(response.get("Attributes", {}))
1412 except ClientError as transition_error:
1413 if self._is_conditional_failure(transition_error):
1414 return None
1415 logger.error("Failed to transition job %s: %s", job_id, transition_error)
1416 raise
1418 def get_job(self, job_id: str) -> dict[str, Any] | None:
1419 """Get a job by ID."""
1420 try:
1421 response = self._table.get_item(Key={"job_id": job_id})
1422 item = response.get("Item")
1423 if not item:
1424 return None
1425 return self._parse_job_item(item)
1426 except ClientError as e:
1427 logger.error(f"Failed to get job {job_id}: {e}")
1428 raise
1430 def list_jobs_page(
1431 self,
1432 target_region: str | None = None,
1433 status: str | None = None,
1434 namespace: str | None = None,
1435 limit: int = 100,
1436 cursor: str | None = None,
1437 ) -> tuple[list[dict[str, Any]], str | None, bool]:
1438 """Return one bounded scan page plus an opaque continuation cursor."""
1439 limit = min(max(int(limit), 1), 1_000)
1440 filters = self._list_filter_identity(target_region, status, namespace)
1441 filter_parts: list[str] = []
1442 values: dict[str, Any] = {}
1443 names: dict[str, str] = {}
1444 if target_region:
1445 filter_parts.append("target_region = :region")
1446 values[":region"] = target_region
1447 if status:
1448 filter_parts.append("#status = :status")
1449 values[":status"] = status
1450 names["#status"] = "status"
1451 if namespace:
1452 filter_parts.append("#namespace = :namespace")
1453 values[":namespace"] = namespace
1454 names["#namespace"] = "namespace"
1456 items: list[dict[str, Any]] = []
1457 evaluated = 0
1458 exclusive_start_key = self._decode_list_cursor(cursor, filters) if cursor else None
1459 next_key: dict[str, Any] | None = None
1460 partial = False
1461 try:
1462 while len(items) < limit and evaluated < _MAX_LIST_EVALUATED_ITEMS:
1463 page_budget = min(
1464 max((limit - len(items)) * 4, 100),
1465 _MAX_LIST_EVALUATED_ITEMS - evaluated,
1466 )
1467 kwargs: dict[str, Any] = {"Limit": page_budget}
1468 if filter_parts:
1469 kwargs["FilterExpression"] = " AND ".join(filter_parts)
1470 kwargs["ExpressionAttributeValues"] = values
1471 if names:
1472 kwargs["ExpressionAttributeNames"] = names
1473 if exclusive_start_key:
1474 kwargs["ExclusiveStartKey"] = exclusive_start_key
1475 response = self._table.scan(**kwargs)
1476 page = [item for item in response.get("Items", []) if isinstance(item, dict)]
1477 remaining = limit - len(items)
1478 selected = page[:remaining]
1479 items.extend(selected)
1480 evaluated += int(response.get("ScannedCount", page_budget))
1482 if len(page) > remaining and selected:
1483 last_job_id = selected[-1].get("job_id")
1484 if isinstance(last_job_id, str) and last_job_id:
1485 next_key = {"job_id": last_job_id}
1486 else:
1487 next_key = response.get("LastEvaluatedKey")
1488 break
1490 response_key = response.get("LastEvaluatedKey")
1491 if not isinstance(response_key, dict) or not response_key:
1492 next_key = None
1493 break
1494 next_key = response_key
1495 exclusive_start_key = response_key
1496 except ClientError as error:
1497 logger.error("Failed to list jobs: %s", error)
1498 raise
1500 if next_key and evaluated >= _MAX_LIST_EVALUATED_ITEMS:
1501 partial = True
1502 logger.warning(
1503 "Job listing reached the %d-item evaluation budget before exhausting the table",
1504 _MAX_LIST_EVALUATED_ITEMS,
1505 )
1506 parsed = [self._parse_job_item(item) for item in items]
1507 parsed.sort(key=lambda job: job.get("submitted_at") or "", reverse=True)
1508 next_cursor = self._encode_list_cursor(next_key, filters) if next_key else None
1509 return parsed, next_cursor, partial
1511 def list_jobs(
1512 self,
1513 target_region: str | None = None,
1514 status: str | None = None,
1515 namespace: str | None = None,
1516 limit: int = 100,
1517 ) -> list[dict[str, Any]]:
1518 """List the first bounded page of matching jobs."""
1519 jobs, _, _ = self.list_jobs_page(
1520 target_region=target_region,
1521 status=status,
1522 namespace=namespace,
1523 limit=limit,
1524 )
1525 return jobs
1527 def get_queued_jobs_for_region(self, region: str, limit: int = 10) -> list[dict[str, Any]]:
1528 """Return the highest-priority queued jobs, FIFO within equal priority."""
1529 try:
1530 items = self._query_region_status(region, JobStatus.QUEUED.value, limit)
1531 return [self._parse_job_item(item) for item in items]
1532 except ClientError as error:
1533 logger.error("Failed to get queued jobs for %s: %s", region, error)
1534 raise
1536 def record_spot_gate_observation(
1537 self,
1538 job_id: str,
1539 *,
1540 observed_price: str,
1541 checked_at: str | None = None,
1542 ) -> bool:
1543 """Persist a spot gate observation on a still-queued job.
1545 Deliberately leaves ``updated_at`` untouched: ``claim_job`` fences on
1546 ``updated_at``, and a gate observation must never invalidate a
1547 concurrent claim attempt or count as queue-state churn. Conditional on
1548 the job still being queued so a late observation cannot decorate a
1549 claimed/terminal record. Returns whether the write happened.
1550 """
1551 try:
1552 self._table.update_item(
1553 Key={"job_id": job_id},
1554 UpdateExpression=(
1555 "SET spot_gate_checked_at = :checked_at, "
1556 "spot_gate_observed_price = :observed_price"
1557 ),
1558 ConditionExpression="attribute_exists(job_id) AND #status = :queued",
1559 ExpressionAttributeNames={"#status": "status"},
1560 ExpressionAttributeValues={
1561 ":checked_at": checked_at or _utc_now_iso(),
1562 ":observed_price": observed_price,
1563 ":queued": JobStatus.QUEUED.value,
1564 },
1565 )
1566 return True
1567 except ClientError as error:
1568 if self._is_conditional_failure(error):
1569 return False
1570 logger.error("Failed to record spot gate observation for %s: %s", job_id, error)
1571 raise
1573 def get_active_jobs_for_region(self, region: str, limit: int = 100) -> list[dict[str, Any]]:
1574 """Return a total-bounded, fair sample of pending and running jobs."""
1575 jobs: list[dict[str, Any]] = []
1576 remaining = limit
1577 statuses = (JobStatus.RUNNING.value, JobStatus.PENDING.value)
1578 try:
1579 for index, status in enumerate(statuses):
1580 statuses_left = len(statuses) - index
1581 allocation = remaining if statuses_left == 1 else max(1, remaining // statuses_left)
1582 page = self._query_region_status(region, status, allocation)
1583 jobs.extend(self._parse_job_item(item) for item in page)
1584 remaining -= len(page)
1585 if remaining <= 0:
1586 break
1587 except ClientError as error:
1588 logger.error("Failed to get active jobs for %s: %s", region, error)
1589 raise
1590 return jobs[:limit]
1592 def requeue_expired_jobs(self, region: str, limit: int = 100) -> int:
1593 """Fence expired claims and return them to the queue for deterministic adoption."""
1594 now = _utc_now_iso()
1595 candidates: list[dict[str, Any]] = []
1596 remaining = limit
1597 statuses = (JobStatus.CLAIMED.value, JobStatus.APPLYING.value)
1598 try:
1599 for index, status in enumerate(statuses):
1600 statuses_left = len(statuses) - index
1601 allocation = remaining if statuses_left == 1 else max(1, remaining // statuses_left)
1602 page = self._query_expired_claims(region, status, now, allocation)
1603 candidates.extend(page)
1604 remaining -= len(page)
1605 if remaining <= 0:
1606 break
1607 except ClientError as error:
1608 logger.error("Failed to find expired jobs for %s: %s", region, error)
1609 raise
1611 candidates.sort(key=lambda item: str(item.get("lease_expires_at") or ""))
1612 recovered = 0
1613 for item in candidates:
1614 if recovered >= limit:
1615 break
1616 lease_expiry = item.get("lease_expires_at")
1617 if lease_expiry is None or str(lease_expiry) > now:
1618 continue
1619 job_id = item.get("job_id")
1620 owner = item.get("claimed_by")
1621 token = item.get("claim_token")
1622 generation = item.get("claim_generation")
1623 expected_status = item.get("status")
1624 expected_updated_at = item.get("updated_at")
1625 if not all(
1626 [job_id, owner, token, generation is not None, expected_status, expected_updated_at]
1627 ):
1628 logger.error("Refusing to recover unfenced queue record %s", job_id or "<missing>")
1629 continue
1631 history = self._history_with(
1632 item,
1633 status=JobStatus.QUEUED.value,
1634 timestamp=now,
1635 message="Expired worker claim fenced and recovered",
1636 )
1637 priority_sort = str(
1638 item.get("priority_sort")
1639 or self._priority_sort_key(
1640 self._legacy_priority(item),
1641 str(item.get("submitted_at") or item.get("updated_at") or now),
1642 str(job_id),
1643 )
1644 )
1645 try:
1646 self._table.update_item(
1647 Key={"job_id": job_id},
1648 UpdateExpression=(
1649 "SET #status = :queued, region_status = :region_status, "
1650 "priority_sort = :priority_sort, work_sort = :priority_sort, "
1651 "updated_at = :now, status_history = :history "
1652 "REMOVE claimed_by, claim_token, lease_expires_at"
1653 ),
1654 ConditionExpression=(
1655 "attribute_exists(job_id) AND #status = :expected AND "
1656 "target_region = :region AND claimed_by = :owner AND "
1657 "claim_token = :token AND claim_generation = :generation AND "
1658 "updated_at = :expected_updated_at AND lease_expires_at <= :now"
1659 ),
1660 ExpressionAttributeNames={"#status": "status"},
1661 ExpressionAttributeValues={
1662 ":queued": JobStatus.QUEUED.value,
1663 ":region_status": self._region_status(region, JobStatus.QUEUED.value),
1664 ":priority_sort": priority_sort,
1665 ":expected": expected_status,
1666 ":region": region,
1667 ":owner": owner,
1668 ":token": token,
1669 ":generation": generation,
1670 ":expected_updated_at": expected_updated_at,
1671 ":now": now,
1672 ":history": history,
1673 },
1674 )
1675 except ClientError as error:
1676 if self._is_conditional_failure(error):
1677 continue
1678 logger.error("Failed to recover expired job %s: %s", job_id, error)
1679 raise
1680 recovered += 1
1681 return recovered
1683 def get_job_count_summary(
1684 self,
1685 max_evaluated: int = _MAX_LIST_EVALUATED_ITEMS,
1686 ) -> tuple[dict[str, dict[str, int]], int, bool]:
1687 """Return bounded region/status counts and whether the result is complete."""
1688 budget = min(max(int(max_evaluated), 1), 100_000)
1689 counts: dict[str, dict[str, int]] = {}
1690 evaluated = 0
1691 exclusive_start_key: dict[str, Any] | None = None
1692 truncated = False
1693 try:
1694 while evaluated < budget:
1695 kwargs: dict[str, Any] = {
1696 "ProjectionExpression": "target_region, #status",
1697 "ExpressionAttributeNames": {"#status": "status"},
1698 "Limit": min(1_000, budget - evaluated),
1699 }
1700 if exclusive_start_key:
1701 kwargs["ExclusiveStartKey"] = exclusive_start_key
1702 response = self._table.scan(**kwargs)
1703 page = response.get("Items", [])
1704 evaluated += int(response.get("ScannedCount", len(page)))
1705 for item in page:
1706 if not isinstance(item, dict):
1707 continue
1708 region = str(item.get("target_region") or "unknown")
1709 item_status = str(item.get("status") or "unknown")
1710 region_counts = counts.setdefault(region, {})
1711 region_counts[item_status] = region_counts.get(item_status, 0) + 1
1712 next_key = response.get("LastEvaluatedKey")
1713 if not isinstance(next_key, dict) or not next_key:
1714 exclusive_start_key = None
1715 break
1716 exclusive_start_key = next_key
1717 truncated = exclusive_start_key is not None
1718 except ClientError as error:
1719 logger.error("Failed to get job counts: %s", error)
1720 raise
1721 return counts, evaluated, truncated
1723 def get_job_counts_by_region(self) -> dict[str, dict[str, int]]:
1724 """Return bounded job counts; use ``get_job_count_summary`` for completeness metadata."""
1725 counts, _, truncated = self.get_job_count_summary()
1726 if truncated:
1727 logger.warning(
1728 "Queue statistics reached the %d-item evaluation budget and are partial",
1729 _MAX_LIST_EVALUATED_ITEMS,
1730 )
1731 return counts
1733 def cancel_job(self, job_id: str, reason: str | None = None) -> bool:
1734 """Cancel only an unclaimed queued job using the same atomic history CAS."""
1735 item = self._get_raw_job(job_id)
1736 if item is None or item.get("status") != JobStatus.QUEUED.value:
1737 return False
1738 now = _utc_now_iso()
1739 history = self._history_with(
1740 item,
1741 status=JobStatus.CANCELLED.value,
1742 timestamp=now,
1743 message=reason or "Cancelled by user",
1744 )
1745 try:
1746 self._table.update_item(
1747 Key={"job_id": job_id},
1748 UpdateExpression=(
1749 "SET #status = :cancelled, region_status = :region_status, "
1750 "updated_at = :now, completed_at = :now, cancelled_at = :now, "
1751 "cancel_reason = :reason, status_history = :history"
1752 ),
1753 ConditionExpression=(
1754 "attribute_exists(job_id) AND #status = :queued AND "
1755 "target_region = :region AND updated_at = :expected_updated_at"
1756 ),
1757 ExpressionAttributeNames={"#status": "status"},
1758 ExpressionAttributeValues={
1759 ":cancelled": JobStatus.CANCELLED.value,
1760 ":queued": JobStatus.QUEUED.value,
1761 ":region_status": self._region_status(
1762 str(item.get("target_region")), JobStatus.CANCELLED.value
1763 ),
1764 ":region": item.get("target_region"),
1765 ":expected_updated_at": item.get("updated_at"),
1766 ":now": now,
1767 ":reason": reason or "Cancelled by user",
1768 ":history": history,
1769 },
1770 )
1771 return True
1772 except ClientError as error:
1773 if self._is_conditional_failure(error):
1774 return False
1775 logger.error("Failed to cancel job %s: %s", job_id, error)
1776 raise
1778 def _parse_job_item(
1779 self,
1780 item: dict[str, Any],
1781 *,
1782 include_internal: bool = False,
1783 ) -> dict[str, Any]:
1784 """Parse a DynamoDB item without exposing reusable claim tokens to APIs."""
1785 parsed = {
1786 "job_id": item.get("job_id"),
1787 "job_name": item.get("job_name"),
1788 "target_region": item.get("target_region"),
1789 "namespace": item.get("namespace"),
1790 "status": item.get("status"),
1791 "priority": int(item.get("priority", 0)),
1792 "manifest": self._decode_json(item.get("manifest"), {}),
1793 "labels": self._decode_json(item.get("labels"), {}),
1794 "submitted_at": item.get("submitted_at"),
1795 "submitted_by": item.get("submitted_by"),
1796 "claimed_by": item.get("claimed_by"),
1797 "claimed_at": item.get("claimed_at"),
1798 "claim_generation": int(item.get("claim_generation", 0)),
1799 "lease_expires_at": item.get("lease_expires_at"),
1800 "completed_at": item.get("completed_at"),
1801 "updated_at": item.get("updated_at"),
1802 "k8s_job_name": item.get("k8s_job_name"),
1803 "k8s_job_namespace": item.get("k8s_job_namespace"),
1804 "k8s_job_uid": item.get("k8s_job_uid"),
1805 "workload_not_created": item.get("workload_not_created"),
1806 "error_message": item.get("error_message"),
1807 "status_history": self._decode_json(item.get("status_history"), []),
1808 }
1809 # Optional spot price gate fields — present only for price-capped
1810 # jobs, so ungated records keep their historical shape.
1811 if item.get("spot_max_price") is not None:
1812 parsed["spot_max_price"] = str(item.get("spot_max_price"))
1813 parsed["spot_instance_type"] = item.get("spot_instance_type")
1814 parsed["spot_gate_checked_at"] = item.get("spot_gate_checked_at")
1815 observed = item.get("spot_gate_observed_price")
1816 parsed["spot_gate_observed_price"] = str(observed) if observed is not None else None
1817 if include_internal:
1818 parsed["claim_token"] = item.get("claim_token")
1819 return parsed
1822# Singleton instances for use in the API
1823_template_store: TemplateStore | None = None
1824_webhook_store: WebhookStore | None = None
1825_job_store: JobStore | None = None
1828def get_template_store() -> TemplateStore:
1829 """Get or create the template store singleton."""
1830 global _template_store
1831 if _template_store is None:
1832 _template_store = TemplateStore()
1833 return _template_store
1836def get_webhook_store() -> WebhookStore:
1837 """Get or create the webhook store singleton."""
1838 global _webhook_store
1839 if _webhook_store is None:
1840 _webhook_store = WebhookStore()
1841 return _webhook_store
1844def get_job_store() -> JobStore:
1845 """Get or create the job store singleton."""
1846 global _job_store
1847 if _job_store is None:
1848 _job_store = JobStore()
1849 return _job_store