OmniSRE

OmniSRE turns DataHub into an autonomous SRE control plane—detecting silent data & ML failures, tracing graph blast radius, and executing reversible self-healing fixes.

Comment

OmniSRE: Autonomous Incident Remediation & Self-Healing for Data & ML Pipelines on DataHub

Inspiration

Modern data and ML pipelines are notorious for failing quietly. When an upstream data provider renames a column, shifts currency units from dollars to cents, or drops half its daily partition volume, traditional test suites often don't help:

  • The "All Green" Blindspot: Standard assertion suites (dbt test, Great Expectations) only validate explicit constraints on data that actually arrives. A pipeline can suffer severe training-serving skew or label leakage (e.g., target features accidentally bleeding into training sets, pushing ROC-AUC to a suspicious $\approx 0.99$), yet all standard schema tests remain green.
  • Triage Fatigue: When an alert does fire downstream, data engineers and on-call SREs spend hours manually tracing dependencies—jumping between dbt run logs, warehouse tables, git commit histories, and Slack threads to assess the blast radius.

We built OmniSRE to turn DataHub from a passive metadata catalog into an active, closed-loop SRE control plane. Grounded in DataHub's column-level lineage graph, OmniSRE catches silent distribution anomalies, isolates root causes to specific upstream commits or schema shifts, executes safe, reversible mitigations, generates AST-level migration pull requests for breaking APIs, and writes post-mortems back to DataHub entity cards.

What It Does

OmniSRE provides end-to-end autonomous reliability engineering across data warehouses, dbt transformations, and MLflow model pipelines:

[ Code / Data / API ] ──► [ 13 Chaos Detectors ] ──► [ DataHub Lineage Graph ]
                                                            │
                                                            ▼
[ Post-Mortem & Writeback ] ◄── [ Reversible Action Journal ] ◄── [ 3-Tier Policy Engine ]

1. 13 Production Chaos Detectors (Catching Silent Failures)

Monitors five distinct failure classes that standard SQL assertions miss:

  • Silent Feed Failures: Upstream pipeline stalls (stale_feed), missing row volumes (volume_collapse), and non-idempotent duplicate batch retries (duplicate_batch).
  • Data Quality & Scaling: 100x unit scaling bugs (unit_bug), null rate spikes (null_spike), subtle distribution drifts (distribution_drift), and upstream schema renames (schema_change).
  • Silent ML Regressions: Training/serving input skew (training_serving_skew), prediction distribution shift (model_drift), and target leakage (label_leakage).
  • Champion Model Degradation: Blocks degraded candidate models from reaching production aliases based on MLflow tracking metrics (training_regression).
  • API Breaking Changes: Detects signature deprecations and breaking changes across upstream dependencies and SaaS webhooks (api_breaking_change).

2. Lineage-Grounded RCA & Blast Radius Computation

Instead of guessing root causes from raw logs, OmniSRE queries DataHub's GraphQL and REST APIs to traverse lineage upstream from the failing model or scoring endpoint back to the culprit dataset, column, or git commit. It then calculates the downstream blast radius and financial exposure in real time based on downstream asset criticality.

3. Safe, Reversible Mitigations with 3-Tier Policy Gating

Every action is executed through an append-only Action Journal (agent/journal.py) with strict inverse operations:

  • AUTO Tier: Low-risk, fully reversible actions (e.g., tagging degraded assets, pausing downstream ingestion jobs, quarantining corrupted partitions) execute autonomously.
  • PR_ONLY Tier: Code refactors and model alias promotions generate verified GitHub Pull Requests with unified diffs.
  • HUMAN_ONLY Tier: High-blast-radius operations require one-click approval via Slack Block Kit interactive buttons and trigger PagerDuty escalations.
  • Time-Machine Rollback: If post-mitigation validation fails, OmniSRE automatically executes the inverse operation from the journal to restore the last known good state.

4. AST-Based Patch Generation for Breaking APIs

When an upstream library or vendor API introduces a breaking signature change, OmniSRE parses the repository using Python's ast module, pinpoints all affected call sites, verifies the proposed patch in a sandbox environment, and opens a GitHub Pull Request with the corrected code.

