Coverage for gco / job_admission.py: 100.00%
344 statements
« prev ^ index » next coverage.py v7.13.5, created at 2026-09-14 22:07 +0000
1"""Pure job-admission policy checks, free of Kubernetes and AWS clients.
3Everything here is a function of a manifest and a policy, nothing else. That
4constraint is the point: the same checks have to run in four places that cannot
5share a client.
7 * the manifest-processor REST service, at submission time
8 * the SQS queue processor, when it drains a queued manifest
9 * ``gco jobs check-policy``, against a policy read back over HTTP from a
10 deployed region
11 * the offline example validator, against caps read out of ``cdk.json``
13Before this module existed the first two shared these helpers by importing them
14from ``gco.services.manifest_processor``, which meant importing ``kubernetes``
15to ask a question about a dict. The offline validator did exactly that. Worse,
16every check that needed deployment state was a *method* reading ``self``, so a
17caller holding a policy document rather than a live processor had no way in and
18the only option was to reimplement the rule. A second implementation of an
19admission rule is a rule that will drift, and it drifts silently -- the copy
20keeps passing while the real gate changes underneath it.
22So the deployment-dependent checks take a :class:`JobValidationPolicy` instead
23of ``self``, and ``ManifestProcessor`` builds one from its own attributes and
24delegates. There is one implementation of each rule, and the offline and
25multi-region callers exercise the same code the cluster does.
26"""
28from __future__ import annotations
30import logging
31from collections.abc import Collection
32from typing import Any, NamedTuple
34from gco.services.structured_logging import sanitize_log_value
36logger = logging.getLogger(__name__)
39# Accelerator resource keys and their corresponding node taint keys. GCO
40# nodepools taint accelerator nodes with these keys (authoritative list:
41# regional_stack._ADDON_NODE_TOLERATIONS), so a job requesting one of these
42# resources must carry a matching toleration or it will never schedule.
43# Taint key == resource key for all three. Kept in sync with the mirror in
44# gco/services/queue_processor.py::ACCELERATOR_TAINTS.
45ACCELERATOR_TAINTS = ("nvidia.com/gpu", "aws.amazon.com/neuron", "vpc.amazonaws.com/efa")
46# Exact pinned group/version for the Kubeflow Trainer v2 TrainJob kind.
47# Kept in lockstep with the kubeflow-trainer chart in
48# lambda/helm-installer/charts.yaml and the extracted runtime manifest in
49# lambda/kubectl-applier-simple/manifests/.
50TRAINJOB_API_VERSION = "trainer.kubeflow.org/v1alpha1"
51# Authoritative resource-kind policy shared by the REST and SQS submission
52# paths. Keep the fallback here so both services fail closed to the same set
53# when ALLOWED_KINDS is not explicitly configured.
54DEFAULT_ALLOWED_KINDS = (
55 "Job",
56 "CronJob",
57 "Deployment",
58 "StatefulSet",
59 "DaemonSet",
60 "Service",
61 "ConfigMap",
62 "Pod",
63 "TrainJob",
64)
65# Authoritative image-source defaults shared by the REST and SQS submission
66# paths. Keep these centralized so missing deployment wiring cannot make either
67# path weaker or let their allowlists drift independently.
68DEFAULT_TRUSTED_REGISTRIES = (
69 "docker.io",
70 "gcr.io",
71 "quay.io",
72 "registry.k8s.io",
73 "k8s.gcr.io",
74 "public.ecr.aws",
75 "nvcr.io",
76 # Org-scoped GHCR prefix (matched via the startswith branch) for the
77 # HuggingFace TGI image shipped in examples/inference-tgi.yaml. Scoped to
78 # the org rather than all of ghcr.io on purpose.
79 "ghcr.io/huggingface",
80)
81DEFAULT_TRUSTED_DOCKERHUB_ORGS = (
82 "nvidia",
83 "pytorch",
84 "rayproject",
85 "tensorflow",
86 "huggingface",
87 "amazon",
88 "bitnami",
89 # Official orgs of the vLLM and SGLang projects — the images shipped in
90 # examples/inference-vllm.yaml and examples/inference-sglang.yaml. Kept in
91 # lockstep with cdk.json job_validation_policy.trusted_dockerhub_orgs (see
92 # tests/test_manifest_processor_extended.py).
93 "vllm",
94 "lmsysorg",
95 "gco",
96)
97# CRUD endpoints accept only these exact built-in, namespaced GVKs. A kind-only
98# allowlist is insufficient because a custom API group can define the same kind
99# name, and cluster-scoped resources must never be reachable through a
100# namespace-shaped user endpoint.
101RESOURCE_API_VERSIONS: dict[str, frozenset[str]] = {
102 "Job": frozenset({"batch/v1"}),
103 "CronJob": frozenset({"batch/v1"}),
104 "Deployment": frozenset({"apps/v1"}),
105 "StatefulSet": frozenset({"apps/v1"}),
106 "DaemonSet": frozenset({"apps/v1"}),
107 "Service": frozenset({"v1"}),
108 "ConfigMap": frozenset({"v1"}),
109 "Pod": frozenset({"v1"}),
110 "TrainJob": frozenset({TRAINJOB_API_VERSION}),
111}
112# Actionable guidance when an allowed kind's CRD/controller is absent from
113# the cluster. Without this, a policy-allowed manifest whose addon is
114# disabled fails with an unactionable "Unknown resource type" — or worse,
115# is accepted and never reconciles.
116ADDON_KIND_HINTS: dict[str, str] = {
117 "TrainJob": (
118 "TrainJob requires the kubeflow-trainer addon; enable "
119 'helm.kubeflow_trainer ("enabled": true) in cdk.json and redeploy '
120 "the regional stack"
121 ),
122}
125def _extract_validation_pod_spec(manifest: dict[str, Any]) -> dict[str, Any]:
126 """Locate the pod spec for image-source validation across resource shapes."""
127 spec = manifest.get("spec", {})
128 pod_spec: Any = {}
129 if "template" in spec:
130 pod_spec = spec.get("template", {}).get("spec", {})
131 elif "jobTemplate" in spec:
132 pod_spec = spec.get("jobTemplate", {}).get("spec", {}).get("template", {}).get("spec", {})
133 elif "containers" in spec:
134 pod_spec = spec
135 return pod_spec if isinstance(pod_spec, dict) else {}
138class TrainJobPodSpecs(NamedTuple):
139 """Every pod-spec-shaped view a TrainJob manifest can cause to run.
141 A TrainJob has no ``spec.template``: its pods come from the referenced
142 ClusterTrainingRuntime, customized by first-class fields
143 (``spec.trainer``) and by arbitrary runtime patches
144 (``spec.runtimePatches[].trainingRuntimeSpec`` — which can nest complete
145 pod specs, including containers, volumes, and security contexts).
146 Validating only the classic single-pod-spec shapes would let every
147 image-trust, security-context, and resource-cap check pass vacuously,
148 so TrainJob validation runs over this decomposition instead.
150 Attributes:
151 trainer: A synthetic pod spec carrying ``spec.trainer``'s image and
152 per-node resources, or ``None`` when neither is set (the runtime
153 defaults then apply — our shipped runtimes pin trusted images).
154 embedded: Every dict carrying a container list found anywhere under
155 ``spec`` (today that means inside ``runtimePatches``); these are
156 live references into the manifest, so security-default injection
157 through them mutates the manifest.
158 num_nodes: ``spec.trainer.numNodes`` (minimum 1) — the replica
159 multiplier for the trainer spec's resource totals.
160 """
162 trainer: dict[str, Any] | None
163 embedded: list[dict[str, Any]]
164 num_nodes: int
167def _collect_embedded_pod_specs(node: Any, found: list[dict[str, Any]]) -> None:
168 """Recursively collect every dict that carries a container list.
170 Walking the whole structure — rather than enumerating known patch paths —
171 means a future TrainJob field that can smuggle a container is validated
172 by default instead of silently skipped.
173 """
174 if isinstance(node, dict):
175 if any(
176 isinstance(node.get(key), list)
177 for key in ("containers", "initContainers", "ephemeralContainers")
178 ):
179 found.append(node)
180 for value in node.values():
181 _collect_embedded_pod_specs(value, found)
182 elif isinstance(node, list):
183 for item in node:
184 _collect_embedded_pod_specs(item, found)
187def extract_trainjob_pod_specs(manifest: dict[str, Any]) -> TrainJobPodSpecs:
188 """Decompose a TrainJob manifest into validatable pod-spec views.
190 Shared by the REST and SQS submission paths (the queue processor imports
191 this) so TrainJob validation semantics cannot drift between them.
192 """
193 spec = manifest.get("spec")
194 spec = spec if isinstance(spec, dict) else {}
195 trainer = spec.get("trainer")
196 trainer = trainer if isinstance(trainer, dict) else {}
198 pseudo: dict[str, Any] | None = None
199 if trainer.get("image") or isinstance(trainer.get("resourcesPerNode"), dict):
200 container: dict[str, Any] = {"name": "trainer"}
201 if trainer.get("image"):
202 container["image"] = trainer["image"]
203 if isinstance(trainer.get("resourcesPerNode"), dict):
204 container["resources"] = {
205 "requests": trainer["resourcesPerNode"].get("requests", {}) or {},
206 "limits": trainer["resourcesPerNode"].get("limits", {}) or {},
207 }
208 pseudo = {"containers": [container]}
210 embedded: list[dict[str, Any]] = []
211 _collect_embedded_pod_specs(spec, embedded)
213 raw_nodes = trainer.get("numNodes")
214 try:
215 num_nodes = max(1, int(raw_nodes)) if raw_nodes is not None else 1
216 except TypeError, ValueError:
217 num_nodes = 1
218 return TrainJobPodSpecs(trainer=pseudo, embedded=embedded, num_nodes=num_nodes)
221def trainjob_validation_pod_specs(manifest: dict[str, Any]) -> list[dict[str, Any]]:
222 """All pod-spec views of a TrainJob, synthetic trainer spec first."""
223 specs = extract_trainjob_pod_specs(manifest)
224 views = [specs.trainer] if specs.trainer is not None else []
225 views.extend(specs.embedded)
226 return views
229def _iter_all_containers(pod_spec: dict[str, Any]) -> list[tuple[str, dict[str, Any]]]:
230 """All (container_type, container) pairs incl. init and ephemeral containers."""
231 result: list[tuple[str, dict[str, Any]]] = []
232 for container in pod_spec.get("containers", []):
233 result.append(("container", container))
234 for container in pod_spec.get("initContainers", []):
235 result.append(("initContainer", container))
236 for container in pod_spec.get("ephemeralContainers", []):
237 result.append(("ephemeralContainer", container))
238 return result
241def _is_trusted_registry_domain(entry: str) -> bool:
242 """True when a registry entry is a domain (dot/colon) rather than a Hub org."""
243 return "." in entry or ":" in entry
246def validate_image_sources(
247 manifest: dict[str, Any],
248 trusted_registries: list[str] | tuple[str, ...] = DEFAULT_TRUSTED_REGISTRIES,
249 trusted_dockerhub_orgs: list[str] | tuple[str, ...] = DEFAULT_TRUSTED_DOCKERHUB_ORGS,
250) -> tuple[bool, str | None]:
251 """Validate container image sources against the trust allowlists.
253 Pure function (no Kubernetes client, no configuration loading) so offline
254 validators — the example-manifest static checks, tests — apply the exact
255 logic the deployed services enforce. ``ManifestProcessor`` delegates here.
257 Matching logic:
258 1. No ``/`` in the image → official Docker Hub image (always allowed)
259 2. First segment contains a dot/colon → registry domain → exact match or
260 org-scoped prefix match against ``trusted_registries``
261 3. Otherwise → Docker Hub org → match against ``trusted_dockerhub_orgs``
263 TrainJob manifests carry no single pod spec; every image they can run —
264 ``spec.trainer.image`` plus any container smuggled in through
265 ``runtimePatches`` — is validated through the TrainJob decomposition.
266 """
267 try:
268 if manifest.get("kind") == "TrainJob":
269 pod_specs = trainjob_validation_pod_specs(manifest)
270 else:
271 pod_specs = [_extract_validation_pod_spec(manifest)]
272 for pod_spec in pod_specs:
273 failure = _untrusted_container_in_pod_spec(
274 pod_spec, trusted_registries, trusted_dockerhub_orgs
275 )
276 if failure is not None:
277 return False, failure
278 return True, None
279 except Exception as e:
280 logger.error(f"Error validating image sources: {e}")
281 return False, f"Image source validation error: {e}"
284def _untrusted_container_in_pod_spec(
285 pod_spec: dict[str, Any],
286 trusted_registries: list[str] | tuple[str, ...],
287 trusted_dockerhub_orgs: list[str] | tuple[str, ...],
288) -> str | None:
289 """Return the failure message for the first untrusted image, or None."""
290 for container_type, container in _iter_all_containers(pod_spec):
291 image = container.get("image", "")
292 if not image:
293 continue
294 is_trusted = False
295 if "/" not in image:
296 is_trusted = True
297 else:
298 first_segment = image.split("/")[0]
299 if _is_trusted_registry_domain(first_segment):
300 for registry in trusted_registries:
301 if first_segment == registry or image.startswith(registry + "/"):
302 is_trusted = True
303 break
304 elif first_segment in trusted_dockerhub_orgs:
305 is_trusted = True
306 if not is_trusted:
307 container_name = container.get("name", "unknown")
308 # image comes from the user-submitted manifest; sanitize it
309 # before logging to prevent log injection / forging (CWE-117).
310 logger.warning("Untrusted image source: %s", sanitize_log_value(image))
311 return f"{container_type} '{container_name}': Untrusted image source '{image}'"
312 return None
315def validate_resource_kind(
316 manifest: dict[str, Any],
317 allowed_kinds: Collection[str] = DEFAULT_ALLOWED_KINDS,
318) -> tuple[bool, str | None]:
319 """Validate a manifest kind against the shared submission allowlist.
321 ``allowed_kinds`` is any collection because the allowlist arrives as a tuple
322 of defaults, a set off a live processor, and a frozenset off a
323 :class:`JobValidationPolicy`; only membership and ordering-for-display are
324 used, so narrowing the type would force callers into pointless conversions.
325 """
326 kind = manifest.get("kind", "")
327 allowed = set(allowed_kinds)
328 if kind not in allowed:
329 return (
330 False,
331 f"Resource kind '{kind}' is not allowed. Allowed kinds: {sorted(allowed)}",
332 )
333 return True, None
336def _positive_quantity(value: Any) -> bool:
337 """True if a K8s resource quantity is present and greater than zero."""
338 if value is None:
339 return False
340 try:
341 return float(value) > 0
342 except TypeError, ValueError:
343 # A non-numeric quantity is still an explicit request.
344 return True
347def _toleration_matches(tolerations: list[dict[str, Any]], taint_key: str) -> bool:
348 """Return True if *tolerations* tolerates the ``<taint_key>=true:NoSchedule`` taint.
350 A toleration matches when its ``key`` equals *taint_key*, its effect is
351 empty (matches all effects) or ``NoSchedule``, and it either uses
352 ``operator: Exists`` or ``operator: Equal`` with ``value: "true"``.
353 Kept in sync with the mirror in queue_processor._toleration_matches.
354 """
355 for tol in tolerations:
356 if not isinstance(tol, dict) or tol.get("key") != taint_key:
357 continue
358 effect = tol.get("effect", "")
359 if effect not in ("", "NoSchedule"):
360 continue
361 operator = tol.get("operator", "Equal")
362 if operator == "Exists":
363 return True
364 if operator == "Equal" and str(tol.get("value")) == "true":
365 return True
366 return False
369# ---------------------------------------------------------------------------
370# Quantity parsing
371# ---------------------------------------------------------------------------
372# Deliberately the same narrow parsers the admission path has always used,
373# rather than gco.resource_governance.parse_k8s_quantity. They round and reject
374# differently at the edges (a bare float CPU, an exponent suffix), and a
375# pre-submit check that parses more permissively than the gate would call a job
376# admissible that the gate then rejects. Matching the gate matters more here
377# than being the better parser.
380def parse_cpu_millicores(cpu_str: str) -> int:
381 """Parse a Kubernetes CPU quantity to millicores."""
382 if not cpu_str:
383 return 0
384 cpu_str = cpu_str.strip()
385 if cpu_str.endswith("m"):
386 return int(cpu_str[:-1])
387 return int(cpu_str) * 1000
390def parse_memory_bytes(memory_str: str) -> int:
391 """Parse a Kubernetes memory quantity to bytes."""
392 if not memory_str:
393 return 0
394 memory_str = memory_str.strip()
395 if memory_str.endswith("Ki"):
396 return int(memory_str[:-2]) * 1024
397 if memory_str.endswith("Mi"):
398 return int(memory_str[:-2]) * 1024 * 1024
399 if memory_str.endswith("Gi"):
400 return int(memory_str[:-2]) * 1024 * 1024 * 1024
401 if memory_str.endswith("Ti"):
402 return int(memory_str[:-2]) * 1024 * 1024 * 1024 * 1024
403 if memory_str.endswith("k"):
404 return int(memory_str[:-1]) * 1000
405 if memory_str.endswith("M"):
406 return int(memory_str[:-1]) * 1000 * 1000
407 if memory_str.endswith("G"):
408 return int(memory_str[:-1]) * 1000 * 1000 * 1000
409 return int(memory_str)
412def extract_pod_spec(manifest: dict[str, Any]) -> dict[str, Any] | None:
413 """Extract the pod spec from a manifest, handling all workload types.
415 Supports:
416 - Deployment / StatefulSet / DaemonSet / ReplicaSet → spec.template.spec
417 - Job → spec.template.spec
418 - CronJob → spec.jobTemplate.spec.template.spec
419 - Bare Pod → spec (when ``containers`` key is present)
421 Returns:
422 The pod spec dict (mutable reference), or ``None`` if the manifest
423 does not contain a recognisable pod spec.
424 """
425 spec = manifest.get("spec")
426 if spec is None or not isinstance(spec, dict):
427 return None
428 kind = manifest.get("kind", "")
429 # CronJob: spec.jobTemplate.spec.template.spec
430 if kind == "CronJob":
431 job_template = spec.get("jobTemplate")
432 if isinstance(job_template, dict):
433 job_spec = job_template.get("spec")
434 if isinstance(job_spec, dict):
435 template = job_spec.get("template")
436 if isinstance(template, dict):
437 pod_spec = template.get("spec")
438 if isinstance(pod_spec, dict):
439 return pod_spec
440 return None
441 # Deployment / StatefulSet / DaemonSet / ReplicaSet / Job:
442 # spec.template.spec
443 if "template" in spec:
444 template = spec.get("template")
445 if isinstance(template, dict):
446 pod_spec = template.get("spec")
447 if isinstance(pod_spec, dict):
448 return pod_spec
449 return None
450 # Bare Pod: spec contains "containers" directly
451 if "containers" in spec:
452 return spec
453 return None
456def requested_accelerators(pod_spec: dict[str, Any]) -> set[str]:
457 """Return the accelerator taint keys any container requests a nonzero
458 quantity of."""
459 requested: set[str] = set()
460 for _container_type, container in _iter_all_containers(pod_spec):
461 resources = container.get("resources", {}) or {}
462 for section in ("requests", "limits"):
463 values = resources.get(section, {}) or {}
464 for taint in ACCELERATOR_TAINTS:
465 if _positive_quantity(values.get(taint)):
466 requested.add(taint)
467 return requested
470def weighted_pod_specs(manifest: dict[str, Any]) -> list[tuple[dict[str, Any], int]]:
471 """Return every pod spec in *manifest* paired with its replica multiplier.
473 A TrainJob runs its trainer spec once per node, so the manifest's total is
474 ``numNodes`` x the per-node request. Counting a 16-node GPU job as one node
475 would make the per-manifest cap meaningless.
476 """
477 specs: list[tuple[dict[str, Any], int]] = []
478 if manifest.get("kind") == "TrainJob":
479 trainjob_specs = extract_trainjob_pod_specs(manifest)
480 if trainjob_specs.trainer is not None:
481 specs.append((trainjob_specs.trainer, trainjob_specs.num_nodes))
482 specs.extend((item, 1) for item in trainjob_specs.embedded)
483 return specs
485 pod_spec = extract_pod_spec(manifest)
486 specs.append((pod_spec if pod_spec is not None else {}, 1))
487 return specs
490# ---------------------------------------------------------------------------
491# The policy a manifest is judged against
492# ---------------------------------------------------------------------------
495class JobValidationPolicy(NamedTuple):
496 """Everything the admission checks need to know about a deployment.
498 Immutable and client-free, so it can come from any of three sources that
499 know progressively less:
501 ``from_processor_attributes``
502 what a live ManifestProcessor enforces right now. Authoritative.
504 ``from_policy_document``
505 the body of ``GET /api/v1/policy`` from a deployed region. Also
506 authoritative -- that endpoint reads the same attributes -- but it
507 arrives over the network, so it may be stale by the age of the response.
509 ``from_cdk_context``
510 ``cdk.json``'s ``job_validation_policy``. **Not** authoritative, and the
511 gap is not hypothetical: CDK appends the project's own ECR registry
512 hostnames to ``trusted_registries`` at synth time, so a deployed region
513 trusts registries that appear nowhere in the file. A live run on
514 2026-08-26 showed two such hostnames in the effective allowlist. The
515 file also says nothing about which commit a region was deployed from.
516 Callers using this source must say so, and must not turn a rejection
517 into a hard failure.
518 """
520 max_cpu_millicores: int
521 max_memory_bytes: int
522 max_gpu_count: int
523 allowed_namespaces: frozenset[str]
524 allowed_kinds: frozenset[str]
525 trusted_registries: tuple[str, ...]
526 trusted_dockerhub_orgs: tuple[str, ...]
527 require_accelerator_toleration: bool
528 security: dict[str, bool]
529 validation_enabled: bool = True
530 yaml_max_depth: int = 50
532 # -- constructors -----------------------------------------------------
533 @classmethod
534 def from_policy_document(cls, document: dict[str, Any]) -> JobValidationPolicy:
535 """Build from a ``GET /api/v1/policy`` response body."""
536 caps = document.get("manifest_caps", {}) or {}
537 defaults = _default_security_flags()
538 security = dict(defaults)
539 security.update(
540 {
541 key: bool(value)
542 for key, value in (document.get("manifest_security_policy", {}) or {}).items()
543 if key in defaults
544 }
545 )
546 return cls(
547 max_cpu_millicores=int(caps.get("max_cpu_millicores", 0)),
548 max_memory_bytes=int(caps.get("max_memory_bytes", 0)),
549 max_gpu_count=int(caps.get("max_gpu_count", 0)),
550 allowed_namespaces=frozenset(document.get("allowed_namespaces", ()) or ()),
551 allowed_kinds=frozenset(document.get("allowed_kinds", ()) or ()),
552 trusted_registries=tuple(sorted(document.get("trusted_registries", ()) or ())),
553 trusted_dockerhub_orgs=tuple(sorted(document.get("trusted_dockerhub_orgs", ()) or ())),
554 require_accelerator_toleration=bool(
555 document.get("require_accelerator_toleration", True)
556 ),
557 security=security,
558 validation_enabled=bool(document.get("validation_enabled", True)),
559 yaml_max_depth=int(document.get("yaml_max_depth", 50)),
560 )
562 @classmethod
563 def from_cdk_context(cls, job_validation_policy: dict[str, Any]) -> JobValidationPolicy:
564 """Build from ``cdk.json``'s ``context.job_validation_policy``.
566 Applies the same fallbacks the service applies when a key is absent, so
567 an unset key reads as the shipped default rather than as zero.
568 """
569 from gco.resource_governance import DEFAULT_MANIFEST_RESOURCE_CAPS
571 cfg = job_validation_policy or {}
572 defaults = _default_security_flags()
573 security = dict(defaults)
574 security.update(
575 {
576 key: bool(value)
577 for key, value in (cfg.get("manifest_security_policy", {}) or {}).items()
578 if key in defaults
579 }
580 )
581 return cls(
582 max_cpu_millicores=parse_cpu_millicores(
583 str(
584 cfg.get(
585 "max_cpu_per_manifest",
586 DEFAULT_MANIFEST_RESOURCE_CAPS["max_cpu_per_manifest"],
587 )
588 )
589 ),
590 max_memory_bytes=parse_memory_bytes(
591 str(
592 cfg.get(
593 "max_memory_per_manifest",
594 DEFAULT_MANIFEST_RESOURCE_CAPS["max_memory_per_manifest"],
595 )
596 )
597 ),
598 max_gpu_count=int(
599 cfg.get(
600 "max_gpu_per_manifest",
601 DEFAULT_MANIFEST_RESOURCE_CAPS["max_gpu_per_manifest"],
602 )
603 ),
604 allowed_namespaces=frozenset(cfg.get("allowed_namespaces", ("gco-jobs",))),
605 allowed_kinds=frozenset(cfg.get("allowed_kinds", DEFAULT_ALLOWED_KINDS)),
606 trusted_registries=tuple(
607 sorted(cfg.get("trusted_registries", DEFAULT_TRUSTED_REGISTRIES))
608 ),
609 trusted_dockerhub_orgs=tuple(
610 sorted(cfg.get("trusted_dockerhub_orgs", DEFAULT_TRUSTED_DOCKERHUB_ORGS))
611 ),
612 require_accelerator_toleration=bool(cfg.get("require_accelerator_toleration", True)),
613 security=security,
614 validation_enabled=bool(cfg.get("validation_enabled", True)),
615 yaml_max_depth=int(cfg.get("yaml_max_depth", 50)),
616 )
619def _default_security_flags() -> dict[str, bool]:
620 """The shipped security-policy defaults, as a plain mutable dict."""
621 from gco.manifest_security_policy import MANIFEST_SECURITY_POLICY_DEFAULTS
623 return dict(MANIFEST_SECURITY_POLICY_DEFAULTS)
626# ---------------------------------------------------------------------------
627# The checks
628# ---------------------------------------------------------------------------
631def check_resource_caps(manifest: dict[str, Any], policy: JobValidationPolicy) -> tuple[bool, str]:
632 """Check a manifest's aggregate CPU / memory / GPU against the caps.
634 Returns ``(is_valid, error_message)``; the message is empty when valid.
635 """
636 try:
637 errors: list[str] = []
638 total_cpu = 0
639 total_memory = 0
640 total_gpu = 0
642 for pod_spec, multiplier in weighted_pod_specs(manifest):
643 for _container_type, container in _iter_all_containers(pod_spec):
644 resources = container.get("resources", {})
645 requests = resources.get("requests", {})
646 limits = resources.get("limits", {})
648 # Use limits if available, otherwise requests.
649 cpu = limits.get("cpu") or requests.get("cpu", "0") # nosec B113 - dict.get()
650 total_cpu += multiplier * parse_cpu_millicores(cpu)
652 memory = limits.get("memory") or requests.get("memory", "0") # nosec B113
653 total_memory += multiplier * parse_memory_bytes(memory)
655 gpu = limits.get("nvidia.com/gpu") or requests.get("nvidia.com/gpu", "0") # nosec B113
656 total_gpu += multiplier * int(gpu)
658 if total_cpu > policy.max_cpu_millicores:
659 logger.warning(f"CPU limit exceeded: {total_cpu}m > {policy.max_cpu_millicores}m")
660 errors.append(f"CPU {total_cpu}m exceeds max {policy.max_cpu_millicores}m")
662 if total_memory > policy.max_memory_bytes:
663 logger.warning(f"Memory limit exceeded: {total_memory} > {policy.max_memory_bytes}")
664 mem_gb = policy.max_memory_bytes / (1024**3)
665 req_gb = total_memory / (1024**3)
666 errors.append(f"Memory {req_gb:.0f}Gi exceeds max {mem_gb:.0f}Gi")
668 if total_gpu > policy.max_gpu_count:
669 logger.warning(f"GPU limit exceeded: {total_gpu} > {policy.max_gpu_count}")
670 errors.append(f"GPU {total_gpu} exceeds max {policy.max_gpu_count}")
672 if errors:
673 hint = (
674 "To raise limits, update resource_quotas in cdk.json "
675 "and redeploy (see examples/README.md#troubleshooting)"
676 )
677 return False, "; ".join(errors) + f". {hint}"
679 return True, ""
681 except Exception as e:
682 logger.error(f"Error validating resource limits: {e}")
683 return False, f"Resource limit validation error: {e}"
686def check_security_context(
687 manifest: dict[str, Any], policy: JobValidationPolicy
688) -> tuple[bool, str | None]:
689 """Check pod- and container-level security settings against the policy.
691 Returns ``(is_valid, error_message)``; the message is ``None`` when valid.
692 """
693 try:
694 flags = policy.security
695 spec = manifest.get("spec", {})
697 # A TrainJob can nest complete pod specs inside runtimePatches, so every
698 # one of them gets the full check — otherwise privileged containers or
699 # hostPath volumes could ride in through a patch.
700 pod_specs: list[dict[str, Any]]
701 if manifest.get("kind") == "TrainJob":
702 pod_specs = trainjob_validation_pod_specs(manifest)
703 else:
704 pod_spec = None
705 if "template" in spec:
706 pod_spec = spec.get("template", {}).get("spec", {})
707 elif "jobTemplate" in spec:
708 pod_spec = (
709 spec.get("jobTemplate", {}).get("spec", {}).get("template", {}).get("spec", {})
710 )
711 elif "containers" in spec:
712 pod_spec = spec
713 pod_specs = [pod_spec] if pod_spec else []
715 for pod_spec in pod_specs:
716 # --- Pod-level checks ---
717 if flags.get("block_host_network") and pod_spec.get("hostNetwork", False):
718 return False, "hostNetwork is not permitted"
720 if flags.get("block_host_pid") and pod_spec.get("hostPID", False):
721 return False, "hostPID is not permitted"
723 if flags.get("block_host_ipc") and pod_spec.get("hostIPC", False):
724 return False, "hostIPC is not permitted"
726 if flags.get("block_host_path"):
727 for volume in pod_spec.get("volumes", []):
728 if volume.get("hostPath") is not None:
729 return False, "hostPath volumes are not permitted"
731 security_context = pod_spec.get("securityContext", {})
732 if flags.get("block_privileged") and security_context.get("privileged", False):
733 return False, "privileged pod security context is not permitted"
735 if flags.get("block_run_as_root"):
736 run_as_user = security_context.get("runAsUser")
737 if run_as_user is not None and run_as_user == 0:
738 return False, "running as root (runAsUser: 0) is not permitted"
740 # --- Container-level checks ---
741 for container_type, container in _iter_all_containers(pod_spec):
742 container_name = container.get("name", "unknown")
743 container_security = container.get("securityContext", {})
744 if flags.get("block_privileged") and container_security.get("privileged", False):
745 return (
746 False,
747 f"{container_type} '{container_name}': privileged containers are not permitted",
748 )
749 if flags.get("block_privilege_escalation") and container_security.get(
750 "allowPrivilegeEscalation", False
751 ):
752 return (
753 False,
754 f"{container_type} '{container_name}': allowPrivilegeEscalation is not permitted",
755 )
757 if flags.get("block_added_capabilities"):
758 added_caps = container_security.get("capabilities", {}).get("add", [])
759 if added_caps:
760 return (
761 False,
762 f"{container_type} '{container_name}': added capabilities are not permitted",
763 )
765 if flags.get("block_run_as_root"):
766 run_as_user = container_security.get("runAsUser")
767 if run_as_user is not None and run_as_user == 0:
768 return (
769 False,
770 f"{container_type} '{container_name}': running as root (runAsUser: 0) is not permitted",
771 )
773 return True, None
775 except Exception as e:
776 logger.error(f"Error validating security context: {e}")
777 return False, f"Security context error: {e}"
780def check_tolerations(
781 manifest: dict[str, Any], policy: JobValidationPolicy | None = None
782) -> tuple[bool, str | None]:
783 """Require accelerator jobs to carry a matching node toleration.
785 GCO nodepools taint accelerator nodes with ``nvidia.com/gpu``,
786 ``aws.amazon.com/neuron``, and ``vpc.amazonaws.com/efa`` (NoSchedule). A pod
787 requesting one of these resources but lacking a matching toleration would
788 stay Pending forever, so we reject it at admission with an actionable
789 message instead.
791 For a TrainJob the accelerator request usually lives in
792 ``spec.trainer.resourcesPerNode`` while the toleration can only be expressed
793 through a ``runtimePatches`` pod spec, so the requested set and the
794 tolerations are each unioned across every pod-spec view before matching.
796 *policy* is accepted for signature symmetry with the other checks and is
797 unused: which taints exist is a property of how GCO builds nodepools, not of
798 per-deployment configuration. Whether this check runs at all is the caller's
799 decision, gated on ``require_accelerator_toleration``.
801 Returns ``(is_valid, error_message)``; the message is ``None`` when valid.
802 """
803 del policy # see docstring
804 if manifest.get("kind") == "TrainJob":
805 pod_specs = trainjob_validation_pod_specs(manifest)
806 hint_example = "examples/kubeflow-trainjob.yaml (GPU variant, via runtimePatches)"
807 else:
808 single = extract_pod_spec(manifest)
809 pod_specs = [single] if single else []
810 hint_example = "examples/gpu-job.yaml"
811 if not pod_specs:
812 return True, None
814 requested: set[str] = set()
815 tolerations: list[dict[str, Any]] = []
816 for pod_spec in pod_specs:
817 requested.update(requested_accelerators(pod_spec))
818 tolerations.extend(pod_spec.get("tolerations", []) or [])
819 if not requested:
820 return True, None
822 for taint in requested:
823 if not _toleration_matches(tolerations, taint):
824 hint = (
825 f"add a matching toleration (e.g. key '{taint}', operator "
826 f"'Exists', effect 'NoSchedule'); see {hint_example}"
827 )
828 return (
829 False,
830 f"Job requests accelerator '{taint}' but no matching "
831 f"toleration for taint {taint}=true:NoSchedule was found. {hint}",
832 )
833 return True, None