GitHub Actions
ci: deploy clean snapshot to HF Space
3a6e663
|
Raw
History Blame Contribute Delete
9.29 kB
metadata
title: OpenEnv Split-Brain Collapse
emoji: 🧠
colorFrom: indigo
colorTo: red
sdk: docker
pinned: false
license: mit

🧠 OpenEnv: Split-Brain Collapse (Three-Datacenter Crisis)

πŸ“– Overview & Motivation

Managing a distributed database cluster spanning multiple datacenters is one of the hardest challenges in Site Reliability Engineering. When network partitions occur between datacenters, competing nodes may each believe they are the authoritative primary β€” a split-brain state that leads to data divergence, replication storms, and cascading deadlocks that require precise, ordered multi-step intervention.

This OpenEnv simulates a three-datacenter enterprise infrastructure under a cascading split-brain crisis. It is a genuine multi-agent RL scenario. An AI team consisting of an orchestrator, a netops specialist, and a dba specialist must collaborate to restore connectivity, elect a single primary, reconcile diverged write logs, and recover the cluster before data loss becomes permanent.


βš™οΈ Architecture: How It Works

This environment is built as a complete, containerized Reinforcement Learning (RL) ecosystem:

  1. The Simulation Engine (agents/split_brain/environment.py): An OpenEnv-compliant multi-agent state machine tracking datacenter health, network link status, replication lag, transaction deadlocks, and enforcing strict logic gates (e.g., you cannot promote a new primary before routing is verified; resolving a deadlock before re-routing re-triggers it).
  2. The Backend API (FastAPI): Exposes standard programmatic RL endpoints (/reset, /step, /state, /health, /tasks) allowing external scripts and evaluators to interact with the environment headlessly.
  3. The Web Dashboard (static/index.html): A dark-mode, multi-panel "SRE Command Center" that auto-discovers agents from GET /agents and lets humans interactively step through the simulation watching the LLM make decisions in real-time.
  4. The Agent (inference.py / OpenAI Client): Uses the deepseek-ai/DeepSeek-R1-Distill-Llama-70B model via the Hugging Face Serverless API. The environment injects context-aware multi-agent prompts (different system prompts for orchestrator, netops, dba) so the LLM knows which role it is playing each step.

🧠 Observation & Action Spaces

The environment strictly adheres to the OpenEnv Pydantic specifications.

Observation Space

The agent receives a SplitBrainObservation detailing the full three-datacenter state:

  • global_health: (Float 0.0 – 1.0) Continuous cluster stability metric across all three DCs.
  • current_actor: Which sub-agent is active this step (orchestrator, netops, or dba).
  • dc1_dc2_connected / dc2_dc3_connected: Network link status between datacenters.
  • network_status: Overall network state (partitioned, degraded, healthy).
  • replication_lag_ms: Database replication lag in milliseconds.
  • primary_db: Which DC currently holds the authoritative primary database.
  • datacenters: List of 3 DatacenterStatus objects β€” status, node count, load %.
  • recent_alerts: System log messages and critical alerts.

Action Space

The active sub-agent issues a SplitBrainAction JSON object:

  • action_type: Command to execute β€” run_diagnostic, update_route, verify_routing, throttle_bandwidth, restore_replica, promote_primary, force_sync, resolve_deadlock, failover_region, delegate, noop.
  • target_id: Resource to target (e.g. dc2_router--dc3_router, replica_dc2).
  • target_agent: Sub-agent to delegate to (netops or dba) β€” only for delegate actions.
  • instruction_payload: Nested action passed to the delegated sub-agent.

🎯 Task Descriptions & Difficulty

The environment features a 5-tier difficulty scale with continuous shaped rewards penalising wrong ordering and rewarding correct multi-step sequences.

  • 🟒 BASIC (The Partition): A single DC1–DC2 link has failed.
    • Expected Sequence: run_diagnostic β†’ update_route β†’ verify_routing.
  • 🟑 STORM (Replication Storm): A thundering-herd burst has saturated all inter-DC links.
    • Expected Sequence: throttle_bandwidth β†’ restore_replica Γ— N β†’ force_sync.
  • πŸ”΄ SPLIT (The Split-Brain): Both inter-DC links are down; two competing primaries have diverged.
    • Expected Sequence: Restore both links, isolate one primary, promote_primary, reconcile all replicas.
  • 🟣 DEADLOCK (Cascading Deadlock): A distributed transaction deadlock has frozen all write paths.
    • Expected Sequence: update_route β†’ verify_routing β†’ resolve_deadlock β†’ force_sync (strict order).
  • ☠️ WIPEOUT (Regional Wipeout): DC3 offline, DC1/DC2 in split-brain. Full regional failover required.
    • Expected Sequence: failover_region β†’ promote_primary β†’ restore all replicas β†’ verify_routing.

