Skip to content

Repository files navigation

🛡️ Aegis: Autonomous Agentic MLOps Command

Aegis is an autonomous, multi-agent operations platform designed to close the loop between anomaly detection, root-cause diagnosis, and remediation across production machine learning stacks (model-serving, GPU infrastructure, and data pipelines).

Aegis acts where it is safe to do so automatically, and coordinates a crisp human-in-the-loop approval handoff where it isn't.


🏗️ Architecture Blueprint

Aegis is composed of four cooperating AI agents built on top of the Google Agent Development Kit (ADK) and powered by Gemini 2.5 Flash:

graph TD
    A[Telemetry Sources] -->|Metric Logs| Sentry[🛡️ Sentry: Monitor]
    Sentry -->|Anomaly JSON| Sleuth[🔍 Sleuth: Diagnose]
    Sleuth -->|Diagnostic JSON| Medic[💊 Medic: Propose]
    
    subgraph "Decision Guardrail"
        Medic --> Command[Remediation Command]
        Command --> Policy{Risk Policy}
        Policy -->|Low Risk| Auto[✓ Auto-Execute]
        Policy -->|Medium/High Risk| Gate[⏳ Approval Gate]
        Gate -->|Approved| Exec[✓ Manual Execute]
        Gate -->|Rejected| Cancel[✕ Cancel]
    end
    
    Auto --> Scribe[✍️ Scribe: Document]
    Exec --> Scribe
    Cancel --> Scribe
    
    Scribe -->|Markdown postmortem| Storage[(ChromaDB Memory & File Docs)]
Loading

The 4 Cooperating Agents

  1. 🛡️ Sentry (Monitoring Agent): Continuously evaluates system telemetry (latency, error rates, GPU utilization, and data drift scores). It computes deviation from historical baselines (with directionality intelligence) and creates structured anomalies.
  2. 🔍 Sleuth (Diagnostic Agent): Correlates metrics, logs, and deployment events across services when Sentry flags an anomaly. It queries ChromaDB vector store for historical occurrences to deliver a root-cause hypothesis with a confidence score.
  3. 💊 Medic (Remediation Agent): Proposes targeted interventions (e.g. pod restarts, model version rollbacks, queue workload rebalancing) and drafts exact command strings.
  4. ✍️ Scribe (Reporting Agent): Generates a complete, publication-quality incident postmortem report in Markdown format capturing timelines, root cause findings, actions taken, and preventative recommendations.

📂 Project Directory Structure

aegis/
├── agents/                      # Core Multi-Agent Logic (ADK)
│   ├── sentry/                  # Monitoring metrics & baseline tools
│   ├── sleuth/                  # Diagnostic log correlation & deploy lookups
│   ├── medic/                   # Remediation proposing & risk tiering policy
│   ├── scribe/                  # Markdown report generator & ChromaDB archiver
│   └── memory/                  # Shared vector store (ChromaDB) & Pydantic schemas
├── backend/                     # FastAPI & gRPC Service Layer
│   ├── api/                     # REST routers (incidents, approvals, metrics)
│   ├── db/                      # SQLite persistence (SQLAlchemy models & sessions)
│   └── grpc_server/             # gRPC streaming server & protobuf schemas
├── simulator/                   # Failure-Injection Harness
│   ├── failure_scenarios/       # Out-of-the-box scenario registries
│   ├── metrics_generator.py     # Realistic telemetry time-series generator
│   └── mock_inference_service.py # FastAPI ML service simulation
├── frontend/                    # Vite + React (TypeScript) Console Dashboard
│   ├── src/components/          # Telemetry charts, trace console, timelines, modals
│   └── src/hooks/               # Real-time Server-Sent Events (SSE) stream hooks
├── infra/                       # Containerized infrastructure
│   ├── Dockerfile.backend       # Multi-stage optimized builder backend container
│   ├── Dockerfile.frontend      # Multi-stage nginx-served frontend container
│   ├── Dockerfile.simulator     # Lightweight simulator container
│   └── docker-compose.yml       # Production-grade service orchestrator
└── tests/                       # Automated testing suite (unit + integration)

⚡ Quick Start

