Spaces:
Sleeping
Sleeping
| """Data models for the SRE Bench environment.""" | |
| from typing import Any, Dict, List, Literal, Optional, Set | |
| from pydantic import BaseModel, Field | |
| try: | |
| from openenv.core.env_server.types import Action, Observation, State | |
| except ImportError: | |
| from openenv.core.env_server.types import Action, Observation, State | |
| ActionType = Literal["restart", "scale_up", "rollback", "investigate", "lookup_runbook", "noop"] | |
| class SREAction(Action): | |
| """ | |
| Action for the SRE Bench environment. | |
| action_type: | |
| - restart: Restart a service (clears crash/memory_leak faults) | |
| - scale_up: Scale up a service (clears latency_spike/high CPU+memory) | |
| - rollback: Roll back a service config (clears config_corruption) | |
| - investigate: Reveal the hidden fault type for a service (costs a step) | |
| - lookup_runbook: Pull the SRE runbook for a service — returns symptom | |
| checklist and recommended fix procedure (costs a step) | |
| - noop: Do nothing | |
| """ | |
| action_type: ActionType = Field(default="noop", description="Type of action") | |
| service: Optional[str] = Field( | |
| default=None, | |
| description="Target service name (required for all actions except noop)", | |
| ) | |
| class ServiceMetrics(BaseModel): | |
| """Public metrics for a single service.""" | |
| status: str = Field(default="healthy", description="healthy | degraded | down") | |
| cpu: float = Field(default=20.0, description="CPU usage %") | |
| memory: float = Field(default=30.0, description="Memory usage %") | |
| error_rate: float = Field(default=0.0, description="Fraction of requests failing (0-1)") | |
| latency_ms: float = Field(default=50.0, description="P99 latency in ms") | |
| restart_count: int = Field(default=0, description="Total restarts this episode") | |
| class SREObservation(Observation): | |
| """Observation returned at each step.""" | |
| services: Dict[str, Dict[str, Any]] = Field( | |
| default_factory=dict, | |
| description=( | |
| "Public metrics per service: status, cpu, memory, error_rate, latency_ms, " | |
| "restart_count, plus trend signals: memory_trend, latency_trend, error_trend " | |
| "(rising | falling | stable)." | |
| ), | |
| ) | |
| alerts: List[str] = Field( | |
| default_factory=list, | |
| description="Recent alert log (last 5 events, with severity prefix CRITICAL/WARNING/INFO)", | |
| ) | |
| step: int = Field(default=0, description="Current step number") | |
| max_steps: int = Field(default=50, description="Episode length") | |
| uptime_score: float = Field( | |
| default=1.0, | |
| description="SLA-weighted uptime score so far (0-1). Final grader score at done=True.", | |
| ) | |
| investigate_result: Optional[Dict[str, Any]] = Field( | |
| default=None, | |
| description="Populated when action_type='investigate'", | |
| ) | |
| runbook_result: Optional[Dict[str, Any]] = Field( | |
| default=None, | |
| description="Populated when action_type='lookup_runbook': symptom checklist + fix procedure", | |
| ) | |
| task_id: str = Field(default="", description="Active task identifier") | |
| task_description: str = Field(default="", description="Human-readable task objective") | |
| incident_severity: Optional[str] = Field( | |
| default=None, | |
| description="P1/P2/P3/OK — severity based on which SLA-critical services are down", | |
| ) | |
| sla_breach_in: Optional[int] = Field( | |
| default=None, | |
| description="For incident_response task: steps remaining before SLA breach penalty kicks in", | |
| ) | |
| class SREState(State): | |
| """Full server-side state (includes hidden fault info).""" | |
| services: Dict[str, Dict[str, Any]] = Field(default_factory=dict) | |
| faults: Dict[str, Optional[str]] = Field( | |
| default_factory=dict, | |
| description="Hidden fault per service: crash | memory_leak | latency_spike | config_corruption | None", | |
| ) | |
| alerts: List[str] = Field(default_factory=list) | |
| step_count: int = Field(default=0) | |
| max_steps: int = Field(default=50) | |
| task_id: str = Field(default="") | |
| total_reward: float = Field(default=0.0) | |
| # SLA-weighted cumulative health (replaces simple cumulative_healthy count) | |
| cumulative_weighted_health: float = Field(default=0.0) | |
| # For cascading_failure grader: tracks which root causes the agent healed | |
| root_causes_healed: List[str] = Field(default_factory=list) | |
| # For single_fault grader: first step when all services were healthy | |
| steps_to_full_health: Optional[int] = Field(default=None) | |
| # Metric history for trend computation: {svc: [{"memory": f, "latency_ms": f, "error_rate": f}, ...]} | |
| metric_history: Dict[str, List[Dict[str, Any]]] = Field(default_factory=dict) | |
| # For incident_response grader: step at which api was restored | |
| api_restored_at: Optional[int] = Field(default=None) | |
| __all__ = ["SREAction", "SREObservation", "SREState", "ServiceMetrics", "ActionType"] | |