Coverage for cli / job_policy.py: 100.00%
167 statements
« prev ^ index » next coverage.py v7.13.5, created at 2026-09-14 22:07 +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 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: list[dict[str, Any]], policy: JobValidationPolicy
275) -> list[AdmissionIssue]:
276 """Evaluate every manifest against one policy."""
277 issues: list[AdmissionIssue] = []
278 with _quiet_admission_logging():
279 for manifest in manifests:
280 if not isinstance(manifest, dict):
281 continue
282 issues.extend(evaluate_manifest(manifest, policy))
283 return issues
286@dataclass(frozen=True)
287class RegionVerdict:
288 """Whether one region would admit the manifests, and why not."""
290 region: str
291 verdict: str
292 issues: list[AdmissionIssue] = field(default_factory=list)
293 reason: str | None = None
294 enforcement_gaps: list[str] = field(default_factory=list)
297def region_verdicts(
298 manifests: list[dict[str, Any]], policies: list[RegionPolicy]
299) -> list[RegionVerdict]:
300 """Judge *manifests* against each region's policy."""
301 verdicts: list[RegionVerdict] = []
302 for entry in policies:
303 if not entry.ok:
304 verdicts.append(
305 RegionVerdict(region=entry.region, verdict=VERDICT_UNKNOWN, reason=entry.reason)
306 )
307 continue
308 assert entry.policy is not None
309 issues = evaluate_manifests(manifests, entry.policy)
310 verdicts.append(
311 RegionVerdict(
312 region=entry.region,
313 verdict=VERDICT_REJECT if issues else VERDICT_ADMIT,
314 issues=issues,
315 enforcement_gaps=entry.enforcement_gaps,
316 )
317 )
318 return verdicts
321# ---------------------------------------------------------------------------
322# Cross-region drift
323# ---------------------------------------------------------------------------
325#: Policy fields compared across regions. ``trusted_registries`` is handled
326#: separately because CDK legitimately varies it per deployment.
327_DRIFT_FIELDS: tuple[str, ...] = (
328 "max_cpu_millicores",
329 "max_memory_bytes",
330 "max_gpu_count",
331 "allowed_namespaces",
332 "allowed_kinds",
333 "trusted_dockerhub_orgs",
334 "require_accelerator_toleration",
335 "validation_enabled",
336 "yaml_max_depth",
337 "security",
338)
341@dataclass(frozen=True)
342class PolicyDrift:
343 """One policy field that is not identical across regions."""
345 field: str
346 values: dict[str, Any]
349def _comparable(value: Any) -> Any:
350 """Normalize a policy value into something hashable and printable."""
351 if isinstance(value, frozenset | set):
352 return tuple(sorted(value))
353 if isinstance(value, dict):
354 return tuple(sorted(value.items()))
355 return value
358def _renderable(value: Any) -> Any:
359 """Turn a normalized value back into something JSON-serializable."""
360 if isinstance(value, tuple):
361 if value and all(isinstance(item, tuple) and len(item) == 2 for item in value):
362 return dict(value)
363 return list(value)
364 return value
367def detect_policy_drift(policies: list[RegionPolicy]) -> list[PolicyDrift]:
368 """Return the policy fields that differ across successfully-read regions.
370 There are no per-region policy overrides -- every region is deployed from
371 the same ``cdk.json`` -- so a field that differs means at least one region
372 was deployed from a different checkout of it. That is worth surfacing
373 because it is otherwise invisible until a manifest that was admitted in one
374 region is rejected in another.
376 ``trusted_registries`` is deliberately excluded from the field list and
377 handled by :func:`registry_drift`, since CDK appends the project's own ECR
378 hostnames at synth time and those legitimately differ.
379 """
380 readable = [entry for entry in policies if entry.ok]
381 if len(readable) < 2:
382 return []
384 drifts: list[PolicyDrift] = []
385 for name in _DRIFT_FIELDS:
386 observed = {entry.region: _comparable(getattr(entry.policy, name)) for entry in readable}
387 if len(set(observed.values())) > 1:
388 drifts.append(
389 PolicyDrift(
390 field=name,
391 values={region: _renderable(value) for region, value in observed.items()},
392 )
393 )
394 return drifts
397def registry_drift(policies: list[RegionPolicy]) -> PolicyDrift | None:
398 """Compare ``trusted_registries`` with the project's ECR hostnames removed.
400 CDK augments the configured allowlist with the project's own ECR registry
401 hostnames, which encode a region, so a raw comparison reports drift on every
402 multi-region deployment. Stripping anything that looks like an ECR hostname
403 leaves the part that came from ``cdk.json``, where a difference is real
404 drift. The stripped entries are not lost -- they are what
405 :func:`ecr_augmentation` reports.
406 """
407 readable = [entry for entry in policies if entry.ok]
408 if len(readable) < 2:
409 return None
411 observed = {
412 entry.region: tuple(
413 sorted(
414 host
415 for host in entry.policy.trusted_registries # type: ignore[union-attr]
416 if not _ECR_HOSTNAME.match(host)
417 )
418 )
419 for entry in readable
420 }
421 if len(set(observed.values())) <= 1:
422 return None
423 return PolicyDrift(
424 field="trusted_registries",
425 values={region: list(value) for region, value in observed.items()},
426 )
429def ecr_augmentation(policies: list[RegionPolicy]) -> dict[str, list[str]]:
430 """Report the ECR hostnames CDK added to each region's allowlist.
432 Informational, not drift: these appear in no ``cdk.json``, which is the
433 concrete reason a locally-computed policy is not authoritative.
434 """
435 return {
436 entry.region: sorted(
437 host
438 for host in entry.policy.trusted_registries # type: ignore[union-attr]
439 if _ECR_HOSTNAME.match(host)
440 )
441 for entry in policies
442 if entry.ok
443 }