Forecast Sentinel

An upstream column changed type. Nothing broke, nothing alerted, and your forecast is now wrong. Forecast Sentinel proves it from DataHub's ML lineage and writes the verdict to your catalog.

Comment

Inspiration

I work in supply chain. The failure that got me here is one I have watched happen: a demand forecast that goes quietly wrong.

Not a crash. An upstream team changes holiday_flag from an integer to a string, because they moved to named holidays. No pipeline fails. No alert fires. The model keeps returning numbers, and planners keep ordering against them. Six weeks later someone notices the forecast error has drifted, and nobody can say when it started or why.

What makes this failure so durable is that every tool in the stack reports green, and each one is correct to do so:

  • Model monitoring watches the endpoint. The endpoint is healthy.
  • Data quality tests watch the table. A string column is a perfectly valid string column.
  • Drift detection watches the prediction distribution. It reports weeks later, after the damage, without saying why.

The link nobody makes is "this column changed" → "that deployed model depends on it". That link is lineage, and DataHub is the one system that already holds it — training data to features to models to deployments. The Production ML Agents challenge is pointed at exactly that gap.

What it does

It answers one question that no dashboard currently answers: this upstream thing changed — is my deployed model still valid?

$ sentinel check "urn:li:mlModel:(urn:li:dataPlatform:mlflow,demand-forecast-v3,PROD)"

╭─ BLOCK — demand-forecast-v3 ────────────────────────────────────────────────╮
│ holiday_flag changed from int to string in the model's primary training      │
│ input; the deployed artefact's learned encoding no longer matches inference  │
│ data, so predictions are silently wrong.                                     │
╰─────────────────────────────────────────────────────────────────────────────╯
lineage: 14 upstream, 6 downstream, depth 3

Risks
  HIGH  Feature encoding mismatch on holiday_flag
    path: raw_sales.holiday_flag → feat_seasonality → demand-forecast-v3

Recommended actions
  1. Stop the nightly batch scoring job before the next 02:00 run.
  2. Pin feat_seasonality to the previous dtype, or retrain on the new encoding.
  3. Notify owners of exec_demand_dashboard — last 3 days of figures suspect.

Written back to DataHub
  ✓ add_tags — tagged urn:li:tag:model-invalidated
  ✓ update_description — updated description with verdict status
  ✓ save_document — saved full findings document

BLOCK verdict — failing with exit code 2.

Two things matter about that output. Exit code 2 means it drops into CI as a gate in front of a scheduled retrain or batch-scoring job. And it writes the verdict back into DataHub — the model gets tagged, its description carries the status, and the full findings document is attached. The next person who opens that model in the catalog sees the finding without being told.

How I built it

The one decision everything else follows from: split the problem in two, because the halves need different tools.

Deterministic half Agentic half
Question What changed? Does it matter?
Implementation Python: bounded BFS + schema diff Claude/Gemini + DataHub MCP tools
Properties Reproducible, free, fast Reasoned, cited, judged

Detecting that a column went from INT to VARCHAR is a diff. Sending a model to do it would be slower, more expensive, and less reliable — the answer would vary between runs. But deciding whether that change invalidates a trained gradient-boosting model requires reading the lineage, understanding how that model consumes that feature, and weighing consequence. That is judgement, and it is the only thing the model is asked for.

The agent is never asked to detect drift. It is handed facts and asked what they mean.

resolve model → walk ML lineage → snapshot training inputs
    → diff vs recorded baseline → [drift?] → agent judges consequence
    → emit_verdict → write back to catalog → exit 2 on BLOCK

It talks to DataHub only through the DataHub MCP server, so the agent calls DataHub's real tools rather than a reimplementation: get_entities, get_lineage, get_lineage_paths_between, list_schema_fields, get_dataset_queries, search for reading, and add_tags, update_description, save_document for writing. Both deployment shapes work from one config layer — a uvx mcp-server-datahub stdio subprocess for self-hosted, or the tenant's streamable-HTTP endpoint for DataHub Cloud.

Stack: Python 3.12, the Anthropic SDK's tool runner with async_mcp_tool bridging the MCP session, a Pydantic verdict contract, a Typer CLI, and a static self-contained HTML report. There is also a Claude Code plugin exposing the same capability conversationally, for investigating rather than gating.

Challenges I ran into

The MCP Python SDK has two incompatible client APIs, and the published docs describe the wrong one. The current docs show a high-level Client facade; anthropic[mcp]'s bridge takes the older 1.x ClientSession. Code written from the documentation fails at the bridge, not at import. I only caught it because a type checker flagged mcp.Client as unknown against the installed package. The fix was to stop trusting docs and introspect the installed library, then pin mcp>=1.8,<2 with the reason in a comment so nobody "modernises" it later.

