openenv-testing/julia_env-pr-170

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        # Extract timeout_s from kwargs for wait_for_ready, with a default100        timeout_s = kwargs.pop('timeout_s', 30.0)101        request_timeout_s = kwargs.pop('request_timeout_s', 15.0)102 103        # 1. Start container with optional kwargs (e.g., env_vars, port)104        base_url = provider.start_container(image, **kwargs)105 106        # 2. Wait for server to be ready with the specified timeout107        provider.wait_for_ready(base_url, timeout_s=timeout_s)108 109        # 3. Create and return client instance with provider reference and request timeout110        return cls(base_url=base_url, request_timeout_s=request_timeout_s, provider=provider)111 112    @classmethod113    def from_hub(cls: Type[EnvClientT], repo_id: str, provider: Optional["ContainerProvider"] = None, **kwargs: Any) -> EnvClientT:114        """115        Create an environment client by pulling from a Hugging Face model hub.116        """117        118        if provider is None:119            provider = LocalDockerProvider()120        121        if "tag" in kwargs:122            tag = kwargs["tag"]123        else:124            tag = "latest"125        126        base_url = f"registry.hf.space/{repo_id.replace('/', '-')}:{tag}"127        128        return cls.from_docker_image(image=base_url, provider=provider)129 130    @abstractmethod131    def _step_payload(self, action: ActT) -> dict:132        """Convert an Action object to the JSON body expected by the env server."""133        raise NotImplementedError134 135    @abstractmethod136    def _parse_result(self, payload: dict) -> StepResult[ObsT]:137        """Convert a JSON response from the env server to StepResult[ObsT]."""138        raise NotImplementedError139 140    @abstractmethod141    def _parse_state(self, payload: dict) -> Any:142        """Convert a JSON response from the state endpoint to a State object."""143        raise NotImplementedError144 145    # ---------- Environment Server Interface Methods ----------146    def reset(self) -> StepResult[ObsT]:147        body: Dict[str, Any] = {}148        # TODO: later:149        # body["seed"] = seed150        # body["episode_id"] = episode_id151        r = self._http.post(152            f"{self._base}/reset",153            json=body,154            headers=self._headers,155            timeout=self._timeout,156        )157        r.raise_for_status()158        return self._parse_result(r.json())159 160    def step(self, action: ActT) -> StepResult[ObsT]:161        body: Dict[str, Any] = {162            "action": self._step_payload(action),163            "timeout_s": int(self._timeout),164        }165        # TODO: later:166        # body["request_id"] = str(uuid.uuid4())167        # body["episode_id"] = current_episode_id168        r = self._http.post(169            f"{self._base}/step",170            json=body,171            headers=self._headers,172            timeout=self._timeout,173        )174        r.raise_for_status()175        return self._parse_result(r.json())176 177    def state(self) -> Any:178        """179        Get the current environment state from the server.180 181        Returns:182            State object with environment state information (e.g., episode_id, step_count)183 184        Example:185            >>> client = EchoEnv.from_docker_image("echo-env:latest")186            >>> result = client.reset()187            >>> state = client.state()188            >>> print(state.episode_id)189            >>> print(state.step_count)190        """191        r = self._http.get(192            f"{self._base}/state",193            headers=self._headers,194            timeout=self._timeout,195        )196        r.raise_for_status()197        return self._parse_state(r.json())198 199    def close(self) -> None:200        """201        Close the environment and clean up resources.202 203        If this client was created via from_docker_image(), this will stop204        and remove the associated container.205        """206        if self._provider is not None:207            self._provider.stop_container()208