ModelGuard: DataHub Production ML Agent

A production ML agent that uses DataHub lineage, metadata and ownership context to detect upstream risks, trace affected models and deployments, and write actionable findings back to DataHub.

Comment

Inspiration

Production machine-learning failures often begin with changes that appear harmless: a missing boundary check, an altered schema, an unexpected null value or a feature transformation that behaves differently at the edge of its valid range.

A conventional continuous integration pipeline can usually tell an engineering team that a test failed or that model performance declined. However, it rarely explains what caused the failure, which upstream data assets were involved, whether there are competing explanations or whether a proposed correction genuinely restores the model.

This gap inspired ModelGuard: DataHub Production ML Agent.

I wanted to explore a practical question:

What if a CI/CD pipeline did not merely detect an ML regression, but traced the likely cause through DataHub, proposed the smallest defensible repair, validated that repair independently and returned a review-ready result?

The objective was not to create an autonomous system that silently changes production code. Instead, I designed ModelGuard to behave like a careful machine-learning platform engineer: collect evidence first, compare competing explanations, act only within explicit boundaries, verify every important claim and keep the final decision with a human reviewer.

DataHub provides the foundation for this approach because model failures rarely exist in isolation. Understanding a regression may require information about datasets, fields, schemas, ownership, quality signals, transformation jobs, ML features, models, deployments and upstream or downstream lineage.

ModelGuard brings this operational context together in one evidence-backed investigation.

What it does

ModelGuard is a metadata-aware production ML agent that handles a failed model-quality gate through six guarded phases:

  1. Detect
  2. Collect context
  3. Diagnose
  4. Propose a repair
  5. Validate
  6. Publish

1. Detect the regression

ModelGuard first compares a candidate model metric with an approved baseline using a deterministic policy.

In the demonstration, the model’s F1 score falls from:

0.842 → 0.771

The decline is:

$$ \Delta F_1 = 0.771 - 0.842 = -0.071 $$

Because this drop exceeds the configured tolerance, ModelGuard stops the pipeline and creates a structured regression artefact.

The decision is made by a reproducible numerical rule. Agent reasoning cannot override the metric gate or decide that a measurable regression should be ignored.

2. Collect DataHub context

ModelGuard then retrieves the metadata needed to understand the failure.

This may include:

  • dataset and model metadata;
  • schema fields and affected variables;
  • ownership information;
  • data-quality signals;
  • upstream and downstream lineage;
  • transformation jobs;
  • ML features;
  • model groups;
  • model deployments; and
  • previous incident information.

The context layer supports three providers:

  • the DataHub Python SDK;
  • the DataHub MCP Server over Streamable HTTP; and
  • deterministic DataHub-shaped fixtures for public evaluation and CI.

The public demonstration uses deterministic fixtures so that judges can reproduce it without credentials. The project has also been verified separately against a running DataHub Core instance through both the SDK and the official self-hosted MCP Server.

3. Diagnose the likely cause

ModelGuard extracts evidence and produces several plausible root-cause hypotheses.

It does not simply accept the first explanation that sounds reasonable. Each hypothesis is evaluated using:

  • temporal proximity;
  • lineage relevance;
  • affected fields;
  • ability to explain the metric decline;
  • quality corroboration;
  • supporting evidence; and
  • counter-evidence.

In the demonstration, the leading hypothesis is:

feature_transformation
confidence score: 1.0000
confidence level: high

The evidence shows that a changed feature calculation produced 37 infinite values. The original source rows remained unchanged, which weakens the broader explanation that the regression was caused entirely by poor source-data quality.

The changed transformation is therefore more convincing because it lies on the relevant lineage path and directly explains the newly invalid model inputs.

When evidence is weak or two explanations are too close, ModelGuard abstains rather than forcing a confident diagnosis.

4. Propose a constrained repair

After a high-confidence diagnosis passes the policy gate, ModelGuard proposes the smallest permitted repair.

The demonstrated correction adds two lines to one cited function:

 def calculate_monthly_spend(
     total_spend: float,
     account_age_months: int,
 ) -> float:
+    if account_age_months <= 0:
+        return 0.0
     return total_spend / account_age_months

The repair system does not have unrestricted authority to rewrite the repository.

It enforces limits on:

  • the permitted file;
  • file type;
  • number of modified files;
  • number of changed lines;
  • approved repair strategy;
  • protected paths; and
  • unsafe tokens or operations.

