Coverage for scripts / example_job_validation / specs.py: 100.00%
57 statements
« prev ^ index » next coverage.py v7.13.5, created at 2026-10-28 16:19 +0000
1"""The per-example validation spec registry.
3One :class:`ExampleSpec` per file under ``examples/`` (pinned by symmetry
4tests against both the directory and the ``gco_mcp`` ``EXAMPLE_METADATA``
5catalog). Each spec answers, for its example:
7* **how** it is submitted — the exact path the docs tell users to use
8 (``submit-direct`` / ``submit-sqs`` / ``dag run`` / ``kubectl apply`` /
9 companion artifact with no live path);
10* **what** infrastructure it needs beyond the stock deploy (optional helm
11 charts via ``helm_enabled_overrides``, optional features via
12 ``feature_enabled_overrides``, EKS Capabilities via
13 ``eks_capabilities_overrides``, GPU/Neuron capacity, special setup drivers);
14* **when** it counts as passed (workload-specific success criteria); and
15* **which** deliberate mutations the harness applies before submission
16 (e.g. replacing a gated HuggingFace model with an ungated one) — every
17 mutation is disclosed in the report.
19Keep this table boring and declarative: ``drivers.py`` interprets it and
20``actions.py`` orchestrates it.
21"""
23from __future__ import annotations
25from dataclasses import dataclass, field
27#: Documented submission paths the drivers know how to execute.
28SUBMIT_DIRECT = "cli-submit-direct" # gco jobs submit-direct <file> -r <region>
29SUBMIT_SQS = "cli-submit-sqs" # gco jobs submit-sqs <file> --region <region>
30SUBMIT_API = "cli-submit-api" # gco jobs submit <file> --region <region>
31DAG_RUN = "cli-dag-run" # gco dag run <file> -r <region>
32KUBECTL_APPLY = "kubectl-apply" # kubectl apply -f <file> (over the SSM tunnel)
33COMPANION = "companion-artifact" # not independently runnable (data/DAG-step file)
35SUBMISSION_PATHS = (SUBMIT_DIRECT, SUBMIT_SQS, SUBMIT_API, DAG_RUN, KUBECTL_APPLY, COMPANION)
37#: Success-criteria kinds the drivers implement.
38JOB_COMPLETES = "job-completes" # batch/v1 Job reaches Complete
39DEPLOYMENT_AVAILABLE = "deployment-available" # Deployment Available + Service endpoints
40RAYCLUSTER_READY = "raycluster-ready" # RayCluster ready: head + minReplicas workers
41VCJOB_COMPLETES = "vcjob-completes" # batch.volcano.sh Job phase Completed
42SCALEDJOB_SCALES = "scaledjob-scales" # KEDA ScaledJob spawns >=1 Job from queue depth
43TRAINJOB_COMPLETES = "trainjob-completes" # trainer.kubeflow.org TrainJob condition Complete
44DAG_SUCCEEDS = "dag-succeeds" # gco dag run exits 0 with all steps completed
45#: argoproj.io Application Synced + Healthy at the applied revision, and every
46#: Job it syncs from Git Complete.
47ARGOCD_APP_HEALTHY = "argocd-app-healthy"
48#: A kro / Crossplane instance's composed Job (named after the instance, in the
49#: instance's namespace) reaches Complete.
50COMPOSED_JOB_COMPLETES = "composed-job-completes"
51#: Every ACK resource reports ACK.ResourceSynced=True (the AWS resource exists).
52ACK_RESOURCE_SYNCED = "ack-resource-synced"
53NONE = "none" # companion artifacts: static checks only
55#: EKS Capability types an example may require (cdk.json ``eks_capabilities``).
56CAPABILITY_TYPES = ("ack", "kro")
59@dataclass(frozen=True)
60class ExampleSpec:
61 """Declarative validation contract for one example file."""
63 #: File stem under ``examples/`` (e.g. ``simple-job`` for simple-job.yaml).
64 name: str
65 #: One of :data:`SUBMISSION_PATHS` — must match the documented usage.
66 submission: str
67 #: One of the success-criteria kinds above.
68 criteria: str
69 #: Optional helm charts that must be force-enabled for this example
70 #: (threaded to CDK as ``helm_enabled_overrides``).
71 helm_overrides: tuple[str, ...] = ()
72 #: Optional infrastructure features that must be force-enabled
73 #: (threaded to CDK as ``feature_enabled_overrides``).
74 feature_overrides: tuple[str, ...] = ()
75 #: EKS Capabilities (one of :data:`CAPABILITY_TYPES`) that must be enabled
76 #: (threaded to CDK as ``eks_capabilities_overrides``).
77 capability_overrides: tuple[str, ...] = ()
78 #: Managed IAM policies the ACK capability role must hold for this
79 #: example (merged into ``eks_capabilities.ack.iam_policy_arns``); only
80 #: valid together with ``capability_overrides=("ack",)``.
81 ack_iam_policy_arns: tuple[str, ...] = ()
82 #: Companion example (file stem, a :data:`COMPANION` spec) the setup driver
83 #: applies before submission and deletes after cleanup — the API
84 #: definition an instance example needs (kro RGD, Crossplane XRD).
85 companion: str = ""
86 #: Accelerator requirement: "" (none), "nvidia", "neuron", or "efa".
87 accelerator: str = ""
88 #: Named setup/teardown driver hook, implemented in ``drivers.py`` and
89 #: listed in its ``KNOWN_SETUP_DRIVERS`` registry, e.g. "keda-demo-queue".
90 setup_driver: str = ""
91 #: Deliberate, report-disclosed manifest mutations applied before
92 #: submission, as (json-path-ish description, replacement) pairs.
93 mutations: dict[str, str] = field(default_factory=dict)
94 #: Per-example completion timeout. GPU examples get longer defaults to
95 #: absorb node provisioning.
96 timeout_seconds: int = 900
97 #: Skip unless the account has usable capacity for this instance family
98 #: (checked via service quotas before submission); empty = never skipped.
99 capacity_quota_code: str = ""
100 #: Human rationale for anything unusual above.
101 notes: str = ""
104#: Sentinel mutation value: remove the targeted entry instead of replacing it.
105REMOVE_VALUE = "__REMOVE__"
107EXAMPLE_SPECS: dict[str, ExampleSpec] = {
108 spec.name: spec
109 for spec in (
110 # --- plain batch jobs over documented CLI paths -------------------
111 ExampleSpec("simple-job", SUBMIT_SQS, JOB_COMPLETES),
112 ExampleSpec(
113 "sqs-job-submission",
114 SUBMIT_SQS,
115 JOB_COMPLETES,
116 accelerator="nvidia",
117 timeout_seconds=1800,
118 notes="two Jobs in one file; the GPU one waits for node provisioning",
119 ),
120 ExampleSpec("model-download-job", KUBECTL_APPLY, JOB_COMPLETES, timeout_seconds=1200),
121 ExampleSpec("efs-output-job", SUBMIT_DIRECT, JOB_COMPLETES),
122 ExampleSpec(
123 "fsx-lustre-job",
124 SUBMIT_DIRECT,
125 JOB_COMPLETES,
126 feature_overrides=("fsx_lustre",),
127 ),
128 ExampleSpec(
129 "aurora-pgvector-job",
130 SUBMIT_DIRECT,
131 JOB_COMPLETES,
132 feature_overrides=("aurora_pgvector",),
133 ),
134 ExampleSpec(
135 "vector-store-search-job",
136 SUBMIT_DIRECT,
137 JOB_COMPLETES,
138 feature_overrides=("vector_store",),
139 setup_driver="vector-demo-corpus",
140 notes=(
141 "read-only search; the setup driver ingests the bundled demo "
142 "corpus (gco vector ingest --demo --wait) so the >=1-hit "
143 "self-assert has something to find"
144 ),
145 ),
146 ExampleSpec(
147 "mlflow-tracking-job",
148 SUBMIT_DIRECT,
149 JOB_COMPLETES,
150 setup_driver="mlflow-ready",
151 notes=(
152 "tracking server ships with the default-on observability "
153 "bundle (no override key needed); the setup driver waits for "
154 "the mlflow Deployment and the gco-jobs CA bundle before the "
155 "client job submits"
156 ),
157 ),
158 ExampleSpec(
159 "valkey-cache-job",
160 SUBMIT_DIRECT,
161 JOB_COMPLETES,
162 feature_overrides=("valkey",),
163 ),
164 ExampleSpec("cluster-shared-bucket-upload-job", SUBMIT_DIRECT, JOB_COMPLETES),
165 ExampleSpec(
166 "regional-shared-bucket-upload-job",
167 SUBMIT_DIRECT,
168 JOB_COMPLETES,
169 notes=(
170 "no overrides needed: the regional bucket, its RW grant, and "
171 "the gco-regional-shared-bucket ConfigMap are all unconditional"
172 ),
173 ),
174 ExampleSpec(
175 "analytics-s3-upload-job",
176 SUBMIT_DIRECT,
177 JOB_COMPLETES,
178 notes=(
179 "validates the cluster-side half only; the Studio-notebook "
180 "reader documented as a prerequisite is out of scope"
181 ),
182 ),
183 ExampleSpec(
184 "analytics-database-export-job",
185 SUBMIT_DIRECT,
186 JOB_COMPLETES,
187 feature_overrides=("aurora_pgvector",),
188 notes="with Aurora forced on, the full export path runs (not the no-op branch)",
189 ),
190 # --- GPU / accelerator jobs ---------------------------------------
191 ExampleSpec(
192 "gpu-job", SUBMIT_SQS, JOB_COMPLETES, accelerator="nvidia", timeout_seconds=1800
193 ),
194 ExampleSpec(
195 "multi-gpu-training",
196 KUBECTL_APPLY,
197 JOB_COMPLETES,
198 accelerator="nvidia",
199 timeout_seconds=2400,
200 notes="indexed Job + headless Service; completion requires all indexes",
201 ),
202 ExampleSpec(
203 "efa-distributed-training",
204 SUBMIT_DIRECT,
205 JOB_COMPLETES,
206 accelerator="efa",
207 capacity_quota_code="L-417A185B",
208 timeout_seconds=2400,
209 notes=(
210 "requires P-family EFA-capable capacity; skipped with evidence "
211 "when the account's Running On-Demand P quota is zero"
212 ),
213 ),
214 ExampleSpec(
215 "inferentia-job",
216 SUBMIT_API,
217 JOB_COMPLETES,
218 accelerator="neuron",
219 capacity_quota_code="L-1945791B",
220 timeout_seconds=1800,
221 ),
222 ExampleSpec(
223 "trainium-job",
224 SUBMIT_API,
225 JOB_COMPLETES,
226 accelerator="neuron",
227 capacity_quota_code="L-2C3B7624",
228 timeout_seconds=1800,
229 ),
230 # --- inference Deployment+Service pairs ---------------------------
231 ExampleSpec(
232 "inference-vllm",
233 SUBMIT_DIRECT,
234 DEPLOYMENT_AVAILABLE,
235 accelerator="nvidia",
236 timeout_seconds=2400,
237 notes="default model (Qwen/Qwen2.5-1.5B-Instruct) is ungated; runs verbatim",
238 ),
239 ExampleSpec(
240 "inference-sglang",
241 SUBMIT_DIRECT,
242 DEPLOYMENT_AVAILABLE,
243 accelerator="nvidia",
244 timeout_seconds=2400,
245 notes="default model (Qwen/Qwen2.5-1.5B-Instruct) is ungated; runs verbatim",
246 ),
247 ExampleSpec(
248 "inference-triton",
249 SUBMIT_DIRECT,
250 DEPLOYMENT_AVAILABLE,
251 accelerator="nvidia",
252 timeout_seconds=2400,
253 notes="empty --model-repository is a valid, live initial state",
254 ),
255 # --- scheduler CRD examples (documented kubectl paths) ------------
256 ExampleSpec(
257 "kueue-job",
258 KUBECTL_APPLY,
259 JOB_COMPLETES,
260 accelerator="nvidia",
261 timeout_seconds=2400,
262 notes=(
263 "applies its own ResourceFlavors/ClusterQueue/LocalQueue plus a "
264 "CPU and a GPU Job; both Jobs must complete and the queue "
265 "objects are deleted afterwards"
266 ),
267 ),
268 ExampleSpec(
269 "volcano-gang-job",
270 KUBECTL_APPLY,
271 VCJOB_COMPLETES,
272 timeout_seconds=1200,
273 ),
274 ExampleSpec(
275 "yunikorn-job",
276 KUBECTL_APPLY,
277 JOB_COMPLETES,
278 helm_overrides=("yunikorn",),
279 accelerator="nvidia",
280 timeout_seconds=2400,
281 notes="three Jobs incl. one GPU and one gang-annotated; all must complete",
282 ),
283 ExampleSpec(
284 "slurm-cluster-job",
285 KUBECTL_APPLY,
286 JOB_COMPLETES,
287 helm_overrides=("slurm",),
288 timeout_seconds=1200,
289 ),
290 ExampleSpec(
291 "ray-cluster",
292 KUBECTL_APPLY,
293 RAYCLUSTER_READY,
294 timeout_seconds=1200,
295 ),
296 ExampleSpec(
297 "kubeflow-trainjob",
298 SUBMIT_SQS,
299 TRAINJOB_COMPLETES,
300 setup_driver="trainer-runtime-ready",
301 timeout_seconds=1800,
302 notes=(
303 "CPU-sized 2-node torchrun all-reduce; the trainer chart is "
304 "on by default (no override key), the setup driver waits for "
305 "the TrainJob CRD and the torch-distributed runtime, and the "
306 "timeout absorbs the multi-GB pytorch image pull on both nodes"
307 ),
308 ),
309 ExampleSpec(
310 "keda-scaled-job",
311 KUBECTL_APPLY,
312 SCALEDJOB_SCALES,
313 setup_driver="keda-demo-queue",
314 timeout_seconds=1200,
315 notes=(
316 "creates a disposable demo queue, grants the KEDA operator "
317 "read-only metric access via a queue policy (the documented "
318 "prerequisite), substitutes the placeholder queueURL, and "
319 "deletes the queue afterwards"
320 ),
321 ),
322 # --- self-managed platform add-ons (helm.argocd, helm.crossplane) --
323 ExampleSpec(
324 "argocd-gitops-job",
325 KUBECTL_APPLY,
326 ARGOCD_APP_HEALTHY,
327 helm_overrides=("argocd",),
328 setup_driver="argocd-revision-pin",
329 notes=(
330 "an Application syncing examples/gitops/hello-job from this "
331 "repository into gco-jobs through the fenced gco-tenants project; "
332 "the setup driver waits for Argo CD and pins targetRevision to the "
333 "commit under validation, so the run syncs exactly the code it tests"
334 ),
335 ),
336 ExampleSpec(
337 "crossplane-batch-api",
338 COMPANION,
339 NONE,
340 notes="XRD + Composition applied by crossplane-batch-job's setup driver",
341 ),
342 ExampleSpec(
343 "crossplane-batch-job",
344 KUBECTL_APPLY,
345 COMPOSED_JOB_COMPLETES,
346 helm_overrides=("crossplane",),
347 setup_driver="crossplane-api",
348 companion="crossplane-batch-api",
349 notes=(
350 "a namespaced BatchJob XR; the setup driver waits for the "
351 "go-templating function, applies the companion XRD + Composition "
352 "and waits for the API to be served; Crossplane composes the Job"
353 ),
354 ),
355 # --- AWS-managed EKS Capabilities (eks_capabilities) ---------------
356 ExampleSpec(
357 "kro-batch-api",
358 COMPANION,
359 NONE,
360 notes="ResourceGraphDefinition applied by kro-batch-job's setup driver",
361 ),
362 ExampleSpec(
363 "kro-batch-job",
364 KUBECTL_APPLY,
365 COMPOSED_JOB_COMPLETES,
366 capability_overrides=("kro",),
367 setup_driver="kro-api",
368 companion="kro-batch-api",
369 notes=(
370 "a BatchJob instance; the setup driver applies the companion "
371 "ResourceGraphDefinition and waits for it to be Active and its API "
372 "served; the kro capability composes the Job"
373 ),
374 ),
375 ExampleSpec(
376 "ack-sqs-queue",
377 KUBECTL_APPLY,
378 ACK_RESOURCE_SYNCED,
379 capability_overrides=("ack",),
380 ack_iam_policy_arns=("arn:aws:iam::aws:policy/AmazonSQSFullAccess",),
381 setup_driver="ack-sqs",
382 notes=(
383 "the ACK capability role gets AmazonSQSFullAccess for this run "
384 "(the documented simple setup); the driver confirms the queue in "
385 "SQS after sync and its deletion after cleanup"
386 ),
387 ),
388 # --- DAG pipeline --------------------------------------------------
389 ExampleSpec("pipeline-dag", DAG_RUN, DAG_SUCCEEDS, timeout_seconds=1800),
390 ExampleSpec(
391 "dag-step-preprocess",
392 COMPANION,
393 NONE,
394 notes="step manifest executed via pipeline-dag",
395 ),
396 ExampleSpec(
397 "dag-step-train", COMPANION, NONE, notes="step manifest executed via pipeline-dag"
398 ),
399 )
400}
403def required_helm_overrides(names: list[str]) -> tuple[str, ...]:
404 """Union of helm overrides needed by the selected examples (sorted)."""
405 needed: set[str] = set()
406 for name in names:
407 needed.update(EXAMPLE_SPECS[name].helm_overrides)
408 return tuple(sorted(needed))
411def required_feature_overrides(names: list[str]) -> tuple[str, ...]:
412 """Union of feature overrides needed by the selected examples (sorted)."""
413 needed: set[str] = set()
414 for name in names:
415 needed.update(EXAMPLE_SPECS[name].feature_overrides)
416 return tuple(sorted(needed))
419def required_capability_overrides(names: list[str]) -> tuple[str, ...]:
420 """Union of EKS Capability types needed by the selected examples (sorted)."""
421 needed: set[str] = set()
422 for name in names:
423 needed.update(EXAMPLE_SPECS[name].capability_overrides)
424 return tuple(sorted(needed))
427def required_capability_settings(names: list[str]) -> dict[str, dict[str, list[str]]]:
428 """Per-type capability settings the selected examples need.
430 Today only ACK carries settings: the union of the managed policies the
431 selected examples name, attached to the ACK capability role.
432 """
433 policies = sorted({arn for name in names for arn in EXAMPLE_SPECS[name].ack_iam_policy_arns})
434 return {"ack": {"iam_policy_arns": policies}} if policies else {}