Spaces:
Sleeping
Sleeping
File size: 7,525 Bytes
7e69b8f | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 | # src/client.py
"""
HTTP client for AquaGuard-RL environment.
Provides a synchronous client that connects to a running AquaGuard-RL
server (local or remote) and provides the standard RL interface:
env.reset(task, seed) → AquaGuardObservation
env.step(action) → AquaGuardObservation
env.state → AquaGuardState
env.close() → stops Docker container if started via from_docker_image()
Usage:
# Connect to a running server
env = AquaGuardEnv("http://localhost:8000")
obs = env.reset(task="baseline", seed=42)
while not obs.done:
action = ...
obs = env.step(action)
env.close()
# Start from Docker image
env = AquaGuardEnv.from_docker_image("aquaguard-env:latest")
obs = env.reset(task="crisis")
"""
from __future__ import annotations
import logging
import time
from typing import Optional
logger = logging.getLogger(__name__)
from models import AquaGuardAction, AquaGuardObservation, AquaGuardState
def _parse_observation(data: dict) -> AquaGuardObservation:
"""Parse server response into AquaGuardObservation.
Handles both flat responses and openenv_core's nested format:
{"observation": {...}, "reward": ..., "done": ...}
"""
if "observation" in data and isinstance(data["observation"], dict):
# Nested format from openenv_core Environment
obs_data = data["observation"]
obs_data["reward"] = data.get("reward")
obs_data["done"] = data.get("done", False)
return AquaGuardObservation(**obs_data)
return AquaGuardObservation(**data)
class AquaGuardEnv:
"""
Synchronous HTTP client for AquaGuard-RL environment.
Implements the standard RL interface (reset/step/state) backed by
HTTP requests to a running FastAPI environment server.
"""
def __init__(self, base_url: str = "http://localhost:8000") -> None:
"""
Initialize client connected to a running environment server.
Args:
base_url: Base URL of the environment server.
"""
try:
import httpx
self._http = httpx.Client(timeout=120.0)
except ImportError:
raise ImportError("httpx required: pip install httpx")
self._base_url = base_url.rstrip("/")
self._container: Optional[str] = None
logger.debug(f"AquaGuardEnv client initialized: {self._base_url}")
@classmethod
def from_docker_image(
cls,
image: str = "aquaguard-env:latest",
port: int = 8000,
timeout: int = 60,
) -> "AquaGuardEnv":
"""
Start AquaGuard-RL in a Docker container and return connected client.
Args:
image: Docker image name/tag.
port: Host port to bind (default 8000).
timeout: Seconds to wait for server startup.
Returns:
Connected AquaGuardEnv client.
Raises:
RuntimeError: If server fails to start within timeout.
"""
import subprocess
import httpx
logger.info(f"Starting Docker container: {image}")
result = subprocess.run(
["docker", "run", "-d", "-p", f"{port}:8000", image],
capture_output=True, text=True, check=True,
)
container_id = result.stdout.strip()
logger.info(f"Container started: {container_id[:12]}")
base_url = f"http://localhost:{port}"
for attempt in range(timeout):
try:
resp = httpx.get(f"{base_url}/health", timeout=2.0)
if resp.status_code == 200:
env = cls(base_url)
env._container = container_id
logger.info(f"Server ready after {attempt + 1}s")
return env
except Exception:
pass
time.sleep(1.0)
# Cleanup failed container
subprocess.run(["docker", "stop", container_id], capture_output=True)
raise RuntimeError(
f"Environment server failed to start within {timeout}s. "
f"Check Docker logs: docker logs {container_id[:12]}"
)
def reset(
self,
task: str = "baseline",
seed: Optional[int] = None,
episode_id: Optional[str] = None,
) -> AquaGuardObservation:
"""
Reset environment to start a new episode.
Args:
task: Task name ('baseline', 'crisis', 'policy_shift', 'climate_shock', 'multi_district').
seed: Optional random seed for reproducibility.
episode_id: Optional explicit episode ID.
Returns:
Initial AquaGuardObservation (step_number=0, reward=None, done=False).
"""
payload: dict = {"task": task}
if seed is not None:
payload["seed"] = seed
if episode_id is not None:
payload["episode_id"] = episode_id
resp = self._http.post(f"{self._base_url}/reset", json=payload)
resp.raise_for_status()
return _parse_observation(resp.json())
def step(self, action: AquaGuardAction) -> AquaGuardObservation:
"""
Execute one policy step (one growing season ~4 months).
Args:
action: AquaGuardAction with crop allocation, water quotas, etc.
Returns:
AquaGuardObservation with updated state, reward, and done flag.
"""
payload = action.model_dump() if hasattr(action, "model_dump") else action.dict()
# Try openenv_core wrapped format first: {"action": {...}}, then flat fallback
resp = self._http.post(f"{self._base_url}/step", json={"action": payload})
if resp.status_code == 422:
# Server may use standalone (flat) format
resp = self._http.post(f"{self._base_url}/step", json=payload)
resp.raise_for_status()
return _parse_observation(resp.json())
@property
def state(self) -> AquaGuardState:
"""Get current episode state metadata."""
resp = self._http.get(f"{self._base_url}/state")
resp.raise_for_status()
return AquaGuardState(**resp.json())
def close(self) -> None:
"""
Close the HTTP client and stop Docker container if started via from_docker_image().
"""
try:
self._http.close()
except Exception:
pass
if self._container:
import subprocess
logger.info(f"Stopping container {self._container[:12]}")
subprocess.run(["docker", "stop", self._container], capture_output=True)
subprocess.run(["docker", "rm", self._container], capture_output=True)
def __enter__(self) -> "AquaGuardEnv":
return self
def __exit__(self, *args) -> None:
self.close()
def health_check(self) -> bool:
"""Check if the server is running and healthy."""
try:
resp = self._http.get(f"{self._base_url}/health", timeout=5.0)
return resp.status_code == 200
except Exception:
return False
def get_info(self) -> dict:
"""Get environment metadata from the server."""
resp = self._http.get(f"{self._base_url}/info")
resp.raise_for_status()
return resp.json() |