DEV Community

Cover image for Can an AI Agent Fix a Failed Data Pipeline? Build a Safe Recovery Workflow
Rahul R
Rahul R

Posted on

Can an AI Agent Fix a Failed Data Pipeline? Build a Safe Recovery Workflow

A data pipeline fails at 2 AM.

An AWS Glue job runs out of memory. Step Functions marks the execution as failed. CloudWatch has the logs, but someone still needs to investigate what happened, identify the cause, decide what to do next, and rerun the job safely.

Now imagine an AI agent helping with that process.

It checks the failed execution, examines relevant logs, identifies a likely cause, recommends a recovery action, and asks for approval before making a production change.

Sounds useful, right?

But there is one important difference between building an AI agent that can explain a failure and building one that can safely recover a production pipeline.

An agent should not get unrestricted access to fix whatever it thinks is broken.

In data engineering, a wrong recovery action can be worse than the original failure. It might duplicate records, overwrite valid data, trigger unnecessary costs, or hide a data quality problem.

So, instead of giving an AI agent complete control, let's design a workflow where it investigates failures, recommends fixes, and executes only approved, validated actions.

1. Why Traditional Pipeline Recovery Takes Time

Consider a typical AWS data pipeline:

Source System
     |
     v
Amazon S3 - Raw Data
     |
     v
AWS Step Functions
     |
     v
AWS Glue Job
     |
     v
S3 - Curated Data
     |
     v
Downstream Analytics
Enter fullscreen mode Exit fullscreen mode

CloudWatch collects execution logs and metrics, while Step Functions manages the workflow.

When a job fails, an engineer usually needs to answer several questions:

  • Which task failed?
  • Was the failure caused by bad input, a schema change, permissions, or insufficient resources?
  • Did the job write any data before failing?
  • Is it safe to retry?
  • Will retrying duplicate records or overwrite existing output?
  • Does the problem require a code change or human intervention?

Some failures are straightforward. Others require checking multiple logs, configurations, and data dependencies.

This is where an AI agent can help—not by replacing the entire recovery process, but by reducing the time spent collecting evidence and identifying the next action.

2. What Should an AI Recovery Agent Actually Do?

I would divide the workflow into four responsibilities.

Observe: Collect the failed execution details, relevant CloudWatch logs, job parameters, and available metrics.

Diagnose: Use an AI model to interpret that evidence and identify likely causes.

Recommend: Suggest an action that matches an approved recovery policy.

Recover: Execute an authorized action only after the required checks and approvals succeed.

The distinction matters because an AI-generated diagnosis is a hypothesis, not proof.

For example, a Glue job might fail because of memory pressure, but blindly increasing the worker count could waste money without fixing the underlying problem. The job might actually be loading too much data into memory because of an inefficient transformation.

A good recovery workflow checks the evidence before acting.

3. Architecture: AI-Assisted Recovery on AWS

Here is one possible architecture using AWS services.

AWS Glue Job Fails
        |
        v
AWS Step Functions
        |
        v
CloudWatch Logs + Job Metadata
        |
        v
Recovery Coordinator
        |
        v
AI Diagnosis
        |
        v
Validate Against Recovery Policy
        |
        v
Human Approval (when required)
        |
        v
Controlled Recovery Action
        |
        v
Retry + Data Quality Checks
        |
        v
Success or Escalation
Enter fullscreen mode Exit fullscreen mode

The recovery coordinator can be implemented using AWS Lambda, while the AI reasoning layer can use an approved model through Amazon Bedrock or another suitable model service.

Step Functions can orchestrate the workflow, including retries, timeouts, approval waits, and escalation paths.

The architecture deliberately separates diagnosis from execution. The model can recommend an action, but deterministic application logic decides whether that action is allowed.

4. Start with Failure Classification

Before asking an AI model to diagnose everything, define a small set of failure categories.

Failure category Possible recovery
Temporary service or network error Retry with bounded backoff
Throttling Retry with backoff and jitter
Invalid input records Quarantine records or request source correction
Schema mismatch Stop and investigate the contract change
Out-of-memory failure Inspect workload and resource metrics
Permission denied Escalate for an authorized IAM fix
Target write conflict Check concurrent writers and retry safety

This table should be treated as a policy starting point, not a universal mapping from error message to solution.

For example, an out-of-memory error may require a code optimization, a partitioning change, or more resources. The correct action depends on evidence from the workload.

