Coverage for scripts / example_job_validation / specs.py: 100.00%
43 statements
« prev ^ index » next coverage.py v7.13.5, created at 2026-09-14 22:07 +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``, GPU/Neuron capacity, special setup drivers);
13* **when** it counts as passed (workload-specific success criteria); and
14* **which** deliberate mutations the harness applies before submission
15 (e.g. replacing a gated HuggingFace model with an ungated one) — every
16 mutation is disclosed in the report.
18Keep this table boring and declarative: the drivers in
19``checks/examples.py`` interpret it.
20"""
22from __future__ import annotations
24from dataclasses import dataclass, field
26#: Documented submission paths the drivers know how to execute.
27SUBMIT_DIRECT = "cli-submit-direct" # gco jobs submit-direct <file> -r <region>
28SUBMIT_SQS = "cli-submit-sqs" # gco jobs submit-sqs <file> --region <region>
29SUBMIT_API = "cli-submit-api" # gco jobs submit <file> --region <region>
30DAG_RUN = "cli-dag-run" # gco dag run <file> -r <region>
31KUBECTL_APPLY = "kubectl-apply" # kubectl apply -f <file> (over the SSM tunnel)
32COMPANION = "companion-artifact" # not independently runnable (data/DAG-step file)
34SUBMISSION_PATHS = (SUBMIT_DIRECT, SUBMIT_SQS, SUBMIT_API, DAG_RUN, KUBECTL_APPLY, COMPANION)
36#: Success-criteria kinds the drivers implement.
37JOB_COMPLETES = "job-completes" # batch/v1 Job reaches Complete
38DEPLOYMENT_AVAILABLE = "deployment-available" # Deployment Available + Service endpoints
39RAYCLUSTER_READY = "raycluster-ready" # RayCluster ready: head + minReplicas workers
40VCJOB_COMPLETES = "vcjob-completes" # batch.volcano.sh Job phase Completed
41SCALEDJOB_SCALES = "scaledjob-scales" # KEDA ScaledJob spawns >=1 Job from queue depth
42TRAINJOB_COMPLETES = "trainjob-completes" # trainer.kubeflow.org TrainJob condition Complete
43DAG_SUCCEEDS = "dag-succeeds" # gco dag run exits 0 with all steps completed
44NONE = "none" # companion artifacts: static checks only
47@dataclass(frozen=True)
48class ExampleSpec:
49 """Declarative validation contract for one example file."""
51 #: File stem under ``examples/`` (e.g. ``simple-job`` for simple-job.yaml).
52 name: str
53 #: One of :data:`SUBMISSION_PATHS` — must match the documented usage.
54 submission: str
55 #: One of the success-criteria kinds above.
56 criteria: str
57 #: Optional helm charts that must be force-enabled for this example
58 #: (threaded to CDK as ``helm_enabled_overrides``).
59 helm_overrides: tuple[str, ...] = ()
60 #: Optional infrastructure features that must be force-enabled
61 #: (threaded to CDK as ``feature_enabled_overrides``).
62 feature_overrides: tuple[str, ...] = ()
63 #: Accelerator requirement: "" (none), "nvidia", "neuron", or "efa".
64 accelerator: str = ""
65 #: Named setup/teardown driver hooks (implemented in checks/examples.py):
66 #: "keda-demo-queue".
67 setup_driver: str = ""
68 #: Deliberate, report-disclosed manifest mutations applied before
69 #: submission, as (json-path-ish description, replacement) pairs.
70 mutations: dict[str, str] = field(default_factory=dict)
71 #: Per-example completion timeout. GPU examples get longer defaults to
72 #: absorb node provisioning.
73 timeout_seconds: int = 900
74 #: Skip unless the account has usable capacity for this instance family
75 #: (checked via service quotas before submission); empty = never skipped.
76 capacity_quota_code: str = ""
77 #: Human rationale for anything unusual above.
78 notes: str = ""
81#: Sentinel mutation value: remove the targeted entry instead of replacing it.
82REMOVE_VALUE = "__REMOVE__"
84_GATED_MODEL_MUTATION_NOTE = (
85 "the manifest's default model is HuggingFace-gated; validation substitutes "
86 "the ungated facebook/opt-125m so the endpoint can become ready without "
87 "credentials — the serving path itself is exercised unchanged"
88)
90EXAMPLE_SPECS: dict[str, ExampleSpec] = {
91 spec.name: spec
92 for spec in (
93 # --- plain batch jobs over documented CLI paths -------------------
94 ExampleSpec("simple-job", SUBMIT_SQS, JOB_COMPLETES),
95 ExampleSpec(
96 "sqs-job-submission",
97 SUBMIT_SQS,
98 JOB_COMPLETES,
99 accelerator="nvidia",
100 timeout_seconds=1800,
101 notes="two Jobs in one file; the GPU one waits for node provisioning",
102 ),
103 ExampleSpec("model-download-job", KUBECTL_APPLY, JOB_COMPLETES, timeout_seconds=1200),
104 ExampleSpec("efs-output-job", SUBMIT_DIRECT, JOB_COMPLETES),
105 ExampleSpec(
106 "fsx-lustre-job",
107 SUBMIT_DIRECT,
108 JOB_COMPLETES,
109 feature_overrides=("fsx_lustre",),
110 ),
111 ExampleSpec(
112 "aurora-pgvector-job",
113 SUBMIT_DIRECT,
114 JOB_COMPLETES,
115 feature_overrides=("aurora_pgvector",),
116 ),
117 ExampleSpec(
118 "vector-store-search-job",
119 SUBMIT_DIRECT,
120 JOB_COMPLETES,
121 feature_overrides=("vector_store",),
122 setup_driver="vector-demo-corpus",
123 notes=(
124 "read-only search; the setup driver ingests the bundled demo "
125 "corpus (gco vector ingest --demo --wait) so the >=1-hit "
126 "self-assert has something to find"
127 ),
128 ),
129 ExampleSpec(
130 "mlflow-tracking-job",
131 SUBMIT_DIRECT,
132 JOB_COMPLETES,
133 setup_driver="mlflow-ready",
134 notes=(
135 "tracking server ships with the default-on observability "
136 "bundle (no override key needed); the setup driver waits for "
137 "the mlflow Deployment before the client job submits"
138 ),
139 ),
140 ExampleSpec(
141 "valkey-cache-job",
142 SUBMIT_DIRECT,
143 JOB_COMPLETES,
144 feature_overrides=("valkey",),
145 ),
146 ExampleSpec("cluster-shared-bucket-upload-job", SUBMIT_DIRECT, JOB_COMPLETES),
147 ExampleSpec(
148 "regional-shared-bucket-upload-job",
149 SUBMIT_DIRECT,
150 JOB_COMPLETES,
151 notes=(
152 "no overrides needed: the regional bucket, its RW grant, and "
153 "the gco-regional-shared-bucket ConfigMap are all unconditional"
154 ),
155 ),
156 ExampleSpec(
157 "analytics-s3-upload-job",
158 SUBMIT_DIRECT,
159 JOB_COMPLETES,
160 notes=(
161 "validates the cluster-side half only; the Studio-notebook "
162 "reader documented as a prerequisite is out of scope"
163 ),
164 ),
165 ExampleSpec(
166 "analytics-database-export-job",
167 SUBMIT_DIRECT,
168 JOB_COMPLETES,
169 feature_overrides=("aurora_pgvector",),
170 notes="with Aurora forced on, the full export path runs (not the no-op branch)",
171 ),
172 # --- GPU / accelerator jobs ---------------------------------------
173 ExampleSpec(
174 "gpu-job", SUBMIT_SQS, JOB_COMPLETES, accelerator="nvidia", timeout_seconds=1800
175 ),
176 ExampleSpec(
177 "multi-gpu-training",
178 KUBECTL_APPLY,
179 JOB_COMPLETES,
180 accelerator="nvidia",
181 timeout_seconds=2400,
182 notes="indexed Job + headless Service; completion requires all indexes",
183 ),
184 ExampleSpec(
185 "efa-distributed-training",
186 SUBMIT_DIRECT,
187 JOB_COMPLETES,
188 accelerator="efa",
189 capacity_quota_code="L-417A185B",
190 timeout_seconds=2400,
191 notes=(
192 "requires P-family EFA-capable capacity; skipped with evidence "
193 "when the account's Running On-Demand P quota is zero"
194 ),
195 ),
196 ExampleSpec(
197 "inferentia-job",
198 SUBMIT_API,
199 JOB_COMPLETES,
200 accelerator="neuron",
201 capacity_quota_code="L-1945791B",
202 timeout_seconds=1800,
203 ),
204 ExampleSpec(
205 "trainium-job",
206 SUBMIT_API,
207 JOB_COMPLETES,
208 accelerator="neuron",
209 capacity_quota_code="L-2C3B7624",
210 timeout_seconds=1800,
211 ),
212 # --- inference Deployment+Service pairs ---------------------------
213 ExampleSpec(
214 "inference-vllm",
215 SUBMIT_DIRECT,
216 DEPLOYMENT_AVAILABLE,
217 accelerator="nvidia",
218 mutations={
219 "Deployment.env.MODEL": "facebook/opt-125m",
220 "Deployment.env.MAX_MODEL_LEN": "2048",
221 },
222 timeout_seconds=2400,
223 notes=_GATED_MODEL_MUTATION_NOTE,
224 ),
225 ExampleSpec(
226 "inference-sglang",
227 SUBMIT_DIRECT,
228 DEPLOYMENT_AVAILABLE,
229 accelerator="nvidia",
230 timeout_seconds=2400,
231 notes="default model (microsoft/Phi-3.5-mini-instruct) is ungated; runs verbatim",
232 ),
233 ExampleSpec(
234 "inference-tgi",
235 SUBMIT_DIRECT,
236 DEPLOYMENT_AVAILABLE,
237 accelerator="nvidia",
238 mutations={
239 "Deployment.env.MODEL_ID": "facebook/opt-125m",
240 # AWQ requires an AWQ-quantized checkpoint; the substitute
241 # model is served unquantized.
242 "Deployment.env.QUANTIZE": REMOVE_VALUE,
243 },
244 timeout_seconds=2400,
245 notes=_GATED_MODEL_MUTATION_NOTE,
246 ),
247 ExampleSpec(
248 "inference-torchserve",
249 SUBMIT_DIRECT,
250 DEPLOYMENT_AVAILABLE,
251 accelerator="nvidia",
252 timeout_seconds=2400,
253 notes="serves from an (empty) EFS model store; readiness with no models is the documented initial state",
254 ),
255 ExampleSpec(
256 "inference-triton",
257 SUBMIT_DIRECT,
258 DEPLOYMENT_AVAILABLE,
259 accelerator="nvidia",
260 timeout_seconds=2400,
261 notes="empty --model-repository is a valid, live initial state",
262 ),
263 # --- scheduler CRD examples (documented kubectl paths) ------------
264 ExampleSpec(
265 "kueue-job",
266 KUBECTL_APPLY,
267 JOB_COMPLETES,
268 accelerator="nvidia",
269 timeout_seconds=2400,
270 notes=(
271 "applies its own ResourceFlavors/ClusterQueue/LocalQueue plus a "
272 "CPU and a GPU Job; both Jobs must complete and the queue "
273 "objects are deleted afterwards"
274 ),
275 ),
276 ExampleSpec(
277 "volcano-gang-job",
278 KUBECTL_APPLY,
279 VCJOB_COMPLETES,
280 timeout_seconds=1200,
281 ),
282 ExampleSpec(
283 "yunikorn-job",
284 KUBECTL_APPLY,
285 JOB_COMPLETES,
286 helm_overrides=("yunikorn",),
287 accelerator="nvidia",
288 timeout_seconds=2400,
289 notes="three Jobs incl. one GPU and one gang-annotated; all must complete",
290 ),
291 ExampleSpec(
292 "slurm-cluster-job",
293 KUBECTL_APPLY,
294 JOB_COMPLETES,
295 helm_overrides=("slurm",),
296 timeout_seconds=1200,
297 ),
298 ExampleSpec(
299 "ray-cluster",
300 KUBECTL_APPLY,
301 RAYCLUSTER_READY,
302 timeout_seconds=1200,
303 ),
304 ExampleSpec(
305 "kubeflow-trainjob",
306 SUBMIT_SQS,
307 TRAINJOB_COMPLETES,
308 setup_driver="trainer-runtime-ready",
309 timeout_seconds=1800,
310 notes=(
311 "CPU-sized 2-node torchrun all-reduce; the trainer chart is "
312 "on by default (no override key), the setup driver waits for "
313 "the TrainJob CRD and the torch-distributed runtime, and the "
314 "timeout absorbs the multi-GB pytorch image pull on both nodes"
315 ),
316 ),
317 ExampleSpec(
318 "keda-scaled-job",
319 KUBECTL_APPLY,
320 SCALEDJOB_SCALES,
321 setup_driver="keda-demo-queue",
322 timeout_seconds=1200,
323 notes=(
324 "creates a disposable demo queue, grants the KEDA operator "
325 "read-only metric access via a queue policy (the documented "
326 "prerequisite), substitutes the placeholder queueURL, and "
327 "deletes the queue afterwards"
328 ),
329 ),
330 # --- DAG pipeline --------------------------------------------------
331 ExampleSpec("pipeline-dag", DAG_RUN, DAG_SUCCEEDS, timeout_seconds=1800),
332 ExampleSpec(
333 "dag-step-preprocess",
334 COMPANION,
335 NONE,
336 notes="step manifest executed via pipeline-dag",
337 ),
338 ExampleSpec(
339 "dag-step-train", COMPANION, NONE, notes="step manifest executed via pipeline-dag"
340 ),
341 )
342}
345def required_helm_overrides(names: list[str]) -> tuple[str, ...]:
346 """Union of helm overrides needed by the selected examples (sorted)."""
347 needed: set[str] = set()
348 for name in names:
349 needed.update(EXAMPLE_SPECS[name].helm_overrides)
350 return tuple(sorted(needed))
353def required_feature_overrides(names: list[str]) -> tuple[str, ...]:
354 """Union of feature overrides needed by the selected examples (sorted)."""
355 needed: set[str] = set()
356 for name in names:
357 needed.update(EXAMPLE_SPECS[name].feature_overrides)
358 return tuple(sorted(needed))