Spaces:
Running
Running
| # Copyright (c) Meta Platforms, Inc. and affiliates. | |
| # All rights reserved. | |
| # | |
| # This source code is licensed under the BSD-style license found in the | |
| # LICENSE file in the root directory of this source tree. | |
| """Kernel Env Environment Client.""" | |
| import os | |
| from typing import Dict | |
| from openenv.core import EnvClient | |
| from websockets.asyncio.client import connect as ws_connect | |
| from openenv.core.client_types import StepResult | |
| from .models import KernelAction, KernelObservation, KernelState | |
| class KernelEnv(EnvClient[KernelAction, KernelObservation, KernelState]): | |
| """ | |
| Client for the Kernel Env Environment. | |
| Example: | |
| >>> with KernelEnv(base_url="http://localhost:8000").sync() as env: | |
| ... result = env.reset(seed=0) | |
| ... print(result.observation.problem) | |
| ... result = env.step(KernelAction(code=my_triton_source)) | |
| ... print(result.observation.stage, result.reward) | |
| Reset kwargs: `task_id="level1/19_ReLU"` to pick a task, `split="eval"` for held-out tasks. | |
| """ | |
| async def connect(self) -> "KernelEnv": | |
| # Modal endpoints are deployed with requires_proxy_auth; send the token if one is configured. | |
| key, secret = os.environ.get("MODAL_PROXY_KEY"), os.environ.get("MODAL_PROXY_SECRET") | |
| if self._ws is not None or not (key and secret): | |
| return await super().connect() | |
| try: | |
| self._ws = await ws_connect( | |
| self._ws_url, | |
| open_timeout=self._connect_timeout, | |
| max_size=self._max_message_size, | |
| additional_headers={"Modal-Key": key, "Modal-Secret": secret}, | |
| ) | |
| except Exception as e: | |
| raise ConnectionError(f"Failed to connect to {self._ws_url}: {e}") from e | |
| return self | |
| def _step_payload(self, action: KernelAction) -> Dict: | |
| return action.model_dump(exclude={"metadata"}) | |
| def _parse_result(self, payload: Dict) -> StepResult[KernelObservation]: | |
| obs = KernelObservation( | |
| **payload.get("observation", {}), | |
| done=payload.get("done", False), | |
| reward=payload.get("reward"), | |
| ) | |
| return StepResult(observation=obs, reward=payload.get("reward"), done=obs.done) | |
| def _parse_state(self, payload: Dict) -> KernelState: | |
| return KernelState(**payload) | |