The DataHub Python SDK cannot express ML lineage. client.lineage.add_lineage accepts dataset, datajob, dashboard and chart URNs — but not mlModel. There is no mlFeatures setter either. So the dataset → feature → model path, which is the exact thing this challenge is about, cannot be built with the high-level SDK. I worked around it in the demo seeder by modelling the feature layer as a materialised dataset and linking the model through a training job.

Nested URNs break naive parsing. A schemaField URN embeds a whole dataset URN, which embeds a dataPlatform URN. split(",") corrupts every one of them. Key arity also varies by entity type across DataHub versions, so the parser tracks parenthesis depth and keeps components positional instead of assuming a shape.

Catalog payloads are inconsistent. Depending on version and platform, keys arrive camelCase or snake_case, lineage entries as objects or bare URN strings, owners as strings or nested dicts. Every parser had to accept all observed shapes and drop what it cannot interpret rather than raising — and every tolerance got a test, so nobody tightens it back later.

Unbounded lineage traversal is a trap. A walk from a central warehouse table enumerates the whole warehouse. Traversal is bounded, and hitting a bound sets a truncated flag that is surfaced in the report — a partial blast radius presented as complete is worse than no blast radius at all.

Hardware. DataHub's quickstart wants 8 GB of RAM allocated to Docker. My machine has 8 GB in total. There is no allocation that satisfies DataHub and leaves macOS able to function, which shaped how the local demo path had to work.

And a small one with a real lesson in it: Pydantic reserves the model_ prefix, so a field called model_urn collides with its protected namespace. I kept the field name and disabled the namespace — DataHub calls the entity an mlModel, and renaming it would have desynced my schema from the catalog's vocabulary for the sake of a warning.

Accomplishments that I'm proud of

Resisting the urge to make it "an AI agent that does everything." The deterministic/agentic split is the design, and it took discipline to keep the diff in code where it belongs.

It writes back. Most catalog integrations read. A read-only lineage report is a nice diagram; persisting the verdict is what makes the next person's search results carry the warning.

Evidence discipline that is structural, not aspirational. Every risk must cite the lineage path it travels and facts retrieved in that session, and anything the agent asserted but could not confirm lands in an explicit unverified_claims field rather than being dropped. An ML-lineage agent that speculates is worse than no agent, because its output gets pasted into an incident channel.

The agent gets read tools only. Write-back happens afterwards in code, from the structured verdict. The model decides what to record; it does not get to decide how many catalog objects to touch.

66 unit tests that need no DataHub instance, covering URN parsing, payload normalisation across every shape I observed, drift diffing, and HTML escaping of catalog strings.

What I learned

Introspect the installed package; don't trust the docs. This cost me the most time and taught me the most. For any fast-moving SDK, dir() and inspect.signature() against what is actually installed beat documentation that describes a different major version.

The right amount of agent is less than you reach for. My first instinct was to let the model do the whole job, discovery through judgement. Pulling detection back into code made the system cheaper, faster, reproducible and easier to test.

Governance findings are not model risks, and a model will happily conflate them. Early runs reported "no owner set on the feature table" at the same severity as "predictions are wrong". I had to write that distinction into the prompt explicitly. A missing owner is a real problem, but it is not evidence that today's forecast is incorrect.

Making the agent state what it could not verify improved usefulness more than any amount of prompt tuning. A verdict that says "I could not confirm whether the pipeline casts this column" is far more actionable than one that quietly assumes an answer.

What's next for Forecast Sentinel

  • Distribution drift. Right now the sentinel reasons about structural change — schema, lineage, ownership — because that is what a catalog knows. Statistical drift needs the data itself, and belongs alongside rather than inside this.
  • DataHub Cloud assertions and incidents. Provision assertions for the gaps it finds, and open an incident on a BLOCK instead of only tagging.
  • Shared baselines. Snapshots are local JSON files today, which is fine for CI with a cached directory. A team deployment wants them in DataHub itself as structured properties, so the baseline lives next to the model.
  • Run it in reverse. The same lineage walk, pointed the other way, becomes a pre-flight check for the data engineer: "this migration will invalidate 3 production models, here are their owners" — before the change ships, not after.
  • Contributing back. Building this surfaced two documentation gaps and one real SDK limitation in DataHub, all worth filing upstream. The MCP client-version ambiguity in particular will cost every Python entrant the same afternoon it cost me.

Built With

Share this project:

Updates