Coverage for gco/services/queue_processor.py: 91.07%
371 statements
« prev ^ index » next coverage.py v7.15.2, created at 2026-07-30 21:22 +0000
1"""
2Queue Processor Service for GCO (Global Capacity Orchestrator on AWS).
4Polls the regional SQS job queue, reads Kubernetes manifests from messages,
5validates them, and applies them to the cluster. Designed to run as a
6short-lived pod managed by a KEDA ScaledJob that scales based on queue depth.
8Each invocation processes a single SQS message (which may contain multiple
9manifests). On success the message is deleted; on failure it returns to the
10queue after the visibility timeout (5 min) and eventually lands in the DLQ
11after 3 failed attempts.
13Message format (produced by `gco jobs submit-sqs`):
14 {
15 "job_id": "abc123",
16 "manifests": [<k8s manifest dicts>],
17 "namespace": "gco-jobs",
18 "priority": 0,
19 "submitted_at": "2026-03-26T12:00:00+00:00"
20 }
22Configuration via environment variables:
23 JOB_QUEUE_URL: SQS queue URL to consume from (required)
24 AWS_REGION: AWS region (default: us-east-1)
25 ALLOWED_NAMESPACES: Comma-separated namespace allowlist
26 (default: gco-jobs)
27 ALLOWED_KINDS: Comma-separated resource-kind allowlist shared
28 with the REST manifest processor
29 MAX_GPU_PER_MANIFEST: Max GPUs summed across all containers
30 (regular + init + ephemeral) (default: 4)
31 MAX_CPU_PER_MANIFEST: Max CPU summed across all containers; accepts
32 K8s suffixes ("500m" or "10" for cores)
33 (default: 10000 millicores = 10 cores)
34 MAX_MEMORY_PER_MANIFEST: Max memory summed across all containers;
35 accepts K8s suffixes ("32Gi", "256Mi") or
36 a bare byte count (default: 32Gi)
37 TRUSTED_REGISTRIES: Comma-separated list of registry domains
38 (e.g. "nvcr.io,public.ecr.aws"). Empty/unset
39 disables the image registry check (fail-open).
40 Keep in sync with
41 cdk.json::job_validation_policy.trusted_registries.
42 TRUSTED_DOCKERHUB_ORGS: Comma-separated list of Docker Hub org names
43 (e.g. "nvidia,pytorch"). Empty/unset disables
44 the check. Keep in sync with
45 cdk.json::job_validation_policy.trusted_dockerhub_orgs.
47Security policy toggles (all default to true except ``BLOCK_RUN_AS_ROOT``
48which defaults to false, matching job_validation_policy.manifest_security_policy
49in cdk.json). Each one controls whether the corresponding pod/container
50setting is rejected; the REST manifest_processor enforces an identical set
51so both submission paths apply the same policy:
53 BLOCK_PRIVILEGED: Reject ``securityContext.privileged: true``
54 on pod or container (default: true)
55 BLOCK_PRIVILEGE_ESCALATION: Reject containers with
56 allowPrivilegeEscalation=true
57 (default: true)
58 BLOCK_HOST_NETWORK: Block pods with hostNetwork=true
59 (default: true)
60 BLOCK_HOST_PID: Block pods with hostPID=true
61 (default: true)
62 BLOCK_HOST_IPC: Block pods with hostIPC=true
63 (default: true)
64 BLOCK_HOST_PATH: Block volumes referencing hostPath
65 (default: true)
66 BLOCK_ADDED_CAPABILITIES: Block containers that add Linux
67 capabilities via securityContext.capabilities.add
68 (default: true)
69 BLOCK_RUN_AS_ROOT: Reject runAsUser: 0 at pod or container
70 level (default: false — many public
71 images still run as root)
72"""
74from __future__ import annotations
76import json
77import logging
78import os
79import sys
80import time
81from typing import Any
83import boto3
84from kubernetes import client, config, dynamic
85from kubernetes.client.rest import ApiException
86from kubernetes.dynamic.exceptions import NotFoundError, ResourceNotFoundError
88from gco.models import ResourceStatus
89from gco.services.manifest_processor import DEFAULT_ALLOWED_KINDS, validate_resource_kind
91logging.basicConfig(
92 level=logging.INFO,
93 format="%(asctime)s %(levelname)s [queue-processor] %(message)s",
94)
95log = logging.getLogger("queue-processor")
98def _parse_cpu_string(cpu_str: str) -> int:
99 """Parse a Kubernetes-style CPU string to millicores.
101 Accepts:
102 - Millicore suffix: "500m" -> 500
103 - Whole cores: "4" -> 4000
104 - Bare millicore counts: "10000" (when > 999) stays as millicores
105 """
106 if not cpu_str:
107 return 0
108 s = cpu_str.strip()
109 if s.endswith("m"):
110 return int(s[:-1])
111 return int(s) * 1000
114def _parse_memory_string(memory_str: str) -> int:
115 """Parse a Kubernetes-style memory string to bytes.
117 Accepts binary suffixes (Ki, Mi, Gi, Ti), decimal suffixes (k, M, G),
118 or a bare byte count.
119 """
120 if not memory_str:
121 return 0
122 s = memory_str.strip()
123 if s.endswith("Ki"):
124 return int(s[:-2]) * 1024
125 if s.endswith("Mi"):
126 return int(s[:-2]) * 1024**2
127 if s.endswith("Gi"):
128 return int(s[:-2]) * 1024**3
129 if s.endswith("Ti"):
130 return int(s[:-2]) * 1024**4
131 if s.endswith("k"): 131 ↛ 132line 131 didn't jump to line 132 because the condition on line 131 was never true
132 return int(s[:-1]) * 1000
133 if s.endswith("M"): 133 ↛ 134line 133 didn't jump to line 134 because the condition on line 133 was never true
134 return int(s[:-1]) * 1000**2
135 if s.endswith("G"):
136 return int(s[:-1]) * 1000**3
137 return int(s)
140# --- Configuration from environment ---
141# These are set by the KEDA ScaledJob manifest (post-helm-sqs-consumer.yaml)
142# and populated from cdk.json queue_processor settings during CDK deploy.
143QUEUE_URL = os.environ.get("JOB_QUEUE_URL", "")
144REGION = os.environ.get("AWS_REGION", os.environ.get("AWS_DEFAULT_REGION", "us-east-1"))
145_allowed_namespaces_env = os.environ.get("ALLOWED_NAMESPACES")
146ALLOWED_NAMESPACES = (
147 {"gco-jobs"}
148 if _allowed_namespaces_env is None
149 else {
150 namespace.strip() for namespace in _allowed_namespaces_env.split(",") if namespace.strip()
151 }
152)
153_allowed_kinds_env = os.environ.get("ALLOWED_KINDS")
154ALLOWED_KINDS = (
155 set(DEFAULT_ALLOWED_KINDS)
156 if _allowed_kinds_env is None
157 else {kind.strip() for kind in _allowed_kinds_env.split(",") if kind.strip()}
158)
159MAX_CPU = _parse_cpu_string(os.environ.get("MAX_CPU_PER_MANIFEST", "10000")) # millicores
160MAX_MEMORY = _parse_memory_string(os.environ.get("MAX_MEMORY_PER_MANIFEST", "32Gi")) # bytes
161MAX_GPU = int(os.environ.get("MAX_GPU_PER_MANIFEST", "4"))
163# Accelerator resource keys and their node taint keys (taint key == resource
164# key for all three). Kept in sync with the mirror in
165# gco/services/manifest_processor.py::ACCELERATOR_TAINTS.
166ACCELERATOR_TAINTS = ("nvidia.com/gpu", "aws.amazon.com/neuron", "vpc.amazonaws.com/efa")
168# Trusted image sources (populated from cdk.json::manifest_processor at deploy time).
169# Comma-separated env vars; empty/unset disables the check (fail-open logged).
170# Keep in sync with gco/services/manifest_processor.py::_validate_image_sources.
171TRUSTED_REGISTRIES = [
172 r.strip() for r in os.environ.get("TRUSTED_REGISTRIES", "").split(",") if r.strip()
173]
174TRUSTED_DOCKERHUB_ORGS = [
175 o.strip() for o in os.environ.get("TRUSTED_DOCKERHUB_ORGS", "").split(",") if o.strip()
176]
179def _env_bool(name: str, default: bool) -> bool:
180 """Parse a boolean environment variable.
182 Empty/unset returns ``default``. Recognized truthy values: "true", "1",
183 "yes", "on" (case-insensitive). Everything else is falsy.
184 """
185 raw = os.environ.get(name)
186 if raw is None or raw == "":
187 return default
188 return raw.strip().lower() in ("true", "1", "yes", "on")
191# Security-policy toggles. Every one of these mirrors an attribute the REST
192# manifest_processor exposes via cdk.json::job_validation_policy.manifest_security_policy.
193# Both submission paths MUST enforce the same policy — an attacker holding
194# sqs:SendMessage on the job queue must not be able to bypass checks the REST
195# path applies. Structural parity is pinned by
196# tests/test_queue_processor.py::TestSecurityPolicyParityWithManifestProcessor.
197BLOCK_PRIVILEGED = _env_bool("BLOCK_PRIVILEGED", True)
198BLOCK_PRIVILEGE_ESCALATION = _env_bool("BLOCK_PRIVILEGE_ESCALATION", True)
199BLOCK_HOST_NETWORK = _env_bool("BLOCK_HOST_NETWORK", True)
200BLOCK_HOST_PID = _env_bool("BLOCK_HOST_PID", True)
201BLOCK_HOST_IPC = _env_bool("BLOCK_HOST_IPC", True)
202BLOCK_HOST_PATH = _env_bool("BLOCK_HOST_PATH", True)
203BLOCK_ADDED_CAPABILITIES = _env_bool("BLOCK_ADDED_CAPABILITIES", True)
204BLOCK_RUN_AS_ROOT = _env_bool("BLOCK_RUN_AS_ROOT", False)
206# Hard-reject accelerator jobs that lack a matching node toleration. Mirrors
207# manifest_processor.require_accelerator_toleration so the SQS path is not a
208# bypass.
209REQUIRE_ACCELERATOR_TOLERATION = _env_bool("REQUIRE_ACCELERATOR_TOLERATION", True)
212def _is_registry_domain(entry: str) -> bool:
213 """True if the entry looks like a registry domain (has '.' or ':')."""
214 return "." in entry or ":" in entry
217def _positive_quantity(value: Any) -> bool:
218 """True if a K8s resource quantity is present and greater than zero."""
219 if value is None:
220 return False
221 try:
222 return float(value) > 0
223 except TypeError, ValueError:
224 return True
227def _toleration_matches(tolerations: list[dict[str, Any]], taint_key: str) -> bool:
228 """True if *tolerations* tolerates the ``<taint_key>=true:NoSchedule`` taint.
230 Matches manifest_processor._toleration_matches: the toleration's ``key``
231 must equal *taint_key*, its effect must be empty or ``NoSchedule``, and it
232 must use ``operator: Exists`` or ``operator: Equal`` with ``value: "true"``.
233 """
234 for tol in tolerations:
235 if not isinstance(tol, dict) or tol.get("key") != taint_key: 235 ↛ 236line 235 didn't jump to line 236 because the condition on line 235 was never true
236 continue
237 effect = tol.get("effect", "")
238 if effect not in ("", "NoSchedule"):
239 continue
240 operator = tol.get("operator", "Equal")
241 if operator == "Exists":
242 return True
243 if operator == "Equal" and str(tol.get("value")) == "true": 243 ↛ 234line 243 didn't jump to line 234 because the condition on line 243 was always true
244 return True
245 return False
248def _requested_accelerators(pod_spec: dict[str, Any]) -> set[str]:
249 """Return the set of accelerator taint keys any container requests."""
250 requested: set[str] = set()
251 for _kind, c in _iter_containers(pod_spec):
252 res = c.get("resources", {}) or {}
253 for section in ("requests", "limits"):
254 values = res.get(section, {}) or {}
255 for taint in ACCELERATOR_TAINTS:
256 if _positive_quantity(values.get(taint)):
257 requested.add(taint)
258 return requested
261def _iter_containers(pod_spec: dict[str, Any]) -> list[tuple[str, dict[str, Any]]]:
262 """Yield (kind, container_dict) for every container, initContainer, and
263 ephemeralContainer in a pod spec."""
264 out: list[tuple[str, dict[str, Any]]] = []
265 for c in pod_spec.get("containers", []) or []:
266 out.append(("container", c))
267 for c in pod_spec.get("initContainers", []) or []:
268 out.append(("initContainer", c))
269 for c in pod_spec.get("ephemeralContainers", []) or []: 269 ↛ 270line 269 didn't jump to line 270 because the loop on line 269 never started
270 out.append(("ephemeralContainer", c))
271 return out
274def _is_image_trusted(image: str) -> bool:
275 """True if the image reference is from a trusted registry or Docker Hub org.
277 Matches the semantics of manifest_processor._validate_image_sources:
278 1. Official Docker Hub images (no '/') are always allowed
279 2. Images with a registry domain (first segment has '.' or ':') must
280 match an entry in TRUSTED_REGISTRIES exactly (or a multi-segment
281 prefix like "public.ecr.aws/lambda")
282 3. Docker Hub images with an org (first segment has no '.' or ':') must
283 match an entry in TRUSTED_DOCKERHUB_ORGS
285 If both allowlists are empty the check is disabled (fail-open, logged).
286 """
287 if not TRUSTED_REGISTRIES and not TRUSTED_DOCKERHUB_ORGS:
288 return True
289 if not image: 289 ↛ 290line 289 didn't jump to line 290 because the condition on line 289 was never true
290 return True
291 if "/" not in image:
292 # Case 1: Official Docker Hub image — always trusted
293 return True
294 first = image.split("/", 1)[0]
295 if _is_registry_domain(first):
296 for registry in TRUSTED_REGISTRIES:
297 if first == registry or image.startswith(registry + "/"):
298 return True
299 return False
300 return first in TRUSTED_DOCKERHUB_ORGS
303def load_k8s() -> None:
304 """Load Kubernetes configuration (in-cluster or local kubeconfig)."""
305 try:
306 config.load_incluster_config()
307 log.info("Loaded in-cluster Kubernetes configuration")
308 except config.ConfigException:
309 config.load_kube_config()
310 log.info("Loaded local kubeconfig")
313def validate_manifest(m: dict[str, Any]) -> tuple[bool, str]:
314 """Validate a manifest before applying it to the cluster.
316 The queue processor mirrors the security checks performed by the REST
317 `manifest_processor` service (``gco/services/manifest_processor.py``)
318 so that the SQS path cannot bypass them. Checks performed:
320 1. **Namespace allowlist** — manifest namespace must be in
321 ``ALLOWED_NAMESPACES`` (from ``ALLOWED_NAMESPACES`` env var,
322 populated from ``cdk.json::job_validation_policy.allowed_namespaces``,
323 shared with the REST manifest_processor).
325 2. **Pod-level security policy** (configurable via cdk.json::
326 job_validation_policy.manifest_security_policy, shared between both
327 services). Rejects ``hostNetwork``, ``hostPID``, ``hostIPC``,
328 ``hostPath`` volumes, privileged pod security context, and
329 (if ``BLOCK_RUN_AS_ROOT``) pod-level ``runAsUser: 0``.
331 3. **Container-level security policy** — for every container kind
332 (regular, init, ephemeral) rejects ``privileged``,
333 ``allowPrivilegeEscalation``, ``capabilities.add``, and (if
334 ``BLOCK_RUN_AS_ROOT``) container-level ``runAsUser: 0``. Iterating
335 every container kind catches the classic "smuggle it via an init
336 container" bypass.
338 4. **Image registry allowlist** — every container's image must come
339 from ``TRUSTED_REGISTRIES`` (registry domains like ``nvcr.io``)
340 or ``TRUSTED_DOCKERHUB_ORGS`` (Docker Hub orgs like ``nvidia``).
341 Official Docker Hub images with no slash are always allowed. When
342 both allowlists are empty the check is disabled. Keep the lists in
343 sync with ``cdk.json::job_validation_policy.trusted_registries`` and
344 ``trusted_dockerhub_orgs`` — CDK wires the same config into both
345 services.
347 5. **Resource caps** — the TOTAL CPU, memory, and GPU across ALL
348 containers (regular + init + ephemeral) must not exceed
349 ``MAX_CPU``, ``MAX_MEMORY``, and ``MAX_GPU``. This matches
350 ``manifest_processor._validate_resource_limits`` — K8s accounts
351 init/ephemeral resources differently at scheduling time, but
352 from an enforcement perspective we sum them so an operator's
353 ``max_*_per_manifest`` budget is a hard cap regardless of where
354 the request is placed.
356 Returns:
357 ``(True, "")`` if the manifest is accepted, otherwise
358 ``(False, reason)`` where ``reason`` is a human-readable string.
359 """
360 kind = m.get("kind")
361 if not kind:
362 return False, "missing 'kind'"
363 api = m.get("apiVersion")
364 if not api:
365 return False, "missing 'apiVersion'"
366 meta = m.get("metadata")
367 if not isinstance(meta, dict) or not meta.get("name"):
368 return False, "missing 'metadata.name'"
369 ns = meta.get("namespace", "gco-jobs")
370 if ns not in ALLOWED_NAMESPACES:
371 return False, f"namespace '{ns}' not in allowed list {ALLOWED_NAMESPACES}"
373 kind_valid, kind_error = validate_resource_kind(m, ALLOWED_KINDS)
374 if not kind_valid:
375 return False, kind_error or "resource kind is not allowed"
377 # Get pod spec for security and resource checks.
378 # Handle multiple resource shapes, matching manifest_processor._get_all_containers:
379 # - Deployments / StatefulSets / ReplicaSets / DaemonSets / Jobs: spec.template.spec
380 # - CronJob: spec.jobTemplate.spec.template.spec
381 # - Pod (bare): spec (has 'containers' directly)
382 spec = m.get("spec", {})
383 pod_spec = None
384 if "template" in spec:
385 pod_spec = spec["template"].get("spec", {})
386 elif "jobTemplate" in spec:
387 pod_spec = spec["jobTemplate"].get("spec", {}).get("template", {}).get("spec", {})
388 elif "containers" in spec:
389 # Plain Pod manifest
390 pod_spec = spec
392 if pod_spec:
393 all_containers = _iter_containers(pod_spec)
395 # --- Accelerator toleration check ---
396 # Mirror manifest_processor._validate_tolerations: a job requesting a
397 # GPU/Neuron/EFA resource must carry a matching toleration or it would
398 # stay Pending forever on tainted accelerator nodes.
399 if REQUIRE_ACCELERATOR_TOLERATION: 399 ↛ 416line 399 didn't jump to line 416 because the condition on line 399 was always true
400 tolerations = pod_spec.get("tolerations", []) or []
401 for taint in _requested_accelerators(pod_spec):
402 if not _toleration_matches(tolerations, taint):
403 hint = (
404 f"add a matching toleration (e.g. key '{taint}', operator "
405 "'Exists', effect 'NoSchedule'); see examples/gpu-job.yaml"
406 )
407 return (
408 False,
409 f"Job requests accelerator '{taint}' but no matching "
410 f"toleration for taint {taint}=true:NoSchedule was found. {hint}",
411 )
413 # --- Pod-level security policy checks ---
414 # Mirror manifest_processor._validate_security_context so the SQS
415 # path enforces the same policy as the REST path.
416 if BLOCK_HOST_NETWORK and pod_spec.get("hostNetwork", False):
417 return False, "hostNetwork is not permitted"
418 if BLOCK_HOST_PID and pod_spec.get("hostPID", False):
419 return False, "hostPID is not permitted"
420 if BLOCK_HOST_IPC and pod_spec.get("hostIPC", False):
421 return False, "hostIPC is not permitted"
422 if BLOCK_HOST_PATH:
423 for volume in pod_spec.get("volumes", []) or []:
424 if volume.get("hostPath") is not None: 424 ↛ 423line 424 didn't jump to line 423 because the condition on line 424 was always true
425 return False, "hostPath volumes are not permitted"
427 pod_security_context = pod_spec.get("securityContext", {}) or {}
428 if BLOCK_PRIVILEGED and pod_security_context.get("privileged", False):
429 return False, "privileged pod security context is not permitted"
430 if BLOCK_RUN_AS_ROOT:
431 pod_run_as_user = pod_security_context.get("runAsUser")
432 if pod_run_as_user is not None and pod_run_as_user == 0:
433 return False, "running as root (runAsUser: 0) is not permitted"
435 # --- Container-level security policy checks ---
436 # Every toggle is applied to every container kind (regular, init,
437 # ephemeral). An init container running as root or with CAP_SYS_ADMIN
438 # has the same blast radius as a regular container running the same
439 # way; there is no reason to give any kind a free pass.
440 for kind, c in all_containers:
441 cname = c.get("name", "unknown")
442 sc = c.get("securityContext", {}) or {}
443 if BLOCK_PRIVILEGED and sc.get("privileged", False):
444 return False, f"{kind} '{cname}': privileged containers are not permitted"
445 if BLOCK_PRIVILEGE_ESCALATION and sc.get("allowPrivilegeEscalation", False):
446 return False, f"{kind} '{cname}': allowPrivilegeEscalation is not permitted"
447 if BLOCK_ADDED_CAPABILITIES:
448 added_caps = (sc.get("capabilities", {}) or {}).get("add", []) or []
449 if added_caps:
450 return False, f"{kind} '{cname}': added capabilities are not permitted"
451 if BLOCK_RUN_AS_ROOT:
452 ras = sc.get("runAsUser")
453 if ras is not None and ras == 0: 453 ↛ 440line 453 didn't jump to line 440 because the condition on line 453 was always true
454 return (
455 False,
456 f"{kind} '{cname}': running as root (runAsUser: 0) is not permitted",
457 )
459 # Enforce image registry allowlist (matches manifest_processor semantics)
460 for kind, c in all_containers:
461 image = c.get("image", "")
462 if not _is_image_trusted(image):
463 cname = c.get("name", "unknown")
464 return (
465 False,
466 f"{kind} '{cname}': untrusted image source '{image}'",
467 )
469 # Enforce resource caps across ALL container kinds.
470 # Sum the resource requests/limits of every container (regular,
471 # init, and ephemeral). This is stricter than the K8s scheduler's
472 # accounting but matches our security intent: an operator's
473 # configured "max CPU/memory/GPU per manifest" is a hard cap on
474 # the total resources a submitter can request regardless of
475 # which container kind carries the request.
476 total_gpu = 0
477 total_cpu = 0
478 total_memory = 0
479 for _kind, c in all_containers:
480 res = c.get("resources", {}) or {}
481 limits = res.get("limits", {}) or {}
482 requests = res.get("requests", {}) or {}
483 gpu = limits.get("nvidia.com/gpu") or requests.get("nvidia.com/gpu", "0") # nosec B113 - dict.get(), not HTTP requests
484 total_gpu += int(gpu)
485 cpu_str = limits.get("cpu") or requests.get("cpu", "0") # nosec B113 - dict.get(), not HTTP requests
486 if isinstance(cpu_str, str) and cpu_str.endswith("m"):
487 total_cpu += int(cpu_str[:-1])
488 else:
489 total_cpu += int(float(cpu_str) * 1000)
490 mem_str = limits.get("memory") or requests.get("memory", "0") # nosec B113 - dict.get(), not HTTP requests
491 if isinstance(mem_str, str): 491 ↛ 501line 491 didn't jump to line 501 because the condition on line 491 was always true
492 if mem_str.endswith("Gi"): 492 ↛ 493line 492 didn't jump to line 493 because the condition on line 492 was never true
493 total_memory += int(float(mem_str[:-2]) * 1024**3)
494 elif mem_str.endswith("Mi"):
495 total_memory += int(float(mem_str[:-2]) * 1024**2)
496 elif mem_str.endswith("Ki"): 496 ↛ 497line 496 didn't jump to line 497 because the condition on line 496 was never true
497 total_memory += int(float(mem_str[:-2]) * 1024)
498 else:
499 total_memory += int(mem_str)
500 else:
501 total_memory += int(mem_str)
503 errors = []
504 if total_gpu > MAX_GPU:
505 errors.append(f"GPU {total_gpu} exceeds max {MAX_GPU}")
506 if total_cpu > MAX_CPU:
507 errors.append(f"CPU {total_cpu}m exceeds max {MAX_CPU}m")
508 if total_memory > MAX_MEMORY: 508 ↛ 509line 508 didn't jump to line 509 because the condition on line 508 was never true
509 errors.append(
510 f"Memory {total_memory / (1024**3):.0f}Gi "
511 f"exceeds max {MAX_MEMORY / (1024**3):.0f}Gi"
512 )
513 if errors:
514 hint = (
515 "To raise limits, update queue_processor in cdk.json "
516 "and redeploy (see examples/README.md#troubleshooting)"
517 )
518 return False, "; ".join(errors) + f". {hint}"
520 return True, ""
523def _extract_pod_spec(manifest: dict[str, Any]) -> dict[str, Any] | None:
524 """Return the pod spec for any supported workload kind, or None.
526 Mirrors manifest_processor._extract_pod_spec so the SQS path and the
527 REST path apply the same injection semantics.
528 """
529 spec = manifest.get("spec")
530 if not isinstance(spec, dict):
531 return None
533 kind = manifest.get("kind", "")
535 # CronJob: spec.jobTemplate.spec.template.spec
536 if kind == "CronJob":
537 job_template = spec.get("jobTemplate")
538 if isinstance(job_template, dict): 538 ↛ 546line 538 didn't jump to line 546 because the condition on line 538 was always true
539 job_spec = job_template.get("spec")
540 if isinstance(job_spec, dict): 540 ↛ 546line 540 didn't jump to line 546 because the condition on line 540 was always true
541 template = job_spec.get("template")
542 if isinstance(template, dict): 542 ↛ 546line 542 didn't jump to line 546 because the condition on line 542 was always true
543 pod_spec = template.get("spec")
544 if isinstance(pod_spec, dict): 544 ↛ 546line 544 didn't jump to line 546 because the condition on line 544 was always true
545 return pod_spec
546 return None
548 # Deployment / StatefulSet / DaemonSet / ReplicaSet / Job: spec.template.spec
549 if "template" in spec:
550 template = spec.get("template")
551 if isinstance(template, dict): 551 ↛ 555line 551 didn't jump to line 555 because the condition on line 551 was always true
552 pod_spec = template.get("spec")
553 if isinstance(pod_spec, dict): 553 ↛ 555line 553 didn't jump to line 555 because the condition on line 553 was always true
554 return pod_spec
555 return None
557 # Bare Pod: spec contains "containers" directly
558 if "containers" in spec: 558 ↛ 561line 558 didn't jump to line 561 because the condition on line 558 was always true
559 return spec
561 return None
564def _inject_security_defaults(manifest: dict[str, Any]) -> dict[str, Any]:
565 """Inject security defaults into a user-submitted manifest in-place.
567 Currently sets ``automountServiceAccountToken: false`` on the pod spec
568 unless the user has explicitly set it either way (uses setdefault).
570 Mirrors manifest_processor._inject_security_defaults so jobs submitted
571 via SQS get the same SA-token-theft protection as those submitted via
572 the REST API.
573 """
574 pod_spec = _extract_pod_spec(manifest)
575 if pod_spec is not None:
576 pod_spec.setdefault("automountServiceAccountToken", False)
577 return manifest
580def apply_manifest(m: dict[str, Any]) -> ResourceStatus:
581 """Apply one prevalidated manifest and return an explicit operation status.
583 Unsupported API resources are failures, never successful skips. This keeps
584 the owning SQS message available for retry and eventual DLQ inspection.
585 """
586 # Inject security defaults BEFORE applying so user pods never
587 # auto-mount the default SA token (T-022 / M-113 parity with the
588 # REST manifest_processor path).
589 _inject_security_defaults(m)
591 api_version = m["apiVersion"]
592 kind = m["kind"]
593 name = m["metadata"]["name"]
594 namespace = m["metadata"].get("namespace", "gco-jobs")
596 def status(result: str, message: str) -> ResourceStatus:
597 return ResourceStatus(
598 api_version=api_version,
599 kind=kind,
600 name=name,
601 namespace=namespace,
602 status=result,
603 message=message,
604 )
606 dyn = dynamic.DynamicClient(client.ApiClient())
607 try:
608 resource = dyn.resources.get(api_version=api_version, kind=kind)
609 except ResourceNotFoundError:
610 return status("failed", f"Unsupported Kubernetes resource {api_version}/{kind}")
612 # For Jobs, delete completed/failed ones first so re-submission works.
613 # Without this, re-submitting the same job name would fail with a 409 conflict
614 # because Kubernetes doesn't allow creating a Job with the same name as an
615 # existing one (even if it's finished).
616 if kind == "Job":
617 try:
618 existing = resource.get(name=name, namespace=namespace)
619 conditions = existing.get("status", {}).get("conditions", [])
620 finished = any(c.get("type") in ("Complete", "Failed") for c in conditions)
621 if finished:
622 log.info("Deleting finished Job %s/%s before re-creation", namespace, name)
623 resource.delete(
624 name=name,
625 namespace=namespace,
626 body=client.V1DeleteOptions(propagation_policy="Background"),
627 )
628 time.sleep(2)
629 except (NotFoundError, ApiException) as e:
630 log.debug("Pre-create lookup for Job %s/%s failed: %s", namespace, name, e)
632 # Create-or-update pattern: try create first, fall back to patch on 409 (conflict).
633 # This is idempotent — safe to retry without side effects.
634 try:
635 if resource.namespaced:
636 resource.create(body=m, namespace=namespace)
637 else:
638 resource.create(body=m)
639 return status("created", "Resource created successfully")
640 except ApiException as e:
641 if e.status == 409:
642 try:
643 if resource.namespaced:
644 resource.patch(body=m, name=name, namespace=namespace)
645 else:
646 resource.patch(body=m, name=name)
647 return status("updated", "Resource updated successfully")
648 except ApiException as patch_err:
649 return status("failed", f"Patch failed: {patch_err.reason}")
650 return status("failed", f"Create failed: {e.reason}")
651 except Exception as e:
652 return status("failed", f"Unexpected apply error: {e}")
655def process_one_message() -> bool:
656 """Receive one SQS message and delete it only after complete success.
658 ``True`` means either the poll was empty or the received message was fully
659 validated, applied, and deleted. ``False`` means the message was not
660 acknowledged (or queue configuration was invalid). Malformed, empty,
661 invalid, unsupported, and apply-failed messages deliberately remain in SQS
662 for visibility-timeout retries and eventual dead-letter-queue handling.
663 """
664 if not QUEUE_URL:
665 log.error("JOB_QUEUE_URL not set")
666 return False
668 sqs = boto3.client("sqs", region_name=REGION)
670 resp = sqs.receive_message(
671 QueueUrl=QUEUE_URL,
672 MaxNumberOfMessages=1,
673 WaitTimeSeconds=5,
674 MessageAttributeNames=["All"],
675 )
677 messages = resp.get("Messages", [])
678 if not messages:
679 log.info("No messages in queue")
680 return True
682 msg = messages[0]
683 receipt = msg.get("ReceiptHandle")
684 if not receipt: 684 ↛ 685line 684 didn't jump to line 685 because the condition on line 684 was never true
685 log.error("Received SQS message without a receipt handle; cannot acknowledge it")
686 return False
688 try:
689 body = json.loads(msg.get("Body", ""))
690 except (json.JSONDecodeError, TypeError) as e:
691 log.error("Malformed SQS message body; retaining for retry/DLQ: %s", e)
692 return False
694 if not isinstance(body, dict): 694 ↛ 695line 694 didn't jump to line 695 because the condition on line 694 was never true
695 log.error("SQS message body must be a JSON object; retaining for retry/DLQ")
696 return False
698 job_id = body.get("job_id", "unknown")
699 manifests = body.get("manifests")
700 if not isinstance(manifests, list) or not manifests:
701 log.error(
702 "Job %s must contain a non-empty manifests list; retaining for retry/DLQ",
703 job_id,
704 )
705 return False
706 if any(not isinstance(manifest, dict) for manifest in manifests): 706 ↛ 707line 706 didn't jump to line 707 because the condition on line 706 was never true
707 log.error("Job %s contains a non-object manifest; retaining for retry/DLQ", job_id)
708 return False
710 log.info("Processing job_id=%s, manifests=%d", job_id, len(manifests))
712 # Validate the entire batch before applying anything. A disallowed resource
713 # later in the message must not leave an earlier resource partially applied.
714 validation_errors: list[tuple[int, str]] = []
715 for i, manifest in enumerate(manifests):
716 try:
717 ok, reason = validate_manifest(manifest)
718 except Exception as e:
719 ok, reason = False, f"validation error: {e}"
720 if not ok:
721 validation_errors.append((i, reason))
722 if validation_errors:
723 for i, reason in validation_errors:
724 log.error(" manifest[%d] validation failed: %s", i, reason)
725 log.error("Job %s failed prevalidation; message will return to queue", job_id)
726 return False
728 failed = False
729 for i, manifest in enumerate(manifests):
730 try:
731 result = apply_manifest(manifest)
732 except Exception as e:
733 log.error(" manifest[%d] apply raised: %s", i, e)
734 failed = True
735 continue
736 log.info(
737 " manifest[%d]: %s %s/%s: %s",
738 i,
739 result.status,
740 result.kind,
741 result.name,
742 result.message or "",
743 )
744 if not result.is_successful():
745 failed = True
747 if failed:
748 # Don't delete the SQS message — it will become visible again after the
749 # visibility timeout and retry. The queue redrive policy eventually
750 # moves it to the DLQ for operator inspection.
751 log.error("Job %s had failures; message will return to queue", job_id)
752 return False
754 sqs.delete_message(QueueUrl=QUEUE_URL, ReceiptHandle=receipt)
755 log.info("Job %s processed successfully", job_id)
756 return True
759def main() -> None:
760 """Entry point for the queue processor."""
761 load_k8s()
762 success = process_one_message()
763 if not success:
764 sys.exit(1)
767if __name__ == "__main__":
768 main()