MeenalSinha/cloud-finops-optimizer
0
1"""2Python client for the Cloud FinOps Optimizer environment.3 4Extends openenv.core.env_client.EnvClient which handles all WebSocket and5HTTP communication. Implements the three required abstract methods:6 7 _step_payload - serialize FinOpsAction to the JSON dict sent over the wire8 _parse_result - deserialize wire response to StepResult9 _parse_state - deserialize wire response to FinOpsState10 11All v2 fields (goal, last_action_error, dependency_graph, cascading_risks,12sla_violations, simulate_result, sla_max_cpu, sla_status, resize_cooldown_steps,13reservation_committed) are fully deserialized so the inference script can use them.14 15Usage (sync):16 17 from client import FinOpsEnv, FinOpsAction18 19 with FinOpsEnv(base_url="http://localhost:7860").sync() as env:20 result = env.reset(task_id="task1")21 result = env.step(FinOpsAction(action_type="terminate", resource_id="ebs-001",22 reasoning="idle 720 hrs, no dependents"))23 print(result.reward, result.done)24 print(result.observation.goal)25 print(result.observation.dependency_graph)26 27Usage (async):28 29 import asyncio30 from client import FinOpsEnv, FinOpsAction31 32 async def main():33 async with FinOpsEnv(base_url="http://localhost:7860") as env:34 result = await env.reset(task_id="task3")35 result = await env.step(FinOpsAction(36 action_type="simulate",37 simulate_action={"action_type": "reserve", "resource_id": "ec2-h01"},38 reasoning="checking reservation safety before committing"39 ))40 print(result.observation.simulate_result.recommendation)41 asyncio.run(main())42"""43 44from __future__ import annotations45 46from typing import Any, Dict, List, Optional47 48from openenv.core.env_client import EnvClient49from openenv.core.client_types import StepResult50 51from models import (52 CloudResource,53 FinOpsAction,54 FinOpsObservation,55 FinOpsState,56 InstanceSize,57 ResourceStatus,58 ResourceType,59 SimulateResult,60 SLAStatus,61)62 63 64class FinOpsEnv(EnvClient[FinOpsAction, FinOpsObservation, FinOpsState]):65 """66 WebSocket/HTTP client for the Cloud FinOps Optimizer environment.67 68 The base class EnvClient handles connection management, WebSocket framing,69 and the sync() wrapper. You only need the three parsing methods below.70 """71 72 # ------------------------------------------------------------------73 # Required override 1: serialize action to wire format74 # ------------------------------------------------------------------75 76 def _step_payload(self, action: FinOpsAction) -> dict:77 """78 Serialize FinOpsAction to the JSON dict sent over the WebSocket.79 All fields must be included so the server receives simulate_action80 and reasoning correctly.81 """82 payload: Dict[str, Any] = {"action_type": action.action_type}83 if action.resource_id is not None:84 payload["resource_id"] = action.resource_id85 if action.target_size is not None:86 payload["target_size"] = action.target_size87 if action.simulate_action is not None:88 payload["simulate_action"] = action.simulate_action89 if action.reasoning is not None:90 payload["reasoning"] = action.reasoning91 return payload92 93 # ------------------------------------------------------------------94 # Required override 2: deserialize step/reset response95 # ------------------------------------------------------------------96 97 def _parse_result(self, payload: dict) -> StepResult:98 """99 Deserialize a step or reset response into a typed StepResult.100 101 The payload from /reset has the observation fields at the top level.102 The payload from /step wraps them under an "observation" key.103 Both shapes are handled via the obs_data fallback.104 """105 obs_data: dict = payload.get("observation") or payload106 107 resources = [108 self._parse_resource(r)109 for r in obs_data.get("resources", [])110 ]111 112 # Deserialize SimulateResult if present (Upgrade 5)113 sim_raw = obs_data.get("simulate_result")114 simulate_result: Optional[SimulateResult] = None115 if sim_raw:116 simulate_result = SimulateResult(117 proposed_action=sim_raw.get("proposed_action", {}),118 projected_cost_per_hour=sim_raw.get("projected_cost_per_hour", 0.0),119 projected_budget_remaining=sim_raw.get("projected_budget_remaining", 0.0),120 projected_reward=sim_raw.get("projected_reward", 0.0),121 projected_sla_violations=sim_raw.get("projected_sla_violations", []),122 cascading_risks=sim_raw.get("cascading_risks", []),123 recommendation=sim_raw.get("recommendation", ""),124 safe_to_apply=sim_raw.get("safe_to_apply", True),125 )126 127 observation = FinOpsObservation(128 done=payload.get("done", False),129 reward=payload.get("reward"),130 resources=resources,131 total_cost_per_hour=obs_data.get("total_cost_per_hour", 0.0),132 budget_per_hour=obs_data.get("budget_per_hour", 0.0),133 budget_remaining=obs_data.get("budget_remaining", 0.0),134 task_id=obs_data.get("task_id", ""),135 task_description=obs_data.get("task_description", ""),136 # v2 fields137 goal=obs_data.get("goal", ""),138 last_action_error=obs_data.get("last_action_error"),139 dependency_graph=obs_data.get("dependency_graph", {}),140 cascading_risks=obs_data.get("cascading_risks", {}),141 sla_violations=obs_data.get("sla_violations", []),142 simulate_result=simulate_result,143 step_count=obs_data.get("step_count", 0),144 max_steps=obs_data.get("max_steps", 20),145 info=obs_data.get("info", {}),146 )147 148 return StepResult(149 observation=observation,150 reward=payload.get("reward"),151 done=payload.get("done", False),152 )153 154 # ------------------------------------------------------------------155 # Required override 3: deserialize state response156 # ------------------------------------------------------------------157 158 def _parse_state(self, payload: dict) -> FinOpsState:159 """Deserialize a /state response into a typed FinOpsState."""160 return FinOpsState(161 episode_id=payload.get("episode_id"),162 step_count=payload.get("step_count", 0),163 task_id=payload.get("task_id", ""),164 total_cost_per_hour=payload.get("total_cost_per_hour", 0.0),165 initial_cost_per_hour=payload.get("initial_cost_per_hour", 0.0),166 budget_per_hour=payload.get("budget_per_hour", 0.0),167 terminated_ids=payload.get("terminated_ids", []),168 reserved_ids=payload.get("reserved_ids", []),169 resize_history=payload.get("resize_history", {}),170 sla_violation_history=payload.get("sla_violation_history", []),171 reasoning_log=payload.get("reasoning_log", []),172 done=payload.get("done", False),173 )174 175 # ------------------------------------------------------------------176 # Helper: parse a single CloudResource dict from the wire177 # ------------------------------------------------------------------178 179 @staticmethod180 def _parse_resource(r: dict) -> CloudResource:181 """182 Deserialize one resource dict including all v2 SLA and temporal fields.183 Defaults mirror the CloudResource model defaults so missing fields184 never cause KeyError or validation failure.185 """186 sla_status_raw = r.get("sla_status", "ok")187 try:188 sla_status = SLAStatus(sla_status_raw)189 except ValueError:190 sla_status = SLAStatus.OK191 192 return CloudResource(193 id=r["id"],194 name=r.get("name", ""),195 type=ResourceType(r["type"]),196 instance_size=InstanceSize(r["instance_size"]) if r.get("instance_size") else None,197 cpu_utilization=r.get("cpu_utilization", 0.0),198 memory_utilization=r.get("memory_utilization", 0.0),199 cost_per_hour=r.get("cost_per_hour", 0.0),200 status=ResourceStatus(r.get("status", "running")),201 critical=r.get("critical", False),202 reserved=r.get("reserved", False),203 idle_hours=r.get("idle_hours", 0),204 dependency_ids=r.get("dependency_ids", []),205 tags=r.get("tags", {}),206 # v2 SLA fields (Upgrade 2)207 sla_max_cpu=r.get("sla_max_cpu", 90.0),208 sla_uptime_pct=r.get("sla_uptime_pct", 99.9),209 sla_status=sla_status,210 # v2 temporal fields (Upgrade 3)211 resize_cooldown_steps=r.get("resize_cooldown_steps", 0),212 reservation_committed=r.get("reservation_committed", False),213 )214