Coverage for scripts / example_job_validation / actions.py: 100.00%
182 statements
« prev ^ index » next coverage.py v7.13.5, created at 2026-10-28 16:19 +0000
1"""The ``examples`` action: run every selected example through its documented path."""
3from __future__ import annotations
5import contextlib
6import math
7import time
8from concurrent.futures import ThreadPoolExecutor
9from pathlib import Path
10from typing import Any
12import yaml
14from scripts.live_release_validation.models import ActionFailure, RunContext
16from . import drivers, kube
17from .drivers import ExampleRunResult, ExampleValidationError
18from .specs import COMPANION, EXAMPLE_SPECS, KUBECTL_APPLY, SCALEDJOB_SCALES
19from .static_checks import examples_dir, parse_example, run_static_checks
22def action_static(ctx: RunContext) -> dict[str, Any]:
23 """Offline checks for the selected examples (also run standalone in CI)."""
24 names = list(getattr(ctx.settings, "selected_examples", ()) or [])
25 findings = run_static_checks(ctx.settings.repo_root, names or None)
26 failed = [finding for finding in findings if not finding.passed]
27 details = {
28 "checked": len(findings),
29 "failed": [
30 {"example": item.example, "check": item.check, "detail": item.detail} for item in failed
31 ],
32 }
33 if failed:
34 raise ActionFailure(
35 f"{len(failed)} static example check(s) failed: {details['failed']}", details
36 )
37 return details
40def _capacity_skip_reason(ctx: RunContext, region: str, quota_code: str) -> str | None:
41 """Return a skip reason when the account has zero quota for the family."""
42 if not quota_code:
43 return None
44 with drivers.BOTO_CLIENT_LOCK:
45 client = ctx.session.client("service-quotas", region_name=region)
46 try:
47 quota = client.get_service_quota(ServiceCode="ec2", QuotaCode=quota_code)
48 value = float(quota["Quota"]["Value"])
49 except Exception as exc: # quota lookup failing must not fail the run
50 return f"quota {quota_code} lookup failed ({type(exc).__name__}); treating as unavailable"
51 if value <= 0:
52 name = quota["Quota"].get("QuotaName", quota_code)
53 return f"account quota '{name}' is {value:g} vCPUs — no capacity for this example"
54 return None
57def _keda_operator_role_arn(kubectl: kube.KubectlRunner) -> str:
58 """Resolve the KEDA operator's IAM role from its service-account annotation."""
59 for namespace in ("keda", "gco-system", "kube-system"):
60 code, out, _ = kubectl(
61 "get",
62 "serviceaccount",
63 "keda-operator",
64 "-n",
65 namespace,
66 "-o",
67 "jsonpath={.metadata.annotations.eks\\.amazonaws\\.com/role-arn}",
68 )
69 if code == 0 and out.strip():
70 return out.strip()
71 raise ExampleValidationError(
72 "KEDA operator service-account role annotation not found in keda/gco-system/kube-system"
73 )
76def _prepare_keda_manifest(parsed: Any, queue_url: str, region: str, manifest_path: Path) -> Path:
77 """Substitute the documented placeholder queue URL with the demo queue."""
78 documents = []
79 for doc in yaml.safe_load_all(manifest_path.read_text(encoding="utf-8")):
80 if doc and doc.get("kind") == "ScaledJob":
81 for trigger in doc["spec"]["triggers"]:
82 metadata = trigger.get("metadata", {})
83 if "queueURL" in metadata:
84 metadata["queueURL"] = queue_url
85 metadata["awsRegion"] = region
86 if doc:
87 documents.append(doc)
88 return drivers.write_temp_manifest(documents, f"-{parsed.name}.yaml")
91def _pin_argocd_revision(parsed: Any, revision: str, manifest_path: Path) -> Path:
92 """Point every Application at the commit under validation (a disclosed mutation).
94 The shipped example tracks ``main``; the run must sync the Git path at
95 exactly the SHA it validates, so a change to the fixture is tested before
96 it merges.
97 """
98 documents = []
99 for doc in yaml.safe_load_all(manifest_path.read_text(encoding="utf-8")):
100 if doc and doc.get("kind") == "Application":
101 doc["spec"]["source"]["targetRevision"] = revision
102 if doc:
103 documents.append(doc)
104 return drivers.write_temp_manifest(documents, f"-{parsed.name}.yaml")
107def _run_one_example(
108 ctx: RunContext,
109 name: str,
110 region: str,
111 kubectl: kube.KubectlRunner,
112) -> ExampleRunResult:
113 spec = EXAMPLE_SPECS[name]
114 parsed = parse_example(ctx.settings.repo_root, name)
115 started = time.monotonic()
117 if spec.submission == COMPANION:
118 return ExampleRunResult(
119 name=name,
120 status="passed",
121 submission=spec.submission,
122 detail=f"companion artifact: {spec.notes}",
123 )
125 skip_reason = _capacity_skip_reason(ctx, region, spec.capacity_quota_code)
126 if skip_reason:
127 return ExampleRunResult(
128 name=name, status="skipped", submission=spec.submission, detail=skip_reason
129 )
131 try:
132 # A tunnel that stalled between examples is reopened before anything
133 # is submitted; one that cannot be reopened fails the example at once
134 # rather than letting its watchers wait out the whole timeout.
135 kube.ensure_tunnel(kubectl)
136 except kube.TunnelUnavailableError as exc:
137 return ExampleRunResult(
138 name=name,
139 status="failed",
140 submission=spec.submission,
141 duration_seconds=time.monotonic() - started,
142 detail=f"the cluster API was unreachable before the example started: {exc}"[:1500],
143 )
145 manifest_path, mutations = drivers.apply_mutations(parsed)
146 evidence: dict[str, Any] = {}
147 keda_queue: drivers.KedaDemoQueue | None = None
148 vector_corpus: drivers.VectorDemoCorpus | None = None
149 companion: drivers.CompanionApi | None = None
150 ack_queues: drivers.AckSqsQueues | None = None
151 try:
152 if spec.setup_driver and spec.setup_driver not in drivers.KNOWN_SETUP_DRIVERS:
153 # Fail closed: a spec naming a driver this dispatcher does not
154 # implement must fail loudly, not run without its precondition
155 # and report an unearned pass.
156 raise ExampleValidationError(
157 f"setup driver {spec.setup_driver!r} is not implemented in actions._run_one_example"
158 )
159 if spec.setup_driver == "keda-demo-queue":
160 role_arn = _keda_operator_role_arn(kubectl)
161 keda_queue = drivers.KedaDemoQueue(
162 session=ctx.session, region=region, run_id=ctx.settings.run_id
163 )
164 evidence["setup"] = keda_queue.create(role_arn)
165 manifest_path = _prepare_keda_manifest(
166 parsed, keda_queue.queue_url, region, manifest_path
167 )
168 mutations["ScaledJob.triggers.queueURL"] = "disposable demo queue for this run"
169 elif spec.setup_driver == "vector-demo-corpus":
170 # The example's documented prerequisite, run verbatim and fully
171 # reverted in the finally below (S3 objects + chunk items).
172 vector_corpus = drivers.VectorDemoCorpus(
173 repo_root=ctx.settings.repo_root, session=ctx.session, region=region
174 )
175 evidence["setup"] = vector_corpus.create()
176 elif spec.setup_driver == "trainer-runtime-ready":
177 # Readiness wait on deploy-time artifacts; nothing to revert.
178 evidence["setup"] = drivers.wait_trainer_runtime_ready(kubectl)
179 elif spec.setup_driver == "mlflow-ready":
180 # Readiness wait; the tracking server may still be rolling out
181 # right after a fresh install, and the CA bundle the client
182 # mounts is published by trust-manager after the post-Helm
183 # pass. Nothing to revert.
184 evidence["setup"] = drivers.wait_mlflow_ready(kubectl)
185 elif spec.setup_driver == "argocd-revision-pin":
186 # Readiness wait on deploy-time artifacts, then the disclosed
187 # revision pin; the Application itself is the example's object.
188 evidence["setup"] = drivers.wait_argocd_ready(kubectl)
189 manifest_path = _pin_argocd_revision(parsed, ctx.settings.expected_sha, manifest_path)
190 mutations["Application.spec.source.targetRevision"] = (
191 f"{ctx.settings.expected_sha} (the commit under validation)"
192 )
193 elif spec.setup_driver in {"kro-api", "crossplane-api"}:
194 # The companion API definition goes first and is deleted after the
195 # instance's cleanup (below, and in the finally on failure).
196 companion = drivers.CompanionApi(
197 flavor=spec.setup_driver.removesuffix("-api"),
198 path=examples_dir(ctx.settings.repo_root) / f"{spec.companion}.yaml",
199 kubectl=kubectl,
200 )
201 evidence["setup"] = companion.create()
202 elif spec.setup_driver == "ack-sqs":
203 ack_queues = drivers.AckSqsQueues(session=ctx.session, region=region)
204 evidence["setup"] = ack_queues.wait_ready(kubectl)
206 evidence["submission"] = drivers.submit_example(
207 parsed, manifest_path, repo_root=ctx.settings.repo_root, region=region, kubectl=kubectl
208 )
209 if spec.criteria in drivers.CRITERIA_WAITERS:
210 evidence["criteria"] = drivers.CRITERIA_WAITERS[spec.criteria](
211 parsed, kubectl, timeout=spec.timeout_seconds
212 )
213 if ack_queues is not None:
214 evidence["aws"] = ack_queues.verify_created(parsed)
215 if spec.criteria == SCALEDJOB_SCALES or spec.submission == KUBECTL_APPLY:
216 evidence["cleanup"] = drivers.cleanup_example(parsed, manifest_path, kubectl)
217 else:
218 # CLI-submitted resources: delete through kubectl as well so quota
219 # headroom is restored for the next example.
220 evidence["cleanup"] = drivers.cleanup_example(parsed, manifest_path, kubectl)
221 if ack_queues is not None:
222 evidence["aws_cleanup"] = ack_queues.verify_deleted(parsed)
223 if companion is not None:
224 evidence["companion_cleanup"] = companion.destroy()
225 return ExampleRunResult(
226 name=name,
227 status="passed",
228 submission=spec.submission,
229 duration_seconds=time.monotonic() - started,
230 mutations=mutations,
231 evidence=evidence,
232 )
233 except ExampleValidationError as exc:
234 with contextlib.suppress(ExampleValidationError):
235 drivers.cleanup_example(parsed, manifest_path, kubectl)
236 return ExampleRunResult(
237 name=name,
238 status="failed",
239 submission=spec.submission,
240 duration_seconds=time.monotonic() - started,
241 detail=str(exc)[:1500],
242 mutations=mutations,
243 evidence=evidence,
244 )
245 finally:
246 if keda_queue is not None:
247 keda_queue.destroy()
248 if vector_corpus is not None:
249 vector_corpus.destroy()
250 if companion is not None and companion.applied:
251 # Failure path: the instance cleanup above ran first.
252 with contextlib.suppress(ExampleValidationError):
253 companion.destroy()
256#: What one example may spend outside its criteria wait: a setup driver's
257#: readiness wait (up to 600 s), the submission (``gco`` allows 600 s, a DAG
258#: run 1800 s), and cleanup (a 300 s delete plus 180 s for derived Jobs).
259_EXAMPLE_OVERHEAD_SECONDS = 1800
262def _bastion_ttl_minutes(names: list[str], workers: int) -> int:
263 """The tunnel bastion's self-termination backstop, sized to the pending examples.
265 The bastion lives exactly as long as the session in the normal path; the
266 TTL only stops an orphan (a harness killed outright). The bastion's
267 two-hour default is shorter than a sequential pass over the catalog, and
268 a bastion that terminates mid-run takes the tunnel with it, so the TTL is
269 the worst case a thread pool can take — the total over the workers plus
270 the longest single example — within the one day a bastion accepts.
271 """
272 from cli import ephemeral_bastion
274 budgets = [EXAMPLE_SPECS[name].timeout_seconds + _EXAMPLE_OVERHEAD_SECONDS for name in names]
275 minutes = math.ceil((sum(budgets) / workers + max(budgets)) / 60)
276 return max(
277 ephemeral_bastion.DEFAULT_TTL_MINUTES, min(ephemeral_bastion.MAX_TTL_MINUTES, minutes)
278 )
281def action_examples(ctx: RunContext) -> dict[str, Any]:
282 """Run the selected examples in parallel inside one cluster session.
284 Every example is self-contained (own workload names, own temp manifest,
285 own cleanup), so all selected examples are submitted at once and each
286 thread drives its example's full documented flow: submit, wait on the
287 success criteria, clean up. GPU node provisioning and image pulls — the
288 dominant wall-clock costs — overlap instead of serializing. Transient
289 ``exceeded quota`` admission rejections while peers hold the namespace
290 quota are expected and retried by the Job controller (the fail-fast in
291 ``drivers`` deliberately exempts them). ``max_parallel_examples``
292 throttles the pool; 0 means all selected examples at once.
293 """
294 selected = list(getattr(ctx.settings, "selected_examples", ()) or EXAMPLE_SPECS)
295 region = ctx.deployment_regions[0]
296 cluster_name = f"{ctx.config.project_name}-{region}"
297 state: dict[str, Any] = ctx.checkpoint.state.setdefault("examples", {})
299 results: dict[str, ExampleRunResult] = {}
300 pending: list[str] = []
301 for name in selected:
302 previous = state.get(name)
303 if isinstance(previous, dict) and previous.get("status") == "passed":
304 results[name] = ExampleRunResult(
305 name=name,
306 status="passed",
307 submission=str(previous.get("submission", "")),
308 detail="checkpoint: already passed in this run",
309 )
310 else:
311 pending.append(name)
313 def run_example(name: str, kubectl: kube.KubectlRunner) -> None:
314 print(f"[example] {name} ({EXAMPLE_SPECS[name].submission}) started")
315 result = _run_one_example(ctx, name, region, kubectl)
316 with ctx.state_lock:
317 results[name] = result
318 state[name] = result.to_dict()
319 ctx.persist()
320 print(f"[example] {name}: {result.status} ({result.duration_seconds:.1f}s)")
322 limit = int(getattr(ctx.settings, "max_parallel_examples", 0) or 0)
323 workers = min(len(pending), limit) if limit > 0 else len(pending)
324 tunnel: dict[str, Any] | None = None
325 if pending:
326 tunnel = {"bastion_ttl_minutes": _bastion_ttl_minutes(pending, workers), "reopens": []}
327 # The run's own kubeconfig in its private report dir, as in the release
328 # harness. On the default kubeconfig a KUBECONFIG exported by the
329 # operator's shell split the two: `aws eks update-kubeconfig` and
330 # kubectl used that file while the tunnel was pinned in
331 # ~/.kube/config, and two live runs timed out on the private endpoint
332 # before their first example. It also leaves ~/.kube/config alone.
333 with kube.cluster_session(
334 ctx.settings.repo_root,
335 cluster_name,
336 region,
337 kubeconfig_path=ctx.settings.kubeconfig_path,
338 bastion_ttl_minutes=tunnel["bastion_ttl_minutes"],
339 tunnel_events=tunnel["reopens"],
340 ) as kubectl:
341 if workers == 1:
342 for name in pending:
343 run_example(name, kubectl)
344 else:
345 with ThreadPoolExecutor(max_workers=workers, thread_name_prefix="example") as pool:
346 futures = {pool.submit(run_example, name, kubectl): name for name in pending}
347 for future, name in futures.items():
348 # Surface unexpected (non-validation) errors with the
349 # example's name; ExampleValidationError is already
350 # converted to a failed result inside the thread.
351 try:
352 future.result()
353 except Exception as exc:
354 raise RuntimeError(f"example {name} crashed: {exc}") from exc
356 ordered = [results[name] for name in selected if name in results]
357 summary = {
358 "region": region,
359 "results": [result.to_dict() for result in ordered],
360 "max_parallel": workers,
361 "passed": sum(1 for item in ordered if item.status == "passed"),
362 "skipped": sum(1 for item in ordered if item.status == "skipped"),
363 "failed": sum(1 for item in ordered if item.status == "failed"),
364 }
365 if tunnel is not None:
366 # Every reopen the tunnel keeper attempted, so a pass that needed one
367 # says so, and the bastion lifetime the run asked for.
368 summary["tunnel"] = tunnel
369 ctx.checkpoint.state["examples_summary"] = summary
370 ctx.persist()
371 if summary["failed"]:
372 failed_names = [item.name for item in ordered if item.status == "failed"]
373 # The summary rides on the exception so the failed action's report
374 # row carries every example's result, as a passing run's does.
375 raise ActionFailure(
376 f"{summary['failed']} example(s) failed: {', '.join(failed_names)} "
377 "(per-example results are in this action's details in the JSON report)",
378 summary,
379 )
380 return summary