Team Ai
Apppublic

tek-wizard/devops-incident-responder

sourceHugging Faceupdated 6mo agoView on Hugging Face
2likes
create_baseline_comparison_notebook.py459 linesDownload Raw Back to server
1from __future__ import annotations2 3import json4from pathlib import Path5 6 7NOTEBOOK_PATH = Path("/home/prateeksingh/Desktop/devops_responder/baseline_comparison.ipynb")8 9 10def md(source: str) -> dict:11    return {12        "cell_type": "markdown",13        "metadata": {},14        "source": [line + "\n" for line in source.strip("\n").split("\n")],15    }16 17 18def code(source: str) -> dict:19    return {20        "cell_type": "code",21        "execution_count": None,22        "metadata": {},23        "outputs": [],24        "source": [line + "\n" for line in source.strip("\n").split("\n")],25    }26 27 28def main() -> None:29    notebook = {30        "cells": [31            md(32                """33# Baseline Comparison — DevOps Incident Responder34 35This notebook compares three policies against the deployed DevOps Incident Responder environment:36 37- `random`38- `signal_guided`: a lightweight evidence-aware baseline that is intentionally imperfect39- `hf_llm`: an external LLM baseline through the Hugging Face router40 41Outputs:42- summary tables43- per-task score breakdown44- `baseline_comparison.png` for the README/blog45                """46            ),47            md("## 1. Install Dependencies"),48            code(49                """50%%capture51!pip install --upgrade requests matplotlib pandas openai52                """53            ),54            md("## 2. Environment And Model Setup"),55            code(56                """57import json58import os59import random60import re61from pathlib import Path62from typing import Any63 64import matplotlib.pyplot as plt65import pandas as pd66import requests67from openai import OpenAI68 69ENV_URL = "https://tek-wizard-devops-incident-responder.hf.space"70HF_MODEL = "openai/gpt-oss-120b:groq"71REQUEST_TIMEOUT = 2072MAX_STEPS = 1573 74 75def _extract_hf_token() -> str | None:76    env_token = os.getenv("HF_TOKEN") or os.getenv("OPENAI_API_KEY")77    if env_token:78        return env_token79 80    candidate = Path("testing/baselineGemini.py")81    if not candidate.exists():82        return None83 84    text = candidate.read_text(encoding="utf-8")85    match = re.search(r'HF_TOKEN\\s*=\\s*"([^"]+)"', text)86    return match.group(1) if match else None87 88 89HF_TOKEN = _extract_hf_token()90 91response = requests.get(f"{ENV_URL}/tasks", timeout=REQUEST_TIMEOUT)92response.raise_for_status()93catalog = response.json()94TASKS = catalog["tasks"]95TASK_IDS = [task["id"] for task in TASKS]96SEED_SETS = catalog["seed_sets"]97 98print("Environment:", ENV_URL)99print("Tasks:", TASK_IDS)100print("Validation seeds:", SEED_SETS["val"])101print("HF model:", HF_MODEL)102print("HF token available:", bool(HF_TOKEN))103                """104            ),105            md("## 3. Environment Helpers"),106            code(107                """108def env_reset(task_id: str, seed: int) -> dict[str, Any]:109    response = requests.post(110        f"{ENV_URL}/reset",111        params={"task_id": task_id, "seed": int(seed)},112        timeout=REQUEST_TIMEOUT,113    )114    response.raise_for_status()115    return response.json()116 117 118def env_step(session_id: str, action: dict[str, Any]) -> dict[str, Any]:119    payload = dict(action)120    payload["session_id"] = session_id121    response = requests.post(f"{ENV_URL}/step", json=payload, timeout=REQUEST_TIMEOUT)122    response.raise_for_status()123    return response.json()124 125 126def env_grade(session_id: str) -> dict[str, Any]:127    response = requests.get(128        f"{ENV_URL}/grader",129        params={"session_id": session_id},130        timeout=REQUEST_TIMEOUT,131    )132    response.raise_for_status()133    return response.json()134 135 136def service_priority(observation: dict[str, Any]) -> list[str]:137    service_status = observation.get("service_status", {})138    scores = []139    for name, service in service_status.items():140        score = (141            3.0 * float(service.get("error_rate", 0.0))142            + 1.7 * float(service.get("latency", 0.0))143            + 1.2 * float(service.get("saturation", 0.0))144            + (0.6 if service.get("status") != "RUNNING" else 0.0)145        )146        scores.append((score, name))147    scores.sort(reverse=True)148    return [name for _, name in scores]149 150 151def combined_text(observation: dict[str, Any]) -> str:152    events = observation.get("recent_events", [])153    facts = observation.get("discovered_facts", [])154    brief = observation.get("incident_brief", "")155    return " ".join([brief, *events, *facts]).lower()156 157 158def seen(history: list[dict[str, str]], command: str, target: str | None = None) -> bool:159    return any(160        item.get("command") == command and (target is None or item.get("target") == target)161        for item in history162    )163                """164            ),165            md("## 4. Policies"),166            code(167                """168def random_policy(observation: dict[str, Any], history: list[dict[str, str]], rng: random.Random) -> dict[str, str]:169    commands = observation["available_commands"]170    targets = ["system", *observation["service_status"].keys()]171    return {"command": rng.choice(commands), "target": rng.choice(targets)}172 173 174def signal_guided_policy(observation: dict[str, Any], history: list[dict[str, str]], rng: random.Random) -> dict[str, str]:175    text = combined_text(observation)176    priority = service_priority(observation)177    top_service = priority[0] if priority else "auth-api"178    queue_depth = float(observation.get("global_signals", {}).get("queue_depth", 0.0))179 180    investigation_actions = []181    if not seen(history, "query_logs", top_service):182        investigation_actions.append({"command": "query_logs", "target": top_service})183    if not seen(history, "get_metrics", top_service):184        investigation_actions.append({"command": "get_metrics", "target": top_service})185    if ("deploy" in text or "rollout" in text or "release" in text) and not seen(history, "get_recent_deploys", "system"):186        investigation_actions.append({"command": "get_recent_deploys", "target": "system"})187    if ("config" in text or "feature" in text or "flag" in text) and not seen(history, "inspect_config", top_service):188        investigation_actions.append({"command": "inspect_config", "target": top_service})189    if ("partition" in text or "zone" in text or "dependency" in text) and not seen(history, "check_dependency", "auth-api"):190        investigation_actions.append({"command": "check_dependency", "target": "auth-api"})191 192    if len(history) < 2 and investigation_actions:193        return investigation_actions[0]194    if investigation_actions and rng.random() < 0.45:195        return investigation_actions[0]196 197    if ("deploy" in text or "serializer" in text or "malformed" in text) and top_service in {"auth-api", "billing-api"}:198        if not seen(history, "rollback_deploy", top_service):199            return {"command": "rollback_deploy", "target": top_service}200 201    if ("shadow_reads" in text or "fanout" in text or "cache miss" in text) and not seen(history, "disable_feature_flag", "auth-api"):202        return {"command": "disable_feature_flag", "target": "auth-api"}203 204    if ("partition" in text or "zone" in text) and seen(history, "check_dependency", "auth-api") and not seen(history, "failover_db", "main-db"):205        return {"command": "failover_db", "target": "main-db"}206 207    if queue_depth > 250 and not seen(history, "drain_queue", "queue"):208        return {"command": "drain_queue", "target": "queue"}209    if queue_depth > 350 and not seen(history, "scale_service", "worker"):210        return {"command": "scale_service", "target": "worker"}211 212    if top_service == "main-db" and not seen(history, "scale_service", "main-db"):213        return {"command": "scale_service", "target": "main-db"}214    if top_service == "main-db" and seen(history, "scale_service", "main-db") and not seen(history, "clear_connections", "main-db"):215        return {"command": "clear_connections", "target": "main-db"}216 217    if top_service == "worker" and not seen(history, "restart_service", "worker"):218        return {"command": "restart_service", "target": "worker"}219 220    smoke_candidates = ["auth-api", "billing-api", "worker"]221    for target in smoke_candidates:222        if target in observation["service_status"] and not seen(history, "run_smoke_test", target):223            return {"command": "run_smoke_test", "target": target}224 225    fallback_pool = [226        {"command": "query_logs", "target": top_service},227        {"command": "get_metrics", "target": top_service},228        {"command": "get_recent_deploys", "target": "system"},229    ]230    return rng.choice(fallback_pool)231 232 233def build_hf_client() -> OpenAI | None:234    if not HF_TOKEN:235        return None236    return OpenAI(base_url="https://router.huggingface.co/v1", api_key=HF_TOKEN)237 238 239def hf_llm_policy(240    observation: dict[str, Any],241    history: list[dict[str, str]],242    rng: random.Random,243    client: OpenAI | None = None,244) -> dict[str, str]:245    if client is None:246        return signal_guided_policy(observation, history, rng)247 248    prompt = f\"\"\"249You are a senior SRE responding to a production incident.250 251Pick exactly one next action as compact JSON:252{{"command":"...","target":"..."}}253 254Rules:255- Return JSON only.256- Use exactly one command and one target.257- Prefer investigation before risky mitigation unless evidence is strong.258 259Observation:260{json.dumps(observation)}261 262Recent history:263{json.dumps(history[-6:])}264\"\"\".strip()265 266    try:267        response = client.chat.completions.create(268            model=HF_MODEL,269            messages=[{"role": "user", "content": prompt}],270            max_tokens=120,271            temperature=0.1,272        )273        content = (response.choices[0].message.content or "").strip()274        if "```" in content:275            content = content.replace("```json", "").replace("```", "").strip()276        match = re.search(r"\\{.*\\}", content, re.DOTALL)277        if match:278            content = match.group(0)279        parsed = json.loads(content)280        return {281            "command": str(parsed.get("command", "query_logs")).strip(),282            "target": str(parsed.get("target", "auth-api")).strip(),283        }284    except Exception:285        return signal_guided_policy(observation, history, rng)286                """287            ),288            md("## 5. Evaluation Loop"),289            code(290                """291def run_episode(policy_name: str, task_id: str, seed: int, client: OpenAI | None = None) -> dict[str, Any]:292    observation = env_reset(task_id, int(seed))293    session_id = observation["session_id"]294    history: list[dict[str, str]] = []295    rng = random.Random(seed)296    parse_failures = 0297 298    for _ in range(MAX_STEPS):299        if policy_name == "random":300            action = random_policy(observation, history, rng)301        elif policy_name == "signal_guided":302            action = signal_guided_policy(observation, history, rng)303        elif policy_name == "hf_llm":304            action = hf_llm_policy(observation, history, rng, client=client)305        else:306            raise ValueError(policy_name)307 308        history.append(action)309 310        try:311            payload = env_step(session_id, action)312        except Exception:313            parse_failures += 1314            break315 316        observation = payload["observation"]317        if payload.get("done", False):318            break319 320    grade = env_grade(session_id)321    return {322        "policy": policy_name,323        "task_id": task_id,324        "seed": int(seed),325        "score": float(grade["score"]),326        "resolved": bool(grade["resolved"]),327        "steps_taken": int(grade["steps_taken"]),328        "unsafe_actions": int(grade["unsafe_actions"]),329        "investigation_coverage": float(grade["investigation_coverage"]),330        "parse_failures": parse_failures,331    }332 333 334def evaluate_policy(policy_name: str, seeds: list[int], client: OpenAI | None = None) -> list[dict[str, Any]]:335    rows = []336    for task_id in TASK_IDS:337        for seed in seeds:338            row = run_episode(policy_name, task_id, int(seed), client=client)339            rows.append(row)340            print(341                f"{policy_name:14s} task={task_id:26s} seed={seed:>3} "342                f"score={row['score']:.4f} resolved={row['resolved']}"343            )344    return rows345                """346            ),347            md("## 6. Run Baseline Comparisons"),348            code(349                """350client = build_hf_client()351 352comparison_rows = []353comparison_rows.extend(evaluate_policy("random", SEED_SETS["val"], client=None))354comparison_rows.extend(evaluate_policy("signal_guided", SEED_SETS["val"], client=None))355comparison_rows.extend(evaluate_policy("hf_llm", SEED_SETS["val"], client=client))356 357results_df = pd.DataFrame(comparison_rows)358results_df.head()359                """360            ),361            md("## 7. Summary Tables"),362            code(363                """364summary_df = (365    results_df.groupby("policy", as_index=False)366    .agg(367        avg_score=("score", "mean"),368        solve_rate=("resolved", "mean"),369        avg_steps=("steps_taken", "mean"),370        avg_unsafe_actions=("unsafe_actions", "mean"),371        avg_investigation_coverage=("investigation_coverage", "mean"),372    )373    .sort_values("avg_score", ascending=False)374)375 376summary_df["avg_score"] = summary_df["avg_score"].round(4)377summary_df["solve_rate"] = summary_df["solve_rate"].round(4)378summary_df["avg_steps"] = summary_df["avg_steps"].round(2)379summary_df["avg_unsafe_actions"] = summary_df["avg_unsafe_actions"].round(2)380summary_df["avg_investigation_coverage"] = summary_df["avg_investigation_coverage"].round(4)381 382print("Policy Summary")383display(summary_df)384 385per_task_df = (386    results_df.groupby(["policy", "task_id"], as_index=False)387    .agg(avg_score=("score", "mean"), solve_rate=("resolved", "mean"))388)389per_task_df["avg_score"] = per_task_df["avg_score"].round(4)390per_task_df["solve_rate"] = per_task_df["solve_rate"].round(4)391 392print("Per-Task Summary")393display(per_task_df)394                """395            ),396            md("## 8. Create README/Blog Plot"),397            code(398                """399plot_df = summary_df.copy()400 401fig, axes = plt.subplots(1, 2, figsize=(13, 5))402 403axes[0].bar(plot_df["policy"], plot_df["avg_score"], color=["#9aa5b1", "#2cb67d", "#ef4565"])404axes[0].set_title("Average Score by Policy")405axes[0].set_ylabel("Average grade score")406axes[0].set_ylim(0, 1.0)407for idx, value in enumerate(plot_df["avg_score"]):408    axes[0].text(idx, value + 0.02, f"{value:.2f}", ha="center", fontsize=10)409 410axes[1].bar(plot_df["policy"], plot_df["solve_rate"], color=["#9aa5b1", "#2cb67d", "#ef4565"])411axes[1].set_title("Solve Rate by Policy")412axes[1].set_ylabel("Fraction of incidents resolved")413axes[1].set_ylim(0, 1.0)414for idx, value in enumerate(plot_df["solve_rate"]):415    axes[1].text(idx, value + 0.02, f"{value:.2f}", ha="center", fontsize=10)416 417fig.suptitle("DevOps Incident Responder Baseline Comparison", fontsize=14)418fig.tight_layout()419plt.savefig("baseline_comparison.png", dpi=160, bbox_inches="tight")420plt.show()421 422print("Saved baseline_comparison.png")423                """424            ),425            md("## 9. Save CSV Outputs"),426            code(427                """428results_df.to_csv("baseline_episode_results.csv", index=False)429summary_df.to_csv("baseline_summary.csv", index=False)430per_task_df.to_csv("baseline_per_task_summary.csv", index=False)431 432print("Saved baseline_episode_results.csv")433print("Saved baseline_summary.csv")434print("Saved baseline_per_task_summary.csv")435                """436            ),437        ],438        "metadata": {439            "kernelspec": {440                "display_name": "Python 3",441                "language": "python",442                "name": "python3",443            },444            "language_info": {445                "name": "python",446                "version": "3.12",447            },448        },449        "nbformat": 4,450        "nbformat_minor": 5,451    }452 453    NOTEBOOK_PATH.write_text(json.dumps(notebook, indent=2) + "\n", encoding="utf-8")454    print(f"Wrote {NOTEBOOK_PATH}")455 456 457if __name__ == "__main__":458    main()459