Coverage for gco / services / queue_processor.py: 100.00%
470 statements
« prev ^ index » next coverage.py v7.13.5, created at 2026-10-28 16:19 +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 — the shared envelope in ``gco.job_envelope``, which
14``gco jobs submit-sqs`` builds and this service reads:
15 {
16 "job_id": "abc123",
17 "manifests": [<k8s manifest dicts>],
18 "namespace": "gco-jobs",
19 "priority": 0,
20 "submitted_at": "2026-03-26T12:00:00+00:00"
21 }
23Configuration via environment variables:
24 JOB_QUEUE_URL: SQS queue URL to consume from (required)
25 AWS_REGION: AWS region (default: us-east-1)
26 ALLOWED_NAMESPACES: Comma-separated namespace allowlist
27 (default: gco-jobs)
28 ALLOWED_KINDS: Comma-separated resource-kind allowlist shared
29 with the REST manifest processor
30 MAX_GPU_PER_MANIFEST: Max GPUs summed across all containers
31 (regular + init + ephemeral) (default: 4)
32 MAX_CPU_PER_MANIFEST: Max CPU summed across all containers; accepts
33 K8s suffixes ("500m" or "10" for cores)
34 (default: 10000 millicores = 10 cores)
35 MAX_MEMORY_PER_MANIFEST: Max memory summed across all containers;
36 accepts K8s suffixes ("32Gi", "256Mi") or
37 a bare byte count (default: 32Gi)
38 TRUSTED_REGISTRIES: Comma-separated list of registry domains
39 (e.g. "nvcr.io,public.ecr.aws"). Empty/unset
40 uses the REST processor's secure defaults.
41 Keep in sync with
42 cdk.json::job_validation_policy.trusted_registries.
43 TRUSTED_DOCKERHUB_ORGS: Comma-separated list of Docker Hub org names
44 (e.g. "nvidia,pytorch"). Empty/unset uses the
45 REST processor's secure defaults. Keep in sync
46 with cdk.json::job_validation_policy.trusted_dockerhub_orgs.
48Security policy toggles (all default to true except ``BLOCK_RUN_AS_ROOT``
49which defaults to false, matching job_validation_policy.manifest_security_policy
50in cdk.json). Each one controls whether the corresponding pod/container
51setting is rejected; the REST manifest_processor enforces an identical set
52so both submission paths apply the same policy:
54 BLOCK_PRIVILEGED: Reject ``securityContext.privileged: true``
55 on pod or container (default: true)
56 BLOCK_PRIVILEGE_ESCALATION: Reject containers with
57 allowPrivilegeEscalation=true
58 (default: true)
59 BLOCK_HOST_NETWORK: Block pods with hostNetwork=true
60 (default: true)
61 BLOCK_HOST_PID: Block pods with hostPID=true
62 (default: true)
63 BLOCK_HOST_IPC: Block pods with hostIPC=true
64 (default: true)
65 BLOCK_HOST_PATH: Block volumes referencing hostPath
66 (default: true)
67 BLOCK_ADDED_CAPABILITIES: Block containers that add Linux
68 capabilities via securityContext.capabilities.add
69 (default: true)
70 BLOCK_RUN_AS_ROOT: Reject runAsUser: 0 at pod or container
71 level (default: false — many public
72 images still run as root)
73"""
75from __future__ import annotations
77import json
78import logging
79import os
80import sys
81import time
82from typing import Any
84import boto3
85from kubernetes import client, config, dynamic
86from kubernetes.client.rest import ApiException
87from kubernetes.dynamic.exceptions import NotFoundError, ResourceNotFoundError
89from gco.job_envelope import DEFAULT_JOB_NAMESPACE
90from gco.manifest_security_policy import parse_boolean_environment
91from gco.models import ResourceStatus
92from gco.resource_governance import DEFAULT_MANIFEST_RESOURCE_CAPS
93from gco.services.manifest_processor import (
94 ADDON_KIND_HINTS,
95 DEFAULT_ALLOWED_KINDS,
96 DEFAULT_TRUSTED_DOCKERHUB_ORGS,
97 DEFAULT_TRUSTED_REGISTRIES,
98 extract_trainjob_pod_specs,
99 validate_resource_kind,
100)
101from gco.services.structured_logging import sanitize_log_value
102from gco.services.template_store import JobStore
104# <pyflowchart-code-diagram> BEGIN - auto-inserted, do not edit
105# Generated at (UTC): 2026-09-18T02:11:36Z
106# Generated from Git commit: b8faa9689385cea16155a285a7f70cf6d488e512
107# Flowchart(s) generated from this file:
108# * ``validate_manifest`` -> ``diagrams/code_diagrams/gco/services/queue_processor.validate_manifest.html``
109# (PNG: ``diagrams/code_diagrams/gco/services/queue_processor.validate_manifest.png``)
110# Regenerate with ``SOURCE_DATE_EPOCH=<unix-seconds> GCO_DIAGRAM_SOURCE_COMMIT=<40-char-sha> python diagrams/generate.py --code-only``.
111# <pyflowchart-code-diagram> END
114logging.basicConfig(
115 level=logging.INFO,
116 format="%(asctime)s %(levelname)s [queue-processor] %(message)s",
117)
118log = logging.getLogger("queue-processor")
121def _parse_cpu_string(cpu_str: str) -> int:
122 """Parse a Kubernetes-style CPU string to millicores.
124 Accepts:
125 - Millicore suffix: "500m" -> 500
126 - Whole cores: "4" -> 4000
127 """
128 if not cpu_str:
129 return 0
130 s = cpu_str.strip()
131 if s.endswith("m"):
132 return int(s[:-1])
133 return int(s) * 1000
136def _parse_memory_string(memory_str: str) -> int:
137 """Parse a Kubernetes-style memory string to bytes.
139 Accepts binary suffixes (Ki, Mi, Gi, Ti), decimal suffixes (k, M, G),
140 or a bare byte count.
141 """
142 if not memory_str:
143 return 0
144 s = memory_str.strip()
145 if s.endswith("Ki"):
146 return int(float(s[:-2]) * 1024)
147 if s.endswith("Mi"):
148 return int(float(s[:-2]) * 1024**2)
149 if s.endswith("Gi"):
150 return int(float(s[:-2]) * 1024**3)
151 if s.endswith("Ti"):
152 return int(float(s[:-2]) * 1024**4)
153 if s.endswith("k"):
154 return int(float(s[:-1]) * 1000)
155 if s.endswith("M"):
156 return int(float(s[:-1]) * 1000**2)
157 if s.endswith("G"):
158 return int(float(s[:-1]) * 1000**3)
159 return int(float(s))
162# --- Configuration from environment ---
163# These are set by the KEDA ScaledJob manifest (post-helm-sqs-consumer.yaml)
164# and populated from cdk.json queue_processor settings during CDK deploy.
165QUEUE_URL = os.environ.get("JOB_QUEUE_URL", "")
166REGION = os.environ.get("AWS_REGION", os.environ.get("AWS_DEFAULT_REGION", "us-east-1"))
167# Jobs-table failure recording. Set by post-helm-sqs-consumer.yaml from the
168# stack's {{JOBS_TABLE_NAME}} replacement; the companion DYNAMODB_REGION env
169# var is consumed by JobStore itself. Empty (an older deployed ScaledJob
170# template) disables recording and preserves the pre-recording behavior.
171JOBS_TABLE_NAME = os.environ.get("JOBS_TABLE_NAME", "")
172# The job queue's redrive policy dead-letters a message after
173# maxReceiveCount failed receives (3 — see the JobQueue definition in
174# gco/stacks/regional_stack.py). The receive that reaches this count is the
175# message's final delivery, so it is the one attempt whose failure is
176# recorded as a terminal FAILED job record. Env override exists for tests
177# and for operators who retune the queue's redrive policy out-of-band.
178FINAL_RECEIVE_COUNT = int(os.environ.get("QP_JOB_QUEUE_MAX_RECEIVE_COUNT", "3"))
179_allowed_namespaces_env = os.environ.get("ALLOWED_NAMESPACES")
180ALLOWED_NAMESPACES = (
181 {"gco-jobs"}
182 if _allowed_namespaces_env is None
183 else {
184 namespace.strip() for namespace in _allowed_namespaces_env.split(",") if namespace.strip()
185 }
186)
187_allowed_kinds_env = os.environ.get("ALLOWED_KINDS")
188ALLOWED_KINDS = (
189 set(DEFAULT_ALLOWED_KINDS)
190 if _allowed_kinds_env is None
191 else {kind.strip() for kind in _allowed_kinds_env.split(",") if kind.strip()}
192)
193# Defaults come from the shared source of truth
194# (gco.resource_governance.DEFAULT_MANIFEST_RESOURCE_CAPS - two full
195# accelerator-node slices) so both submission front doors and the deployed
196# cdk.json values tell one story. The old inline fallback here was "10000",
197# which this parser reads as 10,000 whole cores - a thousandfold looser than
198# the REST processor's fallback of the same era.
199MAX_CPU = _parse_cpu_string(
200 os.environ.get(
201 "MAX_CPU_PER_MANIFEST",
202 str(DEFAULT_MANIFEST_RESOURCE_CAPS["max_cpu_per_manifest"]),
203 )
204) # millicores
205MAX_MEMORY = _parse_memory_string(
206 os.environ.get(
207 "MAX_MEMORY_PER_MANIFEST",
208 str(DEFAULT_MANIFEST_RESOURCE_CAPS["max_memory_per_manifest"]),
209 )
210) # bytes
211MAX_GPU = int(
212 os.environ.get(
213 "MAX_GPU_PER_MANIFEST",
214 str(DEFAULT_MANIFEST_RESOURCE_CAPS["max_gpu_per_manifest"]),
215 )
216)
218# Accelerator resource keys and their node taint keys (taint key == resource
219# key for all three). Kept in sync with the mirror in
220# gco/services/manifest_processor.py::ACCELERATOR_TAINTS.
221ACCELERATOR_TAINTS = ("nvidia.com/gpu", "aws.amazon.com/neuron", "vpc.amazonaws.com/efa")
223# Trusted image sources (populated from cdk.json::manifest_processor at deploy time).
224# Empty/unset values use the same secure defaults as ManifestProcessor so the
225# SQS path cannot bypass REST image-source validation after a wiring error.
226TRUSTED_REGISTRIES = [
227 r.strip() for r in os.environ.get("TRUSTED_REGISTRIES", "").split(",") if r.strip()
228] or list(DEFAULT_TRUSTED_REGISTRIES)
229TRUSTED_DOCKERHUB_ORGS = [
230 o.strip() for o in os.environ.get("TRUSTED_DOCKERHUB_ORGS", "").split(",") if o.strip()
231] or list(DEFAULT_TRUSTED_DOCKERHUB_ORGS)
234# Security-policy toggles. Every one of these mirrors an attribute the REST
235# manifest_processor exposes via cdk.json::job_validation_policy.manifest_security_policy.
236# Both submission paths MUST enforce the same policy — an attacker holding
237# sqs:SendMessage on the job queue must not be able to bypass checks the REST
238# path applies. Structural parity is pinned by
239# tests/test_queue_processor.py::TestSecurityPolicyParityWithManifestProcessor.
240BLOCK_PRIVILEGED = parse_boolean_environment("BLOCK_PRIVILEGED", True)
241BLOCK_PRIVILEGE_ESCALATION = parse_boolean_environment("BLOCK_PRIVILEGE_ESCALATION", True)
242BLOCK_HOST_NETWORK = parse_boolean_environment("BLOCK_HOST_NETWORK", True)
243BLOCK_HOST_PID = parse_boolean_environment("BLOCK_HOST_PID", True)
244BLOCK_HOST_IPC = parse_boolean_environment("BLOCK_HOST_IPC", True)
245BLOCK_HOST_PATH = parse_boolean_environment("BLOCK_HOST_PATH", True)
246BLOCK_ADDED_CAPABILITIES = parse_boolean_environment("BLOCK_ADDED_CAPABILITIES", True)
247BLOCK_RUN_AS_ROOT = parse_boolean_environment("BLOCK_RUN_AS_ROOT", False)
249# Hard-reject accelerator jobs that lack a matching node toleration. Mirrors
250# manifest_processor.require_accelerator_toleration so the SQS path is not a
251# bypass.
252REQUIRE_ACCELERATOR_TOLERATION = parse_boolean_environment("REQUIRE_ACCELERATOR_TOLERATION", True)
255def _is_registry_domain(entry: str) -> bool:
256 """True if the entry looks like a registry domain (has '.' or ':')."""
257 return "." in entry or ":" in entry
260def _positive_quantity(value: Any) -> bool:
261 """True if a K8s resource quantity is present and greater than zero."""
262 if value is None:
263 return False
264 try:
265 return float(value) > 0
266 except TypeError, ValueError:
267 return True
270def _toleration_matches(tolerations: list[dict[str, Any]], taint_key: str) -> bool:
271 """True if *tolerations* tolerates the ``<taint_key>=true:NoSchedule`` taint.
273 Matches manifest_processor._toleration_matches: the toleration's ``key``
274 must equal *taint_key*, its effect must be empty or ``NoSchedule``, and it
275 must use ``operator: Exists`` or ``operator: Equal`` with ``value: "true"``.
276 """
277 for tol in tolerations:
278 if not isinstance(tol, dict) or tol.get("key") != taint_key:
279 continue
280 effect = tol.get("effect", "")
281 if effect not in ("", "NoSchedule"):
282 continue
283 operator = tol.get("operator", "Equal")
284 if operator == "Exists":
285 return True
286 if operator == "Equal" and str(tol.get("value")) == "true":
287 return True
288 return False
291def _requested_accelerators(pod_spec: dict[str, Any]) -> set[str]:
292 """Return the set of accelerator taint keys any container requests."""
293 requested: set[str] = set()
294 for _kind, c in _iter_containers(pod_spec):
295 res = c.get("resources", {}) or {}
296 for section in ("requests", "limits"):
297 values = res.get(section, {}) or {}
298 for taint in ACCELERATOR_TAINTS:
299 if _positive_quantity(values.get(taint)):
300 requested.add(taint)
301 return requested
304def _iter_containers(pod_spec: dict[str, Any]) -> list[tuple[str, dict[str, Any]]]:
305 """Yield (kind, container_dict) for every container, initContainer, and
306 ephemeralContainer in a pod spec."""
307 out: list[tuple[str, dict[str, Any]]] = []
308 for c in pod_spec.get("containers", []) or []:
309 out.append(("container", c))
310 for c in pod_spec.get("initContainers", []) or []:
311 out.append(("initContainer", c))
312 for c in pod_spec.get("ephemeralContainers", []) or []:
313 out.append(("ephemeralContainer", c))
314 return out
317def _is_image_trusted(image: str) -> bool:
318 """True if the image reference is from a trusted registry or Docker Hub org.
320 Matches the semantics of manifest_processor._validate_image_sources:
321 1. Official Docker Hub images (no '/') are always allowed
322 2. Images with a registry domain (first segment has '.' or ':') must
323 match an entry in TRUSTED_REGISTRIES exactly (or a multi-segment
324 prefix like "public.ecr.aws/lambda")
325 3. Docker Hub images with an org (first segment has no '.' or ':') must
326 match an entry in TRUSTED_DOCKERHUB_ORGS
328 Empty or missing environment allowlists use the REST processor's secure
329 defaults; they never disable image-source validation.
330 """
331 if not image:
332 return True
333 if "/" not in image:
334 # Case 1: Official Docker Hub image — always trusted
335 return True
336 first = image.split("/", 1)[0]
337 if _is_registry_domain(first):
338 for registry in TRUSTED_REGISTRIES:
339 if first == registry or image.startswith(registry + "/"):
340 return True
341 return False
342 return first in TRUSTED_DOCKERHUB_ORGS
345# The standard in-cluster credential paths the kubernetes client reads. The
346# worker's ServiceAccount sets ``automountServiceAccountToken: false``, so
347# these exist ONLY because post-helm-sqs-consumer.yaml projects the
348# ``kubernetes-api-token`` volume at this exact mount point.
349_SERVICEACCOUNT_DIR = "/var/run/secrets/kubernetes.io/serviceaccount"
350_SERVICEACCOUNT_TOKEN_PATH = f"{_SERVICEACCOUNT_DIR}/token"
352# Mirrors gco.services.central_queue_worker._MAX_ERROR_LENGTH — every error
353# string persisted to the jobs table is bounded the same way on both paths.
354_MAX_ERROR_LENGTH = 2_000
357def _bounded_error(value: object) -> str:
358 """Bound user/runtime error text before persisting it in DynamoDB."""
359 text = str(value)
360 return text if len(text) <= _MAX_ERROR_LENGTH else f"{text[:_MAX_ERROR_LENGTH]}...[truncated]"
363class KubernetesConfigurationError(RuntimeError):
364 """The process cannot construct a Kubernetes client at all.
366 Distinguishes a config/environment failure (the pod can never apply
367 anything; retrying the same message cannot succeed) from a per-message
368 failure (malformed body, validation, apply errors — which retain the
369 message for visibility-timeout retry and the DLQ).
370 """
373def _incluster_failure_detail(error: Exception) -> str:
374 """Name the actual broken precondition instead of the client's generic error."""
375 service_account = (
376 os.environ.get("SERVICE_ACCOUNT_NAME") or "<unknown - SERVICE_ACCOUNT_NAME unset>"
377 )
378 if not os.path.exists(_SERVICEACCOUNT_TOKEN_PATH):
379 return (
380 f"no Kubernetes API token at {_SERVICEACCOUNT_TOKEN_PATH}. This pod runs as "
381 f"ServiceAccount {service_account!r}, which sets automountServiceAccountToken: "
382 "false, so the token exists only when the workload projects one at "
383 f"{_SERVICEACCOUNT_DIR} (the kubernetes-api-token projected volume in "
384 "post-helm-sqs-consumer.yaml). Redeploy so the base-phase ServiceAccount "
385 "hardening and the post-Helm token projection are from the same release."
386 )
387 return (
388 f"a Kubernetes API token exists at {_SERVICEACCOUNT_TOKEN_PATH} but in-cluster "
389 f"configuration still failed for ServiceAccount {service_account!r}: {error}"
390 )
393def load_k8s() -> None:
394 """Load Kubernetes configuration (in-cluster, or local kubeconfig off-cluster).
396 Inside a pod (``KUBERNETES_SERVICE_HOST`` set) an in-cluster failure is
397 terminal and raises :class:`KubernetesConfigurationError` naming the
398 missing token path and ServiceAccount — it must NOT fall through to
399 ``load_kube_config()``, whose unguarded ``ConfigException`` used to kill
400 the pod before its first SQS receive and crash-loop invisibly.
401 """
402 try:
403 config.load_incluster_config()
404 log.info("Loaded in-cluster Kubernetes configuration")
405 return
406 except config.ConfigException as incluster_error:
407 if os.environ.get("KUBERNETES_SERVICE_HOST"):
408 raise KubernetesConfigurationError(
409 _incluster_failure_detail(incluster_error)
410 ) from incluster_error
411 try:
412 config.load_kube_config()
413 log.info("Loaded local kubeconfig")
414 except config.ConfigException as kubeconfig_error:
415 raise KubernetesConfigurationError(
416 "no Kubernetes credentials available: not running in a cluster "
417 f"(KUBERNETES_SERVICE_HOST unset) and no local kubeconfig: {kubeconfig_error}"
418 ) from kubeconfig_error
421def validate_manifest(m: dict[str, Any]) -> tuple[bool, str]:
422 """Validate a manifest before applying it to the cluster.
424 The queue processor mirrors the security checks performed by the REST
425 `manifest_processor` service (``gco/services/manifest_processor.py``)
426 so that the SQS path cannot bypass them. Checks performed:
428 1. **Namespace allowlist** — manifest namespace must be in
429 ``ALLOWED_NAMESPACES`` (from ``ALLOWED_NAMESPACES`` env var,
430 populated from ``cdk.json::job_validation_policy.allowed_namespaces``,
431 shared with the REST manifest_processor).
433 2. **Pod-level security policy** (configurable via cdk.json::
434 job_validation_policy.manifest_security_policy, shared between both
435 services). Rejects ``hostNetwork``, ``hostPID``, ``hostIPC``,
436 ``hostPath`` volumes, privileged pod security context, and
437 (if ``BLOCK_RUN_AS_ROOT``) pod-level ``runAsUser: 0``.
439 3. **Container-level security policy** — for every container kind
440 (regular, init, ephemeral) rejects ``privileged``,
441 ``allowPrivilegeEscalation``, ``capabilities.add``, and (if
442 ``BLOCK_RUN_AS_ROOT``) container-level ``runAsUser: 0``. Iterating
443 every container kind catches the classic "smuggle it via an init
444 container" bypass.
446 4. **Image registry allowlist** — every container's image must come
447 from ``TRUSTED_REGISTRIES`` (registry domains like ``nvcr.io``)
448 or ``TRUSTED_DOCKERHUB_ORGS`` (Docker Hub orgs like ``nvidia``).
449 Official Docker Hub images with no slash are always allowed. Empty or
450 missing allowlists use the shared secure defaults. Keep explicit lists
451 in sync with ``cdk.json::job_validation_policy.trusted_registries`` and
452 ``trusted_dockerhub_orgs`` — CDK wires the same config into both
453 services.
455 5. **Resource caps** — the TOTAL CPU, memory, and GPU across ALL
456 containers (regular + init + ephemeral) must not exceed
457 ``MAX_CPU``, ``MAX_MEMORY``, and ``MAX_GPU``. This matches
458 ``manifest_processor._validate_resource_limits`` — K8s accounts
459 init/ephemeral resources differently at scheduling time, but
460 from an enforcement perspective we sum them so an operator's
461 ``max_*_per_manifest`` budget is a hard cap regardless of where
462 the request is placed.
464 A TrainJob has no single pod spec: checks 2-5 run over its decomposition
465 (synthetic ``spec.trainer`` view plus every pod spec embedded in
466 ``runtimePatches`` — see ``manifest_processor.TrainJobPodSpecs``), with
467 the trainer view's resources counted once per ``numNodes`` for check 5
468 and accelerator tolerations unioned across views for the toleration
469 check.
471 Returns:
472 ``(True, "")`` if the manifest is accepted, otherwise
473 ``(False, reason)`` where ``reason`` is a human-readable string.
474 """
475 kind = m.get("kind")
476 if not kind:
477 return False, "missing 'kind'"
478 api = m.get("apiVersion")
479 if not api:
480 return False, "missing 'apiVersion'"
481 meta = m.get("metadata")
482 if not isinstance(meta, dict) or not meta.get("name"):
483 return False, "missing 'metadata.name'"
484 ns = meta.get("namespace", "gco-jobs")
485 if ns not in ALLOWED_NAMESPACES:
486 return False, f"namespace '{ns}' not in allowed list {ALLOWED_NAMESPACES}"
488 kind_valid, kind_error = validate_resource_kind(m, ALLOWED_KINDS)
489 if not kind_valid:
490 return False, kind_error or "resource kind is not allowed"
492 # Get pod spec(s) for security and resource checks, each with a replica
493 # multiplier for resource-cap accounting.
494 # Handle multiple resource shapes, matching manifest_processor._get_all_containers:
495 # - Deployments / StatefulSets / ReplicaSets / DaemonSets / Jobs: spec.template.spec
496 # - CronJob: spec.jobTemplate.spec.template.spec
497 # - Pod (bare): spec (has 'containers' directly)
498 # - TrainJob: synthetic spec.trainer view weighted by numNodes, plus every
499 # pod spec embedded under spec (runtimePatches) weighted 1 — see
500 # manifest_processor.TrainJobPodSpecs for why this decomposition exists.
501 weighted_pod_specs: list[tuple[dict[str, Any], int]] = []
502 if kind == "TrainJob":
503 trainjob_specs = extract_trainjob_pod_specs(m)
504 if trainjob_specs.trainer is not None:
505 weighted_pod_specs.append((trainjob_specs.trainer, trainjob_specs.num_nodes))
506 weighted_pod_specs.extend((item, 1) for item in trainjob_specs.embedded)
507 toleration_hint_example = (
508 "examples/kubeflow-trainjob.yaml (GPU variant, via runtimePatches)"
509 )
510 else:
511 pod_spec = _extract_pod_spec(m)
512 workload_kinds = {
513 "Job",
514 "CronJob",
515 "Deployment",
516 "StatefulSet",
517 "DaemonSet",
518 "ReplicaSet",
519 "Pod",
520 }
521 if kind in workload_kinds and pod_spec is None:
522 return False, f"{kind} manifest must contain a valid pod spec"
523 if pod_spec is not None:
524 weighted_pod_specs.append((pod_spec, 1))
525 toleration_hint_example = "examples/gpu-job.yaml"
527 if weighted_pod_specs:
528 # --- Accelerator toleration check ---
529 # Mirror manifest_processor._validate_tolerations: a job requesting a
530 # GPU/Neuron/EFA resource must carry a matching toleration or it would
531 # stay Pending forever on tainted accelerator nodes. For a TrainJob the
532 # request usually lives in spec.trainer.resourcesPerNode while the
533 # toleration can only be expressed through a runtimePatches pod spec,
534 # so requests and tolerations are each unioned across every view
535 # before matching.
536 if REQUIRE_ACCELERATOR_TOLERATION:
537 requested_taints: set[str] = set()
538 tolerations: list[dict[str, Any]] = []
539 for pod_spec, _multiplier in weighted_pod_specs:
540 requested_taints.update(_requested_accelerators(pod_spec))
541 tolerations.extend(pod_spec.get("tolerations", []) or [])
542 for taint in requested_taints:
543 if not _toleration_matches(tolerations, taint):
544 hint = (
545 f"add a matching toleration (e.g. key '{taint}', operator "
546 f"'Exists', effect 'NoSchedule'); see {toleration_hint_example}"
547 )
548 return (
549 False,
550 f"Job requests accelerator '{taint}' but no matching "
551 f"toleration for taint {taint}=true:NoSchedule was found. {hint}",
552 )
554 for pod_spec, _multiplier in weighted_pod_specs:
555 # --- Pod-level security policy checks ---
556 # Mirror manifest_processor._validate_security_context so the SQS
557 # path enforces the same policy as the REST path.
558 if BLOCK_HOST_NETWORK and pod_spec.get("hostNetwork", False):
559 return False, "hostNetwork is not permitted"
560 if BLOCK_HOST_PID and pod_spec.get("hostPID", False):
561 return False, "hostPID is not permitted"
562 if BLOCK_HOST_IPC and pod_spec.get("hostIPC", False):
563 return False, "hostIPC is not permitted"
564 if BLOCK_HOST_PATH:
565 for volume in pod_spec.get("volumes", []) or []:
566 if volume.get("hostPath") is not None:
567 return False, "hostPath volumes are not permitted"
569 pod_security_context = pod_spec.get("securityContext", {}) or {}
570 if BLOCK_PRIVILEGED and pod_security_context.get("privileged", False):
571 return False, "privileged pod security context is not permitted"
572 if BLOCK_RUN_AS_ROOT:
573 pod_run_as_user = pod_security_context.get("runAsUser")
574 if pod_run_as_user is not None and pod_run_as_user == 0:
575 return False, "running as root (runAsUser: 0) is not permitted"
577 # --- Container-level security policy checks ---
578 # Every toggle is applied to every container kind (regular, init,
579 # ephemeral). An init container running as root or with CAP_SYS_ADMIN
580 # has the same blast radius as a regular container running the same
581 # way; there is no reason to give any kind a free pass.
582 for container_type, c in _iter_containers(pod_spec):
583 cname = c.get("name", "unknown")
584 sc = c.get("securityContext", {}) or {}
585 if BLOCK_PRIVILEGED and sc.get("privileged", False):
586 return (
587 False,
588 f"{container_type} '{cname}': privileged containers are not permitted",
589 )
590 if BLOCK_PRIVILEGE_ESCALATION and sc.get("allowPrivilegeEscalation", False):
591 return (
592 False,
593 f"{container_type} '{cname}': allowPrivilegeEscalation is not permitted",
594 )
595 if BLOCK_ADDED_CAPABILITIES:
596 added_caps = (sc.get("capabilities", {}) or {}).get("add", []) or []
597 if added_caps:
598 return (
599 False,
600 f"{container_type} '{cname}': added capabilities are not permitted",
601 )
602 if BLOCK_RUN_AS_ROOT:
603 ras = sc.get("runAsUser")
604 if ras is not None and ras == 0:
605 return (
606 False,
607 f"{container_type} '{cname}': running as root (runAsUser: 0) is not permitted",
608 )
610 # Enforce image registry allowlist (matches manifest_processor semantics)
611 for container_type, c in _iter_containers(pod_spec):
612 image = c.get("image", "")
613 if not _is_image_trusted(image):
614 cname = c.get("name", "unknown")
615 return (
616 False,
617 f"{container_type} '{cname}': untrusted image source '{image}'",
618 )
620 # Enforce resource caps across ALL container kinds and pod-spec views.
621 # Sum the resource requests/limits of every container (regular,
622 # init, and ephemeral), scaled by each view's replica multiplier
623 # (a TrainJob runs its trainer spec once per node — counting a
624 # 16-node GPU job as one node would make the cap meaningless).
625 # This is stricter than the K8s scheduler's accounting but matches
626 # our security intent: an operator's configured "max CPU/memory/GPU
627 # per manifest" is a hard cap on the total resources a submitter
628 # can request regardless of which container kind carries the request.
629 total_gpu = 0
630 total_cpu = 0
631 total_memory = 0
632 for pod_spec, multiplier in weighted_pod_specs:
633 for _container_type, c in _iter_containers(pod_spec):
634 res = c.get("resources", {}) or {}
635 limits = res.get("limits", {}) or {}
636 requests = res.get("requests", {}) or {}
637 gpu = limits.get("nvidia.com/gpu") or requests.get("nvidia.com/gpu", "0") # nosec B113 # dict.get(), not HTTP requests
638 total_gpu += multiplier * int(gpu)
639 cpu_str = limits.get("cpu") or requests.get("cpu", "0") # nosec B113 # dict.get(), not HTTP requests
640 if isinstance(cpu_str, str) and cpu_str.endswith("m"):
641 total_cpu += multiplier * int(cpu_str[:-1])
642 else:
643 total_cpu += multiplier * int(float(cpu_str) * 1000)
644 mem_str = limits.get("memory") or requests.get("memory", "0") # nosec B113 # dict.get(), not HTTP requests
645 mem_bytes = _parse_memory_string(str(mem_str))
646 total_memory += multiplier * mem_bytes
648 errors = []
649 if total_gpu > MAX_GPU:
650 errors.append(f"GPU {total_gpu} exceeds max {MAX_GPU}")
651 if total_cpu > MAX_CPU:
652 errors.append(f"CPU {total_cpu}m exceeds max {MAX_CPU}m")
653 if total_memory > MAX_MEMORY:
654 errors.append(
655 f"Memory {total_memory / (1024**3):.0f}Gi "
656 f"exceeds max {MAX_MEMORY / (1024**3):.0f}Gi"
657 )
658 if errors:
659 hint = (
660 "To raise limits, update queue_processor in cdk.json "
661 "and redeploy (see examples/README.md#troubleshooting)"
662 )
663 return False, "; ".join(errors) + f". {hint}"
665 return True, ""
668def _extract_pod_spec(manifest: dict[str, Any]) -> dict[str, Any] | None:
669 """Return the pod spec for any supported workload kind, or None.
671 Mirrors manifest_processor._extract_pod_spec so the SQS path and the
672 REST path apply the same injection semantics.
673 """
674 spec = manifest.get("spec")
675 if not isinstance(spec, dict):
676 return None
678 kind = manifest.get("kind", "")
680 # CronJob: spec.jobTemplate.spec.template.spec
681 if kind == "CronJob":
682 job_template = spec.get("jobTemplate")
683 if isinstance(job_template, dict):
684 job_spec = job_template.get("spec")
685 if isinstance(job_spec, dict):
686 template = job_spec.get("template")
687 if isinstance(template, dict):
688 pod_spec = template.get("spec")
689 if isinstance(pod_spec, dict):
690 return pod_spec
691 return None
693 # Deployment / StatefulSet / DaemonSet / ReplicaSet / Job: spec.template.spec
694 if "template" in spec:
695 template = spec.get("template")
696 if isinstance(template, dict):
697 pod_spec = template.get("spec")
698 if isinstance(pod_spec, dict):
699 return pod_spec
700 return None
702 # Bare Pod: spec contains "containers" directly
703 if "containers" in spec:
704 return spec
706 return None
709def _inject_security_defaults(manifest: dict[str, Any]) -> dict[str, Any]:
710 """Inject security defaults into a user-submitted manifest in-place.
712 Currently sets ``automountServiceAccountToken: false`` on the pod spec
713 unless the user has explicitly set it either way (uses setdefault).
715 Mirrors manifest_processor._inject_security_defaults so jobs submitted
716 via SQS get the same SA-token-theft protection as those submitted via
717 the REST API.
719 For a TrainJob the default is injected into every pod spec embedded in
720 ``runtimePatches`` (live references into the manifest, so setdefault
721 mutates it in place); the base pod template comes from the shipped
722 ClusterTrainingRuntime, which already disables the token.
723 """
724 if manifest.get("kind") == "TrainJob":
725 for embedded in extract_trainjob_pod_specs(manifest).embedded:
726 embedded.setdefault("automountServiceAccountToken", False)
727 return manifest
728 pod_spec = _extract_pod_spec(manifest)
729 if pod_spec is not None:
730 pod_spec.setdefault("automountServiceAccountToken", False)
731 return manifest
734def _is_job_finished(job_resource: dict[str, Any]) -> bool:
735 """Return whether a Kubernetes Job has a true terminal condition."""
736 status = job_resource.get("status", {})
737 conditions = status.get("conditions") or [] if isinstance(status, dict) else []
738 return any(
739 isinstance(condition, dict)
740 and condition.get("type") in ("Complete", "Failed")
741 and condition.get("status") == "True"
742 for condition in conditions
743 )
746def apply_manifest(m: dict[str, Any]) -> ResourceStatus:
747 """Apply one prevalidated manifest and return an explicit operation status.
749 Unsupported API resources are failures, never successful skips. This keeps
750 the owning SQS message available for retry and eventual DLQ inspection.
751 """
752 # Inject security defaults BEFORE applying so user pods never
753 # auto-mount the default SA token (T-022 / M-113 parity with the
754 # REST manifest_processor path).
755 _inject_security_defaults(m)
757 api_version = m["apiVersion"]
758 kind = m["kind"]
759 name = m["metadata"]["name"]
760 namespace = m["metadata"].get("namespace", "gco-jobs")
762 def status(result: str, message: str) -> ResourceStatus:
763 return ResourceStatus(
764 api_version=api_version,
765 kind=kind,
766 name=name,
767 namespace=namespace,
768 status=result,
769 message=message,
770 )
772 dyn = dynamic.DynamicClient(client.ApiClient())
773 try:
774 resource = dyn.resources.get(api_version=api_version, kind=kind)
775 except ResourceNotFoundError:
776 # A policy-allowed kind whose addon is not installed gets the
777 # actionable remedy appended; the stable prefix is kept for DLQ
778 # triage tooling.
779 addon_hint = ADDON_KIND_HINTS.get(kind)
780 detail = f" ({addon_hint})" if addon_hint else ""
781 return status("failed", f"Unsupported Kubernetes resource {api_version}/{kind}{detail}")
783 # For Jobs, delete completed/failed ones first so re-submission works.
784 # Without this, re-submitting the same job name would fail with a 409 conflict
785 # because Kubernetes doesn't allow creating a Job with the same name as an
786 # existing one (even if it's finished).
787 if kind == "Job":
788 try:
789 existing = resource.get(name=name, namespace=namespace)
790 if _is_job_finished(existing):
791 log.info("Deleting finished Job %s/%s before re-creation", namespace, name)
792 resource.delete(
793 name=name,
794 namespace=namespace,
795 body=client.V1DeleteOptions(propagation_policy="Background"),
796 )
797 time.sleep(2)
798 except (NotFoundError, ApiException) as e:
799 log.debug("Pre-create lookup for Job %s/%s failed: %s", namespace, name, e)
801 # Create-or-update pattern: try create first, fall back to patch on 409 (conflict).
802 # This is idempotent — safe to retry without side effects.
803 try:
804 if resource.namespaced:
805 resource.create(body=m, namespace=namespace)
806 else:
807 resource.create(body=m)
808 return status("created", "Resource created successfully")
809 except ApiException as e:
810 if e.status == 409:
811 try:
812 if resource.namespaced:
813 resource.patch(body=m, name=name, namespace=namespace)
814 else:
815 resource.patch(body=m, name=name)
816 return status("updated", "Resource updated successfully")
817 except ApiException as patch_err:
818 return status("failed", f"Patch failed: {patch_err.reason}")
819 return status("failed", f"Create failed: {e.reason}")
820 except Exception as e:
821 return status("failed", f"Unexpected apply error: {e}")
824def process_one_message() -> bool:
825 """Receive one SQS message and delete it only after complete success.
827 ``True`` means either the poll was empty or the received message was fully
828 validated, applied, and deleted. ``False`` means the message was not
829 acknowledged (or queue configuration was invalid). Malformed, empty,
830 invalid, unsupported, and apply-failed messages deliberately remain in SQS
831 for visibility-timeout retries and eventual dead-letter-queue handling.
832 """
833 if not QUEUE_URL:
834 log.error("JOB_QUEUE_URL not set")
835 return False
837 sqs = boto3.client("sqs", region_name=REGION)
839 resp = sqs.receive_message(
840 QueueUrl=QUEUE_URL,
841 MaxNumberOfMessages=1,
842 WaitTimeSeconds=5,
843 MessageAttributeNames=["All"],
844 AttributeNames=["ApproximateReceiveCount"],
845 )
847 messages = resp.get("Messages", [])
848 if not messages:
849 log.info("No messages in queue")
850 return True
852 msg = messages[0]
853 receipt = msg.get("ReceiptHandle")
854 if not receipt:
855 log.error("Received SQS message without a receipt handle; cannot acknowledge it")
856 return False
858 try:
859 body = json.loads(msg.get("Body", ""))
860 except (json.JSONDecodeError, TypeError) as e:
861 log.error("Malformed SQS message body; retaining for retry/DLQ: %s", e)
862 return False
864 if not isinstance(body, dict):
865 log.error("SQS message body must be a JSON object; retaining for retry/DLQ")
866 return False
868 job_id = body.get("job_id", "unknown")
869 manifests = body.get("manifests")
870 if not isinstance(manifests, list) or not manifests:
871 log.error(
872 "Job %s must contain a non-empty manifests list; retaining for retry/DLQ",
873 job_id,
874 )
875 return False
876 if any(not isinstance(manifest, dict) for manifest in manifests):
877 log.error("Job %s contains a non-object manifest; retaining for retry/DLQ", job_id)
878 return False
880 log.info("Processing job_id=%s, manifests=%d", job_id, len(manifests))
882 # Validate the entire batch before applying anything. A disallowed resource
883 # later in the message must not leave an earlier resource partially applied.
884 validation_errors: list[tuple[int, str]] = []
885 for i, manifest in enumerate(manifests):
886 try:
887 ok, reason = validate_manifest(manifest)
888 except Exception as e:
889 ok, reason = False, f"validation error: {e}"
890 if not ok:
891 validation_errors.append((i, reason))
892 if validation_errors:
893 for i, reason in validation_errors:
894 log.error(" manifest[%d] validation failed: %s", i, reason)
895 log.error("Job %s failed prevalidation; message will return to queue", job_id)
896 _record_failure_on_final_receive(
897 msg,
898 body,
899 job_id,
900 message="SQS job failed prevalidation",
901 error="; ".join(f"manifest[{i}]: {reason}" for i, reason in validation_errors),
902 )
903 return False
905 failed = False
906 failure_details: list[str] = []
907 for i, manifest in enumerate(manifests):
908 try:
909 result = apply_manifest(manifest)
910 except Exception as e:
911 log.error(" manifest[%d] apply raised: %s", i, e)
912 failed = True
913 failure_details.append(f"manifest[{i}] apply raised: {e}")
914 continue
915 log.info(
916 " manifest[%d]: %s %s/%s: %s",
917 i,
918 result.status,
919 result.kind,
920 result.name,
921 result.message or "",
922 )
923 if not result.is_successful():
924 failed = True
925 failure_details.append(
926 f"manifest[{i}] {result.kind}/{result.name}: {result.message or result.status}"
927 )
929 if failed:
930 # Don't delete the SQS message — it will become visible again after the
931 # visibility timeout and retry. The queue redrive policy eventually
932 # moves it to the DLQ for operator inspection.
933 log.error("Job %s had failures; message will return to queue", job_id)
934 _record_failure_on_final_receive(
935 msg,
936 body,
937 job_id,
938 message="SQS job could not be applied to Kubernetes",
939 error="; ".join(failure_details),
940 )
941 return False
943 sqs.delete_message(QueueUrl=QUEUE_URL, ReceiptHandle=receipt)
944 log.info("Job %s processed successfully", job_id)
945 return True
948def _receive_count(msg: dict[str, Any]) -> int:
949 """The message's ApproximateReceiveCount, defaulting to 1 when unreadable."""
950 raw = (msg.get("Attributes") or {}).get("ApproximateReceiveCount", "1")
951 try:
952 return max(int(raw), 1)
953 except TypeError, ValueError:
954 return 1
957def _record_job_failure(
958 job_id: str,
959 *,
960 namespace: str,
961 error: str,
962 message: str,
963 priority: int = 0,
964 submitted_at: str | None = None,
965) -> bool:
966 """Best-effort terminal FAILED record for an SQS job; never raises.
968 Recording failures must not change SQS retention semantics — the queue
969 and its DLQ remain the source of truth when DynamoDB is unreachable, so
970 every failure here is swallowed after logging.
971 """
972 if not JOBS_TABLE_NAME:
973 log.error(
974 "JOBS_TABLE_NAME not set; job %s failure will not be recorded",
975 sanitize_log_value(job_id),
976 )
977 return False
978 try:
979 store = JobStore(table_name=JOBS_TABLE_NAME)
980 created = store.record_job_failure(
981 job_id,
982 target_region=REGION,
983 namespace=namespace,
984 error=_bounded_error(error),
985 message=message,
986 priority=priority,
987 submitted_at=submitted_at,
988 )
989 except Exception:
990 log.exception("Unable to record FAILED for job %s", sanitize_log_value(job_id))
991 return False
992 if created:
993 log.info("Recorded FAILED job record for %s", sanitize_log_value(job_id))
994 else:
995 log.info(
996 "Job %s already has a centralized queue record; leaving it untouched",
997 sanitize_log_value(job_id),
998 )
999 return created
1002def _record_failure_on_final_receive(
1003 msg: dict[str, Any],
1004 body: dict[str, Any],
1005 job_id: str,
1006 *,
1007 message: str,
1008 error: str,
1009) -> None:
1010 """Persist a FAILED record when this receive is the message's last delivery.
1012 Earlier receives keep today's retain-for-retry behavior untouched — a
1013 transient failure that succeeds on retry must not leave a terminal
1014 record. On the final delivery (the receive that exhausts the queue's
1015 redrive maxReceiveCount) the message is about to dead-letter, so its
1016 job is recorded FAILED with the bounded failure detail.
1017 """
1018 if _receive_count(msg) < FINAL_RECEIVE_COUNT:
1019 return
1020 priority = body.get("priority", 0)
1021 if not isinstance(priority, int):
1022 priority = 0
1023 submitted_at = body.get("submitted_at")
1024 _record_job_failure(
1025 job_id,
1026 namespace=str(body.get("namespace", DEFAULT_JOB_NAMESPACE)),
1027 error=error,
1028 message=message,
1029 priority=priority,
1030 submitted_at=submitted_at if isinstance(submitted_at, str) else None,
1031 )
1034def drain_one_message_after_config_failure(reason: str) -> bool:
1035 """Convert one queued message into a visible FAILED record, then delete it.
1037 A pod that cannot construct a Kubernetes client can never apply
1038 anything, and the failure happens before the first SQS receive — so
1039 the receive count never increments, the redrive policy never fires,
1040 and KEDA restarts the loop forever with nothing recorded. That is the
1041 exact shape of the SQS submission-path outage this module is being
1042 hardened against. Draining one message per failed pod records the
1043 terminal failure where operators and submitters look (the jobs table)
1044 while emptying the queue at the same bounded rate KEDA scales pods.
1046 The message is deleted ONLY after its ``job_id`` was recorded FAILED.
1047 A message whose ``job_id`` cannot be parsed — or whose record could not
1048 be written — is retained for the normal visibility-timeout/DLQ path.
1049 Returns whether one message was drained with a record.
1050 """
1051 if not QUEUE_URL:
1052 log.error("JOB_QUEUE_URL not set")
1053 return False
1054 if not JOBS_TABLE_NAME:
1055 log.error("JOBS_TABLE_NAME not set; cannot record failures, leaving the queue untouched")
1056 return False
1058 sqs = boto3.client("sqs", region_name=REGION)
1059 resp = sqs.receive_message(
1060 QueueUrl=QUEUE_URL,
1061 MaxNumberOfMessages=1,
1062 WaitTimeSeconds=5,
1063 MessageAttributeNames=["All"],
1064 AttributeNames=["ApproximateReceiveCount"],
1065 )
1066 messages = resp.get("Messages", [])
1067 if not messages:
1068 log.info("No messages in queue to drain")
1069 return False
1071 msg = messages[0]
1072 receipt = msg.get("ReceiptHandle")
1073 try:
1074 body = json.loads(msg.get("Body", ""))
1075 except json.JSONDecodeError, TypeError:
1076 body = None
1077 job_id = body.get("job_id") if isinstance(body, dict) else None
1078 if not receipt or not isinstance(job_id, str) or not job_id:
1079 log.error("Cannot identify the job in the queued message; retaining it for the DLQ")
1080 return False
1082 priority = body.get("priority", 0)
1083 submitted_at = body.get("submitted_at")
1084 recorded = _record_job_failure(
1085 job_id,
1086 namespace=str(body.get("namespace", DEFAULT_JOB_NAMESPACE)),
1087 error=f"queue processor could not initialize Kubernetes credentials: {reason}",
1088 message="Queue processor configuration failure",
1089 priority=priority if isinstance(priority, int) else 0,
1090 submitted_at=submitted_at if isinstance(submitted_at, str) else None,
1091 )
1092 if not recorded:
1093 return False
1095 sqs.delete_message(QueueUrl=QUEUE_URL, ReceiptHandle=receipt)
1096 log.error(
1097 "Job %s recorded FAILED and drained after a configuration failure",
1098 sanitize_log_value(job_id),
1099 )
1100 return True
1103def main() -> None:
1104 """Entry point for the queue processor.
1106 Configuration failures are terminal and loud: the error is emitted in a
1107 single structured line, one queued message is drained into a FAILED job
1108 record so the queue depth and job status both move, and the process
1109 exits nonzero. Per-message failures keep their retain-for-retry
1110 semantics via :func:`process_one_message`.
1111 """
1112 try:
1113 load_k8s()
1114 except KubernetesConfigurationError as error:
1115 drained = drain_one_message_after_config_failure(str(error))
1116 log.error(
1117 "terminal=config-failure drained_with_record=%s detail=%s",
1118 drained,
1119 error,
1120 )
1121 sys.exit(1)
1122 success = process_one_message()
1123 if not success:
1124 sys.exit(1)
1127if __name__ == "__main__":
1128 main()