MLOps Firefighter - OpenEnv environment

Dockerfile

@@ -0,0 +1,18 @@
+FROM python:3.11-slim
+
+RUN useradd -m -u 1000 user
+USER user
+ENV PATH="/home/user/.local/bin:$PATH"
+ENV PYTHONUNBUFFERED=1
+ENV PORT=7860
+
+WORKDIR /app
+
+COPY --chown=user ./requirements.txt requirements.txt
+RUN pip install --no-cache-dir --upgrade -r requirements.txt
+
+COPY --chown=user . /app
+
+EXPOSE 7860
+
+CMD ["python", "server/app.py"]

README.md

@@ -1,12 +1,292 @@
 ---
-title: Mlops Firefighter
-emoji: 👀
-colorFrom: blue
-colorTo: red
+title: MLOps Firefighter
+colorFrom: red
+colorTo: orange
 sdk: docker
+app_port: 7860
+tags:
+  - openenv
+  - rl
+  - mlops
+  - production-ml
 pinned: false
-license: mit
-short_description: AI agents act as on-call MLOps engineers
+license: bsd-3-clause
 ---
 
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+# MLOps Firefighter — OpenEnv Environment
+
+**Debug and fix ML models failing in production.**
+
+An OpenEnv-compliant reinforcement learning environment where AI agents act as on-call MLOps engineers. When a production ML model starts misbehaving, the agent must diagnose the root cause and apply the correct fix, just like a real engineer would at 3 AM.
+
+---
+
+## Why This Environment?
+
+Every ML team eventually faces the "model broke in prod" moment. Data drifts. Thresholds get misconfigured during deployments. Training pipelines get poisoned. These are real incidents that cost companies millions of dollars and require skilled human reasoning to resolve.
+
+This environment captures that challenge:
+- **Real-world task**: MLOps incident response is performed daily by thousands of engineers
+- **Rich reasoning required**: agents must investigate before acting, weigh evidence, and avoid destructive actions
+- **Novel domain**: no existing OpenEnv covers production ML debugging
+- **Meaningful difficulty progression**: from a simple config error to a subtle adversarial data poisoning attack
+
+---
+
+## Environment Overview
+
+### How It Works
+
+1. **An incident fires**: the agent receives alerts about a failing ML model (e.g., precision dropped, complaints spiked)
+2. **The agent investigates**: run diagnostics like checking metrics, querying logs, inspecting data distributions
+3. **The agent identifies the root cause**: data drift? threshold misconfiguration? poisoned training data?
+4. **The agent applies a fix**: rollback the model, adjust thresholds, fix the data pipeline, add guardrails
+5. **The agent submits a diagnosis**: declare the root cause and close the incident
+
+The grader evaluates: diagnostic thoroughness (30%), diagnosis accuracy (30%), remediation correctness (25%), and efficiency (15%).
+
+### Reward Shaping
+
+Rewards provide signal throughout the episode, not just at the end:
+
+| Action | Reward |
+|--------|--------|
+| Required diagnostic (first time) | +0.3 |
+| Useful but non-critical diagnostic | +0.1 |
+| Redundant diagnostic (already done) | +0.05 |
+| Correct remediation (after investigation) | +1.0 |
+| Hasty remediation (without diagnosis) | −0.5 |
+| Wrong remediation | −0.3 |
+| Final score bonus (on submit_diagnosis) | 0–5.0 scaled by grader |
+| Timeout penalty | Reduced final multiplier |
+
+---
+
+## Action Space
+
+Actions the agent can take, specified via `action_type` string:
+
+### Diagnostic Actions (investigation)
+| Action | Description |
+|--------|-------------|
+| `inspect_metrics` | View model performance dashboard (accuracy, precision, recall, latency) |
+| `query_logs` | Search production logs for deployment events, errors, warnings |
+| `check_data_dist` | Compare training vs. serving data distributions, detect drift |
+| `check_feature_importance` | Examine which features the model relies on |
+| `run_prediction_sample` | Test the model on samples with known ground-truth labels |
+| `check_infrastructure` | Check CPU, memory, GPU, latency — rule out infra issues |
+| `check_upstream_pipeline` | Inspect data pipeline health, connectors, schema changes |
+
+### Remediation Actions (fixes)
+| Action | Parameters | Description |
+|--------|-----------|-------------|
+| `rollback_model` | `target_version: str` | Revert to a previous model version |
+| `adjust_threshold` | `new_threshold: float` | Tune the decision threshold |
+| `retrain_model` | — | Trigger model retraining on corrected data |
+| `fix_data_pipeline` | — | Repair data ingestion / feature pipeline |
+| `scale_infrastructure` | — | Add compute resources |
+| `add_feature_guard` | — | Add input validation and monitoring guardrails |
+
+### Episode Control
+| Action | Parameters | Description |
+|--------|-----------|-------------|
+| `submit_diagnosis` | `root_cause: str, summary: str` | Declare root cause and end the episode |
+
+---
+
+## Observation Space
+
+After each action, the agent receives an `MLOpsObservation` with:
+
+| Field | Type | Description |
+|-------|------|-------------|
+| `done` | bool | Whether the episode has ended |
+| `reward` | float | Reward for the last action |
+| `step_number` | int | Current step (starts at 0) |
+| `max_steps` | int | Steps before timeout |
+| `task_id` | str | Task identifier |
+| `task_description` | str | Natural language incident description |
+| `alerts` | list[Alert] | Active production alerts with severity, metric, value |
+| `model_info` | ModelInfo | Deployed model name, version, framework, endpoint |
+| `action_result` | str | Detailed textual result of the last action |
+| `action_success` | bool | Whether the action executed successfully |
+| `diagnostics_gathered` | list[str] | Summary of all diagnostics collected so far |
+| `available_actions` | list[str] | Valid action types |
+
+---
+
+## Tasks
+
+### Task 1: Threshold Misconfiguration (Easy)
+**Scenario**: A fraud detection model's precision dropped from 0.94 to 0.61 after redeployment. The decision threshold was accidentally changed from 0.55 to 0.30.
+
+- **Root cause**: Threshold misconfiguration during deployment
+- **Expected fix**: Adjust threshold back to ~0.55
+- **Key signal**: Metrics show threshold changed; prediction samples confirm threshold is the lever
+- **Max steps**: 15
+
+### Task 2: Data Drift with Stale Feature Pipeline (Medium)
+**Scenario**: A loan default model's AUC degraded from 0.91 to 0.74. A new credit bureau connector was onboarded 10 days ago, introducing null values and unit changes (dollars → cents) in key features.
+
+- **Root causes**: Data drift + broken feature pipeline
+- **Expected fixes**: Fix data pipeline + retrain model
+- **Key signals**: Distribution comparison shows 100x scale change in income; upstream pipeline shows connector issues
+- **Max steps**: 20
+
+### Task 3: Silent Model Regression with Data Poisoning (Hard)
+**Scenario**: A content moderation model has 0.96 overall accuracy (looks fine!), but hate speech and violence recall collapsed to ~40%. An automated retraining pipeline ingested a poisoned crowd-source batch that systematically mislabeled harmful content as "safe." Aggregate metrics didn't catch it because the safe class dominates traffic (92%).
+
+- **Root causes**: Training data poisoning / label corruption
+- **Expected fixes**: Roll back to safe version + add per-class recall guardrails
+- **Key signals**: Per-class metrics reveal the hidden regression; data audit shows contaminated batch
+- **Difficulty**: Requires the agent to look beyond aggregate metrics and think about class imbalance
+- **Max steps**: 25
+
+### Expected Difficulty
+| Task | Difficulty | Baseline Score | Notes |
+|------|-----------|---------------|-------|
+| Threshold Misconfiguration | Easy | 1.000 | Single clear root cause |
+| Data Drift | Medium | 0.970 | Multiple issues to identify |
+| Silent Regression | Hard | 0.970 | Requires looking past surface metrics |
+
+---
+
+## Setup & Usage
+
+### Quick Start (Local Python)
+
+```bash
+# Clone and install
+git clone <repo-url>
+cd mlops_firefighter
+pip install -r requirements.txt
+
+# Run the server
+python server/app.py
+# → Server running at http://localhost:7860
+
+# In another terminal, test it
+curl http://localhost:7860/health
+curl http://localhost:7860/tasks
+curl -X POST http://localhost:7860/reset -H "Content-Type: application/json" -d '{"task_id": "task_threshold_drift"}'
+curl -X POST http://localhost:7860/step -H "Content-Type: application/json" -d '{"action_type": "inspect_metrics"}'
+```
+
+### Docker
+
+```bash
+# Build
+docker build -t mlops-firefighter .
+
+# Run
+docker run -p 7860:7860 mlops-firefighter
+
+# Test
+curl http://localhost:7860/health
+```
+
+### Run Baseline Inference (requires OpenAI API key)
+
+```bash
+export OPENAI_API_KEY=sk-...
+python baseline_inference.py --base-url http://localhost:7860
+```
+
+### Run the Built-in Baseline (no API key needed)
+
+```bash
+curl -X POST http://localhost:7860/baseline
+```
+
+### Run Tests
+
+```bash
+python tests/test_environment.py
+```
+
+---
+
+## API Endpoints
+
+| Endpoint | Method | Description |
+|----------|--------|-------------|
+| `/health` | GET | Health check — returns `{"status": "healthy"}` |
+| `/reset` | POST | Start new episode. Body: `{"task_id": "..."}` |
+| `/step` | POST | Take action. Body: `{"action_type": "...", "parameters": {...}}` |
+| `/state` | GET | Current environment state |
+| `/tasks` | GET | List all tasks with action schema |
+| `/grader` | POST | Get grader score for completed episode |
+| `/baseline` | POST | Run scripted baseline on all tasks, return scores |
+
+---
+
+## Grader Scoring Breakdown
+
+Each episode is scored 0.0–1.0 based on four components:
+
+| Component | Weight | Description |
+|-----------|--------|-------------|
+| Diagnostic thoroughness | 30% | Did the agent run the key diagnostics? |
+| Diagnosis accuracy | 30% | Did the agent identify the correct root cause? |
+| Remediation accuracy | 25% | Did the agent apply the right fix(es)? Penalizes wrong fixes. |
+| Efficiency | 15% | How many steps did the agent use? Fewer = better. |
+
+---
+
+## Baseline Scores
+
+Using the built-in scripted baseline (perfect knowledge):
+
+| Task | Score | Steps |
+|------|-------|-------|
+| task_threshold_drift (Easy) | 1.000 | 4 |
+| task_data_drift (Medium) | 0.970 | 6 |
+| task_silent_regression (Hard) | 0.970 | 7 |
+| **Average** | **0.980** | — |
+
+These represent near-optimal play. An LLM agent without perfect knowledge will score lower, especially on the hard task where aggregate metrics are misleading.
+
+---
+
+## Project Structure
+
+```
+mlops_firefighter/
+├── models.py                  # Pydantic Action/Observation models
+├── tasks.py                   # Task definitions + grader functions
+├── openenv.yaml               # OpenEnv manifest
+├── requirements.txt           # Python dependencies
+├── pyproject.toml             # Package config
+├── Dockerfile                 # Container image
+├── baseline_inference.py      # LLM baseline script (OpenAI API)
+├── README.md                  # This file
+├── server/
+│   ├── __init__.py
+│   ├── app.py                 # FastAPI server with all endpoints
+│   └── environment.py         # Core environment logic
+└── tests/
+    └── test_environment.py    # Comprehensive test suite (13 tests)
+```
+
+---
+
+## OpenEnv Spec Compliance
+
+- ✅ Typed `Action` and `Observation` Pydantic models
+- ✅ `step(action)` → returns observation, reward, done, info
+- ✅ `reset()` → returns initial observation
+- ✅ `state()` → returns current state
+- ✅ `openenv.yaml` with metadata
+- ✅ 3 tasks with difficulty progression (easy/medium/hard)
+- ✅ Programmatic graders scoring 0.0–1.0
+- ✅ Meaningful reward shaping (not just sparse end-of-episode)
+- ✅ Baseline inference script using OpenAI API
+- ✅ Dockerfile that builds and runs
+- ✅ `/health`, `/tasks`, `/grader`, `/baseline` endpoints
+
+---
+
+## License
+
+BSD-3-Clause
+

