Coverage for gco / services / queue_processor.py: 100.00%

1""" 

2Queue Processor Service for GCO (Global Capacity Orchestrator on AWS). 

3 

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. 

7 

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. 

12 

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 } 

21 

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 uses the REST processor's secure defaults. 

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 uses the 

44 REST processor's secure defaults. Keep in sync 

45 with cdk.json::job_validation_policy.trusted_dockerhub_orgs. 

46 

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: 

52 

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""" 

73 

74from __future__ import annotations 

75 

76import json 

77import logging 

78import os 

79import sys 

80import time 

81from typing import Any 

82 

83import boto3 

84from kubernetes import client, config, dynamic 

85from kubernetes.client.rest import ApiException 

86from kubernetes.dynamic.exceptions import NotFoundError, ResourceNotFoundError 

87 

88from gco.manifest_security_policy import parse_boolean_environment 

89from gco.models import ResourceStatus 

90from gco.resource_governance import DEFAULT_MANIFEST_RESOURCE_CAPS 

91from gco.services.manifest_processor import ( 

92 ADDON_KIND_HINTS, 

93 DEFAULT_ALLOWED_KINDS, 

94 DEFAULT_TRUSTED_DOCKERHUB_ORGS, 

95 DEFAULT_TRUSTED_REGISTRIES, 

96 extract_trainjob_pod_specs, 

97 validate_resource_kind, 

98) 

99from gco.services.structured_logging import sanitize_log_value 

100from gco.services.template_store import JobStore 

101 

102# <pyflowchart-code-diagram> BEGIN - auto-inserted, do not edit 

103# Generated at (UTC): 2026-09-05T22:58:10Z 

104# Generated from Git commit: 745b3fa3a9af9380bfe2797a5d9716fe8ce3a557 

105# Flowchart(s) generated from this file: 

106# * ``validate_manifest`` -> ``diagrams/code_diagrams/gco/services/queue_processor.validate_manifest.html`` 

107# (PNG: ``diagrams/code_diagrams/gco/services/queue_processor.validate_manifest.png``) 

108# Regenerate with ``SOURCE_DATE_EPOCH=<unix-seconds> GCO_DIAGRAM_SOURCE_COMMIT=<40-char-sha> python diagrams/generate.py --code-only``. 

109# <pyflowchart-code-diagram> END 

110 

111 

112logging.basicConfig( 

113 level=logging.INFO, 

114 format="%(asctime)s %(levelname)s [queue-processor] %(message)s", 

115) 

116log = logging.getLogger("queue-processor") 

117 

118 

119def _parse_cpu_string(cpu_str: str) -> int: 

120 """Parse a Kubernetes-style CPU string to millicores. 

121 

122 Accepts: 

123 - Millicore suffix: "500m" -> 500 

124 - Whole cores: "4" -> 4000 

125 """ 

126 if not cpu_str: 

127 return 0 

128 s = cpu_str.strip() 

129 if s.endswith("m"): 

130 return int(s[:-1]) 

131 return int(s) * 1000 

132 

133 

134def _parse_memory_string(memory_str: str) -> int: 

135 """Parse a Kubernetes-style memory string to bytes. 

136 

137 Accepts binary suffixes (Ki, Mi, Gi, Ti), decimal suffixes (k, M, G), 

138 or a bare byte count. 

139 """ 

140 if not memory_str: 

141 return 0 

142 s = memory_str.strip() 

143 if s.endswith("Ki"): 

144 return int(float(s[:-2]) * 1024) 

145 if s.endswith("Mi"): 

146 return int(float(s[:-2]) * 1024**2) 

147 if s.endswith("Gi"): 

148 return int(float(s[:-2]) * 1024**3) 

149 if s.endswith("Ti"): 

150 return int(float(s[:-2]) * 1024**4) 

151 if s.endswith("k"): 

152 return int(float(s[:-1]) * 1000) 

153 if s.endswith("M"): 

154 return int(float(s[:-1]) * 1000**2) 

155 if s.endswith("G"): 

156 return int(float(s[:-1]) * 1000**3) 

157 return int(float(s)) 

158 

159 

160# --- Configuration from environment --- 

161# These are set by the KEDA ScaledJob manifest (post-helm-sqs-consumer.yaml) 

162# and populated from cdk.json queue_processor settings during CDK deploy. 

163QUEUE_URL = os.environ.get("JOB_QUEUE_URL", "") 

164REGION = os.environ.get("AWS_REGION", os.environ.get("AWS_DEFAULT_REGION", "us-east-1")) 

165# Jobs-table failure recording. Set by post-helm-sqs-consumer.yaml from the 

166# stack's {{JOBS_TABLE_NAME}} replacement; the companion DYNAMODB_REGION env 

167# var is consumed by JobStore itself. Empty (an older deployed ScaledJob 

168# template) disables recording and preserves the pre-recording behavior. 

169JOBS_TABLE_NAME = os.environ.get("JOBS_TABLE_NAME", "") 