Similarly, a permission error should not cause an agent to automatically grant broader IAM permissions.

5. Build a Small Failure Diagnosis Component

Let's begin with a simple Python component that prepares failure information for an AI model.

In a production system, the recovery coordinator would retrieve these values from trusted AWS APIs and relevant log records.

failure_context = {
    "pipeline": "customer_ingestion",
    "service": "AWS Glue",
    "job_run_id": "jr_example",
    "failed_stage": "transform_customer_data",
    "error_type": "RESOURCE_EXHAUSTED",
    "error_message": "Executor lost during a large shuffle",
    "retry_count": 1,
    "target_write_started": False
}
Enter fullscreen mode Exit fullscreen mode

These values are illustrative. The actual error type and write status must be derived from execution records and pipeline instrumentation.

Next, create a prompt that asks the model to analyze the evidence without giving it authority to execute a fix.

diagnosis_prompt = f"""
You are assisting with AWS data pipeline incident diagnosis.

Analyze the supplied failure context.

Return:
1. Likely cause
2. Evidence supporting the diagnosis
3. Alternative explanations
4. Recommended next investigation
5. Suggested recovery action
6. Risks of retrying
7. Whether human approval is required

Do not claim certainty when evidence is insufficient.
Do not invent logs, metrics, or execution results.
Do not execute recovery actions.

Failure context:
{failure_context}
"""
Enter fullscreen mode Exit fullscreen mode

This is a starting point for the diagnosis layer, not a complete production integration with a model API.

For a real implementation, pass structured, size-limited evidence to the model, validate the returned response against a schema, and avoid including secrets, credentials, or unnecessary sensitive data in prompts.

Logs should also be treated as untrusted input. A malicious record or error message could contain instructions designed to manipulate the model. Those instructions must never override the recovery policy.

6. Make the Model Return Structured Recommendations

Free-form text is useful for explanations, but it is a poor interface for an automated recovery controller.

Instead, require a structured response such as:

{
  "category": "resource_pressure",
  "confidence": 0.86,
  "evidence": [
    "Executor failures occurred during a large shuffle"
  ],
  "recommended_action": "inspect_spark_metrics",
  "requires_approval": true
}
Enter fullscreen mode Exit fullscreen mode

This is an illustrative response, not a real diagnosis or a measured confidence score.

The application should validate the response before doing anything with it.

For example:

  • Is the failure category recognized?
  • Is the action on an explicit allowlist?
  • Does the recommendation match the evidence?
  • Is the action allowed for this environment?
  • Could the action overwrite data or increase costs?
  • Is approval mandatory?

The confidence value is only a model-generated estimate. It is not a calibrated probability that the recommendation is correct, and it must never be the only basis for a production action.

7. Put a Deterministic Policy Between AI and Execution

This is the most important part of the design.

The AI agent should not directly execute arbitrary shell commands, modify IAM policies, change infrastructure, or launch unrestricted production jobs.

Instead, expose a small set of predefined actions.

For example:

ALLOWED_ACTIONS = {
    "inspect_spark_metrics",
    "retry_transient_failure",
    "request_human_approval",
    "escalate_to_engineer"
}


def validate_recommendation(recommendation):
    action = recommendation.get("recommended_action")

    if action not in ALLOWED_ACTIONS:
        return {
            "approved": False,
            "reason": "Action is not allowlisted"
        }

    if action == "retry_transient_failure":
        return {
            "approved": False,
            "reason": "Retry safety checks are still required"
        }

    return {
        "approved": True,
        "action": action
    }
Enter fullscreen mode Exit fullscreen mode

Notice that even an allowlisted retry is not automatically authorized. The example intentionally leaves retry authorization to a separate safety check.

Before retrying, the controller should verify that the error is retryable, the retry budget has not been exhausted, and the target write is safe to repeat.

This separation makes the system easier to test and audit.

8. Design Retries Without Creating Duplicate Data

Suppose a Glue job writes 100,000 records to S3, but the workflow fails before recording that the write completed.

The recovery process sees a failed execution and starts the job again.

What happens to those 100,000 records?

The answer depends on the target design.

An append-only write may create duplicate records. An overwrite may destroy previously valid output. A transactional table with a suitable merge operation may support safer retries, but only if the write logic and transaction semantics are designed correctly.