baseline_inference.py

@@ -0,0 +1,245 @@
+#!/usr/bin/env python3
+# Copyright (c) 2026 MLOps Firefighter Contributors
+# Licensed under the BSD-3-Clause License
+
+"""
+Baseline inference script for the MLOps Firefighter environment.
+
+Uses the OpenAI API client to run a model (e.g. GPT-4o) against all 3 tasks,
+producing a reproducible baseline score.
+
+Usage:
+    export OPENAI_API_KEY=sk-...
+    python baseline_inference.py [--base-url http://localhost:7860]
+"""
+
+from __future__ import annotations
+
+import argparse
+import json
+import os
+import sys
+
+import requests
+
+try:
+    from openai import OpenAI
+except ImportError:
+    print("Install openai: pip install openai")
+    sys.exit(1)
+
+
+# ── Configuration ────────────────────────────────────────────────────────────
+
+DEFAULT_ENV_URL = "http://localhost:7860"
+MODEL = os.environ.get("OPENAI_MODEL", "gpt-4o")
+
+SYSTEM_PROMPT = """\
+You are an expert MLOps engineer on-call. A production ML model is failing and
+you must diagnose the root cause and fix it.
+
+You interact with the environment by choosing actions. Each action returns
+diagnostic information or applies a remediation.
+
+STRATEGY:
+1. First, investigate: run diagnostics to understand what's wrong
+2. Identify the root cause from the evidence
+3. Apply the correct remediation
+4. Submit your diagnosis with a clear root cause label and summary
+
+Available actions (use exact strings):
+- inspect_metrics: View model performance metrics dashboard
+- query_logs: Search production logs for anomalies
+- check_data_dist: Compare training vs serving data distributions
+- check_feature_importance: Examine feature weights and importance
+- run_prediction_sample: Test model on sample inputs with known labels
+- check_infrastructure: Check latency, memory, GPU, compute
+- check_upstream_pipeline: Inspect data pipeline health
+- rollback_model: Revert to a previous model version (params: target_version)
+- adjust_threshold: Tune decision threshold (params: new_threshold)
+- retrain_model: Trigger model retraining
+- fix_data_pipeline: Repair data ingestion issues
+- scale_infrastructure: Add compute resources
+- add_feature_guard: Add input validation / guardrails
+- submit_diagnosis: Declare root cause and close (params: root_cause, summary)
+
+Respond with ONLY valid JSON: {"action_type": "...", "parameters": {...}}
+"""
+
+
+def call_env(base_url: str, endpoint: str, method: str = "POST", data: dict | None = None) -> dict:
+    """Call an environment endpoint."""
+    url = f"{base_url}{endpoint}"
+    if method == "GET":
+        resp = requests.get(url, timeout=30)
+    else:
+        resp = requests.post(url, json=data or {}, timeout=30)
+    resp.raise_for_status()
+    return resp.json()
+
+
+def run_agent_on_task(client: OpenAI, base_url: str, task_id: str, task_name: str) -> dict:
+    """Run the LLM agent on a single task."""
+    print(f"\n{'='*60}")
+    print(f"  Task: {task_name} ({task_id})")
+    print(f"{'='*60}")
+
+    # Reset environment
+    reset_result = call_env(base_url, "/reset", data={"task_id": task_id})
+    obs = reset_result["observation"]
+    print(f"  Incident: {obs['task_description'][:100]}...")
+
+    messages = [
+        {"role": "system", "content": SYSTEM_PROMPT},
+        {
+            "role": "user",
+            "content": (
+                f"INCIDENT ALERT:\n{obs['task_description']}\n\n"
+                f"ALERTS:\n{json.dumps(obs['alerts'], indent=2)}\n\n"
+                f"MODEL INFO:\n{json.dumps(obs['model_info'], indent=2)}\n\n"
+                f"Available actions: {obs['available_actions']}\n\n"
+                "What is your first action? Respond with JSON only."
+            ),
+        },
+    ]
+
+    total_reward = 0.0
+    steps = 0
+    max_steps = obs.get("max_steps", 20)
+
+    while steps < max_steps:
+        # Get LLM decision
+        try:
+            response = client.chat.completions.create(
+                model=MODEL,
+                messages=messages,
+                temperature=0.2,
+                max_tokens=500,
+            )
+            llm_output = response.choices[0].message.content.strip()
+        except Exception as e:
+            print(f"  LLM error: {e}")
+            break
+
+        # Parse action from LLM
+        try:
+            # Handle markdown code blocks
+            if "```" in llm_output:
+                llm_output = llm_output.split("```")[1]
+                if llm_output.startswith("json"):
+                    llm_output = llm_output[4:]
+                llm_output = llm_output.strip()
+            action_data = json.loads(llm_output)
+        except json.JSONDecodeError:
+            # Try to extract JSON from the response
+            import re
+            match = re.search(r'\{[^}]+\}', llm_output)
+            if match:
+                try:
+                    action_data = json.loads(match.group())
+                except json.JSONDecodeError:
+                    print(f"  Failed to parse LLM output: {llm_output[:100]}")
+                    break
+            else:
+                print(f"  Failed to parse LLM output: {llm_output[:100]}")
+                break
+
+        action_type = action_data.get("action_type", "")
+        parameters = action_data.get("parameters", {})
+
+        # Step environment
+        step_result = call_env(
+            base_url, "/step",
+            data={"action_type": action_type, "parameters": parameters},
+        )
+        obs = step_result["observation"]
+        reward = step_result["reward"]
+        done = step_result["done"]
+        total_reward += reward
+        steps += 1
+
+        print(f"  Step {steps}: {action_type} → reward={reward:.2f}")
+
+        if done:
+            print(f"  Episode done. Total reward: {total_reward:.2f}")
+            break
+
+        # Feed result back to LLM
+        messages.append({"role": "assistant", "content": llm_output})
+        messages.append({
+            "role": "user",
+            "content": (
+                f"ACTION RESULT:\n{obs['action_result']}\n\n"
+                f"Step {obs['step_number']}/{obs['max_steps']}\n"
+                f"Diagnostics gathered so far: {len(obs['diagnostics_gathered'])}\n\n"
+                "What is your next action? Respond with JSON only."
+            ),
+        })
+
+    # Get grader score
+    grader_result = call_env(base_url, "/grader", data={})
+    score = grader_result.get("score", 0.0)
+    breakdown = grader_result.get("breakdown", {})
+
+    print(f"  Grader score: {score:.3f}")
+    print(f"  Breakdown: {json.dumps(breakdown, indent=2)}")
+
+    return {
+        "task_id": task_id,
+        "task_name": task_name,
+        "score": score,
+        "breakdown": breakdown,
+        "steps": steps,
+        "total_reward": round(total_reward, 3),
+    }
+
+
+def main():
+    parser = argparse.ArgumentParser(description="MLOps Firefighter Baseline")
+    parser.add_argument("--base-url", default=DEFAULT_ENV_URL, help="Environment URL")
+    args = parser.parse_args()
+
+    api_key = os.environ.get("OPENAI_API_KEY")
+    if not api_key:
+        print("ERROR: Set OPENAI_API_KEY environment variable")
+        sys.exit(1)
+
+    client = OpenAI(api_key=api_key)
+
+    # Get tasks
+    tasks_info = call_env(args.base_url, "/tasks", method="GET")
+    tasks = tasks_info["tasks"]
+
+    print("\n" + "=" * 60)
+    print("  MLOps Firefighter — Baseline Inference")
+    print(f"  Model: {MODEL}")
+    print(f"  Environment: {args.base_url}")
+    print(f"  Tasks: {len(tasks)}")
+    print("=" * 60)
+
+    results = []
+    for task in tasks:
+        result = run_agent_on_task(
+            client, args.base_url, task["task_id"], task["name"]
+        )
+        results.append(result)
+
+    # Summary
+    print("\n" + "=" * 60)
+    print("  BASELINE RESULTS SUMMARY")
+    print("=" * 60)
+    for r in results:
+        print(f"  [{r['task_id']}] {r['task_name']}")
+        print(f"    Score: {r['score']:.3f} | Steps: {r['steps']} | Reward: {r['total_reward']}")
+    avg = sum(r["score"] for r in results) / len(results) if results else 0
+    print(f"\n  Average Score: {avg:.3f}")
+    print("=" * 60)
+
+    # Write results to file
+    with open("baseline_results.json", "w") as f:
+        json.dump({"results": results, "average_score": round(avg, 3)}, f, indent=2)
+    print("  Results saved to baseline_results.json")
+
+
+if __name__ == "__main__":
+    main()

models.py