170# The job queue's redrive policy dead-letters a message after 

171# maxReceiveCount failed receives (3 — see the JobQueue definition in 

172# gco/stacks/regional_stack.py). The receive that reaches this count is the 

173# message's final delivery, so it is the one attempt whose failure is 

174# recorded as a terminal FAILED job record. Env override exists for tests 

175# and for operators who retune the queue's redrive policy out-of-band. 

176FINAL_RECEIVE_COUNT = int(os.environ.get("QP_JOB_QUEUE_MAX_RECEIVE_COUNT", "3")) 

177_allowed_namespaces_env = os.environ.get("ALLOWED_NAMESPACES") 

178ALLOWED_NAMESPACES = ( 

179 {"gco-jobs"} 

180 if _allowed_namespaces_env is None 

181 else { 

182 namespace.strip() for namespace in _allowed_namespaces_env.split(",") if namespace.strip() 

183 } 

184) 

185_allowed_kinds_env = os.environ.get("ALLOWED_KINDS") 

186ALLOWED_KINDS = ( 

187 set(DEFAULT_ALLOWED_KINDS) 

188 if _allowed_kinds_env is None 

189 else {kind.strip() for kind in _allowed_kinds_env.split(",") if kind.strip()} 

190) 

191# Defaults come from the shared source of truth 

192# (gco.resource_governance.DEFAULT_MANIFEST_RESOURCE_CAPS - two full 

193# accelerator-node slices) so both submission front doors and the deployed 

194# cdk.json values tell one story. The old inline fallback here was "10000", 

195# which this parser reads as 10,000 whole cores - a thousandfold looser than 

196# the REST processor's fallback of the same era. 

197MAX_CPU = _parse_cpu_string( 

198 os.environ.get( 

199 "MAX_CPU_PER_MANIFEST", 

200 str(DEFAULT_MANIFEST_RESOURCE_CAPS["max_cpu_per_manifest"]), 

201 ) 

202) # millicores 

203MAX_MEMORY = _parse_memory_string( 

204 os.environ.get( 

205 "MAX_MEMORY_PER_MANIFEST", 

206 str(DEFAULT_MANIFEST_RESOURCE_CAPS["max_memory_per_manifest"]), 

207 ) 

208) # bytes 

209MAX_GPU = int( 

210 os.environ.get( 

211 "MAX_GPU_PER_MANIFEST", 

212 str(DEFAULT_MANIFEST_RESOURCE_CAPS["max_gpu_per_manifest"]), 

213 ) 

214) 

215 

216# Accelerator resource keys and their node taint keys (taint key == resource 

217# key for all three). Kept in sync with the mirror in 

218# gco/services/manifest_processor.py::ACCELERATOR_TAINTS. 

219ACCELERATOR_TAINTS = ("nvidia.com/gpu", "aws.amazon.com/neuron", "vpc.amazonaws.com/efa") 

220 

221# Trusted image sources (populated from cdk.json::manifest_processor at deploy time). 

222# Empty/unset values use the same secure defaults as ManifestProcessor so the 

223# SQS path cannot bypass REST image-source validation after a wiring error. 

224TRUSTED_REGISTRIES = [ 

225 r.strip() for r in os.environ.get("TRUSTED_REGISTRIES", "").split(",") if r.strip() 

226] or list(DEFAULT_TRUSTED_REGISTRIES) 

227TRUSTED_DOCKERHUB_ORGS = [ 

228 o.strip() for o in os.environ.get("TRUSTED_DOCKERHUB_ORGS", "").split(",") if o.strip() 

229] or list(DEFAULT_TRUSTED_DOCKERHUB_ORGS) 

230 

231 

232# Security-policy toggles. Every one of these mirrors an attribute the REST 

233# manifest_processor exposes via cdk.json::job_validation_policy.manifest_security_policy. 

234# Both submission paths MUST enforce the same policy — an attacker holding 

235# sqs:SendMessage on the job queue must not be able to bypass checks the REST 

236# path applies. Structural parity is pinned by 

237# tests/test_queue_processor.py::TestSecurityPolicyParityWithManifestProcessor. 

238BLOCK_PRIVILEGED = parse_boolean_environment("BLOCK_PRIVILEGED", True) 

239BLOCK_PRIVILEGE_ESCALATION = parse_boolean_environment("BLOCK_PRIVILEGE_ESCALATION", True) 

240BLOCK_HOST_NETWORK = parse_boolean_environment("BLOCK_HOST_NETWORK", True) 

241BLOCK_HOST_PID = parse_boolean_environment("BLOCK_HOST_PID", True) 

242BLOCK_HOST_IPC = parse_boolean_environment("BLOCK_HOST_IPC", True) 

243BLOCK_HOST_PATH = parse_boolean_environment("BLOCK_HOST_PATH", True) 

244BLOCK_ADDED_CAPABILITIES = parse_boolean_environment("BLOCK_ADDED_CAPABILITIES", True) 