1. Prerequisites & Environment Setup

Ensure you have Python 3.11+ and Node.js 20+ installed. Clone the repository and configure your environment key:

# Clone the repository
git clone https://github.com/Debddj/Aegis.git
cd aegis

# Create and configure .env in project root
cp .env.example .env
# Edit .env and paste your GOOGLE_API_KEY

Using uv for lightning-fast environment setup:

# Install uv if missing
pip install uv

# Initialize environment and install dependencies
uv venv
source .venv/bin/activate  # Or .venv\Scripts\activate on Windows
uv pip install -e ".[dev]"

2. Run the Stack (Docker Compose)

Aegis is fully containerized. Start the entire system in one command using the -p aegis namespace to prevent conflicts with other projects named infra:

docker compose -p aegis -f infra/docker-compose.yml up -d

This launches the following services in detached mode:

  • Simulator on http://localhost:8100 (Mock ML service)
  • FastAPI Backend on http://localhost:8000 (Core REST APIs + SSE Event Streams)
  • gRPC Server on port 50051 (Streaming metrics/traces)
  • Vite React Console on http://localhost:5173 (Obsidian-themed operator dashboard)

You can check container status and logs with:

# Get container status and port mappings
docker compose -p aegis -f infra/docker-compose.yml ps

# Follow logs from the backend
docker logs -f aegis-backend-1

🧪 Testing and Quality Gates

Aegis enforces a strict quality and lint standard. To run the complete verification suite:

# Run unit and integration tests (Real Gemini API validation)
.venv/Scripts/pytest -v

# Run the Ruff linter check
.venv/Scripts/ruff check .

🎮 Failure Injection Demo Guide

You can test Aegis's autonomous loops using the built-in CLI to inject simulated production failures:

# Inject a latency spike (caused by bad model rollout)
python -m simulator.cli inject latency_spike

# Inject GPU Out-of-Memory / memory fragmentation
python -m simulator.cli inject gpu_oom

# Inject data distribution drift (PSI exceeds baseline)
python -m simulator.cli inject data_drift

# Inject cascading multi-system collapse
python -m simulator.cli inject cascading_failure

# Reset simulator back to normal healthy baseline
python -m simulator.cli inject reset

Watch the Command Console

Once a failure is injected:

  1. Open http://localhost:5173/ in your browser.
  2. Watch the Telemetry Streams area chart spike instantly.
  3. Observe the Live Agent Execution Trace chain light up sequentially as SentrySleuthMedicScribe execute.
  4. Review Scribe's Markdown Postmortem detailing the incident context, diagnosis, and remediation results.
  5. If the proposed action is Medium/High Risk (e.g. model rollback), Aegis halts and triggers the Human Approval Gate Modal asking for operator permission before executing. Low-risk actions (e.g. pod restarts) are resolved autonomously.

🛡️ Security & Remediation Guardrails

To prevent unauthorized or destructive model mutations, Aegis implements a strict risk policy:

  • Active Background Monitoring: An automated loop checks metrics in the simulator every 5 seconds. If an anomaly threshold is breached (severity $\ge 0.70$), the multi-agent pipeline is run autonomously.
  • Auto-execution is restricted to deterministic low-risk commands (e.g. restart_pod, rebalance_queue).
  • All medium/high risk commands (e.g. rollback_model, scale_down_cluster) are locked behind an operator approval gate.
  • Administrative Protection: The /api/agents/approve/{incident_id} route enforces a Bearer Token validator checking requests against the AEGIS_ADMIN_TOKEN config (defaults to aegis-secure-token-2026).
  • Remediation Sandboxing: Remediation commands are stripped from the LLM tool context to prevent jailbreaking; they are parsed and executed programmatically by the orchestrator only after passing the security gate.

About

Autonomous multi-agent system for ML infrastructure reliability. Detects, diagnoses, and remediates production ML incidents through a coordinated Sentry → Sleuth → Medic → Scribe agent pipeline. Built with Google ADK + Gemini, FastAPI, and a live failure-injection demo harness.

Topics

Resources

Stars

Watchers

Forks

Releases

Packages

Contributors

Languages