5. Multi-Stakeholder Incident Reporting & DataHub Writeback

  • Engineers: Receive exact AST diffs, stack traces, and targeted dbt run commands.
  • SREs / Incident Commanders: Receive MTTR tracking, circuit breaker status, and journal rollback triggers.
  • Business Stakeholders: Receive financial risk metrics and SLA impact summaries.
  • DataHub Graph: Emits structured incident summaries and post-mortems directly into DataHub dataset properties, ML model cards, and operational tags.

Architecture & How We Built It

OmniSRE is built as a modular system deployed across Google Cloud Platform and GitHub Actions:

                       ┌─────────────────────────────────────────────────────────────┐
                       │                     Google Cloud Platform                   │
                       │                                                             │
 [ SREs / Users ] ───► │  [ Cloud Run: Web Console ]                                 │
                       │  (Next.js 16 + React 19 + Standalone Turbopack)             │
                       │         │                                                   │
                       │         │ GraphQL & SSE Streaming                           │
                       │         ▼                                                   │
 [ Webhooks:    ] ───► │  [ Cloud Run: OmniSRE API & Agent Engine ]                  │
 dbt / Airflow /       │  (FastAPI + Strawberry GraphQL + APScheduler)               │
 GitHub / SaaS         │         │                 │                │                │
                       │         ▼                 ▼                ▼                │
                       │  [ Secret Manager ]  [ Cloud SQL /    [ Cloud Storage /     │
                       │  (API keys, tokens)   Supabase DB]     Artifact Registry]   │
                       │                           │                                 │
                       │                           ▼                                 │
                       │  [ Compute Engine: DataHub VM ] ◄──► [ Cloud Run: MLflow ]  │
                       │  (GMS Port 8080 + UI 9002)            (Port 5000)           │
                       └─────────────────────────────────────────────────────────────┘

Extensible Plugin Registry

OmniSRE exposes a declarative plugin system (agent/registry.py) structured around 4 decorator interfaces:

# 1. Anomaly Detection
@detector
class FreshnessDetector(Detector): ...

# 2. Evidence Gathering
@probe
class ColumnLineageProbe(Probe): ...

# 3. Reversible Remediation
@actuator(ActionType.QUARANTINE_ROWS)
class QuarantineActuator(Actuator): ...

# 4. Post-Mitigation Verification
@check(CheckType.DBT_ASSERTIONS)
class DbtCheckRunner(CheckRunner): ...

Core Stack

  • Metadata & Knowledge Graph: DataHub (GMS, Elasticsearch index, MySQL aspect store, REST.li / GraphQL APIs, Metadata Change Proposals).
  • Backend Engine: Python 3.11, FastAPI, Strawberry GraphQL, Uvicorn, APScheduler.
  • Data Warehouse & Transformations: DuckDB, dbt (22 custom assertion suites).
  • ML Infrastructure: MLflow Tracking Server, Scikit-Learn, Pandas, NumPy.
  • Multi-Tier Model Routing (OpenRouter):
    • OPENROUTER_MODEL_CODE (cohere/north-mini-code:free): Generates syntax-verified AST migrations and PR patches.
    • OPENROUTER_MODEL_REASONING (nvidia/nemotron-3-ultra-550b-a55b:free): Synthesizes causal evidence bundles into root cause explanations.
    • OPENROUTER_MODEL_CHAT (google/gemma-4-31b-it:free): Powers the interactive SRE incident copilot.
    • OPENROUTER_MODEL_FAST (nvidia/nemotron-nano-9b-v2:free): Low-latency incident triage and classification.
    • Deterministic fallback engine guarantees core remediation functions even during LLM API outages.
  • Frontend Web Console: Next.js 16 (App Router), React 19, Tailwind CSS, Lucide Icons, Supabase Auth.
  • CI/CD & Cloud Infrastructure: Google Cloud Run, Compute Engine, Secret Manager, and 5 GitHub Actions workflows authenticated via Google Workload Identity Federation (WIF).

Live Demo & Testing

Live Cloud Deployment

The full stack is deployed and available to test on Google Cloud Platform: 👉 https://sentinel-web-55utjgkpwq-uc.a.run.app