245BLOCK_RUN_AS_ROOT = parse_boolean_environment("BLOCK_RUN_AS_ROOT", False) 

246 

247# Hard-reject accelerator jobs that lack a matching node toleration. Mirrors 

248# manifest_processor.require_accelerator_toleration so the SQS path is not a 

249# bypass. 

250REQUIRE_ACCELERATOR_TOLERATION = parse_boolean_environment("REQUIRE_ACCELERATOR_TOLERATION", True) 

251 

252 

253def _is_registry_domain(entry: str) -> bool: 

254 """True if the entry looks like a registry domain (has '.' or ':').""" 

255 return "." in entry or ":" in entry 

256 

257 

258def _positive_quantity(value: Any) -> bool: 

259 """True if a K8s resource quantity is present and greater than zero.""" 

260 if value is None: 

261 return False 

262 try: 

263 return float(value) > 0 

264 except TypeError, ValueError: 

265 return True 

266 

267 

268def _toleration_matches(tolerations: list[dict[str, Any]], taint_key: str) -> bool: 

269 """True if *tolerations* tolerates the ``<taint_key>=true:NoSchedule`` taint. 

270 

271 Matches manifest_processor._toleration_matches: the toleration's ``key`` 

272 must equal *taint_key*, its effect must be empty or ``NoSchedule``, and it 

273 must use ``operator: Exists`` or ``operator: Equal`` with ``value: "true"``. 

274 """ 

275 for tol in tolerations: 

276 if not isinstance(tol, dict) or tol.get("key") != taint_key: 

277 continue 

278 effect = tol.get("effect", "") 

279 if effect not in ("", "NoSchedule"): 

280 continue 

281 operator = tol.get("operator", "Equal") 

282 if operator == "Exists": 

283 return True 

284 if operator == "Equal" and str(tol.get("value")) == "true": 

285 return True 

286 return False 

287 

288 

289def _requested_accelerators(pod_spec: dict[str, Any]) -> set[str]: 

290 """Return the set of accelerator taint keys any container requests.""" 

291 requested: set[str] = set() 

292 for _kind, c in _iter_containers(pod_spec): 

293 res = c.get("resources", {}) or {} 

294 for section in ("requests", "limits"): 

295 values = res.get(section, {}) or {} 

296 for taint in ACCELERATOR_TAINTS: 

297 if _positive_quantity(values.get(taint)): 

298 requested.add(taint) 

299 return requested 

300 

301 

302def _iter_containers(pod_spec: dict[str, Any]) -> list[tuple[str, dict[str, Any]]]: 

303 """Yield (kind, container_dict) for every container, initContainer, and 

304 ephemeralContainer in a pod spec.""" 

305 out: list[tuple[str, dict[str, Any]]] = [] 

306 for c in pod_spec.get("containers", []) or []: 

307 out.append(("container", c)) 

308 for c in pod_spec.get("initContainers", []) or []: 

309 out.append(("initContainer", c)) 

310 for c in pod_spec.get("ephemeralContainers", []) or []: 

311 out.append(("ephemeralContainer", c)) 

312 return out 

313 

314 

315def _is_image_trusted(image: str) -> bool: 

316 """True if the image reference is from a trusted registry or Docker Hub org. 

317 

318 Matches the semantics of manifest_processor._validate_image_sources: 

319 1. Official Docker Hub images (no '/') are always allowed 

320 2. Images with a registry domain (first segment has '.' or ':') must 

321 match an entry in TRUSTED_REGISTRIES exactly (or a multi-segment 

322 prefix like "public.ecr.aws/lambda") 

323 3. Docker Hub images with an org (first segment has no '.' or ':') must 

324 match an entry in TRUSTED_DOCKERHUB_ORGS 

325 

326 Empty or missing environment allowlists use the REST processor's secure 

327 defaults; they never disable image-source validation. 

328 """ 

329 if not image: 

330 return True 

331 if "/" not in image: 

332 # Case 1: Official Docker Hub image — always trusted 

333 return True 

334 first = image.split("/", 1)[0] 

335 if _is_registry_domain(first): 

336 for registry in TRUSTED_REGISTRIES: 

337 if first == registry or image.startswith(registry + "/"): 

338 return True 

339 return False 

340 return first in TRUSTED_DOCKERHUB_ORGS 

341 

342 

343# The standard in-cluster credential paths the kubernetes client reads. The 

344# worker's ServiceAccount sets ``automountServiceAccountToken: false``, so 

345# these exist ONLY because post-helm-sqs-consumer.yaml projects the 

346# ``kubernetes-api-token`` volume at this exact mount point. 

347_SERVICEACCOUNT_DIR = "/var/run/secrets/kubernetes.io/serviceaccount" 

348_SERVICEACCOUNT_TOKEN_PATH = f"{_SERVICEACCOUNT_DIR}/token" 

349 

350# Mirrors gco.services.central_queue_worker._MAX_ERROR_LENGTH — every error 