Protected areas include ModelGuard’s internal implementation, CI configuration, scripts and policy files.

The current repair catalogue contains one deliberately narrow strategy, guarded_division, implemented using Python abstract syntax tree analysis.

5. Validate independently

The proposed patch is never applied directly to the original source workspace.

ModelGuard creates a temporary copy and then:

  • applies the proposed patch;
  • compiles the changed code;
  • runs allow-listed test commands;
  • repeats the original model evaluation;
  • checks the post-repair data;
  • verifies the expected repair identifier; and
  • compares workspace hashes.

In the demonstrated incident:

Targeted tests:       3 passed
Invalid values:       37 → 0
F1 score:             0.771 → 0.842
Source workspace:     unchanged

The patch is withheld when tests fail, the metric is not restored or the source workspace changes unexpectedly.

6. Publish once

After successful validation, ModelGuard produces a review-ready GitHub output and records the incident lifecycle in DataHub.

Stable publication identifiers and hidden delivery markers make this process idempotent. Repeating the same workflow returns noop instead of creating another comment or duplicate incident.

The deterministic demonstration reports:

GitHub first/repeat:    created / noop
DataHub first/repeat:   raised_and_resolved / noop

The live DataHub verification also confirmed idempotent incident write-back. Its final committed evidence reused the existing resolved incident and returned noop rather than creating a duplicate.

The complete controlled loop is:

$$ \text{Detect} \rightarrow \text{Contextualise} \rightarrow \text{Diagnose} \rightarrow \text{Repair} \rightarrow \text{Validate} \rightarrow \text{Publish} $$

ModelGuard never approves or merges its own proposed repair. Human review remains the final authority.

How we built it

I built ModelGuard as a Python 3.11+ command-line application with typed data structures, explicit exit statuses and stable JSON contracts between each phase.

This architecture makes every decision inspectable and allows each stage to run independently within a CI environment.

Deterministic evaluation

The evaluation layer performs the model-performance check before agent-style diagnosis begins.

This separation is important because a language model or reasoning component should not decide whether a measurable regression occurred when a reproducible numerical policy can make that decision more reliably.

Provider-neutral DataHub context

The context layer uses a shared internal representation.

It can retrieve information through:

  • the DataHub Python SDK;
  • the DataHub MCP Server; and
  • deterministic fixtures.

The SDK and MCP integrations expose information through different response structures. ModelGuard normalises them into the same context snapshot while retaining provider provenance, original URNs and diagnostic warnings.

This allows the diagnosis engine to work without becoming tightly coupled to one DataHub access method.

Live ML metadata graph

For live verification, ModelGuard creates an idempotent metadata graph containing:

raw customer dataset
        ↓
feature-engineering data job
        ↓
feature dataset
        ↓
training data job
        ↓
training dataset
        ↓
ML features and feature table
        ↓
churn-model-v3
        ↓
churn-api-prod

The graph uses stable URNs so that repeated runs update the same assets rather than creating randomly named duplicates.

The definitive live harness:

  • checks whether DataHub GMS is available;
  • reuses an existing local instance when healthy;
  • starts DataHub quickstart when necessary;
  • loads the ML metadata graph;
  • collects training-data and model context through the SDK;
  • starts the official self-hosted DataHub MCP Server 0.6.0;
  • repeats context and lineage collection through MCP;
  • verifies supported lineage and deployment relationships;
  • performs live incident write-back twice;
  • checks idempotency; and
  • promotes reviewed, sanitised evidence to examples/.

The harness never runs datahub docker nuke, so it does not automatically delete existing DataHub data.

Diagnosis engine

The diagnosis engine separates:

  • evidence extraction;
  • hypothesis generation;
  • hypothesis ranking; and
  • confidence policy.

Evidence receives stable identifiers, allowing every conclusion to point back to the observations supporting it.

The ranking process considers both supporting and contradictory information. This makes the diagnosis more defensible than selecting an explanation solely because it has some supporting signals.

Constrained repair generation

I deliberately rejected an unrestricted “rewrite the code” design.

ModelGuard currently uses a named guarded_division strategy supported by Python AST analysis.

This approach is narrower than general code generation, but it provides clearer guarantees about:

  • what may be changed;
  • where it may be changed;
  • why the change is relevant;
  • how large the patch may be; and
  • how the proposed result will be tested.

