Spaces:
Sleeping
Sleeping
| """ | |
| RepairEnv V2 — Main OpenEnv-compliant environment class. | |
| HARD CONSTRAINTS ENFORCED: | |
| - Hidden state: Observations only contain symptoms (logs + metrics). | |
| - Strict step order: Action → Cascade → Degrade → Recovery Ticks → Symptoms → Reward | |
| - Delayed recovery: repair_root_cause takes 2 steps to fully heal. | |
| - Inspection cost: 3 free, then -0.1 reward per inspect. | |
| """ | |
| from __future__ import annotations | |
| import copy | |
| import random | |
| from typing import Any, Dict, List, Optional, Tuple | |
| from models_v2 import ( | |
| Action, | |
| ActionType, | |
| ComponentMetrics, | |
| ComponentStatus, | |
| ComponentTarget, | |
| EpisodeResultV2, | |
| InspectResult, | |
| SymptomObservation, | |
| Reward, | |
| ACTION_COSTS | |
| ) | |
| from symptom_engine import SymptomEngine | |
| # --------------------------------------------------------------------------- | |
| # Constants | |
| # --------------------------------------------------------------------------- | |
| SEED = 42 | |
| STEP_PENALTY = -0.2 | |
| INSPECT_PENALTY = -0.1 | |
| MAX_FREE_INSPECTIONS = 3 | |
| COMPONENTS = ["database", "api", "cache", "queue"] | |
| # --------------------------------------------------------------------------- | |
| # RepairEnvV2 | |
| # --------------------------------------------------------------------------- | |
| class RepairEnvV2: | |
| def __init__(self, seed: int = SEED) -> None: | |
| """OpenEnv-compatible constructor. Accepts optional seed (default 42).""" | |
| self._init_seed = seed | |
| self._task_config: Optional[Dict[str, Any]] = None | |
| # Hidden State | |
| self._true_state: Dict[str, Dict] = {} | |
| self._root_causes: List[str] = [] | |
| self._recovery_cooldowns: Dict[str, int] = {} | |
| self._step_count: int = 0 | |
| self._max_steps: int = 0 | |
| self._task_id: str = "" | |
| self._done: bool = False | |
| self._total_reward: float = 0.0 | |
| self._inspections_used: int = 0 | |
| self._fixed_symptom_before_root: bool = False | |
| self._root_cause_fixed: bool = False | |
| self._rng: random.Random = random.Random(seed) | |
| self._symptom_engine: SymptomEngine = SymptomEngine(seed) | |
| # ------------------------------------------------------------------ | |
| # Public API | |
| # ------------------------------------------------------------------ | |
| def reset( | |
| self, | |
| task_config: Optional[Dict[str, Any]] = None, | |
| task_id: str = "easy", | |
| ) -> SymptomObservation: | |
| """ | |
| Reset the environment. | |
| OpenEnv calls this with no arguments → defaults to 'easy' task. | |
| server_v2.py calls it with (task_config=config, task_id=task_id) → still works. | |
| """ | |
| # If no task_config provided, load from tasks_v2 by task_id | |
| if task_config is None: | |
| from tasks_v2 import TASKS | |
| if task_id not in TASKS: | |
| task_id = "easy" | |
| task_config = TASKS[task_id] | |
| self._task_config = copy.deepcopy(task_config) | |
| self._task_id = task_id | |
| self._max_steps = self._task_config["max_steps"] | |
| self._step_count = 0 | |
| self._done = False | |
| self._total_reward = 0.0 | |
| self._inspections_used = 0 | |
| self._fixed_symptom_before_root = False | |
| self._root_causes_fixed: Dict[str, bool] = {} | |
| # Initialize Hidden State | |
| root = self._task_config["root_cause"] | |
| self._root_causes = root if isinstance(root, list) else [root] | |
| self._root_causes_fixed = {r: False for r in self._root_causes} | |
| self._true_state = copy.deepcopy(self._task_config["initial_state"]) | |
| self._recovery_cooldowns = {c: 0 for c in COMPONENTS} | |
| # Ensures rng is deterministic per run | |
| seed = self._task_config.get("seed", SEED) | |
| self._rng.seed(seed) | |
| self._symptom_engine = SymptomEngine(seed) | |
| return self._generate_symptoms() | |
| def step(self, action: Action) -> Tuple[SymptomObservation, Reward, bool, Dict[str, Any]]: | |
| """ | |
| Hard Rule #4: Strict Step Order | |
| 1. Action | |
| 2. Cascading | |
| 3. Degradation | |
| 4. Recovery Ticks | |
| 5. Symptoms (Metrics + Logs) | |
| 6. Reward computation | |
| """ | |
| if self._done: | |
| raise RuntimeError("Episode is done. Call reset().") | |
| # Snapshot for reward computation | |
| old_health = sum(1 for c in self._true_state.values() if c["status"] == "healthy") | |
| # 1. Action | |
| action_reward_modifier, inspect_info = self._apply_action(action) | |
| # 2. Cascading (Side effects driven by root cause) | |
| self._apply_cascading_effects() | |
| # 3. Degradation (Passive worsening) | |
| self._apply_degradation() | |
| # 4. Recovery Ticks | |
| self._tick_recovery_cooldowns(action) | |
| # 5. Generate Symptoms | |
| obs = self._generate_symptoms() | |
| self._step_count += 1 | |
| # Check terminal condition | |
| all_healthy = all(c["status"] == "healthy" for c in self._true_state.values()) | |
| if all_healthy or self._step_count >= self._max_steps: | |
| self._done = True | |
| # 6. Compute Reward | |
| new_health = sum(1 for c in self._true_state.values() if c["status"] == "healthy") | |
| reward = self._compute_reward( | |
| action, old_health, new_health, | |
| all_healthy, action_reward_modifier | |
| ) | |
| self._total_reward = round(self._total_reward + reward.value, 4) | |
| # Info dict (transparency for grading, completely hidden from agent observation) | |
| info = { | |
| "step": self._step_count, | |
| "true_state_snapshot": copy.deepcopy(self._true_state), | |
| "root_causes": self._root_causes, | |
| "cooldowns": self._recovery_cooldowns.copy(), | |
| "inspect_result": inspect_info.model_dump() if inspect_info else None, | |
| "reward_reason": reward.reason | |
| } | |
| return obs, reward, self._done, info | |
| def get_episode_result(self) -> EpisodeResultV2: | |
| final_health = sum(1 for c in self._true_state.values() if c["status"] == "healthy") | |
| health_ratio = final_health / 4.0 | |
| steps = self._step_count | |
| optimal = self._task_config.get("optimal_steps", 5) | |
| efficiency = max(0.0, min(1.0, 1.0 - (steps - optimal) * 0.05)) | |
| all_roots_fixed = all(self._root_causes_fixed.values()) | |
| correct_root_fix_ratio = 1.0 if (all_roots_fixed and not self._fixed_symptom_before_root) else 0.0 | |
| score = (0.5 * health_ratio) + (0.3 * efficiency) + (0.2 * correct_root_fix_ratio) | |
| return EpisodeResultV2( | |
| task_id=self._task_id, | |
| steps_taken=steps, | |
| total_reward=self._total_reward, | |
| score=round(score, 4), | |
| success=(health_ratio == 1.0), | |
| final_health_ratio=health_ratio, | |
| efficiency=efficiency, | |
| correct_root_fix_ratio=correct_root_fix_ratio | |
| ) | |
| # ------------------------------------------------------------------ | |
| # Private Implementation | |
| # ------------------------------------------------------------------ | |
| def _apply_action(self, action: Action) -> Tuple[float, Optional[InspectResult]]: | |
| """Handles both repair and inspect actions. Returns (reward_mod, inspect_result).""" | |
| target = action.target | |
| atype = action.action_type | |
| # Apply strict action costs | |
| reward_modifier = ACTION_COSTS.get(atype, 0.0) | |
| inspect_res = None | |
| if atype.startswith("inspect_"): | |
| self._inspections_used += 1 | |
| if self._inspections_used > MAX_FREE_INSPECTIONS: | |
| reward_modifier += INSPECT_PENALTY | |
| investigating_root = (target in self._root_causes) | |
| # If it's the root cause, high confidence. | |
| if investigating_root: | |
| conf = round(self._rng.uniform(0.75, 1.0), 2) | |
| hint = f"Critical failures traced back to {target}. High correlation with active symptoms." | |
| is_likely = True | |
| else: | |
| # Symptom. Low confidence, but might randomly say True (false positive). | |
| conf = round(self._rng.uniform(0.2, 0.55), 2) | |
| is_likely = self._rng.random() < 0.2 # 20% false positive | |
| hint = f"Anomalies detected in {target}, but could be upstream cascade." | |
| if self._true_state[target]["status"] == "healthy": | |
| conf = round(self._rng.uniform(0.8, 1.0), 2) | |
| is_likely = False | |
| hint = f"{target} appears healthy. No immediate faults found." | |
| inspect_res = InspectResult( | |
| component=target, | |
| hint=hint, | |
| confidence=conf, | |
| is_likely_root_cause=is_likely | |
| ) | |
| elif atype == ActionType.repair_database and target == "database": | |
| # Delayed recovery logic. | |
| if self._true_state["database"]["status"] != "healthy": | |
| if "database" in self._root_causes: | |
| self._recovery_cooldowns["database"] = 2 | |
| else: | |
| self._recovery_cooldowns["database"] = 1 | |
| elif atype == ActionType.restart_service: | |
| # Short term symptom fix. Will be overwritten by cascade if root still broken. | |
| if self._true_state[target]["status"] == "down": | |
| self._true_state[target]["status"] = "degraded" | |
| elif self._true_state[target]["status"] == "degraded": | |
| self._true_state[target]["status"] = "healthy" | |
| active_roots = [r for r in self._root_causes if self._true_state[r]["status"] != "healthy"] | |
| if target not in self._root_causes and active_roots: | |
| # They are fixing a symptom before the root cause | |
| self._fixed_symptom_before_root = True | |
| if target in self._root_causes: | |
| self._recovery_cooldowns[target] = 1 # Restarts are slightly faster but less reliable than dedicated tools | |
| # Other actions omitted for brevity, mapping easily to V1 equivalents where needed. | |
| elif atype == ActionType.clear_cache and target == "cache": | |
| if self._true_state["cache"]["status"] != "healthy": | |
| self._true_state["cache"]["status"] = "healthy" | |
| return reward_modifier, inspect_res | |
| def _apply_cascading_effects(self): | |
| """Root-cause driven cascades. Overwrites temporary symptom fixes.""" | |
| if "database" in self._root_causes and self._true_state["database"]["status"] != "healthy": | |
| # DB brings down the API | |
| self._true_state["api"]["status"] = "degraded" | |
| self._true_state["queue"]["status"] = "degraded" | |
| if "cache" in self._root_causes and self._true_state["cache"]["status"] != "healthy": | |
| # Cache pressure ruins queue | |
| self._true_state["queue"]["status"] = "degraded" | |
| def _apply_degradation(self): | |
| """Passive worsening.""" | |
| pass # Currently handled implicitly by SymptomEngine metric generation, but placeholder for hard state drops | |
| def _tick_recovery_cooldowns(self, action: Action): | |
| for c in COMPONENTS: | |
| if self._recovery_cooldowns[c] > 0: | |
| # Don't tick the cooldown if we *just* started it this exact step. | |
| if action.target == c and action.action_type in [ActionType.repair_database, ActionType.restart_service]: | |
| continue | |
| self._recovery_cooldowns[c] -= 1 | |
| if self._recovery_cooldowns[c] == 0: | |
| self._true_state[c]["status"] = "healthy" | |
| if c in self._root_causes: | |
| self._root_causes_fixed[c] = True | |
| def _generate_symptoms(self) -> SymptomObservation: | |
| metrics = self._symptom_engine.generate_metrics(self._true_state) | |
| logs = self._symptom_engine.generate_logs(self._true_state, self._step_count) | |
| inspections = max(0, MAX_FREE_INSPECTIONS - self._inspections_used) | |
| return SymptomObservation( | |
| metrics=metrics, | |
| logs=logs, | |
| step_count=self._step_count, | |
| max_steps=self._max_steps, | |
| inspections_remaining=inspections | |
| ) | |
| def _compute_reward(self, | |
| action: Action, | |
| old_health: int, new_health: int, | |
| all_healthy: bool, | |
| modifier: float) -> Reward: | |
| base = STEP_PENALTY + modifier | |
| reason = [] | |
| if modifier < 0: | |
| reason.append(f"Inspection cost ({modifier})") | |
| if action.action_type == ActionType.no_op: | |
| return Reward(value=base - 0.5, reason="Useless actions: no_op") | |
| # Root cause fix detection is delayed by cooldown, so we check if cooldown was just set | |
| # (meaning the correct action was taken this turn) | |
| root_cooldown_started = ( | |
| action.target in self._root_causes | |
| and self._recovery_cooldowns[action.target] > 0 | |
| and action.action_type in [ActionType.repair_database, ActionType.restart_service, ActionType.clear_cache] | |
| ) | |
| if root_cooldown_started: | |
| base += 1.5 | |
| reason.append(f"Root cause addressed ({action.target})") | |
| elif new_health > old_health: | |
| base += 0.5 | |
| reason.append("Symptom improved") | |
| elif action.target not in self._root_causes and not action.action_type.startswith("inspect"): | |
| # They tried to fix a symptom that is immediately cascaded away (or just wrong) | |
| active_roots = [r for r in self._root_causes if self._true_state[r]["status"] != "healthy"] | |
| if active_roots: | |
| base -= 1.0 | |
| reason.append("Wrong component targeted (root cause still active)") | |
| else: | |
| base -= 0.5 | |
| reason.append("Component action had no net effect") | |
| if all_healthy: | |
| base += 5.0 | |
| reason.append("Terminal Bonus: All healthy") | |
| return Reward(value=round(base, 4), reason=" | ".join(reason) if reason else "Step tick") | |