@@ -0,0 +1,117 @@
+# Copyright (c) 2026 MLOps Firefighter Contributors
+# Licensed under the BSD-3-Clause License
+
+"""
+Typed Pydantic models for the MLOps Firefighter environment.
+Defines Action, Observation, and related types for the OpenEnv spec.
+"""
+
+from __future__ import annotations
+
+from enum import Enum
+from typing import Any, Optional
+
+from pydantic import BaseModel, Field
+
+
+# ── Action Space ─────────────────────────────────────────────────────────────
+
+class ActionType(str, Enum):
+    """All actions an agent can take to diagnose and fix a production ML model."""
+
+    # Diagnostic actions
+    INSPECT_METRICS = "inspect_metrics"           # View model performance metrics
+    QUERY_LOGS = "query_logs"                     # Search production logs
+    CHECK_DATA_DISTRIBUTION = "check_data_dist"   # Compare train vs. serving data
+    CHECK_FEATURE_IMPORTANCE = "check_feature_importance"  # Examine feature weights
+    RUN_PREDICTION_SAMPLE = "run_prediction_sample"  # Test model on sample inputs
+    CHECK_INFRASTRUCTURE = "check_infrastructure"  # Check latency, memory, GPU
+    CHECK_UPSTREAM_PIPELINE = "check_upstream_pipeline"  # Inspect data pipeline health
+
+    # Remediation actions
+    ROLLBACK_MODEL = "rollback_model"             # Revert to previous model version
+    ADJUST_THRESHOLD = "adjust_threshold"         # Tune decision threshold
+    RETRAIN_MODEL = "retrain_model"               # Trigger retraining
+    FIX_DATA_PIPELINE = "fix_data_pipeline"       # Repair data ingestion issue
+    SCALE_INFRASTRUCTURE = "scale_infrastructure" # Add compute resources
+    ADD_FEATURE_GUARD = "add_feature_guard"       # Add input validation / guardrail
+    SUBMIT_DIAGNOSIS = "submit_diagnosis"         # Declare root cause & close episode
+
+
+class MLOpsAction(BaseModel):
+    """An action taken by the agent in the MLOps Firefighter environment."""
+
+    action_type: ActionType = Field(
+        ..., description="The type of action to perform"
+    )
+    parameters: dict[str, Any] = Field(
+        default_factory=dict,
+        description=(
+            "Action-specific parameters. E.g. for adjust_threshold: "
+            "{'new_threshold': 0.6}. For submit_diagnosis: "
+            "{'root_cause': 'data_drift', 'summary': '...'}."
+        ),
+    )
+
+
+# ── Observation Space ────────────────────────────────────────────────────────
+
+class AlertSeverity(str, Enum):
+    CRITICAL = "critical"
+    HIGH = "high"
+    MEDIUM = "medium"
+    LOW = "low"
+
+
+class Alert(BaseModel):
+    """A production alert that triggered the incident."""
+    severity: AlertSeverity
+    message: str
+    metric_name: str
+    current_value: float
+    threshold: float
+    timestamp: str
+
+
+class ModelInfo(BaseModel):
+    """Information about the deployed model."""
+    model_name: str
+    model_version: str
+    deployed_at: str
+    framework: str
+    endpoint: str
+    previous_versions: list[str] = Field(default_factory=list)
+
+
+class MLOpsObservation(BaseModel):
+    """What the agent sees after each action."""
+
+    # Episode metadata
+    done: bool = Field(False, description="Whether the episode has ended")
+    reward: float = Field(0.0, description="Reward for the last action")
+    step_number: int = Field(0, description="Current step in the episode")
+    max_steps: int = Field(20, description="Maximum steps before timeout")
+
+    # Incident context (always visible)
+    task_id: str = Field("", description="Task identifier")
+    task_description: str = Field("", description="Natural language task description")
+    alerts: list[Alert] = Field(default_factory=list, description="Active alerts")
+    model_info: ModelInfo | None = Field(None, description="Deployed model details")
+
+    # Action result (populated after each step)
+    action_result: str = Field(
+        "", description="Textual result of the last action taken"
+    )
+    action_success: bool = Field(True, description="Whether the action executed OK")
+
+    # Accumulated diagnostic context
+    diagnostics_gathered: list[str] = Field(
+        default_factory=list,
+        description="Summary of diagnostics the agent has collected so far",
+    )
+
+    # Hints / guidance
+    available_actions: list[str] = Field(
+        default_factory=list,
+        description="List of valid action types the agent can take",
+    )

openenv.yaml

@@ -0,0 +1,60 @@
+name: mlops_firefighter
+version: "1.0.0"
+description: >
+  MLOps Firefighter: An OpenEnv environment where AI agents debug and fix
+  ML models failing in production. Agents diagnose root causes (data drift,
+  threshold misconfiguration, training data poisoning) and apply the correct
+  remediation. 3 tasks from easy to hard with programmatic graders.
+
+author: "MLOps Firefighter Contributors"
+license: "BSD-3-Clause"
+tags:
+  - openenv
+  - rl
+  - mlops
+  - production-ml
+  - debugging
+  - real-world
+
+environment:
+  entrypoint: server/app.py
+  port: 7860
+  python_version: "3.11"
+
+  action_model: models.MLOpsAction
+  observation_model: models.MLOpsObservation
+
+  endpoints:
+    reset: /reset
+    step: /step
+    state: /state
+    health: /health
+    tasks: /tasks
+    grader: /grader
+    baseline: /baseline
+
+tasks:
+  - id: task_threshold_drift
+    name: "Threshold Misconfiguration After Redeployment"
+    difficulty: easy
+    description: >
+      Fraud detection model precision dropped after redeployment.
+      Agent must diagnose threshold misconfiguration and adjust it.
+    max_steps: 15
+
+  - id: task_data_drift
+    name: "Data Drift with Stale Feature Pipeline"
+    difficulty: medium
+    description: >
+      Loan default model degraded due to upstream data pipeline change.
+      Agent must identify feature drift/corruption and fix the pipeline + retrain.
+    max_steps: 20
+
+  - id: task_silent_regression
+    name: "Silent Model Regression with Adversarial Inputs"
+    difficulty: hard
+    description: >
+      Content moderation model looks fine on aggregate metrics but is missing
+      hate/violence due to poisoned training data. Agent must find the silent
+      regression, roll back safely, and add guards.
+    max_steps: 25

pyproject.toml

@@ -0,0 +1,24 @@
+[build-system]
+requires = ["setuptools>=68.0", "wheel"]
+build-backend = "setuptools.build_meta"
+
+[project]
+name = "mlops-firefighter"
+version = "1.0.0"
+description = "OpenEnv environment: Debug and fix ML models failing in production"
+readme = "README.md"
+license = {text = "BSD-3-Clause"}
+requires-python = ">=3.10"
+dependencies = [
+    "fastapi>=0.104.0",
+    "uvicorn[standard]>=0.24.0",
+    "pydantic>=2.5.0",
+    "requests>=2.31.0",
+]
+
+[project.optional-dependencies]
+baseline = ["openai>=1.0.0"]
+dev = ["pytest>=7.0", "httpx>=0.25.0"]
+
+[tool.setuptools.packages.find]
+include = ["mlops_firefighter*", "server*"]

requirements.txt

@@ -0,0 +1,5 @@
+fastapi>=0.104.0
+uvicorn[standard]>=0.24.0
+pydantic>=2.5.0
+requests>=2.31.0
+openai>=1.0.0

server/__init__.py

@@ -0,0 +1 @@
+# Server package

server/app.py

@@ -0,0 +1,259 @@
+# Copyright (c) 2026 MLOps Firefighter Contributors
+# Licensed under the BSD-3-Clause License
+
+"""
+FastAPI server for the MLOps Firefighter OpenEnv environment.
+
+Endpoints:
+  POST /reset        — Start a new episode
+  POST /step         — Take an action
+  GET  /state        — Current environment state
+  GET  /health       — Health check
+  GET  /tasks        — List tasks and action schema
+  POST /grader       — Score a completed episode
+  POST /baseline     — Run baseline inference on all tasks
+"""
+
+from __future__ import annotations
+
+import json
+import os
+import sys
+import traceback
+from typing import Any
+
+from fastapi import FastAPI, HTTPException
+from fastapi.middleware.cors import CORSMiddleware
+from pydantic import BaseModel, Field
+
+# Add parent dir to path for imports
+sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+
+from models import ActionType, MLOpsAction, MLOpsObservation
+from tasks import ALL_TASKS, grade_episode
+
+from environment import MLOpsFirefighterEnvironment
+
+# ── App Setup ────────────────────────────────────────────────────────────────
+
+app = FastAPI(
+    title="MLOps Firefighter — OpenEnv Environment",
+    description=(
+        "An AI agent environment for debugging and fixing ML models in production. "
+        "The agent diagnoses root causes of model failures (data drift, threshold "
+        "misconfiguration, training data poisoning) and applies the correct fix."
+    ),
+    version="1.0.0",
+)
+
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=["*"],
+    allow_credentials=True,
+    allow_methods=["*"],
+    allow_headers=["*"],
+)
+
+# Per-session environments (simple in-memory for single container)
+_environments: dict[str, MLOpsFirefighterEnvironment] = {}
+_default_env = MLOpsFirefighterEnvironment()
+
+
+def _get_env(session_id: str | None = None) -> MLOpsFirefighterEnvironment:
+    if session_id and session_id in _environments:
+        return _environments[session_id]
+    return _default_env
+
+
+# ── Request / Response Models ────────────────────────────────────────────────
+
+class ResetRequest(BaseModel):
+    task_id: str | None = Field(None, description="Task to load (optional)")
+    session_id: str | None = Field(None, description="Session ID (optional)")
+
+
+class StepRequest(BaseModel):
+    action_type: str = Field(..., description="Action type to perform")
+    parameters: dict[str, Any] = Field(default_factory=dict)
+    session_id: str | None = Field(None, description="Session ID (optional)")
+
+
+class GraderRequest(BaseModel):
+    session_id: str | None = Field(None)
+
+
+# ── Endpoints ────────────────────────────────────────────────────────────────
+
+@app.get("/health")
+async def health():
+    return {"status": "healthy"}
+
+
+@app.post("/reset")
+async def reset(req: ResetRequest) -> dict:
+    env = _get_env(req.session_id)
+    if req.session_id:
+        _environments[req.session_id] = env
+
+    obs = env.reset(task_id=req.task_id)
+    return {
+        "observation": obs.model_dump(),
+        "reward": obs.reward,
+        "done": obs.done,
+        "info": {"episode_id": env.state()["episode_id"]},
+    }
+
+
+@app.post("/step")
+async def step(req: StepRequest) -> dict:
+    env = _get_env(req.session_id)
+
+    try:
+        action_type = ActionType(req.action_type)
+    except ValueError:
+        raise HTTPException(
+            status_code=400,
+            detail=f"Invalid action_type: '{req.action_type}'. "
+            f"Valid actions: {[a.value for a in ActionType]}",
+        )
+
+    action = MLOpsAction(action_type=action_type, parameters=req.parameters)
+    obs = env.step(action)
+
+    return {
+        "observation": obs.model_dump(),
+        "reward": obs.reward,
+        "done": obs.done,
+        "info": env.state(),
+    }
+
+
+@app.get("/state")
+async def state(session_id: str | None = None) -> dict:
+    env = _get_env(session_id)
+    return env.state()
+
+
+@app.get("/tasks")
+async def list_tasks() -> dict:
+    """List all tasks with their action schema."""
+    tasks = []
+    for tid, t in ALL_TASKS.items():
+        tasks.append({
+            "task_id": t.task_id,
+            "name": t.name,
+            "difficulty": t.difficulty,
+            "description": t.description,
+            "max_steps": t.max_steps,
+        })
+
+    action_schema = {
+        "action_type": {
+            "type": "string",
+            "required": True,
+            "enum": [a.value for a in ActionType],
+            "description": "The action to perform",
+        },
+        "parameters": {
+            "type": "object",
+            "required": False,
+            "description": (
+                "Action-specific parameters. For adjust_threshold: "
+                "{'new_threshold': float}. For submit_diagnosis: "
+                "{'root_cause': str, 'summary': str}. For rollback_model: "
+                "{'target_version': str}."
+            ),
+        },
+    }
+
+    return {"tasks": tasks, "action_schema": action_schema}
+
+
+@app.post("/grader")
+async def grader(req: GraderRequest) -> dict:
+    """Return grader score for a completed episode."""
+    env = _get_env(req.session_id)
+    st = env.state()
+
+    if not st["done"]:
+        raise HTTPException(
+            status_code=400,
+            detail="Episode not complete. Finish the episode first.",
+        )
+
+    if st["grader_result"]:
+        return {
+            "score": st["grader_result"]["total"],
+            "breakdown": st["grader_result"],
+            "task_id": st["task_id"],
+        }
+
+    return {
+        "score": 0.0,
+        "breakdown": {},
+        "task_id": st["task_id"],
+        "message": "No grader result available.",
+    }
+
+
+@app.post("/baseline")
+async def run_baseline() -> dict:
+    """Run a scripted baseline agent on all 3 tasks and return scores."""
+    results = {}
+
+    for task_id, task_def in ALL_TASKS.items():
+        env = MLOpsFirefighterEnvironment()
+        env.reset(task_id=task_id)
+
+        # Baseline strategy: run all required diagnostics, then apply
+        # correct remediation, then submit diagnosis
+        for diag in task_def.required_diagnostics:
+            action = MLOpsAction(action_type=diag, parameters={})
+            env.step(action)
+
+        # Apply correct remediations
+        for rem in task_def.correct_remediations:
+            params = {}
+            if rem == ActionType.ADJUST_THRESHOLD:
+                params = {"new_threshold": task_def.extra_state.get("optimal_threshold", 0.5)}
+            elif rem == ActionType.ROLLBACK_MODEL:
+                params = {"target_version": task_def.model_info.previous_versions[0]}
+            action = MLOpsAction(action_type=rem, parameters=params)
+            env.step(action)
+
+        # Submit diagnosis
+        action = MLOpsAction(
+            action_type=ActionType.SUBMIT_DIAGNOSIS,
+            parameters={
+                "root_cause": task_def.root_causes[0],
+                "summary": f"Baseline diagnosis for {task_def.name}",
+            },
+        )
+        env.step(action)
+
+        st = env.state()
+        results[task_id] = {
+            "task_name": task_def.name,
+            "difficulty": task_def.difficulty,
+            "score": st["grader_result"]["total"] if st["grader_result"] else 0.0,
+            "breakdown": st["grader_result"],
+            "steps_taken": st["step_count"],
+        }
+
+    # Average score
+    scores = [r["score"] for r in results.values()]
+    avg = sum(scores) / len(scores) if scores else 0.0
+
+    return {
+        "baseline_results": results,
+        "average_score": round(avg, 3),
+    }
+
+
+# ── Main ─────────────────────────────────────────────────────────────────────
+
+if __name__ == "__main__":
+    import uvicorn
+
+    port = int(os.environ.get("PORT", 7860))
+    uvicorn.run(app, host="0.0.0.0", port=port)