351# string persisted to the jobs table is bounded the same way on both paths. 

352_MAX_ERROR_LENGTH = 2_000 

353 

354 

355def _bounded_error(value: object) -> str: 

356 """Bound user/runtime error text before persisting it in DynamoDB.""" 

357 text = str(value) 

358 return text if len(text) <= _MAX_ERROR_LENGTH else f"{text[:_MAX_ERROR_LENGTH]}...[truncated]" 

359 

360 

361class KubernetesConfigurationError(RuntimeError): 

362 """The process cannot construct a Kubernetes client at all. 

363 

364 Distinguishes a config/environment failure (the pod can never apply 

365 anything; retrying the same message cannot succeed) from a per-message 

366 failure (malformed body, validation, apply errors — which retain the 

367 message for visibility-timeout retry and the DLQ). 

368 """ 

369 

370 

371def _incluster_failure_detail(error: Exception) -> str: 

372 """Name the actual broken precondition instead of the client's generic error.""" 

373 service_account = ( 

374 os.environ.get("SERVICE_ACCOUNT_NAME") or "<unknown - SERVICE_ACCOUNT_NAME unset>" 

375 ) 

376 if not os.path.exists(_SERVICEACCOUNT_TOKEN_PATH): 

377 return ( 

378 f"no Kubernetes API token at {_SERVICEACCOUNT_TOKEN_PATH}. This pod runs as " 

379 f"ServiceAccount {service_account!r}, which sets automountServiceAccountToken: " 

380 "false, so the token exists only when the workload projects one at " 

381 f"{_SERVICEACCOUNT_DIR} (the kubernetes-api-token projected volume in " 

382 "post-helm-sqs-consumer.yaml). Redeploy so the base-phase ServiceAccount " 

383 "hardening and the post-Helm token projection are from the same release." 

384 ) 

385 return ( 

386 f"a Kubernetes API token exists at {_SERVICEACCOUNT_TOKEN_PATH} but in-cluster " 

387 f"configuration still failed for ServiceAccount {service_account!r}: {error}" 

388 ) 

389 

390 

391def load_k8s() -> None: 

392 """Load Kubernetes configuration (in-cluster, or local kubeconfig off-cluster). 

393 

394 Inside a pod (``KUBERNETES_SERVICE_HOST`` set) an in-cluster failure is 

395 terminal and raises :class:`KubernetesConfigurationError` naming the 

396 missing token path and ServiceAccount — it must NOT fall through to 

397 ``load_kube_config()``, whose unguarded ``ConfigException`` used to kill 

398 the pod before its first SQS receive and crash-loop invisibly. 

399 """ 

400 try: 

401 config.load_incluster_config() 

402 log.info("Loaded in-cluster Kubernetes configuration") 

403 return 

404 except config.ConfigException as incluster_error: 

405 if os.environ.get("KUBERNETES_SERVICE_HOST"): 

406 raise KubernetesConfigurationError( 

407 _incluster_failure_detail(incluster_error) 

408 ) from incluster_error 

409 try: 

410 config.load_kube_config() 

411 log.info("Loaded local kubeconfig") 

412 except config.ConfigException as kubeconfig_error: 

413 raise KubernetesConfigurationError( 

414 "no Kubernetes credentials available: not running in a cluster " 

415 f"(KUBERNETES_SERVICE_HOST unset) and no local kubeconfig: {kubeconfig_error}" 

416 ) from kubeconfig_error 

417 

418 

419def validate_manifest(m: dict[str, Any]) -> tuple[bool, str]: 

