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