server/environment.py

@@ -0,0 +1,334 @@
+# Copyright (c) 2026 MLOps Firefighter Contributors
+# Licensed under the BSD-3-Clause License
+
+"""
+Core environment logic for the MLOps Firefighter.
+
+Implements the OpenEnv Environment interface: reset(), step(), state().
+"""
+
+from __future__ import annotations
+
+import uuid
+from typing import Any
+
+from models import (
+    ActionType,
+    AlertSeverity,
+    MLOpsAction,
+    MLOpsObservation,
+)
+from tasks import ALL_TASKS, TaskDefinition, grade_episode
+
+
+class MLOpsFirefighterEnvironment:
+    """
+    Simulates an ML model failing in production.
+
+    The agent must diagnose the root cause and apply the correct fix
+    through a sequence of diagnostic and remediation actions.
+    """
+
+    def __init__(self) -> None:
+        self._episode_id: str = ""
+        self._step_count: int = 0
+        self._task: TaskDefinition | None = None
+        self._done: bool = False
+        self._cumulative_reward: float = 0.0
+
+        # Tracking for grading
+        self._actions_taken: list[dict[str, Any]] = []
+        self._diagnostics_gathered: list[str] = []
+        self._diagnosis_submitted: dict[str, Any] | None = None
+        self._remediations_applied: list[str] = []
+        self._last_grader_result: dict | None = None
+
+    # ── OpenEnv Interface ────────────────────────────────────────────────
+
+    def reset(self, task_id: str | None = None) -> MLOpsObservation:
+        """Initialize a new incident episode."""
+        self._episode_id = str(uuid.uuid4())
+        self._step_count = 0
+        self._done = False
+        self._cumulative_reward = 0.0
+        self._actions_taken = []
+        self._diagnostics_gathered = []
+        self._diagnosis_submitted = None
+        self._remediations_applied = []
+        self._last_grader_result = None
+
+        # Pick task
+        if task_id and task_id in ALL_TASKS:
+            self._task = ALL_TASKS[task_id]
+        else:
+            # Default to easy task
+            self._task = list(ALL_TASKS.values())[0]
+
+        return MLOpsObservation(
+            done=False,
+            reward=0.0,
+            step_number=0,
+            max_steps=self._task.max_steps,
+            task_id=self._task.task_id,
+            task_description=self._task.description,
+            alerts=self._task.alerts,
+            model_info=self._task.model_info,
+            action_result="Incident assigned to you. Begin investigation.",
+            action_success=True,
+            diagnostics_gathered=[],
+            available_actions=[a.value for a in ActionType],
+        )
+
+    def step(self, action: MLOpsAction) -> MLOpsObservation:
+        """Execute an action and return the resulting observation."""
+        if self._task is None:
+            return MLOpsObservation(
+                done=True,
+                reward=-1.0,
+                action_result="Error: No task loaded. Call reset() first.",
+                action_success=False,
+            )
+
+        if self._done:
+            return MLOpsObservation(
+                done=True,
+                reward=0.0,
+                step_number=self._step_count,
+                max_steps=self._task.max_steps,
+                task_id=self._task.task_id,
+                task_description=self._task.description,
+                action_result="Episode already ended.",
+                action_success=False,
+                diagnostics_gathered=self._diagnostics_gathered,
+                available_actions=[],
+            )
+
+        self._step_count += 1
+        self._actions_taken.append({
+            "action_type": action.action_type.value,
+            "parameters": action.parameters,
+            "step": self._step_count,
+        })
+
+        # Process the action
+        result_text, reward, action_success = self._process_action(action)
+
+        # Check for episode end
+        if action.action_type == ActionType.SUBMIT_DIAGNOSIS:
+            self._done = True
+            # Run grader for final reward
+            final_score, breakdown = grade_episode(
+                task=self._task,
+                actions_taken=self._actions_taken,
+                diagnosis_submitted=self._diagnosis_submitted,
+                remediation_applied=self._remediations_applied,
+                total_steps=self._step_count,
+            )
+            self._last_grader_result = breakdown
+            # Scale final reward: bonus for good diagnosis
+            reward = final_score * 5.0  # 0–5 range for final step
+            result_text += f"\n\n── EPISODE COMPLETE ──\nFinal Score: {final_score:.3f}\nBreakdown: {breakdown}"
+
+        # Check timeout
+        if self._step_count >= self._task.max_steps and not self._done:
+            self._done = True
+            final_score, breakdown = grade_episode(
+                task=self._task,
+                actions_taken=self._actions_taken,
+                diagnosis_submitted=self._diagnosis_submitted,
+                remediation_applied=self._remediations_applied,
+                total_steps=self._step_count,
+            )
+            self._last_grader_result = breakdown
+            reward = final_score * 3.0  # Lower multiplier for timeout
+            result_text += (
+                f"\n\n── TIMEOUT — Episode ended (max {self._task.max_steps} steps) ──\n"
+                f"Score: {final_score:.3f}\nBreakdown: {breakdown}"
+            )
+
+        self._cumulative_reward += reward
+
+        return MLOpsObservation(
+            done=self._done,
+            reward=round(reward, 4),
+            step_number=self._step_count,
+            max_steps=self._task.max_steps,
+            task_id=self._task.task_id,
+            task_description=self._task.description,
+            alerts=self._task.alerts,
+            model_info=self._task.model_info,
+            action_result=result_text,
+            action_success=action_success,
+            diagnostics_gathered=self._diagnostics_gathered,
+            available_actions=(
+                [a.value for a in ActionType] if not self._done else []
+            ),
+        )
+
+    def state(self) -> dict[str, Any]:
+        """Return current environment state."""
+        return {
+            "episode_id": self._episode_id,
+            "step_count": self._step_count,
+            "task_id": self._task.task_id if self._task else None,
+            "done": self._done,
+            "cumulative_reward": round(self._cumulative_reward, 4),
+            "actions_taken_count": len(self._actions_taken),
+            "diagnostics_gathered": self._diagnostics_gathered,
+            "remediations_applied": self._remediations_applied,
+            "diagnosis_submitted": self._diagnosis_submitted is not None,
+            "grader_result": self._last_grader_result,
+        }
+
+    # ── Action Processing ────────────────────────────────────────────────
+
+    def _process_action(self, action: MLOpsAction) -> tuple[str, float, bool]:
+        """
+        Process an action and return (result_text, step_reward, success).
+
+        Reward shaping:
+          - Useful diagnostic: +0.3
+          - Redundant diagnostic: +0.05
+          - Correct remediation: +1.0
+          - Wrong remediation: -0.5
+          - Destructive action without diagnosis: -1.0
+          - Submit diagnosis: handled in step() via grader
+        """
+        task = self._task
+        assert task is not None
+
+        at = action.action_type
+
+        # ── Diagnostic actions ───────────────────────────────────────────
+        if at in (
+            ActionType.INSPECT_METRICS,
+            ActionType.QUERY_LOGS,
+            ActionType.CHECK_DATA_DISTRIBUTION,
+            ActionType.CHECK_FEATURE_IMPORTANCE,
+            ActionType.RUN_PREDICTION_SAMPLE,
+            ActionType.CHECK_INFRASTRUCTURE,
+            ActionType.CHECK_UPSTREAM_PIPELINE,
+        ):
+            result_text = task.diagnostic_results.get(
+                at, "No additional information available for this diagnostic."
+            )
+            # Reward: useful if it's a required diagnostic we haven't done
+            diag_label = f"[{at.value}] {result_text[:80]}..."
+            if diag_label in self._diagnostics_gathered:
+                # Redundant — already gathered
+                reward = 0.05
+                result_text = "(You already ran this diagnostic.)\n\n" + result_text
+            else:
+                self._diagnostics_gathered.append(diag_label)
+                if at in task.required_diagnostics:
+                    reward = 0.3  # Useful diagnostic
+                else:
+                    reward = 0.1  # Valid but not critical
+            return result_text, reward, True
+
+        # ── Remediation actions ──────────────────────────────────────────
+        if at in (
+            ActionType.ROLLBACK_MODEL,
+            ActionType.ADJUST_THRESHOLD,
+            ActionType.RETRAIN_MODEL,
+            ActionType.FIX_DATA_PIPELINE,
+            ActionType.SCALE_INFRASTRUCTURE,
+            ActionType.ADD_FEATURE_GUARD,
+        ):
+            return self._process_remediation(at, action.parameters)
+
+        # ── Submit diagnosis ─────────────────────────────────────────────
+        if at == ActionType.SUBMIT_DIAGNOSIS:
+            self._diagnosis_submitted = action.parameters
+            root_cause = action.parameters.get("root_cause", "not specified")
+            summary = action.parameters.get("summary", "no summary")
+            return (
+                f"Diagnosis submitted.\n"
+                f"  Root cause: {root_cause}\n"
+                f"  Summary: {summary}",
+                0.0,  # Actual reward computed in step() via grader
+                True,
+            )
+
+        return "Unknown action type.", -0.2, False
+
+    def _process_remediation(
+        self, action_type: ActionType, params: dict[str, Any]
+    ) -> tuple[str, float, bool]:
+        """Process a remediation action."""
+        task = self._task
+        assert task is not None
+
+        self._remediations_applied.append(action_type.value)
+
+        is_correct = action_type in task.correct_remediations
+        has_diagnosed = len(self._diagnostics_gathered) >= 2
+
+        # Penalize hasty remediation without diagnosis
+        if not has_diagnosed:
+            return (
+                f"⚠ WARNING: Applying {action_type.value} without sufficient "
+                f"diagnostic investigation. This is risky in production.\n"
+                f"Action applied, but confidence is low.",
+                -0.5,
+                True,
+            )
+
+        if is_correct:
+            # Specific feedback per remediation
+            feedback = self._get_remediation_feedback(action_type, params)
+            return feedback, 1.0, True
+        else:
+            return (
+                f"Applied {action_type.value}, but this doesn't address the "
+                f"root cause. The issue persists. Consider more investigation.",
+                -0.3,
+                True,
+            )
+
+    def _get_remediation_feedback(
+        self, action_type: ActionType, params: dict[str, Any]
+    ) -> str:
+        """Generate specific feedback for correct remediations."""
+        task = self._task
+        assert task is not None
+
+        if action_type == ActionType.ADJUST_THRESHOLD:
+            new_t = params.get("new_threshold", "unspecified")
+            optimal = task.extra_state.get("optimal_threshold", 0.5)
+            return (
+                f"Threshold adjusted to {new_t}.\n"
+                f"Optimal threshold was {optimal}.\n"
+                f"Precision recovering. False positive rate decreasing."
+            )
+        elif action_type == ActionType.ROLLBACK_MODEL:
+            target = params.get("target_version", "previous")
+            return (
+                f"Model rolled back to {target}.\n"
+                f"Previous model restored. Monitoring metrics.\n"
+                f"Harmful content detection recovering."
+            )
+        elif action_type == ActionType.FIX_DATA_PIPELINE:
+            return (
+                "Data pipeline fix initiated.\n"
+                "  - Credit bureau connector patched\n"
+                "  - Unit conversion (cents→dollars) applied to annual_income\n"
+                "  - Null handling added for credit_utilization_ratio\n"
+                "  Pipeline revalidation in progress."
+            )
+        elif action_type == ActionType.RETRAIN_MODEL:
+            return (
+                "Model retraining triggered with corrected data.\n"
+                "  - Stale features refreshed\n"
+                "  - Estimated completion: 2 hours\n"
+                "  - Will auto-deploy after quality gate passes."
+            )
+        elif action_type == ActionType.ADD_FEATURE_GUARD:
+            return (
+                "Feature guard / input validation added.\n"
+                "  - Per-class recall monitoring enabled\n"
+                "  - Crowd-source batch validation gate added\n"
+                "  - Anomalous label distributions will trigger review."
+            )
+        else:
+            return f"{action_type.value} applied successfully."