420 """Validate a manifest before applying it to the cluster. 

421 

422 The queue processor mirrors the security checks performed by the REST 

423 `manifest_processor` service (``gco/services/manifest_processor.py``) 

424 so that the SQS path cannot bypass them. Checks performed: 

425 

426 1. **Namespace allowlist** — manifest namespace must be in 

427 ``ALLOWED_NAMESPACES`` (from ``ALLOWED_NAMESPACES`` env var, 

428 populated from ``cdk.json::job_validation_policy.allowed_namespaces``, 

429 shared with the REST manifest_processor). 

430 

431 2. **Pod-level security policy** (configurable via cdk.json:: 

432 job_validation_policy.manifest_security_policy, shared between both 

433 services). Rejects ``hostNetwork``, ``hostPID``, ``hostIPC``, 

434 ``hostPath`` volumes, privileged pod security context, and 

435 (if ``BLOCK_RUN_AS_ROOT``) pod-level ``runAsUser: 0``. 

436 

437 3. **Container-level security policy** — for every container kind 

438 (regular, init, ephemeral) rejects ``privileged``, 

439 ``allowPrivilegeEscalation``, ``capabilities.add``, and (if 

440 ``BLOCK_RUN_AS_ROOT``) container-level ``runAsUser: 0``. Iterating 

441 every container kind catches the classic "smuggle it via an init 

442 container" bypass. 

443 

444 4. **Image registry allowlist** — every container's image must come 

445 from ``TRUSTED_REGISTRIES`` (registry domains like ``nvcr.io``) 

446 or ``TRUSTED_DOCKERHUB_ORGS`` (Docker Hub orgs like ``nvidia``). 

447 Official Docker Hub images with no slash are always allowed. Empty or 

448 missing allowlists use the shared secure defaults. Keep explicit lists 

449 in sync with ``cdk.json::job_validation_policy.trusted_registries`` and 

450 ``trusted_dockerhub_orgs`` — CDK wires the same config into both 

451 services. 

452 

453 5. **Resource caps** — the TOTAL CPU, memory, and GPU across ALL 

454 containers (regular + init + ephemeral) must not exceed 

455 ``MAX_CPU``, ``MAX_MEMORY``, and ``MAX_GPU``. This matches 

456 ``manifest_processor._validate_resource_limits`` — K8s accounts 

457 init/ephemeral resources differently at scheduling time, but 

458 from an enforcement perspective we sum them so an operator's 

459 ``max_*_per_manifest`` budget is a hard cap regardless of where 

460 the request is placed. 

461 

462 A TrainJob has no single pod spec: checks 2-5 run over its decomposition 

463 (synthetic ``spec.trainer`` view plus every pod spec embedded in 

464 ``runtimePatches`` — see ``manifest_processor.TrainJobPodSpecs``), with 

465 the trainer view's resources counted once per ``numNodes`` for check 5 

466 and accelerator tolerations unioned across views for the toleration 

467 check. 

468 

469 Returns: 

470 ``(True, "")`` if the manifest is accepted, otherwise 

471 ``(False, reason)`` where ``reason`` is a human-readable string. 

472 """ 

473 kind = m.get("kind") 

474 if not kind: 

475 return False, "missing 'kind'" 

476 api = m.get("apiVersion") 

477 if not api: 

478 return False, "missing 'apiVersion'" 

479 meta = m.get("metadata") 

480 if not isinstance(meta, dict) or not meta.get("name"): 

481 return False, "missing 'metadata.name'" 

482 ns = meta.get("namespace", "gco-jobs") 

483 if ns not in ALLOWED_NAMESPACES: 

484 return False, f"namespace '{ns}' not in allowed list {ALLOWED_NAMESPACES}" 

485 

486 kind_valid, kind_error = validate_resource_kind(m, ALLOWED_KINDS) 

487 if not kind_valid: 

488 return False, kind_error or "resource kind is not allowed" 

489 

490 # Get pod spec(s) for security and resource checks, each with a replica 

491 # multiplier for resource-cap accounting. 

492 # Handle multiple resource shapes, matching manifest_processor._get_all_containers: 

493 # - Deployments / StatefulSets / ReplicaSets / DaemonSets / Jobs: spec.template.spec 

494 # - CronJob: spec.jobTemplate.spec.template.spec 

495 # - Pod (bare): spec (has 'containers' directly) 

496 # - TrainJob: synthetic spec.trainer view weighted by numNodes, plus every 

497 # pod spec embedded under spec (runtimePatches) weighted 1 — see 

498 # manifest_processor.TrainJobPodSpecs for why this decomposition exists. 

499 weighted_pod_specs: list[tuple[dict[str, Any], int]] = [] 

500 if kind == "TrainJob": 

501 trainjob_specs = extract_trainjob_pod_specs(m) 

502 if trainjob_specs.trainer is not None: 

503 weighted_pod_specs.append((trainjob_specs.trainer, trainjob_specs.num_nodes)) 

504 weighted_pod_specs.extend((item, 1) for item in trainjob_specs.embedded) 

505 toleration_hint_example = ( 

506 "examples/kubeflow-trainjob.yaml (GPU variant, via runtimePatches)" 

507 ) 

508 else: 

509 pod_spec = _extract_pod_spec(m) 

510 workload_kinds = { 

511 "Job", 

512 "CronJob", 

513 "Deployment", 

514 "StatefulSet", 

515 "DaemonSet", 

516 "ReplicaSet", 

517 "Pod", 

518 } 

519 if kind in workload_kinds and pod_spec is None: 

520 return False, f"{kind} manifest must contain a valid pod spec" 

521 if pod_spec is not None: 

522 weighted_pod_specs.append((pod_spec, 1)) 

523 toleration_hint_example = "examples/gpu-job.yaml" 

524 

525 if weighted_pod_specs: 

526 # --- Accelerator toleration check --- 

527 # Mirror manifest_processor._validate_tolerations: a job requesting a 

528 # GPU/Neuron/EFA resource must carry a matching toleration or it would 

529 # stay Pending forever on tainted accelerator nodes. For a TrainJob the 

530 # request usually lives in spec.trainer.resourcesPerNode while the 

531 # toleration can only be expressed through a runtimePatches pod spec, 

532 # so requests and tolerations are each unioned across every view 

533 # before matching. 

534 if REQUIRE_ACCELERATOR_TOLERATION: 

535 requested_taints: set[str] = set() 

