Coverage for cli / job_policy.py: 100.00%
167 statements
« prev ^ index » next coverage.py v7.13.5, created at 2026-10-28 16:19 +0000
1"""Read deployed job-validation policy and judge manifests against it.
3``GET /api/v1/policy`` reports what a region enforces. This module turns that
4into answers to questions the submission path cannot answer on its own, because
5each needs more than one region or needs an answer before anything is submitted:
7 * *which regions would admit this job* -- fan the policy read out and evaluate
8 the same manifest against each. A 32-GPU job is admissible in one region and
9 over-cap in another, and today you find that out by submitting.
10 * *do the regions still agree* -- there are no per-region policy overrides, so
11 any field that differs across regions is a region deployed from a different
12 checkout. That is invisible until a job that worked yesterday is rejected.
13 * *will this be admitted here* -- an advisory pre-submit check, so a rejection
14 costs a local round trip instead of a queue round trip.
16Everything is advisory. The authoritative gate is the cluster, and this reads a
17snapshot over a network; a check that blocks on its own opinion would refuse
18valid jobs whenever it is wrong or merely stale. So the callers render findings
19and exit 0 unless the user opts into a failing exit code.
21The checks themselves are not reimplemented here -- they come from
22:mod:`gco.job_admission`, which is the same code the manifest processor runs.
23"""
25from __future__ import annotations
27import concurrent.futures
28import contextlib
29import logging
30import re
31from collections.abc import Iterable, Iterator
32from dataclasses import dataclass, field
33from typing import Any
35from gco.job_admission import (
36 JobValidationPolicy,
37 check_resource_caps,
38 check_security_context,
39 check_tolerations,
40 validate_image_sources,
41 validate_resource_kind,
42)
44logger = logging.getLogger(__name__)
46#: Per-region fetch outcomes.
47FETCH_OK = "ok"
48FETCH_UNREACHABLE = "unreachable"
49FETCH_ERROR = "error"
51#: Admissibility verdicts. ``unknown`` is distinct from ``reject`` on purpose:
52#: a region whose policy could not be read has not refused anything, and
53#: collapsing the two would report a network failure as a policy violation.
54VERDICT_ADMIT = "admit"
55VERDICT_REJECT = "reject"
56VERDICT_UNKNOWN = "unknown"
58#: Checks are named so a caller can tell which layer objected. These match the
59#: order the manifest processor applies them in.
60CHECK_KIND = "kind"
61CHECK_NAMESPACE = "namespace"
62CHECK_IMAGES = "images"
63CHECK_CAPS = "resource_caps"
64CHECK_SECURITY = "security_context"
65CHECK_TOLERATIONS = "tolerations"
67#: An AWS ECR registry hostname, e.g. 123456789012.dkr.ecr.us-east-2.amazonaws.com
68#: (also matching the .cn and ISO partition suffixes).
69_ECR_HOSTNAME = re.compile(r"^\d{12}\.dkr\.ecr\.[a-z0-9-]+\.amazonaws\.com(?:\.cn)?$")
71# How long to wait on one region's policy read before giving up on it. A fan-out
72# has to bound the slowest region or one unreachable region stalls the whole
73# answer; the per-region timeout inside the HTTP client is 30s, and this leaves
74# room for its retries without letting a hung region hold everything.
75FETCH_TIMEOUT_SECONDS = 45
78# ---------------------------------------------------------------------------
79# Fetching
80# ---------------------------------------------------------------------------
83@dataclass(frozen=True)
84class RegionPolicy:
85 """One region's policy read, successful or not."""
87 region: str
88 status: str
89 policy: JobValidationPolicy | None = None
90 document: dict[str, Any] = field(default_factory=dict)
91 cluster_enforcement: dict[str, Any] = field(default_factory=dict)
92 reason: str | None = None
94 @property
95 def ok(self) -> bool:
96 return self.status == FETCH_OK and self.policy is not None
98 @property
99 def enforcement_gaps(self) -> list[str]:
100 """Namespaces whose live ResourceQuota / LimitRange could not be read.
102 A gap here means the answer covers only the front-door caps: the
103 manifest can clear those and still be rejected at pod creation. Callers
104 surface this rather than letting an ``admit`` verdict imply more
105 confidence than it has.
106 """
107 return sorted(
108 namespace
109 for namespace, layer in (self.cluster_enforcement or {}).items()
110 if isinstance(layer, dict) and layer.get("status") != "ok"
111 )
114def fetch_region_policy(aws_client: Any, region: str) -> RegionPolicy:
115 """Read one region's policy, converting any failure into a status."""
116 try:
117 document = aws_client.get_job_validation_policy(region=region)
118 except Exception as e:
119 # A region with no regional API bridge raises RuntimeError naming that;
120 # anything else is genuinely unexpected. Both are non-fatal here.
121 message = str(e)
122 status = FETCH_UNREACHABLE if "not deployed" in message else FETCH_ERROR
123 logger.debug("Policy read failed for %s: %s", region, e)
124 return RegionPolicy(region=region, status=status, reason=f"{type(e).__name__}: {e}")
126 policy_document = document.get("policy", {}) or {}
127 if not policy_document:
128 return RegionPolicy(
129 region=region,
130 status=FETCH_ERROR,
131 document=document,
132 reason="the response carried no policy object",
133 )
135 return RegionPolicy(
136 region=region,
137 status=FETCH_OK,
138 policy=JobValidationPolicy.from_policy_document(policy_document),
139 document=policy_document,
140 cluster_enforcement=document.get("cluster_enforcement", {}) or {},
141 )
144def fetch_region_policies(aws_client: Any, regions: list[str]) -> list[RegionPolicy]:
145 """Read every region's policy concurrently, preserving *regions* order.
147 Concurrent because each read costs a CloudFormation describe plus an API
148 call and they are independent; serial fan-out over a handful of regions is
149 slow enough that people stop running the check.
150 """
151 if not regions:
152 return []
153 if len(regions) == 1:
154 return [fetch_region_policy(aws_client, regions[0])]
156 results: dict[str, RegionPolicy] = {}
157 with concurrent.futures.ThreadPoolExecutor(max_workers=min(len(regions), 8)) as pool:
158 futures = {
159 pool.submit(fetch_region_policy, aws_client, region): region for region in regions
160 }
161 for future in concurrent.futures.as_completed(futures, timeout=None):
162 region = futures[future]
163 try:
164 results[region] = future.result(timeout=FETCH_TIMEOUT_SECONDS)
165 except Exception as e: # pragma: no cover - defensive
166 results[region] = RegionPolicy(
167 region=region, status=FETCH_ERROR, reason=f"{type(e).__name__}: {e}"
168 )
169 return [
170 results.get(region, RegionPolicy(region=region, status=FETCH_ERROR, reason="no result"))
171 for region in regions
172 ]
175# ---------------------------------------------------------------------------
176# Judging a manifest
177# ---------------------------------------------------------------------------
180@dataclass(frozen=True)
181class AdmissionIssue:
182 """One reason a manifest would be rejected."""
184 check: str
185 message: str
186 manifest: str | None = None
189def manifest_label(manifest: dict[str, Any]) -> str:
190 """A short ``Kind/name`` label for messages."""
191 kind = manifest.get("kind") or "?"
192 name = (manifest.get("metadata") or {}).get("name") or "<unnamed>"
193 return f"{kind}/{name}"
196def evaluate_manifest(
197 manifest: dict[str, Any], policy: JobValidationPolicy
198) -> list[AdmissionIssue]:
199 """Return every reason *policy* would reject *manifest*.
201 Every check runs -- the cluster short-circuits on the first failure, but a
202 caller fixing a manifest wants the whole list rather than one round trip per
203 problem. The trade-off is that a manifest failing the kind check also
204 reports whatever else is wrong with it, which is more information than the
205 server would give.
206 """
207 label = manifest_label(manifest)
208 issues: list[AdmissionIssue] = []
210 if not policy.validation_enabled:
211 return issues
213 ok, message = validate_resource_kind(manifest, policy.allowed_kinds)
214 if not ok:
215 issues.append(AdmissionIssue(CHECK_KIND, message or "kind is not allowed", label))
217 namespace = (manifest.get("metadata") or {}).get("namespace")
218 if namespace and namespace not in policy.allowed_namespaces:
219 issues.append(
220 AdmissionIssue(
221 CHECK_NAMESPACE,
222 f"namespace '{namespace}' is not in the allowlist "
223 f"({', '.join(sorted(policy.allowed_namespaces))})",
224 label,
225 )
226 )
228 ok, message = validate_image_sources(
229 manifest,
230 trusted_registries=list(policy.trusted_registries),
231 trusted_dockerhub_orgs=list(policy.trusted_dockerhub_orgs),
232 )
233 if not ok:
234 issues.append(AdmissionIssue(CHECK_IMAGES, message or "untrusted image source", label))
236 ok, caps_message = check_resource_caps(manifest, policy)
237 if not ok:
238 issues.append(AdmissionIssue(CHECK_CAPS, caps_message, label))
240 ok, message = check_security_context(manifest, policy)
241 if not ok:
242 issues.append(AdmissionIssue(CHECK_SECURITY, message or "security policy violation", label))
244 if policy.require_accelerator_toleration:
245 ok, message = check_tolerations(manifest)
246 if not ok:
247 issues.append(AdmissionIssue(CHECK_TOLERATIONS, message or "missing toleration", label))
249 return issues
252@contextlib.contextmanager
253def _quiet_admission_logging() -> Iterator[None]:
254 """Silence gco.job_admission's warnings for the duration of a check.
256 Those warnings are an audit trail in the service -- a rejected submission
257 should leave a record of why. In a CLI pre-check they are noise: the caller
258 is about to be shown the same information as formatted findings, so the log
259 line duplicates it and interleaves with the report.
261 Single-threaded by construction: evaluation runs after the concurrent policy
262 fetch has joined, so this never races another region's evaluation.
263 """
264 admission_logger = logging.getLogger("gco.job_admission")
265 previous = admission_logger.level
266 admission_logger.setLevel(logging.ERROR)
267 try:
268 yield
269 finally:
270 admission_logger.setLevel(previous)
273def evaluate_manifests(
274 manifests: Iterable[object], policy: JobValidationPolicy
275) -> list[AdmissionIssue]:
276 """Evaluate every manifest against one policy.
278 ``manifests`` is whatever the YAML loader produced. A scalar or list
279 document mixed into a directory scan is not a manifest and is skipped
280 rather than crashing the batch; the parameter type says as much.
281 """
282 issues: list[AdmissionIssue] = []
283 with _quiet_admission_logging():
284 for manifest in manifests:
285 if not isinstance(manifest, dict):
286 continue
287 issues.extend(evaluate_manifest(manifest, policy))
288 return issues
291@dataclass(frozen=True)
292class RegionVerdict:
293 """Whether one region would admit the manifests, and why not."""
295 region: str
296 verdict: str
297 issues: list[AdmissionIssue] = field(default_factory=list)
298 reason: str | None = None
299 enforcement_gaps: list[str] = field(default_factory=list)
302def region_verdicts(
303 manifests: list[dict[str, Any]], policies: list[RegionPolicy]
304) -> list[RegionVerdict]:
305 """Judge *manifests* against each region's policy."""
306 verdicts: list[RegionVerdict] = []
307 for entry in policies:
308 if not entry.ok:
309 verdicts.append(
310 RegionVerdict(region=entry.region, verdict=VERDICT_UNKNOWN, reason=entry.reason)
311 )
312 continue
313 assert entry.policy is not None
314 issues = evaluate_manifests(manifests, entry.policy)
315 verdicts.append(
316 RegionVerdict(
317 region=entry.region,
318 verdict=VERDICT_REJECT if issues else VERDICT_ADMIT,
319 issues=issues,
320 enforcement_gaps=entry.enforcement_gaps,
321 )
322 )
323 return verdicts
326# ---------------------------------------------------------------------------
327# Cross-region drift
328# ---------------------------------------------------------------------------
330#: Policy fields compared across regions. ``trusted_registries`` is handled
331#: separately because CDK legitimately varies it per deployment.
332_DRIFT_FIELDS: tuple[str, ...] = (
333 "max_cpu_millicores",
334 "max_memory_bytes",
335 "max_gpu_count",
336 "allowed_namespaces",
337 "allowed_kinds",
338 "trusted_dockerhub_orgs",
339 "require_accelerator_toleration",
340 "validation_enabled",
341 "yaml_max_depth",
342 "security",
343)
346@dataclass(frozen=True)
347class PolicyDrift:
348 """One policy field that is not identical across regions."""
350 field: str
351 values: dict[str, Any]
354def _comparable(value: Any) -> Any:
355 """Normalize a policy value into something hashable and printable."""
356 if isinstance(value, frozenset | set):
357 return tuple(sorted(value))
358 if isinstance(value, dict):
359 return tuple(sorted(value.items()))
360 return value
363def _renderable(value: Any) -> Any:
364 """Turn a normalized value back into something JSON-serializable."""
365 if isinstance(value, tuple):
366 if value and all(isinstance(item, tuple) and len(item) == 2 for item in value):
367 return dict(value)
368 return list(value)
369 return value
372def detect_policy_drift(policies: list[RegionPolicy]) -> list[PolicyDrift]:
373 """Return the policy fields that differ across successfully-read regions.
375 There are no per-region policy overrides -- every region is deployed from
376 the same ``cdk.json`` -- so a field that differs means at least one region
377 was deployed from a different checkout of it. That is worth surfacing
378 because it is otherwise invisible until a manifest that was admitted in one
379 region is rejected in another.
381 ``trusted_registries`` is deliberately excluded from the field list and
382 handled by :func:`registry_drift`, since CDK appends the project's own ECR
383 hostnames at synth time and those legitimately differ.
384 """
385 readable = [entry for entry in policies if entry.ok]
386 if len(readable) < 2:
387 return []
389 drifts: list[PolicyDrift] = []
390 for name in _DRIFT_FIELDS:
391 observed = {entry.region: _comparable(getattr(entry.policy, name)) for entry in readable}
392 if len(set(observed.values())) > 1:
393 drifts.append(
394 PolicyDrift(
395 field=name,
396 values={region: _renderable(value) for region, value in observed.items()},
397 )
398 )
399 return drifts
402def registry_drift(policies: list[RegionPolicy]) -> PolicyDrift | None:
403 """Compare ``trusted_registries`` with the project's ECR hostnames removed.
405 CDK augments the configured allowlist with the project's own ECR registry
406 hostnames, which encode a region, so a raw comparison reports drift on every
407 multi-region deployment. Stripping anything that looks like an ECR hostname
408 leaves the part that came from ``cdk.json``, where a difference is real
409 drift. The stripped entries are not lost -- they are what
410 :func:`ecr_augmentation` reports.
411 """
412 readable = [entry for entry in policies if entry.ok]
413 if len(readable) < 2:
414 return None
416 observed = {
417 entry.region: tuple(
418 sorted(
419 host
420 for host in entry.policy.trusted_registries # type: ignore[union-attr]
421 if not _ECR_HOSTNAME.match(host)
422 )
423 )
424 for entry in readable
425 }
426 if len(set(observed.values())) <= 1:
427 return None
428 return PolicyDrift(
429 field="trusted_registries",
430 values={region: list(value) for region, value in observed.items()},
431 )
434def ecr_augmentation(policies: list[RegionPolicy]) -> dict[str, list[str]]:
435 """Report the ECR hostnames CDK added to each region's allowlist.
437 Informational, not drift: these appear in no ``cdk.json``, which is the
438 concrete reason a locally-computed policy is not authoritative.
439 """
440 return {
441 entry.region: sorted(
442 host
443 for host in entry.policy.trusted_registries # type: ignore[union-attr]
444 if _ECR_HOSTNAME.match(host)
445 )
446 for entry in policies
447 if entry.ok
448 }