tasks.py

@@ -0,0 +1,526 @@
+# Copyright (c) 2026 MLOps Firefighter Contributors
+# Licensed under the BSD-3-Clause License
+
+"""
+Task definitions for the MLOps Firefighter environment.
+
+Each task defines:
+  - A production ML incident scenario
+  - The ground-truth root cause(s)
+  - Required diagnostic steps (for partial credit)
+  - Correct remediation action(s)
+  - A grader function that scores agent performance 0.0–1.0
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from typing import Any
+
+from models import (
+    ActionType,
+    Alert,
+    AlertSeverity,
+    ModelInfo,
+)
+
+
+# ── Task Data Structure ─────────────────────────────────────────────────────
+
+@dataclass
+class TaskDefinition:
+    task_id: str
+    name: str
+    difficulty: str  # "easy", "medium", "hard"
+    description: str
+    alerts: list[Alert]
+    model_info: ModelInfo
+    root_causes: list[str]  # acceptable root cause labels
+    required_diagnostics: list[ActionType]  # expected investigation steps
+    correct_remediations: list[ActionType]  # accepted fix actions
+    diagnostic_results: dict[ActionType, str]  # what each diagnostic reveals
+    max_steps: int = 20
+    extra_state: dict[str, Any] = field(default_factory=dict)
+
+
+# ── Task 1: Easy — Threshold Misconfiguration ───────────────────────────────
+
+TASK_EASY = TaskDefinition(
+    task_id="task_threshold_drift",
+    name="Threshold Misconfiguration After Redeployment",
+    difficulty="easy",
+    description=(
+        "INCIDENT: The fraud detection model's precision has dropped from 0.94 to 0.61 "
+        "over the last 2 hours. False positive rate spiked 3x. The model was redeployed "
+        "yesterday with a new version. No changes to training data or features were made. "
+        "Investigate and fix the issue."
+    ),
+    alerts=[
+        Alert(
+            severity=AlertSeverity.HIGH,
+            message="Precision drop detected on fraud-detection-v3",
+            metric_name="precision",
+            current_value=0.61,
+            threshold=0.85,
+            timestamp="2026-03-27T10:15:00Z",
+        ),
+        Alert(
+            severity=AlertSeverity.MEDIUM,
+            message="False positive rate above acceptable range",
+            metric_name="false_positive_rate",
+            current_value=0.18,
+            threshold=0.05,
+            timestamp="2026-03-27T10:17:00Z",
+        ),
+    ],
+    model_info=ModelInfo(
+        model_name="fraud-detection",
+        model_version="v3.1.0",
+        deployed_at="2026-03-26T14:00:00Z",
+        framework="XGBoost",
+        endpoint="/api/v1/predict/fraud",
+        previous_versions=["v3.0.2", "v2.9.1"],
+    ),
+    root_causes=["threshold_misconfiguration", "threshold_too_low", "bad_threshold"],
+    required_diagnostics=[
+        ActionType.INSPECT_METRICS,
+        ActionType.RUN_PREDICTION_SAMPLE,
+    ],
+    correct_remediations=[ActionType.ADJUST_THRESHOLD],
+    diagnostic_results={
+        ActionType.INSPECT_METRICS: (
+            "METRICS DASHBOARD:\n"
+            "  Model: fraud-detection v3.1.0\n"
+            "  Precision: 0.61 (was 0.94 on v3.0.2)\n"
+            "  Recall: 0.99 (was 0.87 on v3.0.2)\n"
+            "  F1: 0.76 (was 0.90 on v3.0.2)\n"
+            "  Decision threshold: 0.30 (was 0.55 on v3.0.2)\n"
+            "  NOTE: Threshold was changed during redeployment.\n"
+            "  Requests/sec: 1,240 (normal)\n"
+            "  P99 latency: 45ms (normal)"
+        ),
+        ActionType.QUERY_LOGS: (
+            "LOG SEARCH RESULTS (last 4 hours):\n"
+            "  [10:00] Deployment v3.1.0 started\n"
+            "  [10:01] Config applied: threshold=0.30 (previous: 0.55)\n"
+            "  [10:02] Health check passed\n"
+            "  [10:15] Alert: precision below 0.85\n"
+            "  No error logs. No OOM events. Pipeline healthy."
+        ),
+        ActionType.CHECK_DATA_DISTRIBUTION: (
+            "DATA DISTRIBUTION COMPARISON:\n"
+            "  Training data: 2.1M samples, fraud_rate=3.2%\n"
+            "  Serving data (last 24h): 148K requests, fraud_rate=3.1%\n"
+            "  Feature distributions: no significant drift detected\n"
+            "  KL divergence: 0.003 (threshold: 0.05)\n"
+            "  All features within expected ranges."
+        ),
+        ActionType.CHECK_FEATURE_IMPORTANCE: (
+            "FEATURE IMPORTANCE (top 5):\n"
+            "  1. transaction_amount: 0.31\n"
+            "  2. merchant_risk_score: 0.22\n"
+            "  3. velocity_1h: 0.18\n"
+            "  4. distance_from_home: 0.14\n"
+            "  5. device_fingerprint_match: 0.09\n"
+            "  No anomalies in feature importance vs. training."
+        ),
+        ActionType.RUN_PREDICTION_SAMPLE: (
+            "PREDICTION SAMPLE TEST (100 known-labeled samples):\n"
+            "  At threshold=0.30: precision=0.62, recall=0.98\n"
+            "  At threshold=0.55: precision=0.93, recall=0.88\n"
+            "  At threshold=0.50: precision=0.91, recall=0.90\n"
+            "  Conclusion: threshold is the primary lever for precision."
+        ),
+        ActionType.CHECK_INFRASTRUCTURE: (
+            "INFRASTRUCTURE STATUS:\n"
+            "  CPU: 34% (normal)\n  Memory: 4.2GB/8GB (normal)\n"
+            "  GPU: N/A (CPU-only model)\n  Latency P99: 45ms (normal)\n"
+            "  No infrastructure issues detected."
+        ),
+        ActionType.CHECK_UPSTREAM_PIPELINE: (
+            "UPSTREAM PIPELINE:\n"
+            "  Feature store: healthy, last refresh 12 min ago\n"
+            "  Data ingestion: nominal, 0 failures in 24h\n"
+            "  Schema validation: passing\n"
+            "  No upstream issues."
+        ),
+    },
+    max_steps=15,
+    extra_state={"optimal_threshold": 0.55},
+)
+
+
+# ── Task 2: Medium — Data Drift + Stale Feature ─────────────────────────────
+
+TASK_MEDIUM = TaskDefinition(
+    task_id="task_data_drift",
+    name="Data Drift with Stale Feature Pipeline",
+    difficulty="medium",
+    description=(
+        "INCIDENT: The loan default prediction model's AUC has degraded from 0.91 to "
+        "0.74 over the past week. Customer complaints about incorrect rejections have "
+        "tripled. The model was last retrained 4 months ago. A new credit bureau data "
+        "provider was onboarded 10 days ago. Diagnose the issue and fix it."
+    ),
+    alerts=[
+        Alert(
+            severity=AlertSeverity.CRITICAL,
+            message="AUC degradation on loan-default-predictor",
+            metric_name="auc_roc",
+            current_value=0.74,
+            threshold=0.85,
+            timestamp="2026-03-27T08:00:00Z",
+        ),
+        Alert(
+            severity=AlertSeverity.HIGH,
+            message="Customer complaint rate above threshold",
+            metric_name="complaint_rate",
+            current_value=0.12,
+            threshold=0.03,
+            timestamp="2026-03-27T09:30:00Z",
+        ),
+    ],
+    model_info=ModelInfo(
+        model_name="loan-default-predictor",
+        model_version="v2.4.0",
+        deployed_at="2025-11-15T10:00:00Z",
+        framework="LightGBM",
+        endpoint="/api/v1/predict/loan-default",
+        previous_versions=["v2.3.1", "v2.2.0"],
+    ),
+    root_causes=[
+        "data_drift",
+        "stale_feature",
+        "feature_pipeline_broken",
+        "data_distribution_shift",
+    ],
+    required_diagnostics=[
+        ActionType.CHECK_DATA_DISTRIBUTION,
+        ActionType.CHECK_UPSTREAM_PIPELINE,
+        ActionType.INSPECT_METRICS,
+    ],
+    correct_remediations=[
+        ActionType.FIX_DATA_PIPELINE,
+        ActionType.RETRAIN_MODEL,
+    ],
+    diagnostic_results={
+        ActionType.INSPECT_METRICS: (
+            "METRICS DASHBOARD:\n"
+            "  Model: loan-default-predictor v2.4.0\n"
+            "  AUC: 0.74 (was 0.91 at deploy, 0.88 last week)\n"
+            "  Accuracy: 0.69 (was 0.84)\n"
+            "  False rejection rate: 14% (was 4%)\n"
+            "  Threshold: 0.50 (unchanged)\n"
+            "  Requests/sec: 320 (normal)\n"
+            "  Degradation trend: gradual decline starting ~10 days ago"
+        ),
+        ActionType.QUERY_LOGS: (
+            "LOG SEARCH RESULTS (last 14 days):\n"
+            "  [Mar 17] New credit bureau connector deployed (v1.2)\n"
+            "  [Mar 17] Feature 'credit_utilization_ratio' source changed\n"
+            "  [Mar 18] Warning: 2,340 null values in 'credit_utilization_ratio'\n"
+            "  [Mar 19-27] Recurring nulls in credit_utilization_ratio (avg 8%/day)\n"
+            "  [Mar 22] Warning: feature 'annual_income' format changed (cents→dollars)\n"
+            "  No deployment events. No model changes."
+        ),
+        ActionType.CHECK_DATA_DISTRIBUTION: (
+            "DATA DISTRIBUTION COMPARISON:\n"
+            "  Training data vs. Serving data (last 7 days):\n"
+            "  ─────────────────────────────────────────────\n"
+            "  credit_utilization_ratio:\n"
+            "    Train: mean=0.34, std=0.21, null_rate=0.1%\n"
+            "    Serve: mean=0.08, std=0.42, null_rate=8.3%  ⚠ DRIFT DETECTED\n"
+            "    KL divergence: 0.87 (threshold: 0.05)\n"
+            "  annual_income:\n"
+            "    Train: mean=65,400, std=28,000\n"
+            "    Serve: mean=6,540,000, std=2,800,000  ⚠ DRIFT DETECTED\n"
+            "    KL divergence: 12.4 (threshold: 0.05)\n"
+            "    NOTE: Values appear 100x larger — possible unit change\n"
+            "  Other features: within normal ranges"
+        ),
+        ActionType.CHECK_FEATURE_IMPORTANCE: (
+            "FEATURE IMPORTANCE (top 5):\n"
+            "  1. credit_utilization_ratio: 0.28  ← AFFECTED\n"
+            "  2. annual_income: 0.22  ← AFFECTED\n"
+            "  3. debt_to_income: 0.19\n"
+            "  4. payment_history_score: 0.15\n"
+            "  5. employment_length: 0.08\n"
+            "  The top 2 features (50% of importance) are drifted."
+        ),
+        ActionType.RUN_PREDICTION_SAMPLE: (
+            "PREDICTION SAMPLE (50 recent applications with known outcomes):\n"
+            "  Model predictions vs. actual outcomes:\n"
+            "  - Correct: 31/50 (62%)\n"
+            "  - False rejections: 14/50 (28%) — most involve high-income applicants\n"
+            "  - Missed defaults: 5/50 (10%)\n"
+            "  Annual income feature appears corrupted (values in cents not dollars)"
+        ),
+        ActionType.CHECK_INFRASTRUCTURE: (
+            "INFRASTRUCTURE STATUS:\n"
+            "  CPU: 22% (normal)\n  Memory: 3.1GB/8GB (normal)\n"
+            "  Latency P99: 38ms (normal)\n"
+            "  No infrastructure issues."
+        ),
+        ActionType.CHECK_UPSTREAM_PIPELINE: (
+            "UPSTREAM PIPELINE STATUS:\n"
+            "  Feature store: HEALTHY (but stale for 2 features)\n"
+            "  Data ingestion: DEGRADED\n"
+            "    - credit_bureau_connector_v1.2: returning nulls 8% of time\n"
+            "    - Schema mismatch: 'annual_income' was USD→now returned in cents\n"
+            "    - Last successful full refresh: 11 days ago\n"
+            "  Other connectors: nominal\n"
+            "  RECOMMENDATION: Fix credit bureau connector, add unit conversion"
+        ),
+    },
+    max_steps=20,
+    extra_state={
+        "requires_pipeline_fix": True,
+        "requires_retrain": True,
+    },
+)
+
+
+# ── Task 3: Hard — Adversarial + Silent Model Regression ────────────────────
+
+TASK_HARD = TaskDefinition(
+    task_id="task_silent_regression",
+    name="Silent Model Regression with Adversarial Inputs",
+    difficulty="hard",
+    description=(
+        "INCIDENT: The content moderation model appears to be working normally — overall "
+        "accuracy is 0.96. However, trust & safety reports show a 5x increase in "
+        "user-reported harmful content slipping through. The model was retrained 3 days ago "
+        "with an automated pipeline. No alerts have fired because aggregate metrics look "
+        "fine. A community moderator escalated this manually. Investigate what's happening "
+        "and fix it. Be careful — hasty actions could take down moderation for millions of "
+        "users."
+    ),
+    alerts=[
+        Alert(
+            severity=AlertSeverity.LOW,
+            message="Manual escalation: T&S reports 5x harmful content increase",
+            metric_name="user_reports_harmful",
+            current_value=847.0,
+            threshold=200.0,
+            timestamp="2026-03-27T06:00:00Z",
+        ),
+    ],
+    model_info=ModelInfo(
+        model_name="content-moderator",
+        model_version="v5.2.0",
+        deployed_at="2026-03-24T03:00:00Z",
+        framework="PyTorch (transformer)",
+        endpoint="/api/v1/moderate/content",
+        previous_versions=["v5.1.3", "v5.0.0", "v4.8.2"],
+    ),
+    root_causes=[
+        "training_data_poisoning",
+        "adversarial_inputs",
+        "label_corruption",
+        "data_poisoning",
+        "silent_regression",
+        "class_specific_regression",
+    ],
+    required_diagnostics=[
+        ActionType.INSPECT_METRICS,
+        ActionType.RUN_PREDICTION_SAMPLE,
+        ActionType.CHECK_DATA_DISTRIBUTION,
+        ActionType.CHECK_UPSTREAM_PIPELINE,
+    ],
+    correct_remediations=[
+        ActionType.ROLLBACK_MODEL,
+        ActionType.ADD_FEATURE_GUARD,
+    ],
+    diagnostic_results={
+        ActionType.INSPECT_METRICS: (
+            "METRICS DASHBOARD:\n"
+            "  Model: content-moderator v5.2.0\n"
+            "  Overall accuracy: 0.96 (looks normal!)\n"
+            "  Overall precision: 0.94, recall: 0.93\n"
+            "  Latency P99: 120ms (normal)\n"
+            "  Requests/sec: 45,000 (normal)\n"
+            "  ── Per-class breakdown ──\n"
+            "  spam:      precision=0.97, recall=0.98 (normal)\n"
+            "  hate:      precision=0.95, recall=0.42 ⚠ (was 0.91)\n"
+            "  violence:  precision=0.93, recall=0.38 ⚠ (was 0.89)\n"
+            "  safe:      precision=0.96, recall=0.99 (normal)\n"
+            "  NOTE: Hate and violence recall collapsed but aggregate looks OK\n"
+            "  because 'safe' class dominates (92% of traffic)."
+        ),
+        ActionType.QUERY_LOGS: (
+            "LOG SEARCH RESULTS:\n"
+            "  [Mar 24 02:30] Auto-retrain pipeline triggered\n"
+            "  [Mar 24 02:45] Training dataset assembled: 2.4M samples\n"
+            "  [Mar 24 02:46] Warning: 34,000 new labels from crowd-source batch #847\n"
+            "  [Mar 24 03:00] Model v5.2.0 deployed (passed aggregate quality gate)\n"
+            "  [Mar 24 03:01] A/B test: v5.2.0 vs v5.1.3 — aggregate acc comparable\n"
+            "  [Mar 25] T&S team flags increase in harmful content reports\n"
+            "  [Mar 26] Reports continue to climb\n"
+            "  [Mar 27] Manual escalation from community moderator"
+        ),
+        ActionType.CHECK_DATA_DISTRIBUTION: (
+            "DATA DISTRIBUTION COMPARISON:\n"
+            "  Training data v5.2.0 vs v5.1.3:\n"
+            "  ─────────────────────────────────────────────\n"
+            "  Total samples: 2.4M (v5.2.0) vs 2.1M (v5.1.3)\n"
+            "  New batch #847: 34,000 samples added\n"
+            "  Label distribution in batch #847:\n"
+            "    safe: 89% | hate: 7% | violence: 4% | spam: 0%\n"
+            "  Cross-check with manual audit (100 random from batch #847):\n"
+            "    Labeled 'safe' but actually hateful: 23/100 ⚠ CONTAMINATION\n"
+            "    Labeled 'safe' but actually violent: 11/100 ⚠ CONTAMINATION\n"
+            "    Correctly labeled: 66/100\n"
+            "  CONCLUSION: Batch #847 contains systematically mislabeled data.\n"
+            "  Hate/violence content was labeled as 'safe', teaching the model\n"
+            "  to pass through harmful content."
+        ),
+        ActionType.CHECK_FEATURE_IMPORTANCE: (
+            "ATTENTION ANALYSIS (transformer model):\n"
+            "  Attention patterns on hate/violence content:\n"
+            "  v5.1.3: High attention on slurs, threats, graphic terms\n"
+            "  v5.2.0: Attention dispersed, less focus on key harmful tokens\n"
+            "  The model appears to have 'unlearned' key harmful patterns\n"
+            "  from the poisoned training batch."
+        ),
+        ActionType.RUN_PREDICTION_SAMPLE: (
+            "PREDICTION SAMPLE (200 curated test samples, ground truth labels):\n"
+            "  Category-level results:\n"
+            "  ── Safe content (100 samples) ──\n"
+            "    Correct: 98/100 (98%) — model correctly passes safe content\n"
+            "  ── Hate speech (50 samples) ──\n"
+            "    Detected: 21/50 (42%) ⚠ — model misses most hate speech\n"
+            "    Missed examples include: coded language, dog-whistles, slurs\n"
+            "  ── Violent content (50 samples) ──\n"
+            "    Detected: 19/50 (38%) ⚠ — model misses most violence\n"
+            "    Missed examples: graphic threats, incitement\n"
+            "  Overall accuracy inflated by safe-class dominance."
+        ),
+        ActionType.CHECK_INFRASTRUCTURE: (
+            "INFRASTRUCTURE STATUS:\n"
+            "  GPU: 4x A100, utilization 67% (normal)\n"
+            "  Memory: 28GB/80GB (normal)\n"
+            "  Latency P99: 120ms (normal)\n"
+            "  Throughput: 45K req/s (normal)\n"
+            "  No infrastructure issues."
+        ),
+        ActionType.CHECK_UPSTREAM_PIPELINE: (
+            "UPSTREAM PIPELINE STATUS:\n"
+            "  Auto-retrain pipeline: HEALTHY (but concern flagged)\n"
+            "  Training data sources:\n"
+            "    - Internal labeled data: 2.1M samples (validated)\n"
+            "    - Crowd-source batch #847: 34K samples (UNVALIDATED) ⚠\n"
+            "  Quality gate: aggregate accuracy only — no per-class checks ⚠\n"
+            "  The pipeline accepted v5.2.0 because overall accuracy was 0.96,\n"
+            "  but it did not check per-class recall.\n"
+            "  RECOMMENDATION: Add per-class recall gates. Remove batch #847.\n"
+            "  Consider rollback to v5.1.3 while investigating."
+        ),
+    },
+    max_steps=25,
+    extra_state={
+        "poisoned_batch": "#847",
+        "safe_rollback_version": "v5.1.3",
+        "requires_guard": True,
+    },
+)
+
+
+# ── Task Registry ────────────────────────────────────────────────────────────
+
+ALL_TASKS: dict[str, TaskDefinition] = {
+    TASK_EASY.task_id: TASK_EASY,
+    TASK_MEDIUM.task_id: TASK_MEDIUM,
+    TASK_HARD.task_id: TASK_HARD,
+}
+
+
+# ── Grader Functions ─────────────────────────────────────────────────────────
+
+def grade_episode(
+    task: TaskDefinition,
+    actions_taken: list[dict],
+    diagnosis_submitted: dict | None,
+    remediation_applied: list[str],
+    total_steps: int,
+) -> tuple[float, dict]:
+    """
+    Grade an agent's performance on a task.
+
+    Returns:
+        (score, breakdown) where score is 0.0–1.0 and breakdown is a dict
+        of component scores for interpretability.
+    """
+    breakdown: dict[str, float] = {}
+
+    # ── 1. Diagnostic thoroughness (30%) ─────────────────────────────────
+    required = set(task.required_diagnostics)
+    performed = set()
+    for a in actions_taken:
+        try:
+            at = ActionType(a["action_type"])
+            if at in required:
+                performed.add(at)
+        except (ValueError, KeyError):
+            pass
+
+    if required:
+        diag_score = len(performed) / len(required)
+    else:
+        diag_score = 1.0
+    breakdown["diagnostic_thoroughness"] = round(diag_score, 3)
+
+    # ── 2. Correct diagnosis (30%) ───────────────────────────────────────
+    diag_submitted_score = 0.0
+    if diagnosis_submitted:
+        root = diagnosis_submitted.get("root_cause", "").lower().strip()
+        for valid in task.root_causes:
+            if valid.lower() in root or root in valid.lower():
+                diag_submitted_score = 1.0
+                break
+        # Partial credit: if they mention a related keyword
+        if diag_submitted_score == 0.0:
+            keywords = set()
+            for rc in task.root_causes:
+                keywords.update(rc.lower().replace("_", " ").split())
+            matches = sum(1 for kw in keywords if kw in root)
+            if matches > 0:
+                diag_submitted_score = min(0.5, matches * 0.2)
+    breakdown["diagnosis_accuracy"] = round(diag_submitted_score, 3)
+
+    # ── 3. Correct remediation (25%) ─────────────────────────────────────
+    correct_rems = set(r.value for r in task.correct_remediations)
+    applied = set(remediation_applied)
+    if correct_rems:
+        matched = len(applied & correct_rems)
+        rem_score = matched / len(correct_rems)
+        # Penalize wrong remediations
+        wrong = applied - correct_rems
+        penalty = len(wrong) * 0.15
+        rem_score = max(0.0, rem_score - penalty)
+    else:
+        rem_score = 1.0 if not applied else 0.5
+    breakdown["remediation_accuracy"] = round(rem_score, 3)
+
+    # ── 4. Efficiency (15%) ──────────────────────────────────────────────
+    if total_steps <= len(task.required_diagnostics) + 2:
+        eff_score = 1.0  # Very efficient
+    elif total_steps <= task.max_steps * 0.5:
+        eff_score = 0.8
+    elif total_steps <= task.max_steps * 0.75:
+        eff_score = 0.5
+    elif total_steps < task.max_steps:
+        eff_score = 0.3
+    else:
+        eff_score = 0.1  # Timed out
+    breakdown["efficiency"] = round(eff_score, 3)
+
+    # ── Weighted total ───────────────────────────────────────────────────
+    total = (
+        0.30 * breakdown["diagnostic_thoroughness"]
+        + 0.30 * breakdown["diagnosis_accuracy"]
+        + 0.25 * breakdown["remediation_accuracy"]
+        + 0.15 * breakdown["efficiency"]
+    )
+    breakdown["total"] = round(total, 3)
+
+    return round(total, 3), breakdown

tests/test_environment.py

@@ -0,0 +1,251 @@
+#!/usr/bin/env python3
+"""Tests for the MLOps Firefighter environments."""
+
+import sys
+import os
+
+sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+
+from models import ActionType, MLOpsAction
+from server.environment import MLOpsFirefighterEnvironment
+from tasks import ALL_TASKS, grade_episode
+
+
+def test_reset_returns_valid_observation():
+    env = MLOpsFirefighterEnvironment()
+    obs = env.reset(task_id="task_threshold_drift")
+    assert obs.done is False
+    assert obs.step_number == 0
+    assert obs.task_id == "task_threshold_drift"
+    assert len(obs.alerts) > 0
+    assert obs.model_info is not None
+    assert len(obs.available_actions) > 0
+    print("✓ test_reset_returns_valid_observation")
+
+
+def test_step_diagnostic_action():
+    env = MLOpsFirefighterEnvironment()
+    env.reset(task_id="task_threshold_drift")
+    action = MLOpsAction(action_type=ActionType.INSPECT_METRICS)
+    obs = env.step(action)
+    assert obs.done is False
+    assert obs.reward > 0  # Useful diagnostic
+    assert obs.step_number == 1
+    assert "METRICS DASHBOARD" in obs.action_result
+    assert len(obs.diagnostics_gathered) == 1
+    print("✓ test_step_diagnostic_action")
+
+
+def test_redundant_diagnostic():
+    env = MLOpsFirefighterEnvironment()
+    env.reset(task_id="task_threshold_drift")
+    action = MLOpsAction(action_type=ActionType.INSPECT_METRICS)
+    env.step(action)
+    obs2 = env.step(action)
+    assert obs2.reward == 0.05  # Redundant
+    assert "already ran" in obs2.action_result.lower()
+    print("✓ test_redundant_diagnostic")
+
+
+def test_hasty_remediation_penalized():
+    env = MLOpsFirefighterEnvironment()
+    env.reset(task_id="task_threshold_drift")
+    # Apply fix without any diagnosis
+    action = MLOpsAction(
+        action_type=ActionType.ADJUST_THRESHOLD,
+        parameters={"new_threshold": 0.55},
+    )
+    obs = env.step(action)
+    assert obs.reward < 0  # Penalized
+    assert "WARNING" in obs.action_result
+    print("✓ test_hasty_remediation_penalized")
+
+
+def test_correct_remediation_after_diagnosis():
+    env = MLOpsFirefighterEnvironment()
+    env.reset(task_id="task_threshold_drift")
+    # Run 2 diagnostics
+    env.step(MLOpsAction(action_type=ActionType.INSPECT_METRICS))
+    env.step(MLOpsAction(action_type=ActionType.RUN_PREDICTION_SAMPLE))
+    # Apply correct fix
+    obs = env.step(MLOpsAction(
+        action_type=ActionType.ADJUST_THRESHOLD,
+        parameters={"new_threshold": 0.55},
+    ))
+    assert obs.reward == 1.0  # Correct remediation
+    print("✓ test_correct_remediation_after_diagnosis")
+
+
+def test_submit_diagnosis_ends_episode():
+    env = MLOpsFirefighterEnvironment()
+    env.reset(task_id="task_threshold_drift")
+    env.step(MLOpsAction(action_type=ActionType.INSPECT_METRICS))
+    env.step(MLOpsAction(action_type=ActionType.RUN_PREDICTION_SAMPLE))
+    env.step(MLOpsAction(
+        action_type=ActionType.ADJUST_THRESHOLD,
+        parameters={"new_threshold": 0.55},
+    ))
+    obs = env.step(MLOpsAction(
+        action_type=ActionType.SUBMIT_DIAGNOSIS,
+        parameters={"root_cause": "threshold_misconfiguration", "summary": "test"},
+    ))
+    assert obs.done is True
+    assert obs.reward > 0
+    print("✓ test_submit_diagnosis_ends_episode")
+
+
+def test_state_tracking():
+    env = MLOpsFirefighterEnvironment()
+    env.reset(task_id="task_threshold_drift")
+    env.step(MLOpsAction(action_type=ActionType.INSPECT_METRICS))
+    st = env.state()
+    assert st["step_count"] == 1
+    assert st["task_id"] == "task_threshold_drift"
+    assert st["done"] is False
+    assert st["actions_taken_count"] == 1
+    print("✓ test_state_tracking")
+
+
+def test_timeout():
+    env = MLOpsFirefighterEnvironment()
+    env.reset(task_id="task_threshold_drift")
+    # Exhaust all steps with useless actions
+    for _ in range(15):
+        obs = env.step(MLOpsAction(action_type=ActionType.CHECK_INFRASTRUCTURE))
+        if obs.done:
+            break
+    assert obs.done is True
+    print("✓ test_timeout")
+
+
+def test_grader_scores_in_range():
+    """Verify all grader scores are between 0.0 and 1.0."""
+    for task_id, task in ALL_TASKS.items():
+        # Perfect run
+        score_perfect, _ = grade_episode(
+            task=task,
+            actions_taken=[{"action_type": d.value} for d in task.required_diagnostics],
+            diagnosis_submitted={"root_cause": task.root_causes[0], "summary": "test"},
+            remediation_applied=[r.value for r in task.correct_remediations],
+            total_steps=len(task.required_diagnostics) + len(task.correct_remediations) + 1,
+        )
+        assert 0.0 <= score_perfect <= 1.0, f"Perfect score out of range: {score_perfect}"
+
+        # Zero run
+        score_zero, _ = grade_episode(
+            task=task,
+            actions_taken=[],
+            diagnosis_submitted=None,
+            remediation_applied=[],
+            total_steps=task.max_steps,
+        )
+        assert 0.0 <= score_zero <= 1.0, f"Zero score out of range: {score_zero}"
+
+    print("✓ test_grader_scores_in_range")
+
+
+def test_all_three_tasks_exist():
+    assert len(ALL_TASKS) >= 3
+    difficulties = {t.difficulty for t in ALL_TASKS.values()}
+    assert "easy" in difficulties
+    assert "medium" in difficulties
+    assert "hard" in difficulties
+    print("✓ test_all_three_tasks_exist")
+
+
+def test_full_episode_easy():
+    """Full integration test: perfect run on easy task."""
+    env = MLOpsFirefighterEnvironment()
+    env.reset(task_id="task_threshold_drift")
+
+    # Diagnostics
+    env.step(MLOpsAction(action_type=ActionType.INSPECT_METRICS))
+    env.step(MLOpsAction(action_type=ActionType.RUN_PREDICTION_SAMPLE))
+
+    # Fix
+    env.step(MLOpsAction(
+        action_type=ActionType.ADJUST_THRESHOLD,
+        parameters={"new_threshold": 0.55},
+    ))
+
+    # Submit
+    obs = env.step(MLOpsAction(
+        action_type=ActionType.SUBMIT_DIAGNOSIS,
+        parameters={"root_cause": "threshold_misconfiguration", "summary": "test"},
+    ))
+
+    assert obs.done is True
+    st = env.state()
+    score = st["grader_result"]["total"]
+    assert score >= 0.85, f"Expected high score for perfect run, got {score}"
+    print(f"✓ test_full_episode_easy (score={score:.3f})")
+
+
+def test_full_episode_medium():
+    """Full integration test: good run on medium task."""
+    env = MLOpsFirefighterEnvironment()
+    env.reset(task_id="task_data_drift")
+
+    env.step(MLOpsAction(action_type=ActionType.INSPECT_METRICS))
+    env.step(MLOpsAction(action_type=ActionType.CHECK_DATA_DISTRIBUTION))
+    env.step(MLOpsAction(action_type=ActionType.CHECK_UPSTREAM_PIPELINE))
+
+    env.step(MLOpsAction(action_type=ActionType.FIX_DATA_PIPELINE))
+    env.step(MLOpsAction(action_type=ActionType.RETRAIN_MODEL))
+
+    obs = env.step(MLOpsAction(
+        action_type=ActionType.SUBMIT_DIAGNOSIS,
+        parameters={"root_cause": "data_drift", "summary": "test"},
+    ))
+
+    assert obs.done is True
+    st = env.state()
+    score = st["grader_result"]["total"]
+    assert score >= 0.80, f"Expected high score, got {score}"
+    print(f"✓ test_full_episode_medium (score={score:.3f})")
+
+
+def test_full_episode_hard():
+    """Full integration test: good run on hard task."""
+    env = MLOpsFirefighterEnvironment()
+    env.reset(task_id="task_silent_regression")
+
+    env.step(MLOpsAction(action_type=ActionType.INSPECT_METRICS))
+    env.step(MLOpsAction(action_type=ActionType.RUN_PREDICTION_SAMPLE))
+    env.step(MLOpsAction(action_type=ActionType.CHECK_DATA_DISTRIBUTION))
+    env.step(MLOpsAction(action_type=ActionType.CHECK_UPSTREAM_PIPELINE))
+
+    env.step(MLOpsAction(
+        action_type=ActionType.ROLLBACK_MODEL,
+        parameters={"target_version": "v5.1.3"},
+    ))
+    env.step(MLOpsAction(action_type=ActionType.ADD_FEATURE_GUARD))
+
+    obs = env.step(MLOpsAction(
+        action_type=ActionType.SUBMIT_DIAGNOSIS,
+        parameters={"root_cause": "training_data_poisoning", "summary": "test"},
+    ))
+
+    assert obs.done is True
+    st = env.state()
+    score = st["grader_result"]["total"]
+    assert score >= 0.80, f"Expected high score, got {score}"
+    print(f"✓ test_full_episode_hard (score={score:.3f})")
+
+
+if __name__ == "__main__":
+    print("\n=== MLOps Firefighter Environment Tests ===\n")
+    test_reset_returns_valid_observation()
+    test_step_diagnostic_action()
+    test_redundant_diagnostic()
+    test_hasty_remediation_penalized()
+    test_correct_remediation_after_diagnosis()
+    test_submit_diagnosis_ends_episode()
+    test_state_tracking()
+    test_timeout()
+    test_grader_scores_in_range()
+    test_all_three_tasks_exist()
+    test_full_episode_easy()
+    test_full_episode_medium()
+    test_full_episode_hard()
+    print("\n=== All tests passed! ===\n")

validate.py

@@ -0,0 +1,219 @@
+#!/usr/bin/env python3
+"""
+Pre-submission validation script for the MLOps Firefighter environment.
+
+Checks all requirements from the hackathon rubric:
+  1. openenv.yaml exists and is valid
+  2. Typed Pydantic models exist
+  3. step()/reset()/state() work correctly
+  4. 3+ tasks with graders
+  5. Grader scores in 0.0–1.0 range
+  6. All required endpoints respond
+  7. Baseline produces scores
+  8. Dockerfile exists
+"""
+
+import json
+import sys
+import os
+import yaml
+
+sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
+sys.path.insert(0, os.path.join(os.path.dirname(os.path.abspath(__file__)), "server"))
+
+PASS = "✅"
+FAIL = "❌"
+results = []
+
+
+def check(name: str, condition: bool, detail: str = ""):
+    status = PASS if condition else FAIL
+    results.append((name, condition))
+    msg = f"  {status} {name}"
+    if detail:
+        msg += f" — {detail}"
+    print(msg)
+    return condition
+
+
+def main():
+    print("\n" + "=" * 60)
+    print("  MLOps Firefighter — Pre-Submission Validator")
+    print("=" * 60 + "\n")
+
+    # 1. openenv.yaml
+    print("[1/8] OpenEnv manifest (openenv.yaml)")
+    yaml_path = os.path.join(os.path.dirname(__file__), "openenv.yaml")
+    has_yaml = os.path.exists(yaml_path)
+    check("openenv.yaml exists", has_yaml)
+    if has_yaml:
+        with open(yaml_path) as f:
+            manifest = yaml.safe_load(f)
+        check("Has name", "name" in manifest)
+        check("Has version", "version" in manifest)
+        check("Has description", "description" in manifest)
+        check("Has tasks", "tasks" in manifest and len(manifest["tasks"]) >= 3)
+        check("Has 'openenv' tag", "openenv" in manifest.get("tags", []))
+
+    # 2. Typed Pydantic models
+    print("\n[2/8] Typed Pydantic models")
+    try:
+        from models import MLOpsAction, MLOpsObservation, ActionType
+        check("MLOpsAction importable", True)
+        check("MLOpsObservation importable", True)
+        check("ActionType enum exists", len(ActionType) >= 10)
+        # Verify they're Pydantic
+        a = MLOpsAction(action_type=ActionType.INSPECT_METRICS)
+        check("MLOpsAction is Pydantic", hasattr(a, "model_dump"))
+    except Exception as e:
+        check("Models import", False, str(e))
+
+    # 3. step()/reset()/state()
+    print("\n[3/8] Environment interface (reset/step/state)")
+    try:
+        from server.environment import MLOpsFirefighterEnvironment
+        env = MLOpsFirefighterEnvironment()
+
+        obs = env.reset(task_id="task_threshold_drift")
+        check("reset() returns observation", obs is not None)
+        check("reset() obs has done=False", obs.done is False)
+        check("reset() obs has step_number=0", obs.step_number == 0)
+
+        from models import MLOpsAction, ActionType
+        obs2 = env.step(MLOpsAction(action_type=ActionType.INSPECT_METRICS))
+        check("step() returns observation", obs2 is not None)
+        check("step() increments step_number", obs2.step_number == 1)
+        check("step() returns reward", isinstance(obs2.reward, float))
+
+        st = env.state()
+        check("state() returns dict", isinstance(st, dict))
+        check("state() has episode_id", "episode_id" in st)
+        check("state() has step_count", "step_count" in st)
+    except Exception as e:
+        check("Environment interface", False, str(e))
+
+    # 4. 3+ tasks
+    print("\n[4/8] Task definitions")
+    try:
+        from tasks import ALL_TASKS
+        check("3+ tasks defined", len(ALL_TASKS) >= 3)
+        difficulties = {t.difficulty for t in ALL_TASKS.values()}
+        check("Has easy task", "easy" in difficulties)
+        check("Has medium task", "medium" in difficulties)
+        check("Has hard task", "hard" in difficulties)
+        for tid, task in ALL_TASKS.items():
+            check(f"Task '{tid}' has root_causes", len(task.root_causes) > 0)
+            check(f"Task '{tid}' has diagnostics", len(task.required_diagnostics) > 0)
+            check(f"Task '{tid}' has remediations", len(task.correct_remediations) > 0)
+    except Exception as e:
+        check("Tasks", False, str(e))
+
+    # 5. Grader scores in range
+    print("\n[5/8] Grader scoring (0.0–1.0)")
+    try:
+        from tasks import grade_episode, ALL_TASKS
+        from models import ActionType
+        for tid, task in ALL_TASKS.items():
+            # Perfect
+            score, bd = grade_episode(
+                task=task,
+                actions_taken=[{"action_type": d.value} for d in task.required_diagnostics],
+                diagnosis_submitted={"root_cause": task.root_causes[0]},
+                remediation_applied=[r.value for r in task.correct_remediations],
+                total_steps=len(task.required_diagnostics) + 2,
+            )
+            check(f"'{tid}' perfect score in [0,1]", 0.0 <= score <= 1.0, f"{score:.3f}")
+
+            # Empty
+            score_z, _ = grade_episode(
+                task=task, actions_taken=[], diagnosis_submitted=None,
+                remediation_applied=[], total_steps=task.max_steps,
+            )
+            check(f"'{tid}' empty score in [0,1]", 0.0 <= score_z <= 1.0, f"{score_z:.3f}")
+
+            # Partial credit varies
+            check(f"'{tid}' grader differentiates", score > score_z, f"perfect={score:.3f} > empty={score_z:.3f}")
+    except Exception as e:
+        check("Grader", False, str(e))
+
+    # 6. All endpoints
+    print("\n[6/8] HTTP endpoints")
+    try:
+        from fastapi.testclient import TestClient
+        from server.app import app
+        client = TestClient(app)
+
+        r = client.get("/health")
+        check("/health returns 200", r.status_code == 200)
+
+        r = client.get("/tasks")
+        check("/tasks returns 200", r.status_code == 200)
+        check("/tasks has action_schema", "action_schema" in r.json())
+
+        r = client.post("/reset", json={"task_id": "task_threshold_drift"})
+        check("/reset returns 200", r.status_code == 200)
+
+        r = client.post("/step", json={"action_type": "inspect_metrics"})
+        check("/step returns 200", r.status_code == 200)
+
+        r = client.get("/state")
+        check("/state returns 200", r.status_code == 200)
+
+        # Complete an episode for grader test
+        client.post("/reset", json={"task_id": "task_threshold_drift"})
+        client.post("/step", json={"action_type": "inspect_metrics"})
+        client.post("/step", json={"action_type": "submit_diagnosis",
+                                    "parameters": {"root_cause": "test", "summary": "t"}})
+        r = client.post("/grader", json={})
+        check("/grader returns 200", r.status_code == 200)
+
+        r = client.post("/baseline")
+        check("/baseline returns 200", r.status_code == 200)
+        check("/baseline has scores", "average_score" in r.json())
+    except Exception as e:
+        check("Endpoints", False, str(e))
+
+    # 7. Baseline produces scores
+    print("\n[7/8] Baseline scoring")
+    try:
+        r = client.post("/baseline")
+        data = r.json()
+        avg = data["average_score"]
+        check("Baseline avg score > 0", avg > 0, f"avg={avg}")
+        for tid, result in data["baseline_results"].items():
+            s = result["score"]
+            check(f"Baseline '{tid}' in [0,1]", 0.0 <= s <= 1.0, f"{s:.3f}")
+    except Exception as e:
+        check("Baseline", False, str(e))
+
+    # 8. Dockerfile exists
+    print("\n[8/8] Dockerfile")
+    df_path = os.path.join(os.path.dirname(__file__), "Dockerfile")
+    check("Dockerfile exists", os.path.exists(df_path))
+    if os.path.exists(df_path):
+        with open(df_path) as f:
+            content = f.read()
+        check("Dockerfile has FROM", "FROM" in content)
+        check("Dockerfile has EXPOSE", "EXPOSE" in content)
+        check("Dockerfile has CMD", "CMD" in content)
+
+    # Summary
+    total = len(results)
+    passed = sum(1 for _, ok in results if ok)
+    failed = total - passed
+
+    print("\n" + "=" * 60)
+    if failed == 0:
+        print(f"  {PASS} ALL {total} CHECKS PASSED — Ready to submit!")
+    else:
+        print(f"  {FAIL} {failed}/{total} checks failed")
+        for name, ok in results:
+            if not ok:
+                print(f"     - {name}")
+    print("=" * 60 + "\n")
+
+    return 0 if failed == 0 else 1
+
+
+if __name__ == "__main__":
+    sys.exit(main())