Test Walkthrough

  1. Trigger a Chaos Fire Drill:
    • Navigate to /incidents.
    • Click Create Incident (Chaos Simulator).
    • Choose any of the 13 scenarios (e.g., schema_change, label_leakage, stale_feed, training_serving_skew).
    • Click Launch Fire Drill to observe the 9-Stage Remediation Stepper and real-time SSE terminal logs stream the detection, RCA, policy evaluation, action execution, and validation phases.
  2. Inspect Blast Radius & Lineage:
    • Navigate to /blast-radius to inspect how an upstream mutation in raw_transactions cascades through features to the fraud_detection_model and downstream scoring services.
  3. Test API Migration & Code Patches:
    • Navigate to /api-health to trigger dependency checks, review breaking changes, and preview generated unified diff PR patches.
  4. Interactive SRE Copilot:
    • Navigate to /chat to query incident context, blast radius estimates, or generate tailored post-mortem reports.

Quantitative Evaluation & Benchmarks

We implemented a reproducible evaluation suite in evaluation/ to benchmark OmniSRE against standard data monitoring approaches:

# Run the evaluation suite
python evaluation/run_all_benchmarks.py

Benchmark Results

Evaluation Suite Baseline Approach OmniSRE Result Key Finding
1. Autonomous Codefix & Migration 12.5% (Passive monitoring) 87.5% pass rate Accurately generates syntax-valid AST patches and verified PR diffs across breaking API signatures.
2. Silent ML Drift & Label Leakage 16.7% (Standard dbt assertions) 91.7% detection rate ($F_1 = 0.92$) Flags subtle distribution shifts, unit scale bugs, and label leakage where standard SQL assertions pass undetected.
3. Multi-Hop Lineage RCA 25.0% (Flat unlinked alerts) 88.9% Top-1 accuracy (94.4% Top-3) Accurately isolates root causes across 4-hop DAGs while compressing cascading alerts by 73.3%.

Challenges We Overcame

  1. Handling Silent Pipeline Failures: Static assertions fail when data isn't rewritten or when values shift within valid ranges (e.g., an upstream pricing bug that doubles transaction values but stays positive and non-null). We addressed this by implementing baseline distribution snapshots (.sentinel/baselines.json) and statistical distance checks that flag deviations before corrupted records reach downstream feature stores.

  2. Cross-Boundary Lineage Resolution: Tracing an issue from a DuckDB raw table through dbt models to an MLflow champion run and a live scoring endpoint required bridging disparate systems. We built a unified Graph Context Tool that queries DataHub's GraphQL API to recursively resolve entity URNs across datasets, data jobs, and ML model entities.

  3. Safe, Reversible Mitigations: Automated remediation carries risk if a fix is applied incorrectly. We implemented an append-only Action Journal backed by SQLite. If a mitigation fails its post-execution invariant check, OmniSRE automatically runs the corresponding inverse operation, reverting the system cleanly.

  4. Dynamic Cross-Service Discovery in Cloud Run: Deploying microservices across Google Cloud Run and Compute Engine required eliminating hardcoded hostnames. We configured the CI/CD deployment workflows to dynamically query GCP Cloud Run endpoints and inject live service URLs during the build and container startup phases.

Accomplishments We're Proud Of

  • Robust Test Suite: Comprehensive coverage across all detectors, probes, actuators, AST diff generators, and GraphQL resolvers.
  • Full Closed-Loop Remediation: From anomaly injection and lineage traversal to policy gating, AST code repair, and DataHub graph writeback.
  • Resilient Multi-Model Routing: Integrated specialized LLMs for code synthesis, causal reasoning, and triage with deterministic offline fallbacks.
  • Keyless Enterprise CI/CD: 5 automated GitHub Actions workflows running with Google Workload Identity Federation (WIF) and Secret Manager, leaving zero credentials stored in code.

What's Next

  • Multi-Warehouse Connectors: Expanding actuators to Snowflake, Google BigQuery, and Databricks Unity Catalog.
  • Slack Socket Mode Bot: Enabling one-click approval workflows directly in Slack channels ([Approve PR], [Extend Quarantine], [Trigger Retrain]).
  • Streaming Pipeline Pre-Emption: Applying predictive drift detection to Kafka topics to trip circuit breakers before bad records land in warehouse tables.

Built With

Share this project:

Updates

Submission history