Team Ai
Apppublic

SidhaGarg/Cloud-DevOps-RLEnv

sourceHugging Faceupdated 6mo agoView on Hugging Face
0likes
cloud_devops_env_environment.py549 linesDownload Raw Back to server
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