πŸ’» Setup & Usage Instructions

You can run this project in three different ways depending on your needs.

Method 1: Hugging Face Spaces (Interactive Web UI)

The easiest way to evaluate the environment is via the public Hugging Face Space. The UI acts as a step-by-step RL debugger.

  1. Access the Dashboard: Open the Hugging Face Space URL.
  2. Select Threat Level: Use the sidebar to choose a task scenario.
  3. Initialize the Environment: Click the πŸ”„ Reset button to boot the simulation.
  4. Deploy the Agent: Click the β–Ά AGENT STEP button. The LLM evaluates the current state and makes exactly one move, routing to the correct sub-agent (orchestrator / netops / dba).
  5. Observe the Triage: Watch the delegation log, step reward, and datacenter health panel update.
  6. Iterate: Continue clicking β–Ά AGENT STEP until the terminal declares CLUSTER RESTORED.

Method 2: Local Docker (Production Simulation)

Run the exact containerized environment that Hugging Face uses, locally on your machine.

  1. Create a .env file in the root directory:
    HF_TOKEN="your_huggingface_token"
    MODEL_NAME="deepseek-ai/DeepSeek-R1-Distill-Llama-70B"
    API_BASE_URL="https://router.huggingface.co/v1"
    SPLIT_BRAIN_MODEL=""  # optional: path to fine-tuned LoRA model
    
  2. Build the Docker image:
    docker build -t split-brain-env .
    
  3. Run the container:
    docker run -p 7860:7860 --env-file .env split-brain-env
    
  4. Open your browser at: http://localhost:7860

Method 3: Local Python (For Developers)

  1. Install the required dependencies:
    pip install -r requirements.txt
    
  2. Ensure your .env file is configured (same as Docker step 1).
  3. Run the Interactive Web Dashboard & API:
    python app.py
    
    (Access the UI via http://127.0.0.1:7860).
  4. Run the Automated Terminal Baseline:
    python inference.py
    

🧬 GRPO Fine-Tuning with Unsloth (LoRA Adapter)

The Problem

When running the Cascading Deadlock task (Task 4) with smaller models like Llama 3.1 8B, the agent gets stuck in an infinite run_diagnostic loop β€” it keeps repeating the same diagnostic action instead of progressing to update_route, verify_routing, and resolve_deadlock. Larger models (70B+) handle this correctly but are expensive to run.

The Solution: GRPO Reinforcement Learning

We used Group Relative Policy Optimization (GRPO) via Unsloth + TRL to fine-tune Llama-3.2-3B-Instruct directly against the OpenEnv reward function. The training:

  1. Feeds the model the exact stuck scenario (post-diagnostic delegation to netops)
  2. Generates multiple candidate actions via sampling
  3. Steps each action through the live SplitBrainEnv.step() function
  4. Rewards correct actions (e.g., update_route) and penalises repeated diagnostics
  5. The model learns to break out of loops and follow multi-step repair sequences

Training Script

The training script is train_unsloth_colab.py β€” designed to run on Google Colab with a free T4 GPU:

# In Google Colab:
!git clone https://github.com/Sayantan181222/openenv-cluster-triage-updated.git
%cd openenv-cluster-triage-updated
!pip install -r requirements.txt
!pip install "unsloth[colab] @ git+https://github.com/unslothai/unsloth.git" trl datasets
!python train_unsloth_colab.py

The trained adapter is saved to openenv-split-brain-lora/ (~93MB LoRA weights).

Using the LoRA Adapter

Run the fine-tuned model against the Split-Brain environment:

python inference_lora.py

This loads Llama-3.2-3B-Instruct + the LoRA adapter and runs a full episode on the cascading_deadlock task, printing a before/after comparison.

Results & Improvement

Metric Base 8B (No LoRA) 3B + GRPO LoRA
Diagnostic Loops 10+ (infinite) ≀1
Reached update_route ❌ Never βœ… Yes
Episode Completion ❌ Timed out βœ… Completes
Model Size 8B parameters 3B parameters

The fine-tuned 3B model outperforms the base 8B model by learning environment-specific action sequences through reinforcement learning.