ariadne/ariadne/services/platform_quality_probe.py

137 lines
5.5 KiB
Python

from __future__ import annotations
from dataclasses import dataclass
import time
from typing import Any
from ..k8s.client import get_json, post_json
from ..settings import settings
from ..utils.logging import get_logger
logger = get_logger(__name__)
@dataclass(frozen=True)
class PlatformQualityProbeResult:
"""Represent one platform-quality probe execution.
Inputs: Kubernetes Job completion details gathered from API polling.
Outputs: a stable status payload for scheduler logs and metrics.
"""
job: str
status: str
class PlatformQualityProbeService:
"""Run the platform quality-suite probe as an Ariadne-owned one-shot Job.
Inputs: scheduler invocations plus settings that define namespace, image,
probe script ConfigMap, and Pushgateway endpoint.
Outputs: structured run status so operators can verify probe freshness
without relying on standalone CronJob ownership.
"""
def _job_payload(self, job_name: str) -> dict[str, Any]:
payload: dict[str, Any] = {
"apiVersion": "batch/v1",
"kind": "Job",
"metadata": {
"name": job_name,
"namespace": settings.platform_quality_probe_namespace,
"labels": {
"app": "platform-quality-suite-probe",
"atlas.bstein.dev/trigger": "ariadne",
},
},
"spec": {
"backoffLimit": 0,
"ttlSecondsAfterFinished": settings.platform_quality_probe_job_ttl_sec,
"template": {
"metadata": {"labels": {"app": "platform-quality-suite-probe"}},
"spec": {
"restartPolicy": "Never",
"containers": [
{
"name": "probe",
"image": settings.platform_quality_probe_image,
"imagePullPolicy": "IfNotPresent",
"command": ["/bin/sh", "/scripts/platform_quality_suite_probe.sh"],
"env": [
{
"name": "PUSHGATEWAY_URL",
"value": settings.platform_quality_probe_pushgateway_url,
},
{
"name": "HTTP_TIMEOUT_SECONDS",
"value": str(settings.platform_quality_probe_http_timeout_sec),
},
],
"volumeMounts": [
{"name": "probe-script", "mountPath": "/scripts", "readOnly": True},
],
}
],
"volumes": [
{
"name": "probe-script",
"configMap": {
"name": settings.platform_quality_probe_script_configmap,
"defaultMode": 365,
},
}
],
},
},
},
}
return payload
def _wait_for_completion(self, job_name: str, timeout_sec: float) -> PlatformQualityProbeResult:
deadline = time.time() + timeout_sec
while time.time() < deadline:
job = get_json(
f"/apis/batch/v1/namespaces/{settings.platform_quality_probe_namespace}/jobs/{job_name}"
)
status = job.get("status") if isinstance(job.get("status"), dict) else {}
if int(status.get("succeeded") or 0) > 0:
return PlatformQualityProbeResult(job=job_name, status="ok")
if int(status.get("failed") or 0) > 0:
return PlatformQualityProbeResult(job=job_name, status="error")
time.sleep(2)
return PlatformQualityProbeResult(job=job_name, status="running")
def run(self, wait: bool = True) -> dict[str, Any]:
"""Launch and optionally wait on the quality-suite probe job.
Inputs: `wait` controls whether the scheduler blocks until completion.
Outputs: job identity and status for metrics/events so Grafana can report
the latest probe outcome.
"""
job_name = f"platform-quality-suite-probe-{int(time.time())}"
created = post_json(
f"/apis/batch/v1/namespaces/{settings.platform_quality_probe_namespace}/jobs",
self._job_payload(job_name),
)
name = created.get("metadata", {}).get("name", job_name)
logger.info(
"platform quality probe job triggered",
extra={"event": "platform_quality_probe_trigger", "job": name},
)
if not wait:
return {"job": name, "status": "queued"}
result = self._wait_for_completion(name, settings.platform_quality_probe_wait_timeout_sec)
if result.status != "ok":
logger.error(
"platform quality probe incomplete",
extra={"event": "platform_quality_probe_incomplete", "job": name, "status": result.status},
)
raise RuntimeError(f"platform quality probe job {name} {result.status}")
return {"job": result.job, "status": result.status}
platform_quality_probe = PlatformQualityProbeService()