Spaces:
Running
Running
File size: 2,351 Bytes
e313085 | 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 | # 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)
|