DataShield Guardian

DataShield Guardian is an autonomous AI agent that protects ML systems from data leakage, PII, drift, and schema failures—tracing risks through DataHub lineage and taking action before models break.

Comment

Inspiration

Machine learning systems rarely fail because of the model alone. The data feeding them can change, drift, contain sensitive information, or accidentally introduce target leakage. The difficult part is understanding how one upstream problem affects downstream features, models, and production deployments.

We were inspired by the idea of combining DataHub's metadata and lineage graph with autonomous AI agents to create a continuous safety layer for ML systems. Instead of waiting for a model to fail, we wanted to build something that could detect risks early, understand their impact, recommend an action, and preserve the investigation as institutional memory.

This led to DataShield Guardian — an autonomous ML data safety, drift, and lineage protection agent.

What it does

DataShield Guardian protects ML systems across the entire lifecycle:

  • Pre-training safety: Detects potential target leakage, PII/sensitive data, deprecated features, schema violations, and data-quality issues.
  • Post-deployment protection: Detects feature drift, schema drift, null-rate changes, freshness problems, and upstream data-quality degradation.
  • Lineage-aware impact analysis: Uses DataHub lineage to trace a problem from its source through datasets and features to affected ML models and deployments.
  • Risk assessment: Combines multiple signals into a transparent ML risk score with explainable severity levels.
  • Autonomous decisions: Depending on the risk, the system can recommend monitoring, investigation, retraining, rollback, or blocking an unsafe training workflow.
  • Incident memory: Records important findings, affected assets, root causes, and recommended actions back into DataHub.

The core workflow is:

Detect → Investigate → Trace → Assess → Decide → Act → Verify → Write Back

For example, when customer_income experiences significant distribution drift, DataShield Guardian can detect the change, trace the feature through DataHub lineage, identify the production models that depend on it, calculate the resulting risk, recommend retraining, and record the incident.

How we built it

We built DataShield Guardian as a full-stack, DataHub-native agentic system.

The architecture consists of:

  • Frontend: Next.js, React, TypeScript, Tailwind CSS, and interactive data visualizations.
  • Backend: Python and FastAPI.
  • Agent orchestration: Specialized components for data safety analysis, drift detection, lineage impact analysis, risk assessment, decision-making, remediation planning, and incident write-back.
  • ML/Data analysis: Python, pandas, NumPy, scikit-learn, and statistical drift detection techniques such as PSI and KS testing.
  • DataHub: Used as the metadata, lineage, governance, and institutional-memory layer.
  • Persistence: Database-backed incident and application state.
  • Infrastructure: Docker-based development and reproducible synthetic ML data.

We created a realistic synthetic ML environment containing source datasets, feature pipelines, ML models, and production deployments. This allowed us to simulate realistic failures such as target leakage, feature drift, schema changes, PII exposure, and null-rate spikes.

Rather than relying on an LLM for everything, deterministic checks are used for calculations such as drift metrics, schema comparisons, and risk scoring. AI is used where it adds value, including investigation summaries, explanations, and remediation recommendations.

Challenges we ran into

One of the biggest challenges was making DataHub a meaningful part of the system rather than simply displaying lineage.

We had to ensure that the agent could actually use metadata and lineage to understand relationships between source data, features, models, and deployments.

Another challenge was distinguishing between detection and useful action. A simple alert such as "feature drift detected" is not enough. The system needed to determine what was affected and why the problem mattered.

We also had to balance autonomous behavior with reliability. Critical decisions such as target leakage detection, drift calculations, and risk scoring should not depend entirely on an LLM. We therefore separated deterministic analysis from AI-powered explanation and recommendations.

Finally, we had to design a realistic demonstration environment where failures could be injected repeatedly and the complete agent workflow could be observed without depending on real production data.

Accomplishments that we're proud of

We are proud of building a system that connects ML safety, monitoring, lineage, agentic reasoning, and governance into one workflow.

In particular, we demonstrated:

  • End-to-end source-to-model lineage analysis.
  • Automated target leakage detection.
  • Potential PII and sensitive-data detection.
  • Real statistical feature-drift detection.
  • Schema and data-quality monitoring.
  • Lineage-based downstream impact analysis.
  • Transparent ML risk scoring.
  • Policy-based training decisions such as ALLOW, WARN, and BLOCK.
  • Actionable remediation recommendations.
  • Incident tracking and audit history.
  • DataHub write-back so that investigations become persistent context for future engineers and agents.
  • A failure simulation environment for repeatable demonstrations.

The biggest accomplishment is the complete feedback loop:

A problem occurs → the agent detects it → traces its impact → assesses the risk → recommends or takes action → records what happened back into the data graph.

What we learned

We learned that effective agentic data systems are not just about adding an LLM to an existing dashboard.

The most valuable part is giving the agent reliable context, tools, and boundaries.

DataHub provides the context layer that allows the agent to understand relationships that would otherwise require engineers to manually investigate. Deterministic ML and data-quality checks provide trustworthy evidence, while AI makes that evidence easier to reason about and act upon.

We also learned that explainability is essential for autonomous systems. Instead of simply showing a risk score, DataShield Guardian provides the evidence behind the decision and identifies the affected assets.

Most importantly, we learned that metadata can become institutional memory. When an incident and its resolution are written back into the data graph, future engineers and agents can build on previous investigations instead of starting from zero.

What's next for DataShield Guardian

Our next goal is to move from intelligent recommendations toward controlled autonomous remediation.

Future versions could include:

  • Automated model retraining with approval workflows.
  • Automated rollback to a last-known-good model.
  • GitHub pull requests for schema and pipeline fixes.
  • Slack and incident-management integrations.
  • Deeper feature-store and MLOps integrations.
  • Model performance monitoring alongside data drift.
  • Policy-as-code for organization-specific ML governance.
  • Historical incident analysis to identify recurring failure patterns.
  • Cross-model risk propagation across larger enterprise lineage graphs.
  • Human-in-the-loop approval for high-impact autonomous actions.

Our long-term vision is for DataShield Guardian to become a continuous reliability layer for ML systems — protecting models from the moment data enters the pipeline until the model is running in production.

Built With

Share this project:

Updates