536 tolerations: list[dict[str, Any]] = [] 

537 for pod_spec, _multiplier in weighted_pod_specs: 

538 requested_taints.update(_requested_accelerators(pod_spec)) 

539 tolerations.extend(pod_spec.get("tolerations", []) or []) 

540 for taint in requested_taints: 

541 if not _toleration_matches(tolerations, taint): 

542 hint = ( 

543 f"add a matching toleration (e.g. key '{taint}', operator " 

544 f"'Exists', effect 'NoSchedule'); see {toleration_hint_example}" 

545 ) 

546 return ( 

547 False, 

548 f"Job requests accelerator '{taint}' but no matching " 

549 f"toleration for taint {taint}=true:NoSchedule was found. {hint}", 

550 ) 

551 

552 for pod_spec, _multiplier in weighted_pod_specs: 

553 # --- Pod-level security policy checks --- 

554 # Mirror manifest_processor._validate_security_context so the SQS 

555 # path enforces the same policy as the REST path. 

556 if BLOCK_HOST_NETWORK and pod_spec.get("hostNetwork", False): 

557 return False, "hostNetwork is not permitted" 

558 if BLOCK_HOST_PID and pod_spec.get("hostPID", False): 

559 return False, "hostPID is not permitted" 

560 if BLOCK_HOST_IPC and pod_spec.get("hostIPC", False): 

561 return False, "hostIPC is not permitted" 

562 if BLOCK_HOST_PATH: 

563 for volume in pod_spec.get("volumes", []) or []: 

564 if volume.get("hostPath") is not None: 

565 return False, "hostPath volumes are not permitted" 

566 

567 pod_security_context = pod_spec.get("securityContext", {}) or {} 

568 if BLOCK_PRIVILEGED and pod_security_context.get("privileged", False): 

569 return False, "privileged pod security context is not permitted" 

570 if BLOCK_RUN_AS_ROOT: 

571 pod_run_as_user = pod_security_context.get("runAsUser") 

572 if pod_run_as_user is not None and pod_run_as_user == 0: 

573 return False, "running as root (runAsUser: 0) is not permitted" 

574 

575 # --- Container-level security policy checks --- 

576 # Every toggle is applied to every container kind (regular, init, 

577 # ephemeral). An init container running as root or with CAP_SYS_ADMIN 

578 # has the same blast radius as a regular container running the same 

579 # way; there is no reason to give any kind a free pass. 

580 for container_type, c in _iter_containers(pod_spec): 

581 cname = c.get("name", "unknown") 

582 sc = c.get("securityContext", {}) or {} 

583 if BLOCK_PRIVILEGED and sc.get("privileged", False): 

584 return ( 

585 False, 

586 f"{container_type} '{cname}': privileged containers are not permitted", 

587 ) 

588 if BLOCK_PRIVILEGE_ESCALATION and sc.get("allowPrivilegeEscalation", False): 

589 return ( 

590 False, 

591 f"{container_type} '{cname}': allowPrivilegeEscalation is not permitted", 

592 ) 

593 if BLOCK_ADDED_CAPABILITIES: 

594 added_caps = (sc.get("capabilities", {}) or {}).get("add", []) or [] 

595 if added_caps: 

596 return ( 

597 False, 

598 f"{container_type} '{cname}': added capabilities are not permitted", 

599 ) 

600 if BLOCK_RUN_AS_ROOT: 

601 ras = sc.get("runAsUser") 

602 if ras is not None and ras == 0: 

603 return ( 

604 False, 

605 f"{container_type} '{cname}': running as root (runAsUser: 0) is not permitted", 

606 ) 

607 

608 # Enforce image registry allowlist (matches manifest_processor semantics) 

609 for container_type, c in _iter_containers(pod_spec): 

610 image = c.get("image", "") 

611 if not _is_image_trusted(image): 

612 cname = c.get("name", "unknown") 

613 return ( 

614 False, 

615 f"{container_type} '{cname}': untrusted image source '{image}'", 

616 ) 

617 

618 # Enforce resource caps across ALL container kinds and pod-spec views. 

619 # Sum the resource requests/limits of every container (regular, 

620 # init, and ephemeral), scaled by each view's replica multiplier 

621 # (a TrainJob runs its trainer spec once per node — counting a 

622 # 16-node GPU job as one node would make the cap meaningless). 

623 # This is stricter than the K8s scheduler's accounting but matches 

624 # our security intent: an operator's configured "max CPU/memory/GPU 

625 # per manifest" is a hard cap on the total resources a submitter 

626 # can request regardless of which container kind carries the request. 

627 total_gpu = 0 

628 total_cpu = 0 

629 total_memory = 0 

630 for pod_spec, multiplier in weighted_pod_specs: 

631 for _container_type, c in _iter_containers(pod_spec): 

632 res = c.get("resources", {}) or {} 

633 limits = res.get("limits", {}) or {} 

634 requests = res.get("requests", {}) or {} 

