openenv-testing/android_env-pr-162
0
1"""2core/runner_env.py3Minimal HTTP-based environment client.4- Talks to a single env worker exposing: POST /reset, POST /step5 6Future hooks (commented below) for:7- episode_id, seed on reset8- request_id on step9- custom headers (auth/trace)10"""11 12from __future__ import annotations13 14from abc import ABC, abstractmethod15from typing import Any, Dict, Generic, Optional, Type, TYPE_CHECKING, TypeVar16 17import requests18 19from .client_types import StepResult20from .containers.runtime import LocalDockerProvider21 22if TYPE_CHECKING:23 from .containers.runtime import ContainerProvider24 25ActT = TypeVar("ActT")26ObsT = TypeVar("ObsT")27EnvClientT = TypeVar("EnvClientT", bound="HTTPEnvClient")28 29 30class HTTPEnvClient(ABC, Generic[ActT, ObsT]):31 def __init__(32 self,33 base_url: str,34 request_timeout_s: float = 15.0,35 default_headers: Optional[Dict[str, str]] = None,36 provider: Optional["ContainerProvider"] = None,37 ):38 self._base = base_url.rstrip("/")39 self._timeout = float(request_timeout_s)40 self._http = requests.Session()41 self._headers = default_headers or {}42 self._provider = provider43 44 @classmethod45 def from_docker_image(46 cls: Type[EnvClientT],47 image: str,48 provider: Optional["ContainerProvider"] = None,49 **kwargs: Any,50 ) -> EnvClientT:51 """52 Create an environment client by spinning up a Docker container locally.53 54 This is a development utility that:55 1. Starts a Docker container from the specified image56 2. Waits for the server to be ready57 3. Creates and returns a client instance connected to the container58 59 Note: The container lifecycle management is left to the user or higher-level60 orchestration. The container will keep running until manually stopped.61 62 Args:63 image: Docker image name to run (e.g., "echo-env:latest")64 provider: Container provider to use (defaults to LocalDockerProvider)65 **kwargs: Additional arguments to pass to provider.start_container()66 (e.g., env_vars, port)67 68 Returns:69 An instance of the client class connected to the running container70 71 Example:72 >>> from envs.coding_env.client import CodingEnv73 >>> from envs.coding_env.models import CodeAction74 >>>75 >>> # Create environment from image76 >>> env = CodingEnv.from_docker_image("coding-env:latest")77 >>>78 >>> # Create environment with custom env vars79 >>> env = CodingEnv.from_docker_image(80 ... "coding-env:latest",81 ... env_vars={"MY_VAR": "value"}82 ... )83 >>>84 >>> # Use the environment85 >>> result = env.reset()86 >>> print(result.observation)87 >>>88 >>> step_result = env.step(CodeAction(code="print('hello')"))89 >>> print(step_result.observation.stdout)90 >>>91 >>> # Cleanup (optional)92 >>> env.close()93 """94 95 # Use default provider if none provided96 if provider is None:97 provider = LocalDockerProvider()98 99 # 1. Start container with optional kwargs (e.g., env_vars, port)100 base_url = provider.start_container(image, **kwargs)101 102 # 2. Wait for server to be ready103 provider.wait_for_ready(base_url)104 105 # 3. Create and return client instance with provider reference106 return cls(base_url=base_url, provider=provider)107 108 @classmethod109 def from_hub(cls: Type[EnvClientT], repo_id: str, provider: Optional["ContainerProvider"] = None, **kwargs: Any) -> EnvClientT:110 """111 Create an environment client by pulling from a Hugging Face model hub.112 """113 114 if provider is None:115 provider = LocalDockerProvider()116 117 if "tag" in kwargs:118 tag = kwargs["tag"]119 else:120 tag = "latest"121 122 base_url = f"registry.hf.space/{repo_id.replace('/', '-')}:{tag}"123 124 return cls.from_docker_image(image=base_url, provider=provider)125 126 @abstractmethod127 def _step_payload(self, action: ActT) -> dict:128 """Convert an Action object to the JSON body expected by the env server."""129 raise NotImplementedError130 131 @abstractmethod132 def _parse_result(self, payload: dict) -> StepResult[ObsT]:133 """Convert a JSON response from the env server to StepResult[ObsT]."""134 raise NotImplementedError135 136 @abstractmethod137 def _parse_state(self, payload: dict) -> Any:138 """Convert a JSON response from the state endpoint to a State object."""139 raise NotImplementedError140 141 # ---------- Environment Server Interface Methods ----------142 def reset(self) -> StepResult[ObsT]:143 body: Dict[str, Any] = {}144 # TODO: later:145 # body["seed"] = seed146 # body["episode_id"] = episode_id147 r = self._http.post(148 f"{self._base}/reset",149 json=body,150 headers=self._headers,151 timeout=self._timeout,152 )153 r.raise_for_status()154 return self._parse_result(r.json())155 156 def step(self, action: ActT) -> StepResult[ObsT]:157 body: Dict[str, Any] = {158 "action": self._step_payload(action),159 "timeout_s": int(self._timeout),160 }161 # TODO: later:162 # body["request_id"] = str(uuid.uuid4())163 # body["episode_id"] = current_episode_id164 r = self._http.post(165 f"{self._base}/step",166 json=body,167 headers=self._headers,168 timeout=self._timeout,169 )170 r.raise_for_status()171 return self._parse_result(r.json())172 173 def state(self) -> Any:174 """175 Get the current environment state from the server.176 177 Returns:178 State object with environment state information (e.g., episode_id, step_count)179 180 Example:181 >>> client = EchoEnv.from_docker_image("echo-env:latest")182 >>> result = client.reset()183 >>> state = client.state()184 >>> print(state.episode_id)185 >>> print(state.step_count)186 """187 r = self._http.get(188 f"{self._base}/state",189 headers=self._headers,190 timeout=self._timeout,191 )192 r.raise_for_status()193 return self._parse_state(r.json())194 195 def close(self) -> None:196 """197 Close the environment and clean up resources.198 199 If this client was created via from_docker_image(), this will stop200 and remove the associated container.201 """202 if self._provider is not None:203 self._provider.stop_container()204