Isolated validation

The validation layer creates a temporary copy of the target workspace.

It performs:

  • patch application;
  • Python compilation;
  • allow-listed tests;
  • metric re-evaluation;
  • data checks; and
  • source-workspace hash verification.

Validation commands use shell=False and an allow-listed Python executable.

GitHub and DataHub publication

The publication layer supports:

  • GitHub review reporting; and
  • DataHub incident lifecycle write-back through GraphQL.

Publication remains a dry run unless explicit apply permission is provided.

Stable identifiers prevent repeated CI executions from creating duplicate outputs.

Public demonstration and CI

I also built a dependency-free interactive demonstration using HTML, CSS and JavaScript and deployed it through GitHub Pages.

The website is a read-only replay of verified fixture artefacts. It is not presented as a browser-based live DataHub client.

GitHub Actions runs:

  • Ruff linting;
  • automated tests;
  • the complete six-phase showcase;
  • result verification;
  • submission-package checks; and
  • live integration contract checks.

The deterministic workflow can be reproduced with one command:

python scripts/run_showcase.py

The definitive live workflow can be run with:

python scripts/run_live_datahub_verified.py \
  --install-mcp-server \
  --promote

Challenges we ran into

Deciding where reasoning should end

The most difficult design challenge was deciding where agent reasoning should end and deterministic control should begin.

A plausible diagnosis does not automatically justify modifying code. A passing unit test does not prove that model performance has recovered. A restored metric does not automatically authorise publication.

I addressed this by separating detection, diagnosis, repair, validation and publication into independent gates.

Each gate produces its own artefact and must satisfy its own policy before the workflow may continue.

Normalising different DataHub providers

The Python SDK, MCP Server and deterministic fixtures expose metadata through different interfaces and response structures.

Without a shared representation, the diagnosis engine would become tightly coupled to one provider.

I created a provider-neutral context snapshot that preserves:

  • normalised entities;
  • schema fields;
  • ownership;
  • quality signals;
  • upstream and downstream assets;
  • original URNs;
  • provider details; and
  • warnings.

Distinguishing lineage from named relationships

One important live-integration challenge involved the model-to-deployment relationship.

The tested DataHub MCP Server get_lineage response exposed lineage edges, but the model-to-deployment association was stored as the named MLModelProperties.deployments relationship.

That relationship was not exposed through get_lineage by the tested combination of DataHub Core 1.5 and MCP Server 0.6.0.

Rather than disguising this difference, the final verifier records the capability boundary explicitly:

  • MCP proves server connectivity;
  • MCP proves training-data context;
  • MCP proves model context;
  • MCP proves upstream and downstream ML lineage; and
  • the exact deployment relationship is verified through live SDK model metadata from the same DataHub graph.

This allowed the project to remain accurate without weakening the supported checks.

Ranking competing explanations

Root-cause analysis required more than collecting supporting evidence.

For example, zero-valued source records might suggest poor source-data quality. However, the harmful behaviour begins only when a changed transformation converts those values into infinite model inputs.

Counter-evidence helped ModelGuard distinguish between unusual source values and the transformation change that made them harmful.

Safe repair generation

A broadly autonomous coding agent could support more failure types, but it would also introduce more risk and make validation harder to interpret.

I chose constrained, named repair strategies instead.

This reduces flexibility in the current version, but it gives reviewers clearer guarantees about what the agent can and cannot change.

Reliable live idempotency

CI workflows can run repeatedly because of retries, new commits or infrastructure interruptions.

Without idempotency, an agent could flood a pull request with repeated comments or create several DataHub incidents for one failure.

Live publication also required accounting for the delay between writing an incident and its becoming visible through DataHub search.

Stable identifiers, visibility checks and delivery markers made repeated publication predictable and auditable.

Accomplishments that we're proud of

I am particularly proud that ModelGuard demonstrates a complete investigation and recovery workflow rather than presenting metadata retrieval, diagnosis or code generation as isolated features.

The deterministic showcase:

  • detects a genuine metric regression;
  • reconstructs relevant DataHub-shaped context;
  • produces competing hypotheses;
  • identifies a high-confidence transformation failure;
  • proposes a two-line correction in one file;
  • validates the patch outside the original workspace;
  • passes three targeted tests;
  • removes all 37 invalid values;
  • restores F1 from 0.771 to 0.842;
  • leaves the source repository unchanged; and
  • publishes idempotently.