A reliable recovery strategy should include:

Idempotency: Repeating the same logical operation should not produce unintended additional effects.

Batch identification: Give each input batch or processing interval a stable identifier.

Checkpoint discipline: Advance progress only when the corresponding output is safely committed.

Bounded retries: Limit attempts and use backoff for transient failures.

Post-recovery validation: Check expected record counts, duplicate keys, data quality rules, and reconciliation results.

A simplified retry policy might look like this:

def should_retry(category, retry_count, max_retries=3):
    retryable_categories = {
        "transient_network_error",
        "service_throttling"
    }

    if category not in retryable_categories:
        return False

    return retry_count < max_retries
Enter fullscreen mode Exit fullscreen mode

This is intentionally conservative. Real error classification should use structured service errors and workflow state, not rely only on text generated by an AI model.

A production retry policy should also include exponential backoff, jitter, execution timeouts, and a clear escalation path.

9. Where Human Approval Fits

Not every action needs the same level of control.

I would group recovery actions into three categories.

Low risk: Collect more logs, inspect metrics, or prepare a diagnosis. These actions can often run automatically under read-only permissions.

Moderate risk: Retry a known transient failure when idempotency checks pass and the retry budget permits it. This can be automated under a narrowly defined policy.

High risk: Change production resources, modify schemas, replay historical data, alter IAM permissions, or overwrite target datasets. These actions should require explicit authorization and, where appropriate, peer review.

Approval should be tied to the exact action, target, and execution context—not a general instruction to let the agent fix the pipeline.

The approval record should capture who approved the action, what was approved, when it was approved, and whether the action's parameters changed afterward.

If the target or action changes, the old approval should no longer authorize the new operation.

10. Prevent the Recovery Agent from Becoming Another Failure Source

An AI recovery system introduces its own operational risks.

A model may misinterpret an error, receive incomplete logs, or produce an unsafe recommendation. Its service may also become unavailable.

Design for these possibilities.

  • Give diagnostic components read-only permissions wherever possible.
  • Use separate, narrowly scoped roles for execution actions.
  • Never expose AWS credentials or secrets to the model.
  • Limit log size and remove unnecessary sensitive fields.
  • Treat all retrieved log content as untrusted data.
  • Validate every proposed action outside the model.
  • Record the evidence, recommendation, policy decision, approval, and execution result.
  • Limit retries, execution duration, and potential resource consumption.
  • Fall back to the normal incident process if the AI service is unavailable.

The recovery mechanism should fail safely. If the model cannot produce a valid diagnosis, the pipeline should follow its established failure and escalation path rather than improvising a fix.

11. Measure Whether AI Actually Helps

Adding an AI agent does not automatically make a data platform more reliable.

Compare the workflow before and after introducing the agent.

Useful metrics include:

  • Mean time to diagnose a pipeline failure.
  • Mean time to recover.
  • Percentage of recommendations accepted by engineers.
  • Percentage of automated retries that succeed.
  • Duplicate or missing records after recovery.
  • Number of unsafe actions blocked by policy.
  • Cost per incident.
  • Number of incidents escalated because evidence was insufficient.

Start in observation mode. Let the agent analyze historical or newly occurring failures without executing changes. Compare its recommendations with the decisions made by engineers.

Only enable limited automation after testing the system against known failures, including misleading logs, malformed model output, repeated failures, and partial target writes.

The objective is not to maximize the number of automatic fixes. It is to reduce recovery time without increasing data correctness or operational risk.

Final Thoughts

AI agents can make data pipeline recovery more efficient, especially when engineers spend too much time gathering logs, comparing execution details, and investigating familiar failure patterns.

But the useful design is not an agent with unrestricted access to production.

It is an agent that gathers evidence, explains likely causes, recommends a bounded action, and works within a recovery policy that engineers can understand and audit.

For AWS data pipelines, Step Functions can coordinate the workflow, CloudWatch can provide execution evidence, Lambda can implement controlled recovery logic, and an AI model can help interpret failures.

Start with diagnosis. Add safe retries only when their behavior is well understood. Require approval for risky changes, and validate the resulting data before declaring recovery complete.

The real achievement is not building an AI agent that can fix everything. It is building one that knows when it has enough evidence to act and when it should stop and ask for help.


What would you automate first in your data platform: failure diagnosis, safe retries, or data quality investigation?

Top comments (0)