635 gpu = limits.get("nvidia.com/gpu") or requests.get("nvidia.com/gpu", "0") # nosec B113 - dict.get(), not HTTP requests 

636 total_gpu += multiplier * int(gpu) 

637 cpu_str = limits.get("cpu") or requests.get("cpu", "0") # nosec B113 - dict.get(), not HTTP requests 

638 if isinstance(cpu_str, str) and cpu_str.endswith("m"): 

639 total_cpu += multiplier * int(cpu_str[:-1]) 

640 else: 

641 total_cpu += multiplier * int(float(cpu_str) * 1000) 

642 mem_str = limits.get("memory") or requests.get("memory", "0") # nosec B113 - dict.get(), not HTTP requests 

643 mem_bytes = _parse_memory_string(str(mem_str)) 

644 total_memory += multiplier * mem_bytes 

645 

646 errors = [] 

647 if total_gpu > MAX_GPU: 

648 errors.append(f"GPU {total_gpu} exceeds max {MAX_GPU}") 

649 if total_cpu > MAX_CPU: 

650 errors.append(f"CPU {total_cpu}m exceeds max {MAX_CPU}m") 

651 if total_memory > MAX_MEMORY: 

652 errors.append( 

653 f"Memory {total_memory / (1024**3):.0f}Gi " 

654 f"exceeds max {MAX_MEMORY / (1024**3):.0f}Gi" 

655 ) 

656 if errors: 

657 hint = ( 

658 "To raise limits, update queue_processor in cdk.json " 

659 "and redeploy (see examples/README.md#troubleshooting)" 

660 ) 

661 return False, "; ".join(errors) + f". {hint}" 

662 

663 return True, "" 

664 

665 

666def _extract_pod_spec(manifest: dict[str, Any]) -> dict[str, Any] | None: 

667 """Return the pod spec for any supported workload kind, or None. 

668 

669 Mirrors manifest_processor._extract_pod_spec so the SQS path and the 

670 REST path apply the same injection semantics. 

671 """ 

672 spec = manifest.get("spec") 

673 if not isinstance(spec, dict): 

674 return None 

675 

676 kind = manifest.get("kind", "") 

677 

678 # CronJob: spec.jobTemplate.spec.template.spec 

679 if kind == "CronJob": 

680 job_template = spec.get("jobTemplate") 

681 if isinstance(job_template, dict): 

682 job_spec = job_template.get("spec") 

683 if isinstance(job_spec, dict): 

684 template = job_spec.get("template") 

685 if isinstance(template, dict): 

686 pod_spec = template.get("spec") 

687 if isinstance(pod_spec, dict): 

688 return pod_spec 

689 return None 

690 

691 # Deployment / StatefulSet / DaemonSet / ReplicaSet / Job: spec.template.spec 

692 if "template" in spec: 

693 template = spec.get("template") 

694 if isinstance(template, dict): 

695 pod_spec = template.get("spec") 

696 if isinstance(pod_spec, dict): 

697 return pod_spec 

698 return None 

699 

700 # Bare Pod: spec contains "containers" directly 

701 if "containers" in spec: 

702 return spec 

703 

704 return None 

705 

706 

707def _inject_security_defaults(manifest: dict[str, Any]) -> dict[str, Any]: 

708 """Inject security defaults into a user-submitted manifest in-place. 

709 

710 Currently sets ``automountServiceAccountToken: false`` on the pod spec 

711 unless the user has explicitly set it either way (uses setdefault). 

712 

713 Mirrors manifest_processor._inject_security_defaults so jobs submitted 

714 via SQS get the same SA-token-theft protection as those submitted via 

715 the REST API. 

716 

717 For a TrainJob the default is injected into every pod spec embedded in 

718 ``runtimePatches`` (live references into the manifest, so setdefault 

719 mutates it in place); the base pod template comes from the shipped 

720 ClusterTrainingRuntime, which already disables the token. 

721 """ 

722 if manifest.get("kind") == "TrainJob": 

723 for embedded in extract_trainjob_pod_specs(manifest).embedded: 

724 embedded.setdefault("automountServiceAccountToken", False) 

725 return manifest 

726 pod_spec = _extract_pod_spec(manifest) 

727 if pod_spec is not None: 

728 pod_spec.setdefault("automountServiceAccountToken", False) 

729 return manifest 

730 

731 

732def _is_job_finished(job_resource: dict[str, Any]) -> bool: 

733 """Return whether a Kubernetes Job has a true terminal condition.""" 

734 status = job_resource.get("status", {}) 

735 conditions = status.get("conditions") or [] if isinstance(status, dict) else [] 

736 return any( 

737 isinstance(condition, dict) 

738 and condition.get("type") in ("Complete", "Failed") 

739 and condition.get("status") == "True" 

740 for condition in conditions 

741 ) 

742 

743 

744def apply_manifest(m: dict[str, Any]) -> ResourceStatus: 