ModelGuard also abstains when evidence is insufficient or ambiguous.

This matters because a trustworthy production agent should be able to say that it does not have enough evidence rather than manufacturing confidence.

Another major accomplishment is the completed live DataHub verification.

The project has now been tested against a running DataHub Core instance using:

  • the DataHub Python SDK;
  • the official self-hosted DataHub MCP Server 0.6.0;
  • Streamable HTTP;
  • live ML metadata and lineage;
  • the MLModelProperties.deployments relationship;
  • GraphQL incident write-back; and
  • repeated publication for idempotency.

The committed live summary confirms that all supported checks passed:

{
  "sdk_provider_verified": true,
  "mcp_provider_verified": true,
  "sdk_ml_lineage_verified": true,
  "mcp_ml_lineage_verified": true,
  "mcp_model_context_verified": true,
  "sdk_model_deployment_link_verified": true,
  "datahub_model_deployment_link_verified": true,
  "live_datahub_writeback_verified": true,
  "live_datahub_writeback_idempotent": true
}

The repository also contains:

  • a public interactive demonstration;
  • a demonstration video;
  • automated CI;
  • comprehensive tests;
  • deterministic sample outputs;
  • sanitised live evidence;
  • judging documentation;
  • testing instructions;
  • a live-verification guide; and
  • a credential-free one-command showcase.

What we learned

The most important lesson was that context and authority should remain separate.

DataHub gives ModelGuard the context needed to understand relationships between datasets, transformations, features, models, owners and deployments.

However, access to rich context should not automatically grant authority to change code or publish an incident.

Deterministic policies must decide whether the available evidence is strong enough to continue.

I also learned that capability claims must match the interface that actually provided the evidence.

The DataHub SDK and MCP Server both contributed valuable information, but they did not expose every relationship in exactly the same way. Recording that distinction made the final verification more trustworthy than forcing both providers to support an identical claim.

Counter-evidence proved to be as important as supporting evidence.

Root-cause analysis becomes more reliable when the agent actively searches for observations that weaken its preferred explanation. This reduces the risk of selecting a convenient hypothesis simply because it has some supporting signals.

A further lesson was that a small, well-defined repair can be more valuable than an impressive but unrestricted code-generation feature.

The two-line correction is useful because ModelGuard can:

  • explain why it is required;
  • identify the relevant function;
  • show the exact patch;
  • test it independently; and
  • prove that it restores the metric policy.

The project also reinforced the importance of reproducibility.

Stable JSON artefacts, deterministic fixtures, idempotent live metadata, committed evidence and one-command workflows make the system easier to inspect, evaluate and debug.

Finally, I learned that a production agent should optimise not only for successful action but also for safe inaction.

Abstaining, withholding an unvalidated patch, preserving the original workspace and avoiding duplicate publication are meaningful outcomes when evidence or permissions are insufficient.

What's next for ModelGuard: DataHub Production ML Agent

The next stage is to expand ModelGuard’s constrained repair catalogue while preserving the same safety principles.

Planned repair strategies include:

  • schema-compatibility corrections;
  • null-handling repairs;
  • category-mapping issues;
  • selected feature-drift responses;
  • type-conversion safeguards; and
  • additional bounded numerical transformations.

I also plan to test the live workflow across a broader matrix of DataHub versions, managed MCP endpoints and authenticated environments.

This will help distinguish stable integration behaviour from version-specific capability boundaries.

Future versions will introduce policy-as-code controls for different organisational risk levels.

For example:

  • a low-risk internal model might permit broader automated validation;
  • a customer-facing production model might require engineering approval; and
  • a healthcare or financial model might require several formal approvals before a repair can be proposed or published.

Signed provenance is another priority.

Generated patches, evidence reports, validation results and publication receipts should be cryptographically linked so that reviewers can confirm which inputs, policies and source revisions produced a recommendation.

Additional integrations could connect validated findings to:

  • pull-request checks;
  • incident-management platforms;
  • model-governance dashboards;
  • deployment approval systems; and
  • metadata-based operational workflows.

The long-term objective is not to remove engineers from production machine-learning decisions.

It is to reduce the time they spend reconstructing failures from scattered logs, disconnected metadata and repeated manual tests.

Built With

Share this project:

Updates