openenv/coding_env

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"""8Synchronous wrapper for async EnvClient.9 10This module provides a SyncEnvClient that wraps an async EnvClient,11allowing synchronous usage while the underlying client uses async I/O.12 13Example:14    >>> from openenv.core import GenericEnvClient15    >>>16    >>> # Create async client and get sync wrapper17    >>> async_client = GenericEnvClient(base_url="http://localhost:8000")18    >>> sync_client = async_client.sync()19    >>>20    >>> # Use synchronous API21    >>> with sync_client:22    ...     result = sync_client.reset()23    ...     result = sync_client.step({"code": "print('hello')"})24"""25 26from __future__ import annotations27 28import asyncio29import concurrent.futures30import inspect31import threading32from typing import Any, Dict, Generic, TYPE_CHECKING, TypeVar33 34from .client_types import StateT, StepResult35 36if TYPE_CHECKING:37    from .env_client import EnvClient38 39ActT = TypeVar("ActT")40ObsT = TypeVar("ObsT")41 42 43class SyncEnvClient(Generic[ActT, ObsT, StateT]):44    """45    Synchronous wrapper around an async EnvClient.46 47    This class provides a synchronous interface to an async EnvClient,48    making it easier to use in synchronous code or to stop async from49    "infecting" the entire call stack.50 51    The wrapper executes async operations on a dedicated background event loop52    so connection state remains bound to a single loop.53 54    Cleanup note:55        For guaranteed resource cleanup, use `with SyncEnvClient(...)` or call56        `close()` explicitly. `__del__` is best-effort only and may not run57        reliably (for example, during interpreter shutdown).58 59    Example:60        >>> # From an async client61        >>> async_client = GenericEnvClient(base_url="http://localhost:8000")62        >>> sync_client = async_client.sync()63        >>>64        >>> # Use synchronous context manager65        >>> with sync_client:66        ...     result = sync_client.reset()67        ...     result = sync_client.step({"action": "test"})68 69    Attributes:70        _async: The wrapped async EnvClient instance71    """72 73    def __init__(self, async_client: "EnvClient[ActT, ObsT, StateT]"):74        """75        Initialize sync wrapper around an async client.76 77        Args:78            async_client: The async EnvClient to wrap79        """80        self._async = async_client81        self._loop: asyncio.AbstractEventLoop | None = None82        self._loop_thread: threading.Thread | None = None83        self._loop_ready = threading.Event()84        self._loop_init_lock = threading.Lock()85        self._async_wrapper_cache: Dict[str, Any] = {}86 87    def _run_loop_forever(self) -> None:88        """Run a dedicated event loop for this sync client."""89        loop = asyncio.new_event_loop()90        self._loop = loop91        asyncio.set_event_loop(loop)92        self._loop_ready.set()93        loop.run_forever()94        loop.close()95 96    def _ensure_loop(self) -> asyncio.AbstractEventLoop:97        """Start background loop thread on first use."""98        if (99            self._loop is not None100            and self._loop_thread101            and self._loop_thread.is_alive()102        ):103            return self._loop104 105        # Protect loop initialization when multiple threads race on first use.106        with self._loop_init_lock:107            if (108                self._loop is not None109                and self._loop_thread110                and self._loop_thread.is_alive()111            ):112                return self._loop113 114            self._loop_ready.clear()115            self._loop_thread = threading.Thread(116                target=self._run_loop_forever,117                name="openenv-sync-client-loop",118                daemon=True,119            )120            self._loop_thread.start()121            if not self._loop_ready.wait(timeout=5):122                raise RuntimeError("Timed out starting sync client event loop")123            assert self._loop is not None124            return self._loop125 126    def _run(self, coro: Any) -> Any:127        """Run coroutine on dedicated loop and block for result."""128        loop = self._ensure_loop()129        future: concurrent.futures.Future[Any] = asyncio.run_coroutine_threadsafe(130            coro, loop131        )132        return future.result()133 134    def _stop_loop(self) -> None:135        """Stop and join background loop thread."""136        loop = self._loop137        thread = self._loop_thread138        if loop is None:139            return140 141        if loop.is_running():142            loop.call_soon_threadsafe(loop.stop)143        if thread is not None:144            thread.join(timeout=5)145 146        self._loop = None147        self._loop_thread = None148 149    @property150    def async_client(self) -> "EnvClient[ActT, ObsT, StateT]":151        """Access the underlying async client."""152        return self._async153 154    def connect(self) -> "SyncEnvClient[ActT, ObsT, StateT]":155        """156        Establish connection to the server.157 158        Returns:159            self for method chaining160        """161        self._run(self._async.connect())162        return self163 164    def disconnect(self) -> None:165        """Close the connection."""166        self._run(self._async.disconnect())167 168    def reset(self, **kwargs: Any) -> StepResult[ObsT]:169        """170        Reset the environment.171 172        Args:173            **kwargs: Optional parameters passed to the environment's reset method174 175        Returns:176            StepResult containing initial observation177        """178        return self._run(self._async.reset(**kwargs))179 180    def step(self, action: ActT, **kwargs: Any) -> StepResult[ObsT]:181        """182        Execute an action in the environment.183 184        Args:185            action: The action to execute186            **kwargs: Optional parameters187 188        Returns:189            StepResult containing observation, reward, and done status190        """191        return self._run(self._async.step(action, **kwargs))192 193    def state(self) -> StateT:194        """195        Get the current environment state.196 197        Returns:198            State object with environment state information199        """200        return self._run(self._async.state())201 202    def close(self) -> None:203        """Close the connection and clean up resources."""204        try:205            self._run(self._async.close())206        finally:207            self._stop_loop()208 209    def __enter__(self) -> "SyncEnvClient[ActT, ObsT, StateT]":210        """Enter context manager, establishing connection."""211        self.connect()212        return self213 214    def __exit__(self, exc_type, exc_val, exc_tb) -> None:215        """Exit context manager, closing connection."""216        self.close()217 218    def __del__(self) -> None:219        """220        Best-effort cleanup for background loop thread.221 222        Do not rely on this for deterministic cleanup; prefer context-manager223        usage or an explicit `close()` call.224        """225        try:226            self._stop_loop()227        except Exception:228            pass229 230    def __getattr__(self, name: str) -> Any:231        """232        Delegate unknown attributes to the async client.233 234        Async methods are wrapped to run on the sync client's dedicated loop.235        """236        attr = getattr(self._async, name)237 238        if inspect.iscoroutinefunction(attr):239            cached = self._async_wrapper_cache.get(name)240            if cached is not None:241                return cached242 243            def sync_wrapper(*args: Any, **kwargs: Any) -> Any:244                method = getattr(self._async, name)245                return self._run(method(*args, **kwargs))246 247            self._async_wrapper_cache[name] = sync_wrapper248            return sync_wrapper249 250        return attr251 252    # Delegate abstract method implementations to the wrapped client253    def _step_payload(self, action: ActT) -> Dict[str, Any]:254        """Delegate to async client's _step_payload."""255        return self._async._step_payload(action)256 257    def _parse_result(self, payload: Dict[str, Any]) -> StepResult[ObsT]:258        """Delegate to async client's _parse_result."""259        return self._async._parse_result(payload)260 261    def _parse_state(self, payload: Dict[str, Any]) -> StateT:262        """Delegate to async client's _parse_state."""263        return self._async._parse_state(payload)264