745 """Apply one prevalidated manifest and return an explicit operation status. 

746 

747 Unsupported API resources are failures, never successful skips. This keeps 

748 the owning SQS message available for retry and eventual DLQ inspection. 

749 """ 

750 # Inject security defaults BEFORE applying so user pods never 

751 # auto-mount the default SA token (T-022 / M-113 parity with the 

752 # REST manifest_processor path). 

753 _inject_security_defaults(m) 

754 

755 api_version = m["apiVersion"] 

756 kind = m["kind"] 

757 name = m["metadata"]["name"] 

758 namespace = m["metadata"].get("namespace", "gco-jobs") 

759 

760 def status(result: str, message: str) -> ResourceStatus: 

761 return ResourceStatus( 

762 api_version=api_version, 

763 kind=kind, 

764 name=name, 

765 namespace=namespace, 

766 status=result, 

767 message=message, 

768 ) 

769 

770 dyn = dynamic.DynamicClient(client.ApiClient()) 

771 try: 

772 resource = dyn.resources.get(api_version=api_version, kind=kind) 

773 except ResourceNotFoundError: 

774 # A policy-allowed kind whose addon is not installed gets the 

775 # actionable remedy appended; the stable prefix is kept for DLQ 

776 # triage tooling. 

777 addon_hint = ADDON_KIND_HINTS.get(kind) 

778 detail = f" ({addon_hint})" if addon_hint else "" 

779 return status("failed", f"Unsupported Kubernetes resource {api_version}/{kind}{detail}") 

780 

781 # For Jobs, delete completed/failed ones first so re-submission works. 

782 # Without this, re-submitting the same job name would fail with a 409 conflict 

783 # because Kubernetes doesn't allow creating a Job with the same name as an 

784 # existing one (even if it's finished). 

785 if kind == "Job": 

786 try: 

787 existing = resource.get(name=name, namespace=namespace) 

788 if _is_job_finished(existing): 

789 log.info("Deleting finished Job %s/%s before re-creation", namespace, name) 

790 resource.delete( 

791 name=name, 

792 namespace=namespace, 

793 body=client.V1DeleteOptions(propagation_policy="Background"), 

794 ) 

795 time.sleep(2) 

796 except (NotFoundError, ApiException) as e: 

797 log.debug("Pre-create lookup for Job %s/%s failed: %s", namespace, name, e) 

798 

799 # Create-or-update pattern: try create first, fall back to patch on 409 (conflict). 

800 # This is idempotent — safe to retry without side effects. 

801 try: 

802 if resource.namespaced: 

803 resource.create(body=m, namespace=namespace) 

804 else: 

805 resource.create(body=m) 

806 return status("created", "Resource created successfully") 

807 except ApiException as e: 

808 if e.status == 409: 

809 try: 

810 if resource.namespaced: 

811 resource.patch(body=m, name=name, namespace=namespace) 

812 else: 

813 resource.patch(body=m, name=name) 

814 return status("updated", "Resource updated successfully") 

815 except ApiException as patch_err: 

816 return status("failed", f"Patch failed: {patch_err.reason}") 

817 return status("failed", f"Create failed: {e.reason}") 

818 except Exception as e: 

819 return status("failed", f"Unexpected apply error: {e}") 

820 

821 

822def process_one_message() -> bool: 

823 """Receive one SQS message and delete it only after complete success. 

824 

825 ``True`` means either the poll was empty or the received message was fully 

826 validated, applied, and deleted. ``False`` means the message was not 

827 acknowledged (or queue configuration was invalid). Malformed, empty, 

828 invalid, unsupported, and apply-failed messages deliberately remain in SQS 

829 for visibility-timeout retries and eventual dead-letter-queue handling. 

830 """ 

831 if not QUEUE_URL: 

832 log.error("JOB_QUEUE_URL not set") 

833 return False 

834 

835 sqs = boto3.client("sqs", region_name=REGION) 

836 

837 resp = sqs.receive_message( 

838 QueueUrl=QUEUE_URL, 

839 MaxNumberOfMessages=1, 

840 WaitTimeSeconds=5, 

841 MessageAttributeNames=["All"], 

842 AttributeNames=["ApproximateReceiveCount"], 

843 ) 

844 

845 messages = resp.get("Messages", []) 

846 if not messages: 

847 log.info("No messages in queue") 

848 return True 

849 

850 msg = messages[0] 

851 receipt = msg.get("ReceiptHandle") 

852 if not receipt: 

853 log.error("Received SQS message without a receipt handle; cannot acknowledge it") 

854 return False 

855 

856 try: 

857 body = json.loads(msg.get("Body", "")) 

858 except (json.JSONDecodeError, TypeError) as e: 

859 log.error("Malformed SQS message body; retaining for retry/DLQ: %s", e) 

860 return False 

861 

862 if not isinstance(body, dict): 

863 log.error("SQS message body must be a JSON object; retaining for retry/DLQ") 

864 return False 

865 

866 job_id = body.get("job_id", "unknown") 

867 manifests = body.get("manifests")