Coverage for scripts / example_job_validation / actions.py: 100.00%
138 statements
« prev ^ index » next coverage.py v7.13.5, created at 2026-09-14 22:07 +0000
1"""The ``examples`` action: run every selected example through its documented path."""
3from __future__ import annotations
5import contextlib
6import time
7from concurrent.futures import ThreadPoolExecutor
8from pathlib import Path
9from typing import Any
11import yaml
13from scripts.live_release_validation.models import RunContext
15from . import drivers, kube
16from .drivers import ExampleRunResult, ExampleValidationError
17from .specs import COMPANION, EXAMPLE_SPECS, KUBECTL_APPLY, SCALEDJOB_SCALES
18from .static_checks import parse_example, run_static_checks
21def action_static(ctx: RunContext) -> dict[str, Any]:
22 """Offline checks for the selected examples (also run standalone in CI)."""
23 names = list(getattr(ctx.settings, "selected_examples", ()) or [])
24 findings = run_static_checks(ctx.settings.repo_root, names or None)
25 failed = [finding for finding in findings if not finding.passed]
26 details = {
27 "checked": len(findings),
28 "failed": [
29 {"example": item.example, "check": item.check, "detail": item.detail} for item in failed
30 ],
31 }
32 if failed:
33 raise RuntimeError(f"{len(failed)} static example check(s) failed: {details['failed']}")
34 return details
37def _capacity_skip_reason(ctx: RunContext, region: str, quota_code: str) -> str | None:
38 """Return a skip reason when the account has zero quota for the family."""
39 if not quota_code:
40 return None
41 with drivers.BOTO_CLIENT_LOCK:
42 client = ctx.session.client("service-quotas", region_name=region)
43 try:
44 quota = client.get_service_quota(ServiceCode="ec2", QuotaCode=quota_code)
45 value = float(quota["Quota"]["Value"])
46 except Exception as exc: # noqa: BLE001 — quota lookup failing must not fail the run
47 return f"quota {quota_code} lookup failed ({type(exc).__name__}); treating as unavailable"
48 if value <= 0:
49 name = quota["Quota"].get("QuotaName", quota_code)
50 return f"account quota '{name}' is {value:g} vCPUs — no capacity for this example"
51 return None
54def _keda_operator_role_arn(kubectl: kube.KubectlRunner) -> str:
55 """Resolve the KEDA operator's IAM role from its service-account annotation."""
56 for namespace in ("keda", "gco-system", "kube-system"):
57 code, out, _ = kubectl(
58 "get",
59 "serviceaccount",
60 "keda-operator",
61 "-n",
62 namespace,
63 "-o",
64 "jsonpath={.metadata.annotations.eks\\.amazonaws\\.com/role-arn}",
65 )
66 if code == 0 and out.strip():
67 return out.strip()
68 raise ExampleValidationError(
69 "KEDA operator service-account role annotation not found in keda/gco-system/kube-system"
70 )
73def _prepare_keda_manifest(parsed: Any, queue_url: str, region: str, manifest_path: Path) -> Path:
74 """Substitute the documented placeholder queue URL with the demo queue."""
75 documents = []
76 for doc in yaml.safe_load_all(manifest_path.read_text(encoding="utf-8")):
77 if doc and doc.get("kind") == "ScaledJob":
78 for trigger in doc["spec"]["triggers"]:
79 metadata = trigger.get("metadata", {})
80 if "queueURL" in metadata:
81 metadata["queueURL"] = queue_url
82 metadata["awsRegion"] = region
83 if doc:
84 documents.append(doc)
85 return drivers.write_temp_manifest(documents, f"-{parsed.name}.yaml")
88def _run_one_example(
89 ctx: RunContext,
90 name: str,
91 region: str,
92 kubectl: kube.KubectlRunner,
93) -> ExampleRunResult:
94 spec = EXAMPLE_SPECS[name]
95 parsed = parse_example(ctx.settings.repo_root, name)
96 started = time.monotonic()
98 if spec.submission == COMPANION:
99 return ExampleRunResult(
100 name=name,
101 status="passed",
102 submission=spec.submission,
103 detail=f"companion artifact: {spec.notes}",
104 )
106 skip_reason = _capacity_skip_reason(ctx, region, spec.capacity_quota_code)
107 if skip_reason:
108 return ExampleRunResult(
109 name=name, status="skipped", submission=spec.submission, detail=skip_reason
110 )
112 manifest_path, mutations = drivers.apply_mutations(parsed)
113 evidence: dict[str, Any] = {}
114 keda_queue: drivers.KedaDemoQueue | None = None
115 vector_corpus: drivers.VectorDemoCorpus | None = None
116 try:
117 if spec.setup_driver and spec.setup_driver not in drivers.KNOWN_SETUP_DRIVERS:
118 # Fail closed: a spec naming a driver this dispatcher does not
119 # implement must fail loudly, not run without its precondition
120 # and report an unearned pass.
121 raise ExampleValidationError(
122 f"setup driver {spec.setup_driver!r} is not implemented in actions._run_one_example"
123 )
124 if spec.setup_driver == "keda-demo-queue":
125 role_arn = _keda_operator_role_arn(kubectl)
126 keda_queue = drivers.KedaDemoQueue(
127 session=ctx.session, region=region, run_id=ctx.settings.run_id
128 )
129 evidence["setup"] = keda_queue.create(role_arn)
130 manifest_path = _prepare_keda_manifest(
131 parsed, keda_queue.queue_url, region, manifest_path
132 )
133 mutations["ScaledJob.triggers.queueURL"] = "disposable demo queue for this run"
134 elif spec.setup_driver == "vector-demo-corpus":
135 # The example's documented prerequisite, run verbatim and fully
136 # reverted in the finally below (S3 objects + chunk items).
137 vector_corpus = drivers.VectorDemoCorpus(
138 repo_root=ctx.settings.repo_root, session=ctx.session, region=region
139 )
140 evidence["setup"] = vector_corpus.create()
141 elif spec.setup_driver == "trainer-runtime-ready":
142 # Readiness wait on deploy-time artifacts; nothing to revert.
143 evidence["setup"] = drivers.wait_trainer_runtime_ready(kubectl)
144 elif spec.setup_driver == "mlflow-ready":
145 # Readiness wait; the tracking server may still be rolling out
146 # right after a fresh install (its PVC lands one applier pass
147 # after the chart). Nothing to revert.
148 evidence["setup"] = drivers.wait_mlflow_ready(kubectl)
150 evidence["submission"] = drivers.submit_example(
151 parsed, manifest_path, repo_root=ctx.settings.repo_root, region=region, kubectl=kubectl
152 )
153 if spec.criteria in drivers.CRITERIA_WAITERS:
154 evidence["criteria"] = drivers.CRITERIA_WAITERS[spec.criteria](
155 parsed, kubectl, timeout=spec.timeout_seconds
156 )
157 if spec.criteria == SCALEDJOB_SCALES or spec.submission == KUBECTL_APPLY:
158 evidence["cleanup"] = drivers.cleanup_example(parsed, manifest_path, kubectl)
159 else:
160 # CLI-submitted resources: delete through kubectl as well so quota
161 # headroom is restored for the next example.
162 evidence["cleanup"] = drivers.cleanup_example(parsed, manifest_path, kubectl)
163 return ExampleRunResult(
164 name=name,
165 status="passed",
166 submission=spec.submission,
167 duration_seconds=time.monotonic() - started,
168 mutations=mutations,
169 evidence=evidence,
170 )
171 except ExampleValidationError as exc:
172 with contextlib.suppress(ExampleValidationError):
173 drivers.cleanup_example(parsed, manifest_path, kubectl)
174 return ExampleRunResult(
175 name=name,
176 status="failed",
177 submission=spec.submission,
178 duration_seconds=time.monotonic() - started,
179 detail=str(exc)[:1500],
180 mutations=mutations,
181 evidence=evidence,
182 )
183 finally:
184 if keda_queue is not None:
185 keda_queue.destroy()
186 if vector_corpus is not None:
187 vector_corpus.destroy()
190def action_examples(ctx: RunContext) -> dict[str, Any]:
191 """Run the selected examples in parallel inside one cluster session.
193 Every example is self-contained (own workload names, own temp manifest,
194 own cleanup), so all selected examples are submitted at once and each
195 thread drives its example's full documented flow: submit, wait on the
196 success criteria, clean up. GPU node provisioning and image pulls — the
197 dominant wall-clock costs — overlap instead of serializing. Transient
198 ``exceeded quota`` admission rejections while peers hold the namespace
199 quota are expected and retried by the Job controller (the fail-fast in
200 ``drivers`` deliberately exempts them). ``max_parallel_examples``
201 throttles the pool; 0 means all selected examples at once.
202 """
203 selected = list(getattr(ctx.settings, "selected_examples", ()) or EXAMPLE_SPECS)
204 region = ctx.deployment_regions[0]
205 cluster_name = f"{ctx.config.project_name}-{region}"
206 state: dict[str, Any] = ctx.checkpoint.state.setdefault("examples", {})
208 results: dict[str, ExampleRunResult] = {}
209 pending: list[str] = []
210 for name in selected:
211 previous = state.get(name)
212 if isinstance(previous, dict) and previous.get("status") == "passed":
213 results[name] = ExampleRunResult(
214 name=name,
215 status="passed",
216 submission=str(previous.get("submission", "")),
217 detail="checkpoint: already passed in this run",
218 )
219 else:
220 pending.append(name)
222 def run_example(name: str, kubectl: kube.KubectlRunner) -> None:
223 print(f"[example] {name} ({EXAMPLE_SPECS[name].submission}) started")
224 result = _run_one_example(ctx, name, region, kubectl)
225 with ctx.state_lock:
226 results[name] = result
227 state[name] = result.to_dict()
228 ctx.persist()
229 print(f"[example] {name}: {result.status} ({result.duration_seconds:.1f}s)")
231 limit = int(getattr(ctx.settings, "max_parallel_examples", 0) or 0)
232 workers = min(len(pending), limit) if limit > 0 else len(pending)
233 if pending:
234 with kube.cluster_session(ctx.settings.repo_root, cluster_name, region) as kubectl:
235 if workers == 1:
236 for name in pending:
237 run_example(name, kubectl)
238 else:
239 with ThreadPoolExecutor(max_workers=workers, thread_name_prefix="example") as pool:
240 futures = {pool.submit(run_example, name, kubectl): name for name in pending}
241 for future, name in futures.items():
242 # Surface unexpected (non-validation) errors with the
243 # example's name; ExampleValidationError is already
244 # converted to a failed result inside the thread.
245 try:
246 future.result()
247 except Exception as exc:
248 raise RuntimeError(f"example {name} crashed: {exc}") from exc
250 ordered = [results[name] for name in selected if name in results]
251 summary = {
252 "region": region,
253 "results": [result.to_dict() for result in ordered],
254 "max_parallel": workers,
255 "passed": sum(1 for item in ordered if item.status == "passed"),
256 "skipped": sum(1 for item in ordered if item.status == "skipped"),
257 "failed": sum(1 for item in ordered if item.status == "failed"),
258 }
259 ctx.checkpoint.state["examples_summary"] = summary
260 ctx.persist()
261 if summary["failed"]:
262 failed_names = [item.name for item in ordered if item.status == "failed"]
263 raise RuntimeError(
264 f"{summary['failed']} example(s) failed: {', '.join(failed_names)} "
265 "(per-example evidence is in the report details)"
266 )
267 return summary