SidhaGarg/Cloud-DevOps-RLEnv
0
1# Copyright (c) Meta Platforms, Inc. and affiliates.2# All rights reserved.3#4# This source code is licensed under the BSD-style license found in the5# LICENSE file in the root directory of this source tree.6 7"""8Cloud Devops Env Environment Implementation.9 10A deterministic mock cloud/devops environment with reward shaping and11anti-farming guardrails for hackathon evaluation.12"""13 14from __future__ import annotations15 16import copy17from uuid import uuid418 19from openenv.core.env_server.interfaces import Environment20from openenv.core.env_server.types import State21 22try:23 from ..models import CloudAction, CloudObservation, CloudState24except ImportError:25 from models import CloudAction, CloudObservation, CloudState26 27 28class CloudDevopsEnvironment(Environment):29 """30 A deterministic mock cloud/devops environment.31 32 Tasks:33 - easy: open port 80 on sg-web34 - medium: inspect noisy API logs, then open port 5432 on sg-db35 - hard: trace 502 from lb-main to i-web2, then restart i-web2 (not i-web1)36 37 Example:38 >>> env = CloudDevopsEnvironment()39 >>> obs = env.reset()40 >>> print(obs.system_health_status) # "CRITICAL"41 >>>42 >>> obs = env.step(CloudAction(command="list_resources"))43 >>> print(obs.output)44 """45 46 # Enable concurrent WebSocket sessions.47 # Set to True if your environment isolates state between instances.48 # When True, multiple WebSocket clients can connect simultaneously, each49 # getting their own environment instance (when using factory mode in app.py).50 SUPPORTS_CONCURRENT_SESSIONS: bool = True51 MAX_STEPS: int = 2052 VALID_TASKS = {"easy", "medium", "hard"}53 ACTION_COST: float = 0.0154 55 def __init__(self, task_name: str = "easy"):56 """Initialize the cloud_devops_env environment."""57 normalized_task = (task_name or "easy").lower()58 if normalized_task not in self.VALID_TASKS:59 raise ValueError(f"Unknown task: {task_name}")60 61 self.task_name = normalized_task62 self._state_data: CloudState | None = None63 self._achievements: set[str] = set()64 65 def _build_noise_resources(self) -> dict[str, dict[str, object]]:66 """Generate deterministic decoy resources to force retrieval and filtering."""67 resources: dict[str, dict[str, object]] = {}68 for i in range(1, 21):69 suffix = f"{i:02d}"70 resources[f"i-backend-{suffix}"] = {71 "type": "Instance",72 "status": "running",73 "logs": (74 "[2026-04-06 17:00:00] INFO node-exporter: "75 "standard metrics reported successfully"76 ),77 }78 resources[f"sg-backend-{suffix}"] = {79 "type": "SecurityGroup",80 "rules": [{"port": 443, "action": "allow"}],81 }82 return resources83 84 def _build_task_resources(self) -> dict[str, dict[str, object]]:85 resources = self._build_noise_resources()86 87 if self.task_name == "easy":88 resources.update(89 {90 "i-web": {"type": "Instance", "status": "running"},91 "sg-web": {92 "type": "SecurityGroup",93 "rules": [{"port": 22, "action": "allow"}],94 },95 }96 )97 return resources98 99 if self.task_name == "medium":100 resources.update(101 {102 "i-api": {103 "type": "Instance",104 "ip_address": "10.0.4.11",105 "status": "running",106 "logs": (107 "[2026-04-06 17:01:22] [CRITICAL] "108 "sqlalchemy.exc.OperationalError: "109 "(psycopg2.OperationalError) connection to server at "110 "'10.0.4.5', port 5432 failed: Connection timed out. "111 "Is the server running and accepting TCP/IP connections?"112 ),113 },114 "i-db": {115 "type": "Instance",116 "ip_address": "10.0.4.5",117 "status": "running",118 },119 "sg-db": {120 "type": "SecurityGroup",121 "rules": [{"port": 22, "action": "allow"}],122 },123 "metadata-svc": {124 "type": "MetadataService",125 "status": "running",126 },127 }128 )129 return resources130 131 resources.update(132 {133 "lb-main": {134 "type": "LoadBalancer",135 "logs": (136 "2026/04/06 17:02:09 [error] 3197#3197: *4189 upstream timed out "137 "(110: Connection timed out) while reading response header from upstream, "138 "client: 10.0.2.14, server: api.prod.local, request: \"GET /checkout HTTP/1.1\", "139 "upstream: \"http://10.0.8.22:8080/checkout\", host: \"api.prod.local\"\n"140 "2026/04/06 17:02:10 [error] 3197#3197: *4190 no live upstreams while "141 "connecting to upstream \"10.0.8.22\""142 ),143 },144 "lb-external": {145 "type": "LoadBalancer",146 "status": "running",147 "logs": "INFO: Edge traffic stable.",148 },149 "i-web1": {150 "type": "Instance",151 "ip_address": "10.0.8.21",152 "status": "running",153 "logs": (154 "[2026-04-06 17:02:11] INFO web-service: readiness probe passed\n"155 "[2026-04-06 17:02:12] INFO jvm: heap usage stable at 42%"156 ),157 },158 "i-web2": {159 "type": "Instance",160 "ip_address": "10.0.8.22",161 "status": "degraded",162 "logs": (163 "kernel: Out of memory: Killed process 12345 (java) total-vm:4194304kB, "164 "anon-rss:3145728kB\n"165 "systemd[1]: web-service.service: Main process exited, code=killed, "166 "status=9/KILL"167 ),168 },169 "sg-web": {170 "type": "SecurityGroup",171 "rules": [{"port": 80, "action": "allow"}],172 },173 "metadata-svc": {174 "type": "MetadataService",175 "status": "running",176 },177 }178 )179 return resources180 181 def _lookup_resource_by_ip(self, ip_address: str) -> str | None:182 if self._state_data is None:183 return None184 for resource_id, data in self._state_data.resources.items():185 if data.get("ip_address") == ip_address:186 return resource_id187 return None188 189 def _apply_cascading_failure(self) -> tuple[float, str]:190 """Simulate system drift in hard mode if root cause is not fixed quickly."""191 if self._state_data is None or self.task_name != "hard":192 return 0.0, ""193 194 state = self._state_data195 if state.is_resolved or state.step_count <= 8:196 return 0.0, ""197 198 lb = state.resources.get("lb-external")199 if not lb:200 return 0.0, ""201 202 if lb.get("status") != "DOWN":203 lb["status"] = "DOWN"204 lb["logs"] = (205 "CRITICAL: Cascading failure triggered after prolonged unresolved OOM incident. "206 "Edge load balancer stopped serving traffic."207 )208 return -0.05, (209 "\nALERT: Cascading failure detected. lb-external is DOWN due to delayed remediation."210 )211 212 return -0.03, ""213 214 def _reward_once(self, achievement: str, points: float) -> float:215 if achievement in self._achievements:216 return 0.0217 self._achievements.add(achievement)218 return points219 220 def _task_objective(self) -> str:221 objectives = {222 "easy": "Restore web access by allowing port 80 on sg-web.",223 "medium": (224 "Restore API to DB connectivity by reading i-api logs, resolving DB IP via "225 "query_metadata, then allowing port 5432 on sg-db."226 ),227 "hard": (228 "Recover checkout path by tracing lb-main upstream IP, resolving it with "229 "query_metadata, inspecting i-web2, and restarting i-web2 safely."230 ),231 }232 return objectives[self.task_name]233 234 def reset(self) -> CloudObservation: # type: ignore[override]235 """Reset the environment to the initial state for the selected task."""236 self._achievements.clear()237 self._state_data = CloudState(238 episode_id=str(uuid4()),239 task_difficulty=self.task_name,240 resources=copy.deepcopy(self._build_task_resources()),241 step_count=0,242 is_resolved=False,243 )244 245 return CloudObservation(246 output=(247 "Environment initialized. System status is currently CRITICAL. "248 "Use 'list_resources' to begin triage."249 ),250 error=None,251 system_health_status="CRITICAL",252 done=False,253 reward=0.0,254 metadata={255 "step_count": 0,256 "resolved": False,257 "task": self.task_name,258 "total_resources": len(self._state_data.resources),259 "objective": self._task_objective(),260 "deterministic": True,261 "max_steps": self.MAX_STEPS,262 "action_cost": self.ACTION_COST,263 "hard_cascade_trigger_step": 8,264 },265 echoed_message="Cloud Devops Env environment ready!",266 message_length=0,267 )268 269 def step(self, action: CloudAction) -> CloudObservation: # type: ignore[override]270 """Execute the agent action and return the next observation."""271 if self._state_data is None:272 self.reset()273 274 assert self._state_data is not None275 state = self._state_data276 277 state.step_count += 1278 reward = -self.ACTION_COST279 reward_breakdown: list[dict[str, object]] = [280 {"event": "action_cost", "delta": -self.ACTION_COST}281 ]282 done = False283 output = ""284 error = None285 termination_reason = "in_progress"286 287 def add_reward(delta: float, event: str) -> None:288 nonlocal reward289 if abs(delta) < 1e-12:290 return291 reward += delta292 reward_breakdown.append({"event": event, "delta": round(float(delta), 4)})293 294 try:295 if action.command == "list_resources":296 res_list = [297 f"{resource_id} ({data['type']})"298 for resource_id, data in sorted(state.resources.items())299 ]300 output = "Available Resources:\n" + "\n".join(res_list)301 302 elif action.command == "describe_resource":303 if not action.resource_id or action.resource_id not in state.resources:304 raise ValueError(f"Resource {action.resource_id} not found.")305 306 output = str(state.resources[action.resource_id])307 308 if self.task_name == "easy" and action.resource_id == "sg-web":309 add_reward(self._reward_once("read_sg", 0.2), "inspect_web_sg")310 elif self.task_name == "medium" and action.resource_id == "sg-db":311 add_reward(self._reward_once("read_sg", 0.2), "inspect_db_sg")312 elif self.task_name == "hard" and action.resource_id == "i-web2":313 add_reward(314 self._reward_once("inspect_target", 0.2),315 "inspect_target_instance",316 )317 318 elif action.command == "view_logs":319 if not action.resource_id:320 raise ValueError("resource_id is required for view_logs.")321 322 res = state.resources.get(action.resource_id)323 if not res:324 raise ValueError(f"Resource {action.resource_id} not found.")325 326 output = str(res.get("logs", "No logs available for this resource."))327 328 if self.task_name == "medium" and action.resource_id == "i-api":329 add_reward(self._reward_once("read_logs", 0.2), "inspect_api_logs")330 elif self.task_name == "hard" and action.resource_id == "lb-main":331 add_reward(self._reward_once("inspect_lb", 0.2), "inspect_lb_logs")332 elif self.task_name == "hard" and action.resource_id == "i-web2":333 add_reward(334 self._reward_once("inspect_target", 0.2),335 "inspect_target_logs",336 )337 338 elif action.command == "query_metadata":339 ip_address = None340 if action.parameters and isinstance(action.parameters, dict):341 ip_address = action.parameters.get("ip_address")342 if not ip_address and action.resource_id:343 ip_address = action.resource_id344 if not ip_address:345 raise ValueError("query_metadata requires parameters.ip_address.")346 347 resource_id = self._lookup_resource_by_ip(str(ip_address))348 if not resource_id:349 raise ValueError(f"No resource found for ip_address={ip_address}")350 351 output = f"Metadata lookup: ip_address={ip_address} resource_id={resource_id}"352 if self.task_name == "medium" and str(ip_address) == "10.0.4.5":353 add_reward(354 self._reward_once("lookup_db_target", 0.2),355 "resolve_db_ip_dependency",356 )357 elif self.task_name == "hard" and str(ip_address) == "10.0.8.22":358 add_reward(359 self._reward_once("lookup_upstream_target", 0.2),360 "resolve_upstream_ip_dependency",361 )362 363 elif action.command == "update_security_group":364 if not action.resource_id:365 raise ValueError("resource_id is required for update_security_group.")366 367 res = state.resources.get(action.resource_id)368 if not res or res.get("type") != "SecurityGroup":369 raise ValueError(f"Invalid Security Group ID: {action.resource_id}")370 if not action.parameters or "port" not in action.parameters:371 raise ValueError("Missing 'port' in parameters.")372 if "action" not in action.parameters:373 raise ValueError("Missing 'action' in parameters. Use 'allow' or 'deny'.")374 375 rule = copy.deepcopy(action.parameters)376 rules = res.get("rules")377 if not isinstance(rules, list):378 raise ValueError(f"Security group {action.resource_id} has invalid rules.")379 port = int(rule["port"])380 rule_action = str(rule.get("action", "")).lower()381 if rule_action not in {"allow", "deny"}:382 raise ValueError(383 "Invalid security-group action. Supported values: 'allow', 'deny'."384 )385 386 rules.append(rule)387 output = f"Successfully updated {action.resource_id} with rule: {rule}"388 if (389 self.task_name == "easy"390 and action.resource_id == "sg-web"391 and port == 80392 and rule_action == "allow"393 ):394 state.is_resolved = True395 add_reward(0.8, "resolve_easy_web_ingress")396 done = True397 termination_reason = "resolved_easy"398 output += "\nSUCCESS: Web server is now accessible!"399 elif (400 self.task_name == "medium"401 and action.resource_id == "sg-db"402 and port == 5432403 and rule_action == "allow"404 ):405 investigated = (406 "read_logs" in self._achievements407 and "lookup_db_target" in self._achievements408 )409 if investigated:410 state.is_resolved = True411 add_reward(0.6, "resolve_medium_db_connectivity")412 done = True413 termination_reason = "resolved_medium"414 output += "\nSUCCESS: Database connection restored!"415 else:416 add_reward(-0.1, "unsafe_change_without_triage")417 output += (418 "\nWARNING: Change applied without incident triage. "419 "Inspect API logs and resolve DB IP via query_metadata before closing the incident."420 )421 elif rule_action == "deny":422 add_reward(-0.1, "deny_rule_during_incident")423 output += "\nWARNING: Deny rule applied during outage remediation."424 425 elif action.command == "restart_service":426 if not action.resource_id:427 raise ValueError("resource_id is required for restart_service.")428 if action.resource_id not in state.resources:429 raise ValueError(f"Resource {action.resource_id} not found.")430 431 output = f"Service on {action.resource_id} restarted."432 433 if self.task_name == "hard":434 if action.resource_id == "i-web2":435 investigated_root_cause = (436 "inspect_lb" in self._achievements437 and "inspect_target" in self._achievements438 and "lookup_upstream_target" in self._achievements439 )440 if investigated_root_cause:441 state.resources["i-web2"]["status"] = "running"442 state.resources["i-web2"][443 "logs"444 ] = "INFO: Restart successful. Memory cleared."445 state.is_resolved = True446 add_reward(0.8, "resolve_hard_upstream_recovery")447 done = True448 termination_reason = "resolved_hard"449 output += "\nSUCCESS: OutOfMemory loop broken. System stable."450 else:451 add_reward(-0.1, "restart_without_root_cause")452 output += (453 "\nWARNING: Restart denied by change policy. "454 "Find failing upstream IP from lb-main, resolve it with query_metadata, and inspect i-web2 first."455 )456 elif action.resource_id == "i-web1":457 add_reward(-0.2, "restart_healthy_node")458 output += (459 "\nWARNING: You restarted a healthy production server! "460 "Users dropped."461 )462 463 elif action.command == "submit_solution":464 if state.is_resolved:465 done = True466 termination_reason = "resolved_submit_solution"467 output = "Solution verified. System is HEALTHY."468 else:469 if self.task_name == "hard":470 # In hard mode, unresolved submission should not abort the run.471 done = False472 add_reward(-0.1, "premature_submit_hard")473 output = (474 "Solution incorrect. Incident is still CRITICAL. "475 "Continue triage and remediation before submitting."476 )477 else:478 done = True479 termination_reason = "incorrect_submit"480 output = "Solution incorrect. System is still CRITICAL."481 482 else:483 raise ValueError(f"Unsupported command: {action.command}")484 485 except Exception as exc:486 error = str(exc)487 output = f"Command Failed: {error}"488 489 cascade_penalty, cascade_msg = self._apply_cascading_failure()490 add_reward(cascade_penalty, "cascading_failure_penalty")491 if cascade_msg:492 output = f"{output}{cascade_msg}" if output else cascade_msg.strip()493 494 if state.step_count >= self.MAX_STEPS and not done:495 done = True496 termination_reason = "max_steps_timeout"497 timeout_suffix = "\nTIMEOUT: Max steps reached."498 output = f"{output}{timeout_suffix}" if output else timeout_suffix.strip()499 500 raw_reward = reward501 reward = max(-1.0, min(1.0, reward))502 if reward != raw_reward:503 reward_breakdown.append(504 {505 "event": "reward_clip",506 "delta": round(float(reward - raw_reward), 4),507 }508 )509 510 lb_external = state.resources.get("lb-external", {})511 if state.is_resolved:512 status = "HEALTHY"513 elif self.task_name == "hard" and lb_external.get("status") == "DOWN":514 status = "DEGRADED"515 else:516 status = "CRITICAL"517 info = {518 "step_count": state.step_count,519 "resolved": state.is_resolved,520 "task": self.task_name,521 "achievements": sorted(self._achievements),522 "total_resources": len(state.resources),523 "action_cost": self.ACTION_COST,524 "objective": self._task_objective(),525 "deterministic": True,526 "max_steps": self.MAX_STEPS,527 "termination_reason": termination_reason if done else "in_progress",528 "reward_breakdown": reward_breakdown,529 }530 531 return CloudObservation(532 output=output,533 error=error,534 system_health_status=status,535 done=done,536 reward=reward,537 metadata=info,538 echoed_message=output,539 message_length=len(output),540 )541 542 @property543 def state(self) -> State:544 """Return hidden environment state for evaluators/debugging."""545 if self._state_data is None:546 self.reset()547 assert self._state_data is not None548 return self._state_data549 