Coverage for scripts / example_job_validation / drivers.py: 100.00%
336 statements
« prev ^ index » next coverage.py v7.13.5, created at 2026-09-14 22:07 +0000
1"""Submission, success-criteria, setup, and cleanup drivers for one example.
3Every driver takes the parsed example plus the run plumbing and returns
4evidence dictionaries for the report. Submission always travels the
5DOCUMENTED path (the real ``gco`` CLI or ``kubectl apply``); any deliberate
6manifest mutation (spec.mutations) is applied to a disclosed temp copy.
7"""
9from __future__ import annotations
11import json
12import subprocess
13import tempfile
14import threading
15import time
16from dataclasses import dataclass, field
17from pathlib import Path
18from typing import Any
20import yaml
22from .kube import KubectlRunner
23from .specs import (
24 DAG_RUN,
25 DEPLOYMENT_AVAILABLE,
26 JOB_COMPLETES,
27 KUBECTL_APPLY,
28 RAYCLUSTER_READY,
29 SCALEDJOB_SCALES,
30 SUBMIT_API,
31 SUBMIT_DIRECT,
32 SUBMIT_SQS,
33 TRAINJOB_COMPLETES,
34 VCJOB_COMPLETES,
35)
36from .static_checks import ParsedExample
38#: boto3 Session.client() is not thread-safe (client creation mutates shared
39#: loader state); every client creation against the run's shared session must
40#: hold this lock when examples run in parallel. The created clients ARE safe
41#: to use concurrently.
42BOTO_CLIENT_LOCK = threading.Lock()
44_POLL_SECONDS = 15
47class ExampleValidationError(RuntimeError):
48 """One example failed its criteria; the message carries the evidence."""
51@dataclass
52class ExampleRunResult:
53 """Evidence for one example's live validation."""
55 name: str
56 status: str # passed | failed | skipped
57 submission: str
58 duration_seconds: float = 0.0
59 detail: str = ""
60 mutations: dict[str, str] = field(default_factory=dict)
61 evidence: dict[str, Any] = field(default_factory=dict)
63 def to_dict(self) -> dict[str, Any]:
64 return {
65 "name": self.name,
66 "status": self.status,
67 "submission": self.submission,
68 "duration_seconds": round(self.duration_seconds, 3),
69 "detail": self.detail,
70 "mutations": self.mutations,
71 "evidence": self.evidence,
72 }
75def _run_cli(args: list[str], repo_root: Path, timeout: int = 600) -> tuple[int, str, str]:
76 result = subprocess.run(args, cwd=repo_root, capture_output=True, text=True, timeout=timeout)
77 return result.returncode, result.stdout, result.stderr
80def write_temp_manifest(documents: list[dict[str, Any]], suffix: str) -> Path:
81 """Write documents to a private temp file and return its path."""
82 fd, name = tempfile.mkstemp(suffix=suffix, text=True)
83 with open(fd, "w", encoding="utf-8") as fh:
84 yaml.safe_dump_all(documents, fh)
85 return Path(name)
88def apply_mutations(parsed: ParsedExample) -> tuple[Path, dict[str, str]]:
89 """Materialize the manifest to submit: verbatim, or a disclosed mutated copy.
91 Mutation keys use the shape ``Deployment.env.NAME`` (replace that env
92 var's value) or ``Deployment.args.--flag`` (replace the argv element
93 following the flag). Anything else is a spec bug and raises.
94 """
95 if not parsed.spec.mutations:
96 return parsed.path, {}
97 from .specs import REMOVE_VALUE
99 documents = [dict(doc) for doc in parsed.documents]
100 for key, replacement in parsed.spec.mutations.items():
101 kind, channel, target = key.split(".", 2)
102 for doc in documents:
103 if doc.get("kind") != kind:
104 continue
105 containers = doc["spec"]["template"]["spec"]["containers"]
106 for container in containers:
107 if channel == "env":
108 env_entries = container.get("env", [])
109 if replacement == REMOVE_VALUE:
110 container["env"] = [
111 entry for entry in env_entries if entry.get("name") != target
112 ]
113 continue
114 for env_entry in env_entries:
115 if env_entry.get("name") == target:
116 env_entry["value"] = replacement
117 elif channel == "args":
118 args = container.get("args", [])
119 for index, value in enumerate(args):
120 if value == target and index + 1 < len(args):
121 args[index + 1] = replacement
122 else:
123 raise ValueError(f"Unsupported mutation channel in {key!r}")
124 return (
125 write_temp_manifest(documents, f"-{parsed.name}.yaml"),
126 dict(parsed.spec.mutations),
127 )
130def submit_example(
131 parsed: ParsedExample,
132 manifest_path: Path,
133 *,
134 repo_root: Path,
135 region: str,
136 kubectl: KubectlRunner,
137) -> dict[str, Any]:
138 """Submit via the documented path; returns submission evidence."""
139 spec = parsed.spec
140 if spec.submission == SUBMIT_DIRECT:
141 args = ["gco", "jobs", "submit-direct", str(manifest_path), "-r", region]
142 elif spec.submission == SUBMIT_SQS:
143 args = ["gco", "jobs", "submit-sqs", str(manifest_path), "--region", region]
144 elif spec.submission == SUBMIT_API:
145 args = ["gco", "jobs", "submit", str(manifest_path), "--region", region]
146 elif spec.submission == DAG_RUN:
147 args = ["gco", "dag", "run", str(manifest_path), "-r", region]
148 elif spec.submission == KUBECTL_APPLY:
149 code, out, err = kubectl("apply", "-f", str(manifest_path))
150 if code != 0:
151 raise ExampleValidationError(f"kubectl apply failed: {err.strip()[:800]}")
152 return {"command": f"kubectl apply -f examples/{parsed.name}.yaml", "output": out.strip()}
153 else:
154 raise ExampleValidationError(f"No live submission for {spec.submission}")
156 timeout = 1800 if spec.submission == DAG_RUN else 600
157 code, out, err = _run_cli(args, repo_root, timeout=timeout)
158 if code != 0:
159 raise ExampleValidationError(
160 f"{' '.join(args[:3])} failed (exit {code}): {(err or out).strip()[:800]}"
161 )
162 return {"command": " ".join(args[:3]) + f" examples/{parsed.name}.yaml", "output": out[-1500:]}
165# --------------------------------------------------------------------------
166# success criteria
167# --------------------------------------------------------------------------
170def _workload_documents(parsed: ParsedExample, kinds: set[str]) -> list[dict[str, Any]]:
171 return [doc for doc in parsed.documents if doc.get("kind") in kinds]
174def _job_status(kubectl: KubectlRunner, namespace: str, name: str) -> tuple[str, str]:
175 code, out, _ = kubectl("get", "job", name, "-n", namespace, "-o", "json")
176 if code != 0:
177 return "missing", ""
178 payload = json.loads(out)
179 for condition in payload.get("status", {}).get("conditions", []) or []:
180 if condition.get("type") == "Complete" and condition.get("status") == "True":
181 return "complete", ""
182 if condition.get("type") == "Failed" and condition.get("status") == "True":
183 return "failed", str(condition.get("message", ""))
184 return "running", ""
187def _pod_diagnostics(kubectl: KubectlRunner, namespace: str, selector: str) -> str:
188 _, out, _ = kubectl(
189 "get",
190 "pods",
191 "-n",
192 namespace,
193 "-l",
194 selector,
195 "-o",
196 "jsonpath={range .items[*]}{.metadata.name}={.status.phase} {end}",
197 )
198 return out.strip()
201def _job_admission_rejection(kubectl: KubectlRunner, namespace: str, name: str) -> str | None:
202 """Return the rejection message when the Job's pods are forbidden.
204 A LimitRange or ResourceQuota rejection never becomes a Job condition:
205 the controller retries pod creation forever, the Job stays podless, and
206 the only signal is ``FailedCreate ... forbidden`` namespace events.
207 Waiting the full example timeout on such a job is pure burn (observed
208 live: example-job validation run ex241-df723811, 40 minutes against the
209 old per-container GPU ceiling) — surface the event message immediately.
210 ResourceQuota rejections (``exceeded quota``) are the one retriable
211 shape: under parallel example submission the namespace quota is
212 transiently full, the Job controller retries pod creation, and the pods
213 land once peers finish. Never fail fast on those; each message is
214 evaluated separately so a transient quota event cannot mask a permanent
215 LimitRange rejection emitted for the same Job.
216 """
217 code, out, _ = kubectl(
218 "get",
219 "events",
220 "-n",
221 namespace,
222 "--field-selector",
223 f"involvedObject.kind=Job,involvedObject.name={name},reason=FailedCreate",
224 "-o",
225 'jsonpath={range .items[*]}{.message}{"\\n"}{end}',
226 timeout=60,
227 )
228 if code != 0:
229 return None
230 for message in out.splitlines():
231 if "forbidden" in message and "exceeded quota" not in message:
232 return message[-600:]
233 return None
236def wait_jobs_complete(
237 parsed: ParsedExample, kubectl: KubectlRunner, *, timeout: int
238) -> dict[str, Any]:
239 """Every batch/v1 Job in the example must reach Complete."""
240 jobs = [
241 ((doc.get("metadata") or {}).get("namespace", "gco-jobs"), doc["metadata"]["name"])
242 for doc in _workload_documents(parsed, {"Job"})
243 ]
244 if not jobs:
245 raise ExampleValidationError("spec says job-completes but the file defines no Jobs")
246 deadline = time.monotonic() + timeout
247 pending = dict.fromkeys(jobs, "unknown")
248 while time.monotonic() < deadline:
249 for namespace, name in jobs:
250 state, message = _job_status(kubectl, namespace, name)
251 pending[(namespace, name)] = state
252 if state == "failed":
253 _, logs, _ = kubectl(
254 "logs", f"job/{name}", "-n", namespace, "--tail", "40", timeout=60
255 )
256 raise ExampleValidationError(
257 f"Job {namespace}/{name} failed: {message} :: last logs: {logs[-800:]}"
258 )
259 rejection = _job_admission_rejection(kubectl, namespace, name)
260 if rejection is not None:
261 raise ExampleValidationError(
262 f"Job {namespace}/{name} pods are rejected at admission and can never "
263 f"run: {rejection}"
264 )
265 if all(state == "complete" for state in pending.values()):
266 return {"jobs": {f"{ns}/{name}": "complete" for (ns, name) in jobs}}
267 time.sleep(_POLL_SECONDS)
268 detail = ", ".join(f"{ns}/{name}={state}" for (ns, name), state in pending.items())
269 pods = _pod_diagnostics(kubectl, jobs[0][0], f"job-name={jobs[0][1]}")
270 raise ExampleValidationError(f"timeout after {timeout}s: {detail}; pods: {pods}")
273def wait_deployment_available(
274 parsed: ParsedExample, kubectl: KubectlRunner, *, timeout: int
275) -> dict[str, Any]:
276 """The Deployment must report Available and its Service must have endpoints."""
277 deployments = _workload_documents(parsed, {"Deployment"})
278 services = _workload_documents(parsed, {"Service"})
279 if not deployments:
280 raise ExampleValidationError("spec says deployment-available but no Deployment found")
281 namespace = deployments[0]["metadata"].get("namespace", "gco-inference")
282 name = deployments[0]["metadata"]["name"]
283 code, _, err = kubectl(
284 "wait",
285 f"deployment/{name}",
286 "-n",
287 namespace,
288 "--for",
289 "condition=Available",
290 f"--timeout={timeout}s",
291 timeout=timeout + 60,
292 )
293 if code != 0:
294 pods = _pod_diagnostics(kubectl, namespace, f"app={name}")
295 _, describe, _ = kubectl("describe", f"deployment/{name}", "-n", namespace, timeout=60)
296 raise ExampleValidationError(
297 f"Deployment {namespace}/{name} never became Available: {err.strip()[:300]}; "
298 f"pods: {pods}; describe tail: {describe[-600:]}"
299 )
300 evidence: dict[str, Any] = {"deployment": f"{namespace}/{name}=Available"}
301 if services:
302 service_name = services[0]["metadata"]["name"]
303 _, endpoints, _ = kubectl(
304 "get",
305 "endpoints",
306 service_name,
307 "-n",
308 namespace,
309 "-o",
310 "jsonpath={.subsets[*].addresses[*].ip}",
311 )
312 if not endpoints.strip():
313 raise ExampleValidationError(f"Service {namespace}/{service_name} has no endpoints")
314 evidence["service_endpoints"] = endpoints.strip()
315 return evidence
318def wait_raycluster_ready(
319 parsed: ParsedExample, kubectl: KubectlRunner, *, timeout: int
320) -> dict[str, Any]:
321 clusters = _workload_documents(parsed, {"RayCluster"})
322 namespace = clusters[0]["metadata"].get("namespace", "gco-jobs")
323 name = clusters[0]["metadata"]["name"]
324 min_workers = int(clusters[0]["spec"]["workerGroupSpecs"][0].get("minReplicas", 1))
325 deadline = time.monotonic() + timeout
326 while time.monotonic() < deadline:
327 code, out, _ = kubectl("get", "raycluster", name, "-n", namespace, "-o", "json")
328 if code == 0:
329 status = json.loads(out).get("status", {})
330 state = str(status.get("state", ""))
331 ready_workers = int(status.get("readyWorkerReplicas", 0) or 0)
332 if state.lower() == "ready" and ready_workers >= min_workers:
333 return {
334 "raycluster": f"{namespace}/{name}",
335 "state": state,
336 "ready_workers": ready_workers,
337 }
338 time.sleep(_POLL_SECONDS)
339 _, describe, _ = kubectl("describe", "raycluster", name, "-n", namespace, timeout=60)
340 raise ExampleValidationError(
341 f"RayCluster {namespace}/{name} not ready after {timeout}s; tail: {describe[-600:]}"
342 )
345def wait_vcjob_completes(
346 parsed: ParsedExample, kubectl: KubectlRunner, *, timeout: int
347) -> dict[str, Any]:
348 jobs = _workload_documents(parsed, {"Job"})
349 volcano_jobs = [
350 doc for doc in jobs if str(doc.get("apiVersion", "")).startswith("batch.volcano")
351 ]
352 namespace = volcano_jobs[0]["metadata"].get("namespace", "gco-jobs")
353 name = volcano_jobs[0]["metadata"]["name"]
354 deadline = time.monotonic() + timeout
355 while time.monotonic() < deadline:
356 code, out, _ = kubectl("get", "vcjob", name, "-n", namespace, "-o", "json")
357 if code == 0:
358 phase = str(json.loads(out).get("status", {}).get("state", {}).get("phase", ""))
359 if phase == "Completed":
360 return {"vcjob": f"{namespace}/{name}", "phase": phase}
361 if phase in {"Failed", "Aborted", "Terminated"}:
362 raise ExampleValidationError(f"vcjob {namespace}/{name} reached phase {phase}")
363 time.sleep(_POLL_SECONDS)
364 raise ExampleValidationError(f"vcjob {namespace}/{name} did not complete within {timeout}s")
367def wait_trainjob_completes(
368 parsed: ParsedExample, kubectl: KubectlRunner, *, timeout: int
369) -> dict[str, Any]:
370 """Kubeflow TrainJob must reach condition Complete (Failed is terminal).
372 Evidence carries the terminal condition plus the per-child-Job counts
373 from status.jobsStatus so the report shows the gang actually ran
374 (numNodes pods succeeded), not merely that a condition flipped.
375 """
376 trainjobs = _workload_documents(parsed, {"TrainJob"})
377 namespace = trainjobs[0]["metadata"].get("namespace", "gco-jobs")
378 name = trainjobs[0]["metadata"]["name"]
379 deadline = time.monotonic() + timeout
380 while time.monotonic() < deadline:
381 code, out, _ = kubectl("get", "trainjob", name, "-n", namespace, "-o", "json")
382 if code == 0:
383 status = json.loads(out).get("status", {}) or {}
384 jobs_status = status.get("jobsStatus", []) or []
385 for condition in status.get("conditions", []) or []:
386 if condition.get("status") != "True":
387 continue
388 if condition.get("type") == "Complete":
389 return {
390 "trainjob": f"{namespace}/{name}",
391 "condition": "Complete",
392 "jobsStatus": jobs_status,
393 }
394 if condition.get("type") == "Failed":
395 raise ExampleValidationError(
396 f"TrainJob {namespace}/{name} reached condition Failed: "
397 f"{condition.get('message', '')}"
398 )
399 time.sleep(_POLL_SECONDS)
400 _, describe, _ = kubectl("describe", "trainjob", name, "-n", namespace, timeout=60)
401 raise ExampleValidationError(
402 f"TrainJob {namespace}/{name} did not complete within {timeout}s; tail: {describe[-600:]}"
403 )
406def wait_scaledjob_scales(
407 parsed: ParsedExample, kubectl: KubectlRunner, *, timeout: int
408) -> dict[str, Any]:
409 """KEDA must spawn at least one Job for the ScaledJob from queue depth."""
410 scaled = _workload_documents(parsed, {"ScaledJob"})
411 namespace = scaled[0]["metadata"].get("namespace", "gco-jobs")
412 name = scaled[0]["metadata"]["name"]
413 deadline = time.monotonic() + timeout
414 while time.monotonic() < deadline:
415 code, out, _ = kubectl(
416 "get",
417 "jobs",
418 "-n",
419 namespace,
420 "-l",
421 f"scaledjob.keda.sh/name={name}",
422 "-o",
423 "jsonpath={.items[*].metadata.name}",
424 )
425 spawned = [item for item in out.split() if item]
426 if code == 0 and spawned:
427 return {"scaledjob": f"{namespace}/{name}", "spawned_jobs": spawned[:5]}
428 time.sleep(_POLL_SECONDS)
429 _, describe, _ = kubectl("describe", "scaledjob", name, "-n", namespace, timeout=60)
430 raise ExampleValidationError(
431 f"ScaledJob {namespace}/{name} spawned no Jobs within {timeout}s; tail: {describe[-600:]}"
432 )
435CRITERIA_WAITERS = {
436 JOB_COMPLETES: wait_jobs_complete,
437 DEPLOYMENT_AVAILABLE: wait_deployment_available,
438 RAYCLUSTER_READY: wait_raycluster_ready,
439 VCJOB_COMPLETES: wait_vcjob_completes,
440 SCALEDJOB_SCALES: wait_scaledjob_scales,
441 TRAINJOB_COMPLETES: wait_trainjob_completes,
442}
445# --------------------------------------------------------------------------
446# cleanup
447# --------------------------------------------------------------------------
450def cleanup_example(
451 parsed: ParsedExample, manifest_path: Path, kubectl: KubectlRunner
452) -> dict[str, Any]:
453 """Delete everything the example created and verify it is gone."""
454 if parsed.spec.submission == DAG_RUN:
455 # A DAG example's file is a pipeline SPEC, not a Kubernetes manifest
456 # (kubectl cannot decode it — observed live in run ex241-4bf01801);
457 # what actually ran on the cluster are the step manifests it names.
458 repo_root = parsed.path.parent.parent
459 deleted: list[str] = []
460 for document in parsed.documents:
461 for step in document.get("steps", []):
462 step_manifest = repo_root / str(step.get("manifest", ""))
463 code, out, err = kubectl(
464 "delete",
465 "-f",
466 str(step_manifest),
467 "--ignore-not-found",
468 "--wait=true",
469 timeout=300,
470 )
471 if code != 0:
472 raise ExampleValidationError(
473 f"cleanup failed for {parsed.name} step "
474 f"{step.get('name', '?')}: {err.strip()[:500]}"
475 )
476 deleted.extend(line for line in out.strip().splitlines() if line)
477 return {"deleted": deleted[:20]}
478 code, out, err = kubectl(
479 "delete", "-f", str(manifest_path), "--ignore-not-found", "--wait=true", timeout=300
480 )
481 if code != 0:
482 raise ExampleValidationError(f"cleanup failed for {parsed.name}: {err.strip()[:500]}")
483 return {"deleted": [line for line in out.strip().splitlines() if line][:20]}
486# --------------------------------------------------------------------------
487# setup drivers (spec.setup_driver)
488# --------------------------------------------------------------------------
491@dataclass
492class KedaDemoQueue:
493 """Disposable SQS queue backing the KEDA scaling demonstration.
495 Implements the example's documented prerequisites: a demo queue with
496 synthetic messages and read-only queue-metric access for the KEDA
497 operator (granted with a queue policy, never by touching IAM roles).
498 """
500 session: Any
501 region: str
502 run_id: str
503 queue_url: str = ""
504 queue_arn: str = ""
506 def _sqs(self) -> Any:
507 # boto3 Session.client() is not thread-safe; examples may run in
508 # parallel threads. The returned client is safe to use concurrently.
509 with BOTO_CLIENT_LOCK:
510 return self.session.client("sqs", region_name=self.region)
512 def create(self, operator_role_arn: str) -> dict[str, Any]:
513 sqs = self._sqs()
514 name = f"gco-keda-demo-{self.run_id}"[:80]
515 self.queue_url = sqs.create_queue(QueueName=name)["QueueUrl"]
516 attrs = sqs.get_queue_attributes(QueueUrl=self.queue_url, AttributeNames=["QueueArn"])
517 self.queue_arn = attrs["Attributes"]["QueueArn"]
518 policy = {
519 "Version": "2012-10-17",
520 "Statement": [
521 {
522 "Sid": "KedaOperatorQueueMetrics",
523 "Effect": "Allow",
524 "Principal": {"AWS": operator_role_arn},
525 "Action": ["sqs:GetQueueAttributes", "sqs:GetQueueUrl"],
526 "Resource": self.queue_arn,
527 }
528 ],
529 }
530 sqs.set_queue_attributes(QueueUrl=self.queue_url, Attributes={"Policy": json.dumps(policy)})
531 for index in range(10):
532 sqs.send_message(QueueUrl=self.queue_url, MessageBody=f"demo-{index}")
533 return {"queue_arn": self.queue_arn, "seeded_messages": 10}
535 def destroy(self) -> None:
536 if self.queue_url:
537 self._sqs().delete_queue(QueueUrl=self.queue_url)
540@dataclass
541class VectorDemoCorpus:
542 """Demo-corpus precondition for the vector-search example, fully reverted.
544 ``create`` runs the example's documented prerequisite verbatim —
545 ``gco vector ingest --demo --wait`` — and records exactly which corpus
546 objects the CLI uploaded. ``destroy`` reverts precisely those: the S3
547 objects (whose upload is what triggered ingestion) and every DynamoDB
548 chunk item whose ``source`` is one of the recorded keys, resolved with
549 the same filtered-Scan shape the CLI's ingest wait uses (the table has
550 no by-source key schema; corpora are document-scale). Scoping deletion
551 to the recorded keys means a pre-existing user corpus in the same table
552 is never touched.
553 """
555 repo_root: Path
556 session: Any
557 region: str
558 uploaded: list[str] = field(default_factory=list)
559 bucket: str = ""
561 def create(self) -> dict[str, Any]:
562 code, out, err = _run_cli(
563 ["gco", "vector", "ingest", "--demo", "--wait", "--output", "json"],
564 self.repo_root,
565 timeout=900,
566 )
567 if code != 0:
568 raise ExampleValidationError(
569 f"gco vector ingest --demo --wait failed (exit {code}): "
570 f"{(err or out).strip()[:800]}"
571 )
572 try:
573 summary = json.loads(out)
574 except ValueError as exc:
575 raise ExampleValidationError(
576 f"gco vector ingest emitted non-JSON output: {out[:400]}"
577 ) from exc
578 self.bucket = str(summary.get("bucket", ""))
579 self.uploaded = [str(key) for key in summary.get("uploaded", [])]
580 if not self.bucket or not self.uploaded:
581 raise ExampleValidationError(
582 f"ingest summary carried no bucket/keys to revert later: {out[:400]}"
583 )
584 return {
585 "command": "gco vector ingest --demo --wait",
586 "bucket": self.bucket,
587 "uploaded": self.uploaded,
588 "chunks_by_source": summary.get("chunks_by_source", {}),
589 }
591 def destroy(self) -> None:
592 if not self.uploaded:
593 return
594 from cli.vector_store import VectorStoreClient
596 client = VectorStoreClient(query_region=self.region)
597 table_name = client._resolve_table_name()
598 bucket_name, bucket_region = client._resolve_bucket()
599 if bucket_name != self.bucket:
600 # Fail loudly rather than delete from a bucket other than the
601 # one create() actually uploaded to.
602 raise ExampleValidationError(
603 f"corpus bucket changed between ingest ({self.bucket}) and "
604 f"revert ({bucket_name}); refusing to delete"
605 )
606 with BOTO_CLIENT_LOCK:
607 dynamodb = self.session.client("dynamodb", region_name=self.region)
608 s3 = self.session.client("s3", region_name=bucket_region)
610 # Chunk items first (their source keys reference the S3 objects),
611 # then the objects themselves. Deleting the objects does NOT
612 # un-ingest — the notification only fires on creates — hence the
613 # explicit item sweep.
614 for key in self.uploaded:
615 doc_ids: list[str] = []
616 scan_kwargs: dict[str, Any] = {
617 "TableName": table_name,
618 "FilterExpression": "#source = :source",
619 "ExpressionAttributeNames": {"#source": "source"},
620 "ExpressionAttributeValues": {":source": {"S": key}},
621 "ProjectionExpression": "doc_id",
622 }
623 while True:
624 page = dynamodb.scan(**scan_kwargs)
625 doc_ids.extend(
626 item["doc_id"]["S"] for item in page.get("Items", []) if "doc_id" in item
627 )
628 last_key = page.get("LastEvaluatedKey")
629 if not last_key:
630 break
631 scan_kwargs["ExclusiveStartKey"] = last_key
632 for start in range(0, len(doc_ids), 25):
633 batch = doc_ids[start : start + 25]
634 dynamodb.batch_write_item(
635 RequestItems={
636 table_name: [
637 {"DeleteRequest": {"Key": {"doc_id": {"S": doc_id}}}}
638 for doc_id in batch
639 ]
640 }
641 )
642 s3.delete_object(Bucket=bucket_name, Key=key)
645def wait_trainer_runtime_ready(kubectl: KubectlRunner, *, timeout: int = 300) -> dict[str, Any]:
646 """Wait until the TrainJob CRD is served and the shipped runtime exists.
648 The trainer chart installs the CRDs and controller; the post-Helm
649 kubectl pass applies the torch-distributed ClusterTrainingRuntime the
650 example's ``runtimeRef`` names. Both are deploy-time artifacts, so this
651 is a readiness wait, not created state — nothing to revert.
652 """
653 deadline = time.monotonic() + timeout
654 last_error = ""
655 while True:
656 code, _, err = kubectl("get", "crd", "trainjobs.trainer.kubeflow.org")
657 if code != 0:
658 last_error = f"TrainJob CRD not present: {err.strip()[:300]}"
659 else:
660 code, out, err = kubectl(
661 "get", "clustertrainingruntime", "torch-distributed", "-o", "json"
662 )
663 if code == 0:
664 runtime = json.loads(out)
665 return {
666 "crd": "trainjobs.trainer.kubeflow.org",
667 "runtime": runtime["metadata"]["name"],
668 "runtime_created": runtime["metadata"].get("creationTimestamp", ""),
669 }
670 last_error = f"torch-distributed runtime not present: {err.strip()[:300]}"
671 if time.monotonic() >= deadline:
672 break
673 time.sleep(_POLL_SECONDS)
674 raise ExampleValidationError(
675 f"Kubeflow Trainer runtime not ready within {timeout}s — is "
676 f"helm.kubeflow_trainer enabled? Last error: {last_error}"
677 )
680def wait_mlflow_ready(kubectl: KubectlRunner, *, timeout: int = 600) -> dict[str, Any]:
681 """Wait until the MLflow tracking server Deployment is Available.
683 The example's client job fails its read-back (or hangs on connect) if
684 it races the server's first rollout — the backend PVC arrives one
685 applier pass after the chart on a fresh install. Readiness wait only;
686 nothing to revert.
687 """
688 deadline = time.monotonic() + timeout
689 last_state = ""
690 while True:
691 code, out, err = kubectl("get", "deployment", "mlflow", "-n", "monitoring", "-o", "json")
692 if code != 0:
693 last_state = f"mlflow Deployment not found: {err.strip()[:300]}"
694 else:
695 payload = json.loads(out)
696 conditions = payload.get("status", {}).get("conditions", []) or []
697 available = any(
698 condition.get("type") == "Available" and condition.get("status") == "True"
699 for condition in conditions
700 )
701 if available:
702 return {
703 "deployment": "monitoring/mlflow",
704 "ready_replicas": payload.get("status", {}).get("readyReplicas", 0),
705 }
706 last_state = f"conditions: {json.dumps(conditions)[:400]}"
707 if time.monotonic() >= deadline:
708 break
709 time.sleep(_POLL_SECONDS)
710 raise ExampleValidationError(
711 f"MLflow tracking server not Available within {timeout}s — is "
712 f"cluster_observability.mlflow enabled? Last state: {last_state}"
713 )
716#: Setup drivers _run_one_example knows how to dispatch, by spec.setup_driver
717#: name. Kept as an explicit registry so a spec naming a driver that does not
718#: exist fails the registry pin test, not a live run.
719KNOWN_SETUP_DRIVERS = frozenset(
720 {
721 "keda-demo-queue",
722 "vector-demo-corpus",
723 "trainer-runtime-ready",
724 "mlflow-ready",
725 }
726)