Cortex

MCP server that lets AI agents monitor and tune neural network training in real time.

Instead of staring at loss curves and manually adjusting hyperparameters, connect an AI agent to your training loop. It watches metrics, detects anomalies, and makes guarded interventions — with safety rails to prevent it from doing more harm than good.

Architecture

┌─────────────────┐     ┌──────────────────┐     ┌───────────────┐
│  Training Loop   │────▶│    Telemetry      │◀────│   AI Agent     │
│                  │     │  (tracker.log)    │     │  (via MCP)     │
│  PyTorch / MLX   │◀────│                  │────▶│                │
│  JAX / any       │     │  ┌────────────┐  │     │  diagnose()    │
│                  │     │  │ Detectors  │  │     │  adjust_param()│
│  tracker.poll()  │     │  │ Guardrails │  │     │  rollback()    │
│  get_override()  │     │  │ Log        │  │     │  save/restore  │
└─────────────────┘     └──────────────────┘     └───────────────┘

Three layers:

1. Telemetry — metrics, config, checkpoints, history
2. Detectors — rule-based anomaly detection producing typed findings (entropy collapse, loss divergence, gradient spikes, eval plateau, loss imbalance)
3. Guardrails — safety constraints on interventions (max % change, cooldowns, rate limits, checkpoint-before-action requirement, full intervention log)

Quick Start

In your training code

from cortex import tracker

tracker.config(total_steps=100000, lr=3e-4, batch_size=64)

@tracker.on_checkpoint
def save(tag):
    torch.save(model.state_dict(), f"checkpoints/{tag}.pt")

@tracker.on_rollback
def rollback(tag):
    model.load_state_dict(torch.load(f"checkpoints/{tag}.pt"))
    return True

for step in range(100000):
    loss = train_step()
    tracker.log(step=step, loss=loss, entropy=ent, grad_norm=gnorm)

    # Process agent commands (checkpoint, rollback, pause/resume)
    tracker.poll()

    # Pick up live param adjustments
    new_lr = tracker.get_override("lr")
    if new_lr is not None:
        optimizer.lr = new_lr

Connect an agent

{
  "mcpServers": {
    "cortex": {
      "command": "cortex"
    }
  }
}

MCP Tools

Observe

Tool	Description
get_status	Step, progress %, ETA, steps/sec, phase
get_metrics	Latest values of all tracked metrics
get_metric_history	Time series for trend analysis
get_config	Hyperparameter configuration
list_metrics	All metric names being tracked

Detect

Tool	Description
diagnose	Run all anomaly detectors, get structured findings with severity and recommendations

Intervene

Tool	Description
adjust_param	Change a hyperparameter (guarded: max % change, cooldown, requires checkpoint)
save_checkpoint	Save model state for later rollback
rollback	Restore model to a checkpoint
pause_training	Pause the loop to analyze and decide
resume_training	Continue after pause

Review

Tool	Description
get_intervention_log	Full history of every change with before/after metrics
get_guardrail_status	Rate limits, cooldowns, intervention count

Detectors

Built-in anomaly detectors that produce structured findings:

Detector	What it catches
entropy_collapse	Policy committing too early (entropy dropping fast while not improving)
loss_divergence	NaN/inf loss, or loss increasing rapidly
loss_imbalance	Value loss drowning out policy loss signal
grad_instability	Gradient norm spikes that precede divergence
eval_plateau	Eval score stopped improving
learning_stall	Loss stopped decreasing

Each finding includes severity, explanation, the specific metric values, and a recommended action.

Guardrails

Every intervention is validated before execution:

• Max % change: Can't change a param by more than 50% at once (configurable)
• Cooldown: 30s minimum between adjustments to the same param
• Rate limit: Max 20 interventions per hour
• Checkpoint required: Must save a checkpoint before making any adjustment
• Full logging: Every intervention recorded with before/after metrics and reason

Install

pip install cortex-mcp

Or from source:

pip install -e .