Upload folder using huggingface_hub
Browse filesThis view is limited to 50 files because it contains too many changes. See raw diff
- Dockerfile +85 -0
- README.md +161 -5
- __init__.py +34 -0
- client.py +107 -0
- discovery.json +9 -0
- envs/chess_env/README.md +144 -0
- envs/chess_env/__init__.py +34 -0
- envs/chess_env/client.py +107 -0
- envs/chess_env/discovery.json +9 -0
- envs/chess_env/models.py +68 -0
- envs/chess_env/openenv.yaml +6 -0
- envs/chess_env/pyproject.toml +31 -0
- envs/chess_env/scripts/verify.sh +60 -0
- envs/chess_env/server/Dockerfile +79 -0
- envs/chess_env/server/__init__.py +11 -0
- envs/chess_env/server/app.py +35 -0
- envs/chess_env/server/chess_environment.py +337 -0
- envs/chess_env/server/rubrics.py +86 -0
- models.py +68 -0
- openenv.yaml +6 -0
- pyproject.toml +31 -0
- scripts/verify.sh +60 -0
- server/Dockerfile +79 -0
- server/__init__.py +11 -0
- server/app.py +35 -0
- server/chess_environment.py +337 -0
- server/rubrics.py +86 -0
- src/__init__.py +3 -0
- src/core/README.md +237 -0
- src/core/__init__.py +77 -0
- src/core/client_types.py +29 -0
- src/core/containers/__init__.py +3 -0
- src/core/containers/images/Dockerfile +60 -0
- src/core/containers/images/README.md +92 -0
- src/core/containers/runtime/__init__.py +26 -0
- src/core/containers/runtime/_server_config.py +53 -0
- src/core/containers/runtime/aca_provider.py +690 -0
- src/core/containers/runtime/daytona_provider.py +547 -0
- src/core/containers/runtime/hf_sandbox_provider.py +342 -0
- src/core/containers/runtime/modal_provider.py +697 -0
- src/core/containers/runtime/novita_provider.py +1216 -0
- src/core/containers/runtime/providers.py +727 -0
- src/core/containers/runtime/uv_provider.py +309 -0
- src/core/containers/test_local_docker_provider.py +260 -0
- src/core/env_client.py +1197 -0
- src/core/env_server/__init__.py +155 -0
- src/core/env_server/_utils.py +10 -0
- src/core/env_server/base_transforms.py +25 -0
- src/core/env_server/exceptions.py +101 -0
- src/core/env_server/gradio_theme.py +142 -0
Dockerfile
ADDED
|
@@ -0,0 +1,85 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 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 the
|
| 5 |
+
# LICENSE file in the root directory of this source tree.
|
| 6 |
+
|
| 7 |
+
# Multi-stage build using openenv-base
|
| 8 |
+
# This Dockerfile is flexible and works for both:
|
| 9 |
+
# - In-repo environments (with local src/core)
|
| 10 |
+
# - Standalone environments (with openenv from pip)
|
| 11 |
+
# The build script (openenv build) handles context detection and sets appropriate build args.
|
| 12 |
+
|
| 13 |
+
ARG BASE_IMAGE=ghcr.io/huggingface/openenv-base:latest
|
| 14 |
+
FROM ghcr.io/huggingface/openenv-base:latest AS builder
|
| 15 |
+
|
| 16 |
+
WORKDIR /app
|
| 17 |
+
|
| 18 |
+
# Build argument to control whether we're building standalone or in-repo
|
| 19 |
+
ARG BUILD_MODE=in-repo
|
| 20 |
+
|
| 21 |
+
# Copy environment code (always at root of build context)
|
| 22 |
+
COPY . /app/env
|
| 23 |
+
|
| 24 |
+
# For in-repo builds, openenv is already in the pyproject.toml dependencies
|
| 25 |
+
# For standalone builds, openenv will be installed from pip via pyproject.toml
|
| 26 |
+
WORKDIR /app/env
|
| 27 |
+
|
| 28 |
+
# Ensure uv is available (for local builds where base image lacks it)
|
| 29 |
+
RUN if ! command -v uv >/dev/null 2>&1; then \
|
| 30 |
+
curl -LsSf https://astral.sh/uv/install.sh | sh && \
|
| 31 |
+
mv /root/.local/bin/uv /usr/local/bin/uv && \
|
| 32 |
+
mv /root/.local/bin/uvx /usr/local/bin/uvx; \
|
| 33 |
+
fi
|
| 34 |
+
|
| 35 |
+
# Install git for building from git repos (build-time only)
|
| 36 |
+
RUN apt-get update && apt-get install -y --no-install-recommends \
|
| 37 |
+
git \
|
| 38 |
+
&& rm -rf /var/lib/apt/lists/*
|
| 39 |
+
|
| 40 |
+
# Install dependencies using uv sync
|
| 41 |
+
# First pass: install dependencies without the project (for better caching)
|
| 42 |
+
# Second pass: install the project itself
|
| 43 |
+
RUN curl -LsSf https://astral.sh/uv/install.sh | sh && \
|
| 44 |
+
install -m 0755 /root/.local/bin/uv /usr/local/bin/uv && \
|
| 45 |
+
install -m 0755 /root/.local/bin/uvx /usr/local/bin/uvx
|
| 46 |
+
|
| 47 |
+
RUN --mount=type=cache,target=/root/.cache/uv \
|
| 48 |
+
if [ -f uv.lock ]; then \
|
| 49 |
+
uv sync --no-install-project --no-editable; \
|
| 50 |
+
else \
|
| 51 |
+
uv sync --no-install-project --no-editable; \
|
| 52 |
+
fi
|
| 53 |
+
|
| 54 |
+
RUN --mount=type=cache,target=/root/.cache/uv \
|
| 55 |
+
if [ -f uv.lock ]; then \
|
| 56 |
+
uv sync --no-editable; \
|
| 57 |
+
else \
|
| 58 |
+
uv sync --no-editable; \
|
| 59 |
+
fi
|
| 60 |
+
|
| 61 |
+
# Final runtime stage
|
| 62 |
+
FROM ghcr.io/huggingface/openenv-base:latest
|
| 63 |
+
|
| 64 |
+
WORKDIR /app
|
| 65 |
+
|
| 66 |
+
# Copy the virtual environment from builder
|
| 67 |
+
COPY --from=builder /app/env/.venv /app/.venv
|
| 68 |
+
|
| 69 |
+
# Copy the environment code
|
| 70 |
+
COPY --from=builder /app/env /app/env
|
| 71 |
+
|
| 72 |
+
# Set PATH to use the virtual environment
|
| 73 |
+
ENV PATH="/app/.venv/bin:$PATH"
|
| 74 |
+
|
| 75 |
+
# Set PYTHONPATH so imports work correctly
|
| 76 |
+
ENV PYTHONPATH="/app/env:$PYTHONPATH"
|
| 77 |
+
|
| 78 |
+
# Health check using Python (more portable than curl/wget)
|
| 79 |
+
HEALTHCHECK --interval=30s --timeout=3s --start-period=5s --retries=3 \
|
| 80 |
+
CMD python -c "import urllib.request; urllib.request.urlopen('http://localhost:8000/health')" || exit 1
|
| 81 |
+
|
| 82 |
+
# Run the FastAPI server from the installed package.
|
| 83 |
+
CMD ["uvicorn", "chess_env.server.app:app", "--host", "0.0.0.0", "--port", "8000"]
|
| 84 |
+
|
| 85 |
+
ENV ENABLE_WEB_INTERFACE=true
|
README.md
CHANGED
|
@@ -1,10 +1,166 @@
|
|
| 1 |
---
|
| 2 |
-
title: Chess
|
| 3 |
-
emoji:
|
| 4 |
-
colorFrom:
|
| 5 |
-
colorTo:
|
| 6 |
sdk: docker
|
| 7 |
pinned: false
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 8 |
---
|
| 9 |
|
| 10 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
---
|
| 2 |
+
title: Chess Environment Server
|
| 3 |
+
emoji: ♟️
|
| 4 |
+
colorFrom: blue
|
| 5 |
+
colorTo: green
|
| 6 |
sdk: docker
|
| 7 |
pinned: false
|
| 8 |
+
app_port: 8000
|
| 9 |
+
base_path: /web
|
| 10 |
+
tags:
|
| 11 |
+
- openenv-0.10.0
|
| 12 |
+
- openenv
|
| 13 |
---
|
| 14 |
|
| 15 |
+
# Chess Environment
|
| 16 |
+
|
| 17 |
+
A chess reinforcement learning environment for OpenEnv, powered by the [moonfish](https://github.com/luccabb/moonfish) chess engine.
|
| 18 |
+
|
| 19 |
+

|
| 20 |
+
|
| 21 |
+
*The `/web` playground after a few moves: the board and one button per legal move.*
|
| 22 |
+
|
| 23 |
+
## Features
|
| 24 |
+
|
| 25 |
+
- **Full chess rules** via python-chess library
|
| 26 |
+
- **Configurable opponent**: moonfish engine, random moves, or self-play
|
| 27 |
+
- **Position evaluation**: Uses moonfish's PSQT-based evaluation
|
| 28 |
+
- **Standard OpenEnv interface**: reset(), step(), state
|
| 29 |
+
|
| 30 |
+
## Quick Start
|
| 31 |
+
|
| 32 |
+
### Using Docker
|
| 33 |
+
|
| 34 |
+
```bash
|
| 35 |
+
# Build the image
|
| 36 |
+
docker build -t chess-env:latest -f envs/chess_env/server/Dockerfile .
|
| 37 |
+
|
| 38 |
+
# Run the server
|
| 39 |
+
docker run -p 8000:8000 chess-env:latest
|
| 40 |
+
```
|
| 41 |
+
|
| 42 |
+
### Using the Client
|
| 43 |
+
|
| 44 |
+
The client is **async by default**:
|
| 45 |
+
|
| 46 |
+
```python
|
| 47 |
+
import asyncio
|
| 48 |
+
from chess_env import ChessEnv, ChessAction
|
| 49 |
+
|
| 50 |
+
async def main():
|
| 51 |
+
async with ChessEnv(base_url="http://localhost:8000") as env:
|
| 52 |
+
# Reset for a new game
|
| 53 |
+
result = await env.reset()
|
| 54 |
+
print(f"Starting position: {result.observation.fen}")
|
| 55 |
+
print(f"Legal moves: {result.observation.legal_moves}")
|
| 56 |
+
|
| 57 |
+
# Make a move
|
| 58 |
+
result = await env.step(ChessAction(move="e2e4"))
|
| 59 |
+
print(f"Reward: {result.reward}, Done: {result.done}")
|
| 60 |
+
|
| 61 |
+
# Play until game ends
|
| 62 |
+
while not result.done:
|
| 63 |
+
move = result.observation.legal_moves[0]
|
| 64 |
+
result = await env.step(ChessAction(move=move))
|
| 65 |
+
|
| 66 |
+
print(f"Game result: {result.observation.result}")
|
| 67 |
+
|
| 68 |
+
asyncio.run(main())
|
| 69 |
+
```
|
| 70 |
+
|
| 71 |
+
For **synchronous usage**, use the `.sync()` wrapper:
|
| 72 |
+
|
| 73 |
+
```python
|
| 74 |
+
from chess_env import ChessEnv, ChessAction
|
| 75 |
+
|
| 76 |
+
with ChessEnv(base_url="http://localhost:8000").sync() as env:
|
| 77 |
+
result = env.reset()
|
| 78 |
+
result = env.step(ChessAction(move="e2e4"))
|
| 79 |
+
print(f"Reward: {result.reward}")
|
| 80 |
+
```
|
| 81 |
+
|
| 82 |
+
## Observation Space
|
| 83 |
+
|
| 84 |
+
| Field | Type | Description |
|
| 85 |
+
|-------|------|-------------|
|
| 86 |
+
| `fen` | str | Board position in FEN notation |
|
| 87 |
+
| `legal_moves` | List[str] | Legal moves in UCI format |
|
| 88 |
+
| `is_check` | bool | Whether current player is in check |
|
| 89 |
+
| `done` | bool | Whether game has ended |
|
| 90 |
+
| `reward` | float | Reward for last action |
|
| 91 |
+
| `result` | str | Game result ("1-0", "0-1", "1/2-1/2") |
|
| 92 |
+
|
| 93 |
+
## Action Space
|
| 94 |
+
|
| 95 |
+
| Field | Type | Description |
|
| 96 |
+
|-------|------|-------------|
|
| 97 |
+
| `move` | str | UCI format move (e.g., "e2e4", "e7e8q") |
|
| 98 |
+
|
| 99 |
+
## Rewards
|
| 100 |
+
|
| 101 |
+
| Outcome | Reward |
|
| 102 |
+
|---------|--------|
|
| 103 |
+
| Win | +1.0 |
|
| 104 |
+
| Loss | -1.0 |
|
| 105 |
+
| Draw | 0.0 |
|
| 106 |
+
| Illegal move | -0.1 |
|
| 107 |
+
|
| 108 |
+
## Configuration
|
| 109 |
+
|
| 110 |
+
The environment supports these configuration options:
|
| 111 |
+
|
| 112 |
+
| Parameter | Default | Description |
|
| 113 |
+
|-----------|---------|-------------|
|
| 114 |
+
| `opponent` | "moonfish" | Opponent type: "moonfish", "random", or None |
|
| 115 |
+
| `opponent_depth` | 2 | Search depth for moonfish opponent |
|
| 116 |
+
| `max_moves` | 500 | Maximum half-moves before draw |
|
| 117 |
+
| `agent_color` | None | Agent color: "white", "black", or None (alternate each episode) |
|
| 118 |
+
| `gamma` | 0.99 | Discount factor for temporal credit assignment |
|
| 119 |
+
|
| 120 |
+
## Temporal Discounting
|
| 121 |
+
|
| 122 |
+
For RL training, the environment computes temporally discounted rewards at episode end. This helps with credit assignment in long games where only the final outcome is known.
|
| 123 |
+
|
| 124 |
+
When an episode ends, the terminal observation's `metadata` includes:
|
| 125 |
+
- `discounted_rewards`: List of discounted rewards for each agent move
|
| 126 |
+
- `gamma`: The discount factor used
|
| 127 |
+
|
| 128 |
+
The formula is `r_t = γ^(T-1-t) × R_final` where:
|
| 129 |
+
- `T` = total agent moves
|
| 130 |
+
- `t` = move index (0-indexed)
|
| 131 |
+
- `R_final` = terminal reward (+1, -1, or 0)
|
| 132 |
+
|
| 133 |
+
Example for a 5-move win with γ=0.99:
|
| 134 |
+
```
|
| 135 |
+
Move 0: 0.99^4 × 1.0 = 0.961
|
| 136 |
+
Move 1: 0.99^3 × 1.0 = 0.970
|
| 137 |
+
Move 2: 0.99^2 × 1.0 = 0.980
|
| 138 |
+
Move 3: 0.99^1 × 1.0 = 0.990
|
| 139 |
+
Move 4: 0.99^0 × 1.0 = 1.000
|
| 140 |
+
```
|
| 141 |
+
|
| 142 |
+
## Links
|
| 143 |
+
|
| 144 |
+
- [moonfish GitHub](https://github.com/luccabb/moonfish)
|
| 145 |
+
- [Play online](https://huggingface.co/spaces/luccabb/moonfish_chess)
|
| 146 |
+
|
| 147 |
+
## Hugging Face Space Deployment
|
| 148 |
+
|
| 149 |
+
This Space is built from OpenEnv environment `chess_env`.
|
| 150 |
+
|
| 151 |
+
- Space URL: `https://huggingface.co/spaces/openenv/chess_env`
|
| 152 |
+
- OpenEnv pinned ref: `0.10.0`
|
| 153 |
+
- Hub tag: `openenv`
|
| 154 |
+
|
| 155 |
+
### Connecting from Code
|
| 156 |
+
|
| 157 |
+
```bash
|
| 158 |
+
pip install git+https://huggingface.co/spaces/openenv/chess_env
|
| 159 |
+
```
|
| 160 |
+
|
| 161 |
+
```python
|
| 162 |
+
from chess_env import ChessEnv
|
| 163 |
+
|
| 164 |
+
with ChessEnv(base_url="https://openenv-chess-env.hf.space").sync() as env:
|
| 165 |
+
result = env.reset()
|
| 166 |
+
```
|
__init__.py
ADDED
|
@@ -0,0 +1,34 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 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 the
|
| 5 |
+
# LICENSE file in the root directory of this source tree.
|
| 6 |
+
|
| 7 |
+
"""
|
| 8 |
+
Chess Environment for OpenEnv.
|
| 9 |
+
|
| 10 |
+
This module provides OpenEnv integration for chess, using the moonfish
|
| 11 |
+
chess engine for position evaluation and opponent play.
|
| 12 |
+
|
| 13 |
+
Example:
|
| 14 |
+
>>> from envs.chess_env import ChessEnv, ChessAction
|
| 15 |
+
>>>
|
| 16 |
+
>>> # Connect to a running server or start via Docker
|
| 17 |
+
>>> env = ChessEnv.from_docker_image("chess-env:latest")
|
| 18 |
+
>>>
|
| 19 |
+
>>> # Reset and interact
|
| 20 |
+
>>> result = env.reset()
|
| 21 |
+
>>> print(result.observation.fen)
|
| 22 |
+
>>> print(result.observation.legal_moves)
|
| 23 |
+
>>>
|
| 24 |
+
>>> result = env.step(ChessAction(move="e2e4"))
|
| 25 |
+
>>> print(result.reward, result.done)
|
| 26 |
+
>>>
|
| 27 |
+
>>> # Cleanup
|
| 28 |
+
>>> env.close()
|
| 29 |
+
"""
|
| 30 |
+
|
| 31 |
+
from .client import ChessEnv
|
| 32 |
+
from .models import ChessAction, ChessObservation, ChessState
|
| 33 |
+
|
| 34 |
+
__all__ = ["ChessEnv", "ChessAction", "ChessObservation", "ChessState"]
|
client.py
ADDED
|
@@ -0,0 +1,107 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 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 the
|
| 5 |
+
# LICENSE file in the root directory of this source tree.
|
| 6 |
+
|
| 7 |
+
"""
|
| 8 |
+
Chess Environment Client.
|
| 9 |
+
|
| 10 |
+
This module provides the client for connecting to a Chess Environment server
|
| 11 |
+
via WebSocket for persistent sessions.
|
| 12 |
+
"""
|
| 13 |
+
|
| 14 |
+
from __future__ import annotations
|
| 15 |
+
|
| 16 |
+
from typing import Any, Dict
|
| 17 |
+
|
| 18 |
+
from openenv.core.client_types import StepResult
|
| 19 |
+
from openenv.core.env_client import EnvClient
|
| 20 |
+
|
| 21 |
+
from .models import ChessAction, ChessObservation, ChessState
|
| 22 |
+
|
| 23 |
+
|
| 24 |
+
class ChessEnv(EnvClient[ChessAction, ChessObservation, ChessState]):
|
| 25 |
+
"""
|
| 26 |
+
Client for Chess Environment.
|
| 27 |
+
|
| 28 |
+
This client maintains a persistent WebSocket connection to the environment
|
| 29 |
+
server, enabling efficient multi-step interactions with lower latency.
|
| 30 |
+
|
| 31 |
+
Uses the moonfish chess engine for opponent moves and position evaluation.
|
| 32 |
+
|
| 33 |
+
Example:
|
| 34 |
+
>>> with ChessEnv(base_url="http://localhost:8000") as client:
|
| 35 |
+
... result = client.reset()
|
| 36 |
+
... print(result.observation.fen)
|
| 37 |
+
... print(result.observation.legal_moves)
|
| 38 |
+
...
|
| 39 |
+
... result = client.step(ChessAction(move="e2e4"))
|
| 40 |
+
... print(result.reward, result.done)
|
| 41 |
+
"""
|
| 42 |
+
|
| 43 |
+
def _step_payload(self, action: ChessAction) -> Dict[str, Any]:
|
| 44 |
+
"""
|
| 45 |
+
Convert ChessAction to JSON payload for step request.
|
| 46 |
+
|
| 47 |
+
Args:
|
| 48 |
+
action: ChessAction instance with UCI move string.
|
| 49 |
+
|
| 50 |
+
Returns:
|
| 51 |
+
Dictionary representation suitable for JSON encoding.
|
| 52 |
+
"""
|
| 53 |
+
return {
|
| 54 |
+
"move": action.move,
|
| 55 |
+
}
|
| 56 |
+
|
| 57 |
+
def _parse_result(self, payload: Dict[str, Any]) -> StepResult[ChessObservation]:
|
| 58 |
+
"""
|
| 59 |
+
Parse server response into StepResult[ChessObservation].
|
| 60 |
+
|
| 61 |
+
Args:
|
| 62 |
+
payload: JSON response from server.
|
| 63 |
+
|
| 64 |
+
Returns:
|
| 65 |
+
StepResult with ChessObservation.
|
| 66 |
+
"""
|
| 67 |
+
obs_data = payload.get("observation", {})
|
| 68 |
+
# `serialize_observation` keeps reward and done on the envelope, not in
|
| 69 |
+
# the observation dict. Reading them from obs_data left every step
|
| 70 |
+
# reporting reward 0.0 and done False.
|
| 71 |
+
reward = payload.get("reward", 0.0)
|
| 72 |
+
done = payload.get("done", False)
|
| 73 |
+
|
| 74 |
+
observation = ChessObservation(
|
| 75 |
+
fen=obs_data.get("fen", ""),
|
| 76 |
+
legal_moves=obs_data.get("legal_moves", []),
|
| 77 |
+
is_check=obs_data.get("is_check", False),
|
| 78 |
+
done=bool(done),
|
| 79 |
+
reward=reward,
|
| 80 |
+
result=obs_data.get("result"),
|
| 81 |
+
metadata=payload.get("metadata", obs_data.get("metadata", {})),
|
| 82 |
+
)
|
| 83 |
+
|
| 84 |
+
return StepResult(
|
| 85 |
+
observation=observation,
|
| 86 |
+
reward=observation.reward,
|
| 87 |
+
done=observation.done,
|
| 88 |
+
metadata=payload.get("metadata"),
|
| 89 |
+
)
|
| 90 |
+
|
| 91 |
+
def _parse_state(self, payload: Dict[str, Any]) -> ChessState:
|
| 92 |
+
"""
|
| 93 |
+
Parse server response into ChessState object.
|
| 94 |
+
|
| 95 |
+
Args:
|
| 96 |
+
payload: JSON response from /state endpoint.
|
| 97 |
+
|
| 98 |
+
Returns:
|
| 99 |
+
ChessState object with environment state information.
|
| 100 |
+
"""
|
| 101 |
+
return ChessState(
|
| 102 |
+
episode_id=payload.get("episode_id", ""),
|
| 103 |
+
fen=payload.get("fen", ""),
|
| 104 |
+
current_player=payload.get("current_player", "white"),
|
| 105 |
+
move_history=payload.get("move_history", []),
|
| 106 |
+
step_count=payload.get("step_count", 0),
|
| 107 |
+
)
|
discovery.json
ADDED
|
@@ -0,0 +1,9 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
{
|
| 2 |
+
"description": "Play chess with legal UCI moves against a Moonfish engine, a random opponent, or in self-play. Observations include the FEN board position, legal moves, and game outcome.",
|
| 3 |
+
"tags": ["openenv", "rl-environment", "chess", "strategy", "board-game", "self-play"],
|
| 4 |
+
"representative_queries": [
|
| 5 |
+
"find a chess environment with a configurable opponent",
|
| 6 |
+
"evaluate move selection from a FEN board position",
|
| 7 |
+
"run chess self-play through an OpenEnv client"
|
| 8 |
+
]
|
| 9 |
+
}
|
envs/chess_env/README.md
ADDED
|
@@ -0,0 +1,144 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
---
|
| 2 |
+
title: Chess Environment Server
|
| 3 |
+
emoji: ♟️
|
| 4 |
+
colorFrom: gray
|
| 5 |
+
colorTo: green
|
| 6 |
+
sdk: docker
|
| 7 |
+
pinned: false
|
| 8 |
+
app_port: 8000
|
| 9 |
+
base_path: /web
|
| 10 |
+
tags:
|
| 11 |
+
- openenv
|
| 12 |
+
---
|
| 13 |
+
|
| 14 |
+
# Chess Environment
|
| 15 |
+
|
| 16 |
+
A chess reinforcement learning environment for OpenEnv, powered by the [moonfish](https://github.com/luccabb/moonfish) chess engine.
|
| 17 |
+
|
| 18 |
+

|
| 19 |
+
|
| 20 |
+
*The `/web` playground after a few moves: the board and one button per legal move.*
|
| 21 |
+
|
| 22 |
+
## Features
|
| 23 |
+
|
| 24 |
+
- **Full chess rules** via python-chess library
|
| 25 |
+
- **Configurable opponent**: moonfish engine, random moves, or self-play
|
| 26 |
+
- **Position evaluation**: Uses moonfish's PSQT-based evaluation
|
| 27 |
+
- **Standard OpenEnv interface**: reset(), step(), state
|
| 28 |
+
|
| 29 |
+
## Quick Start
|
| 30 |
+
|
| 31 |
+
### Using Docker
|
| 32 |
+
|
| 33 |
+
```bash
|
| 34 |
+
# Build the image
|
| 35 |
+
docker build -t chess-env:latest -f envs/chess_env/server/Dockerfile .
|
| 36 |
+
|
| 37 |
+
# Run the server
|
| 38 |
+
docker run -p 8000:8000 chess-env:latest
|
| 39 |
+
```
|
| 40 |
+
|
| 41 |
+
### Using the Client
|
| 42 |
+
|
| 43 |
+
The client is **async by default**:
|
| 44 |
+
|
| 45 |
+
```python
|
| 46 |
+
import asyncio
|
| 47 |
+
from chess_env import ChessEnv, ChessAction
|
| 48 |
+
|
| 49 |
+
async def main():
|
| 50 |
+
async with ChessEnv(base_url="http://localhost:8000") as env:
|
| 51 |
+
# Reset for a new game
|
| 52 |
+
result = await env.reset()
|
| 53 |
+
print(f"Starting position: {result.observation.fen}")
|
| 54 |
+
print(f"Legal moves: {result.observation.legal_moves}")
|
| 55 |
+
|
| 56 |
+
# Make a move
|
| 57 |
+
result = await env.step(ChessAction(move="e2e4"))
|
| 58 |
+
print(f"Reward: {result.reward}, Done: {result.done}")
|
| 59 |
+
|
| 60 |
+
# Play until game ends
|
| 61 |
+
while not result.done:
|
| 62 |
+
move = result.observation.legal_moves[0]
|
| 63 |
+
result = await env.step(ChessAction(move=move))
|
| 64 |
+
|
| 65 |
+
print(f"Game result: {result.observation.result}")
|
| 66 |
+
|
| 67 |
+
asyncio.run(main())
|
| 68 |
+
```
|
| 69 |
+
|
| 70 |
+
For **synchronous usage**, use the `.sync()` wrapper:
|
| 71 |
+
|
| 72 |
+
```python
|
| 73 |
+
from chess_env import ChessEnv, ChessAction
|
| 74 |
+
|
| 75 |
+
with ChessEnv(base_url="http://localhost:8000").sync() as env:
|
| 76 |
+
result = env.reset()
|
| 77 |
+
result = env.step(ChessAction(move="e2e4"))
|
| 78 |
+
print(f"Reward: {result.reward}")
|
| 79 |
+
```
|
| 80 |
+
|
| 81 |
+
## Observation Space
|
| 82 |
+
|
| 83 |
+
| Field | Type | Description |
|
| 84 |
+
|-------|------|-------------|
|
| 85 |
+
| `fen` | str | Board position in FEN notation |
|
| 86 |
+
| `legal_moves` | List[str] | Legal moves in UCI format |
|
| 87 |
+
| `is_check` | bool | Whether current player is in check |
|
| 88 |
+
| `done` | bool | Whether game has ended |
|
| 89 |
+
| `reward` | float | Reward for last action |
|
| 90 |
+
| `result` | str | Game result ("1-0", "0-1", "1/2-1/2") |
|
| 91 |
+
|
| 92 |
+
## Action Space
|
| 93 |
+
|
| 94 |
+
| Field | Type | Description |
|
| 95 |
+
|-------|------|-------------|
|
| 96 |
+
| `move` | str | UCI format move (e.g., "e2e4", "e7e8q") |
|
| 97 |
+
|
| 98 |
+
## Rewards
|
| 99 |
+
|
| 100 |
+
| Outcome | Reward |
|
| 101 |
+
|---------|--------|
|
| 102 |
+
| Win | +1.0 |
|
| 103 |
+
| Loss | -1.0 |
|
| 104 |
+
| Draw | 0.0 |
|
| 105 |
+
| Illegal move | -0.1 |
|
| 106 |
+
|
| 107 |
+
## Configuration
|
| 108 |
+
|
| 109 |
+
The environment supports these configuration options:
|
| 110 |
+
|
| 111 |
+
| Parameter | Default | Description |
|
| 112 |
+
|-----------|---------|-------------|
|
| 113 |
+
| `opponent` | "moonfish" | Opponent type: "moonfish", "random", or None |
|
| 114 |
+
| `opponent_depth` | 2 | Search depth for moonfish opponent |
|
| 115 |
+
| `max_moves` | 500 | Maximum half-moves before draw |
|
| 116 |
+
| `agent_color` | None | Agent color: "white", "black", or None (alternate each episode) |
|
| 117 |
+
| `gamma` | 0.99 | Discount factor for temporal credit assignment |
|
| 118 |
+
|
| 119 |
+
## Temporal Discounting
|
| 120 |
+
|
| 121 |
+
For RL training, the environment computes temporally discounted rewards at episode end. This helps with credit assignment in long games where only the final outcome is known.
|
| 122 |
+
|
| 123 |
+
When an episode ends, the terminal observation's `metadata` includes:
|
| 124 |
+
- `discounted_rewards`: List of discounted rewards for each agent move
|
| 125 |
+
- `gamma`: The discount factor used
|
| 126 |
+
|
| 127 |
+
The formula is `r_t = γ^(T-1-t) × R_final` where:
|
| 128 |
+
- `T` = total agent moves
|
| 129 |
+
- `t` = move index (0-indexed)
|
| 130 |
+
- `R_final` = terminal reward (+1, -1, or 0)
|
| 131 |
+
|
| 132 |
+
Example for a 5-move win with γ=0.99:
|
| 133 |
+
```
|
| 134 |
+
Move 0: 0.99^4 × 1.0 = 0.961
|
| 135 |
+
Move 1: 0.99^3 × 1.0 = 0.970
|
| 136 |
+
Move 2: 0.99^2 × 1.0 = 0.980
|
| 137 |
+
Move 3: 0.99^1 × 1.0 = 0.990
|
| 138 |
+
Move 4: 0.99^0 × 1.0 = 1.000
|
| 139 |
+
```
|
| 140 |
+
|
| 141 |
+
## Links
|
| 142 |
+
|
| 143 |
+
- [moonfish GitHub](https://github.com/luccabb/moonfish)
|
| 144 |
+
- [Play online](https://huggingface.co/spaces/luccabb/moonfish_chess)
|
envs/chess_env/__init__.py
ADDED
|
@@ -0,0 +1,34 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 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 the
|
| 5 |
+
# LICENSE file in the root directory of this source tree.
|
| 6 |
+
|
| 7 |
+
"""
|
| 8 |
+
Chess Environment for OpenEnv.
|
| 9 |
+
|
| 10 |
+
This module provides OpenEnv integration for chess, using the moonfish
|
| 11 |
+
chess engine for position evaluation and opponent play.
|
| 12 |
+
|
| 13 |
+
Example:
|
| 14 |
+
>>> from envs.chess_env import ChessEnv, ChessAction
|
| 15 |
+
>>>
|
| 16 |
+
>>> # Connect to a running server or start via Docker
|
| 17 |
+
>>> env = ChessEnv.from_docker_image("chess-env:latest")
|
| 18 |
+
>>>
|
| 19 |
+
>>> # Reset and interact
|
| 20 |
+
>>> result = env.reset()
|
| 21 |
+
>>> print(result.observation.fen)
|
| 22 |
+
>>> print(result.observation.legal_moves)
|
| 23 |
+
>>>
|
| 24 |
+
>>> result = env.step(ChessAction(move="e2e4"))
|
| 25 |
+
>>> print(result.reward, result.done)
|
| 26 |
+
>>>
|
| 27 |
+
>>> # Cleanup
|
| 28 |
+
>>> env.close()
|
| 29 |
+
"""
|
| 30 |
+
|
| 31 |
+
from .client import ChessEnv
|
| 32 |
+
from .models import ChessAction, ChessObservation, ChessState
|
| 33 |
+
|
| 34 |
+
__all__ = ["ChessEnv", "ChessAction", "ChessObservation", "ChessState"]
|
envs/chess_env/client.py
ADDED
|
@@ -0,0 +1,107 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 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 the
|
| 5 |
+
# LICENSE file in the root directory of this source tree.
|
| 6 |
+
|
| 7 |
+
"""
|
| 8 |
+
Chess Environment Client.
|
| 9 |
+
|
| 10 |
+
This module provides the client for connecting to a Chess Environment server
|
| 11 |
+
via WebSocket for persistent sessions.
|
| 12 |
+
"""
|
| 13 |
+
|
| 14 |
+
from __future__ import annotations
|
| 15 |
+
|
| 16 |
+
from typing import Any, Dict
|
| 17 |
+
|
| 18 |
+
from openenv.core.client_types import StepResult
|
| 19 |
+
from openenv.core.env_client import EnvClient
|
| 20 |
+
|
| 21 |
+
from .models import ChessAction, ChessObservation, ChessState
|
| 22 |
+
|
| 23 |
+
|
| 24 |
+
class ChessEnv(EnvClient[ChessAction, ChessObservation, ChessState]):
|
| 25 |
+
"""
|
| 26 |
+
Client for Chess Environment.
|
| 27 |
+
|
| 28 |
+
This client maintains a persistent WebSocket connection to the environment
|
| 29 |
+
server, enabling efficient multi-step interactions with lower latency.
|
| 30 |
+
|
| 31 |
+
Uses the moonfish chess engine for opponent moves and position evaluation.
|
| 32 |
+
|
| 33 |
+
Example:
|
| 34 |
+
>>> with ChessEnv(base_url="http://localhost:8000") as client:
|
| 35 |
+
... result = client.reset()
|
| 36 |
+
... print(result.observation.fen)
|
| 37 |
+
... print(result.observation.legal_moves)
|
| 38 |
+
...
|
| 39 |
+
... result = client.step(ChessAction(move="e2e4"))
|
| 40 |
+
... print(result.reward, result.done)
|
| 41 |
+
"""
|
| 42 |
+
|
| 43 |
+
def _step_payload(self, action: ChessAction) -> Dict[str, Any]:
|
| 44 |
+
"""
|
| 45 |
+
Convert ChessAction to JSON payload for step request.
|
| 46 |
+
|
| 47 |
+
Args:
|
| 48 |
+
action: ChessAction instance with UCI move string.
|
| 49 |
+
|
| 50 |
+
Returns:
|
| 51 |
+
Dictionary representation suitable for JSON encoding.
|
| 52 |
+
"""
|
| 53 |
+
return {
|
| 54 |
+
"move": action.move,
|
| 55 |
+
}
|
| 56 |
+
|
| 57 |
+
def _parse_result(self, payload: Dict[str, Any]) -> StepResult[ChessObservation]:
|
| 58 |
+
"""
|
| 59 |
+
Parse server response into StepResult[ChessObservation].
|
| 60 |
+
|
| 61 |
+
Args:
|
| 62 |
+
payload: JSON response from server.
|
| 63 |
+
|
| 64 |
+
Returns:
|
| 65 |
+
StepResult with ChessObservation.
|
| 66 |
+
"""
|
| 67 |
+
obs_data = payload.get("observation", {})
|
| 68 |
+
# `serialize_observation` keeps reward and done on the envelope, not in
|
| 69 |
+
# the observation dict. Reading them from obs_data left every step
|
| 70 |
+
# reporting reward 0.0 and done False.
|
| 71 |
+
reward = payload.get("reward", 0.0)
|
| 72 |
+
done = payload.get("done", False)
|
| 73 |
+
|
| 74 |
+
observation = ChessObservation(
|
| 75 |
+
fen=obs_data.get("fen", ""),
|
| 76 |
+
legal_moves=obs_data.get("legal_moves", []),
|
| 77 |
+
is_check=obs_data.get("is_check", False),
|
| 78 |
+
done=bool(done),
|
| 79 |
+
reward=reward,
|
| 80 |
+
result=obs_data.get("result"),
|
| 81 |
+
metadata=payload.get("metadata", obs_data.get("metadata", {})),
|
| 82 |
+
)
|
| 83 |
+
|
| 84 |
+
return StepResult(
|
| 85 |
+
observation=observation,
|
| 86 |
+
reward=observation.reward,
|
| 87 |
+
done=observation.done,
|
| 88 |
+
metadata=payload.get("metadata"),
|
| 89 |
+
)
|
| 90 |
+
|
| 91 |
+
def _parse_state(self, payload: Dict[str, Any]) -> ChessState:
|
| 92 |
+
"""
|
| 93 |
+
Parse server response into ChessState object.
|
| 94 |
+
|
| 95 |
+
Args:
|
| 96 |
+
payload: JSON response from /state endpoint.
|
| 97 |
+
|
| 98 |
+
Returns:
|
| 99 |
+
ChessState object with environment state information.
|
| 100 |
+
"""
|
| 101 |
+
return ChessState(
|
| 102 |
+
episode_id=payload.get("episode_id", ""),
|
| 103 |
+
fen=payload.get("fen", ""),
|
| 104 |
+
current_player=payload.get("current_player", "white"),
|
| 105 |
+
move_history=payload.get("move_history", []),
|
| 106 |
+
step_count=payload.get("step_count", 0),
|
| 107 |
+
)
|
envs/chess_env/discovery.json
ADDED
|
@@ -0,0 +1,9 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
{
|
| 2 |
+
"description": "Play chess with legal UCI moves against a Moonfish engine, a random opponent, or in self-play. Observations include the FEN board position, legal moves, and game outcome.",
|
| 3 |
+
"tags": ["openenv", "rl-environment", "chess", "strategy", "board-game", "self-play"],
|
| 4 |
+
"representative_queries": [
|
| 5 |
+
"find a chess environment with a configurable opponent",
|
| 6 |
+
"evaluate move selection from a FEN board position",
|
| 7 |
+
"run chess self-play through an OpenEnv client"
|
| 8 |
+
]
|
| 9 |
+
}
|
envs/chess_env/models.py
ADDED
|
@@ -0,0 +1,68 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 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 the
|
| 5 |
+
# LICENSE file in the root directory of this source tree.
|
| 6 |
+
|
| 7 |
+
"""
|
| 8 |
+
Data models for Chess Environment.
|
| 9 |
+
|
| 10 |
+
This module defines the Action, Observation, and State types for chess games
|
| 11 |
+
via the OpenEnv interface. Uses the moonfish chess engine for move search
|
| 12 |
+
and position evaluation.
|
| 13 |
+
"""
|
| 14 |
+
|
| 15 |
+
from __future__ import annotations
|
| 16 |
+
|
| 17 |
+
from typing import List, Optional
|
| 18 |
+
|
| 19 |
+
from openenv.core.env_server import Action, Observation, State
|
| 20 |
+
from pydantic import Field
|
| 21 |
+
|
| 22 |
+
|
| 23 |
+
class ChessAction(Action):
|
| 24 |
+
"""
|
| 25 |
+
Action for Chess environment.
|
| 26 |
+
|
| 27 |
+
Attributes:
|
| 28 |
+
move: UCI format move string (e.g., "e2e4", "e7e8q" for promotion).
|
| 29 |
+
"""
|
| 30 |
+
|
| 31 |
+
move: str
|
| 32 |
+
|
| 33 |
+
|
| 34 |
+
class ChessObservation(Observation):
|
| 35 |
+
"""
|
| 36 |
+
Observation for Chess environment.
|
| 37 |
+
|
| 38 |
+
Attributes:
|
| 39 |
+
fen: Board position in FEN notation.
|
| 40 |
+
legal_moves: List of legal moves in UCI format.
|
| 41 |
+
is_check: Whether the current player is in check.
|
| 42 |
+
done: Whether the game is over.
|
| 43 |
+
reward: Reward for the last action.
|
| 44 |
+
result: Game result string if game is over (e.g., "1-0", "0-1", "1/2-1/2").
|
| 45 |
+
"""
|
| 46 |
+
|
| 47 |
+
fen: str = ""
|
| 48 |
+
legal_moves: List[str] = Field(default_factory=list)
|
| 49 |
+
is_check: bool = False
|
| 50 |
+
result: Optional[str] = None
|
| 51 |
+
|
| 52 |
+
|
| 53 |
+
class ChessState(State):
|
| 54 |
+
"""
|
| 55 |
+
State for Chess environment.
|
| 56 |
+
|
| 57 |
+
Attributes:
|
| 58 |
+
episode_id: Unique ID for the current game.
|
| 59 |
+
fen: Current board position in FEN notation.
|
| 60 |
+
current_player: "white" or "black".
|
| 61 |
+
move_history: List of moves played in UCI format.
|
| 62 |
+
step_count: Number of half-moves played.
|
| 63 |
+
"""
|
| 64 |
+
|
| 65 |
+
fen: str = "rnbqkbnr/pppppppp/8/8/8/8/PPPPPPPP/RNBQKBNR w KQkq - 0 1"
|
| 66 |
+
current_player: str = "white"
|
| 67 |
+
move_history: List[str] = Field(default_factory=list)
|
| 68 |
+
step_count: int = 0
|
envs/chess_env/openenv.yaml
ADDED
|
@@ -0,0 +1,6 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
spec_version: 1
|
| 2 |
+
name: chess_env
|
| 3 |
+
type: space
|
| 4 |
+
runtime: fastapi
|
| 5 |
+
app: chess_env.server.app:app
|
| 6 |
+
port: 8000
|
envs/chess_env/pyproject.toml
ADDED
|
@@ -0,0 +1,31 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
[build-system]
|
| 2 |
+
requires = ["setuptools>=45", "wheel"]
|
| 3 |
+
build-backend = "setuptools.build_meta"
|
| 4 |
+
|
| 5 |
+
[project]
|
| 6 |
+
name = "openenv-chess-env"
|
| 7 |
+
version = "0.1.0"
|
| 8 |
+
description = "Chess Environment for OpenEnv"
|
| 9 |
+
requires-python = ">=3.10"
|
| 10 |
+
dependencies = [
|
| 11 |
+
"openenv @ git+https://github.com/huggingface/OpenEnv.git@v0.10.0",
|
| 12 |
+
"fastapi>=0.115.0",
|
| 13 |
+
"pydantic>=2.0.0",
|
| 14 |
+
"uvicorn[standard]>=0.24.0",
|
| 15 |
+
"requests>=2.31.0",
|
| 16 |
+
"python-chess>=1.999",
|
| 17 |
+
"moonfish>=0.0.1",
|
| 18 |
+
]
|
| 19 |
+
|
| 20 |
+
[project.optional-dependencies]
|
| 21 |
+
dev = [
|
| 22 |
+
"pytest>=8.0.0",
|
| 23 |
+
"pytest-cov>=4.0.0",
|
| 24 |
+
]
|
| 25 |
+
|
| 26 |
+
[project.scripts]
|
| 27 |
+
server = "chess_env.server.app:main"
|
| 28 |
+
|
| 29 |
+
[tool.setuptools]
|
| 30 |
+
packages = ["chess_env", "chess_env.server"]
|
| 31 |
+
package-dir = { "chess_env" = ".", "chess_env.server" = "server" }
|
envs/chess_env/scripts/verify.sh
ADDED
|
@@ -0,0 +1,60 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
#!/bin/bash
|
| 2 |
+
# Verification script for chess_env
|
| 3 |
+
# Run from repository root: bash envs/chess_env/scripts/verify.sh
|
| 4 |
+
|
| 5 |
+
set -e
|
| 6 |
+
|
| 7 |
+
echo "=== Chess Environment Verification ==="
|
| 8 |
+
echo
|
| 9 |
+
|
| 10 |
+
# Colors for output
|
| 11 |
+
GREEN='\033[0;32m'
|
| 12 |
+
RED='\033[0;31m'
|
| 13 |
+
NC='\033[0m' # No Color
|
| 14 |
+
|
| 15 |
+
check_pass() {
|
| 16 |
+
echo -e "${GREEN}✓ $1${NC}"
|
| 17 |
+
}
|
| 18 |
+
|
| 19 |
+
check_fail() {
|
| 20 |
+
echo -e "${RED}✗ $1${NC}"
|
| 21 |
+
exit 1
|
| 22 |
+
}
|
| 23 |
+
|
| 24 |
+
# 1. Run unit tests
|
| 25 |
+
echo "1. Running unit tests..."
|
| 26 |
+
if PYTHONPATH=src:. uv run pytest tests/envs/test_chess_environment.py -v; then
|
| 27 |
+
check_pass "Unit tests passed"
|
| 28 |
+
else
|
| 29 |
+
check_fail "Unit tests failed"
|
| 30 |
+
fi
|
| 31 |
+
echo
|
| 32 |
+
|
| 33 |
+
# 2. Run linting
|
| 34 |
+
echo "2. Running linting..."
|
| 35 |
+
if uv run ruff check envs/chess_env/ tests/envs/test_chess_environment.py; then
|
| 36 |
+
check_pass "Linting passed"
|
| 37 |
+
else
|
| 38 |
+
check_fail "Linting failed"
|
| 39 |
+
fi
|
| 40 |
+
echo
|
| 41 |
+
|
| 42 |
+
# 3. Check formatting
|
| 43 |
+
echo "3. Checking formatting..."
|
| 44 |
+
if uv run ruff format envs/chess_env/ tests/envs/test_chess_environment.py --check; then
|
| 45 |
+
check_pass "Formatting passed"
|
| 46 |
+
else
|
| 47 |
+
check_fail "Formatting failed"
|
| 48 |
+
fi
|
| 49 |
+
echo
|
| 50 |
+
|
| 51 |
+
# 4. Verify imports
|
| 52 |
+
echo "4. Verifying module imports..."
|
| 53 |
+
if PYTHONPATH=src:. uv run python -c "from envs.chess_env import ChessEnv, ChessAction, ChessObservation, ChessState; print(' Imports: ChessEnv, ChessAction, ChessObservation, ChessState')"; then
|
| 54 |
+
check_pass "Module imports successful"
|
| 55 |
+
else
|
| 56 |
+
check_fail "Module imports failed"
|
| 57 |
+
fi
|
| 58 |
+
echo
|
| 59 |
+
|
| 60 |
+
echo "=== All checks passed! ==="
|
envs/chess_env/server/Dockerfile
ADDED
|
@@ -0,0 +1,79 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 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 the
|
| 5 |
+
# LICENSE file in the root directory of this source tree.
|
| 6 |
+
|
| 7 |
+
# Multi-stage build using openenv-base
|
| 8 |
+
# This Dockerfile is flexible and works for both:
|
| 9 |
+
# - In-repo environments (with local src/core)
|
| 10 |
+
# - Standalone environments (with openenv from pip)
|
| 11 |
+
# The build script (openenv build) handles context detection and sets appropriate build args.
|
| 12 |
+
|
| 13 |
+
ARG BASE_IMAGE=ghcr.io/huggingface/openenv-base:latest
|
| 14 |
+
FROM ${BASE_IMAGE} AS builder
|
| 15 |
+
|
| 16 |
+
WORKDIR /app
|
| 17 |
+
|
| 18 |
+
# Build argument to control whether we're building standalone or in-repo
|
| 19 |
+
ARG BUILD_MODE=in-repo
|
| 20 |
+
|
| 21 |
+
# Copy environment code (always at root of build context)
|
| 22 |
+
COPY . /app/env
|
| 23 |
+
|
| 24 |
+
# For in-repo builds, openenv is already in the pyproject.toml dependencies
|
| 25 |
+
# For standalone builds, openenv will be installed from pip via pyproject.toml
|
| 26 |
+
WORKDIR /app/env
|
| 27 |
+
|
| 28 |
+
# Ensure uv is available (for local builds where base image lacks it)
|
| 29 |
+
RUN if ! command -v uv >/dev/null 2>&1; then \
|
| 30 |
+
curl -LsSf https://astral.sh/uv/install.sh | sh && \
|
| 31 |
+
mv /root/.local/bin/uv /usr/local/bin/uv && \
|
| 32 |
+
mv /root/.local/bin/uvx /usr/local/bin/uvx; \
|
| 33 |
+
fi
|
| 34 |
+
|
| 35 |
+
# Install git for building from git repos (build-time only)
|
| 36 |
+
RUN apt-get update && apt-get install -y --no-install-recommends \
|
| 37 |
+
git \
|
| 38 |
+
&& rm -rf /var/lib/apt/lists/*
|
| 39 |
+
|
| 40 |
+
# Install dependencies using uv sync
|
| 41 |
+
# First pass: install dependencies without the project (for better caching)
|
| 42 |
+
# Second pass: install the project itself
|
| 43 |
+
RUN --mount=type=cache,target=/root/.cache/uv \
|
| 44 |
+
if [ -f uv.lock ]; then \
|
| 45 |
+
uv sync --frozen --no-install-project --no-editable; \
|
| 46 |
+
else \
|
| 47 |
+
uv sync --no-install-project --no-editable; \
|
| 48 |
+
fi
|
| 49 |
+
|
| 50 |
+
RUN --mount=type=cache,target=/root/.cache/uv \
|
| 51 |
+
if [ -f uv.lock ]; then \
|
| 52 |
+
uv sync --frozen --no-editable; \
|
| 53 |
+
else \
|
| 54 |
+
uv sync --no-editable; \
|
| 55 |
+
fi
|
| 56 |
+
|
| 57 |
+
# Final runtime stage
|
| 58 |
+
FROM ${BASE_IMAGE}
|
| 59 |
+
|
| 60 |
+
WORKDIR /app
|
| 61 |
+
|
| 62 |
+
# Copy the virtual environment from builder
|
| 63 |
+
COPY --from=builder /app/env/.venv /app/.venv
|
| 64 |
+
|
| 65 |
+
# Copy the environment code
|
| 66 |
+
COPY --from=builder /app/env /app/env
|
| 67 |
+
|
| 68 |
+
# Set PATH to use the virtual environment
|
| 69 |
+
ENV PATH="/app/.venv/bin:$PATH"
|
| 70 |
+
|
| 71 |
+
# Set PYTHONPATH so imports work correctly
|
| 72 |
+
ENV PYTHONPATH="/app/env:$PYTHONPATH"
|
| 73 |
+
|
| 74 |
+
# Health check using Python (more portable than curl/wget)
|
| 75 |
+
HEALTHCHECK --interval=30s --timeout=3s --start-period=5s --retries=3 \
|
| 76 |
+
CMD python -c "import urllib.request; urllib.request.urlopen('http://localhost:8000/health')" || exit 1
|
| 77 |
+
|
| 78 |
+
# Run the FastAPI server from the installed package.
|
| 79 |
+
CMD ["uvicorn", "chess_env.server.app:app", "--host", "0.0.0.0", "--port", "8000"]
|
envs/chess_env/server/__init__.py
ADDED
|
@@ -0,0 +1,11 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 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 the
|
| 5 |
+
# LICENSE file in the root directory of this source tree.
|
| 6 |
+
|
| 7 |
+
"""Chess Environment server module."""
|
| 8 |
+
|
| 9 |
+
from .chess_environment import ChessEnvironment
|
| 10 |
+
|
| 11 |
+
__all__ = ["ChessEnvironment"]
|
envs/chess_env/server/app.py
ADDED
|
@@ -0,0 +1,35 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 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 the
|
| 5 |
+
# LICENSE file in the root directory of this source tree.
|
| 6 |
+
|
| 7 |
+
"""FastAPI application for the Chess Environment."""
|
| 8 |
+
|
| 9 |
+
import os
|
| 10 |
+
|
| 11 |
+
from openenv.core.env_server import create_app
|
| 12 |
+
|
| 13 |
+
from ..models import ChessAction, ChessObservation, ChessState
|
| 14 |
+
from .chess_environment import ChessEnvironment
|
| 15 |
+
|
| 16 |
+
# Create the FastAPI app
|
| 17 |
+
# Pass the class (factory) instead of an instance for WebSocket session support
|
| 18 |
+
app = create_app(
|
| 19 |
+
ChessEnvironment,
|
| 20 |
+
ChessAction,
|
| 21 |
+
ChessObservation,
|
| 22 |
+
env_name="chess_env",
|
| 23 |
+
state_cls=ChessState,
|
| 24 |
+
max_concurrent_envs=int(os.getenv("MAX_CONCURRENT_ENVS", "8")),
|
| 25 |
+
)
|
| 26 |
+
|
| 27 |
+
|
| 28 |
+
def main():
|
| 29 |
+
import uvicorn
|
| 30 |
+
|
| 31 |
+
uvicorn.run(app, host="0.0.0.0", port=8000)
|
| 32 |
+
|
| 33 |
+
|
| 34 |
+
if __name__ == "__main__":
|
| 35 |
+
main()
|
envs/chess_env/server/chess_environment.py
ADDED
|
@@ -0,0 +1,337 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 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 the
|
| 5 |
+
# LICENSE file in the root directory of this source tree.
|
| 6 |
+
|
| 7 |
+
"""
|
| 8 |
+
Chess Environment for OpenEnv using moonfish.
|
| 9 |
+
|
| 10 |
+
This module provides a chess environment for reinforcement learning,
|
| 11 |
+
using python-chess for game logic and moonfish for position evaluation
|
| 12 |
+
and opponent play.
|
| 13 |
+
"""
|
| 14 |
+
|
| 15 |
+
import random
|
| 16 |
+
import threading
|
| 17 |
+
import uuid
|
| 18 |
+
from typing import Any, Dict, List, Optional, Tuple
|
| 19 |
+
|
| 20 |
+
import chess
|
| 21 |
+
from moonfish.lib import search_move
|
| 22 |
+
from moonfish.psqt import board_evaluation, BOARD_EVALUATION_CACHE
|
| 23 |
+
from openenv.core.env_server import Environment
|
| 24 |
+
|
| 25 |
+
from ..models import ChessAction, ChessObservation, ChessState
|
| 26 |
+
from .rubrics import ChessWinLossRubric
|
| 27 |
+
|
| 28 |
+
# moonfish keeps one evaluation cache per process, so sessions take turns with it.
|
| 29 |
+
_MOONFISH_LOCK = threading.Lock()
|
| 30 |
+
|
| 31 |
+
|
| 32 |
+
class ChessEnvironment(Environment):
|
| 33 |
+
"""
|
| 34 |
+
Chess environment implementing the OpenEnv interface.
|
| 35 |
+
|
| 36 |
+
Uses python-chess for game logic and moonfish for position evaluation.
|
| 37 |
+
Designed for RL training where an agent plays as one color against
|
| 38 |
+
an opponent (which can be random, moonfish engine, or self-play).
|
| 39 |
+
"""
|
| 40 |
+
|
| 41 |
+
SUPPORTS_CONCURRENT_SESSIONS = True
|
| 42 |
+
|
| 43 |
+
def __init__(
|
| 44 |
+
self,
|
| 45 |
+
opponent: str = "moonfish",
|
| 46 |
+
opponent_depth: int = 2,
|
| 47 |
+
max_moves: int = 500,
|
| 48 |
+
agent_color: str | None = None,
|
| 49 |
+
gamma: float = 0.99,
|
| 50 |
+
):
|
| 51 |
+
"""
|
| 52 |
+
Initialize the chess environment.
|
| 53 |
+
|
| 54 |
+
Args:
|
| 55 |
+
opponent: Opponent type - "moonfish", "random", or None (self-play)
|
| 56 |
+
opponent_depth: Search depth when using moonfish as opponent
|
| 57 |
+
max_moves: Maximum half-moves before draw
|
| 58 |
+
agent_color: Which color the agent plays - "white", "black", or None (alternate each episode)
|
| 59 |
+
gamma: Discount factor for temporal credit assignment (0-1)
|
| 60 |
+
"""
|
| 61 |
+
super().__init__(rubric=ChessWinLossRubric(gamma=gamma))
|
| 62 |
+
self._opponent = opponent
|
| 63 |
+
self._opponent_depth = opponent_depth
|
| 64 |
+
self._max_moves = max_moves
|
| 65 |
+
self._agent_color_setting = agent_color
|
| 66 |
+
self._gamma = gamma
|
| 67 |
+
self._board = None
|
| 68 |
+
self._state = None
|
| 69 |
+
self._agent_color = chess.WHITE
|
| 70 |
+
self._agent_move_count = 0 # Track agent moves for discounting
|
| 71 |
+
self.reset()
|
| 72 |
+
|
| 73 |
+
def reset(self, fen: str = None):
|
| 74 |
+
"""
|
| 75 |
+
Initialize a new chess game episode.
|
| 76 |
+
|
| 77 |
+
Args:
|
| 78 |
+
fen: Optional starting position in FEN notation.
|
| 79 |
+
|
| 80 |
+
Returns:
|
| 81 |
+
Initial observation of the board state.
|
| 82 |
+
"""
|
| 83 |
+
self._reset_rubric()
|
| 84 |
+
# moonfish's module-level evaluation cache is never trimmed.
|
| 85 |
+
with _MOONFISH_LOCK:
|
| 86 |
+
BOARD_EVALUATION_CACHE.clear()
|
| 87 |
+
|
| 88 |
+
if fen:
|
| 89 |
+
self._board = chess.Board(fen)
|
| 90 |
+
else:
|
| 91 |
+
self._board = chess.Board()
|
| 92 |
+
|
| 93 |
+
# Generate episode ID once
|
| 94 |
+
episode_id = str(uuid.uuid4())
|
| 95 |
+
|
| 96 |
+
# Determine agent color
|
| 97 |
+
if self._agent_color_setting == "white":
|
| 98 |
+
self._agent_color = chess.WHITE
|
| 99 |
+
elif self._agent_color_setting == "black":
|
| 100 |
+
self._agent_color = chess.BLACK
|
| 101 |
+
elif self._agent_color_setting is None:
|
| 102 |
+
# Alternate each episode based on episode_id hash
|
| 103 |
+
self._agent_color = (
|
| 104 |
+
chess.WHITE if hash(episode_id) % 2 == 0 else chess.BLACK
|
| 105 |
+
)
|
| 106 |
+
else:
|
| 107 |
+
self._agent_color = chess.WHITE
|
| 108 |
+
|
| 109 |
+
self._state = ChessState(
|
| 110 |
+
episode_id=episode_id,
|
| 111 |
+
step_count=0,
|
| 112 |
+
current_player="white" if self._board.turn else "black",
|
| 113 |
+
fen=self._board.fen(),
|
| 114 |
+
move_history=[],
|
| 115 |
+
)
|
| 116 |
+
self._agent_move_count = 0
|
| 117 |
+
|
| 118 |
+
# If agent plays Black and opponent is configured, opponent moves first
|
| 119 |
+
if self._opponent is not None and self._agent_color == chess.BLACK:
|
| 120 |
+
self._make_opponent_move()
|
| 121 |
+
|
| 122 |
+
# Check if starting position is already terminal (e.g., custom FEN with checkmate/stalemate)
|
| 123 |
+
_, done = self._calculate_reward_and_done()
|
| 124 |
+
return self._make_observation(done=done)
|
| 125 |
+
|
| 126 |
+
def step(self, action: ChessAction):
|
| 127 |
+
"""
|
| 128 |
+
Execute a chess move and return the resulting state.
|
| 129 |
+
|
| 130 |
+
Args:
|
| 131 |
+
action: The move to make in UCI format (e.g., "e2e4").
|
| 132 |
+
|
| 133 |
+
Returns:
|
| 134 |
+
Observation with reward and done flag.
|
| 135 |
+
"""
|
| 136 |
+
if self._board is None:
|
| 137 |
+
raise RuntimeError("Environment not initialized. Call reset() first.")
|
| 138 |
+
|
| 139 |
+
# Parse the move
|
| 140 |
+
try:
|
| 141 |
+
move = chess.Move.from_uci(action.move)
|
| 142 |
+
except ValueError:
|
| 143 |
+
return self._make_observation(reward=-0.1, done=False)
|
| 144 |
+
|
| 145 |
+
# Check if move is legal
|
| 146 |
+
if move not in self._board.legal_moves:
|
| 147 |
+
return self._make_observation(reward=-0.1, done=False)
|
| 148 |
+
|
| 149 |
+
# Execute the move
|
| 150 |
+
self._board.push(move)
|
| 151 |
+
self._state.step_count += 1
|
| 152 |
+
self._agent_move_count += 1
|
| 153 |
+
self._state.move_history.append(action.move)
|
| 154 |
+
self._state.current_player = "white" if self._board.turn else "black"
|
| 155 |
+
self._state.fen = self._board.fen()
|
| 156 |
+
|
| 157 |
+
# Calculate reward and check for game end
|
| 158 |
+
reward, done = self._calculate_reward_and_done()
|
| 159 |
+
|
| 160 |
+
# If game not over and opponent is configured, make opponent move
|
| 161 |
+
if not done and self._opponent is not None:
|
| 162 |
+
self._make_opponent_move()
|
| 163 |
+
reward, done = self._calculate_reward_and_done()
|
| 164 |
+
|
| 165 |
+
obs = self._make_observation(reward=reward, done=done)
|
| 166 |
+
|
| 167 |
+
# Feed to rubric for trajectory accumulation
|
| 168 |
+
self._apply_rubric(action, obs)
|
| 169 |
+
|
| 170 |
+
# Note: obs.reward is kept from _calculate_reward_and_done() rather than
|
| 171 |
+
# overridden by the rubric. The rubric here is a TrajectoryRubric that
|
| 172 |
+
# returns 0.0 until done then the final score — the inline reward logic
|
| 173 |
+
# already provides the correct per-step reward including penalties for
|
| 174 |
+
# invalid/illegal moves. The rubric's value is in compute_step_rewards()
|
| 175 |
+
# for training infrastructure, not in replacing obs.reward.
|
| 176 |
+
|
| 177 |
+
return obs
|
| 178 |
+
|
| 179 |
+
def _make_observation(self, reward: float = 0.0, done: bool = False):
|
| 180 |
+
"""Build observation from current board state."""
|
| 181 |
+
legal_moves = [move.uci() for move in self._board.legal_moves]
|
| 182 |
+
|
| 183 |
+
result = None
|
| 184 |
+
if done:
|
| 185 |
+
result = self._get_result_string()
|
| 186 |
+
|
| 187 |
+
with _MOONFISH_LOCK:
|
| 188 |
+
evaluation = board_evaluation(self._board)
|
| 189 |
+
metadata = {
|
| 190 |
+
"evaluation": evaluation,
|
| 191 |
+
"fullmove_number": self._board.fullmove_number,
|
| 192 |
+
"halfmove_clock": self._board.halfmove_clock,
|
| 193 |
+
}
|
| 194 |
+
|
| 195 |
+
# Compute discounted rewards for all agent moves when episode ends
|
| 196 |
+
if done and self._agent_move_count > 0:
|
| 197 |
+
discounted_rewards = self._compute_discounted_rewards(reward)
|
| 198 |
+
metadata["discounted_rewards"] = discounted_rewards
|
| 199 |
+
metadata["gamma"] = self._gamma
|
| 200 |
+
|
| 201 |
+
return ChessObservation(
|
| 202 |
+
fen=self._board.fen(),
|
| 203 |
+
legal_moves=legal_moves,
|
| 204 |
+
is_check=self._board.is_check(),
|
| 205 |
+
done=done,
|
| 206 |
+
reward=reward,
|
| 207 |
+
result=result,
|
| 208 |
+
metadata=metadata,
|
| 209 |
+
)
|
| 210 |
+
|
| 211 |
+
def _calculate_reward_and_done(self):
|
| 212 |
+
"""Calculate reward and check if episode is done."""
|
| 213 |
+
if self._board.is_checkmate():
|
| 214 |
+
winner = not self._board.turn
|
| 215 |
+
if winner == self._agent_color:
|
| 216 |
+
return 1.0, True
|
| 217 |
+
else:
|
| 218 |
+
return -1.0, True
|
| 219 |
+
|
| 220 |
+
if self._board.is_stalemate():
|
| 221 |
+
return 0.0, True
|
| 222 |
+
|
| 223 |
+
if self._board.is_insufficient_material():
|
| 224 |
+
return 0.0, True
|
| 225 |
+
|
| 226 |
+
if self._board.is_fifty_moves():
|
| 227 |
+
return 0.0, True
|
| 228 |
+
|
| 229 |
+
if self._board.is_repetition(3):
|
| 230 |
+
return 0.0, True
|
| 231 |
+
|
| 232 |
+
if self._state.step_count >= self._max_moves:
|
| 233 |
+
return 0.0, True
|
| 234 |
+
|
| 235 |
+
return 0.0, False
|
| 236 |
+
|
| 237 |
+
def _get_result_string(self):
|
| 238 |
+
"""Get the game result as a string."""
|
| 239 |
+
if self._board.is_checkmate():
|
| 240 |
+
return "1-0" if not self._board.turn else "0-1"
|
| 241 |
+
return "1/2-1/2"
|
| 242 |
+
|
| 243 |
+
def _compute_discounted_rewards(self, terminal_reward: float) -> list[float]:
|
| 244 |
+
"""
|
| 245 |
+
Compute temporally discounted rewards for all agent moves.
|
| 246 |
+
|
| 247 |
+
Uses exponential discounting: r_t = γ^(T-1-t) * R_final
|
| 248 |
+
where T is total agent moves, t is move index, R_final is terminal reward.
|
| 249 |
+
|
| 250 |
+
Earlier moves get less credit, later moves get more credit.
|
| 251 |
+
This helps with credit assignment in long games.
|
| 252 |
+
|
| 253 |
+
Args:
|
| 254 |
+
terminal_reward: The final reward (+1 win, -1 loss, 0 draw)
|
| 255 |
+
|
| 256 |
+
Returns:
|
| 257 |
+
List of discounted rewards, one per agent move
|
| 258 |
+
"""
|
| 259 |
+
T = self._agent_move_count
|
| 260 |
+
discounted_rewards = []
|
| 261 |
+
for t in range(T):
|
| 262 |
+
# γ^(T-1-t) means last move (t=T-1) gets γ^0 = 1.0
|
| 263 |
+
# First move (t=0) gets γ^(T-1)
|
| 264 |
+
discount = self._gamma ** (T - 1 - t)
|
| 265 |
+
discounted_rewards.append(discount * terminal_reward)
|
| 266 |
+
return discounted_rewards
|
| 267 |
+
|
| 268 |
+
def _make_opponent_move(self):
|
| 269 |
+
"""Make a move for the opponent using configured strategy."""
|
| 270 |
+
if not self._board.legal_moves:
|
| 271 |
+
return
|
| 272 |
+
|
| 273 |
+
if self._opponent == "moonfish":
|
| 274 |
+
with _MOONFISH_LOCK:
|
| 275 |
+
move = search_move(self._board, depth=self._opponent_depth)
|
| 276 |
+
elif self._opponent == "random":
|
| 277 |
+
move = random.choice(list(self._board.legal_moves))
|
| 278 |
+
else:
|
| 279 |
+
return
|
| 280 |
+
|
| 281 |
+
self._board.push(move)
|
| 282 |
+
self._state.step_count += 1
|
| 283 |
+
self._state.move_history.append(move.uci())
|
| 284 |
+
self._state.current_player = "white" if self._board.turn else "black"
|
| 285 |
+
self._state.fen = self._board.fen()
|
| 286 |
+
|
| 287 |
+
def web_actions(
|
| 288 |
+
self, observation: Dict[str, Any]
|
| 289 |
+
) -> List[Tuple[str, Dict[str, Any]]]:
|
| 290 |
+
"""The legal moves (UCI) as buttons, sorted so moves of the same piece sit together."""
|
| 291 |
+
return [
|
| 292 |
+
(move, {"move": move})
|
| 293 |
+
for move in sorted(observation.get("legal_moves") or [])
|
| 294 |
+
]
|
| 295 |
+
|
| 296 |
+
def render_web(self, observation: Dict[str, Any]) -> Optional[str]:
|
| 297 |
+
"""Draw the board from the FEN, white at the bottom, the king in check highlighted."""
|
| 298 |
+
if "fen" not in observation:
|
| 299 |
+
return None
|
| 300 |
+
board = chess.Board(observation["fen"])
|
| 301 |
+
checked = board.king(board.turn) if observation["is_check"] else None
|
| 302 |
+
cells = []
|
| 303 |
+
for rank in range(7, -1, -1):
|
| 304 |
+
cells.append(f'<span style="opacity:.6;font-size:12px">{rank + 1}</span>')
|
| 305 |
+
for file in range(8):
|
| 306 |
+
square = chess.square(file, rank)
|
| 307 |
+
piece = board.piece_at(square)
|
| 308 |
+
# Fixed board colours, so a1 stays dark in both themes.
|
| 309 |
+
if square == checked:
|
| 310 |
+
bg = "var(--color-accent)"
|
| 311 |
+
elif (file + rank) % 2:
|
| 312 |
+
bg = "hsl(37 50% 84%)"
|
| 313 |
+
else:
|
| 314 |
+
bg = "hsl(28 33% 55%)"
|
| 315 |
+
# U+FE0E keeps the pawn as a text glyph instead of an emoji.
|
| 316 |
+
glyph = piece.unicode_symbol() + "\ufe0e" if piece else ""
|
| 317 |
+
cells.append(
|
| 318 |
+
f'<span style="background:{bg};color:black">{glyph}</span>'
|
| 319 |
+
)
|
| 320 |
+
cells.append("<span></span>")
|
| 321 |
+
cells += [
|
| 322 |
+
f'<span style="opacity:.6;font-size:12px;line-height:20px">{f}</span>'
|
| 323 |
+
for f in "abcdefgh"
|
| 324 |
+
]
|
| 325 |
+
return (
|
| 326 |
+
'<div role="img" aria-label="Chess board" style="display:inline-grid;'
|
| 327 |
+
"grid-template-columns:16px repeat(8,40px);grid-template-rows:repeat(8,40px) 20px;"
|
| 328 |
+
"line-height:40px;text-align:center;font-size:30px;color:var(--body-text-color);padding:10px;"
|
| 329 |
+
'border:1px solid var(--border-color-primary);border-radius:10px;background:var(--background-fill-secondary)">'
|
| 330 |
+
+ "".join(cells)
|
| 331 |
+
+ "</div>"
|
| 332 |
+
)
|
| 333 |
+
|
| 334 |
+
@property
|
| 335 |
+
def state(self):
|
| 336 |
+
"""Return the current episode state."""
|
| 337 |
+
return self._state
|
envs/chess_env/server/rubrics.py
ADDED
|
@@ -0,0 +1,86 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 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 the
|
| 5 |
+
# LICENSE file in the root directory of this source tree.
|
| 6 |
+
|
| 7 |
+
"""Chess-specific rubrics for reward computation.
|
| 8 |
+
|
| 9 |
+
Provides ChessWinLossRubric that extends ExponentialDiscountingTrajectoryRubric
|
| 10 |
+
to compute temporally-discounted rewards based on chess game outcomes.
|
| 11 |
+
|
| 12 |
+
See RFC 004 for rubric design: rfcs/004-rubrics.md
|
| 13 |
+
"""
|
| 14 |
+
|
| 15 |
+
from typing import Any, List, Tuple
|
| 16 |
+
|
| 17 |
+
try:
|
| 18 |
+
from openenv.core.rubrics.trajectory import ExponentialDiscountingTrajectoryRubric
|
| 19 |
+
except ModuleNotFoundError:
|
| 20 |
+
|
| 21 |
+
class ExponentialDiscountingTrajectoryRubric:
|
| 22 |
+
"""Compatibility fallback when the installed core lacks rubrics."""
|
| 23 |
+
|
| 24 |
+
def __init__(self, gamma: float = 0.99, intermediate_reward: float = 0.0):
|
| 25 |
+
self.gamma = gamma
|
| 26 |
+
self.intermediate_reward = intermediate_reward
|
| 27 |
+
self._trajectory: List[Tuple[Any, Any]] = []
|
| 28 |
+
|
| 29 |
+
def __call__(self, action: Any, observation: Any) -> float:
|
| 30 |
+
self._trajectory.append((action, observation))
|
| 31 |
+
if getattr(observation, "done", False):
|
| 32 |
+
return self.score_trajectory(self._trajectory)
|
| 33 |
+
return self.intermediate_reward
|
| 34 |
+
|
| 35 |
+
def reset(self) -> None:
|
| 36 |
+
self._trajectory = []
|
| 37 |
+
|
| 38 |
+
def compute_step_rewards(self) -> List[float]:
|
| 39 |
+
if not self._trajectory:
|
| 40 |
+
return []
|
| 41 |
+
|
| 42 |
+
final_score = self.score_trajectory(self._trajectory)
|
| 43 |
+
total_steps = len(self._trajectory)
|
| 44 |
+
return [
|
| 45 |
+
self.gamma ** (total_steps - 1 - step_index) * final_score
|
| 46 |
+
for step_index in range(total_steps)
|
| 47 |
+
]
|
| 48 |
+
|
| 49 |
+
|
| 50 |
+
class ChessWinLossRubric(ExponentialDiscountingTrajectoryRubric):
|
| 51 |
+
"""Score chess game based on win/loss/draw outcome with temporal discounting.
|
| 52 |
+
|
| 53 |
+
Per-step reward: r_t = gamma^(T-1-t) * R_final
|
| 54 |
+
|
| 55 |
+
Terminal rewards:
|
| 56 |
+
- Win: +1.0
|
| 57 |
+
- Loss: -1.0
|
| 58 |
+
- Draw: 0.0
|
| 59 |
+
|
| 60 |
+
Note: Unlike the base class convention of 0.0-1.0 scores, chess uses
|
| 61 |
+
-1.0 for losses to match standard chess RL conventions.
|
| 62 |
+
|
| 63 |
+
Usage:
|
| 64 |
+
rubric = ChessWinLossRubric(gamma=0.99)
|
| 65 |
+
rubric.reset()
|
| 66 |
+
for action, obs in episode:
|
| 67 |
+
reward = rubric(action, obs) # 0.0 until done
|
| 68 |
+
step_rewards = rubric.compute_step_rewards()
|
| 69 |
+
"""
|
| 70 |
+
|
| 71 |
+
def score_trajectory(self, trajectory: List[Tuple[Any, Any]]) -> float:
|
| 72 |
+
"""Score based on game outcome from final observation.
|
| 73 |
+
|
| 74 |
+
Reads the reward from the terminal observation, which encodes:
|
| 75 |
+
+1.0 for agent win, -1.0 for agent loss, 0.0 for draw.
|
| 76 |
+
|
| 77 |
+
Args:
|
| 78 |
+
trajectory: List of (action, observation) tuples.
|
| 79 |
+
|
| 80 |
+
Returns:
|
| 81 |
+
Terminal reward: +1.0 (win), -1.0 (loss), 0.0 (draw).
|
| 82 |
+
"""
|
| 83 |
+
if not trajectory:
|
| 84 |
+
return 0.0
|
| 85 |
+
_, final_obs = trajectory[-1]
|
| 86 |
+
return getattr(final_obs, "reward", 0.0)
|
models.py
ADDED
|
@@ -0,0 +1,68 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 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 the
|
| 5 |
+
# LICENSE file in the root directory of this source tree.
|
| 6 |
+
|
| 7 |
+
"""
|
| 8 |
+
Data models for Chess Environment.
|
| 9 |
+
|
| 10 |
+
This module defines the Action, Observation, and State types for chess games
|
| 11 |
+
via the OpenEnv interface. Uses the moonfish chess engine for move search
|
| 12 |
+
and position evaluation.
|
| 13 |
+
"""
|
| 14 |
+
|
| 15 |
+
from __future__ import annotations
|
| 16 |
+
|
| 17 |
+
from typing import List, Optional
|
| 18 |
+
|
| 19 |
+
from openenv.core.env_server import Action, Observation, State
|
| 20 |
+
from pydantic import Field
|
| 21 |
+
|
| 22 |
+
|
| 23 |
+
class ChessAction(Action):
|
| 24 |
+
"""
|
| 25 |
+
Action for Chess environment.
|
| 26 |
+
|
| 27 |
+
Attributes:
|
| 28 |
+
move: UCI format move string (e.g., "e2e4", "e7e8q" for promotion).
|
| 29 |
+
"""
|
| 30 |
+
|
| 31 |
+
move: str
|
| 32 |
+
|
| 33 |
+
|
| 34 |
+
class ChessObservation(Observation):
|
| 35 |
+
"""
|
| 36 |
+
Observation for Chess environment.
|
| 37 |
+
|
| 38 |
+
Attributes:
|
| 39 |
+
fen: Board position in FEN notation.
|
| 40 |
+
legal_moves: List of legal moves in UCI format.
|
| 41 |
+
is_check: Whether the current player is in check.
|
| 42 |
+
done: Whether the game is over.
|
| 43 |
+
reward: Reward for the last action.
|
| 44 |
+
result: Game result string if game is over (e.g., "1-0", "0-1", "1/2-1/2").
|
| 45 |
+
"""
|
| 46 |
+
|
| 47 |
+
fen: str = ""
|
| 48 |
+
legal_moves: List[str] = Field(default_factory=list)
|
| 49 |
+
is_check: bool = False
|
| 50 |
+
result: Optional[str] = None
|
| 51 |
+
|
| 52 |
+
|
| 53 |
+
class ChessState(State):
|
| 54 |
+
"""
|
| 55 |
+
State for Chess environment.
|
| 56 |
+
|
| 57 |
+
Attributes:
|
| 58 |
+
episode_id: Unique ID for the current game.
|
| 59 |
+
fen: Current board position in FEN notation.
|
| 60 |
+
current_player: "white" or "black".
|
| 61 |
+
move_history: List of moves played in UCI format.
|
| 62 |
+
step_count: Number of half-moves played.
|
| 63 |
+
"""
|
| 64 |
+
|
| 65 |
+
fen: str = "rnbqkbnr/pppppppp/8/8/8/8/PPPPPPPP/RNBQKBNR w KQkq - 0 1"
|
| 66 |
+
current_player: str = "white"
|
| 67 |
+
move_history: List[str] = Field(default_factory=list)
|
| 68 |
+
step_count: int = 0
|
openenv.yaml
ADDED
|
@@ -0,0 +1,6 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
spec_version: 1
|
| 2 |
+
name: chess_env
|
| 3 |
+
type: space
|
| 4 |
+
runtime: fastapi
|
| 5 |
+
app: chess_env.server.app:app
|
| 6 |
+
port: 8000
|
pyproject.toml
ADDED
|
@@ -0,0 +1,31 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
[build-system]
|
| 2 |
+
requires = ["setuptools>=45", "wheel"]
|
| 3 |
+
build-backend = "setuptools.build_meta"
|
| 4 |
+
|
| 5 |
+
[project]
|
| 6 |
+
name = "openenv-chess-env"
|
| 7 |
+
version = "0.1.0"
|
| 8 |
+
description = "Chess Environment for OpenEnv"
|
| 9 |
+
requires-python = ">=3.10"
|
| 10 |
+
dependencies = [
|
| 11 |
+
"openenv @ git+https://github.com/huggingface/OpenEnv.git@v0.10.0",
|
| 12 |
+
"fastapi>=0.115.0",
|
| 13 |
+
"pydantic>=2.0.0",
|
| 14 |
+
"uvicorn[standard]>=0.24.0",
|
| 15 |
+
"requests>=2.31.0",
|
| 16 |
+
"python-chess>=1.999",
|
| 17 |
+
"moonfish>=0.0.1",
|
| 18 |
+
]
|
| 19 |
+
|
| 20 |
+
[project.optional-dependencies]
|
| 21 |
+
dev = [
|
| 22 |
+
"pytest>=8.0.0",
|
| 23 |
+
"pytest-cov>=4.0.0",
|
| 24 |
+
]
|
| 25 |
+
|
| 26 |
+
[project.scripts]
|
| 27 |
+
server = "chess_env.server.app:main"
|
| 28 |
+
|
| 29 |
+
[tool.setuptools]
|
| 30 |
+
packages = ["chess_env", "chess_env.server"]
|
| 31 |
+
package-dir = { "chess_env" = ".", "chess_env.server" = "server" }
|
scripts/verify.sh
ADDED
|
@@ -0,0 +1,60 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
#!/bin/bash
|
| 2 |
+
# Verification script for chess_env
|
| 3 |
+
# Run from repository root: bash envs/chess_env/scripts/verify.sh
|
| 4 |
+
|
| 5 |
+
set -e
|
| 6 |
+
|
| 7 |
+
echo "=== Chess Environment Verification ==="
|
| 8 |
+
echo
|
| 9 |
+
|
| 10 |
+
# Colors for output
|
| 11 |
+
GREEN='\033[0;32m'
|
| 12 |
+
RED='\033[0;31m'
|
| 13 |
+
NC='\033[0m' # No Color
|
| 14 |
+
|
| 15 |
+
check_pass() {
|
| 16 |
+
echo -e "${GREEN}✓ $1${NC}"
|
| 17 |
+
}
|
| 18 |
+
|
| 19 |
+
check_fail() {
|
| 20 |
+
echo -e "${RED}✗ $1${NC}"
|
| 21 |
+
exit 1
|
| 22 |
+
}
|
| 23 |
+
|
| 24 |
+
# 1. Run unit tests
|
| 25 |
+
echo "1. Running unit tests..."
|
| 26 |
+
if PYTHONPATH=src:. uv run pytest tests/envs/test_chess_environment.py -v; then
|
| 27 |
+
check_pass "Unit tests passed"
|
| 28 |
+
else
|
| 29 |
+
check_fail "Unit tests failed"
|
| 30 |
+
fi
|
| 31 |
+
echo
|
| 32 |
+
|
| 33 |
+
# 2. Run linting
|
| 34 |
+
echo "2. Running linting..."
|
| 35 |
+
if uv run ruff check envs/chess_env/ tests/envs/test_chess_environment.py; then
|
| 36 |
+
check_pass "Linting passed"
|
| 37 |
+
else
|
| 38 |
+
check_fail "Linting failed"
|
| 39 |
+
fi
|
| 40 |
+
echo
|
| 41 |
+
|
| 42 |
+
# 3. Check formatting
|
| 43 |
+
echo "3. Checking formatting..."
|
| 44 |
+
if uv run ruff format envs/chess_env/ tests/envs/test_chess_environment.py --check; then
|
| 45 |
+
check_pass "Formatting passed"
|
| 46 |
+
else
|
| 47 |
+
check_fail "Formatting failed"
|
| 48 |
+
fi
|
| 49 |
+
echo
|
| 50 |
+
|
| 51 |
+
# 4. Verify imports
|
| 52 |
+
echo "4. Verifying module imports..."
|
| 53 |
+
if PYTHONPATH=src:. uv run python -c "from envs.chess_env import ChessEnv, ChessAction, ChessObservation, ChessState; print(' Imports: ChessEnv, ChessAction, ChessObservation, ChessState')"; then
|
| 54 |
+
check_pass "Module imports successful"
|
| 55 |
+
else
|
| 56 |
+
check_fail "Module imports failed"
|
| 57 |
+
fi
|
| 58 |
+
echo
|
| 59 |
+
|
| 60 |
+
echo "=== All checks passed! ==="
|
server/Dockerfile
ADDED
|
@@ -0,0 +1,79 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 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 the
|
| 5 |
+
# LICENSE file in the root directory of this source tree.
|
| 6 |
+
|
| 7 |
+
# Multi-stage build using openenv-base
|
| 8 |
+
# This Dockerfile is flexible and works for both:
|
| 9 |
+
# - In-repo environments (with local src/core)
|
| 10 |
+
# - Standalone environments (with openenv from pip)
|
| 11 |
+
# The build script (openenv build) handles context detection and sets appropriate build args.
|
| 12 |
+
|
| 13 |
+
ARG BASE_IMAGE=ghcr.io/huggingface/openenv-base:latest
|
| 14 |
+
FROM ${BASE_IMAGE} AS builder
|
| 15 |
+
|
| 16 |
+
WORKDIR /app
|
| 17 |
+
|
| 18 |
+
# Build argument to control whether we're building standalone or in-repo
|
| 19 |
+
ARG BUILD_MODE=in-repo
|
| 20 |
+
|
| 21 |
+
# Copy environment code (always at root of build context)
|
| 22 |
+
COPY . /app/env
|
| 23 |
+
|
| 24 |
+
# For in-repo builds, openenv is already in the pyproject.toml dependencies
|
| 25 |
+
# For standalone builds, openenv will be installed from pip via pyproject.toml
|
| 26 |
+
WORKDIR /app/env
|
| 27 |
+
|
| 28 |
+
# Ensure uv is available (for local builds where base image lacks it)
|
| 29 |
+
RUN if ! command -v uv >/dev/null 2>&1; then \
|
| 30 |
+
curl -LsSf https://astral.sh/uv/install.sh | sh && \
|
| 31 |
+
mv /root/.local/bin/uv /usr/local/bin/uv && \
|
| 32 |
+
mv /root/.local/bin/uvx /usr/local/bin/uvx; \
|
| 33 |
+
fi
|
| 34 |
+
|
| 35 |
+
# Install git for building from git repos (build-time only)
|
| 36 |
+
RUN apt-get update && apt-get install -y --no-install-recommends \
|
| 37 |
+
git \
|
| 38 |
+
&& rm -rf /var/lib/apt/lists/*
|
| 39 |
+
|
| 40 |
+
# Install dependencies using uv sync
|
| 41 |
+
# First pass: install dependencies without the project (for better caching)
|
| 42 |
+
# Second pass: install the project itself
|
| 43 |
+
RUN --mount=type=cache,target=/root/.cache/uv \
|
| 44 |
+
if [ -f uv.lock ]; then \
|
| 45 |
+
uv sync --frozen --no-install-project --no-editable; \
|
| 46 |
+
else \
|
| 47 |
+
uv sync --no-install-project --no-editable; \
|
| 48 |
+
fi
|
| 49 |
+
|
| 50 |
+
RUN --mount=type=cache,target=/root/.cache/uv \
|
| 51 |
+
if [ -f uv.lock ]; then \
|
| 52 |
+
uv sync --frozen --no-editable; \
|
| 53 |
+
else \
|
| 54 |
+
uv sync --no-editable; \
|
| 55 |
+
fi
|
| 56 |
+
|
| 57 |
+
# Final runtime stage
|
| 58 |
+
FROM ${BASE_IMAGE}
|
| 59 |
+
|
| 60 |
+
WORKDIR /app
|
| 61 |
+
|
| 62 |
+
# Copy the virtual environment from builder
|
| 63 |
+
COPY --from=builder /app/env/.venv /app/.venv
|
| 64 |
+
|
| 65 |
+
# Copy the environment code
|
| 66 |
+
COPY --from=builder /app/env /app/env
|
| 67 |
+
|
| 68 |
+
# Set PATH to use the virtual environment
|
| 69 |
+
ENV PATH="/app/.venv/bin:$PATH"
|
| 70 |
+
|
| 71 |
+
# Set PYTHONPATH so imports work correctly
|
| 72 |
+
ENV PYTHONPATH="/app/env:$PYTHONPATH"
|
| 73 |
+
|
| 74 |
+
# Health check using Python (more portable than curl/wget)
|
| 75 |
+
HEALTHCHECK --interval=30s --timeout=3s --start-period=5s --retries=3 \
|
| 76 |
+
CMD python -c "import urllib.request; urllib.request.urlopen('http://localhost:8000/health')" || exit 1
|
| 77 |
+
|
| 78 |
+
# Run the FastAPI server from the installed package.
|
| 79 |
+
CMD ["uvicorn", "chess_env.server.app:app", "--host", "0.0.0.0", "--port", "8000"]
|
server/__init__.py
ADDED
|
@@ -0,0 +1,11 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 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 the
|
| 5 |
+
# LICENSE file in the root directory of this source tree.
|
| 6 |
+
|
| 7 |
+
"""Chess Environment server module."""
|
| 8 |
+
|
| 9 |
+
from .chess_environment import ChessEnvironment
|
| 10 |
+
|
| 11 |
+
__all__ = ["ChessEnvironment"]
|
server/app.py
ADDED
|
@@ -0,0 +1,35 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 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 the
|
| 5 |
+
# LICENSE file in the root directory of this source tree.
|
| 6 |
+
|
| 7 |
+
"""FastAPI application for the Chess Environment."""
|
| 8 |
+
|
| 9 |
+
import os
|
| 10 |
+
|
| 11 |
+
from openenv.core.env_server import create_app
|
| 12 |
+
|
| 13 |
+
from ..models import ChessAction, ChessObservation, ChessState
|
| 14 |
+
from .chess_environment import ChessEnvironment
|
| 15 |
+
|
| 16 |
+
# Create the FastAPI app
|
| 17 |
+
# Pass the class (factory) instead of an instance for WebSocket session support
|
| 18 |
+
app = create_app(
|
| 19 |
+
ChessEnvironment,
|
| 20 |
+
ChessAction,
|
| 21 |
+
ChessObservation,
|
| 22 |
+
env_name="chess_env",
|
| 23 |
+
state_cls=ChessState,
|
| 24 |
+
max_concurrent_envs=int(os.getenv("MAX_CONCURRENT_ENVS", "8")),
|
| 25 |
+
)
|
| 26 |
+
|
| 27 |
+
|
| 28 |
+
def main():
|
| 29 |
+
import uvicorn
|
| 30 |
+
|
| 31 |
+
uvicorn.run(app, host="0.0.0.0", port=8000)
|
| 32 |
+
|
| 33 |
+
|
| 34 |
+
if __name__ == "__main__":
|
| 35 |
+
main()
|
server/chess_environment.py
ADDED
|
@@ -0,0 +1,337 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 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 the
|
| 5 |
+
# LICENSE file in the root directory of this source tree.
|
| 6 |
+
|
| 7 |
+
"""
|
| 8 |
+
Chess Environment for OpenEnv using moonfish.
|
| 9 |
+
|
| 10 |
+
This module provides a chess environment for reinforcement learning,
|
| 11 |
+
using python-chess for game logic and moonfish for position evaluation
|
| 12 |
+
and opponent play.
|
| 13 |
+
"""
|
| 14 |
+
|
| 15 |
+
import random
|
| 16 |
+
import threading
|
| 17 |
+
import uuid
|
| 18 |
+
from typing import Any, Dict, List, Optional, Tuple
|
| 19 |
+
|
| 20 |
+
import chess
|
| 21 |
+
from moonfish.lib import search_move
|
| 22 |
+
from moonfish.psqt import board_evaluation, BOARD_EVALUATION_CACHE
|
| 23 |
+
from openenv.core.env_server import Environment
|
| 24 |
+
|
| 25 |
+
from ..models import ChessAction, ChessObservation, ChessState
|
| 26 |
+
from .rubrics import ChessWinLossRubric
|
| 27 |
+
|
| 28 |
+
# moonfish keeps one evaluation cache per process, so sessions take turns with it.
|
| 29 |
+
_MOONFISH_LOCK = threading.Lock()
|
| 30 |
+
|
| 31 |
+
|
| 32 |
+
class ChessEnvironment(Environment):
|
| 33 |
+
"""
|
| 34 |
+
Chess environment implementing the OpenEnv interface.
|
| 35 |
+
|
| 36 |
+
Uses python-chess for game logic and moonfish for position evaluation.
|
| 37 |
+
Designed for RL training where an agent plays as one color against
|
| 38 |
+
an opponent (which can be random, moonfish engine, or self-play).
|
| 39 |
+
"""
|
| 40 |
+
|
| 41 |
+
SUPPORTS_CONCURRENT_SESSIONS = True
|
| 42 |
+
|
| 43 |
+
def __init__(
|
| 44 |
+
self,
|
| 45 |
+
opponent: str = "moonfish",
|
| 46 |
+
opponent_depth: int = 2,
|
| 47 |
+
max_moves: int = 500,
|
| 48 |
+
agent_color: str | None = None,
|
| 49 |
+
gamma: float = 0.99,
|
| 50 |
+
):
|
| 51 |
+
"""
|
| 52 |
+
Initialize the chess environment.
|
| 53 |
+
|
| 54 |
+
Args:
|
| 55 |
+
opponent: Opponent type - "moonfish", "random", or None (self-play)
|
| 56 |
+
opponent_depth: Search depth when using moonfish as opponent
|
| 57 |
+
max_moves: Maximum half-moves before draw
|
| 58 |
+
agent_color: Which color the agent plays - "white", "black", or None (alternate each episode)
|
| 59 |
+
gamma: Discount factor for temporal credit assignment (0-1)
|
| 60 |
+
"""
|
| 61 |
+
super().__init__(rubric=ChessWinLossRubric(gamma=gamma))
|
| 62 |
+
self._opponent = opponent
|
| 63 |
+
self._opponent_depth = opponent_depth
|
| 64 |
+
self._max_moves = max_moves
|
| 65 |
+
self._agent_color_setting = agent_color
|
| 66 |
+
self._gamma = gamma
|
| 67 |
+
self._board = None
|
| 68 |
+
self._state = None
|
| 69 |
+
self._agent_color = chess.WHITE
|
| 70 |
+
self._agent_move_count = 0 # Track agent moves for discounting
|
| 71 |
+
self.reset()
|
| 72 |
+
|
| 73 |
+
def reset(self, fen: str = None):
|
| 74 |
+
"""
|
| 75 |
+
Initialize a new chess game episode.
|
| 76 |
+
|
| 77 |
+
Args:
|
| 78 |
+
fen: Optional starting position in FEN notation.
|
| 79 |
+
|
| 80 |
+
Returns:
|
| 81 |
+
Initial observation of the board state.
|
| 82 |
+
"""
|
| 83 |
+
self._reset_rubric()
|
| 84 |
+
# moonfish's module-level evaluation cache is never trimmed.
|
| 85 |
+
with _MOONFISH_LOCK:
|
| 86 |
+
BOARD_EVALUATION_CACHE.clear()
|
| 87 |
+
|
| 88 |
+
if fen:
|
| 89 |
+
self._board = chess.Board(fen)
|
| 90 |
+
else:
|
| 91 |
+
self._board = chess.Board()
|
| 92 |
+
|
| 93 |
+
# Generate episode ID once
|
| 94 |
+
episode_id = str(uuid.uuid4())
|
| 95 |
+
|
| 96 |
+
# Determine agent color
|
| 97 |
+
if self._agent_color_setting == "white":
|
| 98 |
+
self._agent_color = chess.WHITE
|
| 99 |
+
elif self._agent_color_setting == "black":
|
| 100 |
+
self._agent_color = chess.BLACK
|
| 101 |
+
elif self._agent_color_setting is None:
|
| 102 |
+
# Alternate each episode based on episode_id hash
|
| 103 |
+
self._agent_color = (
|
| 104 |
+
chess.WHITE if hash(episode_id) % 2 == 0 else chess.BLACK
|
| 105 |
+
)
|
| 106 |
+
else:
|
| 107 |
+
self._agent_color = chess.WHITE
|
| 108 |
+
|
| 109 |
+
self._state = ChessState(
|
| 110 |
+
episode_id=episode_id,
|
| 111 |
+
step_count=0,
|
| 112 |
+
current_player="white" if self._board.turn else "black",
|
| 113 |
+
fen=self._board.fen(),
|
| 114 |
+
move_history=[],
|
| 115 |
+
)
|
| 116 |
+
self._agent_move_count = 0
|
| 117 |
+
|
| 118 |
+
# If agent plays Black and opponent is configured, opponent moves first
|
| 119 |
+
if self._opponent is not None and self._agent_color == chess.BLACK:
|
| 120 |
+
self._make_opponent_move()
|
| 121 |
+
|
| 122 |
+
# Check if starting position is already terminal (e.g., custom FEN with checkmate/stalemate)
|
| 123 |
+
_, done = self._calculate_reward_and_done()
|
| 124 |
+
return self._make_observation(done=done)
|
| 125 |
+
|
| 126 |
+
def step(self, action: ChessAction):
|
| 127 |
+
"""
|
| 128 |
+
Execute a chess move and return the resulting state.
|
| 129 |
+
|
| 130 |
+
Args:
|
| 131 |
+
action: The move to make in UCI format (e.g., "e2e4").
|
| 132 |
+
|
| 133 |
+
Returns:
|
| 134 |
+
Observation with reward and done flag.
|
| 135 |
+
"""
|
| 136 |
+
if self._board is None:
|
| 137 |
+
raise RuntimeError("Environment not initialized. Call reset() first.")
|
| 138 |
+
|
| 139 |
+
# Parse the move
|
| 140 |
+
try:
|
| 141 |
+
move = chess.Move.from_uci(action.move)
|
| 142 |
+
except ValueError:
|
| 143 |
+
return self._make_observation(reward=-0.1, done=False)
|
| 144 |
+
|
| 145 |
+
# Check if move is legal
|
| 146 |
+
if move not in self._board.legal_moves:
|
| 147 |
+
return self._make_observation(reward=-0.1, done=False)
|
| 148 |
+
|
| 149 |
+
# Execute the move
|
| 150 |
+
self._board.push(move)
|
| 151 |
+
self._state.step_count += 1
|
| 152 |
+
self._agent_move_count += 1
|
| 153 |
+
self._state.move_history.append(action.move)
|
| 154 |
+
self._state.current_player = "white" if self._board.turn else "black"
|
| 155 |
+
self._state.fen = self._board.fen()
|
| 156 |
+
|
| 157 |
+
# Calculate reward and check for game end
|
| 158 |
+
reward, done = self._calculate_reward_and_done()
|
| 159 |
+
|
| 160 |
+
# If game not over and opponent is configured, make opponent move
|
| 161 |
+
if not done and self._opponent is not None:
|
| 162 |
+
self._make_opponent_move()
|
| 163 |
+
reward, done = self._calculate_reward_and_done()
|
| 164 |
+
|
| 165 |
+
obs = self._make_observation(reward=reward, done=done)
|
| 166 |
+
|
| 167 |
+
# Feed to rubric for trajectory accumulation
|
| 168 |
+
self._apply_rubric(action, obs)
|
| 169 |
+
|
| 170 |
+
# Note: obs.reward is kept from _calculate_reward_and_done() rather than
|
| 171 |
+
# overridden by the rubric. The rubric here is a TrajectoryRubric that
|
| 172 |
+
# returns 0.0 until done then the final score — the inline reward logic
|
| 173 |
+
# already provides the correct per-step reward including penalties for
|
| 174 |
+
# invalid/illegal moves. The rubric's value is in compute_step_rewards()
|
| 175 |
+
# for training infrastructure, not in replacing obs.reward.
|
| 176 |
+
|
| 177 |
+
return obs
|
| 178 |
+
|
| 179 |
+
def _make_observation(self, reward: float = 0.0, done: bool = False):
|
| 180 |
+
"""Build observation from current board state."""
|
| 181 |
+
legal_moves = [move.uci() for move in self._board.legal_moves]
|
| 182 |
+
|
| 183 |
+
result = None
|
| 184 |
+
if done:
|
| 185 |
+
result = self._get_result_string()
|
| 186 |
+
|
| 187 |
+
with _MOONFISH_LOCK:
|
| 188 |
+
evaluation = board_evaluation(self._board)
|
| 189 |
+
metadata = {
|
| 190 |
+
"evaluation": evaluation,
|
| 191 |
+
"fullmove_number": self._board.fullmove_number,
|
| 192 |
+
"halfmove_clock": self._board.halfmove_clock,
|
| 193 |
+
}
|
| 194 |
+
|
| 195 |
+
# Compute discounted rewards for all agent moves when episode ends
|
| 196 |
+
if done and self._agent_move_count > 0:
|
| 197 |
+
discounted_rewards = self._compute_discounted_rewards(reward)
|
| 198 |
+
metadata["discounted_rewards"] = discounted_rewards
|
| 199 |
+
metadata["gamma"] = self._gamma
|
| 200 |
+
|
| 201 |
+
return ChessObservation(
|
| 202 |
+
fen=self._board.fen(),
|
| 203 |
+
legal_moves=legal_moves,
|
| 204 |
+
is_check=self._board.is_check(),
|
| 205 |
+
done=done,
|
| 206 |
+
reward=reward,
|
| 207 |
+
result=result,
|
| 208 |
+
metadata=metadata,
|
| 209 |
+
)
|
| 210 |
+
|
| 211 |
+
def _calculate_reward_and_done(self):
|
| 212 |
+
"""Calculate reward and check if episode is done."""
|
| 213 |
+
if self._board.is_checkmate():
|
| 214 |
+
winner = not self._board.turn
|
| 215 |
+
if winner == self._agent_color:
|
| 216 |
+
return 1.0, True
|
| 217 |
+
else:
|
| 218 |
+
return -1.0, True
|
| 219 |
+
|
| 220 |
+
if self._board.is_stalemate():
|
| 221 |
+
return 0.0, True
|
| 222 |
+
|
| 223 |
+
if self._board.is_insufficient_material():
|
| 224 |
+
return 0.0, True
|
| 225 |
+
|
| 226 |
+
if self._board.is_fifty_moves():
|
| 227 |
+
return 0.0, True
|
| 228 |
+
|
| 229 |
+
if self._board.is_repetition(3):
|
| 230 |
+
return 0.0, True
|
| 231 |
+
|
| 232 |
+
if self._state.step_count >= self._max_moves:
|
| 233 |
+
return 0.0, True
|
| 234 |
+
|
| 235 |
+
return 0.0, False
|
| 236 |
+
|
| 237 |
+
def _get_result_string(self):
|
| 238 |
+
"""Get the game result as a string."""
|
| 239 |
+
if self._board.is_checkmate():
|
| 240 |
+
return "1-0" if not self._board.turn else "0-1"
|
| 241 |
+
return "1/2-1/2"
|
| 242 |
+
|
| 243 |
+
def _compute_discounted_rewards(self, terminal_reward: float) -> list[float]:
|
| 244 |
+
"""
|
| 245 |
+
Compute temporally discounted rewards for all agent moves.
|
| 246 |
+
|
| 247 |
+
Uses exponential discounting: r_t = γ^(T-1-t) * R_final
|
| 248 |
+
where T is total agent moves, t is move index, R_final is terminal reward.
|
| 249 |
+
|
| 250 |
+
Earlier moves get less credit, later moves get more credit.
|
| 251 |
+
This helps with credit assignment in long games.
|
| 252 |
+
|
| 253 |
+
Args:
|
| 254 |
+
terminal_reward: The final reward (+1 win, -1 loss, 0 draw)
|
| 255 |
+
|
| 256 |
+
Returns:
|
| 257 |
+
List of discounted rewards, one per agent move
|
| 258 |
+
"""
|
| 259 |
+
T = self._agent_move_count
|
| 260 |
+
discounted_rewards = []
|
| 261 |
+
for t in range(T):
|
| 262 |
+
# γ^(T-1-t) means last move (t=T-1) gets γ^0 = 1.0
|
| 263 |
+
# First move (t=0) gets γ^(T-1)
|
| 264 |
+
discount = self._gamma ** (T - 1 - t)
|
| 265 |
+
discounted_rewards.append(discount * terminal_reward)
|
| 266 |
+
return discounted_rewards
|
| 267 |
+
|
| 268 |
+
def _make_opponent_move(self):
|
| 269 |
+
"""Make a move for the opponent using configured strategy."""
|
| 270 |
+
if not self._board.legal_moves:
|
| 271 |
+
return
|
| 272 |
+
|
| 273 |
+
if self._opponent == "moonfish":
|
| 274 |
+
with _MOONFISH_LOCK:
|
| 275 |
+
move = search_move(self._board, depth=self._opponent_depth)
|
| 276 |
+
elif self._opponent == "random":
|
| 277 |
+
move = random.choice(list(self._board.legal_moves))
|
| 278 |
+
else:
|
| 279 |
+
return
|
| 280 |
+
|
| 281 |
+
self._board.push(move)
|
| 282 |
+
self._state.step_count += 1
|
| 283 |
+
self._state.move_history.append(move.uci())
|
| 284 |
+
self._state.current_player = "white" if self._board.turn else "black"
|
| 285 |
+
self._state.fen = self._board.fen()
|
| 286 |
+
|
| 287 |
+
def web_actions(
|
| 288 |
+
self, observation: Dict[str, Any]
|
| 289 |
+
) -> List[Tuple[str, Dict[str, Any]]]:
|
| 290 |
+
"""The legal moves (UCI) as buttons, sorted so moves of the same piece sit together."""
|
| 291 |
+
return [
|
| 292 |
+
(move, {"move": move})
|
| 293 |
+
for move in sorted(observation.get("legal_moves") or [])
|
| 294 |
+
]
|
| 295 |
+
|
| 296 |
+
def render_web(self, observation: Dict[str, Any]) -> Optional[str]:
|
| 297 |
+
"""Draw the board from the FEN, white at the bottom, the king in check highlighted."""
|
| 298 |
+
if "fen" not in observation:
|
| 299 |
+
return None
|
| 300 |
+
board = chess.Board(observation["fen"])
|
| 301 |
+
checked = board.king(board.turn) if observation["is_check"] else None
|
| 302 |
+
cells = []
|
| 303 |
+
for rank in range(7, -1, -1):
|
| 304 |
+
cells.append(f'<span style="opacity:.6;font-size:12px">{rank + 1}</span>')
|
| 305 |
+
for file in range(8):
|
| 306 |
+
square = chess.square(file, rank)
|
| 307 |
+
piece = board.piece_at(square)
|
| 308 |
+
# Fixed board colours, so a1 stays dark in both themes.
|
| 309 |
+
if square == checked:
|
| 310 |
+
bg = "var(--color-accent)"
|
| 311 |
+
elif (file + rank) % 2:
|
| 312 |
+
bg = "hsl(37 50% 84%)"
|
| 313 |
+
else:
|
| 314 |
+
bg = "hsl(28 33% 55%)"
|
| 315 |
+
# U+FE0E keeps the pawn as a text glyph instead of an emoji.
|
| 316 |
+
glyph = piece.unicode_symbol() + "\ufe0e" if piece else ""
|
| 317 |
+
cells.append(
|
| 318 |
+
f'<span style="background:{bg};color:black">{glyph}</span>'
|
| 319 |
+
)
|
| 320 |
+
cells.append("<span></span>")
|
| 321 |
+
cells += [
|
| 322 |
+
f'<span style="opacity:.6;font-size:12px;line-height:20px">{f}</span>'
|
| 323 |
+
for f in "abcdefgh"
|
| 324 |
+
]
|
| 325 |
+
return (
|
| 326 |
+
'<div role="img" aria-label="Chess board" style="display:inline-grid;'
|
| 327 |
+
"grid-template-columns:16px repeat(8,40px);grid-template-rows:repeat(8,40px) 20px;"
|
| 328 |
+
"line-height:40px;text-align:center;font-size:30px;color:var(--body-text-color);padding:10px;"
|
| 329 |
+
'border:1px solid var(--border-color-primary);border-radius:10px;background:var(--background-fill-secondary)">'
|
| 330 |
+
+ "".join(cells)
|
| 331 |
+
+ "</div>"
|
| 332 |
+
)
|
| 333 |
+
|
| 334 |
+
@property
|
| 335 |
+
def state(self):
|
| 336 |
+
"""Return the current episode state."""
|
| 337 |
+
return self._state
|
server/rubrics.py
ADDED
|
@@ -0,0 +1,86 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 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 the
|
| 5 |
+
# LICENSE file in the root directory of this source tree.
|
| 6 |
+
|
| 7 |
+
"""Chess-specific rubrics for reward computation.
|
| 8 |
+
|
| 9 |
+
Provides ChessWinLossRubric that extends ExponentialDiscountingTrajectoryRubric
|
| 10 |
+
to compute temporally-discounted rewards based on chess game outcomes.
|
| 11 |
+
|
| 12 |
+
See RFC 004 for rubric design: rfcs/004-rubrics.md
|
| 13 |
+
"""
|
| 14 |
+
|
| 15 |
+
from typing import Any, List, Tuple
|
| 16 |
+
|
| 17 |
+
try:
|
| 18 |
+
from openenv.core.rubrics.trajectory import ExponentialDiscountingTrajectoryRubric
|
| 19 |
+
except ModuleNotFoundError:
|
| 20 |
+
|
| 21 |
+
class ExponentialDiscountingTrajectoryRubric:
|
| 22 |
+
"""Compatibility fallback when the installed core lacks rubrics."""
|
| 23 |
+
|
| 24 |
+
def __init__(self, gamma: float = 0.99, intermediate_reward: float = 0.0):
|
| 25 |
+
self.gamma = gamma
|
| 26 |
+
self.intermediate_reward = intermediate_reward
|
| 27 |
+
self._trajectory: List[Tuple[Any, Any]] = []
|
| 28 |
+
|
| 29 |
+
def __call__(self, action: Any, observation: Any) -> float:
|
| 30 |
+
self._trajectory.append((action, observation))
|
| 31 |
+
if getattr(observation, "done", False):
|
| 32 |
+
return self.score_trajectory(self._trajectory)
|
| 33 |
+
return self.intermediate_reward
|
| 34 |
+
|
| 35 |
+
def reset(self) -> None:
|
| 36 |
+
self._trajectory = []
|
| 37 |
+
|
| 38 |
+
def compute_step_rewards(self) -> List[float]:
|
| 39 |
+
if not self._trajectory:
|
| 40 |
+
return []
|
| 41 |
+
|
| 42 |
+
final_score = self.score_trajectory(self._trajectory)
|
| 43 |
+
total_steps = len(self._trajectory)
|
| 44 |
+
return [
|
| 45 |
+
self.gamma ** (total_steps - 1 - step_index) * final_score
|
| 46 |
+
for step_index in range(total_steps)
|
| 47 |
+
]
|
| 48 |
+
|
| 49 |
+
|
| 50 |
+
class ChessWinLossRubric(ExponentialDiscountingTrajectoryRubric):
|
| 51 |
+
"""Score chess game based on win/loss/draw outcome with temporal discounting.
|
| 52 |
+
|
| 53 |
+
Per-step reward: r_t = gamma^(T-1-t) * R_final
|
| 54 |
+
|
| 55 |
+
Terminal rewards:
|
| 56 |
+
- Win: +1.0
|
| 57 |
+
- Loss: -1.0
|
| 58 |
+
- Draw: 0.0
|
| 59 |
+
|
| 60 |
+
Note: Unlike the base class convention of 0.0-1.0 scores, chess uses
|
| 61 |
+
-1.0 for losses to match standard chess RL conventions.
|
| 62 |
+
|
| 63 |
+
Usage:
|
| 64 |
+
rubric = ChessWinLossRubric(gamma=0.99)
|
| 65 |
+
rubric.reset()
|
| 66 |
+
for action, obs in episode:
|
| 67 |
+
reward = rubric(action, obs) # 0.0 until done
|
| 68 |
+
step_rewards = rubric.compute_step_rewards()
|
| 69 |
+
"""
|
| 70 |
+
|
| 71 |
+
def score_trajectory(self, trajectory: List[Tuple[Any, Any]]) -> float:
|
| 72 |
+
"""Score based on game outcome from final observation.
|
| 73 |
+
|
| 74 |
+
Reads the reward from the terminal observation, which encodes:
|
| 75 |
+
+1.0 for agent win, -1.0 for agent loss, 0.0 for draw.
|
| 76 |
+
|
| 77 |
+
Args:
|
| 78 |
+
trajectory: List of (action, observation) tuples.
|
| 79 |
+
|
| 80 |
+
Returns:
|
| 81 |
+
Terminal reward: +1.0 (win), -1.0 (loss), 0.0 (draw).
|
| 82 |
+
"""
|
| 83 |
+
if not trajectory:
|
| 84 |
+
return 0.0
|
| 85 |
+
_, final_obs = trajectory[-1]
|
| 86 |
+
return getattr(final_obs, "reward", 0.0)
|
src/__init__.py
ADDED
|
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# SPDX-License-Identifier: BSD-3-Clause
|
| 2 |
+
|
| 3 |
+
"""EnvTorch: Standardized agentic execution environments."""
|
src/core/README.md
ADDED
|
@@ -0,0 +1,237 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# OpenEnv: Agentic Execution Environments
|
| 2 |
+
|
| 3 |
+
An e2e framework for creating, deploying and using isolated execution environments for agentic RL training, built using Gymnasium style simple APIs. OpenEnv provides a standard for interacting with agentic execution environments via simple Gymnasium style APIs - step(), reset(), state(). Users of agentic execution environments can interact with the environment during RL training loops using these simple APIs.
|
| 4 |
+
|
| 5 |
+
In addition to making it easier for researchers and RL framework writers, we also provide tools for environment creators making it easier for them to create richer environments and make them available over familiar protocols like HTTP and packaged using canonical technologies like docker. Environment creators can use the OpenEnv framework to create environments that are isolated, secure, and easy to deploy and use.
|
| 6 |
+
|
| 7 |
+
> [!NOTE]
|
| 8 |
+
> OpenEnv is in early development, see the note in the [project README](../../../README.md) before contributing. The [README](../../../README.md#community-support--acknowledgments) lists the technical committee that governs the project, and the [charter](../../../GOVERNANCE.md) explains how it works.
|
| 9 |
+
|
| 10 |
+
## Overview
|
| 11 |
+
`openenv.core` provides the foundational building blocks for creating and interacting with containerized environments over HTTP. It enables you to build agent environments that can be deployed as Docker containers and accessed via a simple HTTP API.
|
| 12 |
+
|
| 13 |
+
# OpenEnv Core
|
| 14 |
+
|
| 15 |
+
Core components for OpenEnv - a framework for building HTTP-based agentic environments.
|
| 16 |
+
|
| 17 |
+
## Features
|
| 18 |
+
|
| 19 |
+
- **EnvClient**: Async-first client for interacting with remote environments
|
| 20 |
+
- **SyncEnvClient**: Synchronous wrapper via `.sync()` for sync codebases
|
| 21 |
+
- **HTTPEnvServer**: FastAPI-based server wrapper for exposing environments over HTTP/WebSocket
|
| 22 |
+
- **Container Providers**: Pluggable architecture for running containers (Docker, Kubernetes, etc.)
|
| 23 |
+
- **Type System**: Strongly-typed Action/Observation/State interfaces
|
| 24 |
+
- **Web Interface**: Optional web UI for interacting with environments
|
| 25 |
+
- **Experimental Harness Helpers**: `openenv.core.harness` provides
|
| 26 |
+
`ResourceSessionFactory`, `StepEnvSessionAdapter`, and `HarnessAdapter` for
|
| 27 |
+
MCP-first training/evaluation harnesses
|
| 28 |
+
|
| 29 |
+
## Experimental Harness Helpers
|
| 30 |
+
|
| 31 |
+
OpenEnv now includes an additive harness-facing layer for cases where a rollout
|
| 32 |
+
driver should interact with environment resources through tools instead of
|
| 33 |
+
calling `reset()` / `step()` directly. These helpers are currently
|
| 34 |
+
experimental while RFC 005 remains in review, so import them from
|
| 35 |
+
`openenv.core.harness`.
|
| 36 |
+
|
| 37 |
+
Core pieces:
|
| 38 |
+
|
| 39 |
+
- `ResourceSessionFactory.create(task, seed, episode_id)` creates an isolated
|
| 40 |
+
per-rollout session
|
| 41 |
+
- `StepEnvSessionAdapter` wraps an existing client and exposes session tools
|
| 42 |
+
- `SessionMCPBridge` exposes a session through MCP-style JSON-RPC
|
| 43 |
+
- `MCPHarnessAdapter` drives white-box rollouts where the trainer still owns
|
| 44 |
+
model sampling, token IDs, and logprobs
|
| 45 |
+
- `CLIHarnessAdapter` supports opaque evaluation harnesses over the same session
|
| 46 |
+
surface
|
| 47 |
+
- `build_harness_rollout_func(...)` builds a TRL-compatible `rollout_func`
|
| 48 |
+
|
| 49 |
+
This layer is additive. Existing environment clients and MCP environments keep
|
| 50 |
+
their current control-plane APIs unchanged.
|
| 51 |
+
|
| 52 |
+
## Installation
|
| 53 |
+
|
| 54 |
+
```bash
|
| 55 |
+
pip install "openenv"
|
| 56 |
+
```
|
| 57 |
+
|
| 58 |
+
For development:
|
| 59 |
+
```bash
|
| 60 |
+
pip install "openenv"
|
| 61 |
+
```
|
| 62 |
+
|
| 63 |
+
## Quick Start
|
| 64 |
+
|
| 65 |
+
### Creating an Environment Client
|
| 66 |
+
|
| 67 |
+
EnvClient is **async by default**. Use `async with` and `await` for all operations:
|
| 68 |
+
|
| 69 |
+
```python
|
| 70 |
+
import asyncio
|
| 71 |
+
from openenv.core import EnvClient, StepResult
|
| 72 |
+
from dataclasses import dataclass
|
| 73 |
+
from typing import Any
|
| 74 |
+
|
| 75 |
+
|
| 76 |
+
@dataclass
|
| 77 |
+
class MyAction:
|
| 78 |
+
text: str
|
| 79 |
+
|
| 80 |
+
|
| 81 |
+
@dataclass
|
| 82 |
+
class MyObservation:
|
| 83 |
+
response: str
|
| 84 |
+
|
| 85 |
+
|
| 86 |
+
class MyEnvClient(EnvClient[MyAction, MyObservation, Any]):
|
| 87 |
+
def _step_payload(self, action: MyAction) -> dict:
|
| 88 |
+
return {"text": action.text}
|
| 89 |
+
|
| 90 |
+
def _parse_result(self, payload: dict) -> StepResult[MyObservation]:
|
| 91 |
+
obs_data = payload["observation"]
|
| 92 |
+
return StepResult(
|
| 93 |
+
observation=MyObservation(**obs_data),
|
| 94 |
+
reward=payload.get("reward"),
|
| 95 |
+
done=payload.get("done", False),
|
| 96 |
+
)
|
| 97 |
+
|
| 98 |
+
def _parse_state(self, payload: dict) -> Any:
|
| 99 |
+
return payload
|
| 100 |
+
|
| 101 |
+
|
| 102 |
+
# Async usage (recommended)
|
| 103 |
+
async def main():
|
| 104 |
+
client = await MyEnvClient.from_docker_image("my-env:latest")
|
| 105 |
+
async with client:
|
| 106 |
+
result = await client.reset()
|
| 107 |
+
step_result = await client.step(MyAction(text="hello"))
|
| 108 |
+
|
| 109 |
+
|
| 110 |
+
asyncio.run(main())
|
| 111 |
+
|
| 112 |
+
# Sync usage (via .sync() wrapper)
|
| 113 |
+
with MyEnvClient(base_url="http://localhost:8000").sync() as client:
|
| 114 |
+
result = client.reset()
|
| 115 |
+
step_result = client.step(MyAction(text="hello"))
|
| 116 |
+
```
|
| 117 |
+
|
| 118 |
+
### Creating an Environment Server
|
| 119 |
+
|
| 120 |
+
```python
|
| 121 |
+
from openenv.core.env_server import Environment, HTTPEnvServer, create_app
|
| 122 |
+
from dataclasses import dataclass
|
| 123 |
+
|
| 124 |
+
|
| 125 |
+
@dataclass
|
| 126 |
+
class MyAction:
|
| 127 |
+
text: str
|
| 128 |
+
|
| 129 |
+
|
| 130 |
+
@dataclass
|
| 131 |
+
class MyObservation:
|
| 132 |
+
response: str
|
| 133 |
+
reward: float = 0.0
|
| 134 |
+
done: bool = False
|
| 135 |
+
|
| 136 |
+
|
| 137 |
+
class MyEnvironment(Environment):
|
| 138 |
+
def reset(self) -> MyObservation:
|
| 139 |
+
return MyObservation(response="Ready")
|
| 140 |
+
|
| 141 |
+
def step(self, action: MyAction) -> MyObservation:
|
| 142 |
+
return MyObservation(response=f"Echo: {action.text}", reward=1.0, done=False)
|
| 143 |
+
|
| 144 |
+
|
| 145 |
+
# Create FastAPI app
|
| 146 |
+
env = MyEnvironment()
|
| 147 |
+
app = create_app(env, MyAction, MyObservation)
|
| 148 |
+
|
| 149 |
+
# Run with: uvicorn module:app --host 0.0.0.0 --port 8000
|
| 150 |
+
```
|
| 151 |
+
|
| 152 |
+
## Container Providers
|
| 153 |
+
|
| 154 |
+
OpenEnv Core supports multiple container providers:
|
| 155 |
+
|
| 156 |
+
### Local Docker Provider
|
| 157 |
+
|
| 158 |
+
```python
|
| 159 |
+
from openenv.core.containers.runtime import LocalDockerProvider
|
| 160 |
+
|
| 161 |
+
provider = LocalDockerProvider()
|
| 162 |
+
base_url = provider.start_container("my-env:latest")
|
| 163 |
+
provider.wait_for_ready(base_url)
|
| 164 |
+
# Use environment...
|
| 165 |
+
provider.stop_container()
|
| 166 |
+
```
|
| 167 |
+
|
| 168 |
+
### Kubernetes Provider (Coming Soon)
|
| 169 |
+
|
| 170 |
+
```python
|
| 171 |
+
from openenv.core.containers.runtime import KubernetesProvider
|
| 172 |
+
|
| 173 |
+
provider = KubernetesProvider(namespace="envs")
|
| 174 |
+
base_url = provider.start_container("my-env:latest")
|
| 175 |
+
# Use environment...
|
| 176 |
+
provider.stop_container()
|
| 177 |
+
```
|
| 178 |
+
|
| 179 |
+
|
| 180 |
+
## API Reference
|
| 181 |
+
|
| 182 |
+
### EnvClient
|
| 183 |
+
|
| 184 |
+
Async base class for environment clients. Key methods:
|
| 185 |
+
|
| 186 |
+
- `async connect()`: Establish WebSocket connection
|
| 187 |
+
- `async reset(**kwargs)`: Reset environment
|
| 188 |
+
- `async step(action)`: Execute action
|
| 189 |
+
- `async state()`: Get current state
|
| 190 |
+
- `async close()`: Close connection and cleanup
|
| 191 |
+
- `sync()`: Return a SyncEnvClient wrapper for synchronous usage
|
| 192 |
+
|
| 193 |
+
Abstract methods to implement:
|
| 194 |
+
- `_step_payload(action)`: Convert action to JSON
|
| 195 |
+
- `_parse_result(payload)`: Parse response to StepResult
|
| 196 |
+
- `_parse_state(payload)`: Parse state response
|
| 197 |
+
|
| 198 |
+
### SyncEnvClient
|
| 199 |
+
|
| 200 |
+
Synchronous wrapper around EnvClient. Use `client.sync()` to get one:
|
| 201 |
+
|
| 202 |
+
```python
|
| 203 |
+
sync_client = async_client.sync()
|
| 204 |
+
with sync_client:
|
| 205 |
+
result = sync_client.reset()
|
| 206 |
+
result = sync_client.step(action)
|
| 207 |
+
```
|
| 208 |
+
|
| 209 |
+
### HTTPEnvServer
|
| 210 |
+
|
| 211 |
+
Server wrapper with these methods:
|
| 212 |
+
|
| 213 |
+
- `register_routes(app)`: Register endpoints on FastAPI app
|
| 214 |
+
- `_deserialize_action(data)`: Convert JSON to Action
|
| 215 |
+
- `_serialize_observation(obs)`: Convert Observation to JSON
|
| 216 |
+
|
| 217 |
+
### Environment Interface
|
| 218 |
+
|
| 219 |
+
Base interface for environment implementations:
|
| 220 |
+
|
| 221 |
+
- `reset()`: Reset environment and return initial observation
|
| 222 |
+
- `step(action)`: Execute action and return observation
|
| 223 |
+
- `state`: Property returning current environment state
|
| 224 |
+
|
| 225 |
+
## License
|
| 226 |
+
|
| 227 |
+
This project is licensed under the BSD 3-Clause License - see the LICENSE file for details.
|
| 228 |
+
|
| 229 |
+
## Contributing
|
| 230 |
+
|
| 231 |
+
Contributions are welcome! Please see the main OpenEnv repository for contribution guidelines.
|
| 232 |
+
|
| 233 |
+
## Links
|
| 234 |
+
|
| 235 |
+
- **Homepage**: https://github.com/huggingface/OpenEnv
|
| 236 |
+
- **Documentation**: https://github.com/huggingface/OpenEnv/blob/main/README.md
|
| 237 |
+
- **Bug Tracker**: https://github.com/huggingface/OpenEnv/issues
|
src/core/__init__.py
ADDED
|
@@ -0,0 +1,77 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# SPDX-License-Identifier: BSD-3-Clause
|
| 2 |
+
|
| 3 |
+
"""Core components for agentic environments."""
|
| 4 |
+
|
| 5 |
+
from __future__ import annotations
|
| 6 |
+
|
| 7 |
+
from importlib import import_module
|
| 8 |
+
from typing import TYPE_CHECKING
|
| 9 |
+
|
| 10 |
+
from . import env_server
|
| 11 |
+
from .env_server import * # noqa: F403
|
| 12 |
+
|
| 13 |
+
if TYPE_CHECKING:
|
| 14 |
+
from .env_client import EnvClient
|
| 15 |
+
from .generic_client import GenericAction, GenericEnvClient
|
| 16 |
+
from .llm_client import (
|
| 17 |
+
AnthropicClient,
|
| 18 |
+
create_llm_client,
|
| 19 |
+
LLMClient,
|
| 20 |
+
LLMResponse,
|
| 21 |
+
OpenAIClient,
|
| 22 |
+
ToolCall,
|
| 23 |
+
)
|
| 24 |
+
from .mcp_client import MCPClientBase, MCPToolClient
|
| 25 |
+
from .sync_client import SyncEnvClient
|
| 26 |
+
|
| 27 |
+
__all__ = [
|
| 28 |
+
"EnvClient",
|
| 29 |
+
"SyncEnvClient",
|
| 30 |
+
"GenericEnvClient",
|
| 31 |
+
"GenericAction",
|
| 32 |
+
"MCPClientBase",
|
| 33 |
+
"MCPToolClient",
|
| 34 |
+
"AnthropicClient",
|
| 35 |
+
"LLMClient",
|
| 36 |
+
"LLMResponse",
|
| 37 |
+
"OpenAIClient",
|
| 38 |
+
"ToolCall",
|
| 39 |
+
"create_llm_client",
|
| 40 |
+
] + env_server.__all__ # type: ignore
|
| 41 |
+
|
| 42 |
+
|
| 43 |
+
_LAZY_ATTRS = {
|
| 44 |
+
"EnvClient": (".env_client", "EnvClient"),
|
| 45 |
+
"SyncEnvClient": (".sync_client", "SyncEnvClient"),
|
| 46 |
+
"GenericEnvClient": (".generic_client", "GenericEnvClient"),
|
| 47 |
+
"GenericAction": (".generic_client", "GenericAction"),
|
| 48 |
+
"MCPClientBase": (".mcp_client", "MCPClientBase"),
|
| 49 |
+
"MCPToolClient": (".mcp_client", "MCPToolClient"),
|
| 50 |
+
"AnthropicClient": (".llm_client", "AnthropicClient"),
|
| 51 |
+
"LLMClient": (".llm_client", "LLMClient"),
|
| 52 |
+
"LLMResponse": (".llm_client", "LLMResponse"),
|
| 53 |
+
"OpenAIClient": (".llm_client", "OpenAIClient"),
|
| 54 |
+
"ToolCall": (".llm_client", "ToolCall"),
|
| 55 |
+
"create_llm_client": (".llm_client", "create_llm_client"),
|
| 56 |
+
}
|
| 57 |
+
|
| 58 |
+
|
| 59 |
+
def __getattr__(name: str):
|
| 60 |
+
if name in _LAZY_ATTRS:
|
| 61 |
+
module_path, attr_name = _LAZY_ATTRS[name]
|
| 62 |
+
module = import_module(module_path, __name__)
|
| 63 |
+
value = getattr(module, attr_name)
|
| 64 |
+
globals()[name] = value
|
| 65 |
+
return value
|
| 66 |
+
|
| 67 |
+
try:
|
| 68 |
+
value = getattr(env_server, name)
|
| 69 |
+
except AttributeError as exc:
|
| 70 |
+
raise AttributeError(f"module {__name__!r} has no attribute {name!r}") from exc
|
| 71 |
+
|
| 72 |
+
globals()[name] = value
|
| 73 |
+
return value
|
| 74 |
+
|
| 75 |
+
|
| 76 |
+
def __dir__() -> list[str]:
|
| 77 |
+
return sorted(set(globals().keys()) | set(__all__))
|
src/core/client_types.py
ADDED
|
@@ -0,0 +1,29 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# Type definitions for EnvTorch
|
| 2 |
+
from dataclasses import dataclass
|
| 3 |
+
from typing import Any, Dict, Generic, Optional, TypeVar
|
| 4 |
+
|
| 5 |
+
# Generic type for observations
|
| 6 |
+
ObsT = TypeVar("ObsT")
|
| 7 |
+
StateT = TypeVar("StateT")
|
| 8 |
+
|
| 9 |
+
|
| 10 |
+
@dataclass
|
| 11 |
+
class StepResult(Generic[ObsT]):
|
| 12 |
+
"""
|
| 13 |
+
Represents the result of one environment step.
|
| 14 |
+
|
| 15 |
+
Attributes:
|
| 16 |
+
observation:
|
| 17 |
+
The environment's observation after the action.
|
| 18 |
+
reward (`float`, *optional*):
|
| 19 |
+
Scalar reward for this step.
|
| 20 |
+
done (`bool`, *optional*, defaults to `False`):
|
| 21 |
+
Whether the episode is finished.
|
| 22 |
+
metadata (`dict`, *optional*):
|
| 23 |
+
Additional metadata returned alongside the observation.
|
| 24 |
+
"""
|
| 25 |
+
|
| 26 |
+
observation: ObsT
|
| 27 |
+
reward: Optional[float] = None
|
| 28 |
+
done: bool = False
|
| 29 |
+
metadata: Optional[Dict[str, Any]] = None
|
src/core/containers/__init__.py
ADDED
|
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# SPDX-License-Identifier: BSD-3-Clause
|
| 2 |
+
|
| 3 |
+
"""Container management for environment servers."""
|
src/core/containers/images/Dockerfile
ADDED
|
@@ -0,0 +1,60 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# SPDX-License-Identifier: BSD-3-Clause
|
| 2 |
+
|
| 3 |
+
#
|
| 4 |
+
# OpenEnv Base Image
|
| 5 |
+
#
|
| 6 |
+
# This is the standard base image for all OpenEnv environment servers.
|
| 7 |
+
# It includes the minimal dependencies needed to run HTTP environment servers
|
| 8 |
+
# and uv for fast dependency management.
|
| 9 |
+
#
|
| 10 |
+
# Build from repo root: docker build -t openenv-base:latest -f src/openenv/core/containers/images/Dockerfile .
|
| 11 |
+
# Tag: docker tag openenv-base:latest openenv-base:0.2.0
|
| 12 |
+
#
|
| 13 |
+
|
| 14 |
+
FROM ghcr.io/astral-sh/uv:0.5.27-python3.11-bookworm-slim AS builder
|
| 15 |
+
|
| 16 |
+
# Set working directory
|
| 17 |
+
WORKDIR /app
|
| 18 |
+
|
| 19 |
+
# Copy core pyproject.toml and lockfile for dependency installation
|
| 20 |
+
COPY pyproject.toml uv.lock* ./
|
| 21 |
+
|
| 22 |
+
# Install core dependencies using uv with cache mount
|
| 23 |
+
RUN --mount=type=cache,target=/root/.cache/uv \
|
| 24 |
+
uv pip install --system -r pyproject.toml
|
| 25 |
+
|
| 26 |
+
# Final runtime stage
|
| 27 |
+
FROM python:3.11-slim
|
| 28 |
+
|
| 29 |
+
# Set metadata
|
| 30 |
+
LABEL maintainer="OpenEnv Team"
|
| 31 |
+
LABEL description="Base image for OpenEnv based environment servers with uv"
|
| 32 |
+
LABEL version="0.2.0"
|
| 33 |
+
|
| 34 |
+
# Install system dependencies
|
| 35 |
+
RUN apt-get update && apt-get install -y --no-install-recommends \
|
| 36 |
+
curl \
|
| 37 |
+
ca-certificates \
|
| 38 |
+
&& rm -rf /var/lib/apt/lists/*
|
| 39 |
+
|
| 40 |
+
# Copy uv from builder
|
| 41 |
+
COPY --from=builder /usr/local/bin/uv /usr/local/bin/uvx /usr/local/bin/
|
| 42 |
+
|
| 43 |
+
# Copy installed Python packages from builder
|
| 44 |
+
COPY --from=builder /usr/local/lib/python3.11/site-packages /usr/local/lib/python3.11/site-packages
|
| 45 |
+
|
| 46 |
+
# Copy console scripts installed by pip (uvicorn, fastapi, etc.)
|
| 47 |
+
COPY --from=builder /usr/local/bin/uvicorn /usr/local/bin/fastapi /usr/local/bin/
|
| 48 |
+
|
| 49 |
+
# Set working directory
|
| 50 |
+
WORKDIR /app
|
| 51 |
+
|
| 52 |
+
# Default environment variables
|
| 53 |
+
ENV PYTHONPATH=/app/src
|
| 54 |
+
ENV PYTHONUNBUFFERED=1
|
| 55 |
+
ENV UV_SYSTEM_PYTHON=1
|
| 56 |
+
|
| 57 |
+
# Default expose port (can be overridden)
|
| 58 |
+
EXPOSE 8000
|
| 59 |
+
|
| 60 |
+
# Note: CMD should be specified in child Dockerfiles
|
src/core/containers/images/README.md
ADDED
|
@@ -0,0 +1,92 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# OpenEnv Base Image
|
| 2 |
+
|
| 3 |
+
Standard base image for all OpenEnv environment servers.
|
| 4 |
+
|
| 5 |
+
## What's Included
|
| 6 |
+
|
| 7 |
+
| Layer | Size | Contents |
|
| 8 |
+
|-------|------|----------|
|
| 9 |
+
| python:3.11-slim | 200 MB | Base Python runtime |
|
| 10 |
+
| + Dependencies | 100 MB | FastAPI, uvicorn, requests |
|
| 11 |
+
| **Total** | **~300 MB** | Ready for environment servers |
|
| 12 |
+
|
| 13 |
+
## Image Sizes
|
| 14 |
+
|
| 15 |
+
```
|
| 16 |
+
openenv-base:latest 300 MB (python + fastapi + uvicorn)
|
| 17 |
+
```
|
| 18 |
+
echo-env:latest 500 MB (python + fastapi + uvicorn + app)
|
| 19 |
+
coding-env:latest 520 MB (python + fastapi + uvicorn + app + tools)
|
| 20 |
+
another-env:latest 510 MB (python + fastapi + uvicorn + app)
|
| 21 |
+
---
|
| 22 |
+
Total: 1.5 GB (with lots of duplication)
|
| 23 |
+
```
|
| 24 |
+
|
| 25 |
+
### With Base Images (✅ Solution)
|
| 26 |
+
```
|
| 27 |
+
openenv-base:latest 300 MB (python + fastapi + uvicorn)
|
| 28 |
+
echo-env:latest 50 MB (app only, uses base)
|
| 29 |
+
coding-env:latest 70 MB (app + tools, uses base)
|
| 30 |
+
another-env:latest 45 MB (app only, uses base)
|
| 31 |
+
---
|
| 32 |
+
Total: 465 MB (base shared, minimal duplication)
|
| 33 |
+
```
|
| 34 |
+
|
| 35 |
+
## Building the Base Image
|
| 36 |
+
|
| 37 |
+
```bash
|
| 38 |
+
# From project root
|
| 39 |
+
docker build -t openenv-base:latest -f src/openenv/core/containers/images/Dockerfile .
|
| 40 |
+
```
|
| 41 |
+
|
| 42 |
+
## Usage in Environment Dockerfiles
|
| 43 |
+
|
| 44 |
+
Each environment Dockerfile should start with:
|
| 45 |
+
|
| 46 |
+
```dockerfile
|
| 47 |
+
FROM openenv-base:latest
|
| 48 |
+
|
| 49 |
+
# Copy only environment-specific files
|
| 50 |
+
COPY src/openenv/core/ /app/src/openenv/core/
|
| 51 |
+
COPY envs/my_env/ /app/envs/my_env/
|
| 52 |
+
|
| 53 |
+
# Run the server
|
| 54 |
+
CMD ["uvicorn", "envs.my_env.server.app:app", "--host", "0.0.0.0", "--port", "8000"]
|
| 55 |
+
```
|
| 56 |
+
|
| 57 |
+
## Base Image Contents
|
| 58 |
+
|
| 59 |
+
- Python 3.11-slim
|
| 60 |
+
- FastAPI >= 0.104.0
|
| 61 |
+
- Uvicorn >= 0.24.0
|
| 62 |
+
- Requests >= 2.25.0
|
| 63 |
+
- curl (for health checks)
|
| 64 |
+
|
| 65 |
+
## Example: Building Echo Environment
|
| 66 |
+
|
| 67 |
+
```bash
|
| 68 |
+
# Step 1: Build base image (do this once)
|
| 69 |
+
docker build -t openenv-base:latest -f src/openenv/core/containers/images/Dockerfile .
|
| 70 |
+
|
| 71 |
+
# Step 2: Build echo environment (uses base)
|
| 72 |
+
docker build -t echo-env:latest -f envs/echo_env/server/Dockerfile .
|
| 73 |
+
|
| 74 |
+
# Step 3: Run echo environment
|
| 75 |
+
docker run -p 8000:8000 echo-env:latest
|
| 76 |
+
```
|
| 77 |
+
|
| 78 |
+
## Updating the Base
|
| 79 |
+
|
| 80 |
+
When dependencies need updating:
|
| 81 |
+
|
| 82 |
+
1. Update `src/openenv/core/containers/images/Dockerfile`
|
| 83 |
+
2. Rebuild base image
|
| 84 |
+
3. Rebuild all environment images (they'll use new base)
|
| 85 |
+
|
| 86 |
+
```bash
|
| 87 |
+
# Update base
|
| 88 |
+
docker build -t openenv-base:latest -f src/openenv/core/containers/images/Dockerfile .
|
| 89 |
+
|
| 90 |
+
# Rebuild environments (they automatically use new base)
|
| 91 |
+
docker build -t echo-env:latest -f envs/echo_env/server/Dockerfile .
|
| 92 |
+
```
|
src/core/containers/runtime/__init__.py
ADDED
|
@@ -0,0 +1,26 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# SPDX-License-Identifier: BSD-3-Clause
|
| 2 |
+
|
| 3 |
+
"""Container runtime providers."""
|
| 4 |
+
|
| 5 |
+
from .providers import (
|
| 6 |
+
ContainerProvider,
|
| 7 |
+
DockerSwarmProvider,
|
| 8 |
+
KubernetesProvider,
|
| 9 |
+
LocalDockerProvider,
|
| 10 |
+
RuntimeProvider,
|
| 11 |
+
)
|
| 12 |
+
from .uv_provider import UVProvider
|
| 13 |
+
|
| 14 |
+
# Note: optional cloud providers that require extra SDKs (e.g.
|
| 15 |
+
# `ACASandboxProvider`, `DaytonaProvider`) are intentionally NOT re-exported
|
| 16 |
+
# here. Import them from their module, e.g.
|
| 17 |
+
# `from openenv.core.containers.runtime.aca_provider import ACASandboxProvider`.
|
| 18 |
+
|
| 19 |
+
__all__ = [
|
| 20 |
+
"ContainerProvider",
|
| 21 |
+
"DockerSwarmProvider",
|
| 22 |
+
"LocalDockerProvider",
|
| 23 |
+
"KubernetesProvider",
|
| 24 |
+
"RuntimeProvider",
|
| 25 |
+
"UVProvider",
|
| 26 |
+
]
|
src/core/containers/runtime/_server_config.py
ADDED
|
@@ -0,0 +1,53 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# SPDX-License-Identifier: BSD-3-Clause
|
| 2 |
+
|
| 3 |
+
"""Shared helpers for reading environment server configuration."""
|
| 4 |
+
|
| 5 |
+
from __future__ import annotations
|
| 6 |
+
|
| 7 |
+
import json
|
| 8 |
+
import re
|
| 9 |
+
|
| 10 |
+
import yaml
|
| 11 |
+
|
| 12 |
+
|
| 13 |
+
def parse_openenv_app_field(yaml_content: str) -> str | None:
|
| 14 |
+
"""Extract the top-level ``app`` value from raw ``openenv.yaml`` content."""
|
| 15 |
+
try:
|
| 16 |
+
data = yaml.safe_load(yaml_content) or {}
|
| 17 |
+
except Exception:
|
| 18 |
+
return None
|
| 19 |
+
|
| 20 |
+
if not isinstance(data, dict):
|
| 21 |
+
return None
|
| 22 |
+
|
| 23 |
+
value = data.get("app")
|
| 24 |
+
if isinstance(value, str):
|
| 25 |
+
value = value.strip()
|
| 26 |
+
return value if value else None
|
| 27 |
+
return None
|
| 28 |
+
|
| 29 |
+
|
| 30 |
+
def parse_dockerfile_cmd(dockerfile_content: str) -> str | None:
|
| 31 |
+
"""Extract the server command from the last Dockerfile ``CMD``."""
|
| 32 |
+
last_cmd: str | None = None
|
| 33 |
+
for line in dockerfile_content.splitlines():
|
| 34 |
+
stripped = line.strip()
|
| 35 |
+
if stripped.startswith("#"):
|
| 36 |
+
continue
|
| 37 |
+
match = re.match(r"CMD\s+(.+)", stripped, flags=re.IGNORECASE)
|
| 38 |
+
if match:
|
| 39 |
+
last_cmd = match.group(1).strip()
|
| 40 |
+
|
| 41 |
+
if last_cmd is None:
|
| 42 |
+
return None
|
| 43 |
+
|
| 44 |
+
if last_cmd.startswith("["):
|
| 45 |
+
try:
|
| 46 |
+
parts = json.loads(last_cmd)
|
| 47 |
+
if isinstance(parts, list) and all(isinstance(p, str) for p in parts):
|
| 48 |
+
return " ".join(parts)
|
| 49 |
+
except (json.JSONDecodeError, TypeError):
|
| 50 |
+
# Invalid exec-form JSON falls back to the raw CMD below.
|
| 51 |
+
pass
|
| 52 |
+
|
| 53 |
+
return last_cmd if last_cmd else None
|
src/core/containers/runtime/aca_provider.py
ADDED
|
@@ -0,0 +1,690 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# SPDX-License-Identifier: BSD-3-Clause
|
| 2 |
+
|
| 3 |
+
"""Azure Container Apps Sandbox provider for OpenEnv environments.
|
| 4 |
+
|
| 5 |
+
Note: the ``RFC 002 security invariant S<n>`` references in this module point to
|
| 6 |
+
the **proposed** "Cloud Sandbox Providers" amendment to RFC 002, which is pending
|
| 7 |
+
review/sign-off by the RFC authors. They describe the security properties this
|
| 8 |
+
provider enforces; treat them as proposed (not yet ratified) until that sign-off.
|
| 9 |
+
"""
|
| 10 |
+
|
| 11 |
+
from __future__ import annotations
|
| 12 |
+
|
| 13 |
+
import os
|
| 14 |
+
import shlex
|
| 15 |
+
import time
|
| 16 |
+
import warnings
|
| 17 |
+
from typing import Any, Dict, Literal, Optional
|
| 18 |
+
|
| 19 |
+
from .providers import ContainerProvider
|
| 20 |
+
|
| 21 |
+
|
| 22 |
+
_DEFAULT_ACA_PORT = 8000
|
| 23 |
+
# 15 minutes matches the conservative auto-stop default used by other cloud
|
| 24 |
+
# sandbox providers (e.g. Daytona). A live RL rollout holds one WebSocket for
|
| 25 |
+
# the whole episode, so an aggressive auto-suspend can silently drop the
|
| 26 |
+
# session between slow steps. See RFC 002 "Cloud Sandbox Providers".
|
| 27 |
+
_DEFAULT_AUTO_SUSPEND_SECONDS = 900
|
| 28 |
+
_VALID_SUSPEND_MODES = {"Memory", "Disk"}
|
| 29 |
+
|
| 30 |
+
|
| 31 |
+
def _require_secure_url(url: str) -> str:
|
| 32 |
+
"""Enforce https/wss transport (RFC 002 security invariant S1).
|
| 33 |
+
|
| 34 |
+
`EnvClient` derives its WebSocket URL from this base URL, so a plaintext
|
| 35 |
+
URL would become a cleartext `ws://` connection. The offending URL is
|
| 36 |
+
deliberately omitted from the error because an anonymous ACA port URL is a
|
| 37 |
+
bearer capability (S2) that must not leak into logs.
|
| 38 |
+
"""
|
| 39 |
+
if not isinstance(url, str) or not url.lower().startswith("https://"):
|
| 40 |
+
raise RuntimeError(
|
| 41 |
+
"ACA sandbox returned a non-HTTPS port URL. OpenEnv requires an "
|
| 42 |
+
"https/wss base_url so EnvClient traffic is encrypted. Refusing to "
|
| 43 |
+
"connect over plaintext."
|
| 44 |
+
)
|
| 45 |
+
return url
|
| 46 |
+
|
| 47 |
+
|
| 48 |
+
def _validate_anonymous_port(anonymous_port: Optional[bool]) -> bool:
|
| 49 |
+
"""Require explicit, public-only ingress opt-in (RFC 002 invariant S2).
|
| 50 |
+
|
| 51 |
+
Public exposure is never implicit (`None` is rejected), and authenticated
|
| 52 |
+
ACA ports are rejected because `EnvClient` has no WebSocket auth mechanism
|
| 53 |
+
yet.
|
| 54 |
+
"""
|
| 55 |
+
if anonymous_port is None:
|
| 56 |
+
raise ValueError(
|
| 57 |
+
"ACASandboxProvider requires explicit anonymous_port=True. "
|
| 58 |
+
"ACA anonymous ports are public ingress; authenticated ACA "
|
| 59 |
+
"ports are not supported by EnvClient until a client-auth RFC "
|
| 60 |
+
"defines WebSocket credentials."
|
| 61 |
+
)
|
| 62 |
+
if not anonymous_port:
|
| 63 |
+
raise ValueError(
|
| 64 |
+
"authenticated ACA sandbox ports are not supported by EnvClient "
|
| 65 |
+
"yet. Use anonymous_port=True only for intentionally exposed "
|
| 66 |
+
"sandbox URLs, or wait for a client-auth RFC."
|
| 67 |
+
)
|
| 68 |
+
return True
|
| 69 |
+
|
| 70 |
+
|
| 71 |
+
def _resolve_required(value: str | None, env_name: str, parameter: str) -> str:
|
| 72 |
+
resolved = value or os.environ.get(env_name)
|
| 73 |
+
if not resolved:
|
| 74 |
+
raise ValueError(
|
| 75 |
+
f"{parameter} is required. Pass it to ACASandboxProvider or set "
|
| 76 |
+
f"the {env_name} environment variable."
|
| 77 |
+
)
|
| 78 |
+
return resolved
|
| 79 |
+
|
| 80 |
+
|
| 81 |
+
def _raise_install_error(exc: ImportError) -> None:
|
| 82 |
+
raise RuntimeError(
|
| 83 |
+
"Azure Container Apps sandbox support requires optional dependencies. "
|
| 84 |
+
"Install them with `pip install openenv[aca]`."
|
| 85 |
+
) from exc
|
| 86 |
+
|
| 87 |
+
|
| 88 |
+
class _DefaultACASandboxAdapter:
|
| 89 |
+
"""Thin adapter over the preview `azure-containerapps-sandbox` SDK.
|
| 90 |
+
|
| 91 |
+
The provider talks to this private adapter instead of spreading SDK details
|
| 92 |
+
through its own logic; tests inject a duck-typed fake in its place.
|
| 93 |
+
"""
|
| 94 |
+
|
| 95 |
+
def __init__(
|
| 96 |
+
self,
|
| 97 |
+
*,
|
| 98 |
+
subscription_id: str,
|
| 99 |
+
resource_group: str,
|
| 100 |
+
sandbox_group: str,
|
| 101 |
+
region: str | None,
|
| 102 |
+
endpoint: str | None,
|
| 103 |
+
credential: Any,
|
| 104 |
+
sdk_kwargs: dict[str, Any],
|
| 105 |
+
):
|
| 106 |
+
try:
|
| 107 |
+
from azure.containerapps.sandbox import (
|
| 108 |
+
endpoint_for_region,
|
| 109 |
+
SandboxGroupClient,
|
| 110 |
+
)
|
| 111 |
+
except ImportError as exc: # pragma: no cover - exercised via provider tests
|
| 112 |
+
_raise_install_error(exc)
|
| 113 |
+
|
| 114 |
+
if credential is None:
|
| 115 |
+
try:
|
| 116 |
+
from azure.identity import DefaultAzureCredential
|
| 117 |
+
except (
|
| 118 |
+
ImportError
|
| 119 |
+
) as exc: # pragma: no cover - exercised via provider tests
|
| 120 |
+
_raise_install_error(exc)
|
| 121 |
+
|
| 122 |
+
credential = DefaultAzureCredential()
|
| 123 |
+
|
| 124 |
+
if endpoint is None:
|
| 125 |
+
if region is None:
|
| 126 |
+
raise ValueError(
|
| 127 |
+
"region is required when endpoint is not provided. Pass it to "
|
| 128 |
+
"ACASandboxProvider or set AZURE_REGION."
|
| 129 |
+
)
|
| 130 |
+
endpoint = endpoint_for_region(region)
|
| 131 |
+
|
| 132 |
+
self._client = SandboxGroupClient(
|
| 133 |
+
endpoint,
|
| 134 |
+
credential,
|
| 135 |
+
subscription_id=subscription_id,
|
| 136 |
+
resource_group=resource_group,
|
| 137 |
+
sandbox_group=sandbox_group,
|
| 138 |
+
**sdk_kwargs,
|
| 139 |
+
)
|
| 140 |
+
|
| 141 |
+
def create_sandbox(
|
| 142 |
+
self,
|
| 143 |
+
*,
|
| 144 |
+
disk: str | None,
|
| 145 |
+
disk_id: str | None,
|
| 146 |
+
env_vars: dict[str, str] | None,
|
| 147 |
+
labels: dict[str, str] | None,
|
| 148 |
+
egress_policy: Any,
|
| 149 |
+
cpu: str,
|
| 150 |
+
memory: str,
|
| 151 |
+
disk_size: str | None,
|
| 152 |
+
auto_suspend_seconds: int,
|
| 153 |
+
auto_suspend_mode: str,
|
| 154 |
+
create_timeout: int,
|
| 155 |
+
polling_interval: int,
|
| 156 |
+
) -> Any:
|
| 157 |
+
kwargs: dict[str, Any] = {
|
| 158 |
+
"disk": disk,
|
| 159 |
+
"disk_id": disk_id,
|
| 160 |
+
"cpu": cpu,
|
| 161 |
+
"memory": memory,
|
| 162 |
+
"disk_size": disk_size,
|
| 163 |
+
"auto_suspend_seconds": auto_suspend_seconds,
|
| 164 |
+
"auto_suspend_mode": auto_suspend_mode,
|
| 165 |
+
"labels": labels,
|
| 166 |
+
"environment": env_vars,
|
| 167 |
+
"egress_policy": egress_policy,
|
| 168 |
+
"polling_timeout": create_timeout,
|
| 169 |
+
"polling_interval": polling_interval,
|
| 170 |
+
}
|
| 171 |
+
kwargs = {key: value for key, value in kwargs.items() if value is not None}
|
| 172 |
+
# `begin_create_sandbox(...).result()` blocks until the sandbox is
|
| 173 |
+
# *Running* and returns a sandbox-scoped `SandboxClient` (not a plain
|
| 174 |
+
# data model). That client exposes exec / add_port / begin_delete
|
| 175 |
+
# directly, which is the handle the rest of this adapter operates on.
|
| 176 |
+
return self._client.begin_create_sandbox(**kwargs).result()
|
| 177 |
+
|
| 178 |
+
def add_port(self, sandbox: Any, *, port: int, anonymous: bool) -> str:
|
| 179 |
+
sandbox_port = sandbox.add_port(port, anonymous=anonymous)
|
| 180 |
+
url = getattr(sandbox_port, "url", None)
|
| 181 |
+
if not url and isinstance(sandbox_port, dict):
|
| 182 |
+
url = sandbox_port.get("url")
|
| 183 |
+
if not url:
|
| 184 |
+
raise RuntimeError(
|
| 185 |
+
"ACA sandbox did not return a URL for the exposed port. "
|
| 186 |
+
"OpenEnv requires a base_url that EnvClient can connect to directly."
|
| 187 |
+
)
|
| 188 |
+
return str(url)
|
| 189 |
+
|
| 190 |
+
def exec(
|
| 191 |
+
self,
|
| 192 |
+
sandbox: Any,
|
| 193 |
+
command: str,
|
| 194 |
+
*,
|
| 195 |
+
working_directory: str | None = None,
|
| 196 |
+
) -> Any:
|
| 197 |
+
return sandbox.exec(command, working_directory=working_directory)
|
| 198 |
+
|
| 199 |
+
def delete_sandbox(self, sandbox: Any) -> None:
|
| 200 |
+
if hasattr(sandbox, "begin_delete"):
|
| 201 |
+
sandbox.begin_delete().result()
|
| 202 |
+
else:
|
| 203 |
+
sandbox.delete()
|
| 204 |
+
|
| 205 |
+
def close(self) -> None:
|
| 206 |
+
self._client.close()
|
| 207 |
+
|
| 208 |
+
|
| 209 |
+
def _source_from_image(image: str) -> tuple[str | None, str | None]:
|
| 210 |
+
if not image or not image.strip():
|
| 211 |
+
raise ValueError("image must be a non-empty ACA disk or disk-id reference")
|
| 212 |
+
image = image.strip()
|
| 213 |
+
|
| 214 |
+
if image.startswith("disk:"):
|
| 215 |
+
disk = image[len("disk:") :].strip()
|
| 216 |
+
if not disk:
|
| 217 |
+
raise ValueError("disk: source must include a disk image name")
|
| 218 |
+
return disk, None
|
| 219 |
+
if image.startswith("disk-id:"):
|
| 220 |
+
disk_id = image[len("disk-id:") :].strip()
|
| 221 |
+
if not disk_id:
|
| 222 |
+
raise ValueError("disk-id: source must include a disk image id")
|
| 223 |
+
return None, disk_id
|
| 224 |
+
|
| 225 |
+
# A bare string is treated as a public disk-image name. Guard the common
|
| 226 |
+
# migration foot-gun: an ACA sandbox source is NOT a Docker/OCI registry
|
| 227 |
+
# image, so a value that looks like one (a registry path with "/", or an
|
| 228 |
+
# "image:tag") is almost certainly a mistake and would otherwise surface as
|
| 229 |
+
# an opaque ACA SDK error. Fail fast with guidance instead.
|
| 230 |
+
if "/" in image or ":" in image:
|
| 231 |
+
raise ValueError(
|
| 232 |
+
f"{image!r} looks like a container/OCI image reference, but "
|
| 233 |
+
"ACASandboxProvider sources are Azure Container Apps sandbox disks, "
|
| 234 |
+
"not Docker registry images. Use a bare public disk name, "
|
| 235 |
+
"'disk:<name>' for a public disk image, or 'disk-id:<id>' for a "
|
| 236 |
+
"private one (building a disk image from a container image is a "
|
| 237 |
+
"separate, out-of-band ACA step)."
|
| 238 |
+
)
|
| 239 |
+
|
| 240 |
+
return image, None
|
| 241 |
+
|
| 242 |
+
|
| 243 |
+
def _exec_stdout(result: Any) -> str:
|
| 244 |
+
for attr in ("stdout", "result", "output"):
|
| 245 |
+
value = getattr(result, attr, None)
|
| 246 |
+
if value is not None:
|
| 247 |
+
return str(value)
|
| 248 |
+
if isinstance(result, dict):
|
| 249 |
+
for key in ("stdout", "result", "output"):
|
| 250 |
+
value = result.get(key)
|
| 251 |
+
if value is not None:
|
| 252 |
+
return str(value)
|
| 253 |
+
return str(result)
|
| 254 |
+
|
| 255 |
+
|
| 256 |
+
class ACASandboxProvider(ContainerProvider):
|
| 257 |
+
"""Container provider backed by Azure Container Apps Sandboxes.
|
| 258 |
+
|
| 259 |
+
`start_container`'s `image` is an ACA sandbox source (`disk:<name>` for a
|
| 260 |
+
public disk image or `disk-id:<id>` for a private one), not a Docker
|
| 261 |
+
registry image. The provider boots an OpenEnv server inside the sandbox,
|
| 262 |
+
exposes it on an anonymous ACA port, and returns an `https://` URL that
|
| 263 |
+
`EnvClient` connects to over `wss://`.
|
| 264 |
+
|
| 265 |
+
The environment runs untrusted code, so the provider is secure by default
|
| 266 |
+
(proposed RFC 002 security invariants): it requires explicit `anonymous_port=True`
|
| 267 |
+
ingress (S2), enforces https/wss transport (S1), offers
|
| 268 |
+
`deny_all_egress()` to block the cloud metadata/IMDS endpoint (S3), and
|
| 269 |
+
never surfaces raw sandbox output unless `surface_server_logs=True` (S4).
|
| 270 |
+
|
| 271 |
+
Validated end to end against a live ACA sandbox group; the underlying
|
| 272 |
+
`azure-containerapps-sandbox` SDK is preview, so pin it and re-validate
|
| 273 |
+
after upgrades (see `tests/test_core/test_aca_provider_integration.py`).
|
| 274 |
+
|
| 275 |
+
`close()` stops the active sandbox *and* releases the underlying SDK client,
|
| 276 |
+
so prefer using the provider as a context manager (or call `close()`
|
| 277 |
+
explicitly) for deterministic cleanup. `close()`/`__enter__`/`__exit__` are
|
| 278 |
+
defined on `ContainerProvider` (default no-op), so a caller holding a bare
|
| 279 |
+
`ContainerProvider` reference can release the client polymorphically.
|
| 280 |
+
|
| 281 |
+
Only one sandbox is active per provider: calling `start_container` again
|
| 282 |
+
before `stop_container()`/`close()` raises `RuntimeError` rather than
|
| 283 |
+
orphaning the running sandbox.
|
| 284 |
+
|
| 285 |
+
```python
|
| 286 |
+
with ACASandboxProvider(
|
| 287 |
+
image="disk:my-env",
|
| 288 |
+
anonymous_port=True,
|
| 289 |
+
...,
|
| 290 |
+
) as provider:
|
| 291 |
+
base_url = provider.start_container(cmd=...)
|
| 292 |
+
...
|
| 293 |
+
# sandbox deleted and SDK client closed on exit
|
| 294 |
+
```
|
| 295 |
+
"""
|
| 296 |
+
|
| 297 |
+
def __init__(
|
| 298 |
+
self,
|
| 299 |
+
*,
|
| 300 |
+
subscription_id: Optional[str] = None,
|
| 301 |
+
resource_group: Optional[str] = None,
|
| 302 |
+
sandbox_group: Optional[str] = None,
|
| 303 |
+
region: Optional[str] = None,
|
| 304 |
+
endpoint: Optional[str] = None,
|
| 305 |
+
credential: Any = None,
|
| 306 |
+
cpu: str = "1000m",
|
| 307 |
+
memory: str = "2048Mi",
|
| 308 |
+
disk_size: Optional[str] = None,
|
| 309 |
+
auto_suspend_seconds: int = _DEFAULT_AUTO_SUSPEND_SECONDS,
|
| 310 |
+
auto_suspend_mode: Literal["Memory", "Disk"] = "Memory",
|
| 311 |
+
labels: Optional[Dict[str, str]] = None,
|
| 312 |
+
egress_policy: Any = None,
|
| 313 |
+
anonymous_port: Literal[True] | None = None,
|
| 314 |
+
image: Optional[str] = None,
|
| 315 |
+
env_vars: Optional[Dict[str, str]] = None,
|
| 316 |
+
cmd: Optional[str] = None,
|
| 317 |
+
working_directory: Optional[str] = None,
|
| 318 |
+
surface_server_logs: bool = False,
|
| 319 |
+
create_timeout: int = 300,
|
| 320 |
+
polling_interval: int = 3,
|
| 321 |
+
sdk_kwargs: Optional[dict[str, Any]] = None,
|
| 322 |
+
_adapter: Any = None,
|
| 323 |
+
):
|
| 324 |
+
if auto_suspend_mode not in _VALID_SUSPEND_MODES:
|
| 325 |
+
raise ValueError("auto_suspend_mode must be 'Memory' or 'Disk'")
|
| 326 |
+
_validate_anonymous_port(anonymous_port)
|
| 327 |
+
|
| 328 |
+
self.subscription_id = subscription_id or os.environ.get(
|
| 329 |
+
"AZURE_SUBSCRIPTION_ID"
|
| 330 |
+
)
|
| 331 |
+
self.resource_group = resource_group or os.environ.get("AZURE_RESOURCE_GROUP")
|
| 332 |
+
self.sandbox_group = sandbox_group or os.environ.get("AZURE_SANDBOX_GROUP")
|
| 333 |
+
self.region = region or os.environ.get("AZURE_REGION")
|
| 334 |
+
self.endpoint = endpoint
|
| 335 |
+
|
| 336 |
+
self.cpu = cpu
|
| 337 |
+
self.memory = memory
|
| 338 |
+
self.disk_size = disk_size
|
| 339 |
+
self.auto_suspend_seconds = auto_suspend_seconds
|
| 340 |
+
self.auto_suspend_mode = auto_suspend_mode
|
| 341 |
+
self.labels = dict(labels or {})
|
| 342 |
+
self.egress_policy = egress_policy
|
| 343 |
+
self.anonymous_port = anonymous_port
|
| 344 |
+
self._image = image
|
| 345 |
+
self._env_vars = dict(env_vars or {}) if env_vars is not None else None
|
| 346 |
+
self.cmd = cmd
|
| 347 |
+
self.working_directory = working_directory
|
| 348 |
+
self.surface_server_logs = surface_server_logs
|
| 349 |
+
self.create_timeout = create_timeout
|
| 350 |
+
self.polling_interval = polling_interval
|
| 351 |
+
|
| 352 |
+
if _adapter is None:
|
| 353 |
+
self._adapter = _DefaultACASandboxAdapter(
|
| 354 |
+
subscription_id=_resolve_required(
|
| 355 |
+
self.subscription_id, "AZURE_SUBSCRIPTION_ID", "subscription_id"
|
| 356 |
+
),
|
| 357 |
+
resource_group=_resolve_required(
|
| 358 |
+
self.resource_group, "AZURE_RESOURCE_GROUP", "resource_group"
|
| 359 |
+
),
|
| 360 |
+
sandbox_group=_resolve_required(
|
| 361 |
+
self.sandbox_group, "AZURE_SANDBOX_GROUP", "sandbox_group"
|
| 362 |
+
),
|
| 363 |
+
region=self.region,
|
| 364 |
+
endpoint=self.endpoint,
|
| 365 |
+
credential=credential,
|
| 366 |
+
sdk_kwargs=dict(sdk_kwargs or {}),
|
| 367 |
+
)
|
| 368 |
+
self._owns_adapter = True
|
| 369 |
+
else:
|
| 370 |
+
self._adapter = _adapter
|
| 371 |
+
self._owns_adapter = False
|
| 372 |
+
|
| 373 |
+
self._sandbox: Any = None
|
| 374 |
+
self._base_url: Optional[str] = None
|
| 375 |
+
self._port = _DEFAULT_ACA_PORT
|
| 376 |
+
self._started_server = False
|
| 377 |
+
# Injected env-var values, used to scrub captured server output before
|
| 378 |
+
# it is ever surfaced in an error (RFC 002 security invariant S4).
|
| 379 |
+
self._redact_values: set[str] = set()
|
| 380 |
+
|
| 381 |
+
@staticmethod
|
| 382 |
+
def deny_all_egress(allow: Optional[list[str]] = None) -> Any:
|
| 383 |
+
"""Build a default-deny ACA `EgressPolicy` (RFC 002 security invariant S3).
|
| 384 |
+
|
| 385 |
+
Untrusted RL/agent code should not be able to exfiltrate data or reach
|
| 386 |
+
the cloud metadata/IMDS endpoint. This returns an `EgressPolicy` whose
|
| 387 |
+
default action is `Deny`, with an optional host allowlist (for example a
|
| 388 |
+
model endpoint or package registry the environment legitimately needs).
|
| 389 |
+
|
| 390 |
+
```python
|
| 391 |
+
provider = ACASandboxProvider(
|
| 392 |
+
anonymous_port=True,
|
| 393 |
+
egress_policy=ACASandboxProvider.deny_all_egress(
|
| 394 |
+
allow=["my-model.openai.azure.com"]
|
| 395 |
+
),
|
| 396 |
+
)
|
| 397 |
+
```
|
| 398 |
+
"""
|
| 399 |
+
try:
|
| 400 |
+
from azure.containerapps.sandbox import EgressHostRule, EgressPolicy
|
| 401 |
+
except ImportError as exc: # pragma: no cover - exercised via provider tests
|
| 402 |
+
_raise_install_error(exc)
|
| 403 |
+
|
| 404 |
+
host_rules = [
|
| 405 |
+
EgressHostRule(pattern=pattern, action="Allow") for pattern in (allow or [])
|
| 406 |
+
]
|
| 407 |
+
return EgressPolicy(default_action="Deny", host_rules=host_rules)
|
| 408 |
+
|
| 409 |
+
def start_container(
|
| 410 |
+
self,
|
| 411 |
+
image: Optional[str] = None,
|
| 412 |
+
port: Optional[int] = None,
|
| 413 |
+
env_vars: Optional[Dict[str, str]] = None,
|
| 414 |
+
**kwargs: Any,
|
| 415 |
+
) -> str:
|
| 416 |
+
"""Start an OpenEnv server in an ACA sandbox and return its base URL.
|
| 417 |
+
|
| 418 |
+
`image` is an ACA sandbox *source*, not a Docker/OCI registry image: a
|
| 419 |
+
bare string or `disk:<name>` for a public disk image, or `disk-id:<id>`
|
| 420 |
+
for a private one. It may be omitted when supplied to the constructor.
|
| 421 |
+
A value that looks like a container image (a registry path or
|
| 422 |
+
`name:tag`) is rejected with guidance. Only port 8000 is supported.
|
| 423 |
+
`**kwargs` accepts per-start overrides (`cmd`, `labels`,
|
| 424 |
+
`egress_policy`, `anonymous_port`); unknown options raise `ValueError`
|
| 425 |
+
so typos cannot silently change sandbox behavior.
|
| 426 |
+
"""
|
| 427 |
+
if self._sandbox is not None:
|
| 428 |
+
raise RuntimeError(
|
| 429 |
+
"ACASandboxProvider already has an active sandbox. Call "
|
| 430 |
+
"stop_container() (or close()) before starting another — a "
|
| 431 |
+
"second start would orphan the running sandbox."
|
| 432 |
+
)
|
| 433 |
+
|
| 434 |
+
bind_port = port or _DEFAULT_ACA_PORT
|
| 435 |
+
if bind_port != _DEFAULT_ACA_PORT:
|
| 436 |
+
raise ValueError(
|
| 437 |
+
f"ACASandboxProvider only supports port {_DEFAULT_ACA_PORT} "
|
| 438 |
+
f"(got {bind_port})."
|
| 439 |
+
)
|
| 440 |
+
|
| 441 |
+
effective_image = image if image is not None else self._image
|
| 442 |
+
if effective_image is None:
|
| 443 |
+
raise ValueError(
|
| 444 |
+
"ACASandboxProvider requires an image. Pass it to the constructor "
|
| 445 |
+
"or start_container()."
|
| 446 |
+
)
|
| 447 |
+
effective_env_vars = self._env_vars if env_vars is None else env_vars
|
| 448 |
+
|
| 449 |
+
disk, disk_id = _source_from_image(effective_image)
|
| 450 |
+
labels = dict(self.labels)
|
| 451 |
+
labels.update(kwargs.pop("labels", {}) or {})
|
| 452 |
+
cmd = kwargs.pop("cmd", None) or self.cmd
|
| 453 |
+
egress_policy = kwargs.pop("egress_policy", self.egress_policy)
|
| 454 |
+
anonymous_port = kwargs.pop("anonymous_port", self.anonymous_port)
|
| 455 |
+
# Re-validate in case anonymous_port was overridden at start time: public
|
| 456 |
+
# exposure must stay an explicit, opt-in decision (S2).
|
| 457 |
+
_validate_anonymous_port(anonymous_port)
|
| 458 |
+
|
| 459 |
+
if kwargs:
|
| 460 |
+
unknown = ", ".join(sorted(kwargs))
|
| 461 |
+
raise ValueError(f"Unsupported ACASandboxProvider start options: {unknown}")
|
| 462 |
+
|
| 463 |
+
# Security (RFC 002 S3): the env server runs untrusted code and ACA's
|
| 464 |
+
# EgressPolicy defaults to default_action="Allow", so an unset policy
|
| 465 |
+
# means unrestricted egress — the workload could exfiltrate data or
|
| 466 |
+
# steal the sandbox managed-identity token from the cloud metadata
|
| 467 |
+
# endpoint. Warn loudly; do not silently allow it.
|
| 468 |
+
if egress_policy is None:
|
| 469 |
+
warnings.warn(
|
| 470 |
+
"ACASandboxProvider is starting a sandbox with unrestricted "
|
| 471 |
+
"egress (no egress_policy). The environment runs untrusted code; "
|
| 472 |
+
"prefer a default-deny policy via "
|
| 473 |
+
"ACASandboxProvider.deny_all_egress(allow=[...]) and ensure the "
|
| 474 |
+
"cloud metadata/IMDS endpoint is blocked.",
|
| 475 |
+
stacklevel=2,
|
| 476 |
+
)
|
| 477 |
+
|
| 478 |
+
# Record injected secret values so captured server output can be scrubbed
|
| 479 |
+
# before it is ever surfaced in an error (security invariant S4).
|
| 480 |
+
self._redact_values = {
|
| 481 |
+
value for value in (effective_env_vars or {}).values() if value
|
| 482 |
+
}
|
| 483 |
+
|
| 484 |
+
# Two-stage cleanup. A create failure created nothing, so just drop the
|
| 485 |
+
# recorded secrets and re-raise (notably it must NOT delete a sandbox
|
| 486 |
+
# this call did not create — e.g. one left over on the provider). Only a
|
| 487 |
+
# failure AFTER the sandbox exists runs stop_container() to delete it.
|
| 488 |
+
try:
|
| 489 |
+
self._sandbox = self._adapter.create_sandbox(
|
| 490 |
+
disk=disk,
|
| 491 |
+
disk_id=disk_id,
|
| 492 |
+
env_vars=dict(effective_env_vars or {}) if effective_env_vars else None,
|
| 493 |
+
labels=labels or None,
|
| 494 |
+
egress_policy=egress_policy,
|
| 495 |
+
cpu=self.cpu,
|
| 496 |
+
memory=self.memory,
|
| 497 |
+
disk_size=self.disk_size,
|
| 498 |
+
auto_suspend_seconds=self.auto_suspend_seconds,
|
| 499 |
+
auto_suspend_mode=self.auto_suspend_mode,
|
| 500 |
+
create_timeout=self.create_timeout,
|
| 501 |
+
polling_interval=self.polling_interval,
|
| 502 |
+
)
|
| 503 |
+
except Exception:
|
| 504 |
+
self._redact_values = set()
|
| 505 |
+
raise
|
| 506 |
+
|
| 507 |
+
try:
|
| 508 |
+
if cmd:
|
| 509 |
+
self._start_server(cmd)
|
| 510 |
+
self._base_url = _require_secure_url(
|
| 511 |
+
self._adapter.add_port(
|
| 512 |
+
self._sandbox,
|
| 513 |
+
port=bind_port,
|
| 514 |
+
anonymous=anonymous_port,
|
| 515 |
+
)
|
| 516 |
+
)
|
| 517 |
+
self._port = bind_port
|
| 518 |
+
except Exception:
|
| 519 |
+
# A cleanup failure here must not mask the original error: swallow
|
| 520 |
+
# any exception from stop_container() so the root cause propagates.
|
| 521 |
+
try:
|
| 522 |
+
self.stop_container()
|
| 523 |
+
except Exception:
|
| 524 |
+
pass
|
| 525 |
+
raise
|
| 526 |
+
|
| 527 |
+
return self._base_url
|
| 528 |
+
|
| 529 |
+
def _start_server(self, cmd: str) -> None:
|
| 530 |
+
# `cmd` is trusted orchestrator configuration, not agent/environment
|
| 531 |
+
# input — callers must not pass agent-controlled text as `cmd`, and any
|
| 532 |
+
# dynamic value interpolated into it must be quoted by the caller. Given
|
| 533 |
+
# that, running it through the subshell below is safe (RFC 002 amendment
|
| 534 |
+
# S5 — no command injection from untrusted sources). `working_directory`
|
| 535 |
+
# is still shlex-quoted defensively.
|
| 536 |
+
if self.working_directory:
|
| 537 |
+
command = f"cd {shlex.quote(self.working_directory)} && {cmd}"
|
| 538 |
+
else:
|
| 539 |
+
command = cmd
|
| 540 |
+
escaped = shlex.quote(command)
|
| 541 |
+
# The adapter's `exec` maps to the ACA SDK's `executeShellCommand`
|
| 542 |
+
# endpoint, which runs its argument through a POSIX shell in the sandbox
|
| 543 |
+
# (the SDK docstring says "Execute a shell command" and tells callers to
|
| 544 |
+
# `shlex.quote()` interpolated input). That shell is what makes the
|
| 545 |
+
# backgrounding (`&`), output redirection, and `$!` PID capture below
|
| 546 |
+
# work; if a future preview SDK changes `exec` to bypass the shell, this
|
| 547 |
+
# start sequence must be revisited (re-validate via the integration test).
|
| 548 |
+
self._adapter.exec(
|
| 549 |
+
self._sandbox,
|
| 550 |
+
f"nohup bash -c {escaped} > /tmp/openenv-server.log 2>&1 & "
|
| 551 |
+
"echo $! > /tmp/openenv-server.pid",
|
| 552 |
+
)
|
| 553 |
+
self._started_server = True
|
| 554 |
+
|
| 555 |
+
def _redact(self, text: str, *, max_chars: int = 2000) -> str:
|
| 556 |
+
"""Scrub injected secret values and bound length before surfacing output.
|
| 557 |
+
|
| 558 |
+
Replaces any injected env-var value with `***` and keeps only the tail.
|
| 559 |
+
This is best-effort (exact-match only), which is why server output is
|
| 560 |
+
withheld entirely unless `surface_server_logs=True` (RFC 002 S4).
|
| 561 |
+
"""
|
| 562 |
+
redacted = text or ""
|
| 563 |
+
for value in self._redact_values:
|
| 564 |
+
redacted = redacted.replace(value, "***")
|
| 565 |
+
if len(redacted) > max_chars:
|
| 566 |
+
redacted = "...(truncated)...\n" + redacted[-max_chars:]
|
| 567 |
+
return redacted
|
| 568 |
+
|
| 569 |
+
def _server_died_message(self) -> str:
|
| 570 |
+
"""Build the startup-crash error, secure by default (RFC 002 S4).
|
| 571 |
+
|
| 572 |
+
Untrusted code can print secrets then force a crash to exfiltrate them
|
| 573 |
+
through the exception (which lands in orchestrator/CI logs), so sandbox
|
| 574 |
+
output is excluded unless `surface_server_logs=True`, in which case a
|
| 575 |
+
best-effort redacted, bounded excerpt is included.
|
| 576 |
+
"""
|
| 577 |
+
base = (
|
| 578 |
+
"ACA sandbox server process died during startup. Server output is "
|
| 579 |
+
"not surfaced to avoid leaking secrets injected into the sandbox; "
|
| 580 |
+
"retrieve /tmp/openenv-server.log from the sandbox out of band, or "
|
| 581 |
+
"construct the provider with surface_server_logs=True to include a "
|
| 582 |
+
"redacted excerpt."
|
| 583 |
+
)
|
| 584 |
+
if not self.surface_server_logs:
|
| 585 |
+
return base
|
| 586 |
+
if self._sandbox is None:
|
| 587 |
+
return base
|
| 588 |
+
log = self._redact(
|
| 589 |
+
_exec_stdout(
|
| 590 |
+
self._adapter.exec(
|
| 591 |
+
self._sandbox,
|
| 592 |
+
"cat /tmp/openenv-server.log 2>/dev/null",
|
| 593 |
+
)
|
| 594 |
+
)
|
| 595 |
+
)
|
| 596 |
+
return (
|
| 597 |
+
"ACA sandbox server process died during startup. The excerpt below "
|
| 598 |
+
"is the sandbox server output with injected secret values redacted "
|
| 599 |
+
"(best-effort); it may still contain secrets the workload printed "
|
| 600 |
+
f"by other means.\nLog (redacted):\n{log}"
|
| 601 |
+
)
|
| 602 |
+
|
| 603 |
+
def wait_for_ready(self, base_url: str, timeout_s: float = 120.0) -> None:
|
| 604 |
+
"""Wait for the ACA-hosted OpenEnv server to answer on `/health`.
|
| 605 |
+
|
| 606 |
+
A `200` on `/health` proves HTTP reachability but not that the exposed
|
| 607 |
+
ACA port proxies the `/ws` WebSocket upgrade `EnvClient` needs; the
|
| 608 |
+
integration test covers the full `wss://` round-trip. If the captured
|
| 609 |
+
server process dies during startup this raises `RuntimeError` (sandbox
|
| 610 |
+
output is withheld unless `surface_server_logs=True`; see RFC 002 S4).
|
| 611 |
+
"""
|
| 612 |
+
# Imported lazily inside the method that needs it, matching the other
|
| 613 |
+
# providers in this package (e.g. DaytonaProvider).
|
| 614 |
+
import requests
|
| 615 |
+
|
| 616 |
+
deadline = time.time() + timeout_s
|
| 617 |
+
health_url = f"{base_url}/health"
|
| 618 |
+
|
| 619 |
+
while time.time() < deadline:
|
| 620 |
+
try:
|
| 621 |
+
response = requests.get(health_url, timeout=5.0)
|
| 622 |
+
if response.status_code == 200:
|
| 623 |
+
return
|
| 624 |
+
except requests.RequestException:
|
| 625 |
+
pass
|
| 626 |
+
|
| 627 |
+
if self._started_server and self._sandbox is not None:
|
| 628 |
+
status = _exec_stdout(
|
| 629 |
+
self._adapter.exec(
|
| 630 |
+
self._sandbox,
|
| 631 |
+
"kill -0 $(cat /tmp/openenv-server.pid) 2>/dev/null "
|
| 632 |
+
"&& echo RUNNING || echo DEAD",
|
| 633 |
+
)
|
| 634 |
+
)
|
| 635 |
+
if "DEAD" in status:
|
| 636 |
+
raise RuntimeError(self._server_died_message())
|
| 637 |
+
|
| 638 |
+
time.sleep(1.0)
|
| 639 |
+
|
| 640 |
+
raise TimeoutError(
|
| 641 |
+
f"ACA sandbox did not become ready within {timeout_s}s. "
|
| 642 |
+
"If the sandbox image starts the server itself, pass cmd= to let "
|
| 643 |
+
"ACASandboxProvider capture /tmp/openenv-server.log for early crash "
|
| 644 |
+
"diagnostics."
|
| 645 |
+
)
|
| 646 |
+
|
| 647 |
+
def stop_container(self) -> None:
|
| 648 |
+
"""Delete the active ACA sandbox."""
|
| 649 |
+
if self._sandbox is None:
|
| 650 |
+
# Still drop any injected secret values recorded by a failed start
|
| 651 |
+
# (e.g. create_sandbox raised before a sandbox handle existed).
|
| 652 |
+
self._redact_values = set()
|
| 653 |
+
return
|
| 654 |
+
|
| 655 |
+
try:
|
| 656 |
+
self._adapter.delete_sandbox(self._sandbox)
|
| 657 |
+
finally:
|
| 658 |
+
self._sandbox = None
|
| 659 |
+
self._base_url = None
|
| 660 |
+
self._started_server = False
|
| 661 |
+
# Drop the previous episode's injected secret values; a new start
|
| 662 |
+
# repopulates them (avoids carrying stale values across sandboxes).
|
| 663 |
+
self._redact_values = set()
|
| 664 |
+
|
| 665 |
+
@property
|
| 666 |
+
def base_url(self) -> str:
|
| 667 |
+
"""URL returned by the last `start_container`."""
|
| 668 |
+
if self._base_url is None:
|
| 669 |
+
raise RuntimeError(
|
| 670 |
+
"ACASandboxProvider has no active base_url. Start the provider "
|
| 671 |
+
"before reading base_url."
|
| 672 |
+
)
|
| 673 |
+
return self._base_url
|
| 674 |
+
|
| 675 |
+
def close(self) -> None:
|
| 676 |
+
"""Stop the active sandbox and close the underlying SDK client.
|
| 677 |
+
|
| 678 |
+
Overrides the base no-op so a caller can release the preview SDK client
|
| 679 |
+
deterministically (also invoked on context-manager exit). The base
|
| 680 |
+
`ContainerProvider` defines `close()`/`__enter__`/`__exit__`, so callers
|
| 681 |
+
holding a bare `ContainerProvider` can release resources polymorphically.
|
| 682 |
+
"""
|
| 683 |
+
self.stop_container()
|
| 684 |
+
if self._owns_adapter:
|
| 685 |
+
self._adapter.close()
|
| 686 |
+
# don't close a second time on a subsequent close()/context exit
|
| 687 |
+
self._owns_adapter = False
|
| 688 |
+
|
| 689 |
+
|
| 690 |
+
__all__ = ["ACASandboxProvider"]
|
src/core/containers/runtime/daytona_provider.py
ADDED
|
@@ -0,0 +1,547 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# SPDX-License-Identifier: BSD-3-Clause
|
| 2 |
+
|
| 3 |
+
"""
|
| 4 |
+
Daytona container provider for running OpenEnv environments in Daytona cloud sandboxes.
|
| 5 |
+
|
| 6 |
+
Requires the ``daytona`` SDK: ``pip install daytona>=0.10``
|
| 7 |
+
"""
|
| 8 |
+
|
| 9 |
+
from __future__ import annotations
|
| 10 |
+
|
| 11 |
+
import os
|
| 12 |
+
import shlex
|
| 13 |
+
import time
|
| 14 |
+
from typing import Any, Callable, Dict, Optional
|
| 15 |
+
|
| 16 |
+
from ._server_config import parse_dockerfile_cmd, parse_openenv_app_field
|
| 17 |
+
from .providers import ContainerProvider
|
| 18 |
+
|
| 19 |
+
|
| 20 |
+
class DaytonaProvider(ContainerProvider):
|
| 21 |
+
"""
|
| 22 |
+
Container provider that runs environments in Daytona cloud sandboxes.
|
| 23 |
+
|
| 24 |
+
Example:
|
| 25 |
+
>>> provider = DaytonaProvider(api_key="your-key")
|
| 26 |
+
>>> image = DaytonaProvider.image_from_dockerfile("envs/echo_env/server/Dockerfile")
|
| 27 |
+
>>> base_url = provider.start_container(image)
|
| 28 |
+
>>> provider.wait_for_ready(base_url)
|
| 29 |
+
>>> provider.stop_container()
|
| 30 |
+
"""
|
| 31 |
+
|
| 32 |
+
_dockerfile_registry: Dict[str, Dict[str, Any]] = {}
|
| 33 |
+
|
| 34 |
+
def __init__(
|
| 35 |
+
self,
|
| 36 |
+
*,
|
| 37 |
+
image: Optional[str] = None,
|
| 38 |
+
env_vars: Optional[Dict[str, str]] = None,
|
| 39 |
+
api_key: Optional[str] = None,
|
| 40 |
+
public: bool = False,
|
| 41 |
+
resources: Optional[Any] = None,
|
| 42 |
+
auto_stop_interval: int = 15,
|
| 43 |
+
target: Optional[str] = None,
|
| 44 |
+
on_snapshot_create_logs: Optional[Callable[[str], None]] = None,
|
| 45 |
+
cmd: Optional[str] = None,
|
| 46 |
+
create_timeout: float = 300,
|
| 47 |
+
):
|
| 48 |
+
"""
|
| 49 |
+
Args:
|
| 50 |
+
image: Docker image, snapshot reference, or ``"dockerfile:<path>"``
|
| 51 |
+
source to use when ``start_container()`` is called without an
|
| 52 |
+
image.
|
| 53 |
+
env_vars: Environment variables to use when ``start_container()`` is
|
| 54 |
+
called without explicit ``env_vars``.
|
| 55 |
+
api_key: Daytona API key. Falls back to ``DAYTONA_API_KEY`` env var.
|
| 56 |
+
public: If True, the sandbox preview is publicly accessible.
|
| 57 |
+
resources: Optional ``daytona.Resources`` instance for CPU/memory.
|
| 58 |
+
auto_stop_interval: Minutes of inactivity before auto-stop (0 disables).
|
| 59 |
+
target: Daytona target region (e.g. "us").
|
| 60 |
+
on_snapshot_create_logs: Callback for snapshot build log lines.
|
| 61 |
+
cmd: Shell command to start the server inside the sandbox.
|
| 62 |
+
create_timeout: Seconds to wait for sandbox creation (default 300).
|
| 63 |
+
Heavy images (e.g. with Playwright/Chromium) may need more.
|
| 64 |
+
"""
|
| 65 |
+
from daytona import Daytona, DaytonaConfig
|
| 66 |
+
|
| 67 |
+
config_kwargs: Dict[str, Any] = {}
|
| 68 |
+
resolved_key = api_key or os.environ.get("DAYTONA_API_KEY")
|
| 69 |
+
if resolved_key:
|
| 70 |
+
config_kwargs["api_key"] = resolved_key
|
| 71 |
+
if target:
|
| 72 |
+
config_kwargs["target"] = target
|
| 73 |
+
|
| 74 |
+
self._daytona = Daytona(DaytonaConfig(**config_kwargs))
|
| 75 |
+
self._image = image
|
| 76 |
+
self._env_vars = env_vars
|
| 77 |
+
self._public = public
|
| 78 |
+
self._resources = resources
|
| 79 |
+
self._auto_stop_interval = auto_stop_interval
|
| 80 |
+
self._on_snapshot_create_logs = on_snapshot_create_logs
|
| 81 |
+
self._cmd = cmd
|
| 82 |
+
self._create_timeout = create_timeout
|
| 83 |
+
self._sandbox: Any = None
|
| 84 |
+
self._preview_url: Optional[str] = None
|
| 85 |
+
|
| 86 |
+
def _discover_server_cmd(self, sandbox: Any, port: int = 8000) -> str:
|
| 87 |
+
"""Discover the server command from ``openenv.yaml`` inside *sandbox*.
|
| 88 |
+
|
| 89 |
+
Finds the file, reads the ``app`` field, and constructs a command
|
| 90 |
+
of the form ``cd <env_root> && python -m uvicorn <app> --host 0.0.0.0 --port <port>``.
|
| 91 |
+
|
| 92 |
+
Raises:
|
| 93 |
+
ValueError: If ``openenv.yaml`` is not found or lacks an ``app`` field.
|
| 94 |
+
"""
|
| 95 |
+
yaml_path = self._find_openenv_yaml(sandbox)
|
| 96 |
+
if yaml_path is None:
|
| 97 |
+
raise ValueError(
|
| 98 |
+
"Could not find openenv.yaml inside the sandbox. "
|
| 99 |
+
"Pass an explicit cmd= to DaytonaProvider or start_container()."
|
| 100 |
+
)
|
| 101 |
+
|
| 102 |
+
cat_resp = sandbox.process.exec(f"cat {shlex.quote(yaml_path)}", timeout=10)
|
| 103 |
+
content = cat_resp.result if hasattr(cat_resp, "result") else str(cat_resp)
|
| 104 |
+
app = self._parse_app_field(content)
|
| 105 |
+
if app is None:
|
| 106 |
+
raise ValueError(
|
| 107 |
+
f"openenv.yaml at {yaml_path} does not contain an 'app' field. "
|
| 108 |
+
"Pass an explicit cmd= to DaytonaProvider or start_container()."
|
| 109 |
+
)
|
| 110 |
+
|
| 111 |
+
# The directory containing openenv.yaml is the env root
|
| 112 |
+
env_root = yaml_path.rsplit("/", 1)[0]
|
| 113 |
+
return (
|
| 114 |
+
f"cd {shlex.quote(env_root)} && "
|
| 115 |
+
f"python -m uvicorn {shlex.quote(app)} --host 0.0.0.0 --port {port}"
|
| 116 |
+
)
|
| 117 |
+
|
| 118 |
+
def _find_openenv_yaml(self, sandbox: Any) -> Optional[str]:
|
| 119 |
+
"""Locate ``openenv.yaml`` inside the sandbox.
|
| 120 |
+
|
| 121 |
+
Tries the modern layout path ``/app/env/openenv.yaml`` first,
|
| 122 |
+
then falls back to a ``find`` command for the old layout.
|
| 123 |
+
"""
|
| 124 |
+
# Fast path: modern Dockerfile layout
|
| 125 |
+
resp = sandbox.process.exec(
|
| 126 |
+
"test -f /app/env/openenv.yaml && echo found", timeout=10
|
| 127 |
+
)
|
| 128 |
+
out = resp.result if hasattr(resp, "result") else str(resp)
|
| 129 |
+
if "found" in (out or ""):
|
| 130 |
+
return "/app/env/openenv.yaml"
|
| 131 |
+
|
| 132 |
+
# Fallback: search for it (redirect stderr so error messages
|
| 133 |
+
# like "No such file or directory" don't get mistaken for paths).
|
| 134 |
+
resp = sandbox.process.exec(
|
| 135 |
+
"find /app -maxdepth 4 -name openenv.yaml -print -quit 2>/dev/null",
|
| 136 |
+
timeout=10,
|
| 137 |
+
)
|
| 138 |
+
path = (resp.result if hasattr(resp, "result") else str(resp) or "").strip()
|
| 139 |
+
if path and path.startswith("/"):
|
| 140 |
+
return path
|
| 141 |
+
|
| 142 |
+
return None
|
| 143 |
+
|
| 144 |
+
@staticmethod
|
| 145 |
+
def _parse_app_field(yaml_content: str) -> Optional[str]:
|
| 146 |
+
"""Extract the ``app`` value from raw openenv.yaml content.
|
| 147 |
+
|
| 148 |
+
Uses PyYAML to handle comments, quotes, and nested keys correctly.
|
| 149 |
+
"""
|
| 150 |
+
return parse_openenv_app_field(yaml_content)
|
| 151 |
+
|
| 152 |
+
@staticmethod
|
| 153 |
+
def _parse_dockerfile_cmd(dockerfile_content: str) -> Optional[str]:
|
| 154 |
+
"""Extract the server command from the last ``CMD`` in a Dockerfile.
|
| 155 |
+
|
| 156 |
+
Handles exec form (``CMD ["prog", "arg"]``) and shell form
|
| 157 |
+
(``CMD prog arg``). When a Dockerfile has multiple ``CMD``
|
| 158 |
+
instructions (e.g. multi-stage builds), the last one wins - same
|
| 159 |
+
semantics as Docker itself. Lines where ``CMD`` appears inside a
|
| 160 |
+
comment are ignored.
|
| 161 |
+
|
| 162 |
+
Returns:
|
| 163 |
+
The command as a single string, or ``None`` if no ``CMD`` found.
|
| 164 |
+
"""
|
| 165 |
+
return parse_dockerfile_cmd(dockerfile_content)
|
| 166 |
+
|
| 167 |
+
@staticmethod
|
| 168 |
+
def strip_buildkit_syntax(dockerfile_content: str) -> str:
|
| 169 |
+
"""Remove BuildKit ``--mount=...`` flags from ``RUN`` instructions.
|
| 170 |
+
|
| 171 |
+
Handles single-line flags, multi-line continuations, and multiple
|
| 172 |
+
``--mount`` flags spread across continuation lines. Only leading
|
| 173 |
+
``--mount`` flags are removed (before the actual command starts).
|
| 174 |
+
|
| 175 |
+
Daytona's ``Image.from_dockerfile`` does not support BuildKit
|
| 176 |
+
``--mount`` syntax. This helper strips the flags so that standard
|
| 177 |
+
Dockerfiles (like the ones generated by ``openenv build``) can
|
| 178 |
+
be used directly.
|
| 179 |
+
"""
|
| 180 |
+
import re
|
| 181 |
+
|
| 182 |
+
def strip_leading_mounts(text: str) -> str:
|
| 183 |
+
remaining = text
|
| 184 |
+
while True:
|
| 185 |
+
match = re.match(r"\s*--mount=\S+\s*", remaining)
|
| 186 |
+
if not match:
|
| 187 |
+
return remaining
|
| 188 |
+
remaining = remaining[match.end() :]
|
| 189 |
+
|
| 190 |
+
lines = dockerfile_content.split("\n")
|
| 191 |
+
result: list[str] = []
|
| 192 |
+
in_run = False
|
| 193 |
+
in_mount_prefix = False
|
| 194 |
+
|
| 195 |
+
for line in lines:
|
| 196 |
+
line_out = line
|
| 197 |
+
run_start = False
|
| 198 |
+
if re.match(r"\s*RUN(\s+|$)", line, flags=re.IGNORECASE):
|
| 199 |
+
in_run = True
|
| 200 |
+
in_mount_prefix = True
|
| 201 |
+
run_start = True
|
| 202 |
+
|
| 203 |
+
if in_run and in_mount_prefix:
|
| 204 |
+
original_ends_with_slash = line_out.rstrip().endswith("\\")
|
| 205 |
+
if run_start:
|
| 206 |
+
match = re.match(r"(\s*RUN\s+)(.*)$", line_out, flags=re.IGNORECASE)
|
| 207 |
+
if match:
|
| 208 |
+
run_prefix, remainder = match.group(1), match.group(2)
|
| 209 |
+
else:
|
| 210 |
+
run_prefix, remainder = line_out, ""
|
| 211 |
+
new_remainder = strip_leading_mounts(remainder)
|
| 212 |
+
line_out = run_prefix + new_remainder
|
| 213 |
+
content_for_check = new_remainder
|
| 214 |
+
else:
|
| 215 |
+
new_remainder = strip_leading_mounts(line_out)
|
| 216 |
+
line_out = new_remainder
|
| 217 |
+
content_for_check = new_remainder
|
| 218 |
+
|
| 219 |
+
if original_ends_with_slash and not line_out.rstrip().endswith("\\"):
|
| 220 |
+
line_out = line_out.rstrip() + " \\"
|
| 221 |
+
|
| 222 |
+
if content_for_check.strip() not in ("", "\\"):
|
| 223 |
+
in_mount_prefix = False
|
| 224 |
+
|
| 225 |
+
if in_run and not line_out.rstrip().endswith("\\"):
|
| 226 |
+
in_run = False
|
| 227 |
+
in_mount_prefix = False
|
| 228 |
+
|
| 229 |
+
result.append(line_out)
|
| 230 |
+
|
| 231 |
+
return "\n".join(result)
|
| 232 |
+
|
| 233 |
+
@classmethod
|
| 234 |
+
def image_from_dockerfile(
|
| 235 |
+
cls,
|
| 236 |
+
dockerfile_path: str,
|
| 237 |
+
context_dir: str | None = None,
|
| 238 |
+
) -> str:
|
| 239 |
+
"""Validate a Dockerfile and return a ``dockerfile:`` URI for
|
| 240 |
+
:meth:`start_container`.
|
| 241 |
+
|
| 242 |
+
Eagerly validates the Dockerfile (existence, COPY sources,
|
| 243 |
+
BuildKit stripping) and stores the processed content in an
|
| 244 |
+
internal registry. The actual ``daytona.Image`` is created
|
| 245 |
+
later inside ``start_container``.
|
| 246 |
+
|
| 247 |
+
Args:
|
| 248 |
+
dockerfile_path: Path to the Dockerfile on disk.
|
| 249 |
+
context_dir: Build context directory. Defaults to the
|
| 250 |
+
Dockerfile's grandparent directory, matching the
|
| 251 |
+
``openenv init`` convention where Dockerfiles live in
|
| 252 |
+
``<env>/server/Dockerfile`` and the build context is
|
| 253 |
+
``<env>/``. Pass explicitly for non-standard layouts
|
| 254 |
+
(e.g. ``context_dir="."`` for repo-root contexts).
|
| 255 |
+
|
| 256 |
+
Returns:
|
| 257 |
+
A ``"dockerfile:<abs_path>"`` string to pass to
|
| 258 |
+
``start_container``.
|
| 259 |
+
|
| 260 |
+
Raises:
|
| 261 |
+
FileNotFoundError: If *dockerfile_path* does not exist.
|
| 262 |
+
ValueError: If *context_dir* is given but does not exist,
|
| 263 |
+
or if COPY sources in the Dockerfile cannot be found
|
| 264 |
+
under the resolved context directory.
|
| 265 |
+
"""
|
| 266 |
+
import pathlib
|
| 267 |
+
import re
|
| 268 |
+
|
| 269 |
+
src = pathlib.Path(dockerfile_path).resolve()
|
| 270 |
+
if not src.is_file():
|
| 271 |
+
raise FileNotFoundError(f"Dockerfile not found: {dockerfile_path}")
|
| 272 |
+
|
| 273 |
+
if context_dir is not None:
|
| 274 |
+
ctx = pathlib.Path(context_dir)
|
| 275 |
+
if not ctx.is_dir():
|
| 276 |
+
raise ValueError(f"context_dir does not exist: {context_dir}")
|
| 277 |
+
else:
|
| 278 |
+
# Default: grandparent of the Dockerfile, matching the
|
| 279 |
+
# openenv init layout (<env>/server/Dockerfile -> <env>/).
|
| 280 |
+
ctx = src.parent.parent
|
| 281 |
+
|
| 282 |
+
content = src.read_text()
|
| 283 |
+
stripped = cls.strip_buildkit_syntax(content)
|
| 284 |
+
|
| 285 |
+
# Validate that COPY sources exist under the context directory.
|
| 286 |
+
# This catches mismatches early (e.g. a Dockerfile expecting repo
|
| 287 |
+
# root as context when we defaulted to the env directory).
|
| 288 |
+
for line in stripped.splitlines():
|
| 289 |
+
m = re.match(r"^\s*COPY\s+(?!--from=)(\S+)\s+", line, re.IGNORECASE)
|
| 290 |
+
if not m:
|
| 291 |
+
continue
|
| 292 |
+
copy_src = m.group(1)
|
| 293 |
+
if copy_src.startswith("/"):
|
| 294 |
+
continue
|
| 295 |
+
resolved = ctx / copy_src
|
| 296 |
+
if not resolved.exists() and not any(ctx.glob(copy_src)):
|
| 297 |
+
raise ValueError(
|
| 298 |
+
f"Dockerfile COPY source '{copy_src}' not found "
|
| 299 |
+
f"under context_dir '{ctx}'. This Dockerfile may "
|
| 300 |
+
f"expect a different build context (e.g. the repo "
|
| 301 |
+
f"root). Pass context_dir explicitly."
|
| 302 |
+
)
|
| 303 |
+
|
| 304 |
+
# Parse CMD from the original Dockerfile so start_container can
|
| 305 |
+
# use it as a fallback when openenv.yaml is unavailable.
|
| 306 |
+
parsed_cmd = cls._parse_dockerfile_cmd(content)
|
| 307 |
+
|
| 308 |
+
cls._dockerfile_registry[str(src)] = {
|
| 309 |
+
"stripped_content": stripped,
|
| 310 |
+
"context_dir": str(ctx),
|
| 311 |
+
"server_cmd": parsed_cmd,
|
| 312 |
+
}
|
| 313 |
+
|
| 314 |
+
return f"dockerfile:{src}"
|
| 315 |
+
|
| 316 |
+
def start_container(
|
| 317 |
+
self,
|
| 318 |
+
image: Optional[str] = None,
|
| 319 |
+
port: Optional[int] = None,
|
| 320 |
+
env_vars: Optional[Dict[str, str]] = None,
|
| 321 |
+
**kwargs: Any,
|
| 322 |
+
) -> str:
|
| 323 |
+
"""
|
| 324 |
+
Create a Daytona sandbox from a Docker image or snapshot.
|
| 325 |
+
|
| 326 |
+
Daytona does not execute the image's CMD (known bug — ENTRYPOINT
|
| 327 |
+
runs, CMD does not). The server command is resolved in order:
|
| 328 |
+
|
| 329 |
+
1. Explicit ``cmd`` passed to the constructor.
|
| 330 |
+
2. ``cmd`` key in ``**kwargs`` (popped before forwarding).
|
| 331 |
+
3. Auto-discovered from ``openenv.yaml`` inside the sandbox.
|
| 332 |
+
4. ``CMD`` parsed from the Dockerfile (when *image* came from
|
| 333 |
+
``image_from_dockerfile``).
|
| 334 |
+
|
| 335 |
+
Args:
|
| 336 |
+
image: Docker image name (e.g. ``"echo-env:latest"``),
|
| 337 |
+
``"snapshot:<name>"`` to create from a pre-built snapshot,
|
| 338 |
+
or ``"dockerfile:<path>"`` returned by
|
| 339 |
+
:meth:`image_from_dockerfile`. May be omitted when supplied
|
| 340 |
+
to the constructor.
|
| 341 |
+
port: Must be ``None`` or ``8000``. Daytona exposes port 8000
|
| 342 |
+
via its preview proxy; other ports raise ``ValueError``.
|
| 343 |
+
env_vars: Environment variables forwarded to the sandbox.
|
| 344 |
+
**kwargs: ``cmd`` (str) to override the server command;
|
| 345 |
+
remaining kwargs passed through to ``Daytona.create()``.
|
| 346 |
+
|
| 347 |
+
Returns:
|
| 348 |
+
HTTPS preview URL for the sandbox (base_url).
|
| 349 |
+
"""
|
| 350 |
+
if port is not None and port != 8000:
|
| 351 |
+
raise ValueError(
|
| 352 |
+
f"DaytonaProvider only supports port 8000 (got {port}). "
|
| 353 |
+
"The Daytona preview proxy routes to port 8000 inside the sandbox."
|
| 354 |
+
)
|
| 355 |
+
|
| 356 |
+
effective_image = image if image is not None else self._image
|
| 357 |
+
if effective_image is None:
|
| 358 |
+
raise ValueError(
|
| 359 |
+
"DaytonaProvider requires an image. Pass it to the constructor "
|
| 360 |
+
"or start_container()."
|
| 361 |
+
)
|
| 362 |
+
effective_env_vars = self._env_vars if env_vars is None else env_vars
|
| 363 |
+
|
| 364 |
+
# Resolve the server command (may be None; discovery happens after
|
| 365 |
+
# sandbox creation when we can inspect the filesystem).
|
| 366 |
+
cmd = kwargs.pop("cmd", None) or self._cmd
|
| 367 |
+
|
| 368 |
+
# CMD parsed from Dockerfile (populated for "dockerfile:" images).
|
| 369 |
+
parsed_cmd: Optional[str] = None
|
| 370 |
+
|
| 371 |
+
# Build creation params
|
| 372 |
+
create_kwargs: Dict[str, Any] = {}
|
| 373 |
+
if effective_env_vars:
|
| 374 |
+
create_kwargs["env_vars"] = effective_env_vars
|
| 375 |
+
if self._public:
|
| 376 |
+
create_kwargs["public"] = True
|
| 377 |
+
if self._auto_stop_interval != 15:
|
| 378 |
+
create_kwargs["auto_stop_interval"] = self._auto_stop_interval
|
| 379 |
+
|
| 380 |
+
if effective_image.startswith("snapshot:"):
|
| 381 |
+
from daytona import CreateSandboxFromSnapshotParams
|
| 382 |
+
|
| 383 |
+
snapshot_name = effective_image[len("snapshot:") :]
|
| 384 |
+
params = CreateSandboxFromSnapshotParams(
|
| 385 |
+
snapshot=snapshot_name, **create_kwargs
|
| 386 |
+
)
|
| 387 |
+
elif effective_image.startswith("dockerfile:"):
|
| 388 |
+
from daytona import CreateSandboxFromImageParams, Image
|
| 389 |
+
|
| 390 |
+
dockerfile_path = effective_image[len("dockerfile:") :]
|
| 391 |
+
meta = self._dockerfile_registry.get(dockerfile_path)
|
| 392 |
+
if meta is None:
|
| 393 |
+
raise ValueError(
|
| 394 |
+
f"No registered Dockerfile metadata for {dockerfile_path}. "
|
| 395 |
+
"Call DaytonaProvider.image_from_dockerfile() first."
|
| 396 |
+
)
|
| 397 |
+
|
| 398 |
+
parsed_cmd = meta.get("server_cmd")
|
| 399 |
+
|
| 400 |
+
# Build the daytona Image from the pre-stripped content.
|
| 401 |
+
import pathlib
|
| 402 |
+
import uuid
|
| 403 |
+
|
| 404 |
+
ctx = pathlib.Path(meta["context_dir"])
|
| 405 |
+
tmp_name = f".daytona-{uuid.uuid4().hex[:8]}.dockerfile"
|
| 406 |
+
tmp_path = ctx / tmp_name
|
| 407 |
+
try:
|
| 408 |
+
tmp_path.write_text(meta["stripped_content"])
|
| 409 |
+
daytona_image = Image.from_dockerfile(str(tmp_path))
|
| 410 |
+
finally:
|
| 411 |
+
tmp_path.unlink(missing_ok=True)
|
| 412 |
+
|
| 413 |
+
img_kwargs: Dict[str, Any] = {
|
| 414 |
+
"image": daytona_image,
|
| 415 |
+
**create_kwargs,
|
| 416 |
+
}
|
| 417 |
+
if self._resources is not None:
|
| 418 |
+
img_kwargs["resources"] = self._resources
|
| 419 |
+
params = CreateSandboxFromImageParams(**img_kwargs)
|
| 420 |
+
else:
|
| 421 |
+
from daytona import CreateSandboxFromImageParams
|
| 422 |
+
|
| 423 |
+
img_kwargs = {"image": effective_image, **create_kwargs}
|
| 424 |
+
if self._resources is not None:
|
| 425 |
+
img_kwargs["resources"] = self._resources
|
| 426 |
+
params = CreateSandboxFromImageParams(**img_kwargs)
|
| 427 |
+
|
| 428 |
+
# Create sandbox
|
| 429 |
+
extra: Dict[str, Any] = dict(kwargs)
|
| 430 |
+
if self._on_snapshot_create_logs is not None:
|
| 431 |
+
extra["on_snapshot_create_logs"] = self._on_snapshot_create_logs
|
| 432 |
+
|
| 433 |
+
self._sandbox = self._daytona.create(
|
| 434 |
+
params, timeout=self._create_timeout, **extra
|
| 435 |
+
)
|
| 436 |
+
|
| 437 |
+
try:
|
| 438 |
+
# Discover server command from openenv.yaml if not explicitly set.
|
| 439 |
+
if cmd is None:
|
| 440 |
+
try:
|
| 441 |
+
cmd = self._discover_server_cmd(self._sandbox)
|
| 442 |
+
except ValueError:
|
| 443 |
+
# Fall back to CMD parsed from Dockerfile (if available).
|
| 444 |
+
if parsed_cmd:
|
| 445 |
+
cmd = parsed_cmd
|
| 446 |
+
else:
|
| 447 |
+
raise
|
| 448 |
+
|
| 449 |
+
# Wrap in bash -c so compound commands (cd ... && uvicorn ...)
|
| 450 |
+
# are handled correctly by nohup. Write PID so we can check
|
| 451 |
+
# if the process crashed later in wait_for_ready().
|
| 452 |
+
escaped_cmd = shlex.quote(cmd)
|
| 453 |
+
self._sandbox.process.exec(
|
| 454 |
+
f"nohup bash -c {escaped_cmd} > /tmp/openenv-server.log 2>&1 &"
|
| 455 |
+
" echo $! > /tmp/openenv-server.pid",
|
| 456 |
+
timeout=10,
|
| 457 |
+
)
|
| 458 |
+
|
| 459 |
+
# Get a signed preview URL for port 8000. The token is
|
| 460 |
+
# embedded in the URL itself so no extra headers are needed.
|
| 461 |
+
signed = self._sandbox.create_signed_preview_url(
|
| 462 |
+
8000, expires_in_seconds=86400
|
| 463 |
+
)
|
| 464 |
+
self._preview_url = signed.url
|
| 465 |
+
except Exception:
|
| 466 |
+
self.stop_container()
|
| 467 |
+
raise
|
| 468 |
+
|
| 469 |
+
return self._preview_url
|
| 470 |
+
|
| 471 |
+
def refresh_preview_url(self) -> str:
|
| 472 |
+
"""Get a fresh signed preview URL (valid for 24h).
|
| 473 |
+
|
| 474 |
+
Daytona signed URLs expire after at most 24 hours. Call this to
|
| 475 |
+
get a new one for long-running sessions. The returned URL points
|
| 476 |
+
to the same sandbox — clients will need to reconnect using it.
|
| 477 |
+
"""
|
| 478 |
+
if self._sandbox is None:
|
| 479 |
+
raise RuntimeError("No active sandbox to refresh URL for.")
|
| 480 |
+
signed = self._sandbox.create_signed_preview_url(8000, expires_in_seconds=86400)
|
| 481 |
+
self._preview_url = signed.url
|
| 482 |
+
return self._preview_url
|
| 483 |
+
|
| 484 |
+
def stop_container(self) -> None:
|
| 485 |
+
"""Delete the Daytona sandbox."""
|
| 486 |
+
if self._sandbox is None:
|
| 487 |
+
return
|
| 488 |
+
|
| 489 |
+
try:
|
| 490 |
+
self._daytona.delete(self._sandbox)
|
| 491 |
+
finally:
|
| 492 |
+
self._sandbox = None
|
| 493 |
+
self._preview_url = None
|
| 494 |
+
|
| 495 |
+
def wait_for_ready(self, base_url: str, timeout_s: float = 120.0) -> None:
|
| 496 |
+
"""
|
| 497 |
+
Poll the /health endpoint until the sandbox is ready.
|
| 498 |
+
|
| 499 |
+
Uses a longer default timeout (120s) than Docker providers because
|
| 500 |
+
Daytona sandboxes may have cold-start latency.
|
| 501 |
+
|
| 502 |
+
Args:
|
| 503 |
+
base_url: Preview URL returned by ``start_container()``.
|
| 504 |
+
timeout_s: Maximum seconds to wait.
|
| 505 |
+
|
| 506 |
+
Raises:
|
| 507 |
+
TimeoutError: If the sandbox doesn't become ready in time.
|
| 508 |
+
RuntimeError: If the server process died (detected via PID check).
|
| 509 |
+
"""
|
| 510 |
+
import requests
|
| 511 |
+
|
| 512 |
+
health_url = f"{base_url}/health"
|
| 513 |
+
|
| 514 |
+
deadline = time.time() + timeout_s
|
| 515 |
+
while time.time() < deadline:
|
| 516 |
+
try:
|
| 517 |
+
response = requests.get(health_url, timeout=5.0)
|
| 518 |
+
if response.status_code == 200:
|
| 519 |
+
return
|
| 520 |
+
except requests.RequestException:
|
| 521 |
+
pass
|
| 522 |
+
|
| 523 |
+
# Early exit: if the server process died, raise immediately
|
| 524 |
+
# instead of waiting for the full health-check timeout.
|
| 525 |
+
if self._sandbox is not None:
|
| 526 |
+
resp = self._sandbox.process.exec(
|
| 527 |
+
"kill -0 $(cat /tmp/openenv-server.pid) 2>/dev/null"
|
| 528 |
+
" && echo RUNNING || echo DEAD",
|
| 529 |
+
timeout=10,
|
| 530 |
+
)
|
| 531 |
+
out = resp.result if hasattr(resp, "result") else str(resp)
|
| 532 |
+
if "DEAD" in (out or ""):
|
| 533 |
+
log_resp = self._sandbox.process.exec(
|
| 534 |
+
"cat /tmp/openenv-server.log 2>/dev/null", timeout=10
|
| 535 |
+
)
|
| 536 |
+
log = (
|
| 537 |
+
log_resp.result
|
| 538 |
+
if hasattr(log_resp, "result")
|
| 539 |
+
else str(log_resp)
|
| 540 |
+
)
|
| 541 |
+
raise RuntimeError(f"Server process died.\nLog:\n{log}")
|
| 542 |
+
|
| 543 |
+
time.sleep(1.0)
|
| 544 |
+
|
| 545 |
+
raise TimeoutError(
|
| 546 |
+
f"Daytona sandbox at {base_url} did not become ready within {timeout_s}s"
|
| 547 |
+
)
|
src/core/containers/runtime/hf_sandbox_provider.py
ADDED
|
@@ -0,0 +1,342 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 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 the
|
| 5 |
+
# LICENSE file in the root directory of this source tree.
|
| 6 |
+
|
| 7 |
+
"""Hugging Face-backed provider for OpenEnv environment servers."""
|
| 8 |
+
|
| 9 |
+
from __future__ import annotations
|
| 10 |
+
|
| 11 |
+
import asyncio
|
| 12 |
+
import hashlib
|
| 13 |
+
import socket
|
| 14 |
+
import threading
|
| 15 |
+
import time
|
| 16 |
+
from contextlib import suppress
|
| 17 |
+
from typing import Any
|
| 18 |
+
|
| 19 |
+
import requests
|
| 20 |
+
from fastapi import FastAPI, Request, Response, WebSocket, WebSocketDisconnect
|
| 21 |
+
from websockets.asyncio.client import connect as ws_connect
|
| 22 |
+
from websockets.exceptions import ConnectionClosed
|
| 23 |
+
|
| 24 |
+
from .providers import ContainerProvider
|
| 25 |
+
|
| 26 |
+
|
| 27 |
+
_DEFAULT_PORT = 8000
|
| 28 |
+
_SERVER_COMMAND = (
|
| 29 |
+
'export SBX_PROXY_DIR="${SBX_PROXY_DIR:-$HOME/.sbx/proxy}"; '
|
| 30 |
+
'mkdir -p "$SBX_PROXY_DIR"; '
|
| 31 |
+
'nohup server >"$HOME/openenv-server.log" 2>&1 &'
|
| 32 |
+
)
|
| 33 |
+
_MAX_WS_MESSAGE_SIZE = 100 * 1024 * 1024
|
| 34 |
+
_HOP_BY_HOP_HEADERS = {
|
| 35 |
+
"connection",
|
| 36 |
+
"content-encoding",
|
| 37 |
+
"content-length",
|
| 38 |
+
"host",
|
| 39 |
+
"keep-alive",
|
| 40 |
+
"proxy-authenticate",
|
| 41 |
+
"proxy-authorization",
|
| 42 |
+
"te",
|
| 43 |
+
"trailer",
|
| 44 |
+
"transfer-encoding",
|
| 45 |
+
"upgrade",
|
| 46 |
+
}
|
| 47 |
+
_POOLS: dict[tuple[str, str], Any] = {}
|
| 48 |
+
_POOL_LOCK = threading.Lock()
|
| 49 |
+
|
| 50 |
+
|
| 51 |
+
def _find_available_port() -> int:
|
| 52 |
+
with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock:
|
| 53 |
+
sock.bind(("127.0.0.1", 0))
|
| 54 |
+
sock.listen(1)
|
| 55 |
+
return sock.getsockname()[1]
|
| 56 |
+
|
| 57 |
+
|
| 58 |
+
def _to_ws_url(url: str) -> str:
|
| 59 |
+
if url.startswith("https://"):
|
| 60 |
+
return "wss://" + url[len("https://") :]
|
| 61 |
+
if url.startswith("http://"):
|
| 62 |
+
return "ws://" + url[len("http://") :]
|
| 63 |
+
return url
|
| 64 |
+
|
| 65 |
+
|
| 66 |
+
def _pool_name(image: str, flavor: str) -> str:
|
| 67 |
+
digest = hashlib.sha1(f"{image}\0{flavor}".encode("utf-8")).hexdigest()[:12]
|
| 68 |
+
return f"openenv-{digest}"
|
| 69 |
+
|
| 70 |
+
|
| 71 |
+
def _get_sandbox_pool_cls() -> Any:
|
| 72 |
+
try:
|
| 73 |
+
from huggingface_hub import SandboxPool
|
| 74 |
+
except ImportError as exc:
|
| 75 |
+
raise RuntimeError(
|
| 76 |
+
"HFSandboxProvider requires a huggingface_hub version with "
|
| 77 |
+
"SandboxPool.create and Sandbox.proxy_url_for support."
|
| 78 |
+
) from exc
|
| 79 |
+
return SandboxPool
|
| 80 |
+
|
| 81 |
+
|
| 82 |
+
def _get_pool(image: str, flavor: str) -> Any:
|
| 83 |
+
key = (image, flavor)
|
| 84 |
+
with _POOL_LOCK:
|
| 85 |
+
pool = _POOLS.get(key)
|
| 86 |
+
if pool is None:
|
| 87 |
+
SandboxPool = _get_sandbox_pool_cls()
|
| 88 |
+
pool = SandboxPool(
|
| 89 |
+
image=image,
|
| 90 |
+
flavor=flavor,
|
| 91 |
+
name=_pool_name(image, flavor),
|
| 92 |
+
)
|
| 93 |
+
_POOLS[key] = pool
|
| 94 |
+
return pool
|
| 95 |
+
|
| 96 |
+
|
| 97 |
+
class _LocalAuthProxy:
|
| 98 |
+
def __init__(self, *, target_url: str, headers: dict[str, str]):
|
| 99 |
+
self.target_url = target_url.rstrip("/")
|
| 100 |
+
self.headers = headers
|
| 101 |
+
self.port = _find_available_port()
|
| 102 |
+
self._server: Any = None
|
| 103 |
+
self._thread: threading.Thread | None = None
|
| 104 |
+
|
| 105 |
+
@property
|
| 106 |
+
def base_url(self) -> str:
|
| 107 |
+
return f"http://127.0.0.1:{self.port}"
|
| 108 |
+
|
| 109 |
+
def start(self) -> str:
|
| 110 |
+
app = FastAPI()
|
| 111 |
+
|
| 112 |
+
@app.api_route(
|
| 113 |
+
"/{path:path}",
|
| 114 |
+
methods=["GET", "POST", "PUT", "PATCH", "DELETE", "OPTIONS"],
|
| 115 |
+
)
|
| 116 |
+
async def proxy_http(path: str, request: Request) -> Response:
|
| 117 |
+
query = request.url.query
|
| 118 |
+
target = f"{self.target_url}/{path}"
|
| 119 |
+
if query:
|
| 120 |
+
target = f"{target}?{query}"
|
| 121 |
+
headers = dict(self.headers)
|
| 122 |
+
content_type = request.headers.get("content-type")
|
| 123 |
+
if content_type is not None:
|
| 124 |
+
headers["Content-Type"] = content_type
|
| 125 |
+
body = await request.body()
|
| 126 |
+
try:
|
| 127 |
+
upstream = await asyncio.to_thread(
|
| 128 |
+
requests.request,
|
| 129 |
+
request.method,
|
| 130 |
+
target,
|
| 131 |
+
data=body,
|
| 132 |
+
headers=headers,
|
| 133 |
+
timeout=60.0,
|
| 134 |
+
allow_redirects=True,
|
| 135 |
+
)
|
| 136 |
+
except requests.RequestException:
|
| 137 |
+
return Response(
|
| 138 |
+
content=b"upstream HF job unreachable",
|
| 139 |
+
status_code=502,
|
| 140 |
+
)
|
| 141 |
+
response_headers = {
|
| 142 |
+
key: value
|
| 143 |
+
for key, value in upstream.headers.items()
|
| 144 |
+
if key.lower() not in _HOP_BY_HOP_HEADERS
|
| 145 |
+
}
|
| 146 |
+
return Response(
|
| 147 |
+
content=upstream.content,
|
| 148 |
+
status_code=upstream.status_code,
|
| 149 |
+
headers=response_headers,
|
| 150 |
+
)
|
| 151 |
+
|
| 152 |
+
@app.websocket("/{path:path}")
|
| 153 |
+
async def proxy_websocket(path: str, websocket: WebSocket) -> None:
|
| 154 |
+
query = websocket.url.query
|
| 155 |
+
target = f"{_to_ws_url(self.target_url)}/{path}"
|
| 156 |
+
if query:
|
| 157 |
+
target = f"{target}?{query}"
|
| 158 |
+
headers = {**self.headers, "accept": "*/*"}
|
| 159 |
+
upstream = await self._connect_upstream_websocket(target, headers)
|
| 160 |
+
await websocket.accept()
|
| 161 |
+
try:
|
| 162 |
+
to_upstream = asyncio.create_task(
|
| 163 |
+
self._client_to_upstream(websocket, upstream)
|
| 164 |
+
)
|
| 165 |
+
to_client = asyncio.create_task(
|
| 166 |
+
self._upstream_to_client(websocket, upstream)
|
| 167 |
+
)
|
| 168 |
+
done, pending = await asyncio.wait(
|
| 169 |
+
{to_upstream, to_client},
|
| 170 |
+
return_when=asyncio.FIRST_COMPLETED,
|
| 171 |
+
)
|
| 172 |
+
for task in pending:
|
| 173 |
+
task.cancel()
|
| 174 |
+
for task in done:
|
| 175 |
+
with suppress(ConnectionClosed, WebSocketDisconnect):
|
| 176 |
+
task.result()
|
| 177 |
+
finally:
|
| 178 |
+
with suppress(Exception):
|
| 179 |
+
await upstream.close()
|
| 180 |
+
|
| 181 |
+
import uvicorn
|
| 182 |
+
|
| 183 |
+
config = uvicorn.Config(
|
| 184 |
+
app,
|
| 185 |
+
host="127.0.0.1",
|
| 186 |
+
port=self.port,
|
| 187 |
+
log_level="warning",
|
| 188 |
+
access_log=False,
|
| 189 |
+
)
|
| 190 |
+
self._server = uvicorn.Server(config)
|
| 191 |
+
self._thread = threading.Thread(target=self._server.run, daemon=True)
|
| 192 |
+
self._thread.start()
|
| 193 |
+
while not self._server.started:
|
| 194 |
+
if not self._thread.is_alive():
|
| 195 |
+
raise RuntimeError("HF sandbox auth proxy failed to start")
|
| 196 |
+
time.sleep(0.05)
|
| 197 |
+
return self.base_url
|
| 198 |
+
|
| 199 |
+
async def _connect_upstream_websocket(
|
| 200 |
+
self, target: str, headers: dict[str, str]
|
| 201 |
+
) -> Any:
|
| 202 |
+
last_error: Exception | None = None
|
| 203 |
+
for _ in range(5):
|
| 204 |
+
try:
|
| 205 |
+
return await ws_connect(
|
| 206 |
+
target,
|
| 207 |
+
additional_headers=headers,
|
| 208 |
+
max_size=_MAX_WS_MESSAGE_SIZE,
|
| 209 |
+
compression=None,
|
| 210 |
+
)
|
| 211 |
+
except Exception as exc:
|
| 212 |
+
last_error = exc
|
| 213 |
+
await asyncio.sleep(0.5)
|
| 214 |
+
assert last_error is not None
|
| 215 |
+
raise last_error
|
| 216 |
+
|
| 217 |
+
async def _client_to_upstream(self, websocket: WebSocket, upstream: Any) -> None:
|
| 218 |
+
# EnvClient sends JSON text frames; binary frames are only relayed downstream.
|
| 219 |
+
async for message in websocket.iter_text():
|
| 220 |
+
await upstream.send(message)
|
| 221 |
+
|
| 222 |
+
async def _upstream_to_client(self, websocket: WebSocket, upstream: Any) -> None:
|
| 223 |
+
async for message in upstream:
|
| 224 |
+
if isinstance(message, bytes):
|
| 225 |
+
await websocket.send_bytes(message)
|
| 226 |
+
else:
|
| 227 |
+
await websocket.send_text(message)
|
| 228 |
+
|
| 229 |
+
def stop(self) -> None:
|
| 230 |
+
if self._server is None or self._thread is None:
|
| 231 |
+
return
|
| 232 |
+
self._server.should_exit = True
|
| 233 |
+
self._thread.join(timeout=5.0)
|
| 234 |
+
self._server = None
|
| 235 |
+
self._thread = None
|
| 236 |
+
|
| 237 |
+
|
| 238 |
+
class HFSandboxProvider(ContainerProvider):
|
| 239 |
+
"""Run an OpenEnv server on Hugging Face infrastructure.
|
| 240 |
+
|
| 241 |
+
The server runs in a `huggingface_hub` sandbox pool, which bills its hosts as
|
| 242 |
+
Hugging Face Jobs. `start_container()` returns a local proxy URL that adds the
|
| 243 |
+
Hugging Face auth headers, so any `EnvClient` can connect to it unchanged.
|
| 244 |
+
|
| 245 |
+
Args:
|
| 246 |
+
image (`str`):
|
| 247 |
+
Image to run, for example `hf.co/spaces/openenv/coding_env`. It must
|
| 248 |
+
provide a `server` command that serves OpenEnv on port 8000.
|
| 249 |
+
env_vars (`dict[str, str]`, *optional*):
|
| 250 |
+
Environment variables for the server.
|
| 251 |
+
flavor (`str`, *optional*, defaults to `"cpu-basic"`):
|
| 252 |
+
Hugging Face hardware flavor of the sandbox hosts.
|
| 253 |
+
"""
|
| 254 |
+
|
| 255 |
+
def __init__(
|
| 256 |
+
self,
|
| 257 |
+
*,
|
| 258 |
+
image: str,
|
| 259 |
+
env_vars: dict[str, str] | None = None,
|
| 260 |
+
flavor: str = "cpu-basic",
|
| 261 |
+
):
|
| 262 |
+
self.image = image
|
| 263 |
+
self.env_vars = env_vars
|
| 264 |
+
self.flavor = flavor
|
| 265 |
+
self._sandbox: Any = None
|
| 266 |
+
self._proxy: _LocalAuthProxy | None = None
|
| 267 |
+
|
| 268 |
+
def start_container(
|
| 269 |
+
self,
|
| 270 |
+
image: str | None = None,
|
| 271 |
+
port: int | None = None,
|
| 272 |
+
env_vars: dict[str, str] | None = None,
|
| 273 |
+
**kwargs: Any,
|
| 274 |
+
) -> str:
|
| 275 |
+
if self._sandbox is not None:
|
| 276 |
+
raise RuntimeError("HFSandboxProvider already has an active service")
|
| 277 |
+
if kwargs:
|
| 278 |
+
unsupported = ", ".join(sorted(kwargs))
|
| 279 |
+
raise ValueError(
|
| 280 |
+
f"HFSandboxProvider does not support start_container kwargs: {unsupported}"
|
| 281 |
+
)
|
| 282 |
+
|
| 283 |
+
if port not in (None, _DEFAULT_PORT):
|
| 284 |
+
raise ValueError(
|
| 285 |
+
f"HFSandboxProvider only supports port {_DEFAULT_PORT} (got {port})."
|
| 286 |
+
)
|
| 287 |
+
|
| 288 |
+
effective_image = self.image if image is None else image
|
| 289 |
+
effective_env = self.env_vars if env_vars is None else env_vars
|
| 290 |
+
pool = _get_pool(effective_image, self.flavor)
|
| 291 |
+
self._sandbox = pool.create(env=effective_env)
|
| 292 |
+
try:
|
| 293 |
+
if not hasattr(self._sandbox, "proxy_url_for") or not hasattr(
|
| 294 |
+
self._sandbox, "proxy_headers"
|
| 295 |
+
):
|
| 296 |
+
raise RuntimeError(
|
| 297 |
+
"HFSandboxProvider requires a huggingface_hub version with "
|
| 298 |
+
"Sandbox.proxy_url_for and Sandbox.proxy_headers support."
|
| 299 |
+
)
|
| 300 |
+
self._sandbox.run(
|
| 301 |
+
_SERVER_COMMAND,
|
| 302 |
+
shell=True,
|
| 303 |
+
)
|
| 304 |
+
self._proxy = _LocalAuthProxy(
|
| 305 |
+
target_url=self._sandbox.proxy_url_for(_DEFAULT_PORT, "/"),
|
| 306 |
+
headers=self._sandbox.proxy_headers,
|
| 307 |
+
)
|
| 308 |
+
return self._proxy.start()
|
| 309 |
+
except Exception:
|
| 310 |
+
self.stop_container()
|
| 311 |
+
raise
|
| 312 |
+
|
| 313 |
+
def stop_container(self) -> None:
|
| 314 |
+
if self._proxy is not None:
|
| 315 |
+
self._proxy.stop()
|
| 316 |
+
self._proxy = None
|
| 317 |
+
if self._sandbox is not None:
|
| 318 |
+
sandbox = self._sandbox
|
| 319 |
+
self._sandbox = None
|
| 320 |
+
with suppress(Exception):
|
| 321 |
+
sandbox.kill()
|
| 322 |
+
|
| 323 |
+
def wait_for_ready(self, base_url: str, timeout_s: float = 120.0) -> None:
|
| 324 |
+
deadline = time.time() + timeout_s
|
| 325 |
+
health_url = f"{base_url}/health"
|
| 326 |
+
while time.time() < deadline:
|
| 327 |
+
try:
|
| 328 |
+
response = requests.get(health_url, timeout=5.0)
|
| 329 |
+
if response.status_code == 200:
|
| 330 |
+
return
|
| 331 |
+
except requests.exceptions.RequestException:
|
| 332 |
+
pass
|
| 333 |
+
time.sleep(1.0)
|
| 334 |
+
raise TimeoutError(
|
| 335 |
+
f"HF sandbox job at {base_url} did not become ready within {timeout_s}s"
|
| 336 |
+
)
|
| 337 |
+
|
| 338 |
+
def close(self) -> None:
|
| 339 |
+
self.stop_container()
|
| 340 |
+
|
| 341 |
+
|
| 342 |
+
__all__ = ["HFSandboxProvider"]
|
src/core/containers/runtime/modal_provider.py
ADDED
|
@@ -0,0 +1,697 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# SPDX-License-Identifier: BSD-3-Clause
|
| 2 |
+
|
| 3 |
+
"""
|
| 4 |
+
Modal container provider for running OpenEnv environments in Modal sandboxes.
|
| 5 |
+
|
| 6 |
+
Requires the ``modal`` SDK: ``pip install modal>=1.4``
|
| 7 |
+
|
| 8 |
+
The provider boots an OpenEnv server inside a Modal sandbox, exposes it on an
|
| 9 |
+
encrypted tunnel, and returns an ``https://`` URL that ``EnvClient`` connects to
|
| 10 |
+
over ``wss://``.
|
| 11 |
+
|
| 12 |
+
Supports both the stable Sandbox API (``modal.Sandbox.create``) and the beta
|
| 13 |
+
Sandbox v2 API (``modal.Sandbox._experimental_create``). Sandbox v2 is opt-in
|
| 14 |
+
via ``use_sandbox_v2=True`` because it is an experimental, private SDK feature;
|
| 15 |
+
see https://modal.com/docs/guide/sandbox-v2 .
|
| 16 |
+
"""
|
| 17 |
+
|
| 18 |
+
from __future__ import annotations
|
| 19 |
+
|
| 20 |
+
import logging
|
| 21 |
+
import shlex
|
| 22 |
+
import time
|
| 23 |
+
from typing import Any, Dict, Optional
|
| 24 |
+
|
| 25 |
+
from ._server_config import parse_dockerfile_cmd, parse_openenv_app_field
|
| 26 |
+
from .providers import ContainerProvider
|
| 27 |
+
|
| 28 |
+
logger = logging.getLogger(__name__)
|
| 29 |
+
|
| 30 |
+
_DEFAULT_MODAL_PORT = 8000
|
| 31 |
+
|
| 32 |
+
# Modal's default per-container resource requests. We pass configured cpu/memory
|
| 33 |
+
# as the *limit* half of Modal's ``(request, limit)`` tuple so they cap usage
|
| 34 |
+
# rather than reserving (and billing for) the full amount. See
|
| 35 |
+
# https://modal.com/docs/guide/resources .
|
| 36 |
+
_DEFAULT_CPU_REQUEST = 0.125
|
| 37 |
+
_DEFAULT_MEMORY_REQUEST = 128
|
| 38 |
+
|
| 39 |
+
|
| 40 |
+
def _require_secure_url(url: str) -> str:
|
| 41 |
+
"""Enforce https/wss transport (RFC 002 security invariant S1).
|
| 42 |
+
|
| 43 |
+
``EnvClient`` derives its WebSocket URL from this base URL, so a plaintext
|
| 44 |
+
URL would become a cleartext ``ws://`` connection. The offending URL is
|
| 45 |
+
deliberately omitted from the error because a Modal tunnel URL is a bearer
|
| 46 |
+
capability that must not leak into logs.
|
| 47 |
+
"""
|
| 48 |
+
if not isinstance(url, str) or not url.lower().startswith("https://"):
|
| 49 |
+
raise RuntimeError(
|
| 50 |
+
"Modal sandbox returned a non-HTTPS tunnel URL. OpenEnv requires an "
|
| 51 |
+
"https/wss base_url so EnvClient traffic is encrypted. Refusing to "
|
| 52 |
+
"connect over plaintext."
|
| 53 |
+
)
|
| 54 |
+
return url
|
| 55 |
+
|
| 56 |
+
|
| 57 |
+
class _DefaultModalAdapter:
|
| 58 |
+
"""Thin adapter over the ``modal`` SDK.
|
| 59 |
+
|
| 60 |
+
The provider talks to this private adapter instead of spreading SDK details
|
| 61 |
+
through its own logic; tests inject a duck-typed fake in its place. Keeping
|
| 62 |
+
the SDK surface here also localizes the feature check for the private
|
| 63 |
+
Sandbox v2 (``_experimental_create``) API.
|
| 64 |
+
"""
|
| 65 |
+
|
| 66 |
+
def __init__(self, *, app_name: str, use_sandbox_v2: bool):
|
| 67 |
+
import modal
|
| 68 |
+
|
| 69 |
+
self._modal = modal
|
| 70 |
+
self._app_name = app_name
|
| 71 |
+
self._use_sandbox_v2 = use_sandbox_v2
|
| 72 |
+
|
| 73 |
+
# Sandbox v2 rides on a private SDK entry point that is not present in
|
| 74 |
+
# every modal release. Fail fast at construction with clear guidance
|
| 75 |
+
# rather than at start_container time with an AttributeError.
|
| 76 |
+
if use_sandbox_v2 and not hasattr(modal.Sandbox, "_experimental_create"):
|
| 77 |
+
raise RuntimeError(
|
| 78 |
+
"use_sandbox_v2=True requires modal.Sandbox._experimental_create, "
|
| 79 |
+
"which is not available in the installed modal SDK. Upgrade modal, "
|
| 80 |
+
"or use the stable Sandbox API (use_sandbox_v2=False). Sandbox v2 "
|
| 81 |
+
"is an experimental feature; contact support@modal.com for access."
|
| 82 |
+
)
|
| 83 |
+
|
| 84 |
+
def image_from_registry(self, tag: str) -> Any:
|
| 85 |
+
return self._modal.Image.from_registry(tag)
|
| 86 |
+
|
| 87 |
+
def image_from_dockerfile(self, dockerfile_path: str, context_dir: str) -> Any:
|
| 88 |
+
return self._modal.Image.from_dockerfile(
|
| 89 |
+
dockerfile_path, context_dir=context_dir
|
| 90 |
+
)
|
| 91 |
+
|
| 92 |
+
def create_sandbox(
|
| 93 |
+
self,
|
| 94 |
+
*,
|
| 95 |
+
image: Any,
|
| 96 |
+
encrypted_ports: list[int],
|
| 97 |
+
timeout: int,
|
| 98 |
+
env: Optional[dict[str, str]],
|
| 99 |
+
extra: dict[str, Any],
|
| 100 |
+
) -> Any:
|
| 101 |
+
app = self._modal.App.lookup(self._app_name, create_if_missing=True)
|
| 102 |
+
kwargs: Dict[str, Any] = {
|
| 103 |
+
"app": app,
|
| 104 |
+
"image": image,
|
| 105 |
+
"encrypted_ports": encrypted_ports,
|
| 106 |
+
"timeout": timeout,
|
| 107 |
+
**extra,
|
| 108 |
+
}
|
| 109 |
+
if env:
|
| 110 |
+
kwargs["env"] = dict(env)
|
| 111 |
+
|
| 112 |
+
# The sandbox is started with a keep-alive entrypoint; the server
|
| 113 |
+
# command is launched via exec afterwards (see ModalProvider).
|
| 114 |
+
if self._use_sandbox_v2:
|
| 115 |
+
return self._modal.Sandbox._experimental_create(
|
| 116 |
+
"sleep", "infinity", **kwargs
|
| 117 |
+
)
|
| 118 |
+
return self._modal.Sandbox.create("sleep", "infinity", **kwargs)
|
| 119 |
+
|
| 120 |
+
def exec(self, sandbox: Any, command: str, *, timeout: int = 10) -> str:
|
| 121 |
+
"""Run *command* through a shell inside *sandbox* and return its stdout."""
|
| 122 |
+
proc = sandbox.exec("bash", "-c", command, timeout=timeout)
|
| 123 |
+
try:
|
| 124 |
+
out = proc.stdout.read()
|
| 125 |
+
except Exception:
|
| 126 |
+
out = ""
|
| 127 |
+
proc.wait()
|
| 128 |
+
return out or ""
|
| 129 |
+
|
| 130 |
+
def tunnel_url(self, sandbox: Any, port: int) -> str:
|
| 131 |
+
return sandbox.tunnels()[port].url
|
| 132 |
+
|
| 133 |
+
def terminate(self, sandbox: Any) -> None:
|
| 134 |
+
sandbox.terminate()
|
| 135 |
+
|
| 136 |
+
|
| 137 |
+
class ModalProvider(ContainerProvider):
|
| 138 |
+
"""
|
| 139 |
+
Container provider that runs environments in Modal sandboxes.
|
| 140 |
+
|
| 141 |
+
``start_container``'s ``image`` is either a registry tag
|
| 142 |
+
(``"echo-env:latest"``) or a ``"dockerfile:<path>"`` reference returned by
|
| 143 |
+
:meth:`image_from_dockerfile`. The server is exposed on an encrypted Modal
|
| 144 |
+
tunnel and the returned ``https://`` URL is what ``EnvClient`` connects to
|
| 145 |
+
over ``wss://``.
|
| 146 |
+
|
| 147 |
+
The environment runs untrusted code, so the provider is secure by default:
|
| 148 |
+
it enforces https/wss transport, treats the tunnel URL as a bearer secret
|
| 149 |
+
(never interpolated into errors), and never surfaces raw sandbox output
|
| 150 |
+
unless ``surface_server_logs=True``.
|
| 151 |
+
|
| 152 |
+
Only one sandbox is active per provider: calling ``start_container`` again
|
| 153 |
+
before ``stop_container()``/``close()`` raises ``RuntimeError`` rather than
|
| 154 |
+
orphaning the running sandbox. ``close()`` (and context-manager exit) stops
|
| 155 |
+
the active sandbox.
|
| 156 |
+
|
| 157 |
+
Example:
|
| 158 |
+
```python
|
| 159 |
+
with ModalProvider(app_name="openenv", cpu=2.0, memory=4096) as provider:
|
| 160 |
+
image = ModalProvider.image_from_dockerfile(
|
| 161 |
+
"envs/echo_env/server/Dockerfile"
|
| 162 |
+
)
|
| 163 |
+
base_url = provider.start_container(image)
|
| 164 |
+
provider.wait_for_ready(base_url)
|
| 165 |
+
# sandbox terminated on exit
|
| 166 |
+
```
|
| 167 |
+
|
| 168 |
+
Sandbox v2 (beta) is opt-in:
|
| 169 |
+
```python
|
| 170 |
+
provider = ModalProvider(app_name="openenv", use_sandbox_v2=True)
|
| 171 |
+
```
|
| 172 |
+
"""
|
| 173 |
+
|
| 174 |
+
_dockerfile_registry: Dict[str, Dict[str, Any]] = {}
|
| 175 |
+
|
| 176 |
+
def __init__(
|
| 177 |
+
self,
|
| 178 |
+
*,
|
| 179 |
+
image: str | None = None,
|
| 180 |
+
env_vars: dict[str, str] | None = None,
|
| 181 |
+
app_name: str = "openenv",
|
| 182 |
+
use_sandbox_v2: bool = False,
|
| 183 |
+
timeout: int = 300,
|
| 184 |
+
cmd: str | None = None,
|
| 185 |
+
cpu: float | None = None,
|
| 186 |
+
memory: int | None = None,
|
| 187 |
+
surface_server_logs: bool = False,
|
| 188 |
+
_adapter: Any = None,
|
| 189 |
+
):
|
| 190 |
+
"""
|
| 191 |
+
Args:
|
| 192 |
+
image (`str`, *optional*):
|
| 193 |
+
Registry image tag or ``"dockerfile:<path>"`` source to use
|
| 194 |
+
when ``start_container()`` is called without an image.
|
| 195 |
+
env_vars (`dict`, *optional*):
|
| 196 |
+
Environment variables to use when ``start_container()`` is
|
| 197 |
+
called without explicit ``env_vars``.
|
| 198 |
+
app_name (`str`, *optional*, defaults to `"openenv"`):
|
| 199 |
+
Modal app name the sandbox is created under. Looked up (and
|
| 200 |
+
created if missing) via ``modal.App.lookup``.
|
| 201 |
+
use_sandbox_v2 (`bool`, *optional*, defaults to `False`):
|
| 202 |
+
When `True`, sandboxes are created via the beta
|
| 203 |
+
``modal.Sandbox._experimental_create`` API (Sandbox v2). This is
|
| 204 |
+
an experimental, private SDK feature and is therefore off by
|
| 205 |
+
default; a feature check fails fast at construction if the
|
| 206 |
+
installed SDK lacks it. See https://modal.com/docs/guide/sandbox-v2 .
|
| 207 |
+
timeout (`int`, *optional*, defaults to `300`):
|
| 208 |
+
Maximum sandbox lifetime in seconds.
|
| 209 |
+
cmd (`str`, *optional*):
|
| 210 |
+
Shell command to start the server inside the sandbox. When
|
| 211 |
+
omitted, the command is auto-discovered from ``openenv.yaml``
|
| 212 |
+
(falling back to the Dockerfile ``CMD``).
|
| 213 |
+
cpu (`float`, *optional*):
|
| 214 |
+
Hard CPU-core limit for the sandbox. Passed as the limit half of
|
| 215 |
+
Modal's ``(request, limit)`` tuple (request stays at Modal's
|
| 216 |
+
default), so it caps usage rather than reserving cores. When
|
| 217 |
+
`None`, Modal's default applies.
|
| 218 |
+
memory (`int`, *optional*):
|
| 219 |
+
Hard memory limit in MiB for the sandbox (containers exceeding it
|
| 220 |
+
are OOM-killed). Passed as the limit half of Modal's
|
| 221 |
+
``(request, limit)`` tuple. When `None`, Modal's default applies.
|
| 222 |
+
surface_server_logs (`bool`, *optional*, defaults to `False`):
|
| 223 |
+
When `False` (default), captured sandbox output is withheld from
|
| 224 |
+
raised errors so secrets the workload printed cannot leak into
|
| 225 |
+
orchestrator/CI logs. When `True`, a best-effort redacted,
|
| 226 |
+
length-bounded excerpt is included in startup-crash errors.
|
| 227 |
+
"""
|
| 228 |
+
self._image = image
|
| 229 |
+
self._env_vars = env_vars
|
| 230 |
+
self._app_name = app_name
|
| 231 |
+
self._use_sandbox_v2 = use_sandbox_v2
|
| 232 |
+
self._timeout = timeout
|
| 233 |
+
self._cmd = cmd
|
| 234 |
+
self._cpu = cpu
|
| 235 |
+
self._memory = memory
|
| 236 |
+
self._surface_server_logs = surface_server_logs
|
| 237 |
+
self._sandbox: Any = None
|
| 238 |
+
self._base_url: str | None = None
|
| 239 |
+
# Injected env-var values, used to scrub captured server output before
|
| 240 |
+
# it is ever surfaced in an error.
|
| 241 |
+
self._redact_values: set[str] = set()
|
| 242 |
+
|
| 243 |
+
if _adapter is None:
|
| 244 |
+
# Import eagerly (inside the adapter) so SDK/configuration errors —
|
| 245 |
+
# including the Sandbox v2 feature check — surface at construction.
|
| 246 |
+
self._adapter: Any = _DefaultModalAdapter(
|
| 247 |
+
app_name=app_name, use_sandbox_v2=use_sandbox_v2
|
| 248 |
+
)
|
| 249 |
+
else:
|
| 250 |
+
self._adapter = _adapter
|
| 251 |
+
|
| 252 |
+
if use_sandbox_v2:
|
| 253 |
+
logger.info(
|
| 254 |
+
"Using Modal Sandbox v2 (experimental). This feature must be "
|
| 255 |
+
"explicitly enabled. Contact support@modal.com to get access."
|
| 256 |
+
)
|
| 257 |
+
|
| 258 |
+
def _discover_server_cmd(self, port: int = _DEFAULT_MODAL_PORT) -> str:
|
| 259 |
+
"""Discover the server command from ``openenv.yaml`` inside the sandbox.
|
| 260 |
+
|
| 261 |
+
Finds the file, reads the ``app`` field, and constructs a command
|
| 262 |
+
of the form ``cd <env_root> && python -m uvicorn <app> --host 0.0.0.0 --port <port>``.
|
| 263 |
+
|
| 264 |
+
Raises:
|
| 265 |
+
ValueError: If ``openenv.yaml`` is not found or lacks an ``app`` field.
|
| 266 |
+
"""
|
| 267 |
+
yaml_path = self._find_openenv_yaml()
|
| 268 |
+
if yaml_path is None:
|
| 269 |
+
raise ValueError(
|
| 270 |
+
"Could not find openenv.yaml inside the sandbox. "
|
| 271 |
+
"Pass an explicit cmd= to ModalProvider or start_container()."
|
| 272 |
+
)
|
| 273 |
+
|
| 274 |
+
content = self._adapter.exec(self._sandbox, f"cat {shlex.quote(yaml_path)}")
|
| 275 |
+
app = self._parse_app_field(content)
|
| 276 |
+
if app is None:
|
| 277 |
+
raise ValueError(
|
| 278 |
+
f"openenv.yaml at {yaml_path} does not contain an 'app' field. "
|
| 279 |
+
"Pass an explicit cmd= to ModalProvider or start_container()."
|
| 280 |
+
)
|
| 281 |
+
|
| 282 |
+
# The directory containing openenv.yaml is the env root
|
| 283 |
+
env_root = yaml_path.rsplit("/", 1)[0]
|
| 284 |
+
return (
|
| 285 |
+
f"cd {shlex.quote(env_root)} && "
|
| 286 |
+
f"python -m uvicorn {shlex.quote(app)} --host 0.0.0.0 --port {port}"
|
| 287 |
+
)
|
| 288 |
+
|
| 289 |
+
def _find_openenv_yaml(self) -> str | None:
|
| 290 |
+
"""Locate ``openenv.yaml`` inside the sandbox.
|
| 291 |
+
|
| 292 |
+
Tries the modern layout path ``/app/env/openenv.yaml`` first,
|
| 293 |
+
then falls back to a ``find`` command for the old layout.
|
| 294 |
+
"""
|
| 295 |
+
# Fast path: modern Dockerfile layout
|
| 296 |
+
out = self._adapter.exec(
|
| 297 |
+
self._sandbox, "test -f /app/env/openenv.yaml && echo found"
|
| 298 |
+
)
|
| 299 |
+
if "found" in (out or ""):
|
| 300 |
+
return "/app/env/openenv.yaml"
|
| 301 |
+
|
| 302 |
+
# Fallback: search for it (redirect stderr so error messages
|
| 303 |
+
# like "No such file or directory" don't get mistaken for paths).
|
| 304 |
+
path = self._adapter.exec(
|
| 305 |
+
self._sandbox,
|
| 306 |
+
"find /app -maxdepth 4 -name openenv.yaml -print -quit 2>/dev/null",
|
| 307 |
+
).strip()
|
| 308 |
+
if path and path.startswith("/"):
|
| 309 |
+
return path
|
| 310 |
+
|
| 311 |
+
return None
|
| 312 |
+
|
| 313 |
+
@staticmethod
|
| 314 |
+
def _parse_app_field(yaml_content: str) -> str | None:
|
| 315 |
+
"""Extract the ``app`` value from raw openenv.yaml content.
|
| 316 |
+
|
| 317 |
+
Uses PyYAML to handle comments, quotes, and nested keys correctly.
|
| 318 |
+
"""
|
| 319 |
+
return parse_openenv_app_field(yaml_content)
|
| 320 |
+
|
| 321 |
+
@staticmethod
|
| 322 |
+
def _parse_dockerfile_cmd(dockerfile_content: str) -> str | None:
|
| 323 |
+
"""Extract the server command from the last ``CMD`` in a Dockerfile.
|
| 324 |
+
|
| 325 |
+
Handles exec form (``CMD ["prog", "arg"]``) and shell form
|
| 326 |
+
(``CMD prog arg``). When a Dockerfile has multiple ``CMD``
|
| 327 |
+
instructions (e.g. multi-stage builds), the last one wins - same
|
| 328 |
+
semantics as Docker itself. Lines where ``CMD`` appears inside a
|
| 329 |
+
comment are ignored.
|
| 330 |
+
|
| 331 |
+
Returns:
|
| 332 |
+
The command as a single string, or ``None`` if no ``CMD`` found.
|
| 333 |
+
"""
|
| 334 |
+
return parse_dockerfile_cmd(dockerfile_content)
|
| 335 |
+
|
| 336 |
+
@classmethod
|
| 337 |
+
def image_from_dockerfile(
|
| 338 |
+
cls,
|
| 339 |
+
dockerfile_path: str,
|
| 340 |
+
context_dir: str | None = None,
|
| 341 |
+
) -> str:
|
| 342 |
+
"""Validate a Dockerfile and return a ``dockerfile:`` URI for
|
| 343 |
+
:meth:`start_container`.
|
| 344 |
+
|
| 345 |
+
Eagerly validates the Dockerfile (existence, COPY sources) and stores
|
| 346 |
+
its content in an internal registry. The actual ``modal.Image`` is
|
| 347 |
+
created later inside ``start_container``.
|
| 348 |
+
|
| 349 |
+
Args:
|
| 350 |
+
dockerfile_path (`str`):
|
| 351 |
+
Path to the Dockerfile on disk.
|
| 352 |
+
context_dir (`str`, *optional*):
|
| 353 |
+
Build context directory. Defaults to the Dockerfile's
|
| 354 |
+
grandparent directory, matching the ``openenv init``
|
| 355 |
+
convention where Dockerfiles live in
|
| 356 |
+
``<env>/server/Dockerfile`` and the build context is
|
| 357 |
+
``<env>/``. Pass explicitly for non-standard layouts
|
| 358 |
+
(e.g. ``context_dir="."`` for repo-root contexts).
|
| 359 |
+
|
| 360 |
+
Returns:
|
| 361 |
+
`str`: A ``"dockerfile:<abs_path>"`` string to pass to
|
| 362 |
+
``start_container``.
|
| 363 |
+
|
| 364 |
+
Raises:
|
| 365 |
+
FileNotFoundError: If *dockerfile_path* does not exist.
|
| 366 |
+
ValueError: If *context_dir* is given but does not exist,
|
| 367 |
+
or if COPY sources in the Dockerfile cannot be found
|
| 368 |
+
under the resolved context directory.
|
| 369 |
+
"""
|
| 370 |
+
import pathlib
|
| 371 |
+
import re
|
| 372 |
+
|
| 373 |
+
src = pathlib.Path(dockerfile_path).resolve()
|
| 374 |
+
if not src.is_file():
|
| 375 |
+
raise FileNotFoundError(f"Dockerfile not found: {dockerfile_path}")
|
| 376 |
+
|
| 377 |
+
if context_dir is not None:
|
| 378 |
+
ctx = pathlib.Path(context_dir)
|
| 379 |
+
if not ctx.is_dir():
|
| 380 |
+
raise ValueError(f"context_dir does not exist: {context_dir}")
|
| 381 |
+
else:
|
| 382 |
+
# Default: grandparent of the Dockerfile, matching the
|
| 383 |
+
# openenv init layout (<env>/server/Dockerfile -> <env>/).
|
| 384 |
+
ctx = src.parent.parent
|
| 385 |
+
|
| 386 |
+
content = src.read_text()
|
| 387 |
+
|
| 388 |
+
# Validate that COPY sources exist under the context directory.
|
| 389 |
+
# This catches mismatches early (e.g. a Dockerfile expecting repo
|
| 390 |
+
# root as context when we defaulted to the env directory).
|
| 391 |
+
for line in content.splitlines():
|
| 392 |
+
m = re.match(r"^\s*COPY\s+(?!--from=)(\S+)\s+", line, re.IGNORECASE)
|
| 393 |
+
if not m:
|
| 394 |
+
continue
|
| 395 |
+
copy_src = m.group(1)
|
| 396 |
+
if copy_src.startswith("/"):
|
| 397 |
+
continue
|
| 398 |
+
resolved = ctx / copy_src
|
| 399 |
+
if not resolved.exists() and not any(ctx.glob(copy_src)):
|
| 400 |
+
raise ValueError(
|
| 401 |
+
f"Dockerfile COPY source '{copy_src}' not found "
|
| 402 |
+
f"under context_dir '{ctx}'. This Dockerfile may "
|
| 403 |
+
f"expect a different build context (e.g. the repo "
|
| 404 |
+
f"root). Pass context_dir explicitly."
|
| 405 |
+
)
|
| 406 |
+
|
| 407 |
+
# Parse CMD from the Dockerfile so start_container can use it as a
|
| 408 |
+
# fallback when openenv.yaml is unavailable.
|
| 409 |
+
parsed_cmd = cls._parse_dockerfile_cmd(content)
|
| 410 |
+
|
| 411 |
+
cls._dockerfile_registry[str(src)] = {
|
| 412 |
+
"dockerfile_path": str(src),
|
| 413 |
+
"context_dir": str(ctx),
|
| 414 |
+
"server_cmd": parsed_cmd,
|
| 415 |
+
}
|
| 416 |
+
|
| 417 |
+
return f"dockerfile:{src}"
|
| 418 |
+
|
| 419 |
+
def _build_image(self, image: str) -> Any:
|
| 420 |
+
"""Build the ``modal.Image`` for *image* (registry tag or dockerfile:)."""
|
| 421 |
+
if image.startswith("dockerfile:"):
|
| 422 |
+
dockerfile_path = image[len("dockerfile:") :]
|
| 423 |
+
meta = self._dockerfile_registry.get(dockerfile_path)
|
| 424 |
+
if meta is None:
|
| 425 |
+
raise ValueError(
|
| 426 |
+
f"No registered Dockerfile metadata for {dockerfile_path}. "
|
| 427 |
+
"Call ModalProvider.image_from_dockerfile() first."
|
| 428 |
+
)
|
| 429 |
+
return self._adapter.image_from_dockerfile(
|
| 430 |
+
meta["dockerfile_path"], meta["context_dir"]
|
| 431 |
+
)
|
| 432 |
+
|
| 433 |
+
# Plain registry tag (e.g. "echo-env:latest").
|
| 434 |
+
return self._adapter.image_from_registry(image)
|
| 435 |
+
|
| 436 |
+
def start_container(
|
| 437 |
+
self,
|
| 438 |
+
image: str | None = None,
|
| 439 |
+
port: int | None = None,
|
| 440 |
+
env_vars: dict[str, str] | None = None,
|
| 441 |
+
**kwargs: Any,
|
| 442 |
+
) -> str:
|
| 443 |
+
"""
|
| 444 |
+
Create a Modal sandbox from a Docker image or Dockerfile.
|
| 445 |
+
|
| 446 |
+
The sandbox is started with a keep-alive process and the server
|
| 447 |
+
command is launched via ``exec`` afterwards, mirroring the discovery
|
| 448 |
+
flow used by other cloud providers. The server command is resolved in
|
| 449 |
+
order:
|
| 450 |
+
|
| 451 |
+
1. Explicit ``cmd`` passed to the constructor.
|
| 452 |
+
2. ``cmd`` key in ``**kwargs`` (popped before forwarding).
|
| 453 |
+
3. Auto-discovered from ``openenv.yaml`` inside the sandbox.
|
| 454 |
+
4. ``CMD`` parsed from the Dockerfile (when *image* came from
|
| 455 |
+
``image_from_dockerfile``).
|
| 456 |
+
|
| 457 |
+
Args:
|
| 458 |
+
image (`str`, *optional*):
|
| 459 |
+
Registry image tag (e.g. ``"echo-env:latest"``) or
|
| 460 |
+
``"dockerfile:<path>"`` returned by
|
| 461 |
+
:meth:`image_from_dockerfile`. May be omitted when supplied to
|
| 462 |
+
the constructor.
|
| 463 |
+
port (`int`, *optional*):
|
| 464 |
+
Must be ``None`` or ``8000``. Modal exposes port 8000 via an
|
| 465 |
+
encrypted tunnel; other ports raise ``ValueError``.
|
| 466 |
+
env_vars (`dict`, *optional*):
|
| 467 |
+
Environment variables forwarded to the sandbox.
|
| 468 |
+
**kwargs:
|
| 469 |
+
``cmd`` (`str`) to override the server command; any remaining
|
| 470 |
+
keyword arguments are forwarded to ``modal.Sandbox.create``.
|
| 471 |
+
|
| 472 |
+
Returns:
|
| 473 |
+
`str`: HTTPS tunnel URL for the sandbox (base_url).
|
| 474 |
+
"""
|
| 475 |
+
if self._sandbox is not None:
|
| 476 |
+
raise RuntimeError(
|
| 477 |
+
"ModalProvider already has an active sandbox. Call "
|
| 478 |
+
"stop_container() (or close()) before starting another — a "
|
| 479 |
+
"second start would orphan the running sandbox."
|
| 480 |
+
)
|
| 481 |
+
|
| 482 |
+
if port is not None and port != _DEFAULT_MODAL_PORT:
|
| 483 |
+
raise ValueError(
|
| 484 |
+
f"ModalProvider only supports port {_DEFAULT_MODAL_PORT} "
|
| 485 |
+
f"(got {port}). The Modal tunnel routes to port "
|
| 486 |
+
f"{_DEFAULT_MODAL_PORT} inside the sandbox."
|
| 487 |
+
)
|
| 488 |
+
|
| 489 |
+
effective_image = image if image is not None else self._image
|
| 490 |
+
if effective_image is None:
|
| 491 |
+
raise ValueError(
|
| 492 |
+
"ModalProvider requires an image. Pass it to the constructor "
|
| 493 |
+
"or start_container()."
|
| 494 |
+
)
|
| 495 |
+
effective_env_vars = self._env_vars if env_vars is None else env_vars
|
| 496 |
+
|
| 497 |
+
# Resolve the server command (may be None; discovery happens after
|
| 498 |
+
# sandbox creation when we can inspect the filesystem).
|
| 499 |
+
cmd = kwargs.pop("cmd", None) or self._cmd
|
| 500 |
+
|
| 501 |
+
# CMD parsed from Dockerfile (populated for "dockerfile:" images).
|
| 502 |
+
parsed_cmd: str | None = None
|
| 503 |
+
if effective_image.startswith("dockerfile:"):
|
| 504 |
+
meta = self._dockerfile_registry.get(effective_image[len("dockerfile:") :])
|
| 505 |
+
if meta is not None:
|
| 506 |
+
parsed_cmd = meta.get("server_cmd")
|
| 507 |
+
|
| 508 |
+
modal_image = self._build_image(effective_image)
|
| 509 |
+
|
| 510 |
+
extra: Dict[str, Any] = dict(kwargs)
|
| 511 |
+
if self._cpu is not None:
|
| 512 |
+
extra["cpu"] = (_DEFAULT_CPU_REQUEST, self._cpu)
|
| 513 |
+
if self._memory is not None:
|
| 514 |
+
extra["memory"] = (_DEFAULT_MEMORY_REQUEST, self._memory)
|
| 515 |
+
|
| 516 |
+
# Record injected secret values so captured server output can be
|
| 517 |
+
# scrubbed before it is ever surfaced in an error.
|
| 518 |
+
self._redact_values = {
|
| 519 |
+
value for value in (effective_env_vars or {}).values() if value
|
| 520 |
+
}
|
| 521 |
+
|
| 522 |
+
# A create failure created nothing, so just drop the recorded secrets
|
| 523 |
+
# and re-raise (the double-start guard above guarantees there is no
|
| 524 |
+
# pre-existing sandbox to delete).
|
| 525 |
+
try:
|
| 526 |
+
self._sandbox = self._adapter.create_sandbox(
|
| 527 |
+
image=modal_image,
|
| 528 |
+
encrypted_ports=[_DEFAULT_MODAL_PORT],
|
| 529 |
+
timeout=self._timeout,
|
| 530 |
+
env=effective_env_vars,
|
| 531 |
+
extra=extra,
|
| 532 |
+
)
|
| 533 |
+
except Exception:
|
| 534 |
+
self._redact_values = set()
|
| 535 |
+
raise
|
| 536 |
+
|
| 537 |
+
try:
|
| 538 |
+
# Discover server command from openenv.yaml if not explicitly set.
|
| 539 |
+
if cmd is None:
|
| 540 |
+
try:
|
| 541 |
+
cmd = self._discover_server_cmd()
|
| 542 |
+
except ValueError:
|
| 543 |
+
# Fall back to CMD parsed from Dockerfile (if available).
|
| 544 |
+
if parsed_cmd:
|
| 545 |
+
cmd = parsed_cmd
|
| 546 |
+
else:
|
| 547 |
+
raise
|
| 548 |
+
|
| 549 |
+
# Launch the server in the background. Write the PID so we can
|
| 550 |
+
# check whether the process crashed in wait_for_ready().
|
| 551 |
+
escaped_cmd = shlex.quote(cmd)
|
| 552 |
+
self._adapter.exec(
|
| 553 |
+
self._sandbox,
|
| 554 |
+
f"nohup bash -c {escaped_cmd} > /tmp/openenv-server.log 2>&1 &"
|
| 555 |
+
" echo $! > /tmp/openenv-server.pid",
|
| 556 |
+
)
|
| 557 |
+
|
| 558 |
+
# Resolve the public tunnel URL for port 8000.
|
| 559 |
+
self._base_url = _require_secure_url(
|
| 560 |
+
self._adapter.tunnel_url(self._sandbox, _DEFAULT_MODAL_PORT)
|
| 561 |
+
)
|
| 562 |
+
except Exception:
|
| 563 |
+
# A cleanup failure here must not mask the original error: swallow
|
| 564 |
+
# any exception from stop_container() so the root cause propagates.
|
| 565 |
+
try:
|
| 566 |
+
self.stop_container()
|
| 567 |
+
except Exception:
|
| 568 |
+
pass
|
| 569 |
+
raise
|
| 570 |
+
|
| 571 |
+
return self._base_url
|
| 572 |
+
|
| 573 |
+
def stop_container(self) -> None:
|
| 574 |
+
"""Terminate the Modal sandbox."""
|
| 575 |
+
if self._sandbox is None:
|
| 576 |
+
# Still drop any injected secret values recorded by a failed start.
|
| 577 |
+
self._redact_values = set()
|
| 578 |
+
return
|
| 579 |
+
|
| 580 |
+
try:
|
| 581 |
+
self._adapter.terminate(self._sandbox)
|
| 582 |
+
finally:
|
| 583 |
+
self._sandbox = None
|
| 584 |
+
self._base_url = None
|
| 585 |
+
self._redact_values = set()
|
| 586 |
+
|
| 587 |
+
def close(self) -> None:
|
| 588 |
+
"""Stop the active sandbox.
|
| 589 |
+
|
| 590 |
+
Overrides the base no-op so a caller holding a bare ``ContainerProvider``
|
| 591 |
+
reference can release the sandbox polymorphically (also invoked on
|
| 592 |
+
context-manager exit). ``ModalProvider`` holds no separate SDK client, so
|
| 593 |
+
this is equivalent to ``stop_container()``.
|
| 594 |
+
"""
|
| 595 |
+
self.stop_container()
|
| 596 |
+
|
| 597 |
+
@property
|
| 598 |
+
def base_url(self) -> str:
|
| 599 |
+
"""URL returned by the last ``start_container``."""
|
| 600 |
+
if self._base_url is None:
|
| 601 |
+
raise RuntimeError(
|
| 602 |
+
"ModalProvider has no active base_url. Start the provider "
|
| 603 |
+
"before reading base_url."
|
| 604 |
+
)
|
| 605 |
+
return self._base_url
|
| 606 |
+
|
| 607 |
+
def _redact(self, text: str, *, max_chars: int = 2000) -> str:
|
| 608 |
+
"""Scrub injected secret values and bound length before surfacing output.
|
| 609 |
+
|
| 610 |
+
Replaces any injected env-var value with `***` and keeps only the tail.
|
| 611 |
+
This is best-effort (exact-match only), which is why server output is
|
| 612 |
+
withheld entirely unless ``surface_server_logs=True``.
|
| 613 |
+
"""
|
| 614 |
+
redacted = text or ""
|
| 615 |
+
for value in self._redact_values:
|
| 616 |
+
redacted = redacted.replace(value, "***")
|
| 617 |
+
if len(redacted) > max_chars:
|
| 618 |
+
redacted = "...(truncated)...\n" + redacted[-max_chars:]
|
| 619 |
+
return redacted
|
| 620 |
+
|
| 621 |
+
def _server_died_message(self) -> str:
|
| 622 |
+
"""Build the startup-crash error, secure by default.
|
| 623 |
+
|
| 624 |
+
Untrusted code can print secrets then force a crash to exfiltrate them
|
| 625 |
+
through the exception (which lands in orchestrator/CI logs), so sandbox
|
| 626 |
+
output is excluded unless ``surface_server_logs=True``, in which case a
|
| 627 |
+
best-effort redacted, bounded excerpt is included.
|
| 628 |
+
"""
|
| 629 |
+
base = (
|
| 630 |
+
"Modal sandbox server process died during startup. Server output is "
|
| 631 |
+
"not surfaced to avoid leaking secrets injected into the sandbox; "
|
| 632 |
+
"retrieve /tmp/openenv-server.log from the sandbox out of band, or "
|
| 633 |
+
"construct the provider with surface_server_logs=True to include a "
|
| 634 |
+
"redacted excerpt."
|
| 635 |
+
)
|
| 636 |
+
if not self._surface_server_logs or self._sandbox is None:
|
| 637 |
+
return base
|
| 638 |
+
|
| 639 |
+
log = self._redact(
|
| 640 |
+
self._adapter.exec(self._sandbox, "cat /tmp/openenv-server.log 2>/dev/null")
|
| 641 |
+
)
|
| 642 |
+
return (
|
| 643 |
+
"Modal sandbox server process died during startup. The excerpt below "
|
| 644 |
+
"is the sandbox server output with injected secret values redacted "
|
| 645 |
+
"(best-effort); it may still contain secrets the workload printed "
|
| 646 |
+
f"by other means.\nLog (redacted):\n{log}"
|
| 647 |
+
)
|
| 648 |
+
|
| 649 |
+
def wait_for_ready(self, base_url: str, timeout_s: float = 120.0) -> None:
|
| 650 |
+
"""
|
| 651 |
+
Poll the /health endpoint until the sandbox is ready.
|
| 652 |
+
|
| 653 |
+
Uses a longer default timeout (120s) than local Docker providers
|
| 654 |
+
because Modal sandboxes may have cold-start latency.
|
| 655 |
+
|
| 656 |
+
Args:
|
| 657 |
+
base_url (`str`):
|
| 658 |
+
Tunnel URL returned by ``start_container()``.
|
| 659 |
+
timeout_s (`float`, *optional*, defaults to `120.0`):
|
| 660 |
+
Maximum seconds to wait.
|
| 661 |
+
|
| 662 |
+
Raises:
|
| 663 |
+
TimeoutError: If the sandbox doesn't become ready in time.
|
| 664 |
+
RuntimeError: If the server process died (detected via PID check).
|
| 665 |
+
"""
|
| 666 |
+
import requests
|
| 667 |
+
|
| 668 |
+
health_url = f"{base_url}/health"
|
| 669 |
+
|
| 670 |
+
deadline = time.time() + timeout_s
|
| 671 |
+
while time.time() < deadline:
|
| 672 |
+
try:
|
| 673 |
+
response = requests.get(health_url, timeout=5.0)
|
| 674 |
+
if response.status_code == 200:
|
| 675 |
+
return
|
| 676 |
+
except requests.RequestException:
|
| 677 |
+
pass
|
| 678 |
+
|
| 679 |
+
# Early exit: if the server process died, raise immediately
|
| 680 |
+
# instead of waiting for the full health-check timeout.
|
| 681 |
+
if self._sandbox is not None:
|
| 682 |
+
out = self._adapter.exec(
|
| 683 |
+
self._sandbox,
|
| 684 |
+
"kill -0 $(cat /tmp/openenv-server.pid) 2>/dev/null"
|
| 685 |
+
" && echo RUNNING || echo DEAD",
|
| 686 |
+
)
|
| 687 |
+
if "DEAD" in (out or ""):
|
| 688 |
+
raise RuntimeError(self._server_died_message())
|
| 689 |
+
|
| 690 |
+
time.sleep(1.0)
|
| 691 |
+
|
| 692 |
+
# The tunnel URL is a bearer capability, so it is deliberately omitted
|
| 693 |
+
# from the timeout error.
|
| 694 |
+
raise TimeoutError(f"Modal sandbox did not become ready within {timeout_s}s.")
|
| 695 |
+
|
| 696 |
+
|
| 697 |
+
__all__ = ["ModalProvider"]
|
src/core/containers/runtime/novita_provider.py
ADDED
|
@@ -0,0 +1,1216 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# SPDX-License-Identifier: BSD-3-Clause
|
| 2 |
+
|
| 3 |
+
"""
|
| 4 |
+
Novita AI Sandbox container provider for running OpenEnv environments.
|
| 5 |
+
|
| 6 |
+
Requires the ``novita-sandbox`` SDK: ``pip install openenv[novita]``
|
| 7 |
+
|
| 8 |
+
The provider boots an OpenEnv server inside a Novita sandbox, exposes port 8000
|
| 9 |
+
on the sandbox's public host, and returns an ``https://`` URL that ``EnvClient``
|
| 10 |
+
connects to over ``wss://``.
|
| 11 |
+
|
| 12 |
+
Note: the ``RFC 002 security invariant S<n>`` references in this module point to
|
| 13 |
+
the **proposed** "Cloud Sandbox Providers" amendment to RFC 002, which is pending
|
| 14 |
+
review/sign-off by the RFC authors. They describe the security properties this
|
| 15 |
+
provider enforces; treat them as proposed (not yet ratified) until that sign-off.
|
| 16 |
+
"""
|
| 17 |
+
|
| 18 |
+
from __future__ import annotations
|
| 19 |
+
|
| 20 |
+
import re
|
| 21 |
+
import shlex
|
| 22 |
+
import time
|
| 23 |
+
from typing import Any, Callable, Dict, List, Optional
|
| 24 |
+
|
| 25 |
+
from ._server_config import parse_openenv_app_field
|
| 26 |
+
from .providers import ContainerProvider
|
| 27 |
+
|
| 28 |
+
_DEFAULT_NOVITA_PORT = 8000
|
| 29 |
+
|
| 30 |
+
# Novita defaults a sandbox to a 300s hard lifetime, which would kill an RL
|
| 31 |
+
# rollout mid-episode. One hour is the documented ceiling for Hobby accounts, so
|
| 32 |
+
# it is the largest default that cannot make ``create`` fail on plan limits;
|
| 33 |
+
# longer-running work passes ``timeout=`` explicitly.
|
| 34 |
+
_DEFAULT_SANDBOX_TIMEOUT_S = 3600
|
| 35 |
+
|
| 36 |
+
# The sandbox is created with this as its start command so the image's own CMD
|
| 37 |
+
# is not launched (see `_launch_server`). `sleep infinity` keeps the sandbox
|
| 38 |
+
# alive with nothing bound to port 8000, leaving the port free for the server
|
| 39 |
+
# this provider starts itself.
|
| 40 |
+
_KEEPALIVE_CMD = "sleep infinity"
|
| 41 |
+
_KEEPALIVE_READY_MS = 5_000
|
| 42 |
+
|
| 43 |
+
# Resource defaults for a template built from a Dockerfile. The SDK's own
|
| 44 |
+
# defaults, restated here so they are visible at the call site.
|
| 45 |
+
_DEFAULT_TEMPLATE_CPU = 2
|
| 46 |
+
_DEFAULT_TEMPLATE_MEMORY_MB = 1024
|
| 47 |
+
|
| 48 |
+
# User the sandbox runs as. Root because the environment installs under /app,
|
| 49 |
+
# which the image creates as root and the Novita parser's default user ("user")
|
| 50 |
+
# cannot write to -- see `build_template`. Matches what a plain `docker run` of
|
| 51 |
+
# the same image gives, which is what the other providers' users observe.
|
| 52 |
+
_SANDBOX_USER = "root"
|
| 53 |
+
|
| 54 |
+
_FROM_RE = re.compile(r"^\s*FROM\s+(?P<rest>.+?)\s*$", re.IGNORECASE)
|
| 55 |
+
_COPY_FROM_RE = re.compile(
|
| 56 |
+
r"^\s*COPY\s+--from=(?P<stage>\S+)\s+(?P<src>\S+)\s+(?P<dst>\S+)\s*$",
|
| 57 |
+
re.IGNORECASE,
|
| 58 |
+
)
|
| 59 |
+
_COPY_FROM_REMAINDER_RE = re.compile(r"^\s*COPY\b.*\s--from(?:=|\s|$)", re.IGNORECASE)
|
| 60 |
+
_USER_RE = re.compile(r"^\s*USER(?:\s+|$)", re.IGNORECASE)
|
| 61 |
+
_RUN_RE = re.compile(r"^\s*RUN\s+(?P<rest>.*)$", re.IGNORECASE)
|
| 62 |
+
_ARG_DEFAULT_RE = re.compile(
|
| 63 |
+
r"^\s*ARG\s+(?P<name>\w+)=(?P<value>\S+)\s*$", re.IGNORECASE
|
| 64 |
+
)
|
| 65 |
+
|
| 66 |
+
|
| 67 |
+
def _strip_mount_flags(content: str) -> str:
|
| 68 |
+
"""Remove supported BuildKit cache mounts from ``RUN`` instructions.
|
| 69 |
+
|
| 70 |
+
Novita's template parser copies a ``RUN`` body verbatim into the build
|
| 71 |
+
command, so a leading ``--mount=type=cache,...`` would be handed to the
|
| 72 |
+
shell as an argument rather than honored as BuildKit syntax. Every in-repo
|
| 73 |
+
OpenEnv Dockerfile uses one to cache ``uv`` downloads, so explicit cache
|
| 74 |
+
mounts are stripped rather than left to fail the build. Other mount types
|
| 75 |
+
affect build semantics and cannot be represented by Novita's parser, so
|
| 76 |
+
they are rejected instead of silently removed.
|
| 77 |
+
"""
|
| 78 |
+
lines = content.split("\n")
|
| 79 |
+
out: List[str] = []
|
| 80 |
+
in_run = False
|
| 81 |
+
for line in lines:
|
| 82 |
+
if _RUN_RE.match(line):
|
| 83 |
+
in_run = True
|
| 84 |
+
if in_run:
|
| 85 |
+
# Strip leading --mount flags from the RUN line itself and from any
|
| 86 |
+
# continuation lines that follow it.
|
| 87 |
+
stripped = line
|
| 88 |
+
prefix = ""
|
| 89 |
+
run_match = _RUN_RE.match(line)
|
| 90 |
+
if run_match:
|
| 91 |
+
prefix = line[: run_match.start("rest")]
|
| 92 |
+
stripped = run_match.group("rest")
|
| 93 |
+
while True:
|
| 94 |
+
mount = re.match(r"\s*--mount=(?P<spec>\S+)(?P<space>\s*)", stripped)
|
| 95 |
+
if not mount:
|
| 96 |
+
break
|
| 97 |
+
spec = mount.group("spec")
|
| 98 |
+
options = spec.split(",")
|
| 99 |
+
types = [
|
| 100 |
+
option.split("=", 1)[1].lower()
|
| 101 |
+
for option in options
|
| 102 |
+
if option.lower().startswith("type=")
|
| 103 |
+
]
|
| 104 |
+
if types != ["cache"]:
|
| 105 |
+
raise ValueError(
|
| 106 |
+
"Novita Dockerfile rewriting only supports explicit "
|
| 107 |
+
"RUN --mount=type=cache options; unsupported mount "
|
| 108 |
+
"types must be handled by `openenv build` and a "
|
| 109 |
+
"registry image."
|
| 110 |
+
)
|
| 111 |
+
stripped = stripped[mount.end() :]
|
| 112 |
+
line = prefix + stripped
|
| 113 |
+
if not line.rstrip().endswith("\\"):
|
| 114 |
+
in_run = False
|
| 115 |
+
out.append(line)
|
| 116 |
+
return "\n".join(out)
|
| 117 |
+
|
| 118 |
+
|
| 119 |
+
def _resolve_from_references(content: str) -> str:
|
| 120 |
+
"""Substitute ``ARG`` defaults into ``FROM`` lines and drop ``--platform``.
|
| 121 |
+
|
| 122 |
+
Two parser limitations, handled together because both corrupt the base
|
| 123 |
+
image reference:
|
| 124 |
+
|
| 125 |
+
- The parser stores a ``FROM ${BASE_IMAGE}`` line verbatim, so the template
|
| 126 |
+
would be built from a literal image named ``${BASE_IMAGE}``. Only global
|
| 127 |
+
ARGs (declared before the first ``FROM``) are in scope for a ``FROM``
|
| 128 |
+
line, which is the form every in-repo Dockerfile uses.
|
| 129 |
+
- ``FROM --platform=linux/amd64 python:3.10-slim`` likewise keeps the flag
|
| 130 |
+
as part of the name. Novita builds for its own platform, so the flag is
|
| 131 |
+
dropped rather than propagated.
|
| 132 |
+
|
| 133 |
+
Substitutions are applied longest-name-first so a short unbraced prefix
|
| 134 |
+
cannot corrupt a longer sibling (e.g. ``$BASE`` must not rewrite
|
| 135 |
+
``$BASE_IMAGE`` / ``${BASE_IMAGE}`` into ``python:3.12_IMAGE``).
|
| 136 |
+
"""
|
| 137 |
+
arg_defaults: Dict[str, str] = {}
|
| 138 |
+
for line in content.split("\n"):
|
| 139 |
+
if _FROM_RE.match(line):
|
| 140 |
+
break
|
| 141 |
+
match = _ARG_DEFAULT_RE.match(line)
|
| 142 |
+
if match:
|
| 143 |
+
arg_defaults[match.group("name")] = match.group("value")
|
| 144 |
+
|
| 145 |
+
# Longest names first: `$BASE` must not rewrite `$BASE_IMAGE` mid-token.
|
| 146 |
+
arg_items = sorted(
|
| 147 |
+
arg_defaults.items(), key=lambda item: len(item[0]), reverse=True
|
| 148 |
+
)
|
| 149 |
+
|
| 150 |
+
out: List[str] = []
|
| 151 |
+
for line in content.split("\n"):
|
| 152 |
+
match = _FROM_RE.match(line)
|
| 153 |
+
if not match:
|
| 154 |
+
out.append(line)
|
| 155 |
+
continue
|
| 156 |
+
|
| 157 |
+
reference = match.group("rest").strip()
|
| 158 |
+
reference = re.sub(r"^(--platform=\S+\s*)+", "", reference).strip()
|
| 159 |
+
for name, value in arg_items:
|
| 160 |
+
reference = reference.replace(f"${{{name}}}", value)
|
| 161 |
+
reference = reference.replace(f"${name}", value)
|
| 162 |
+
out.append(f"FROM {reference}")
|
| 163 |
+
return "\n".join(out)
|
| 164 |
+
|
| 165 |
+
|
| 166 |
+
def _split_stages(content: str) -> List[Dict[str, Any]]:
|
| 167 |
+
"""Split a Dockerfile into its ``FROM``-delimited stages."""
|
| 168 |
+
stages: List[Dict[str, Any]] = []
|
| 169 |
+
current: Optional[Dict[str, Any]] = None
|
| 170 |
+
for line in content.split("\n"):
|
| 171 |
+
match = _FROM_RE.match(line)
|
| 172 |
+
if match:
|
| 173 |
+
rest = match.group("rest").strip()
|
| 174 |
+
alias_match = re.match(r"(?P<base>.+?)\s+[Aa][Ss]\s+(?P<alias>\S+)$", rest)
|
| 175 |
+
if alias_match:
|
| 176 |
+
base = alias_match.group("base").strip()
|
| 177 |
+
alias = alias_match.group("alias")
|
| 178 |
+
else:
|
| 179 |
+
base, alias = rest, None
|
| 180 |
+
current = {"base": base, "alias": alias, "body": []}
|
| 181 |
+
stages.append(current)
|
| 182 |
+
elif current is not None:
|
| 183 |
+
current["body"].append(line)
|
| 184 |
+
return stages
|
| 185 |
+
|
| 186 |
+
|
| 187 |
+
def _dockerfile_declares_user(content: str) -> bool:
|
| 188 |
+
"""Return whether *content* contains a Dockerfile ``USER`` instruction."""
|
| 189 |
+
return any(_USER_RE.match(line) for line in content.splitlines())
|
| 190 |
+
|
| 191 |
+
|
| 192 |
+
def _reject_unflattened_copy_from(content: str) -> None:
|
| 193 |
+
"""Reject COPY instructions with ``--from`` left after flattening."""
|
| 194 |
+
if any(_COPY_FROM_REMAINDER_RE.match(line) for line in content.splitlines()):
|
| 195 |
+
raise ValueError(
|
| 196 |
+
"Dockerfile contains a COPY --from instruction that Novita cannot "
|
| 197 |
+
"flatten. Use a simple `COPY --from=<stage> <src> <dst>` form or "
|
| 198 |
+
"build the image with `openenv build` and pass the registry image."
|
| 199 |
+
)
|
| 200 |
+
|
| 201 |
+
|
| 202 |
+
def _flatten_multistage(content: str) -> str:
|
| 203 |
+
"""Collapse a multi-stage Dockerfile into a single stage.
|
| 204 |
+
|
| 205 |
+
Novita's parser rejects multi-stage Dockerfiles outright, and that is the
|
| 206 |
+
layout every in-repo OpenEnv environment uses: a ``builder`` stage that runs
|
| 207 |
+
``uv sync`` into ``/app/env/.venv``, then a runtime stage that copies the
|
| 208 |
+
result across. Both stages name the *same* base image, so the build can be
|
| 209 |
+
replayed linearly in one stage -- the builder's ``RUN`` steps execute, then
|
| 210 |
+
each ``COPY --from=builder`` becomes an in-place ``cp``.
|
| 211 |
+
|
| 212 |
+
A copy whose source and destination are equal is dropped: after flattening,
|
| 213 |
+
the file is already at its destination, and Novita rejects an absolute
|
| 214 |
+
source path.
|
| 215 |
+
|
| 216 |
+
Raises:
|
| 217 |
+
ValueError: If the stages do not share one base image, a copy pulls from
|
| 218 |
+
a stage other than the builder, or a ``COPY --from`` cannot be
|
| 219 |
+
flattened.
|
| 220 |
+
"""
|
| 221 |
+
stages = _split_stages(content)
|
| 222 |
+
if len(stages) < 2:
|
| 223 |
+
_reject_unflattened_copy_from(content)
|
| 224 |
+
return content
|
| 225 |
+
|
| 226 |
+
bases = {stage["base"] for stage in stages}
|
| 227 |
+
if len(bases) != 1:
|
| 228 |
+
raise ValueError(
|
| 229 |
+
"Novita templates cannot express a multi-stage Dockerfile whose "
|
| 230 |
+
f"stages use different base images: {sorted(bases)}. Build the "
|
| 231 |
+
"image with `openenv build` and push it with "
|
| 232 |
+
"`openenv push --registry`, then pass the registry reference to "
|
| 233 |
+
"NovitaSandboxProvider(image=...)."
|
| 234 |
+
)
|
| 235 |
+
|
| 236 |
+
builder_alias = stages[0]["alias"]
|
| 237 |
+
if builder_alias is None:
|
| 238 |
+
raise ValueError(
|
| 239 |
+
"The first stage of a multi-stage Dockerfile must be named with "
|
| 240 |
+
"`AS <alias>` so its output can be referenced after flattening."
|
| 241 |
+
)
|
| 242 |
+
|
| 243 |
+
out: List[str] = [f"FROM {stages[0]['base']}"]
|
| 244 |
+
for index, stage in enumerate(stages):
|
| 245 |
+
for line in stage["body"]:
|
| 246 |
+
copy_match = _COPY_FROM_RE.match(line)
|
| 247 |
+
if copy_match is None:
|
| 248 |
+
out.append(line)
|
| 249 |
+
continue
|
| 250 |
+
if index == 0:
|
| 251 |
+
# A copy within the builder stage is already valid as written.
|
| 252 |
+
out.append(line)
|
| 253 |
+
continue
|
| 254 |
+
source_stage = copy_match.group("stage")
|
| 255 |
+
if source_stage != builder_alias:
|
| 256 |
+
raise ValueError(
|
| 257 |
+
f"COPY --from={source_stage} does not reference the builder "
|
| 258 |
+
f"stage ({builder_alias}), so it cannot be flattened into a "
|
| 259 |
+
"single-stage template."
|
| 260 |
+
)
|
| 261 |
+
src = copy_match.group("src")
|
| 262 |
+
dst = copy_match.group("dst")
|
| 263 |
+
if src == dst:
|
| 264 |
+
continue
|
| 265 |
+
out.append(
|
| 266 |
+
f"RUN mkdir -p $(dirname {shlex.quote(dst)}) && "
|
| 267 |
+
f"cp -a {shlex.quote(src)} {shlex.quote(dst)}"
|
| 268 |
+
)
|
| 269 |
+
|
| 270 |
+
flattened = "\n".join(out)
|
| 271 |
+
_reject_unflattened_copy_from(flattened)
|
| 272 |
+
return flattened
|
| 273 |
+
|
| 274 |
+
|
| 275 |
+
def _prepare_dockerfile(content: str) -> str:
|
| 276 |
+
"""Rewrite an OpenEnv Dockerfile into what Novita's template parser accepts.
|
| 277 |
+
|
| 278 |
+
Applies, in order: BuildKit ``--mount`` stripping, ``ARG``/``--platform``
|
| 279 |
+
resolution in ``FROM``, then multi-stage flattening. Each transform
|
| 280 |
+
compensates for a parser limitation rather than changing the build's
|
| 281 |
+
meaning.
|
| 282 |
+
"""
|
| 283 |
+
return _flatten_multistage(_resolve_from_references(_strip_mount_flags(content)))
|
| 284 |
+
|
| 285 |
+
|
| 286 |
+
def _require_secure_url(url: str) -> str:
|
| 287 |
+
"""Enforce https/wss transport (RFC 002 security invariant S1).
|
| 288 |
+
|
| 289 |
+
``EnvClient`` derives its WebSocket URL from this base URL, so a plaintext
|
| 290 |
+
URL would become a cleartext ``ws://`` connection. The offending URL is
|
| 291 |
+
deliberately omitted from the error: a sandbox host is reachable only
|
| 292 |
+
through the account that owns it, so the address is treated as
|
| 293 |
+
account-scoped rather than pasted into logs.
|
| 294 |
+
"""
|
| 295 |
+
if (
|
| 296 |
+
not isinstance(url, str)
|
| 297 |
+
or not url.lower().startswith("https://")
|
| 298 |
+
or "://" in url[len("https://") :]
|
| 299 |
+
):
|
| 300 |
+
raise RuntimeError(
|
| 301 |
+
"Novita sandbox returned a non-HTTPS host URL. OpenEnv requires an "
|
| 302 |
+
"https/wss base_url so EnvClient traffic is encrypted. Refusing to "
|
| 303 |
+
"connect over plaintext."
|
| 304 |
+
)
|
| 305 |
+
|
| 306 |
+
# With NOVITA_DEBUG set, the SDK's `get_host` returns `localhost:<port>`
|
| 307 |
+
# instead of the sandbox host. Adding the scheme would yield a URL that
|
| 308 |
+
# looks valid and silently targets whatever is listening on the caller's
|
| 309 |
+
# machine, so reject it explicitly rather than connect to the wrong thing.
|
| 310 |
+
hostname = url.split("://", 1)[1].split("/", 1)[0].rsplit(":", 1)[0]
|
| 311 |
+
if hostname in ("localhost", "127.0.0.1", "::1"):
|
| 312 |
+
raise RuntimeError(
|
| 313 |
+
"Novita sandbox returned a loopback host URL, which means the SDK is "
|
| 314 |
+
"running with NOVITA_DEBUG enabled. Debug mode bypasses the sandbox "
|
| 315 |
+
"host and would connect to the local machine instead; unset "
|
| 316 |
+
"NOVITA_DEBUG to reach the sandbox."
|
| 317 |
+
)
|
| 318 |
+
|
| 319 |
+
return url
|
| 320 |
+
|
| 321 |
+
|
| 322 |
+
def _require_bare_host(host: str) -> str:
|
| 323 |
+
"""Require the SDK host value to contain no URL scheme."""
|
| 324 |
+
if not isinstance(host, str) or not host or "://" in host:
|
| 325 |
+
raise RuntimeError(
|
| 326 |
+
"Novita sandbox returned an invalid host. Expected a bare hostname "
|
| 327 |
+
"without a URL scheme."
|
| 328 |
+
)
|
| 329 |
+
return host
|
| 330 |
+
|
| 331 |
+
|
| 332 |
+
def _raise_install_error(exc: ImportError) -> None:
|
| 333 |
+
raise RuntimeError(
|
| 334 |
+
"Novita sandbox support requires optional dependencies. "
|
| 335 |
+
"Install them with `pip install openenv[novita]`."
|
| 336 |
+
) from exc
|
| 337 |
+
|
| 338 |
+
|
| 339 |
+
class _DefaultNovitaAdapter:
|
| 340 |
+
"""Thin adapter over the ``novita-sandbox`` SDK.
|
| 341 |
+
|
| 342 |
+
The provider talks to this private adapter instead of spreading SDK details
|
| 343 |
+
through its own logic; tests inject a duck-typed fake in its place. Keeping
|
| 344 |
+
the SDK surface here is what makes the provider testable without a Novita
|
| 345 |
+
account and localizes API churn (RFC 002 implementation-hygiene invariant).
|
| 346 |
+
"""
|
| 347 |
+
|
| 348 |
+
def __init__(
|
| 349 |
+
self,
|
| 350 |
+
*,
|
| 351 |
+
api_key: Optional[str],
|
| 352 |
+
domain: Optional[str],
|
| 353 |
+
):
|
| 354 |
+
try:
|
| 355 |
+
from novita_sandbox import Novita
|
| 356 |
+
except ImportError as exc: # pragma: no cover - exercised via provider tests
|
| 357 |
+
_raise_install_error(exc)
|
| 358 |
+
|
| 359 |
+
# `api_key`/`domain` fall back to NOVITA_API_KEY / NOVITA_DOMAIN inside
|
| 360 |
+
# the SDK's ConnectionConfig, so `None` is the documented "use the env
|
| 361 |
+
# var" spelling rather than a missing value.
|
| 362 |
+
self._novita = Novita(api_key=api_key, domain=domain)
|
| 363 |
+
|
| 364 |
+
def create_sandbox(
|
| 365 |
+
self,
|
| 366 |
+
*,
|
| 367 |
+
image: Optional[str],
|
| 368 |
+
template: Optional[str],
|
| 369 |
+
env_vars: Optional[Dict[str, str]],
|
| 370 |
+
timeout: int,
|
| 371 |
+
metadata: Optional[Dict[str, str]],
|
| 372 |
+
secure: Optional[bool],
|
| 373 |
+
) -> Any:
|
| 374 |
+
# Exactly one source: a registry image (resolved and cached by the SDK
|
| 375 |
+
# itself) or an already-built template id from `image_from_dockerfile`.
|
| 376 |
+
kwargs: Dict[str, Any] = {
|
| 377 |
+
"timeout": timeout,
|
| 378 |
+
}
|
| 379 |
+
|
| 380 |
+
if template is not None:
|
| 381 |
+
kwargs["template"] = template
|
| 382 |
+
else:
|
| 383 |
+
from novita_sandbox import wait_for_timeout
|
| 384 |
+
|
| 385 |
+
# `image` is an OCI reference; the SDK resolves it to a template
|
| 386 |
+
# itself (Template.from_image -> Template.build) and caches the
|
| 387 |
+
# result by image fingerprint, so repeated starts do not rebuild.
|
| 388 |
+
#
|
| 389 |
+
# The `build` block exists solely to pin the template's start
|
| 390 |
+
# command to a keepalive. `from_image` defaults to
|
| 391 |
+
# `inherit_config=True`, which restores the image's ENV/WORKDIR
|
| 392 |
+
# *and* its ENTRYPOINT/CMD as the template start command -- for an
|
| 393 |
+
# OpenEnv image that CMD is the uvicorn server, so leaving it in
|
| 394 |
+
# place would bind port 8000 before this provider launches its own
|
| 395 |
+
# copy. Setting `start_cmd` here wins over the inherited value (the
|
| 396 |
+
# caller always overrides the image). The SDK attempts to restore
|
| 397 |
+
# the image's ENV, but callers must use explicit paths when a
|
| 398 |
+
# registry config cannot be read during OCI template conversion.
|
| 399 |
+
ready_cmd = wait_for_timeout(_KEEPALIVE_READY_MS)
|
| 400 |
+
kwargs["image"] = image
|
| 401 |
+
kwargs["build"] = {
|
| 402 |
+
"cmd": _KEEPALIVE_CMD,
|
| 403 |
+
# Sandbox.create fingerprints this mapping with json.dumps
|
| 404 |
+
# before converting it into a TemplateBuilder. Keep the
|
| 405 |
+
# serialized command here; Template.set_start_cmd accepts the
|
| 406 |
+
# string as well as the ReadyCmd helper.
|
| 407 |
+
"ready_cmd": ready_cmd.get_cmd(),
|
| 408 |
+
}
|
| 409 |
+
|
| 410 |
+
if env_vars:
|
| 411 |
+
kwargs["envs"] = dict(env_vars)
|
| 412 |
+
if metadata:
|
| 413 |
+
kwargs["metadata"] = dict(metadata)
|
| 414 |
+
if secure is not None:
|
| 415 |
+
kwargs["secure"] = secure
|
| 416 |
+
|
| 417 |
+
return self._novita.sandbox.create(**kwargs)
|
| 418 |
+
|
| 419 |
+
def exec(
|
| 420 |
+
self,
|
| 421 |
+
sandbox: Any,
|
| 422 |
+
command: str,
|
| 423 |
+
*,
|
| 424 |
+
timeout: float = 10,
|
| 425 |
+
background: bool = False,
|
| 426 |
+
user: Optional[str] = None,
|
| 427 |
+
) -> str:
|
| 428 |
+
"""Run *command* through a login shell inside *sandbox*, returning stdout.
|
| 429 |
+
|
| 430 |
+
The SDK runs commands as ``/bin/bash -l -c <command>`` and raises
|
| 431 |
+
``CommandExitException`` on a non-zero exit. The discovery and liveness
|
| 432 |
+
probes below are *expected* to exit non-zero ("no openenv.yaml", "process
|
| 433 |
+
is dead"), so the exit code is swallowed and stdout returned; callers
|
| 434 |
+
distinguish states by the marker the command echoes, not by the status.
|
| 435 |
+
"""
|
| 436 |
+
from novita_sandbox import CommandExitException
|
| 437 |
+
|
| 438 |
+
run_kwargs: Dict[str, Any] = {
|
| 439 |
+
"timeout": timeout,
|
| 440 |
+
"background": background,
|
| 441 |
+
}
|
| 442 |
+
if user is not None:
|
| 443 |
+
run_kwargs["user"] = user
|
| 444 |
+
|
| 445 |
+
try:
|
| 446 |
+
result = sandbox.commands.run(command, **run_kwargs)
|
| 447 |
+
except CommandExitException as exc:
|
| 448 |
+
return exc.stdout or ""
|
| 449 |
+
except ImportError: # pragma: no cover - SDK always exports this
|
| 450 |
+
result = sandbox.commands.run(command, **run_kwargs)
|
| 451 |
+
|
| 452 |
+
return getattr(result, "stdout", "") or ""
|
| 453 |
+
|
| 454 |
+
def host(self, sandbox: Any, port: int) -> str:
|
| 455 |
+
"""Return the public host (no scheme) exposing *port* on *sandbox*."""
|
| 456 |
+
return str(sandbox.get_host(port))
|
| 457 |
+
|
| 458 |
+
def kill(self, sandbox: Any) -> None:
|
| 459 |
+
sandbox.kill()
|
| 460 |
+
|
| 461 |
+
def build_template(
|
| 462 |
+
self,
|
| 463 |
+
*,
|
| 464 |
+
dockerfile_content: str,
|
| 465 |
+
context_dir: str,
|
| 466 |
+
name: str,
|
| 467 |
+
cpu_count: int,
|
| 468 |
+
memory_mb: int,
|
| 469 |
+
on_build_logs: Optional[Callable[[Any], None]],
|
| 470 |
+
) -> str:
|
| 471 |
+
"""Build a Novita template from Dockerfile *content* and return its id.
|
| 472 |
+
|
| 473 |
+
Uses ``Template.from_dockerfile`` to translate the Dockerfile into
|
| 474 |
+
builder instructions, pins the template's start command to a keepalive
|
| 475 |
+
(so the image's own ``CMD`` does not bind port 8000 before the provider
|
| 476 |
+
launches the server), and builds it. ``file_context_path`` is what makes
|
| 477 |
+
relative ``COPY`` sources resolve against the build context.
|
| 478 |
+
|
| 479 |
+
Returns:
|
| 480 |
+
`str`: The built template id, to pass to ``Sandbox.create``.
|
| 481 |
+
"""
|
| 482 |
+
from novita_sandbox import Template, wait_for_timeout
|
| 483 |
+
|
| 484 |
+
template = Template(file_context_path=context_dir)
|
| 485 |
+
builder = template.from_dockerfile(dockerfile_content)
|
| 486 |
+
|
| 487 |
+
# The parser rewrites USER to "user" when the Dockerfile declares none
|
| 488 |
+
# (dockerfile_parser.py: `if not user_changed: set_user("user")`). For an
|
| 489 |
+
# OpenEnv image that is wrong: `WORKDIR /app`, created by root, stays
|
| 490 |
+
# root-owned, so the server and every command the environment later runs
|
| 491 |
+
# in-process (`TB2_MODE=local`, and any env doing the same) fail with
|
| 492 |
+
# "Permission denied" on the install tree. Preserve an explicit Dockerfile
|
| 493 |
+
# USER, though, so the provider does not override the image author's
|
| 494 |
+
# requested runtime identity.
|
| 495 |
+
if not _dockerfile_declares_user(dockerfile_content):
|
| 496 |
+
builder.set_user(_SANDBOX_USER)
|
| 497 |
+
|
| 498 |
+
# `set_start_cmd` is the only way to pin the start command and it
|
| 499 |
+
# requires a readiness check, so the keepalive gets a short timer. The
|
| 500 |
+
# provider's own `wait_for_ready` does the real health polling.
|
| 501 |
+
builder.set_start_cmd(_KEEPALIVE_CMD, wait_for_timeout(_KEEPALIVE_READY_MS))
|
| 502 |
+
|
| 503 |
+
build_kwargs: Dict[str, Any] = {
|
| 504 |
+
"cpu_count": cpu_count,
|
| 505 |
+
"memory_mb": memory_mb,
|
| 506 |
+
}
|
| 507 |
+
if on_build_logs is not None:
|
| 508 |
+
build_kwargs["on_build_logs"] = on_build_logs
|
| 509 |
+
|
| 510 |
+
# Build through the configured namespace so the SDK merges the
|
| 511 |
+
# api_key/domain passed to Novita(...) into Template.build's connection
|
| 512 |
+
# config. Calling Template.build directly would fall back to an empty
|
| 513 |
+
# ConnectionConfig when those values are not present in the environment.
|
| 514 |
+
info = self._novita.template.build(builder, name, **build_kwargs)
|
| 515 |
+
return str(info.template_id)
|
| 516 |
+
|
| 517 |
+
|
| 518 |
+
class NovitaSandboxProvider(ContainerProvider):
|
| 519 |
+
"""
|
| 520 |
+
Container provider that runs environments in Novita AI sandboxes.
|
| 521 |
+
|
| 522 |
+
``start_container`` accepts either form of source:
|
| 523 |
+
|
| 524 |
+
- A Docker/OCI registry reference (e.g. ``"ghcr.io/org/env:tag"``), which is
|
| 525 |
+
the form the ``ContainerProvider`` contract specifies. The SDK resolves it
|
| 526 |
+
into a Novita template internally and caches it by image fingerprint.
|
| 527 |
+
- A ``"template:<id>"`` reference returned by
|
| 528 |
+
:meth:`image_from_dockerfile`, which builds a template from a local
|
| 529 |
+
Dockerfile.
|
| 530 |
+
|
| 531 |
+
:meth:`image_from_dockerfile` rewrites the Dockerfile for Novita's template
|
| 532 |
+
parser, which does not accept multi-stage build definitions. The rewrite is
|
| 533 |
+
mechanical and preserves the build's meaning: BuildKit ``--mount`` flags are
|
| 534 |
+
stripped, ``ARG``/``--platform`` in ``FROM`` lines are resolved, and a
|
| 535 |
+
two-stage build whose stages share one base image is replayed as a single
|
| 536 |
+
stage (see `_prepare_dockerfile`). A Dockerfile that does not fit those
|
| 537 |
+
rules raises `ValueError` with the registry route as the alternative.
|
| 538 |
+
|
| 539 |
+
The environment runs untrusted code, so the provider is secure by default: it
|
| 540 |
+
enforces https/wss transport (S1) and never surfaces raw sandbox output
|
| 541 |
+
unless ``surface_server_logs=True`` (S4).
|
| 542 |
+
|
| 543 |
+
Only one sandbox is active per provider: calling ``start_container`` again
|
| 544 |
+
before ``stop_container()``/``close()`` raises ``RuntimeError`` rather than
|
| 545 |
+
orphaning the running sandbox.
|
| 546 |
+
|
| 547 |
+
Examples:
|
| 548 |
+
|
| 549 |
+
```python
|
| 550 |
+
# From a pre-built registry image
|
| 551 |
+
with NovitaSandboxProvider(image="ghcr.io/org/echo-env:latest") as provider:
|
| 552 |
+
base_url = provider.start_container()
|
| 553 |
+
provider.wait_for_ready(base_url)
|
| 554 |
+
|
| 555 |
+
# From a local Dockerfile (builds a template on first use)
|
| 556 |
+
image = NovitaSandboxProvider.image_from_dockerfile(
|
| 557 |
+
"envs/echo_env/server/Dockerfile"
|
| 558 |
+
)
|
| 559 |
+
with NovitaSandboxProvider(image=image) as provider:
|
| 560 |
+
base_url = provider.start_container()
|
| 561 |
+
# sandbox killed on exit
|
| 562 |
+
```
|
| 563 |
+
"""
|
| 564 |
+
|
| 565 |
+
# Prepared Dockerfile content, keyed by "template:<abs_path>". Populated by
|
| 566 |
+
# `image_from_dockerfile` so `start_container` can build the template later,
|
| 567 |
+
# mirroring DaytonaProvider / ModalProvider.
|
| 568 |
+
_dockerfile_registry: Dict[str, Dict[str, Any]] = {}
|
| 569 |
+
|
| 570 |
+
@classmethod
|
| 571 |
+
def image_from_dockerfile(
|
| 572 |
+
cls,
|
| 573 |
+
dockerfile_path: str,
|
| 574 |
+
context_dir: Optional[str] = None,
|
| 575 |
+
) -> str:
|
| 576 |
+
"""Validate a Dockerfile and return a ``template:`` reference for
|
| 577 |
+
:meth:`start_container`.
|
| 578 |
+
|
| 579 |
+
Eagerly validates the Dockerfile (existence, COPY sources) and stores
|
| 580 |
+
the rewritten content in an internal registry. The Novita template is
|
| 581 |
+
built later, inside ``start_container``, by the adapter's
|
| 582 |
+
``build_template`` — building needs credentials and network, which this
|
| 583 |
+
class-level helper deliberately does not.
|
| 584 |
+
|
| 585 |
+
The rewrite compensates for Novita's Dockerfile parser, which rejects
|
| 586 |
+
multi-stage build definitions (the layout every in-repo OpenEnv
|
| 587 |
+
environment uses). BuildKit ``--mount`` flags are stripped, ``ARG`` and
|
| 588 |
+
``--platform`` in ``FROM`` lines are resolved, and a two-stage build
|
| 589 |
+
whose stages share one base image is replayed as a single stage.
|
| 590 |
+
|
| 591 |
+
Args:
|
| 592 |
+
dockerfile_path (`str`):
|
| 593 |
+
Path to the Dockerfile on disk.
|
| 594 |
+
context_dir (`str`, *optional*):
|
| 595 |
+
Build context directory, used to resolve relative ``COPY``
|
| 596 |
+
sources. Defaults to the Dockerfile's grandparent directory,
|
| 597 |
+
matching the `openenv init` convention where Dockerfiles live in
|
| 598 |
+
`<env>/server/Dockerfile` and the build context is `<env>/`.
|
| 599 |
+
|
| 600 |
+
Returns:
|
| 601 |
+
`str`: A `"template:<abs_path>"` reference to pass to
|
| 602 |
+
`start_container` or the constructor's `image`.
|
| 603 |
+
|
| 604 |
+
Raises:
|
| 605 |
+
FileNotFoundError: If *dockerfile_path* does not exist.
|
| 606 |
+
ValueError: If *context_dir* does not exist, if COPY sources cannot
|
| 607 |
+
be found under the resolved context directory, or if the
|
| 608 |
+
Dockerfile cannot be expressed as a Novita template (see
|
| 609 |
+
`_prepare_dockerfile`).
|
| 610 |
+
|
| 611 |
+
Examples:
|
| 612 |
+
|
| 613 |
+
```python
|
| 614 |
+
image = NovitaSandboxProvider.image_from_dockerfile(
|
| 615 |
+
"envs/echo_env/server/Dockerfile"
|
| 616 |
+
)
|
| 617 |
+
provider = NovitaSandboxProvider(image=image)
|
| 618 |
+
base_url = provider.start_container()
|
| 619 |
+
```
|
| 620 |
+
"""
|
| 621 |
+
import pathlib
|
| 622 |
+
|
| 623 |
+
src = pathlib.Path(dockerfile_path).resolve()
|
| 624 |
+
if not src.is_file():
|
| 625 |
+
raise FileNotFoundError(f"Dockerfile not found: {dockerfile_path}")
|
| 626 |
+
|
| 627 |
+
if context_dir is not None:
|
| 628 |
+
ctx = pathlib.Path(context_dir)
|
| 629 |
+
if not ctx.is_dir():
|
| 630 |
+
raise ValueError(f"context_dir does not exist: {context_dir}")
|
| 631 |
+
else:
|
| 632 |
+
# Default: grandparent of the Dockerfile, matching the openenv init
|
| 633 |
+
# layout (<env>/server/Dockerfile -> <env>/).
|
| 634 |
+
ctx = src.parent.parent
|
| 635 |
+
|
| 636 |
+
prepared = _prepare_dockerfile(src.read_text())
|
| 637 |
+
|
| 638 |
+
# Validate that COPY sources exist under the context directory. This
|
| 639 |
+
# catches mismatches early (e.g. a Dockerfile expecting the repo root as
|
| 640 |
+
# context when we defaulted to the env directory).
|
| 641 |
+
for line in prepared.splitlines():
|
| 642 |
+
match = re.match(r"^\s*COPY\s+(?!--from=)(\S+)\s+", line, re.IGNORECASE)
|
| 643 |
+
if not match:
|
| 644 |
+
continue
|
| 645 |
+
copy_src = match.group(1)
|
| 646 |
+
if copy_src.startswith("/"):
|
| 647 |
+
continue
|
| 648 |
+
resolved = ctx / copy_src
|
| 649 |
+
if not resolved.exists() and not any(ctx.glob(copy_src)):
|
| 650 |
+
raise ValueError(
|
| 651 |
+
f"Dockerfile COPY source '{copy_src}' not found under "
|
| 652 |
+
f"context_dir '{ctx}'. This Dockerfile may expect a "
|
| 653 |
+
"different build context (e.g. the repo root). Pass "
|
| 654 |
+
"context_dir explicitly."
|
| 655 |
+
)
|
| 656 |
+
|
| 657 |
+
cls._dockerfile_registry[str(src)] = {
|
| 658 |
+
"content": prepared,
|
| 659 |
+
"context_dir": str(ctx),
|
| 660 |
+
}
|
| 661 |
+
|
| 662 |
+
return f"template:{src}"
|
| 663 |
+
|
| 664 |
+
def _resolve_template_ref(self, reference: str) -> str:
|
| 665 |
+
"""Build the template for a ``template:<path>`` reference; return its id."""
|
| 666 |
+
path = reference[len("template:") :]
|
| 667 |
+
meta = self._dockerfile_registry.get(path)
|
| 668 |
+
if meta is None:
|
| 669 |
+
raise ValueError(
|
| 670 |
+
f"No registered Dockerfile metadata for {path}. Call "
|
| 671 |
+
"NovitaSandboxProvider.image_from_dockerfile() first."
|
| 672 |
+
)
|
| 673 |
+
|
| 674 |
+
return self._adapter.build_template(
|
| 675 |
+
dockerfile_content=meta["content"],
|
| 676 |
+
context_dir=meta["context_dir"],
|
| 677 |
+
name=self._template_name or self._default_template_name(path),
|
| 678 |
+
cpu_count=self._cpu_count,
|
| 679 |
+
memory_mb=self._memory_mb,
|
| 680 |
+
on_build_logs=self._on_build_logs,
|
| 681 |
+
)
|
| 682 |
+
|
| 683 |
+
@staticmethod
|
| 684 |
+
def _default_template_name(dockerfile_path: str) -> str:
|
| 685 |
+
"""Derive a stable template name from the Dockerfile's env directory.
|
| 686 |
+
|
| 687 |
+
The env directory keeps the template recognizable in the Novita
|
| 688 |
+
dashboard; the content hash keeps two different Dockerfiles from
|
| 689 |
+
colliding on one name, while a rebuild of identical content still hits
|
| 690 |
+
the build cache.
|
| 691 |
+
"""
|
| 692 |
+
import hashlib
|
| 693 |
+
import pathlib
|
| 694 |
+
|
| 695 |
+
src = pathlib.Path(dockerfile_path)
|
| 696 |
+
env_dir = src.parent.parent.name or src.parent.name
|
| 697 |
+
slug = re.sub(r"[^a-z0-9-]+", "-", env_dir.lower().replace("_", "-"))
|
| 698 |
+
digest = hashlib.sha256(src.read_bytes()).hexdigest()[:8]
|
| 699 |
+
return f"openenv-{slug.strip('-')}-{digest}"
|
| 700 |
+
|
| 701 |
+
def __init__(
|
| 702 |
+
self,
|
| 703 |
+
*,
|
| 704 |
+
image: Optional[str] = None,
|
| 705 |
+
env_vars: Optional[Dict[str, str]] = None,
|
| 706 |
+
api_key: Optional[str] = None,
|
| 707 |
+
domain: Optional[str] = None,
|
| 708 |
+
timeout: int = _DEFAULT_SANDBOX_TIMEOUT_S,
|
| 709 |
+
metadata: Optional[Dict[str, str]] = None,
|
| 710 |
+
secure: Optional[bool] = None,
|
| 711 |
+
cmd: Optional[str] = None,
|
| 712 |
+
working_directory: Optional[str] = None,
|
| 713 |
+
surface_server_logs: bool = False,
|
| 714 |
+
cpu_count: int = _DEFAULT_TEMPLATE_CPU,
|
| 715 |
+
memory_mb: int = _DEFAULT_TEMPLATE_MEMORY_MB,
|
| 716 |
+
template_name: Optional[str] = None,
|
| 717 |
+
on_build_logs: Optional[Callable[[Any], None]] = None,
|
| 718 |
+
_adapter: Any = None,
|
| 719 |
+
):
|
| 720 |
+
"""
|
| 721 |
+
Args:
|
| 722 |
+
image (`str`, *optional*):
|
| 723 |
+
Registry image reference (e.g. `"ghcr.io/org/env:latest"`) or a
|
| 724 |
+
`"template:<id>"` reference from
|
| 725 |
+
[`~image_from_dockerfile`], to use when `start_container()` is
|
| 726 |
+
called without an image.
|
| 727 |
+
env_vars (`dict`, *optional*):
|
| 728 |
+
Environment variables to use when `start_container()` is called
|
| 729 |
+
without explicit `env_vars`.
|
| 730 |
+
api_key (`str`, *optional*):
|
| 731 |
+
Novita API key. Falls back to the `NOVITA_API_KEY` environment
|
| 732 |
+
variable.
|
| 733 |
+
domain (`str`, *optional*):
|
| 734 |
+
Novita region domain. Falls back to the `NOVITA_DOMAIN`
|
| 735 |
+
environment variable, then to the SDK default (`us-phx-1`).
|
| 736 |
+
Features marked modern-domain-only (snapshots, secrets) are
|
| 737 |
+
unavailable on the legacy `us-virginia-1` domain.
|
| 738 |
+
timeout (`int`, *optional*, defaults to `3600`):
|
| 739 |
+
Hard sandbox lifetime in seconds. Counts down from creation
|
| 740 |
+
regardless of activity and always fires, so it must exceed the
|
| 741 |
+
longest expected episode.
|
| 742 |
+
metadata (`dict`, *optional*):
|
| 743 |
+
Arbitrary key/value labels on the sandbox. The `idle_timeout`
|
| 744 |
+
key is meaningful to Novita and takes seconds as a **string**
|
| 745 |
+
(e.g. `{"idle_timeout": "900"}`); omitting it disables the idle
|
| 746 |
+
timer, which is the conservative default for RL rollouts — an
|
| 747 |
+
automatic transition can drop a live WebSocket mid-episode
|
| 748 |
+
(RFC 002 invariant 6).
|
| 749 |
+
secure (`bool`, *optional*):
|
| 750 |
+
Whether envd requires its access token. `None` uses the SDK
|
| 751 |
+
default (secured on modern domains). This governs the sandbox
|
| 752 |
+
control plane, not the exposed server port.
|
| 753 |
+
cmd (`str`, *optional*):
|
| 754 |
+
Shell command to start the server inside the sandbox. When
|
| 755 |
+
omitted, the command is auto-discovered from `openenv.yaml`.
|
| 756 |
+
working_directory (`str`, *optional*):
|
| 757 |
+
Directory to `cd` into before running `cmd`.
|
| 758 |
+
surface_server_logs (`bool`, *optional*, defaults to `False`):
|
| 759 |
+
When `False` (default), captured sandbox output is withheld from
|
| 760 |
+
raised errors so secrets the workload printed cannot leak into
|
| 761 |
+
orchestrator/CI logs. When `True`, a best-effort redacted,
|
| 762 |
+
length-bounded excerpt is included in startup-crash errors (S4).
|
| 763 |
+
cpu_count (`int`, *optional*, defaults to `2`):
|
| 764 |
+
vCPU cores for sandboxes built from a Dockerfile via
|
| 765 |
+
[`~image_from_dockerfile`]. Ignored for registry images, whose
|
| 766 |
+
resources come from the image's own template.
|
| 767 |
+
memory_mb (`int`, *optional*, defaults to `1024`):
|
| 768 |
+
Memory in MiB for sandboxes built from a Dockerfile. Ignored for
|
| 769 |
+
registry images.
|
| 770 |
+
template_name (`str`, *optional*):
|
| 771 |
+
Template name to register when building from a Dockerfile.
|
| 772 |
+
Defaults to a name derived from the Dockerfile's directory plus a
|
| 773 |
+
content hash, so re-running with the same Dockerfile reuses the
|
| 774 |
+
build cache.
|
| 775 |
+
on_build_logs (`Callable`, *optional*):
|
| 776 |
+
Callback receiving build-log entries while a Dockerfile template
|
| 777 |
+
is built. Off by default: build logs are not redacted, so they
|
| 778 |
+
are withheld rather than printed (RFC 002 S4). Pass a callable to
|
| 779 |
+
surface them.
|
| 780 |
+
_adapter (`Any`, *optional*):
|
| 781 |
+
Injection seam for tests; a duck-typed fake replaces the SDK.
|
| 782 |
+
"""
|
| 783 |
+
self._image = image
|
| 784 |
+
self._env_vars = env_vars
|
| 785 |
+
self._timeout = timeout
|
| 786 |
+
self._metadata = metadata
|
| 787 |
+
self._secure = secure
|
| 788 |
+
self._cmd = cmd
|
| 789 |
+
self._working_directory = working_directory
|
| 790 |
+
self.surface_server_logs = surface_server_logs
|
| 791 |
+
self._cpu_count = cpu_count
|
| 792 |
+
self._memory_mb = memory_mb
|
| 793 |
+
self._template_name = template_name
|
| 794 |
+
self._on_build_logs = on_build_logs
|
| 795 |
+
|
| 796 |
+
self._sandbox: Any = None
|
| 797 |
+
self._base_url: Optional[str] = None
|
| 798 |
+
self._server_user: Optional[str] = None
|
| 799 |
+
# Injected env-var values, used to scrub captured server output before
|
| 800 |
+
# it is ever surfaced in an error.
|
| 801 |
+
self._redact_values: set[str] = set()
|
| 802 |
+
|
| 803 |
+
if _adapter is None:
|
| 804 |
+
# Build the client eagerly so credential/region problems surface at
|
| 805 |
+
# construction rather than mid-episode, matching ModalProvider.
|
| 806 |
+
self._adapter: Any = _DefaultNovitaAdapter(api_key=api_key, domain=domain)
|
| 807 |
+
else:
|
| 808 |
+
self._adapter = _adapter
|
| 809 |
+
|
| 810 |
+
def _discover_server_cmd(self, port: int = _DEFAULT_NOVITA_PORT) -> str:
|
| 811 |
+
"""Discover the server command from ``openenv.yaml`` inside the sandbox.
|
| 812 |
+
|
| 813 |
+
Finds the file, reads the ``app`` field, and constructs a command of the
|
| 814 |
+
form ``cd <env_root> && python -m uvicorn <app> --host 0.0.0.0 --port <port>``.
|
| 815 |
+
|
| 816 |
+
Raises:
|
| 817 |
+
ValueError: If ``openenv.yaml`` is not found or lacks an ``app`` field.
|
| 818 |
+
"""
|
| 819 |
+
yaml_path = self._find_openenv_yaml()
|
| 820 |
+
if yaml_path is None:
|
| 821 |
+
raise ValueError(
|
| 822 |
+
"Could not find openenv.yaml inside the sandbox. Pass an "
|
| 823 |
+
"explicit cmd= to NovitaSandboxProvider or start_container()."
|
| 824 |
+
)
|
| 825 |
+
|
| 826 |
+
content = self._adapter.exec(self._sandbox, f"cat {shlex.quote(yaml_path)}")
|
| 827 |
+
app = self._parse_app_field(content)
|
| 828 |
+
if app is None:
|
| 829 |
+
raise ValueError(
|
| 830 |
+
f"openenv.yaml at {yaml_path} does not contain an 'app' field. "
|
| 831 |
+
"Pass an explicit cmd= to NovitaSandboxProvider or start_container()."
|
| 832 |
+
)
|
| 833 |
+
|
| 834 |
+
# The directory containing openenv.yaml is the env root
|
| 835 |
+
env_root = yaml_path.rsplit("/", 1)[0]
|
| 836 |
+
return (
|
| 837 |
+
f"cd {shlex.quote(env_root)} && "
|
| 838 |
+
f"python -m uvicorn {shlex.quote(app)} --host 0.0.0.0 --port {port}"
|
| 839 |
+
)
|
| 840 |
+
|
| 841 |
+
def _find_openenv_yaml(self) -> Optional[str]:
|
| 842 |
+
"""Locate ``openenv.yaml`` inside the sandbox.
|
| 843 |
+
|
| 844 |
+
Tries the modern layout path ``/app/env/openenv.yaml`` first, then falls
|
| 845 |
+
back to a ``find`` command for the old layout.
|
| 846 |
+
"""
|
| 847 |
+
# Fast path: modern Dockerfile layout
|
| 848 |
+
out = self._adapter.exec(
|
| 849 |
+
self._sandbox, "test -f /app/env/openenv.yaml && echo found"
|
| 850 |
+
)
|
| 851 |
+
if "found" in (out or ""):
|
| 852 |
+
return "/app/env/openenv.yaml"
|
| 853 |
+
|
| 854 |
+
# Fallback: search for it (redirect stderr so error messages like
|
| 855 |
+
# "No such file or directory" don't get mistaken for paths).
|
| 856 |
+
path = self._adapter.exec(
|
| 857 |
+
self._sandbox,
|
| 858 |
+
"find /app -maxdepth 4 -name openenv.yaml -print -quit 2>/dev/null",
|
| 859 |
+
).strip()
|
| 860 |
+
if path and path.startswith("/"):
|
| 861 |
+
return path
|
| 862 |
+
|
| 863 |
+
return None
|
| 864 |
+
|
| 865 |
+
@staticmethod
|
| 866 |
+
def _parse_app_field(yaml_content: str) -> Optional[str]:
|
| 867 |
+
"""Extract the ``app`` value from raw openenv.yaml content.
|
| 868 |
+
|
| 869 |
+
Uses PyYAML to handle comments, quotes, and nested keys correctly.
|
| 870 |
+
"""
|
| 871 |
+
return parse_openenv_app_field(yaml_content)
|
| 872 |
+
|
| 873 |
+
def start_container(
|
| 874 |
+
self,
|
| 875 |
+
image: Optional[str] = None,
|
| 876 |
+
port: Optional[int] = None,
|
| 877 |
+
env_vars: Optional[Dict[str, str]] = None,
|
| 878 |
+
**kwargs: Any,
|
| 879 |
+
) -> str:
|
| 880 |
+
"""
|
| 881 |
+
Create a Novita sandbox and start an OpenEnv server inside it.
|
| 882 |
+
|
| 883 |
+
Args:
|
| 884 |
+
image (`str`, *optional*):
|
| 885 |
+
Either a registry image reference (e.g.
|
| 886 |
+
`"ghcr.io/org/env:latest"`) or a `"template:<path>"` reference
|
| 887 |
+
returned by [`~image_from_dockerfile`]. May be omitted when
|
| 888 |
+
supplied to the constructor.
|
| 889 |
+
port (`int`, *optional*):
|
| 890 |
+
Must be `None` or `8000`. Novita exposes port 8000 on the
|
| 891 |
+
sandbox host; other ports raise `ValueError`.
|
| 892 |
+
env_vars (`dict`, *optional*):
|
| 893 |
+
Environment variables forwarded to the sandbox, overriding the
|
| 894 |
+
constructor's.
|
| 895 |
+
**kwargs:
|
| 896 |
+
`cmd` (`str`) to override the server command. For a registry
|
| 897 |
+
OCI image, pass the command explicitly together with any
|
| 898 |
+
runtime environment the image needs: Novita may not be able to
|
| 899 |
+
restore the image's ``ENV`` while converting it to a template.
|
| 900 |
+
For example, the tbench2 image requires:
|
| 901 |
+
|
| 902 |
+
``cmd="cd /app/env && PYTHONPATH=/app/env "
|
| 903 |
+
"/app/.venv/bin/python -m uvicorn server.app:app "
|
| 904 |
+
"--host 0.0.0.0 --port 8000"``
|
| 905 |
+
|
| 906 |
+
and any additional values can be supplied through
|
| 907 |
+
``env_vars``. The provider runs this command as ``root`` for a
|
| 908 |
+
registry image; a ``template:<id>`` keeps the template's
|
| 909 |
+
configured user. Prefer
|
| 910 |
+
``NovitaSandboxProvider.image_from_dockerfile(...)`` when a
|
| 911 |
+
local Dockerfile is available; its build preserves the
|
| 912 |
+
environment and startup configuration in the generated
|
| 913 |
+
template. Unknown options raise `ValueError` so typos cannot
|
| 914 |
+
silently change sandbox behavior.
|
| 915 |
+
|
| 916 |
+
Returns:
|
| 917 |
+
`str`: HTTPS sandbox URL for the exposed port (base_url).
|
| 918 |
+
|
| 919 |
+
Raises:
|
| 920 |
+
RuntimeError: If a sandbox is already active on this provider.
|
| 921 |
+
ValueError: If no image is available, the port is unsupported, or an
|
| 922 |
+
unknown option is passed.
|
| 923 |
+
"""
|
| 924 |
+
if self._sandbox is not None:
|
| 925 |
+
raise RuntimeError(
|
| 926 |
+
"NovitaSandboxProvider already has an active sandbox. Call "
|
| 927 |
+
"stop_container() (or close()) before starting another — a "
|
| 928 |
+
"second start would orphan the running sandbox."
|
| 929 |
+
)
|
| 930 |
+
|
| 931 |
+
if port is not None and port != _DEFAULT_NOVITA_PORT:
|
| 932 |
+
raise ValueError(
|
| 933 |
+
f"NovitaSandboxProvider only supports port {_DEFAULT_NOVITA_PORT} "
|
| 934 |
+
f"(got {port}). The sandbox host routes to port "
|
| 935 |
+
f"{_DEFAULT_NOVITA_PORT} inside the sandbox."
|
| 936 |
+
)
|
| 937 |
+
|
| 938 |
+
effective_image = image if image is not None else self._image
|
| 939 |
+
if effective_image is None:
|
| 940 |
+
raise ValueError(
|
| 941 |
+
"NovitaSandboxProvider requires an image. Pass it to the "
|
| 942 |
+
"constructor or start_container()."
|
| 943 |
+
)
|
| 944 |
+
effective_env_vars = self._env_vars if env_vars is None else env_vars
|
| 945 |
+
|
| 946 |
+
cmd = kwargs.pop("cmd", None) or self._cmd
|
| 947 |
+
# `AutoEnv.from_env()` forwards this unconditionally (default 30.0) via
|
| 948 |
+
# `from_docker_image` -> `_bootstrap_container`, so rejecting it would
|
| 949 |
+
# make the entire AutoEnv path unusable for this provider. It is
|
| 950 |
+
# meaningless here: `_bootstrap_container` calls `wait_for_ready(base_url)`
|
| 951 |
+
# without a timeout, so no provider ever sees the value. Accept and drop
|
| 952 |
+
# it rather than fail; unknown options still raise, so a typo cannot
|
| 953 |
+
# silently change sandbox behavior.
|
| 954 |
+
kwargs.pop("wait_timeout", None)
|
| 955 |
+
if kwargs:
|
| 956 |
+
unknown = ", ".join(sorted(kwargs))
|
| 957 |
+
raise ValueError(
|
| 958 |
+
f"Unsupported NovitaSandboxProvider start options: {unknown}"
|
| 959 |
+
)
|
| 960 |
+
|
| 961 |
+
# A "template:<path>" reference builds a Novita template from a local
|
| 962 |
+
# Dockerfile; anything else is a registry image the SDK resolves itself.
|
| 963 |
+
# The build runs before `create_sandbox` so a build failure never leaves
|
| 964 |
+
# a sandbox behind, and outside the `_redact_values` window below so its
|
| 965 |
+
# recorded secrets do not outlive a failed build.
|
| 966 |
+
if effective_image.startswith("template:"):
|
| 967 |
+
template_id: Optional[str] = self._resolve_template_ref(effective_image)
|
| 968 |
+
sandbox_image: Optional[str] = None
|
| 969 |
+
else:
|
| 970 |
+
template_id = None
|
| 971 |
+
sandbox_image = effective_image
|
| 972 |
+
|
| 973 |
+
# Record injected secret values so captured server output can be scrubbed
|
| 974 |
+
# before it is ever surfaced in an error (security invariant S4).
|
| 975 |
+
self._redact_values = {
|
| 976 |
+
value for value in (effective_env_vars or {}).values() if value
|
| 977 |
+
}
|
| 978 |
+
|
| 979 |
+
# A create failure created nothing, so just drop the recorded secrets and
|
| 980 |
+
# re-raise (the single-sandbox guard above guarantees there is no
|
| 981 |
+
# pre-existing sandbox to delete).
|
| 982 |
+
try:
|
| 983 |
+
self._sandbox = self._adapter.create_sandbox(
|
| 984 |
+
image=sandbox_image,
|
| 985 |
+
template=template_id,
|
| 986 |
+
env_vars=effective_env_vars,
|
| 987 |
+
timeout=self._timeout,
|
| 988 |
+
metadata=self._metadata,
|
| 989 |
+
secure=self._secure,
|
| 990 |
+
)
|
| 991 |
+
except Exception:
|
| 992 |
+
self._redact_values = set()
|
| 993 |
+
raise
|
| 994 |
+
|
| 995 |
+
try:
|
| 996 |
+
# Discovery runs after creation, when the filesystem can be read.
|
| 997 |
+
if cmd is None:
|
| 998 |
+
cmd = self._discover_server_cmd()
|
| 999 |
+
|
| 1000 |
+
# Registry images can lose their Dockerfile USER/ownership settings
|
| 1001 |
+
# while Novita converts them to a cached OCI template. Run the
|
| 1002 |
+
# provider-managed server command as root for that path. A template
|
| 1003 |
+
# keeps its configured user, including an explicit non-root USER.
|
| 1004 |
+
self._server_user = _SANDBOX_USER if sandbox_image is not None else None
|
| 1005 |
+
self._launch_server(
|
| 1006 |
+
cmd,
|
| 1007 |
+
user=self._server_user,
|
| 1008 |
+
)
|
| 1009 |
+
host = _require_bare_host(
|
| 1010 |
+
self._adapter.host(self._sandbox, _DEFAULT_NOVITA_PORT)
|
| 1011 |
+
)
|
| 1012 |
+
self._base_url = _require_secure_url(f"https://{host}")
|
| 1013 |
+
except Exception:
|
| 1014 |
+
# A cleanup failure here must not mask the original error: swallow
|
| 1015 |
+
# any exception from stop_container() so the root cause propagates.
|
| 1016 |
+
try:
|
| 1017 |
+
self.stop_container()
|
| 1018 |
+
except Exception:
|
| 1019 |
+
pass
|
| 1020 |
+
raise
|
| 1021 |
+
|
| 1022 |
+
return self._base_url
|
| 1023 |
+
|
| 1024 |
+
def _launch_server(self, cmd: str, *, user: Optional[str] = None) -> None:
|
| 1025 |
+
"""Start the OpenEnv server inside the sandbox as a background process.
|
| 1026 |
+
|
| 1027 |
+
This is the one place the provider assumes something about the image, so
|
| 1028 |
+
it is kept to a single method.
|
| 1029 |
+
|
| 1030 |
+
The server is launched through the SDK's background command support and
|
| 1031 |
+
its PID is written to ``/tmp/openenv-server.pid``, so ``wait_for_ready``
|
| 1032 |
+
can tell "still booting" apart from "crashed" instead of waiting out
|
| 1033 |
+
the full timeout.
|
| 1034 |
+
The sandbox was created with a ``sleep infinity`` start command rather
|
| 1035 |
+
than the image's own CMD (see ``_DefaultNovitaAdapter.create_sandbox``),
|
| 1036 |
+
which is what leaves port 8000 free for this process.
|
| 1037 |
+
|
| 1038 |
+
An alternative is to let the platform own it: pass the server command as
|
| 1039 |
+
the template's start command via the SDK's ``build={"cmd": ...}`` and
|
| 1040 |
+
drop this method. That trades the PID-based crash detection
|
| 1041 |
+
below for platform-managed lifecycle and readiness. It is not the
|
| 1042 |
+
default because a start command can only be set when the image is first
|
| 1043 |
+
resolved into a template, so it would not apply to an already-cached
|
| 1044 |
+
image.
|
| 1045 |
+
|
| 1046 |
+
``cmd`` is trusted orchestrator configuration, not agent or environment
|
| 1047 |
+
input — callers must not pass agent-controlled text, and any dynamic
|
| 1048 |
+
value interpolated into it must be quoted by the caller (S5). The
|
| 1049 |
+
sandbox SDK runs commands as ``/bin/bash -l -c <command>``; the command
|
| 1050 |
+
itself owns the PID file and log redirection while ``background=True``
|
| 1051 |
+
keeps the remote process alive after this call.
|
| 1052 |
+
"""
|
| 1053 |
+
if self._working_directory:
|
| 1054 |
+
command = f"cd {shlex.quote(self._working_directory)} && {cmd}"
|
| 1055 |
+
else:
|
| 1056 |
+
command = cmd
|
| 1057 |
+
|
| 1058 |
+
launch_command = (
|
| 1059 |
+
"echo $$ > /tmp/openenv-server.pid; "
|
| 1060 |
+
f"exec sh -c {shlex.quote(command)} "
|
| 1061 |
+
"> /tmp/openenv-server.log 2>&1"
|
| 1062 |
+
)
|
| 1063 |
+
escaped = shlex.quote(launch_command)
|
| 1064 |
+
self._adapter.exec(
|
| 1065 |
+
self._sandbox,
|
| 1066 |
+
f"sh -c {escaped}",
|
| 1067 |
+
timeout=0,
|
| 1068 |
+
background=True,
|
| 1069 |
+
user=user,
|
| 1070 |
+
)
|
| 1071 |
+
|
| 1072 |
+
def stop_container(self) -> None:
|
| 1073 |
+
"""Kill the active Novita sandbox.
|
| 1074 |
+
|
| 1075 |
+
`kill()` permanently removes the sandbox: it cannot be resumed or
|
| 1076 |
+
reconnected, and any snapshot must have been taken beforehand.
|
| 1077 |
+
"""
|
| 1078 |
+
if self._sandbox is None:
|
| 1079 |
+
# Still drop any injected secret values recorded by a failed start.
|
| 1080 |
+
self._redact_values = set()
|
| 1081 |
+
return
|
| 1082 |
+
|
| 1083 |
+
sandbox = self._sandbox
|
| 1084 |
+
try:
|
| 1085 |
+
self._adapter.kill(sandbox)
|
| 1086 |
+
finally:
|
| 1087 |
+
self._sandbox = None
|
| 1088 |
+
self._base_url = None
|
| 1089 |
+
self._server_user = None
|
| 1090 |
+
self._redact_values = set()
|
| 1091 |
+
|
| 1092 |
+
def close(self) -> None:
|
| 1093 |
+
"""Stop the active sandbox and release provider-held resources.
|
| 1094 |
+
|
| 1095 |
+
Overrides the base no-op so a caller holding a bare `ContainerProvider`
|
| 1096 |
+
reference can release the sandbox polymorphically (also invoked on
|
| 1097 |
+
context-manager exit). The Novita client is stateless HTTP per call, so
|
| 1098 |
+
this is equivalent to `stop_container()`.
|
| 1099 |
+
"""
|
| 1100 |
+
self.stop_container()
|
| 1101 |
+
|
| 1102 |
+
@property
|
| 1103 |
+
def base_url(self) -> str:
|
| 1104 |
+
"""URL returned by the last `start_container`."""
|
| 1105 |
+
if self._base_url is None:
|
| 1106 |
+
raise RuntimeError(
|
| 1107 |
+
"NovitaSandboxProvider has no active base_url. Start the "
|
| 1108 |
+
"provider before reading base_url."
|
| 1109 |
+
)
|
| 1110 |
+
return self._base_url
|
| 1111 |
+
|
| 1112 |
+
def _redact(self, text: str, *, max_chars: int = 2000) -> str:
|
| 1113 |
+
"""Scrub injected secret values and bound length before surfacing output.
|
| 1114 |
+
|
| 1115 |
+
Replaces any injected env-var value with `***` and keeps only the tail.
|
| 1116 |
+
This is best-effort (exact-match only), which is why server output is
|
| 1117 |
+
withheld entirely unless `surface_server_logs=True` (RFC 002 S4).
|
| 1118 |
+
"""
|
| 1119 |
+
redacted = text or ""
|
| 1120 |
+
for value in self._redact_values:
|
| 1121 |
+
redacted = redacted.replace(value, "***")
|
| 1122 |
+
if len(redacted) > max_chars:
|
| 1123 |
+
redacted = "...(truncated)...\n" + redacted[-max_chars:]
|
| 1124 |
+
return redacted
|
| 1125 |
+
|
| 1126 |
+
def _server_died_message(self) -> str:
|
| 1127 |
+
"""Build the startup-crash error, secure by default (RFC 002 S4).
|
| 1128 |
+
|
| 1129 |
+
Untrusted code can print secrets then force a crash to exfiltrate them
|
| 1130 |
+
through the exception (which lands in orchestrator/CI logs), so sandbox
|
| 1131 |
+
output is excluded unless `surface_server_logs=True`, in which case a
|
| 1132 |
+
best-effort redacted, bounded excerpt is included.
|
| 1133 |
+
"""
|
| 1134 |
+
base = (
|
| 1135 |
+
"Novita sandbox server process died during startup. Server output is "
|
| 1136 |
+
"not surfaced to avoid leaking secrets injected into the sandbox; "
|
| 1137 |
+
"retrieve /tmp/openenv-server.log from the sandbox out of band, or "
|
| 1138 |
+
"construct the provider with surface_server_logs=True to include a "
|
| 1139 |
+
"redacted excerpt."
|
| 1140 |
+
)
|
| 1141 |
+
if not self.surface_server_logs:
|
| 1142 |
+
return base
|
| 1143 |
+
if self._sandbox is None:
|
| 1144 |
+
return base
|
| 1145 |
+
|
| 1146 |
+
log = self._redact(
|
| 1147 |
+
self._adapter.exec(
|
| 1148 |
+
self._sandbox,
|
| 1149 |
+
"cat /tmp/openenv-server.log 2>/dev/null",
|
| 1150 |
+
user=self._server_user,
|
| 1151 |
+
)
|
| 1152 |
+
)
|
| 1153 |
+
return (
|
| 1154 |
+
"Novita sandbox server process died during startup. The excerpt below "
|
| 1155 |
+
"is the sandbox server output with injected secret values redacted "
|
| 1156 |
+
"(best-effort); it may still contain secrets the workload printed "
|
| 1157 |
+
f"by other means.\nLog (redacted):\n{log}"
|
| 1158 |
+
)
|
| 1159 |
+
|
| 1160 |
+
def wait_for_ready(self, base_url: str, timeout_s: float = 120.0) -> None:
|
| 1161 |
+
"""
|
| 1162 |
+
Poll the /health endpoint until the sandbox is ready.
|
| 1163 |
+
|
| 1164 |
+
Uses a longer default timeout (120s) than local Docker providers because
|
| 1165 |
+
a Novita sandbox is created from an image that may need to be resolved
|
| 1166 |
+
and built on first use.
|
| 1167 |
+
|
| 1168 |
+
A `200` on `/health` proves HTTP reachability but **not** that the
|
| 1169 |
+
exposed sandbox host proxies the `/ws` WebSocket upgrade `EnvClient`
|
| 1170 |
+
needs; that requires a real `wss://` round-trip (RFC 002 invariant 2).
|
| 1171 |
+
|
| 1172 |
+
Args:
|
| 1173 |
+
base_url (`str`):
|
| 1174 |
+
Sandbox URL returned by `start_container()`.
|
| 1175 |
+
timeout_s (`float`, *optional*, defaults to `120.0`):
|
| 1176 |
+
Maximum seconds to wait.
|
| 1177 |
+
|
| 1178 |
+
Raises:
|
| 1179 |
+
TimeoutError: If the sandbox doesn't become ready in time.
|
| 1180 |
+
RuntimeError: If the server process died (detected via PID check).
|
| 1181 |
+
"""
|
| 1182 |
+
# Imported lazily inside the method that needs it, matching the other
|
| 1183 |
+
# providers in this package (e.g. DaytonaProvider).
|
| 1184 |
+
import requests
|
| 1185 |
+
|
| 1186 |
+
deadline = time.time() + timeout_s
|
| 1187 |
+
health_url = f"{base_url}/health"
|
| 1188 |
+
|
| 1189 |
+
while time.time() < deadline:
|
| 1190 |
+
try:
|
| 1191 |
+
response = requests.get(health_url, timeout=5.0)
|
| 1192 |
+
if response.status_code == 200:
|
| 1193 |
+
return
|
| 1194 |
+
except requests.RequestException:
|
| 1195 |
+
pass
|
| 1196 |
+
|
| 1197 |
+
# Early exit: if the server process died, raise immediately instead
|
| 1198 |
+
# of waiting for the full health-check timeout.
|
| 1199 |
+
if self._sandbox is not None:
|
| 1200 |
+
out = self._adapter.exec(
|
| 1201 |
+
self._sandbox,
|
| 1202 |
+
"if test -s /tmp/openenv-server.pid; then "
|
| 1203 |
+
"kill -0 $(cat /tmp/openenv-server.pid) 2>/dev/null "
|
| 1204 |
+
"&& echo RUNNING || echo DEAD; "
|
| 1205 |
+
"else echo STARTING; fi",
|
| 1206 |
+
user=self._server_user,
|
| 1207 |
+
)
|
| 1208 |
+
if "DEAD" in (out or ""):
|
| 1209 |
+
raise RuntimeError(self._server_died_message())
|
| 1210 |
+
|
| 1211 |
+
time.sleep(1.0)
|
| 1212 |
+
|
| 1213 |
+
raise TimeoutError(f"Novita sandbox did not become ready within {timeout_s}s.")
|
| 1214 |
+
|
| 1215 |
+
|
| 1216 |
+
__all__ = ["NovitaSandboxProvider"]
|
src/core/containers/runtime/providers.py
ADDED
|
@@ -0,0 +1,727 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# SPDX-License-Identifier: BSD-3-Clause
|
| 2 |
+
|
| 3 |
+
"""
|
| 4 |
+
Container provider abstractions for running environment servers.
|
| 5 |
+
|
| 6 |
+
This module provides a pluggable architecture for different container providers
|
| 7 |
+
(local Docker, Kubernetes, cloud providers, etc.) to be used with EnvClient.
|
| 8 |
+
"""
|
| 9 |
+
|
| 10 |
+
from __future__ import annotations
|
| 11 |
+
|
| 12 |
+
from abc import ABC, abstractmethod
|
| 13 |
+
from typing import Any, Dict, Optional, Sequence, TypeVar
|
| 14 |
+
|
| 15 |
+
_ContainerProviderT = TypeVar("_ContainerProviderT", bound="ContainerProvider")
|
| 16 |
+
|
| 17 |
+
|
| 18 |
+
class ContainerProvider(ABC):
|
| 19 |
+
"""
|
| 20 |
+
Abstract base class for container providers.
|
| 21 |
+
|
| 22 |
+
Providers implement this interface to support different container platforms:
|
| 23 |
+
- LocalDockerProvider: Runs containers on local Docker daemon
|
| 24 |
+
- KubernetesProvider: Runs containers in Kubernetes cluster
|
| 25 |
+
- FargateProvider: Runs containers on AWS Fargate
|
| 26 |
+
- CloudRunProvider: Runs containers on Google Cloud Run
|
| 27 |
+
|
| 28 |
+
The provider manages a single container lifecycle and provides the base URL
|
| 29 |
+
for connecting to it.
|
| 30 |
+
|
| 31 |
+
Examples:
|
| 32 |
+
|
| 33 |
+
```python
|
| 34 |
+
provider = LocalDockerProvider()
|
| 35 |
+
base_url = provider.start_container("echo-env:latest")
|
| 36 |
+
print(base_url) # http://localhost:8000
|
| 37 |
+
# Use the environment via base_url
|
| 38 |
+
provider.stop_container()
|
| 39 |
+
```
|
| 40 |
+
"""
|
| 41 |
+
|
| 42 |
+
@abstractmethod
|
| 43 |
+
def start_container(
|
| 44 |
+
self,
|
| 45 |
+
image: str,
|
| 46 |
+
port: Optional[int] = None,
|
| 47 |
+
env_vars: Optional[Dict[str, str]] = None,
|
| 48 |
+
**kwargs: Any,
|
| 49 |
+
) -> str:
|
| 50 |
+
"""
|
| 51 |
+
Start a container from the specified image.
|
| 52 |
+
|
| 53 |
+
Args:
|
| 54 |
+
image (`str`):
|
| 55 |
+
Provider-specific container *source* identifier. For
|
| 56 |
+
container-based providers this is a registry image name (e.g.
|
| 57 |
+
`"echo-env:latest"`); other providers may map it to a
|
| 58 |
+
provider-specific source (see the provider's documentation).
|
| 59 |
+
port (`int`, *optional*):
|
| 60 |
+
Port to expose. If `None`, the provider chooses.
|
| 61 |
+
env_vars (`dict`, *optional*):
|
| 62 |
+
Environment variables to pass to container.
|
| 63 |
+
**kwargs:
|
| 64 |
+
Provider-specific options.
|
| 65 |
+
|
| 66 |
+
Returns:
|
| 67 |
+
`str`: Base URL to connect to the container (e.g., `"http://localhost:8000"`).
|
| 68 |
+
|
| 69 |
+
Raises:
|
| 70 |
+
RuntimeError: If container fails to start.
|
| 71 |
+
"""
|
| 72 |
+
pass
|
| 73 |
+
|
| 74 |
+
@abstractmethod
|
| 75 |
+
def stop_container(self) -> None:
|
| 76 |
+
"""
|
| 77 |
+
Stop and remove the running container.
|
| 78 |
+
|
| 79 |
+
This cleans up the container that was started by start_container().
|
| 80 |
+
"""
|
| 81 |
+
pass
|
| 82 |
+
|
| 83 |
+
@abstractmethod
|
| 84 |
+
def wait_for_ready(self, base_url: str, timeout_s: float = 30.0) -> None:
|
| 85 |
+
"""
|
| 86 |
+
Wait for the container to be ready to accept requests.
|
| 87 |
+
|
| 88 |
+
This typically polls the /health endpoint until it returns 200.
|
| 89 |
+
|
| 90 |
+
Args:
|
| 91 |
+
base_url (`str`):
|
| 92 |
+
Base URL of the container.
|
| 93 |
+
timeout_s (`float`, *optional*, defaults to `30.0`):
|
| 94 |
+
Maximum time to wait in seconds.
|
| 95 |
+
|
| 96 |
+
Raises:
|
| 97 |
+
TimeoutError: If container doesn't become ready in time.
|
| 98 |
+
"""
|
| 99 |
+
pass
|
| 100 |
+
|
| 101 |
+
def close(self) -> None:
|
| 102 |
+
"""
|
| 103 |
+
Release provider-held resources (e.g. SDK clients, connections).
|
| 104 |
+
|
| 105 |
+
Defaults to a no-op so existing providers are unaffected. Providers that
|
| 106 |
+
hold external resources beyond the container itself (such as a cloud SDK
|
| 107 |
+
client) override this to release them; it is also invoked on context-
|
| 108 |
+
manager exit. Lightweight providers need not override it.
|
| 109 |
+
"""
|
| 110 |
+
pass
|
| 111 |
+
|
| 112 |
+
def __enter__(self: _ContainerProviderT) -> _ContainerProviderT:
|
| 113 |
+
return self
|
| 114 |
+
|
| 115 |
+
def __exit__(self, exc_type, exc, tb) -> Optional[bool]:
|
| 116 |
+
self.close()
|
| 117 |
+
return None
|
| 118 |
+
|
| 119 |
+
|
| 120 |
+
class LocalDockerProvider(ContainerProvider):
|
| 121 |
+
"""
|
| 122 |
+
Container provider for local Docker daemon.
|
| 123 |
+
|
| 124 |
+
This provider runs containers on the local machine using Docker.
|
| 125 |
+
Useful for development and testing.
|
| 126 |
+
|
| 127 |
+
Examples:
|
| 128 |
+
|
| 129 |
+
```python
|
| 130 |
+
provider = LocalDockerProvider()
|
| 131 |
+
base_url = provider.start_container("echo-env:latest")
|
| 132 |
+
# Container running on http://localhost:<random-port>
|
| 133 |
+
provider.stop_container()
|
| 134 |
+
```
|
| 135 |
+
"""
|
| 136 |
+
|
| 137 |
+
def __init__(self):
|
| 138 |
+
"""Initialize the local Docker provider."""
|
| 139 |
+
self._container_id: Optional[str] = None
|
| 140 |
+
self._container_name: Optional[str] = None
|
| 141 |
+
|
| 142 |
+
# Check if Docker is available
|
| 143 |
+
import subprocess
|
| 144 |
+
|
| 145 |
+
try:
|
| 146 |
+
subprocess.run(
|
| 147 |
+
["docker", "version"],
|
| 148 |
+
check=True,
|
| 149 |
+
capture_output=True,
|
| 150 |
+
timeout=5,
|
| 151 |
+
)
|
| 152 |
+
except (
|
| 153 |
+
subprocess.CalledProcessError,
|
| 154 |
+
FileNotFoundError,
|
| 155 |
+
subprocess.TimeoutExpired,
|
| 156 |
+
):
|
| 157 |
+
raise RuntimeError(
|
| 158 |
+
"Docker is not available. Please install Docker Desktop or Docker Engine."
|
| 159 |
+
)
|
| 160 |
+
|
| 161 |
+
def start_container(
|
| 162 |
+
self,
|
| 163 |
+
image: str,
|
| 164 |
+
port: Optional[int] = None,
|
| 165 |
+
env_vars: Optional[Dict[str, str]] = None,
|
| 166 |
+
**kwargs: Any,
|
| 167 |
+
) -> str:
|
| 168 |
+
"""
|
| 169 |
+
Start a Docker container locally.
|
| 170 |
+
|
| 171 |
+
Args:
|
| 172 |
+
image (`str`):
|
| 173 |
+
Docker image name.
|
| 174 |
+
port (`int`, *optional*):
|
| 175 |
+
Port to expose. If `None`, finds an available port.
|
| 176 |
+
env_vars (`dict`, *optional*):
|
| 177 |
+
Environment variables for the container.
|
| 178 |
+
**kwargs:
|
| 179 |
+
Additional Docker run options.
|
| 180 |
+
|
| 181 |
+
Returns:
|
| 182 |
+
`str`: Base URL to connect to the container.
|
| 183 |
+
"""
|
| 184 |
+
import subprocess
|
| 185 |
+
import time
|
| 186 |
+
|
| 187 |
+
# Find available port if not specified
|
| 188 |
+
if port is None:
|
| 189 |
+
port = self._find_available_port()
|
| 190 |
+
|
| 191 |
+
# Generate container name
|
| 192 |
+
self._container_name = self._generate_container_name(image)
|
| 193 |
+
|
| 194 |
+
# Build docker run command
|
| 195 |
+
cmd = [
|
| 196 |
+
"docker",
|
| 197 |
+
"run",
|
| 198 |
+
"-d", # Detached
|
| 199 |
+
"--name",
|
| 200 |
+
self._container_name,
|
| 201 |
+
"-p",
|
| 202 |
+
f"{port}:8000", # Map port
|
| 203 |
+
]
|
| 204 |
+
|
| 205 |
+
# Add environment variables
|
| 206 |
+
if env_vars:
|
| 207 |
+
for key, value in env_vars.items():
|
| 208 |
+
cmd.extend(["-e", f"{key}={value}"])
|
| 209 |
+
|
| 210 |
+
# Add image
|
| 211 |
+
cmd.append(image)
|
| 212 |
+
|
| 213 |
+
# Run container
|
| 214 |
+
try:
|
| 215 |
+
result = subprocess.run(cmd, capture_output=True, text=True, check=True)
|
| 216 |
+
self._container_id = result.stdout.strip()
|
| 217 |
+
except subprocess.CalledProcessError as e:
|
| 218 |
+
error_msg = f"Failed to start Docker container.\nCommand: {' '.join(cmd)}\nExit code: {e.returncode}\nStderr: {e.stderr}\nStdout: {e.stdout}"
|
| 219 |
+
raise RuntimeError(error_msg) from e
|
| 220 |
+
|
| 221 |
+
# Wait a moment for container to start
|
| 222 |
+
time.sleep(1)
|
| 223 |
+
|
| 224 |
+
base_url = f"http://localhost:{port}"
|
| 225 |
+
return base_url
|
| 226 |
+
|
| 227 |
+
def stop_container(self) -> None:
|
| 228 |
+
"""
|
| 229 |
+
Stop and remove the Docker container.
|
| 230 |
+
"""
|
| 231 |
+
if self._container_id is None:
|
| 232 |
+
return
|
| 233 |
+
|
| 234 |
+
import subprocess
|
| 235 |
+
|
| 236 |
+
try:
|
| 237 |
+
# Stop container
|
| 238 |
+
subprocess.run(
|
| 239 |
+
["docker", "stop", self._container_id],
|
| 240 |
+
capture_output=True,
|
| 241 |
+
check=True,
|
| 242 |
+
timeout=10,
|
| 243 |
+
)
|
| 244 |
+
|
| 245 |
+
# Remove container
|
| 246 |
+
subprocess.run(
|
| 247 |
+
["docker", "rm", self._container_id],
|
| 248 |
+
capture_output=True,
|
| 249 |
+
check=True,
|
| 250 |
+
timeout=10,
|
| 251 |
+
)
|
| 252 |
+
except subprocess.CalledProcessError:
|
| 253 |
+
# Container might already be stopped/removed
|
| 254 |
+
pass
|
| 255 |
+
finally:
|
| 256 |
+
self._container_id = None
|
| 257 |
+
self._container_name = None
|
| 258 |
+
|
| 259 |
+
def wait_for_ready(self, base_url: str, timeout_s: float = 30.0) -> None:
|
| 260 |
+
"""
|
| 261 |
+
Wait for container to be ready by polling /health endpoint.
|
| 262 |
+
|
| 263 |
+
Args:
|
| 264 |
+
base_url (`str`):
|
| 265 |
+
Base URL of the container.
|
| 266 |
+
timeout_s (`float`, *optional*, defaults to `30.0`):
|
| 267 |
+
Maximum time to wait in seconds.
|
| 268 |
+
|
| 269 |
+
Raises:
|
| 270 |
+
TimeoutError: If container doesn't become ready.
|
| 271 |
+
"""
|
| 272 |
+
import time
|
| 273 |
+
|
| 274 |
+
import requests
|
| 275 |
+
|
| 276 |
+
start_time = time.time()
|
| 277 |
+
health_url = f"{base_url}/health"
|
| 278 |
+
|
| 279 |
+
# Bypass proxy for localhost to avoid proxy issues
|
| 280 |
+
proxies = {"http": None, "https": None}
|
| 281 |
+
|
| 282 |
+
while time.time() - start_time < timeout_s:
|
| 283 |
+
try:
|
| 284 |
+
response = requests.get(health_url, timeout=2.0, proxies=proxies)
|
| 285 |
+
if response.status_code == 200:
|
| 286 |
+
return
|
| 287 |
+
except requests.RequestException:
|
| 288 |
+
pass
|
| 289 |
+
|
| 290 |
+
time.sleep(0.5)
|
| 291 |
+
|
| 292 |
+
raise TimeoutError(
|
| 293 |
+
f"Container at {base_url} did not become ready within {timeout_s}s"
|
| 294 |
+
)
|
| 295 |
+
|
| 296 |
+
def _find_available_port(self) -> int:
|
| 297 |
+
"""
|
| 298 |
+
Find an available port on localhost.
|
| 299 |
+
|
| 300 |
+
Returns:
|
| 301 |
+
`int`: An available port number.
|
| 302 |
+
"""
|
| 303 |
+
import socket
|
| 304 |
+
|
| 305 |
+
with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
|
| 306 |
+
s.bind(("", 0))
|
| 307 |
+
s.listen(1)
|
| 308 |
+
port = s.getsockname()[1]
|
| 309 |
+
return port
|
| 310 |
+
|
| 311 |
+
def _generate_container_name(self, image: str) -> str:
|
| 312 |
+
"""
|
| 313 |
+
Generate a unique container name based on image name and timestamp.
|
| 314 |
+
|
| 315 |
+
Args:
|
| 316 |
+
image (`str`):
|
| 317 |
+
Docker image name.
|
| 318 |
+
|
| 319 |
+
Returns:
|
| 320 |
+
`str`: A unique container name.
|
| 321 |
+
"""
|
| 322 |
+
import time
|
| 323 |
+
|
| 324 |
+
clean_image = image.split("/")[-1].split(":")[0]
|
| 325 |
+
timestamp = int(time.time() * 1000)
|
| 326 |
+
return f"{clean_image}-{timestamp}"
|
| 327 |
+
|
| 328 |
+
|
| 329 |
+
class DockerSwarmProvider(ContainerProvider):
|
| 330 |
+
"""
|
| 331 |
+
Container provider that uses Docker Swarm services for local concurrency.
|
| 332 |
+
|
| 333 |
+
This provider creates a replicated Swarm service backed by the local Docker
|
| 334 |
+
engine. The built-in load-balancer fans requests across the replicas,
|
| 335 |
+
allowing multiple container instances to run concurrently on the developer
|
| 336 |
+
workstation (mirroring the workflow described in the Docker stack docs).
|
| 337 |
+
"""
|
| 338 |
+
|
| 339 |
+
def __init__(
|
| 340 |
+
self,
|
| 341 |
+
*,
|
| 342 |
+
auto_init_swarm: bool = True,
|
| 343 |
+
overlay_network: Optional[str] = None,
|
| 344 |
+
):
|
| 345 |
+
"""
|
| 346 |
+
Args:
|
| 347 |
+
auto_init_swarm (`bool`, *optional*, defaults to `True`):
|
| 348 |
+
Whether to call `docker swarm init` when Swarm is not active.
|
| 349 |
+
Otherwise, the user must manually initialize Swarm.
|
| 350 |
+
overlay_network (`str`, *optional*):
|
| 351 |
+
Overlay network name for the service. When provided, the network
|
| 352 |
+
is created with `docker network create --driver overlay --attachable`
|
| 353 |
+
if it does not already exist.
|
| 354 |
+
"""
|
| 355 |
+
self._service_name: Optional[str] = None
|
| 356 |
+
self._service_id: Optional[str] = None
|
| 357 |
+
self._published_port: Optional[int] = None
|
| 358 |
+
self._overlay_network = overlay_network
|
| 359 |
+
self._auto_init_swarm = auto_init_swarm
|
| 360 |
+
|
| 361 |
+
self._ensure_docker_available()
|
| 362 |
+
self._ensure_swarm_initialized()
|
| 363 |
+
if self._overlay_network:
|
| 364 |
+
self._ensure_overlay_network(self._overlay_network)
|
| 365 |
+
|
| 366 |
+
def start_container(
|
| 367 |
+
self,
|
| 368 |
+
image: str,
|
| 369 |
+
port: Optional[int] = None,
|
| 370 |
+
env_vars: Optional[Dict[str, str]] = None,
|
| 371 |
+
**kwargs: Any,
|
| 372 |
+
) -> str:
|
| 373 |
+
"""
|
| 374 |
+
Start (or scale) a Swarm service for the given image.
|
| 375 |
+
|
| 376 |
+
Args:
|
| 377 |
+
image (`str`):
|
| 378 |
+
Docker image name.
|
| 379 |
+
port (`int`, *optional*):
|
| 380 |
+
Port to expose. If `None`, finds an available port.
|
| 381 |
+
env_vars (`dict`, *optional*):
|
| 382 |
+
Environment variables for the container.
|
| 383 |
+
replicas (`int`, *optional*, defaults to `2`):
|
| 384 |
+
Number of container replicas.
|
| 385 |
+
cpu_limit (`float` or `str`, *optional*):
|
| 386 |
+
CPU limit passed to `--limit-cpu`.
|
| 387 |
+
memory_limit (`str`, *optional*):
|
| 388 |
+
Memory limit passed to `--limit-memory`.
|
| 389 |
+
constraints (`Sequence[str]`, *optional*):
|
| 390 |
+
Placement constraints.
|
| 391 |
+
labels (`dict`, *optional*):
|
| 392 |
+
Service labels.
|
| 393 |
+
command (`Sequence[str]` or `str`, *optional*):
|
| 394 |
+
Override container command.
|
| 395 |
+
|
| 396 |
+
Returns:
|
| 397 |
+
`str`: Base URL to connect to the service.
|
| 398 |
+
"""
|
| 399 |
+
import shlex
|
| 400 |
+
import subprocess
|
| 401 |
+
import time
|
| 402 |
+
|
| 403 |
+
allowed_kwargs = {
|
| 404 |
+
"replicas",
|
| 405 |
+
"cpu_limit",
|
| 406 |
+
"memory_limit",
|
| 407 |
+
"constraints",
|
| 408 |
+
"labels",
|
| 409 |
+
"command",
|
| 410 |
+
}
|
| 411 |
+
unknown = set(kwargs) - allowed_kwargs
|
| 412 |
+
if unknown:
|
| 413 |
+
raise ValueError(f"Unsupported kwargs for DockerSwarmProvider: {unknown}")
|
| 414 |
+
|
| 415 |
+
replicas = int(kwargs.get("replicas", 2))
|
| 416 |
+
cpu_limit = kwargs.get("cpu_limit")
|
| 417 |
+
memory_limit = kwargs.get("memory_limit")
|
| 418 |
+
constraints: Optional[Sequence[str]] = kwargs.get("constraints")
|
| 419 |
+
labels: Optional[Dict[str, str]] = kwargs.get("labels")
|
| 420 |
+
command_override = kwargs.get("command")
|
| 421 |
+
|
| 422 |
+
if port is None:
|
| 423 |
+
port = self._find_available_port()
|
| 424 |
+
|
| 425 |
+
self._service_name = self._generate_service_name(image)
|
| 426 |
+
self._published_port = port
|
| 427 |
+
|
| 428 |
+
cmd = [
|
| 429 |
+
"docker",
|
| 430 |
+
"service",
|
| 431 |
+
"create",
|
| 432 |
+
"--detach",
|
| 433 |
+
"--name",
|
| 434 |
+
self._service_name,
|
| 435 |
+
"--replicas",
|
| 436 |
+
str(max(1, replicas)),
|
| 437 |
+
"--publish",
|
| 438 |
+
f"{port}:8000",
|
| 439 |
+
]
|
| 440 |
+
|
| 441 |
+
if self._overlay_network:
|
| 442 |
+
cmd.extend(["--network", self._overlay_network])
|
| 443 |
+
|
| 444 |
+
if env_vars:
|
| 445 |
+
for key, value in env_vars.items():
|
| 446 |
+
cmd.extend(["--env", f"{key}={value}"])
|
| 447 |
+
|
| 448 |
+
if cpu_limit is not None:
|
| 449 |
+
cmd.extend(["--limit-cpu", str(cpu_limit)])
|
| 450 |
+
|
| 451 |
+
if memory_limit is not None:
|
| 452 |
+
cmd.extend(["--limit-memory", str(memory_limit)])
|
| 453 |
+
|
| 454 |
+
if constraints:
|
| 455 |
+
for constraint in constraints:
|
| 456 |
+
cmd.extend(["--constraint", constraint])
|
| 457 |
+
|
| 458 |
+
if labels:
|
| 459 |
+
for key, value in labels.items():
|
| 460 |
+
cmd.extend(["--label", f"{key}={value}"])
|
| 461 |
+
|
| 462 |
+
cmd.append(image)
|
| 463 |
+
|
| 464 |
+
if command_override:
|
| 465 |
+
if isinstance(command_override, str):
|
| 466 |
+
cmd.extend(shlex.split(command_override))
|
| 467 |
+
else:
|
| 468 |
+
cmd.extend(command_override)
|
| 469 |
+
|
| 470 |
+
try:
|
| 471 |
+
result = subprocess.run(
|
| 472 |
+
cmd,
|
| 473 |
+
capture_output=True,
|
| 474 |
+
text=True,
|
| 475 |
+
check=True,
|
| 476 |
+
)
|
| 477 |
+
self._service_id = result.stdout.strip()
|
| 478 |
+
except subprocess.CalledProcessError as e:
|
| 479 |
+
error_msg = (
|
| 480 |
+
"Failed to start Docker Swarm service.\n"
|
| 481 |
+
f"Command: {' '.join(cmd)}\n"
|
| 482 |
+
f"Exit code: {e.returncode}\n"
|
| 483 |
+
f"Stdout: {e.stdout}\n"
|
| 484 |
+
f"Stderr: {e.stderr}"
|
| 485 |
+
)
|
| 486 |
+
raise RuntimeError(error_msg) from e
|
| 487 |
+
|
| 488 |
+
# Give Swarm a brief moment to schedule the tasks.
|
| 489 |
+
time.sleep(1.0)
|
| 490 |
+
|
| 491 |
+
return f"http://localhost:{port}"
|
| 492 |
+
|
| 493 |
+
def stop_container(self) -> None:
|
| 494 |
+
"""
|
| 495 |
+
Remove the Swarm service (and keep the Swarm manager running).
|
| 496 |
+
"""
|
| 497 |
+
if not self._service_name:
|
| 498 |
+
return
|
| 499 |
+
|
| 500 |
+
import subprocess
|
| 501 |
+
|
| 502 |
+
try:
|
| 503 |
+
subprocess.run(
|
| 504 |
+
["docker", "service", "rm", self._service_name],
|
| 505 |
+
capture_output=True,
|
| 506 |
+
check=True,
|
| 507 |
+
timeout=10,
|
| 508 |
+
)
|
| 509 |
+
except subprocess.CalledProcessError:
|
| 510 |
+
# Service may already be gone; ignore.
|
| 511 |
+
pass
|
| 512 |
+
finally:
|
| 513 |
+
self._service_name = None
|
| 514 |
+
self._service_id = None
|
| 515 |
+
self._published_port = None
|
| 516 |
+
|
| 517 |
+
def wait_for_ready(self, base_url: str, timeout_s: float = 30.0) -> None:
|
| 518 |
+
"""
|
| 519 |
+
Wait for at least one replica to become healthy by polling /health.
|
| 520 |
+
|
| 521 |
+
With Swarm's load balancer, requests round-robin across replicas,
|
| 522 |
+
so this only verifies that at least one replica is responding. Some
|
| 523 |
+
replicas may still be starting when this returns.
|
| 524 |
+
"""
|
| 525 |
+
import time
|
| 526 |
+
|
| 527 |
+
import requests
|
| 528 |
+
|
| 529 |
+
deadline = time.time() + timeout_s
|
| 530 |
+
health_url = f"{base_url}/health"
|
| 531 |
+
|
| 532 |
+
# Bypass proxy for localhost to avoid proxy issues
|
| 533 |
+
proxies = {"http": None, "https": None}
|
| 534 |
+
|
| 535 |
+
while time.time() < deadline:
|
| 536 |
+
try:
|
| 537 |
+
response = requests.get(health_url, timeout=2.0, proxies=proxies)
|
| 538 |
+
if response.status_code == 200:
|
| 539 |
+
return
|
| 540 |
+
except requests.RequestException:
|
| 541 |
+
pass
|
| 542 |
+
|
| 543 |
+
time.sleep(0.5)
|
| 544 |
+
|
| 545 |
+
raise TimeoutError(
|
| 546 |
+
f"Swarm service at {base_url} did not become ready within {timeout_s}s"
|
| 547 |
+
)
|
| 548 |
+
|
| 549 |
+
def _ensure_docker_available(self) -> None:
|
| 550 |
+
import subprocess
|
| 551 |
+
|
| 552 |
+
try:
|
| 553 |
+
subprocess.run(
|
| 554 |
+
["docker", "version"],
|
| 555 |
+
check=True,
|
| 556 |
+
capture_output=True,
|
| 557 |
+
timeout=5,
|
| 558 |
+
)
|
| 559 |
+
except (
|
| 560 |
+
subprocess.CalledProcessError,
|
| 561 |
+
FileNotFoundError,
|
| 562 |
+
subprocess.TimeoutExpired,
|
| 563 |
+
) as exc:
|
| 564 |
+
raise RuntimeError(
|
| 565 |
+
"Docker is not available. Please install Docker Desktop or Docker Engine."
|
| 566 |
+
) from exc
|
| 567 |
+
|
| 568 |
+
def _ensure_swarm_initialized(self) -> None:
|
| 569 |
+
import subprocess
|
| 570 |
+
|
| 571 |
+
try:
|
| 572 |
+
result = subprocess.run(
|
| 573 |
+
["docker", "info", "--format", "{{.Swarm.LocalNodeState}}"],
|
| 574 |
+
capture_output=True,
|
| 575 |
+
text=True,
|
| 576 |
+
check=True,
|
| 577 |
+
timeout=5,
|
| 578 |
+
)
|
| 579 |
+
state = result.stdout.strip().lower()
|
| 580 |
+
if state == "active":
|
| 581 |
+
return
|
| 582 |
+
except subprocess.CalledProcessError:
|
| 583 |
+
state = "unknown"
|
| 584 |
+
|
| 585 |
+
if not self._auto_init_swarm:
|
| 586 |
+
raise RuntimeError(
|
| 587 |
+
f"Docker Swarm is not active (state={state}). Enable Swarm manually or pass auto_init_swarm=True."
|
| 588 |
+
)
|
| 589 |
+
|
| 590 |
+
try:
|
| 591 |
+
subprocess.run(
|
| 592 |
+
["docker", "swarm", "init"],
|
| 593 |
+
check=True,
|
| 594 |
+
capture_output=True,
|
| 595 |
+
timeout=10,
|
| 596 |
+
)
|
| 597 |
+
except subprocess.CalledProcessError as e:
|
| 598 |
+
raise RuntimeError("Failed to initialize Docker Swarm") from e
|
| 599 |
+
|
| 600 |
+
def _ensure_overlay_network(self, network: str) -> None:
|
| 601 |
+
import subprocess
|
| 602 |
+
|
| 603 |
+
inspect = subprocess.run(
|
| 604 |
+
["docker", "network", "inspect", network],
|
| 605 |
+
capture_output=True,
|
| 606 |
+
text=True,
|
| 607 |
+
check=False,
|
| 608 |
+
)
|
| 609 |
+
if inspect.returncode == 0:
|
| 610 |
+
return
|
| 611 |
+
|
| 612 |
+
try:
|
| 613 |
+
subprocess.run(
|
| 614 |
+
[
|
| 615 |
+
"docker",
|
| 616 |
+
"network",
|
| 617 |
+
"create",
|
| 618 |
+
"--driver",
|
| 619 |
+
"overlay",
|
| 620 |
+
"--attachable",
|
| 621 |
+
network,
|
| 622 |
+
],
|
| 623 |
+
check=True,
|
| 624 |
+
capture_output=True,
|
| 625 |
+
timeout=10,
|
| 626 |
+
)
|
| 627 |
+
except subprocess.CalledProcessError as e:
|
| 628 |
+
raise RuntimeError(f"Failed to create overlay network '{network}'") from e
|
| 629 |
+
|
| 630 |
+
def _find_available_port(self) -> int:
|
| 631 |
+
import socket
|
| 632 |
+
|
| 633 |
+
with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
|
| 634 |
+
s.bind(("", 0))
|
| 635 |
+
s.listen(1)
|
| 636 |
+
port = s.getsockname()[1]
|
| 637 |
+
return port
|
| 638 |
+
|
| 639 |
+
def _generate_service_name(self, image: str) -> str:
|
| 640 |
+
import time
|
| 641 |
+
|
| 642 |
+
clean_image = image.split("/")[-1].split(":")[0]
|
| 643 |
+
timestamp = int(time.time() * 1000)
|
| 644 |
+
return f"{clean_image}-swarm-{timestamp}"
|
| 645 |
+
|
| 646 |
+
|
| 647 |
+
class KubernetesProvider(ContainerProvider):
|
| 648 |
+
"""
|
| 649 |
+
Planned container provider for Kubernetes clusters.
|
| 650 |
+
|
| 651 |
+
Not yet implemented: this is a placeholder for the planned Kubernetes
|
| 652 |
+
backend and does not implement the abstract `ContainerProvider` methods, so
|
| 653 |
+
it cannot be instantiated. Use `LocalDockerProvider`, `DockerSwarmProvider`,
|
| 654 |
+
`DaytonaProvider`, or `ACASandboxProvider` instead.
|
| 655 |
+
"""
|
| 656 |
+
|
| 657 |
+
pass
|
| 658 |
+
|
| 659 |
+
|
| 660 |
+
class RuntimeProvider(ABC):
|
| 661 |
+
"""
|
| 662 |
+
Abstract base class for runtime providers that are not container providers.
|
| 663 |
+
Providers implement this interface to support different runtime platforms:
|
| 664 |
+
- UVProvider: Runs environments via `uv run`
|
| 665 |
+
|
| 666 |
+
The provider manages a single runtime lifecycle and provides the base URL
|
| 667 |
+
for connecting to it.
|
| 668 |
+
|
| 669 |
+
Examples:
|
| 670 |
+
|
| 671 |
+
```python
|
| 672 |
+
provider = UVProvider(project_path="/path/to/env")
|
| 673 |
+
base_url = provider.start()
|
| 674 |
+
print(base_url) # http://localhost:8000
|
| 675 |
+
provider.stop()
|
| 676 |
+
```
|
| 677 |
+
"""
|
| 678 |
+
|
| 679 |
+
@abstractmethod
|
| 680 |
+
def start(
|
| 681 |
+
self,
|
| 682 |
+
port: Optional[int] = None,
|
| 683 |
+
env_vars: Optional[Dict[str, str]] = None,
|
| 684 |
+
**kwargs: Any,
|
| 685 |
+
) -> str:
|
| 686 |
+
"""
|
| 687 |
+
Start the runtime.
|
| 688 |
+
|
| 689 |
+
Args:
|
| 690 |
+
port (`int`, *optional*):
|
| 691 |
+
Port to expose. If `None`, the provider chooses.
|
| 692 |
+
env_vars (`dict`, *optional*):
|
| 693 |
+
Environment variables for the runtime.
|
| 694 |
+
**kwargs:
|
| 695 |
+
Additional runtime options.
|
| 696 |
+
|
| 697 |
+
Returns:
|
| 698 |
+
`str`: Base URL to connect to the runtime.
|
| 699 |
+
"""
|
| 700 |
+
|
| 701 |
+
@abstractmethod
|
| 702 |
+
def stop(self) -> None:
|
| 703 |
+
"""
|
| 704 |
+
Stop the runtime.
|
| 705 |
+
"""
|
| 706 |
+
pass
|
| 707 |
+
|
| 708 |
+
@abstractmethod
|
| 709 |
+
def wait_for_ready(self, timeout_s: float = 30.0) -> None:
|
| 710 |
+
"""
|
| 711 |
+
Wait for the runtime to be ready to accept requests.
|
| 712 |
+
"""
|
| 713 |
+
pass
|
| 714 |
+
|
| 715 |
+
def __enter__(self) -> "RuntimeProvider":
|
| 716 |
+
"""
|
| 717 |
+
Enter the runtime provider.
|
| 718 |
+
"""
|
| 719 |
+
self.start()
|
| 720 |
+
return self
|
| 721 |
+
|
| 722 |
+
def __exit__(self, exc_type, exc, tb) -> None:
|
| 723 |
+
"""
|
| 724 |
+
Exit the runtime provider.
|
| 725 |
+
"""
|
| 726 |
+
self.stop()
|
| 727 |
+
return False
|
src/core/containers/runtime/uv_provider.py
ADDED
|
@@ -0,0 +1,309 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""Providers for launching ASGI applications via ``uv run``."""
|
| 2 |
+
|
| 3 |
+
from __future__ import annotations
|
| 4 |
+
|
| 5 |
+
import os
|
| 6 |
+
import shutil
|
| 7 |
+
import socket
|
| 8 |
+
import subprocess
|
| 9 |
+
import tempfile
|
| 10 |
+
import time
|
| 11 |
+
from typing import Dict, Optional
|
| 12 |
+
|
| 13 |
+
import requests
|
| 14 |
+
|
| 15 |
+
from .providers import RuntimeProvider
|
| 16 |
+
|
| 17 |
+
GIT_URL_PREFIX = "git+"
|
| 18 |
+
|
| 19 |
+
|
| 20 |
+
def _check_uv_installed() -> None:
|
| 21 |
+
try:
|
| 22 |
+
subprocess.check_output(["uv", "--version"])
|
| 23 |
+
except FileNotFoundError as exc:
|
| 24 |
+
raise RuntimeError(
|
| 25 |
+
"`uv` executable not found. Install uv from https://docs.astral.sh and ensure it is on PATH."
|
| 26 |
+
) from exc
|
| 27 |
+
|
| 28 |
+
|
| 29 |
+
def _clone_git_project(git_url: str, timeout_s: float) -> str:
|
| 30 |
+
"""Clone a `git+<url>` spec to a temp dir and return its local path.
|
| 31 |
+
|
| 32 |
+
`uv run --project` only discovers a project in a local directory (per
|
| 33 |
+
`uv run --help`) -- it has no support for remote git sources. Callers
|
| 34 |
+
that accept a `git+...` `project_path` (e.g. `EnvClient.from_env`)
|
| 35 |
+
must therefore clone it themselves before handing a path to `uv run`.
|
| 36 |
+
"""
|
| 37 |
+
repo_url = git_url[len(GIT_URL_PREFIX) :]
|
| 38 |
+
clone_dir = tempfile.mkdtemp(prefix="openenv-uv-clone-")
|
| 39 |
+
try:
|
| 40 |
+
subprocess.run(
|
| 41 |
+
["git", "clone", "--depth", "1", repo_url, clone_dir],
|
| 42 |
+
check=True,
|
| 43 |
+
capture_output=True,
|
| 44 |
+
text=True,
|
| 45 |
+
# A hung/slow remote would otherwise block start() (and the
|
| 46 |
+
# caller's own readiness timeout) indefinitely.
|
| 47 |
+
timeout=timeout_s,
|
| 48 |
+
)
|
| 49 |
+
except FileNotFoundError as exc:
|
| 50 |
+
shutil.rmtree(clone_dir, ignore_errors=True)
|
| 51 |
+
raise RuntimeError(
|
| 52 |
+
f"`git` executable not found; required to clone project_path {git_url!r}."
|
| 53 |
+
) from exc
|
| 54 |
+
except subprocess.CalledProcessError as exc:
|
| 55 |
+
shutil.rmtree(clone_dir, ignore_errors=True)
|
| 56 |
+
raise RuntimeError(f"Failed to clone {repo_url!r}: {exc.stderr}") from exc
|
| 57 |
+
except subprocess.TimeoutExpired as exc:
|
| 58 |
+
shutil.rmtree(clone_dir, ignore_errors=True)
|
| 59 |
+
raise RuntimeError(
|
| 60 |
+
f"Timed out cloning {repo_url!r} after {timeout_s:.1f}s"
|
| 61 |
+
) from exc
|
| 62 |
+
return clone_dir
|
| 63 |
+
|
| 64 |
+
|
| 65 |
+
def _find_free_port() -> int:
|
| 66 |
+
with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock:
|
| 67 |
+
sock.bind(("", 0))
|
| 68 |
+
sock.listen(1)
|
| 69 |
+
return sock.getsockname()[1]
|
| 70 |
+
|
| 71 |
+
|
| 72 |
+
def _create_uv_command(
|
| 73 |
+
*,
|
| 74 |
+
host: str,
|
| 75 |
+
port: int,
|
| 76 |
+
reload: bool,
|
| 77 |
+
workers: int,
|
| 78 |
+
app: str,
|
| 79 |
+
project_path: str,
|
| 80 |
+
) -> list[str]:
|
| 81 |
+
command: list[str] = ["uv", "run", "--isolated", "--project", project_path]
|
| 82 |
+
|
| 83 |
+
command.append("--")
|
| 84 |
+
command.extend(
|
| 85 |
+
[
|
| 86 |
+
"uvicorn",
|
| 87 |
+
app,
|
| 88 |
+
"--host",
|
| 89 |
+
host,
|
| 90 |
+
"--port",
|
| 91 |
+
str(port),
|
| 92 |
+
"--workers",
|
| 93 |
+
str(workers),
|
| 94 |
+
]
|
| 95 |
+
)
|
| 96 |
+
|
| 97 |
+
if reload:
|
| 98 |
+
command.append("--reload")
|
| 99 |
+
|
| 100 |
+
return command
|
| 101 |
+
|
| 102 |
+
|
| 103 |
+
def _poll_health(health_url: str, timeout_s: float) -> None:
|
| 104 |
+
"""Poll a health endpoint until it returns HTTP 200 or times out."""
|
| 105 |
+
|
| 106 |
+
deadline = time.time() + timeout_s
|
| 107 |
+
while time.time() < deadline:
|
| 108 |
+
try:
|
| 109 |
+
timeout = max(0.0001, min(deadline - time.time(), 2.0))
|
| 110 |
+
response = requests.get(health_url, timeout=timeout)
|
| 111 |
+
if response.status_code == 200:
|
| 112 |
+
return
|
| 113 |
+
except requests.RequestException:
|
| 114 |
+
# Server not accepting connections yet. Fall through to the sleep
|
| 115 |
+
# below rather than `continue`-ing: a refused connection returns
|
| 116 |
+
# immediately, so retrying without a pause busy-spins a CPU core
|
| 117 |
+
# for the whole timeout window while the server boots.
|
| 118 |
+
pass
|
| 119 |
+
|
| 120 |
+
time.sleep(0.5)
|
| 121 |
+
|
| 122 |
+
raise TimeoutError(f"Server did not become ready within {timeout_s:.1f} seconds")
|
| 123 |
+
|
| 124 |
+
|
| 125 |
+
class UVProvider(RuntimeProvider):
|
| 126 |
+
"""
|
| 127 |
+
RuntimeProvider implementation backed by ``uv run``.
|
| 128 |
+
|
| 129 |
+
Args:
|
| 130 |
+
project_path (`str`):
|
| 131 |
+
Local path to a uv project (passed to `uv run --project`), or a
|
| 132 |
+
`git+<url>` spec that is cloned to a temp directory on `start()`.
|
| 133 |
+
app (`str`, *optional*, defaults to `"server.app:app"`):
|
| 134 |
+
ASGI application path for uvicorn.
|
| 135 |
+
host (`str`, *optional*, defaults to `"0.0.0.0"`):
|
| 136 |
+
Host interface to bind to.
|
| 137 |
+
reload (`bool`, *optional*, defaults to `False`):
|
| 138 |
+
Whether to enable uvicorn's reload mode.
|
| 139 |
+
env_vars (`dict`, *optional*):
|
| 140 |
+
Environment variables to pass through to the spawned process.
|
| 141 |
+
context_timeout_s (`float`, *optional*, defaults to `60.0`):
|
| 142 |
+
How long to wait for the environment to become ready.
|
| 143 |
+
|
| 144 |
+
Examples:
|
| 145 |
+
|
| 146 |
+
```python
|
| 147 |
+
provider = UVProvider(project_path="/path/to/env")
|
| 148 |
+
base_url = provider.start()
|
| 149 |
+
print(base_url) # http://localhost:8000
|
| 150 |
+
# Use the environment via base_url
|
| 151 |
+
provider.stop()
|
| 152 |
+
```
|
| 153 |
+
"""
|
| 154 |
+
|
| 155 |
+
def __init__(
|
| 156 |
+
self,
|
| 157 |
+
*,
|
| 158 |
+
project_path: str,
|
| 159 |
+
app: str = "server.app:app",
|
| 160 |
+
host: str = "0.0.0.0",
|
| 161 |
+
reload: bool = False,
|
| 162 |
+
env_vars: Optional[Dict[str, str]] = None,
|
| 163 |
+
context_timeout_s: float = 60.0,
|
| 164 |
+
):
|
| 165 |
+
"""Initialize the UVProvider."""
|
| 166 |
+
# `os.path.abspath` would mangle a `git+<url>` spec (e.g. collapsing
|
| 167 |
+
# `https://` to `https:/`) and is meaningless for one anyway -- only
|
| 168 |
+
# resolve it for genuine local paths. Git specs are cloned in `start()`.
|
| 169 |
+
self.project_path = (
|
| 170 |
+
project_path
|
| 171 |
+
if project_path.startswith(GIT_URL_PREFIX)
|
| 172 |
+
else os.path.abspath(project_path)
|
| 173 |
+
)
|
| 174 |
+
self.app = app
|
| 175 |
+
self.host = host
|
| 176 |
+
self.reload = reload
|
| 177 |
+
self.env_vars = env_vars
|
| 178 |
+
self.context_timeout_s = context_timeout_s
|
| 179 |
+
_check_uv_installed()
|
| 180 |
+
self._process = None
|
| 181 |
+
self._base_url = None
|
| 182 |
+
self._clone_dir: Optional[str] = None
|
| 183 |
+
|
| 184 |
+
def start(
|
| 185 |
+
self,
|
| 186 |
+
port: Optional[int] = None,
|
| 187 |
+
env_vars: Optional[Dict[str, str]] = None,
|
| 188 |
+
workers: int = 1,
|
| 189 |
+
**_: Dict[str, str],
|
| 190 |
+
) -> str:
|
| 191 |
+
"""
|
| 192 |
+
Start the environment via `uv run`.
|
| 193 |
+
|
| 194 |
+
Args:
|
| 195 |
+
port (`int`, *optional*):
|
| 196 |
+
The port to bind the environment to.
|
| 197 |
+
env_vars (`dict`, *optional*):
|
| 198 |
+
Environment variables to pass to the environment.
|
| 199 |
+
workers (`int`, *optional*, defaults to `1`):
|
| 200 |
+
The number of workers to use.
|
| 201 |
+
|
| 202 |
+
Returns:
|
| 203 |
+
`str`: Base URL of the environment.
|
| 204 |
+
|
| 205 |
+
Raises:
|
| 206 |
+
RuntimeError: If the environment is already running.
|
| 207 |
+
"""
|
| 208 |
+
if self._process is not None and self._process.poll() is None:
|
| 209 |
+
raise RuntimeError("UVProvider is already running")
|
| 210 |
+
|
| 211 |
+
bind_port = port or _find_free_port()
|
| 212 |
+
|
| 213 |
+
# A previous start() may have left a clone dir behind if the spawned
|
| 214 |
+
# process later died on its own -- the "already running" guard above
|
| 215 |
+
# only checks liveness, so a restart in that state would otherwise
|
| 216 |
+
# overwrite self._clone_dir and leak the old directory.
|
| 217 |
+
if self._clone_dir is not None:
|
| 218 |
+
shutil.rmtree(self._clone_dir, ignore_errors=True)
|
| 219 |
+
self._clone_dir = None
|
| 220 |
+
|
| 221 |
+
project_path = self.project_path
|
| 222 |
+
if project_path.startswith(GIT_URL_PREFIX):
|
| 223 |
+
self._clone_dir = _clone_git_project(project_path, self.context_timeout_s)
|
| 224 |
+
project_path = self._clone_dir
|
| 225 |
+
|
| 226 |
+
command = _create_uv_command(
|
| 227 |
+
host=self.host,
|
| 228 |
+
port=bind_port,
|
| 229 |
+
reload=self.reload,
|
| 230 |
+
workers=workers,
|
| 231 |
+
app=self.app,
|
| 232 |
+
project_path=project_path,
|
| 233 |
+
)
|
| 234 |
+
|
| 235 |
+
env = os.environ.copy()
|
| 236 |
+
|
| 237 |
+
if self.env_vars:
|
| 238 |
+
env.update(self.env_vars)
|
| 239 |
+
if env_vars:
|
| 240 |
+
env.update(env_vars)
|
| 241 |
+
|
| 242 |
+
try:
|
| 243 |
+
self._process = subprocess.Popen(command, env=env)
|
| 244 |
+
except OSError as exc:
|
| 245 |
+
if self._clone_dir is not None:
|
| 246 |
+
shutil.rmtree(self._clone_dir, ignore_errors=True)
|
| 247 |
+
self._clone_dir = None
|
| 248 |
+
raise RuntimeError(f"Failed to launch `uv run`: {exc}") from exc
|
| 249 |
+
|
| 250 |
+
client_host = "127.0.0.1" if self.host in {"0.0.0.0", "::"} else self.host
|
| 251 |
+
self._base_url = f"http://{client_host}:{bind_port}"
|
| 252 |
+
return self._base_url
|
| 253 |
+
|
| 254 |
+
def wait_for_ready(self, timeout_s: float | None = None) -> None:
|
| 255 |
+
"""
|
| 256 |
+
Wait for the environment to become ready.
|
| 257 |
+
|
| 258 |
+
Args:
|
| 259 |
+
timeout_s (`float`, *optional*):
|
| 260 |
+
Maximum time in seconds to wait for the environment to become
|
| 261 |
+
ready. Defaults to the provider's `context_timeout_s`.
|
| 262 |
+
|
| 263 |
+
Raises:
|
| 264 |
+
RuntimeError: If the environment is not running.
|
| 265 |
+
TimeoutError: If the environment does not become ready within the timeout.
|
| 266 |
+
"""
|
| 267 |
+
if timeout_s is None:
|
| 268 |
+
timeout_s = self.context_timeout_s
|
| 269 |
+
|
| 270 |
+
if self._process and self._process.poll() is not None:
|
| 271 |
+
code = self._process.returncode
|
| 272 |
+
raise RuntimeError(f"uv process exited prematurely with code {code}")
|
| 273 |
+
|
| 274 |
+
_poll_health(f"{self._base_url}/health", timeout_s=timeout_s)
|
| 275 |
+
|
| 276 |
+
def stop(self) -> None:
|
| 277 |
+
"""
|
| 278 |
+
Stop the environment.
|
| 279 |
+
"""
|
| 280 |
+
if self._process is not None:
|
| 281 |
+
if self._process.poll() is None:
|
| 282 |
+
self._process.terminate()
|
| 283 |
+
try:
|
| 284 |
+
self._process.wait(timeout=10.0)
|
| 285 |
+
except subprocess.TimeoutExpired:
|
| 286 |
+
self._process.kill()
|
| 287 |
+
self._process.wait(timeout=5.0)
|
| 288 |
+
|
| 289 |
+
self._process = None
|
| 290 |
+
self._base_url = None
|
| 291 |
+
|
| 292 |
+
if self._clone_dir is not None:
|
| 293 |
+
shutil.rmtree(self._clone_dir, ignore_errors=True)
|
| 294 |
+
self._clone_dir = None
|
| 295 |
+
|
| 296 |
+
@property
|
| 297 |
+
def base_url(self) -> str:
|
| 298 |
+
"""
|
| 299 |
+
The base URL of the environment.
|
| 300 |
+
|
| 301 |
+
Returns:
|
| 302 |
+
`str`: Base URL of the running environment.
|
| 303 |
+
|
| 304 |
+
Raises:
|
| 305 |
+
RuntimeError: If the environment has not been started.
|
| 306 |
+
"""
|
| 307 |
+
if self._base_url is None:
|
| 308 |
+
raise RuntimeError("UVProvider has not been started")
|
| 309 |
+
return self._base_url
|
src/core/containers/test_local_docker_provider.py
ADDED
|
@@ -0,0 +1,260 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
#!/usr/bin/env python3
|
| 2 |
+
"""
|
| 3 |
+
End-to-end test for LocalDockerProvider.
|
| 4 |
+
|
| 5 |
+
This script tests the complete flow:
|
| 6 |
+
1. Start a container using LocalDockerProvider
|
| 7 |
+
2. Wait for it to be ready
|
| 8 |
+
3. Make HTTP requests to test the environment
|
| 9 |
+
4. Clean up the container
|
| 10 |
+
"""
|
| 11 |
+
|
| 12 |
+
import sys
|
| 13 |
+
from pathlib import Path
|
| 14 |
+
|
| 15 |
+
# Add src to path
|
| 16 |
+
sys.path.insert(0, str(Path(__file__).parent.parent.parent))
|
| 17 |
+
|
| 18 |
+
import requests
|
| 19 |
+
from openenv.core.containers.runtime import LocalDockerProvider
|
| 20 |
+
|
| 21 |
+
|
| 22 |
+
# TODO: Remove this test or make it a functional test since this will be tested in e2e test for echo env
|
| 23 |
+
def test_local_docker_provider():
|
| 24 |
+
"""Test LocalDockerProvider end-to-end."""
|
| 25 |
+
print("=" * 60)
|
| 26 |
+
print("LocalDockerProvider End-to-End Test")
|
| 27 |
+
print("=" * 60)
|
| 28 |
+
print()
|
| 29 |
+
|
| 30 |
+
provider = None
|
| 31 |
+
|
| 32 |
+
try:
|
| 33 |
+
# Step 1: Create provider
|
| 34 |
+
print("Step 1: Creating LocalDockerProvider...")
|
| 35 |
+
provider = LocalDockerProvider()
|
| 36 |
+
print("✓ Provider created\n")
|
| 37 |
+
|
| 38 |
+
# Step 2: Start container
|
| 39 |
+
print("Step 2: Starting echo-env container...")
|
| 40 |
+
base_url = provider.start_container("echo-env:latest")
|
| 41 |
+
print(f"✓ Container started at: {base_url}")
|
| 42 |
+
if provider._container_id:
|
| 43 |
+
print(f" Container ID: {provider._container_id[:12]}...")
|
| 44 |
+
if provider._container_name:
|
| 45 |
+
print(f" Container name: {provider._container_name}\n")
|
| 46 |
+
|
| 47 |
+
# Step 3: Wait for ready
|
| 48 |
+
print("Step 3: Waiting for container to be ready...")
|
| 49 |
+
provider.wait_for_ready(base_url, timeout_s=30.0)
|
| 50 |
+
print("✓ Container is ready!\n")
|
| 51 |
+
|
| 52 |
+
# Step 4: Test health endpoint
|
| 53 |
+
print("Step 4: Testing /health endpoint...")
|
| 54 |
+
response = requests.get(f"{base_url}/health")
|
| 55 |
+
print(f" Status: {response.status_code}")
|
| 56 |
+
print(f" Response: {response.json()}")
|
| 57 |
+
assert response.status_code == 200
|
| 58 |
+
assert response.json()["status"] == "healthy"
|
| 59 |
+
print("✓ Health check passed\n")
|
| 60 |
+
|
| 61 |
+
# Step 5: Test reset endpoint
|
| 62 |
+
print("Step 5: Testing /reset endpoint...")
|
| 63 |
+
response = requests.post(
|
| 64 |
+
f"{base_url}/reset",
|
| 65 |
+
json={},
|
| 66 |
+
headers={"Content-Type": "application/json"},
|
| 67 |
+
)
|
| 68 |
+
print(f" Status: {response.status_code}")
|
| 69 |
+
data = response.json()
|
| 70 |
+
print(f" Message: {data['observation']['echoed_message']}")
|
| 71 |
+
print(f" Reward: {data['reward']}")
|
| 72 |
+
print(f" Done: {data['done']}")
|
| 73 |
+
assert response.status_code == 200
|
| 74 |
+
assert data["observation"]["echoed_message"] == "Echo environment ready!"
|
| 75 |
+
print("✓ Reset test passed\n")
|
| 76 |
+
|
| 77 |
+
# Step 6: Test step endpoint
|
| 78 |
+
print("Step 6: Testing /step endpoint...")
|
| 79 |
+
response = requests.post(
|
| 80 |
+
f"{base_url}/step",
|
| 81 |
+
json={"action": {"message": "Hello from LocalDockerProvider!"}},
|
| 82 |
+
headers={"Content-Type": "application/json"},
|
| 83 |
+
)
|
| 84 |
+
print(f" Status: {response.status_code}")
|
| 85 |
+
data = response.json()
|
| 86 |
+
print(f" Echoed: {data['observation']['echoed_message']}")
|
| 87 |
+
print(f" Length: {data['observation']['message_length']}")
|
| 88 |
+
print(f" Reward: {data['reward']}")
|
| 89 |
+
assert response.status_code == 200
|
| 90 |
+
assert (
|
| 91 |
+
data["observation"]["echoed_message"] == "Hello from LocalDockerProvider!"
|
| 92 |
+
)
|
| 93 |
+
assert data["observation"]["message_length"] == 31
|
| 94 |
+
print("✓ Step test passed\n")
|
| 95 |
+
|
| 96 |
+
# Step 7: Test state endpoint
|
| 97 |
+
print("Step 7: Testing /state endpoint...")
|
| 98 |
+
response = requests.get(f"{base_url}/state")
|
| 99 |
+
print(f" Status: {response.status_code}")
|
| 100 |
+
data = response.json()
|
| 101 |
+
print(f" Episode ID: {data['episode_id']}")
|
| 102 |
+
print(f" Step count: {data['step_count']}")
|
| 103 |
+
assert response.status_code == 200
|
| 104 |
+
assert data["step_count"] == 1 # One step from above
|
| 105 |
+
print("✓ State test passed\n")
|
| 106 |
+
|
| 107 |
+
# Step 8: Multiple steps
|
| 108 |
+
print("Step 8: Testing multiple steps...")
|
| 109 |
+
for i in range(3):
|
| 110 |
+
response = requests.post(
|
| 111 |
+
f"{base_url}/step",
|
| 112 |
+
json={"action": {"message": f"Message {i + 1}"}},
|
| 113 |
+
headers={"Content-Type": "application/json"},
|
| 114 |
+
)
|
| 115 |
+
assert response.status_code == 200
|
| 116 |
+
print(f" Step {i + 1}: ✓")
|
| 117 |
+
|
| 118 |
+
# Check state updated
|
| 119 |
+
response = requests.get(f"{base_url}/state")
|
| 120 |
+
data = response.json()
|
| 121 |
+
assert data["step_count"] == 4 # 1 + 3 more steps
|
| 122 |
+
print(f" Final step count: {data['step_count']}")
|
| 123 |
+
print("✓ Multiple steps test passed\n")
|
| 124 |
+
|
| 125 |
+
print("=" * 60)
|
| 126 |
+
print("✓ All tests passed!")
|
| 127 |
+
print("=" * 60)
|
| 128 |
+
print()
|
| 129 |
+
|
| 130 |
+
return True
|
| 131 |
+
|
| 132 |
+
except Exception as e:
|
| 133 |
+
print(f"\n❌ Test failed: {e}")
|
| 134 |
+
import traceback
|
| 135 |
+
|
| 136 |
+
traceback.print_exc()
|
| 137 |
+
return False
|
| 138 |
+
|
| 139 |
+
finally:
|
| 140 |
+
# Step 9: Cleanup
|
| 141 |
+
if provider is not None:
|
| 142 |
+
print("\nStep 9: Cleaning up container...")
|
| 143 |
+
try:
|
| 144 |
+
provider.stop_container()
|
| 145 |
+
print("✓ Container stopped and removed\n")
|
| 146 |
+
except Exception as e:
|
| 147 |
+
print(f"⚠️ Cleanup warning: {e}\n")
|
| 148 |
+
|
| 149 |
+
|
| 150 |
+
def test_provider_with_custom_port():
|
| 151 |
+
"""Test provider with custom port."""
|
| 152 |
+
print("=" * 60)
|
| 153 |
+
print("LocalDockerProvider with Custom Port Test")
|
| 154 |
+
print("=" * 60)
|
| 155 |
+
print()
|
| 156 |
+
|
| 157 |
+
provider = None
|
| 158 |
+
|
| 159 |
+
try:
|
| 160 |
+
provider = LocalDockerProvider()
|
| 161 |
+
|
| 162 |
+
print("Starting container on custom port 8123...")
|
| 163 |
+
base_url = provider.start_container("echo-env:latest", port=8123)
|
| 164 |
+
print(f"✓ Started at: {base_url}")
|
| 165 |
+
assert ":8123" in base_url
|
| 166 |
+
|
| 167 |
+
print("Waiting for ready...")
|
| 168 |
+
provider.wait_for_ready(base_url)
|
| 169 |
+
print("✓ Ready!")
|
| 170 |
+
|
| 171 |
+
print("Testing health...")
|
| 172 |
+
response = requests.get(f"{base_url}/health")
|
| 173 |
+
assert response.status_code == 200
|
| 174 |
+
print("✓ Health check passed")
|
| 175 |
+
|
| 176 |
+
print("\n✓ Custom port test passed!\n")
|
| 177 |
+
return True
|
| 178 |
+
|
| 179 |
+
except Exception as e:
|
| 180 |
+
print(f"\n❌ Test failed: {e}")
|
| 181 |
+
return False
|
| 182 |
+
|
| 183 |
+
finally:
|
| 184 |
+
if provider is not None:
|
| 185 |
+
provider.stop_container()
|
| 186 |
+
print("✓ Cleaned up\n")
|
| 187 |
+
|
| 188 |
+
|
| 189 |
+
def test_provider_with_env_vars():
|
| 190 |
+
"""Test provider with environment variables."""
|
| 191 |
+
print("=" * 60)
|
| 192 |
+
print("LocalDockerProvider with Environment Variables Test")
|
| 193 |
+
print("=" * 60)
|
| 194 |
+
print()
|
| 195 |
+
|
| 196 |
+
provider = None
|
| 197 |
+
|
| 198 |
+
try:
|
| 199 |
+
provider = LocalDockerProvider()
|
| 200 |
+
|
| 201 |
+
print("Starting container with environment variables...")
|
| 202 |
+
base_url = provider.start_container(
|
| 203 |
+
"echo-env:latest", env_vars={"DEBUG": "true", "LOG_LEVEL": "info"}
|
| 204 |
+
)
|
| 205 |
+
print(f"✓ Started at: {base_url}")
|
| 206 |
+
|
| 207 |
+
print("Waiting for ready...")
|
| 208 |
+
provider.wait_for_ready(base_url)
|
| 209 |
+
print("✓ Ready!")
|
| 210 |
+
|
| 211 |
+
print("Testing health...")
|
| 212 |
+
response = requests.get(f"{base_url}/health")
|
| 213 |
+
assert response.status_code == 200
|
| 214 |
+
print("✓ Health check passed")
|
| 215 |
+
|
| 216 |
+
print("\n✓ Environment variables test passed!\n")
|
| 217 |
+
return True
|
| 218 |
+
|
| 219 |
+
except Exception as e:
|
| 220 |
+
print(f"\n❌ Test failed: {e}")
|
| 221 |
+
return False
|
| 222 |
+
|
| 223 |
+
finally:
|
| 224 |
+
if provider is not None:
|
| 225 |
+
provider.stop_container()
|
| 226 |
+
print("✓ Cleaned up\n")
|
| 227 |
+
|
| 228 |
+
|
| 229 |
+
if __name__ == "__main__":
|
| 230 |
+
print()
|
| 231 |
+
print("🐳 LocalDockerProvider Test Suite")
|
| 232 |
+
print()
|
| 233 |
+
|
| 234 |
+
results = []
|
| 235 |
+
|
| 236 |
+
# Run basic test
|
| 237 |
+
results.append(("Basic End-to-End", test_local_docker_provider()))
|
| 238 |
+
|
| 239 |
+
# Run custom port test
|
| 240 |
+
results.append(("Custom Port", test_provider_with_custom_port()))
|
| 241 |
+
|
| 242 |
+
# Run environment variables test
|
| 243 |
+
results.append(("Environment Variables", test_provider_with_env_vars()))
|
| 244 |
+
|
| 245 |
+
# Summary
|
| 246 |
+
print("=" * 60)
|
| 247 |
+
print("Test Summary")
|
| 248 |
+
print("=" * 60)
|
| 249 |
+
for name, passed in results:
|
| 250 |
+
status = "✓ PASSED" if passed else "✗ FAILED"
|
| 251 |
+
print(f"{name:25} {status}")
|
| 252 |
+
print("=" * 60)
|
| 253 |
+
|
| 254 |
+
all_passed = all(result for _, result in results)
|
| 255 |
+
if all_passed:
|
| 256 |
+
print("\n🎉 All tests passed!")
|
| 257 |
+
exit(0)
|
| 258 |
+
else:
|
| 259 |
+
print("\n❌ Some tests failed")
|
| 260 |
+
exit(1)
|
src/core/env_client.py
ADDED
|
@@ -0,0 +1,1197 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# SPDX-License-Identifier: BSD-3-Clause
|
| 2 |
+
|
| 3 |
+
"""
|
| 4 |
+
Environment client for persistent sessions.
|
| 5 |
+
|
| 6 |
+
This module provides a WebSocket-based client that maintains a persistent connection
|
| 7 |
+
to an environment server, enabling efficient multi-step interactions without
|
| 8 |
+
the overhead of HTTP request/response cycles.
|
| 9 |
+
|
| 10 |
+
The client is async by default. For synchronous usage, use the `.sync()` method
|
| 11 |
+
to get a `SyncEnvClient` wrapper.
|
| 12 |
+
|
| 13 |
+
Examples:
|
| 14 |
+
|
| 15 |
+
Async usage:
|
| 16 |
+
|
| 17 |
+
```python
|
| 18 |
+
async with GenericEnvClient(base_url="ws://localhost:8000") as env:
|
| 19 |
+
result = await env.reset()
|
| 20 |
+
result = await env.step({"code": "print('hello')"})
|
| 21 |
+
```
|
| 22 |
+
|
| 23 |
+
Sync usage via `.sync()` wrapper:
|
| 24 |
+
|
| 25 |
+
```python
|
| 26 |
+
env = GenericEnvClient(base_url="ws://localhost:8000").sync()
|
| 27 |
+
with env:
|
| 28 |
+
result = env.reset()
|
| 29 |
+
result = env.step({"code": "print('hello')"})
|
| 30 |
+
```
|
| 31 |
+
"""
|
| 32 |
+
|
| 33 |
+
from __future__ import annotations
|
| 34 |
+
|
| 35 |
+
import asyncio
|
| 36 |
+
import inspect
|
| 37 |
+
import ipaddress
|
| 38 |
+
import json
|
| 39 |
+
import os
|
| 40 |
+
import time
|
| 41 |
+
from abc import ABC, abstractmethod
|
| 42 |
+
from collections.abc import Coroutine
|
| 43 |
+
from contextlib import suppress
|
| 44 |
+
from typing import Any, Callable, Dict, Generic, Optional, Type, TYPE_CHECKING, TypeVar
|
| 45 |
+
from urllib.parse import urlsplit
|
| 46 |
+
|
| 47 |
+
from .client_types import StateT, StepResult
|
| 48 |
+
from .containers.runtime import LocalDockerProvider, UVProvider
|
| 49 |
+
from .utils import convert_to_ws_url, hf_space_subdomain
|
| 50 |
+
|
| 51 |
+
if TYPE_CHECKING:
|
| 52 |
+
from websockets.asyncio.client import ClientConnection
|
| 53 |
+
|
| 54 |
+
from .containers.runtime import ContainerProvider, RuntimeProvider
|
| 55 |
+
from .sync_client import SyncEnvClient
|
| 56 |
+
|
| 57 |
+
from websockets.asyncio.client import connect as ws_connect
|
| 58 |
+
from websockets.exceptions import ConnectionClosed
|
| 59 |
+
from websockets.protocol import State
|
| 60 |
+
|
| 61 |
+
ActT = TypeVar("ActT")
|
| 62 |
+
ObsT = TypeVar("ObsT")
|
| 63 |
+
EnvClientT = TypeVar("EnvClientT", bound="EnvClient")
|
| 64 |
+
ResultT = TypeVar("ResultT")
|
| 65 |
+
|
| 66 |
+
_VALID_CLIENT_MODES = ("simulation", "production")
|
| 67 |
+
|
| 68 |
+
|
| 69 |
+
class _AutoAsyncResult(Generic[ResultT]):
|
| 70 |
+
"""Awaitable result that can also be resolved by synchronous access."""
|
| 71 |
+
|
| 72 |
+
def __init__(
|
| 73 |
+
self,
|
| 74 |
+
client: "EnvClient[Any, Any, Any]",
|
| 75 |
+
coro_factory: Callable[[], Any],
|
| 76 |
+
):
|
| 77 |
+
self._client = client
|
| 78 |
+
self._coro_factory = coro_factory
|
| 79 |
+
self._resolved = False
|
| 80 |
+
self._result: ResultT | None = None
|
| 81 |
+
|
| 82 |
+
def __await__(self):
|
| 83 |
+
async def _await_result():
|
| 84 |
+
self._client._claim_execution_mode("async")
|
| 85 |
+
return await self._coro_factory()
|
| 86 |
+
|
| 87 |
+
return _await_result().__await__()
|
| 88 |
+
|
| 89 |
+
def _sync_result(self) -> ResultT:
|
| 90 |
+
if not self._resolved:
|
| 91 |
+
self._result = self._client._run_sync(self._coro_factory)
|
| 92 |
+
self._resolved = True
|
| 93 |
+
return self._result
|
| 94 |
+
|
| 95 |
+
def __getattr__(self, name: str) -> Any:
|
| 96 |
+
return getattr(self._sync_result(), name)
|
| 97 |
+
|
| 98 |
+
def __bool__(self) -> bool:
|
| 99 |
+
return bool(self._sync_result())
|
| 100 |
+
|
| 101 |
+
def __repr__(self) -> str:
|
| 102 |
+
if self._resolved:
|
| 103 |
+
return repr(self._result)
|
| 104 |
+
return f"<{type(self).__name__} pending>"
|
| 105 |
+
|
| 106 |
+
|
| 107 |
+
class _BootstrapResult(Coroutine, Generic[EnvClientT]):
|
| 108 |
+
"""Bootstrap handle returned by the client factory methods.
|
| 109 |
+
|
| 110 |
+
Returned by [`~openenv.core.EnvClient.from_docker_image`] and
|
| 111 |
+
[`~openenv.core.EnvClient.from_env`]. Resolving the handle is what actually
|
| 112 |
+
starts the container / Space and connects the WebSocket, so a single factory
|
| 113 |
+
call serves both execution modes:
|
| 114 |
+
|
| 115 |
+
- `await handle` connects on the running event loop and returns the connected
|
| 116 |
+
async client (unchanged async behavior).
|
| 117 |
+
- `handle.sync()` connects on the sync background loop and returns a
|
| 118 |
+
`SyncEnvClient`, mirroring the instance-level [`~openenv.core.EnvClient.sync`].
|
| 119 |
+
|
| 120 |
+
The handle is a full coroutine (it implements `send` / `throw` / `close`), so
|
| 121 |
+
it is a drop-in for the previous `async def` factories: `asyncio.run(...)`,
|
| 122 |
+
`run_async_safely(...)`, and bare `await` all accept it, in addition to the
|
| 123 |
+
new `.sync()` chain. Bootstrap is lazy — the underlying coroutine (and the
|
| 124 |
+
container/Space start) is created only when the handle is driven or `.sync()`
|
| 125 |
+
is called.
|
| 126 |
+
"""
|
| 127 |
+
|
| 128 |
+
def __init__(self, bootstrap: Callable[[], EnvClientT]):
|
| 129 |
+
self._bootstrap = bootstrap
|
| 130 |
+
self._coro: Optional[Coroutine[Any, Any, EnvClientT]] = None
|
| 131 |
+
self._used = False
|
| 132 |
+
|
| 133 |
+
def _consume(self) -> EnvClientT:
|
| 134 |
+
"""Run the bootstrap exactly once; a second resolve is a programming error.
|
| 135 |
+
|
| 136 |
+
Mirrors native coroutine semantics (a coroutine cannot be awaited twice)
|
| 137 |
+
so that a handle re-driven via a second `.sync()`, or `await` followed by
|
| 138 |
+
`.sync()`, raises instead of silently starting a second container/Space.
|
| 139 |
+
"""
|
| 140 |
+
if self._used:
|
| 141 |
+
raise RuntimeError(
|
| 142 |
+
"This bootstrap handle has already been resolved; call the "
|
| 143 |
+
"factory again to start a new environment."
|
| 144 |
+
)
|
| 145 |
+
self._used = True
|
| 146 |
+
return self._bootstrap()
|
| 147 |
+
|
| 148 |
+
async def _resolve_async(self) -> EnvClientT:
|
| 149 |
+
client = self._consume()
|
| 150 |
+
try:
|
| 151 |
+
await client.connect()
|
| 152 |
+
except asyncio.CancelledError:
|
| 153 |
+
# _connect_async releases the provider on ordinary failures, but
|
| 154 |
+
# cancellation bypasses its `except Exception`. No client reaches
|
| 155 |
+
# the caller, so release the provider here before propagating the
|
| 156 |
+
# original cancellation. A repeated cancellation or a stop failure
|
| 157 |
+
# during cleanup must not replace it.
|
| 158 |
+
with suppress(Exception, asyncio.CancelledError):
|
| 159 |
+
await client.close()
|
| 160 |
+
raise
|
| 161 |
+
return client
|
| 162 |
+
|
| 163 |
+
def _ensure_coro(self) -> "Coroutine[Any, Any, EnvClientT]":
|
| 164 |
+
if self._coro is None:
|
| 165 |
+
self._coro = self._resolve_async()
|
| 166 |
+
return self._coro
|
| 167 |
+
|
| 168 |
+
def __await__(self):
|
| 169 |
+
return self._ensure_coro().__await__()
|
| 170 |
+
|
| 171 |
+
def send(self, value: Any) -> Any:
|
| 172 |
+
return self._ensure_coro().send(value)
|
| 173 |
+
|
| 174 |
+
def throw(self, *args: Any, **kwargs: Any) -> Any:
|
| 175 |
+
return self._ensure_coro().throw(*args, **kwargs)
|
| 176 |
+
|
| 177 |
+
def close(self) -> None:
|
| 178 |
+
if self._coro is not None:
|
| 179 |
+
self._coro.close()
|
| 180 |
+
|
| 181 |
+
def sync(self) -> "SyncEnvClient":
|
| 182 |
+
client = self._consume()
|
| 183 |
+
try:
|
| 184 |
+
client._run_sync(client._connect_async)
|
| 185 |
+
except Exception:
|
| 186 |
+
# _consume() already started the provider. On the async path
|
| 187 |
+
# _connect_async releases it, but here the failure may be the sync
|
| 188 |
+
# loop setup itself (before _connect_async runs), so stop the
|
| 189 |
+
# provider directly rather than routing through the broken loop.
|
| 190 |
+
client._stop_provider_best_effort()
|
| 191 |
+
raise
|
| 192 |
+
return client.sync()
|
| 193 |
+
|
| 194 |
+
def __repr__(self) -> str:
|
| 195 |
+
return f"<{type(self).__name__} pending (await, run, or call .sync())>"
|
| 196 |
+
|
| 197 |
+
|
| 198 |
+
def _normalize_mode(mode: Optional[str]) -> str:
|
| 199 |
+
"""Resolve and validate the client communication mode."""
|
| 200 |
+
raw_mode = (
|
| 201 |
+
os.environ.get("OPENENV_CLIENT_MODE", "simulation") if mode is None else mode
|
| 202 |
+
)
|
| 203 |
+
normalized_mode = raw_mode.lower()
|
| 204 |
+
if normalized_mode not in _VALID_CLIENT_MODES:
|
| 205 |
+
raise ValueError(
|
| 206 |
+
f"Invalid mode: '{normalized_mode}'. Must be 'simulation' or 'production'. "
|
| 207 |
+
f"Set via constructor parameter or OPENENV_CLIENT_MODE environment variable."
|
| 208 |
+
)
|
| 209 |
+
return normalized_mode
|
| 210 |
+
|
| 211 |
+
|
| 212 |
+
def _is_localhost_ws_url(ws_url: str) -> bool:
|
| 213 |
+
"""Return True when the WebSocket URL targets the local loopback interface.
|
| 214 |
+
|
| 215 |
+
The hostname is parsed from the URL so that only the actual host is matched.
|
| 216 |
+
Substring matching is avoided because remote hosts such as
|
| 217 |
+
``my-localhost-proxy.example.com`` or ``127.0.0.1.example.com`` must not be
|
| 218 |
+
treated as local.
|
| 219 |
+
"""
|
| 220 |
+
hostname = urlsplit(ws_url).hostname
|
| 221 |
+
if hostname is None:
|
| 222 |
+
return False
|
| 223 |
+
if hostname == "localhost":
|
| 224 |
+
return True
|
| 225 |
+
try:
|
| 226 |
+
return ipaddress.ip_address(hostname).is_loopback
|
| 227 |
+
except ValueError:
|
| 228 |
+
return False
|
| 229 |
+
|
| 230 |
+
|
| 231 |
+
def _required_start_container_parameters(provider: Any) -> list[str]:
|
| 232 |
+
"""Return required arguments for a bound provider.start_container()."""
|
| 233 |
+
try:
|
| 234 |
+
signature = inspect.signature(provider.start_container)
|
| 235 |
+
except (TypeError, ValueError):
|
| 236 |
+
return []
|
| 237 |
+
return [
|
| 238 |
+
name
|
| 239 |
+
for name, parameter in signature.parameters.items()
|
| 240 |
+
if parameter.default is inspect.Parameter.empty
|
| 241 |
+
and parameter.kind
|
| 242 |
+
in (
|
| 243 |
+
inspect.Parameter.POSITIONAL_ONLY,
|
| 244 |
+
inspect.Parameter.POSITIONAL_OR_KEYWORD,
|
| 245 |
+
inspect.Parameter.KEYWORD_ONLY,
|
| 246 |
+
)
|
| 247 |
+
]
|
| 248 |
+
|
| 249 |
+
|
| 250 |
+
async def _best_effort_close(ws: ClientConnection) -> None:
|
| 251 |
+
"""Close a socket without letting a slow handshake or failure propagate.
|
| 252 |
+
|
| 253 |
+
Scheduled as a background task (never awaited directly) so a dropped
|
| 254 |
+
socket's close handshake can't hold up the caller that triggered the
|
| 255 |
+
drop -- see `EnvClient._receive()`. `EnvClient._close_async()` awaits
|
| 256 |
+
any still-running instance of this task before it returns, so a real
|
| 257 |
+
`close()` call does wait for the handshake; only the caller that
|
| 258 |
+
happened to trigger the drop is spared. `CancelledError` is caught
|
| 259 |
+
alongside `Exception` because that awaiting is exactly what would
|
| 260 |
+
otherwise cancel this mid-handshake if a caller called `close()`
|
| 261 |
+
concurrently.
|
| 262 |
+
"""
|
| 263 |
+
try:
|
| 264 |
+
await ws.close()
|
| 265 |
+
except (Exception, asyncio.CancelledError):
|
| 266 |
+
pass # Best effort
|
| 267 |
+
|
| 268 |
+
|
| 269 |
+
async def _best_effort_graceful_close(ws: ClientConnection) -> None:
|
| 270 |
+
"""Notify the server, then close a socket without propagating failures."""
|
| 271 |
+
try:
|
| 272 |
+
await ws.send(json.dumps({"type": "close"}))
|
| 273 |
+
except (Exception, asyncio.CancelledError):
|
| 274 |
+
pass # Best effort
|
| 275 |
+
await _best_effort_close(ws)
|
| 276 |
+
|
| 277 |
+
|
| 278 |
+
class EnvClient(ABC, Generic[ActT, ObsT, StateT]):
|
| 279 |
+
"""
|
| 280 |
+
Async environment client for persistent sessions.
|
| 281 |
+
|
| 282 |
+
This client maintains a persistent WebSocket connection to an environment
|
| 283 |
+
server, enabling efficient multi-step interactions. Each client instance
|
| 284 |
+
corresponds to a dedicated environment session on the server.
|
| 285 |
+
|
| 286 |
+
The client is async by default. For synchronous usage, use the `.sync()`
|
| 287 |
+
method to get a `SyncEnvClient` wrapper.
|
| 288 |
+
|
| 289 |
+
Features:
|
| 290 |
+
- Lower latency for sequential interactions
|
| 291 |
+
- Session state is maintained server-side
|
| 292 |
+
- Better suited for long-running episodes
|
| 293 |
+
- Async by default for modern Python async/await patterns
|
| 294 |
+
|
| 295 |
+
Examples:
|
| 296 |
+
|
| 297 |
+
Async usage:
|
| 298 |
+
|
| 299 |
+
```python
|
| 300 |
+
from envs.coding_env.client import CodingEnv
|
| 301 |
+
|
| 302 |
+
async with CodingEnv(base_url="ws://localhost:8000") as env:
|
| 303 |
+
result = await env.reset(seed=42)
|
| 304 |
+
while not result.done:
|
| 305 |
+
action = agent.predict(result.observation)
|
| 306 |
+
result = await env.step(action)
|
| 307 |
+
```
|
| 308 |
+
|
| 309 |
+
Sync usage via `.sync()` wrapper:
|
| 310 |
+
|
| 311 |
+
```python
|
| 312 |
+
env = CodingEnv(base_url="ws://localhost:8000").sync()
|
| 313 |
+
with env:
|
| 314 |
+
result = env.reset(seed=42)
|
| 315 |
+
result = env.step(action)
|
| 316 |
+
```
|
| 317 |
+
"""
|
| 318 |
+
|
| 319 |
+
def __init__(
|
| 320 |
+
self,
|
| 321 |
+
base_url: Optional[str] = None,
|
| 322 |
+
connect_timeout_s: float = 10.0,
|
| 323 |
+
message_timeout_s: float = 60.0,
|
| 324 |
+
max_message_size_mb: float = 100.0,
|
| 325 |
+
websocket_ping_interval_s: Optional[float] = 20.0,
|
| 326 |
+
websocket_ping_timeout_s: Optional[float] = 20.0,
|
| 327 |
+
provider: Optional["ContainerProvider | RuntimeProvider"] = None,
|
| 328 |
+
mode: Optional[str] = None,
|
| 329 |
+
headers: Optional[Dict[str, str]] = None,
|
| 330 |
+
):
|
| 331 |
+
"""
|
| 332 |
+
Initialize environment client.
|
| 333 |
+
|
| 334 |
+
Args:
|
| 335 |
+
base_url (`str`, *optional*):
|
| 336 |
+
Base URL of the environment server (http:// or ws://). Will be converted to
|
| 337 |
+
ws:// if http:// is provided. May be omitted when the provider
|
| 338 |
+
has enough constructor state to start itself.
|
| 339 |
+
connect_timeout_s (`float`, *optional*, defaults to `10.0`):
|
| 340 |
+
Timeout for establishing WebSocket connection.
|
| 341 |
+
message_timeout_s (`float`, *optional*, defaults to `60.0`):
|
| 342 |
+
Timeout for receiving responses to messages.
|
| 343 |
+
max_message_size_mb (`float`, *optional*, defaults to `100.0`):
|
| 344 |
+
Maximum WebSocket message size in megabytes. Default 100MB to handle large
|
| 345 |
+
observations (screenshots, DOM, etc.).
|
| 346 |
+
websocket_ping_interval_s (`float` or `None`, *optional*, defaults to `20.0`):
|
| 347 |
+
WebSocket keepalive ping interval. Pass `None` to disable.
|
| 348 |
+
websocket_ping_timeout_s (`float` or `None`, *optional*, defaults to `20.0`):
|
| 349 |
+
WebSocket keepalive pong timeout. Pass `None` to disable.
|
| 350 |
+
provider (`ContainerProvider` or `RuntimeProvider`, *optional*):
|
| 351 |
+
Container/runtime provider for lifecycle management.
|
| 352 |
+
mode (`str`, *optional*):
|
| 353 |
+
Communication mode: `'simulation'` for Gym-style API (default) or
|
| 354 |
+
`'production'` for MCP JSON-RPC protocol. Can also be set via the
|
| 355 |
+
`OPENENV_CLIENT_MODE` environment variable. Constructor parameter takes
|
| 356 |
+
precedence over environment variable. Case-insensitive.
|
| 357 |
+
headers (`dict[str, str]`, *optional*):
|
| 358 |
+
Extra HTTP headers sent when opening the connection, e.g.
|
| 359 |
+
`{"Authorization": f"Bearer {hf_token}"}` for a private Hugging Face Space.
|
| 360 |
+
"""
|
| 361 |
+
if base_url is None and provider is None:
|
| 362 |
+
raise ValueError("EnvClient requires either base_url or provider.")
|
| 363 |
+
|
| 364 |
+
# Store mode (use object.__setattr__ to bypass immutability)
|
| 365 |
+
object.__setattr__(self, "_mode", _normalize_mode(mode))
|
| 366 |
+
|
| 367 |
+
self._base_url: Optional[str] = None
|
| 368 |
+
self._ws_url: Optional[str] = None
|
| 369 |
+
self._connect_timeout = connect_timeout_s
|
| 370 |
+
self._message_timeout = message_timeout_s
|
| 371 |
+
self._max_message_size = int(
|
| 372 |
+
max_message_size_mb * 1024 * 1024
|
| 373 |
+
) # Convert MB to bytes
|
| 374 |
+
self._websocket_ping_interval_s = websocket_ping_interval_s
|
| 375 |
+
self._websocket_ping_timeout_s = websocket_ping_timeout_s
|
| 376 |
+
self._headers = headers
|
| 377 |
+
self._provider = provider
|
| 378 |
+
self._provider_stopped = False
|
| 379 |
+
self._provider_cleanup_pending = False
|
| 380 |
+
self._start_provider_on_connect = base_url is None
|
| 381 |
+
self._child_clients: list[EnvClient[Any, Any, Any]] = []
|
| 382 |
+
self._ws: Optional[ClientConnection] = None
|
| 383 |
+
self._execution_mode: Optional[str] = None
|
| 384 |
+
self._sync_client: Optional["SyncEnvClient[ActT, ObsT, StateT]"] = None
|
| 385 |
+
self._ws_loop: Optional[asyncio.AbstractEventLoop] = None
|
| 386 |
+
# Strong references for fire-and-forget socket closes (see _receive),
|
| 387 |
+
# so the task isn't garbage-collected mid-close. Discarded on done.
|
| 388 |
+
self._pending_close_tasks: set[asyncio.Task] = set()
|
| 389 |
+
if base_url is not None:
|
| 390 |
+
self._set_base_url(base_url)
|
| 391 |
+
|
| 392 |
+
@property
|
| 393 |
+
def base_url(self) -> Optional[str]:
|
| 394 |
+
"""Public read-only URL of the environment server this client targets.
|
| 395 |
+
|
| 396 |
+
Returns the normalized `http(s)`/`ws(s)` base URL (no trailing slash),
|
| 397 |
+
or `None` when the client was created with a provider but has not yet
|
| 398 |
+
started its container (the URL is assigned lazily on `connect()`).
|
| 399 |
+
|
| 400 |
+
This is the public counterpart to the private `self._base_url`; the
|
| 401 |
+
previously-public `base_url` attribute was dropped in the `core`
|
| 402 |
+
refactor, leaving sync consumers with no way to read the URL back.
|
| 403 |
+
"""
|
| 404 |
+
return self._base_url
|
| 405 |
+
|
| 406 |
+
def _set_base_url(self, base_url: str) -> None:
|
| 407 |
+
self._base_url = base_url.rstrip("/")
|
| 408 |
+
ws_url = convert_to_ws_url(base_url)
|
| 409 |
+
self._ws_url = f"{ws_url}/ws"
|
| 410 |
+
|
| 411 |
+
def _start_provider_if_needed(self) -> None:
|
| 412 |
+
# A missing URL does not mean the previous resource was released.
|
| 413 |
+
# Retry its cleanup before start can overwrite the provider's handle.
|
| 414 |
+
if self._provider_cleanup_pending:
|
| 415 |
+
if not self._start_provider_on_connect:
|
| 416 |
+
raise RuntimeError(
|
| 417 |
+
"Provider cleanup is pending for this client with an existing "
|
| 418 |
+
"base URL. Retry close() to finish cleanup, then create a new "
|
| 419 |
+
"client instead of reconnecting to the old URL."
|
| 420 |
+
)
|
| 421 |
+
self._stop_provider()
|
| 422 |
+
if self._ws_url is not None:
|
| 423 |
+
return
|
| 424 |
+
if self._provider is None:
|
| 425 |
+
raise RuntimeError("EnvClient has no base URL or provider.")
|
| 426 |
+
if hasattr(self._provider, "start_container"):
|
| 427 |
+
required_parameters = _required_start_container_parameters(self._provider)
|
| 428 |
+
if required_parameters:
|
| 429 |
+
required = ", ".join(required_parameters)
|
| 430 |
+
raise ValueError(
|
| 431 |
+
f"{type(self._provider).__name__} does not support "
|
| 432 |
+
"provider-owned startup because start_container() requires "
|
| 433 |
+
f"{required}. Start the provider manually and pass base_url, "
|
| 434 |
+
"or configure a provider with a constructor-owned image/source."
|
| 435 |
+
)
|
| 436 |
+
self._provider_stopped = False
|
| 437 |
+
base_url = self._provider.start_container()
|
| 438 |
+
self._provider.wait_for_ready(base_url)
|
| 439 |
+
elif hasattr(self._provider, "start"):
|
| 440 |
+
self._provider_stopped = False
|
| 441 |
+
base_url = self._provider.start()
|
| 442 |
+
self._provider.wait_for_ready()
|
| 443 |
+
else:
|
| 444 |
+
raise TypeError("provider must define start_container() or start().")
|
| 445 |
+
self._set_base_url(base_url)
|
| 446 |
+
|
| 447 |
+
def _create_session_client(self) -> "EnvClient[Any, Any, Any]":
|
| 448 |
+
# Match _start_provider_if_needed's startup condition. A provider with
|
| 449 |
+
# an existing URL may already be serving the parent or other children.
|
| 450 |
+
starts_provider_here = self._provider is not None and self._ws_url is None
|
| 451 |
+
try:
|
| 452 |
+
self._start_provider_if_needed()
|
| 453 |
+
if self._base_url is None:
|
| 454 |
+
raise RuntimeError("EnvClient has no base URL.")
|
| 455 |
+
|
| 456 |
+
signature = inspect.signature(type(self))
|
| 457 |
+
accepts_kwargs = any(
|
| 458 |
+
parameter.kind == inspect.Parameter.VAR_KEYWORD
|
| 459 |
+
for parameter in signature.parameters.values()
|
| 460 |
+
)
|
| 461 |
+
candidate_kwargs = {
|
| 462 |
+
"base_url": self._base_url,
|
| 463 |
+
"connect_timeout_s": self._connect_timeout,
|
| 464 |
+
"message_timeout_s": self._message_timeout,
|
| 465 |
+
"max_message_size_mb": self._max_message_size / (1024 * 1024),
|
| 466 |
+
"websocket_ping_interval_s": self._websocket_ping_interval_s,
|
| 467 |
+
"websocket_ping_timeout_s": self._websocket_ping_timeout_s,
|
| 468 |
+
"mode": self._mode,
|
| 469 |
+
"headers": self._headers,
|
| 470 |
+
}
|
| 471 |
+
constructor_kwargs = {}
|
| 472 |
+
for name, value in candidate_kwargs.items():
|
| 473 |
+
if accepts_kwargs or name in signature.parameters:
|
| 474 |
+
constructor_kwargs[name] = value
|
| 475 |
+
|
| 476 |
+
return type(self)(**constructor_kwargs)
|
| 477 |
+
except Exception:
|
| 478 |
+
if starts_provider_here and not self._provider_cleanup_pending:
|
| 479 |
+
self._stop_provider_best_effort()
|
| 480 |
+
raise
|
| 481 |
+
|
| 482 |
+
async def new_session(self) -> "EnvClient[Any, Any, Any]":
|
| 483 |
+
"""
|
| 484 |
+
Create and connect a new session against the same environment server.
|
| 485 |
+
|
| 486 |
+
Returns:
|
| 487 |
+
`EnvClient`: A connected child client of the same concrete type.
|
| 488 |
+
|
| 489 |
+
The child session is tracked by this parent and closed when the parent
|
| 490 |
+
is closed. Server-side capacity still applies: when the server is at
|
| 491 |
+
`MAX_CONCURRENT_ENVS`, opening the child WebSocket can fail and is
|
| 492 |
+
surfaced as a connection error.
|
| 493 |
+
"""
|
| 494 |
+
client = self._create_session_client()
|
| 495 |
+
await client.connect()
|
| 496 |
+
self._child_clients.append(client)
|
| 497 |
+
return client
|
| 498 |
+
|
| 499 |
+
def __setattr__(self, name: str, value: Any) -> None:
|
| 500 |
+
"""Prevent modification of _mode after initialization."""
|
| 501 |
+
if name == "_mode" and hasattr(self, "_mode"):
|
| 502 |
+
raise AttributeError("Cannot modify mode after initialization")
|
| 503 |
+
super().__setattr__(name, value)
|
| 504 |
+
|
| 505 |
+
def _claim_execution_mode(self, mode: str) -> None:
|
| 506 |
+
"""Lock the client to sync or async execution on first use."""
|
| 507 |
+
if self._execution_mode is None:
|
| 508 |
+
self._execution_mode = mode
|
| 509 |
+
elif self._execution_mode != mode:
|
| 510 |
+
raise RuntimeError(
|
| 511 |
+
f"EnvClient is already being used in {self._execution_mode} mode. "
|
| 512 |
+
"Create a separate client instance when mixing sync and async code."
|
| 513 |
+
)
|
| 514 |
+
|
| 515 |
+
def _run_sync(
|
| 516 |
+
self,
|
| 517 |
+
coro_factory: Callable[[], Any],
|
| 518 |
+
*,
|
| 519 |
+
allow_async_handoff: bool = False,
|
| 520 |
+
) -> Any:
|
| 521 |
+
"""Run an async operation through the sync wrapper."""
|
| 522 |
+
if allow_async_handoff and self._execution_mode == "async":
|
| 523 |
+
self._execution_mode = "sync"
|
| 524 |
+
else:
|
| 525 |
+
self._claim_execution_mode("sync")
|
| 526 |
+
if self._sync_client is None:
|
| 527 |
+
self._sync_client = self.sync()
|
| 528 |
+
return self._sync_client._run(coro_factory())
|
| 529 |
+
|
| 530 |
+
def _dispatch(self, coro_factory: Callable[[], Any]) -> Any:
|
| 531 |
+
"""Return an awaitable in async code and a concrete result in sync code."""
|
| 532 |
+
try:
|
| 533 |
+
running_loop = asyncio.get_running_loop()
|
| 534 |
+
except RuntimeError:
|
| 535 |
+
running_loop = None
|
| 536 |
+
|
| 537 |
+
if self._execution_mode == "sync":
|
| 538 |
+
sync_loop = (
|
| 539 |
+
self._sync_client._loop if self._sync_client is not None else None
|
| 540 |
+
)
|
| 541 |
+
if running_loop is sync_loop:
|
| 542 |
+
return coro_factory()
|
| 543 |
+
return self._run_sync(coro_factory)
|
| 544 |
+
if running_loop is not None:
|
| 545 |
+
if self._execution_mode == "async":
|
| 546 |
+
return coro_factory()
|
| 547 |
+
return _AutoAsyncResult(self, coro_factory)
|
| 548 |
+
|
| 549 |
+
return self._run_sync(coro_factory, allow_async_handoff=True)
|
| 550 |
+
|
| 551 |
+
def connect(self) -> Any:
|
| 552 |
+
return self._dispatch(self._connect_async)
|
| 553 |
+
|
| 554 |
+
async def _connect_async(self) -> "EnvClient":
|
| 555 |
+
"""
|
| 556 |
+
Establish WebSocket connection to the server.
|
| 557 |
+
|
| 558 |
+
Returns:
|
| 559 |
+
self for method chaining
|
| 560 |
+
|
| 561 |
+
Raises:
|
| 562 |
+
ConnectionError: If connection cannot be established
|
| 563 |
+
"""
|
| 564 |
+
if self._ws is not None:
|
| 565 |
+
if self._ws.state in (State.CLOSING, State.CLOSED):
|
| 566 |
+
# Closed by the far end: a keepalive timeout, a tunnel dropping
|
| 567 |
+
# the socket, the server restarting. The object is still not
|
| 568 |
+
# None, so without this check it stays cached and every later
|
| 569 |
+
# call raises `ConnectionClosed` for the rest of the process's
|
| 570 |
+
# life. Only a demonstrably closed socket is dropped, so one
|
| 571 |
+
# still CONNECTING is left alone.
|
| 572 |
+
self._ws = None
|
| 573 |
+
self._ws_loop = None
|
| 574 |
+
elif self._ws_loop is asyncio.get_running_loop():
|
| 575 |
+
return self
|
| 576 |
+
# Connected from a different event loop than the one running
|
| 577 |
+
# now -- e.g. `client = await Client.from_env(...)` inside
|
| 578 |
+
# `asyncio.run(...)`, then `client.sync()` drives every later
|
| 579 |
+
# call on `SyncEnvClient`'s own dedicated background loop. The
|
| 580 |
+
# websocket object is bound to internals of the original loop,
|
| 581 |
+
# which is typically already closed by the time we get here, so
|
| 582 |
+
# it cannot be reused (or even cleanly closed) from this loop.
|
| 583 |
+
# Drop the stale reference and reconnect fresh below rather than
|
| 584 |
+
# silently no-op-ing onto a dead connection.
|
| 585 |
+
self._ws = None
|
| 586 |
+
self._ws_loop = None
|
| 587 |
+
|
| 588 |
+
# A timed-out request drops its socket immediately but closes it in the
|
| 589 |
+
# background so the timeout itself remains prompt. Wait for that close
|
| 590 |
+
# before opening a replacement: the old server-side session continues
|
| 591 |
+
# occupying a capacity slot until the close handshake finishes, and
|
| 592 |
+
# many environments allow only one session.
|
| 593 |
+
await self._drain_pending_close_tasks()
|
| 594 |
+
|
| 595 |
+
try:
|
| 596 |
+
self._start_provider_if_needed()
|
| 597 |
+
except Exception:
|
| 598 |
+
# A failed cleanup retry must propagate without attempting stop
|
| 599 |
+
# again. For a new startup failure, preserve its original error.
|
| 600 |
+
if not self._provider_cleanup_pending:
|
| 601 |
+
with suppress(Exception):
|
| 602 |
+
await self.close()
|
| 603 |
+
raise
|
| 604 |
+
|
| 605 |
+
assert self._ws_url is not None
|
| 606 |
+
|
| 607 |
+
# Disable the proxy for localhost connections via the per-connection
|
| 608 |
+
# `proxy` argument rather than mutating the process-global NO_PROXY
|
| 609 |
+
# env var: concurrent connect() calls (e.g. asyncio.gather over many
|
| 610 |
+
# env clients) would otherwise race on os.environ and leak state.
|
| 611 |
+
connect_kwargs: Dict[str, Any] = {}
|
| 612 |
+
if _is_localhost_ws_url(self._ws_url):
|
| 613 |
+
connect_kwargs["proxy"] = None
|
| 614 |
+
|
| 615 |
+
try:
|
| 616 |
+
self._ws = await ws_connect(
|
| 617 |
+
self._ws_url,
|
| 618 |
+
open_timeout=self._connect_timeout,
|
| 619 |
+
max_size=self._max_message_size,
|
| 620 |
+
ping_interval=self._websocket_ping_interval_s,
|
| 621 |
+
ping_timeout=self._websocket_ping_timeout_s,
|
| 622 |
+
additional_headers=self._headers,
|
| 623 |
+
**connect_kwargs,
|
| 624 |
+
)
|
| 625 |
+
self._ws_loop = asyncio.get_running_loop()
|
| 626 |
+
except Exception as e:
|
| 627 |
+
await self.close()
|
| 628 |
+
raise ConnectionError(f"Failed to connect to {self._ws_url}: {e}") from e
|
| 629 |
+
|
| 630 |
+
return self
|
| 631 |
+
|
| 632 |
+
def disconnect(self) -> Any:
|
| 633 |
+
return self._dispatch(self._disconnect_async)
|
| 634 |
+
|
| 635 |
+
def _schedule_socket_close(
|
| 636 |
+
self, ws: ClientConnection, *, notify_server: bool = False
|
| 637 |
+
) -> asyncio.Task[None]:
|
| 638 |
+
"""Schedule and track a socket close until its handshake finishes."""
|
| 639 |
+
close_coro = (
|
| 640 |
+
_best_effort_graceful_close(ws) if notify_server else _best_effort_close(ws)
|
| 641 |
+
)
|
| 642 |
+
close_task = asyncio.create_task(close_coro)
|
| 643 |
+
self._pending_close_tasks.add(close_task)
|
| 644 |
+
close_task.add_done_callback(self._pending_close_tasks.discard)
|
| 645 |
+
return close_task
|
| 646 |
+
|
| 647 |
+
async def _disconnect_async(self) -> None:
|
| 648 |
+
"""Close the WebSocket connection."""
|
| 649 |
+
if self._ws is not None:
|
| 650 |
+
ws = self._ws
|
| 651 |
+
ws_loop = self._ws_loop
|
| 652 |
+
# Detach first so cancellation during the close handshake cannot
|
| 653 |
+
# leave a stale socket cached for a later operation.
|
| 654 |
+
self._ws = None
|
| 655 |
+
self._ws_loop = None
|
| 656 |
+
same_loop = ws_loop is asyncio.get_running_loop()
|
| 657 |
+
if same_loop:
|
| 658 |
+
# Keep ownership of the handshake if this caller is cancelled.
|
| 659 |
+
# A later close/reconnect drains the task before tearing down
|
| 660 |
+
# the loop or opening a replacement server session.
|
| 661 |
+
close_task = self._schedule_socket_close(ws, notify_server=True)
|
| 662 |
+
await asyncio.shield(close_task)
|
| 663 |
+
|
| 664 |
+
async def _drain_pending_close_tasks(self) -> None:
|
| 665 |
+
"""Wait for background socket closes owned by the current event loop.
|
| 666 |
+
|
| 667 |
+
Shielding keeps cancellation of the caller from cancelling the close
|
| 668 |
+
tasks themselves. This matters both before reconnecting, when the old
|
| 669 |
+
server session must release its capacity slot, and during explicit
|
| 670 |
+
client shutdown.
|
| 671 |
+
"""
|
| 672 |
+
loop = asyncio.get_running_loop()
|
| 673 |
+
tasks = [
|
| 674 |
+
task
|
| 675 |
+
for task in tuple(self._pending_close_tasks)
|
| 676 |
+
if not task.done() and task.get_loop() is loop
|
| 677 |
+
]
|
| 678 |
+
if tasks:
|
| 679 |
+
await asyncio.shield(asyncio.gather(*tasks, return_exceptions=True))
|
| 680 |
+
|
| 681 |
+
async def _ensure_connected(self) -> None:
|
| 682 |
+
"""Ensure WebSocket connection is established on the current loop.
|
| 683 |
+
|
| 684 |
+
Always delegates to `_connect_async()` rather than pre-checking
|
| 685 |
+
`self._ws is None`: `_connect_async()` itself is the one that knows
|
| 686 |
+
whether an existing `_ws` is reusable (same event loop) or stale (a
|
| 687 |
+
different one, e.g. from a prior `from_env()` call now being driven
|
| 688 |
+
through `.sync()`'s own loop).
|
| 689 |
+
"""
|
| 690 |
+
await self._connect_async()
|
| 691 |
+
|
| 692 |
+
async def _send(self, message: Dict[str, Any]) -> None:
|
| 693 |
+
"""Send a message over the WebSocket."""
|
| 694 |
+
await self._ensure_connected()
|
| 695 |
+
assert self._ws is not None
|
| 696 |
+
await self._ws.send(json.dumps(message))
|
| 697 |
+
|
| 698 |
+
async def _receive(self) -> Dict[str, Any]:
|
| 699 |
+
"""Receive the response on the same WebSocket that carried the request.
|
| 700 |
+
|
| 701 |
+
Reconnecting here would wait on a fresh socket where the request was
|
| 702 |
+
never sent. The next complete operation may reconnect in `_send()`.
|
| 703 |
+
"""
|
| 704 |
+
assert self._ws is not None
|
| 705 |
+
try:
|
| 706 |
+
raw = await asyncio.wait_for(self._ws.recv(), timeout=self._message_timeout)
|
| 707 |
+
except (asyncio.TimeoutError, asyncio.CancelledError):
|
| 708 |
+
# The server may still write the response for this request to the
|
| 709 |
+
# socket after we give up waiting on it. If we left the socket
|
| 710 |
+
# open, the next call's `_receive()` would read that stale frame
|
| 711 |
+
# and silently pair it with an unrelated request -- and every
|
| 712 |
+
# response after that would be shifted by one for the life of
|
| 713 |
+
# the connection. Drop the socket so the next `_send()` opens a
|
| 714 |
+
# fresh one instead of reusing a desynced one.
|
| 715 |
+
#
|
| 716 |
+
# An outer `asyncio.wait_for` around the whole call (e.g. a
|
| 717 |
+
# caller-imposed deadline on step()) cancels this await with
|
| 718 |
+
# CancelledError, not our own message_timeout's TimeoutError --
|
| 719 |
+
# same desync, different exception, so both are caught here.
|
| 720 |
+
ws = self._ws
|
| 721 |
+
self._ws = None
|
| 722 |
+
self._ws_loop = None
|
| 723 |
+
# Fire-and-forget: `close()` waits out the library's close
|
| 724 |
+
# handshake (`close_timeout`, 10s by default) if the server is
|
| 725 |
+
# slow to ack, and *awaiting* that here would hold up this
|
| 726 |
+
# exception -- so an outer `asyncio.wait_for(..., timeout=50ms)`
|
| 727 |
+
# would actually block for up to 10s before its deadline was
|
| 728 |
+
# honored. Scheduling it lets the exception propagate
|
| 729 |
+
# immediately while the close still happens in the background.
|
| 730 |
+
self._schedule_socket_close(ws)
|
| 731 |
+
raise
|
| 732 |
+
return json.loads(raw)
|
| 733 |
+
|
| 734 |
+
async def _send_and_receive(self, message: Dict[str, Any]) -> Dict[str, Any]:
|
| 735 |
+
"""Send a message and wait for response."""
|
| 736 |
+
try:
|
| 737 |
+
await self._send(message)
|
| 738 |
+
except ConnectionClosed:
|
| 739 |
+
# A refused session (e.g. at capacity) sends an error frame then closes. If
|
| 740 |
+
# the close beats our send, read the queued frame (recv() re-raises if none).
|
| 741 |
+
pass
|
| 742 |
+
response = await self._receive()
|
| 743 |
+
|
| 744 |
+
# Check for error response
|
| 745 |
+
if response.get("type") == "error":
|
| 746 |
+
error_data = response.get("data", {})
|
| 747 |
+
raise RuntimeError(
|
| 748 |
+
f"Server error: {error_data.get('message', 'Unknown error')} "
|
| 749 |
+
f"(code: {error_data.get('code', 'UNKNOWN')})"
|
| 750 |
+
)
|
| 751 |
+
|
| 752 |
+
return response
|
| 753 |
+
|
| 754 |
+
@classmethod
|
| 755 |
+
def _bootstrap_container(
|
| 756 |
+
cls: Type[EnvClientT],
|
| 757 |
+
image: str,
|
| 758 |
+
provider: Optional["ContainerProvider"] = None,
|
| 759 |
+
**kwargs: Any,
|
| 760 |
+
) -> EnvClientT:
|
| 761 |
+
"""Start a Docker container and build an *unconnected* client for it.
|
| 762 |
+
|
| 763 |
+
Invoked lazily by the `from_docker_image` bootstrap handle when it is
|
| 764 |
+
awaited or `.sync()`'d: the container start / readiness wait is plain
|
| 765 |
+
blocking code, so the only thing that differs between the async and sync
|
| 766 |
+
resolution paths is how the WebSocket is connected (awaited vs. run sync).
|
| 767 |
+
"""
|
| 768 |
+
if provider is None:
|
| 769 |
+
provider = LocalDockerProvider()
|
| 770 |
+
|
| 771 |
+
try:
|
| 772 |
+
# Start container, wait for readiness, then build the client.
|
| 773 |
+
base_url = provider.start_container(image, **kwargs)
|
| 774 |
+
provider.wait_for_ready(base_url)
|
| 775 |
+
return cls(base_url=base_url, provider=provider)
|
| 776 |
+
except Exception:
|
| 777 |
+
# No EnvClient exists yet for the caller to close(), so release the
|
| 778 |
+
# container here if start / readiness / construction fails.
|
| 779 |
+
provider.stop_container()
|
| 780 |
+
raise
|
| 781 |
+
|
| 782 |
+
@classmethod
|
| 783 |
+
def from_docker_image(
|
| 784 |
+
cls: Type[EnvClientT],
|
| 785 |
+
image: str,
|
| 786 |
+
provider: Optional["ContainerProvider"] = None,
|
| 787 |
+
**kwargs: Any,
|
| 788 |
+
) -> "_BootstrapResult[EnvClientT]":
|
| 789 |
+
"""
|
| 790 |
+
Create an environment client by spinning up a Docker container.
|
| 791 |
+
|
| 792 |
+
Returns a bootstrap handle so the same call works from both async and
|
| 793 |
+
synchronous code: `await` it for the connected async client, or chain
|
| 794 |
+
`.sync()` for a connected `SyncEnvClient` (e.g. a TRL GRPO rollout loop
|
| 795 |
+
that cannot `await`). Bootstrap is lazy — the container starts when the
|
| 796 |
+
handle is resolved.
|
| 797 |
+
|
| 798 |
+
Args:
|
| 799 |
+
image (`str`):
|
| 800 |
+
Docker image name to run (e.g., `"coding-env:latest"`).
|
| 801 |
+
provider (`ContainerProvider`, *optional*):
|
| 802 |
+
Container provider to use. Defaults to `LocalDockerProvider`.
|
| 803 |
+
**kwargs:
|
| 804 |
+
Additional arguments to pass to `provider.start_container()`.
|
| 805 |
+
|
| 806 |
+
Returns:
|
| 807 |
+
`_BootstrapResult`: `await` for a connected async client, or call
|
| 808 |
+
`.sync()` for a connected `SyncEnvClient`.
|
| 809 |
+
|
| 810 |
+
Examples:
|
| 811 |
+
|
| 812 |
+
```python
|
| 813 |
+
# Async
|
| 814 |
+
env = await MyEnv.from_docker_image("coding-env:latest")
|
| 815 |
+
|
| 816 |
+
# Sync
|
| 817 |
+
env = MyEnv.from_docker_image("coding-env:latest").sync()
|
| 818 |
+
result = env.reset()
|
| 819 |
+
```
|
| 820 |
+
"""
|
| 821 |
+
return _BootstrapResult(
|
| 822 |
+
lambda: cls._bootstrap_container(image, provider, **kwargs)
|
| 823 |
+
)
|
| 824 |
+
|
| 825 |
+
@classmethod
|
| 826 |
+
def _bootstrap_env(
|
| 827 |
+
cls: Type[EnvClientT],
|
| 828 |
+
repo_id: str,
|
| 829 |
+
*,
|
| 830 |
+
use_docker: bool = True,
|
| 831 |
+
provider: Optional["ContainerProvider | RuntimeProvider"] = None,
|
| 832 |
+
**provider_kwargs: Any,
|
| 833 |
+
) -> EnvClientT:
|
| 834 |
+
"""Start a HF Space (docker or uv) and build an *unconnected* client.
|
| 835 |
+
|
| 836 |
+
Invoked lazily by the `from_env` bootstrap handle; see
|
| 837 |
+
`_bootstrap_container` for the same split rationale. On a startup
|
| 838 |
+
failure the spawned process (and, for a git+ `project_path`, the temp
|
| 839 |
+
clone directory) is released before re-raising; a later *connection*
|
| 840 |
+
failure is cleaned up by `_connect_async` via `close()`.
|
| 841 |
+
"""
|
| 842 |
+
# Extract start args that apply to both providers
|
| 843 |
+
start_args = {}
|
| 844 |
+
for key in ("port", "env_vars", "workers"):
|
| 845 |
+
if key in provider_kwargs:
|
| 846 |
+
start_args[key] = provider_kwargs.pop(key)
|
| 847 |
+
|
| 848 |
+
if use_docker:
|
| 849 |
+
# Docker mode: pull from HF registry
|
| 850 |
+
docker_provider = provider or LocalDockerProvider()
|
| 851 |
+
tag = provider_kwargs.pop("tag", "latest")
|
| 852 |
+
image = f"registry.hf.space/{hf_space_subdomain(repo_id)}:{tag}"
|
| 853 |
+
try:
|
| 854 |
+
base_url = docker_provider.start_container(
|
| 855 |
+
image, **start_args, **provider_kwargs
|
| 856 |
+
)
|
| 857 |
+
docker_provider.wait_for_ready(base_url)
|
| 858 |
+
return cls(base_url=base_url, provider=docker_provider)
|
| 859 |
+
except Exception:
|
| 860 |
+
# No EnvClient exists yet for the caller to close(), so release
|
| 861 |
+
# the container here if start / readiness / construction fails.
|
| 862 |
+
docker_provider.stop_container()
|
| 863 |
+
raise
|
| 864 |
+
else:
|
| 865 |
+
# UV mode: clone and run with uv
|
| 866 |
+
if provider is None:
|
| 867 |
+
uv_kwargs = dict(provider_kwargs)
|
| 868 |
+
project_path = uv_kwargs.pop("project_path", None)
|
| 869 |
+
if project_path is None:
|
| 870 |
+
project_path = f"git+https://huggingface.co/spaces/{repo_id}"
|
| 871 |
+
|
| 872 |
+
provider = UVProvider(project_path=project_path, **uv_kwargs)
|
| 873 |
+
else:
|
| 874 |
+
if provider_kwargs:
|
| 875 |
+
raise ValueError(
|
| 876 |
+
"provider_kwargs cannot be used when supplying a provider instance"
|
| 877 |
+
)
|
| 878 |
+
|
| 879 |
+
try:
|
| 880 |
+
context_timeout_s = getattr(provider, "context_timeout_s", None)
|
| 881 |
+
deadline = (
|
| 882 |
+
time.monotonic() + context_timeout_s
|
| 883 |
+
if context_timeout_s is not None
|
| 884 |
+
else None
|
| 885 |
+
)
|
| 886 |
+
base_url = provider.start(**start_args)
|
| 887 |
+
if deadline is None:
|
| 888 |
+
provider.wait_for_ready()
|
| 889 |
+
else:
|
| 890 |
+
provider.wait_for_ready(
|
| 891 |
+
timeout_s=max(0.0, deadline - time.monotonic())
|
| 892 |
+
)
|
| 893 |
+
return cls(base_url=base_url, provider=provider)
|
| 894 |
+
except Exception:
|
| 895 |
+
# No EnvClient exists yet for the caller to close(), so this is
|
| 896 |
+
# the only chance to release the spawned process and (for a
|
| 897 |
+
# git+ project_path) the temp clone directory. Covers start,
|
| 898 |
+
# readiness, and client construction (e.g. an invalid mode).
|
| 899 |
+
provider.stop()
|
| 900 |
+
raise
|
| 901 |
+
|
| 902 |
+
@classmethod
|
| 903 |
+
def from_env(
|
| 904 |
+
cls: Type[EnvClientT],
|
| 905 |
+
repo_id: str,
|
| 906 |
+
*,
|
| 907 |
+
use_docker: bool = True,
|
| 908 |
+
provider: Optional["ContainerProvider | RuntimeProvider"] = None,
|
| 909 |
+
**provider_kwargs: Any,
|
| 910 |
+
) -> "_BootstrapResult[EnvClientT]":
|
| 911 |
+
"""
|
| 912 |
+
Create a client from a Hugging Face Space.
|
| 913 |
+
|
| 914 |
+
Returns a bootstrap handle: `await` it for the connected async client, or
|
| 915 |
+
chain `.sync()` for a connected `SyncEnvClient`. Bootstrap is lazy — the
|
| 916 |
+
Space starts when the handle is resolved.
|
| 917 |
+
|
| 918 |
+
Args:
|
| 919 |
+
repo_id (`str`):
|
| 920 |
+
Hugging Face space identifier `{org}/{space}`.
|
| 921 |
+
use_docker (`bool`, *optional*, defaults to `True`):
|
| 922 |
+
When `True`, pull from the HF registry and launch via `LocalDockerProvider`.
|
| 923 |
+
When `False`, run the space locally with `UVProvider`.
|
| 924 |
+
provider (`ContainerProvider` or `RuntimeProvider`, *optional*):
|
| 925 |
+
Provider instance to reuse. Must be a `ContainerProvider` when
|
| 926 |
+
`use_docker=True` and a `RuntimeProvider` otherwise.
|
| 927 |
+
**provider_kwargs:
|
| 928 |
+
Additional keyword arguments forwarded to either the container provider's
|
| 929 |
+
`start_container` (docker) or to the `UVProvider` constructor/start (uv).
|
| 930 |
+
When `use_docker=False`, the `project_path` argument can be used to override
|
| 931 |
+
the default git URL (`git+https://huggingface.co/spaces/{repo_id}`).
|
| 932 |
+
|
| 933 |
+
Returns:
|
| 934 |
+
`_BootstrapResult`: `await` for a connected async client, or call
|
| 935 |
+
`.sync()` for a connected `SyncEnvClient`.
|
| 936 |
+
|
| 937 |
+
Examples:
|
| 938 |
+
|
| 939 |
+
```python
|
| 940 |
+
# Async: pull and run from HF Docker registry
|
| 941 |
+
env = await MyEnv.from_env("openenv/echo-env")
|
| 942 |
+
|
| 943 |
+
# Sync: chain .sync()
|
| 944 |
+
env = MyEnv.from_env("openenv/echo-env").sync()
|
| 945 |
+
|
| 946 |
+
# Run locally with UV (clones the space)
|
| 947 |
+
env = await MyEnv.from_env("openenv/echo-env", use_docker=False)
|
| 948 |
+
|
| 949 |
+
# Run from a local checkout
|
| 950 |
+
env = await MyEnv.from_env(
|
| 951 |
+
"openenv/echo-env",
|
| 952 |
+
use_docker=False,
|
| 953 |
+
project_path="/path/to/local/checkout"
|
| 954 |
+
)
|
| 955 |
+
```
|
| 956 |
+
"""
|
| 957 |
+
return _BootstrapResult(
|
| 958 |
+
lambda: cls._bootstrap_env(
|
| 959 |
+
repo_id,
|
| 960 |
+
use_docker=use_docker,
|
| 961 |
+
provider=provider,
|
| 962 |
+
**provider_kwargs,
|
| 963 |
+
)
|
| 964 |
+
)
|
| 965 |
+
|
| 966 |
+
@abstractmethod
|
| 967 |
+
def _step_payload(self, action: ActT) -> Dict[str, Any]:
|
| 968 |
+
"""Convert an Action object to the JSON data expected by the env server."""
|
| 969 |
+
raise NotImplementedError
|
| 970 |
+
|
| 971 |
+
@abstractmethod
|
| 972 |
+
def _parse_result(self, payload: Dict[str, Any]) -> StepResult[ObsT]:
|
| 973 |
+
"""Convert a JSON response from the env server to StepResult[ObsT]."""
|
| 974 |
+
raise NotImplementedError
|
| 975 |
+
|
| 976 |
+
@abstractmethod
|
| 977 |
+
def _parse_state(self, payload: Dict[str, Any]) -> StateT:
|
| 978 |
+
"""Convert a JSON response from the state endpoint to a State object."""
|
| 979 |
+
raise NotImplementedError
|
| 980 |
+
|
| 981 |
+
def reset(self, **kwargs: Any) -> Any:
|
| 982 |
+
return self._dispatch(lambda: self._reset_async(**kwargs))
|
| 983 |
+
|
| 984 |
+
async def _reset_async(self, **kwargs: Any) -> StepResult[ObsT]:
|
| 985 |
+
"""
|
| 986 |
+
Reset the environment with optional parameters.
|
| 987 |
+
|
| 988 |
+
Args:
|
| 989 |
+
**kwargs:
|
| 990 |
+
Optional parameters passed to the environment's reset method.
|
| 991 |
+
|
| 992 |
+
Returns:
|
| 993 |
+
StepResult containing initial observation
|
| 994 |
+
"""
|
| 995 |
+
message = {
|
| 996 |
+
"type": "reset",
|
| 997 |
+
"data": kwargs,
|
| 998 |
+
}
|
| 999 |
+
response = await self._send_and_receive(message)
|
| 1000 |
+
return self._parse_result(response.get("data", {}))
|
| 1001 |
+
|
| 1002 |
+
def step(self, action: ActT, **kwargs: Any) -> Any:
|
| 1003 |
+
return self._dispatch(lambda: self._step_async(action, **kwargs))
|
| 1004 |
+
|
| 1005 |
+
async def _step_async(self, action: ActT, **kwargs: Any) -> StepResult[ObsT]:
|
| 1006 |
+
"""
|
| 1007 |
+
Execute an action in the environment.
|
| 1008 |
+
|
| 1009 |
+
Args:
|
| 1010 |
+
action:
|
| 1011 |
+
The action to execute.
|
| 1012 |
+
**kwargs:
|
| 1013 |
+
Optional parameters (currently ignored).
|
| 1014 |
+
|
| 1015 |
+
Returns:
|
| 1016 |
+
StepResult containing observation, reward, and done status
|
| 1017 |
+
"""
|
| 1018 |
+
message = {
|
| 1019 |
+
"type": "step",
|
| 1020 |
+
"data": self._step_payload(action),
|
| 1021 |
+
}
|
| 1022 |
+
response = await self._send_and_receive(message)
|
| 1023 |
+
return self._parse_result(response.get("data", {}))
|
| 1024 |
+
|
| 1025 |
+
def state(self) -> Any:
|
| 1026 |
+
return self._dispatch(self._state_async)
|
| 1027 |
+
|
| 1028 |
+
async def _state_async(self) -> StateT:
|
| 1029 |
+
"""
|
| 1030 |
+
Get the current environment state from the server.
|
| 1031 |
+
|
| 1032 |
+
Returns:
|
| 1033 |
+
State object with environment state information
|
| 1034 |
+
"""
|
| 1035 |
+
message = {"type": "state"}
|
| 1036 |
+
response = await self._send_and_receive(message)
|
| 1037 |
+
return self._parse_state(response.get("data", {}))
|
| 1038 |
+
|
| 1039 |
+
def close(self) -> Any:
|
| 1040 |
+
return self._dispatch(self._close_async)
|
| 1041 |
+
|
| 1042 |
+
async def _close_child_clients(self) -> asyncio.CancelledError | None:
|
| 1043 |
+
"""Close every captured child while deferring caller cancellation."""
|
| 1044 |
+
children = list(self._child_clients)
|
| 1045 |
+
self._child_clients.clear()
|
| 1046 |
+
close_tasks = []
|
| 1047 |
+
deferred_cancellation: asyncio.CancelledError | None = None
|
| 1048 |
+
|
| 1049 |
+
for child in children:
|
| 1050 |
+
try:
|
| 1051 |
+
close_tasks.append(asyncio.ensure_future(child.close()))
|
| 1052 |
+
except asyncio.CancelledError as exc:
|
| 1053 |
+
if deferred_cancellation is None:
|
| 1054 |
+
deferred_cancellation = exc
|
| 1055 |
+
except Exception:
|
| 1056 |
+
pass # Best effort
|
| 1057 |
+
|
| 1058 |
+
if close_tasks:
|
| 1059 |
+
close_group = asyncio.gather(*close_tasks, return_exceptions=True)
|
| 1060 |
+
while not close_group.done():
|
| 1061 |
+
try:
|
| 1062 |
+
await asyncio.shield(close_group)
|
| 1063 |
+
except asyncio.CancelledError as exc:
|
| 1064 |
+
if deferred_cancellation is None:
|
| 1065 |
+
deferred_cancellation = exc
|
| 1066 |
+
|
| 1067 |
+
for result in close_group.result():
|
| 1068 |
+
if (
|
| 1069 |
+
isinstance(result, asyncio.CancelledError)
|
| 1070 |
+
and deferred_cancellation is None
|
| 1071 |
+
):
|
| 1072 |
+
deferred_cancellation = result
|
| 1073 |
+
|
| 1074 |
+
return deferred_cancellation
|
| 1075 |
+
|
| 1076 |
+
async def _close_async(self) -> None:
|
| 1077 |
+
"""
|
| 1078 |
+
Close the WebSocket connection and clean up resources.
|
| 1079 |
+
|
| 1080 |
+
If this client was created via from_docker_image() or from_env(),
|
| 1081 |
+
this will also stop and remove the associated container/process.
|
| 1082 |
+
"""
|
| 1083 |
+
deferred_cancellation: asyncio.CancelledError | None = None
|
| 1084 |
+
try:
|
| 1085 |
+
try:
|
| 1086 |
+
try:
|
| 1087 |
+
child_cancellation = await self._close_child_clients()
|
| 1088 |
+
except asyncio.CancelledError as exc:
|
| 1089 |
+
child_cancellation = exc
|
| 1090 |
+
if child_cancellation is not None:
|
| 1091 |
+
deferred_cancellation = child_cancellation
|
| 1092 |
+
|
| 1093 |
+
# A real close waits out backgrounded closes, while shielding
|
| 1094 |
+
# their socket handshakes from cancellation.
|
| 1095 |
+
try:
|
| 1096 |
+
await self._drain_pending_close_tasks()
|
| 1097 |
+
except asyncio.CancelledError as exc:
|
| 1098 |
+
if deferred_cancellation is None:
|
| 1099 |
+
deferred_cancellation = exc
|
| 1100 |
+
finally:
|
| 1101 |
+
# Run even when child or pending-close cleanup is cancelled. A
|
| 1102 |
+
# client may already have reconnected, and that current socket
|
| 1103 |
+
# must not remain cached or open during teardown.
|
| 1104 |
+
try:
|
| 1105 |
+
await self._disconnect_async()
|
| 1106 |
+
except asyncio.CancelledError as exc:
|
| 1107 |
+
if deferred_cancellation is None:
|
| 1108 |
+
deferred_cancellation = exc
|
| 1109 |
+
finally:
|
| 1110 |
+
self._stop_provider()
|
| 1111 |
+
|
| 1112 |
+
if deferred_cancellation is not None:
|
| 1113 |
+
raise deferred_cancellation
|
| 1114 |
+
|
| 1115 |
+
def _stop_provider(self) -> None:
|
| 1116 |
+
"""Stop the provider once, retaining it for retries and later startup."""
|
| 1117 |
+
try:
|
| 1118 |
+
provider = self._provider
|
| 1119 |
+
if provider is None or self._provider_stopped:
|
| 1120 |
+
return
|
| 1121 |
+
|
| 1122 |
+
if hasattr(provider, "stop_container"):
|
| 1123 |
+
stop = provider.stop_container
|
| 1124 |
+
elif hasattr(provider, "stop"):
|
| 1125 |
+
stop = provider.stop
|
| 1126 |
+
else:
|
| 1127 |
+
return
|
| 1128 |
+
|
| 1129 |
+
self._provider_cleanup_pending = True
|
| 1130 |
+
stop()
|
| 1131 |
+
# Only a successful stop discharges our cleanup responsibility.
|
| 1132 |
+
self._provider_cleanup_pending = False
|
| 1133 |
+
self._provider_stopped = True
|
| 1134 |
+
finally:
|
| 1135 |
+
if self._start_provider_on_connect:
|
| 1136 |
+
self._base_url = None
|
| 1137 |
+
self._ws_url = None
|
| 1138 |
+
|
| 1139 |
+
def _stop_provider_best_effort(self) -> None:
|
| 1140 |
+
"""Stop the underlying provider directly, ignoring any errors.
|
| 1141 |
+
|
| 1142 |
+
Releases a started container/process when there is no connected client
|
| 1143 |
+
to `close()` through the normal path — e.g. sync bootstrap setup fails
|
| 1144 |
+
after the provider started but before the connection is established, so
|
| 1145 |
+
routing cleanup through the (possibly broken) sync loop is not an option.
|
| 1146 |
+
"""
|
| 1147 |
+
with suppress(Exception):
|
| 1148 |
+
self._stop_provider()
|
| 1149 |
+
|
| 1150 |
+
async def __aenter__(self) -> "EnvClient":
|
| 1151 |
+
"""Enter async context manager, ensuring connection is established."""
|
| 1152 |
+
self._claim_execution_mode("async")
|
| 1153 |
+
await self.connect()
|
| 1154 |
+
return self
|
| 1155 |
+
|
| 1156 |
+
async def __aexit__(self, exc_type, exc_val, exc_tb) -> None:
|
| 1157 |
+
"""Exit async context manager, closing connection."""
|
| 1158 |
+
await self.close()
|
| 1159 |
+
|
| 1160 |
+
def __enter__(self) -> "EnvClient":
|
| 1161 |
+
"""Enter sync context manager, ensuring connection is established."""
|
| 1162 |
+
self._run_sync(self._connect_async)
|
| 1163 |
+
return self
|
| 1164 |
+
|
| 1165 |
+
def __exit__(self, exc_type, exc_val, exc_tb) -> None:
|
| 1166 |
+
"""Exit sync context manager, closing connection."""
|
| 1167 |
+
self._run_sync(self._close_async)
|
| 1168 |
+
|
| 1169 |
+
def sync(self) -> "SyncEnvClient":
|
| 1170 |
+
"""
|
| 1171 |
+
Return a synchronous wrapper around this async client.
|
| 1172 |
+
|
| 1173 |
+
Use this method when you need synchronous access to the environment
|
| 1174 |
+
without async/await syntax. This is useful for:
|
| 1175 |
+
- Integration with synchronous codebases
|
| 1176 |
+
- Interactive/REPL usage
|
| 1177 |
+
- Stopping async from "infecting" the call stack
|
| 1178 |
+
|
| 1179 |
+
Returns:
|
| 1180 |
+
SyncEnvClient wrapper that provides synchronous methods
|
| 1181 |
+
|
| 1182 |
+
Examples:
|
| 1183 |
+
|
| 1184 |
+
```python
|
| 1185 |
+
async_client = GenericEnvClient(base_url="http://localhost:8000")
|
| 1186 |
+
sync_client = async_client.sync()
|
| 1187 |
+
|
| 1188 |
+
with sync_client:
|
| 1189 |
+
result = sync_client.reset()
|
| 1190 |
+
result = sync_client.step({"code": "print('hello')"})
|
| 1191 |
+
```
|
| 1192 |
+
"""
|
| 1193 |
+
from .sync_client import SyncEnvClient
|
| 1194 |
+
|
| 1195 |
+
if self._sync_client is None:
|
| 1196 |
+
self._sync_client = SyncEnvClient(self)
|
| 1197 |
+
return self._sync_client
|
src/core/env_server/__init__.py
ADDED
|
@@ -0,0 +1,155 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# SPDX-License-Identifier: BSD-3-Clause
|
| 2 |
+
|
| 3 |
+
"""Core environment interfaces and types."""
|
| 4 |
+
|
| 5 |
+
from .base_transforms import CompositeTransform, NullTransform
|
| 6 |
+
from .exceptions import (
|
| 7 |
+
ConcurrencyConfigurationError,
|
| 8 |
+
EnvironmentFactoryError,
|
| 9 |
+
OpenEnvError,
|
| 10 |
+
SessionCapacityError,
|
| 11 |
+
SessionCreationError,
|
| 12 |
+
SessionNotFoundError,
|
| 13 |
+
)
|
| 14 |
+
from .http_server import create_app, create_fastapi_app, HTTPEnvServer
|
| 15 |
+
from .interfaces import Environment, Message, ModelTokenizer, TaskProvider, Transform
|
| 16 |
+
|
| 17 |
+
try:
|
| 18 |
+
from .mcp_environment import MCPEnvironment
|
| 19 |
+
except ModuleNotFoundError:
|
| 20 |
+
MCPEnvironment = None # type: ignore[assignment]
|
| 21 |
+
|
| 22 |
+
from .mcp_types import (
|
| 23 |
+
CallToolAction,
|
| 24 |
+
CallToolObservation,
|
| 25 |
+
JsonRpcError,
|
| 26 |
+
# JSON-RPC types
|
| 27 |
+
JsonRpcErrorCode,
|
| 28 |
+
JsonRpcRequest,
|
| 29 |
+
JsonRpcResponse,
|
| 30 |
+
ListToolsAction,
|
| 31 |
+
ListToolsObservation,
|
| 32 |
+
McpMethod,
|
| 33 |
+
RESERVED_TOOL_NAMES,
|
| 34 |
+
Tool,
|
| 35 |
+
ToolError,
|
| 36 |
+
ToolErrorType,
|
| 37 |
+
WSMCPMessage,
|
| 38 |
+
WSMCPResponse,
|
| 39 |
+
)
|
| 40 |
+
from .route_config import GetEndpointConfig
|
| 41 |
+
from .serialization import (
|
| 42 |
+
deserialize_action,
|
| 43 |
+
deserialize_action_with_preprocessing,
|
| 44 |
+
serialize_observation,
|
| 45 |
+
)
|
| 46 |
+
from .types import (
|
| 47 |
+
Action,
|
| 48 |
+
BaseMessage,
|
| 49 |
+
ConcurrencyConfig,
|
| 50 |
+
GetTaskRangeRequest,
|
| 51 |
+
GetTaskRequest,
|
| 52 |
+
HealthResponse,
|
| 53 |
+
HealthStatus,
|
| 54 |
+
ListTasksRequest,
|
| 55 |
+
NumTasksRequest,
|
| 56 |
+
Observation,
|
| 57 |
+
SchemaResponse,
|
| 58 |
+
ServerCapacityStatus,
|
| 59 |
+
ServerMode,
|
| 60 |
+
SessionInfo,
|
| 61 |
+
State,
|
| 62 |
+
WSCloseMessage,
|
| 63 |
+
WSErrorCode,
|
| 64 |
+
WSErrorResponse,
|
| 65 |
+
WSIncomingMessage,
|
| 66 |
+
WSObservationResponse,
|
| 67 |
+
WSResetMessage,
|
| 68 |
+
WSStateMessage,
|
| 69 |
+
WSStateResponse,
|
| 70 |
+
WSStepMessage,
|
| 71 |
+
)
|
| 72 |
+
|
| 73 |
+
try:
|
| 74 |
+
from .web_interface import create_web_interface_app, WebInterfaceManager
|
| 75 |
+
except ModuleNotFoundError:
|
| 76 |
+
create_web_interface_app = None # type: ignore[assignment]
|
| 77 |
+
WebInterfaceManager = None # type: ignore[assignment]
|
| 78 |
+
|
| 79 |
+
__all__ = [
|
| 80 |
+
# Core interfaces
|
| 81 |
+
"Environment",
|
| 82 |
+
"Transform",
|
| 83 |
+
"TaskProvider",
|
| 84 |
+
"Message",
|
| 85 |
+
"ModelTokenizer",
|
| 86 |
+
# Types
|
| 87 |
+
"Action",
|
| 88 |
+
"Observation",
|
| 89 |
+
"State",
|
| 90 |
+
"SchemaResponse",
|
| 91 |
+
"HealthResponse",
|
| 92 |
+
"ListTasksRequest",
|
| 93 |
+
"NumTasksRequest",
|
| 94 |
+
"GetTaskRequest",
|
| 95 |
+
"GetTaskRangeRequest",
|
| 96 |
+
# Enums
|
| 97 |
+
"HealthStatus",
|
| 98 |
+
"ServerMode",
|
| 99 |
+
"WSErrorCode",
|
| 100 |
+
# WebSocket message types
|
| 101 |
+
"BaseMessage",
|
| 102 |
+
"WSIncomingMessage",
|
| 103 |
+
"WSResetMessage",
|
| 104 |
+
"WSStepMessage",
|
| 105 |
+
"WSStateMessage",
|
| 106 |
+
"WSCloseMessage",
|
| 107 |
+
"WSObservationResponse",
|
| 108 |
+
"WSStateResponse",
|
| 109 |
+
"WSErrorResponse",
|
| 110 |
+
# Concurrency types
|
| 111 |
+
"ConcurrencyConfig",
|
| 112 |
+
"ServerCapacityStatus",
|
| 113 |
+
"SessionInfo",
|
| 114 |
+
# Exceptions
|
| 115 |
+
"OpenEnvError",
|
| 116 |
+
"ConcurrencyConfigurationError",
|
| 117 |
+
"SessionCapacityError",
|
| 118 |
+
"SessionNotFoundError",
|
| 119 |
+
"SessionCreationError",
|
| 120 |
+
"EnvironmentFactoryError",
|
| 121 |
+
# Base transforms
|
| 122 |
+
"CompositeTransform",
|
| 123 |
+
"NullTransform",
|
| 124 |
+
# HTTP Server
|
| 125 |
+
"HTTPEnvServer",
|
| 126 |
+
"create_app",
|
| 127 |
+
"create_fastapi_app",
|
| 128 |
+
# Web Interface
|
| 129 |
+
"create_web_interface_app",
|
| 130 |
+
"WebInterfaceManager",
|
| 131 |
+
# Serialization utilities
|
| 132 |
+
"deserialize_action",
|
| 133 |
+
"deserialize_action_with_preprocessing",
|
| 134 |
+
"serialize_observation",
|
| 135 |
+
# Route configuration
|
| 136 |
+
"GetEndpointConfig",
|
| 137 |
+
# MCP types
|
| 138 |
+
"Tool",
|
| 139 |
+
"ToolError",
|
| 140 |
+
"ToolErrorType",
|
| 141 |
+
"ListToolsAction",
|
| 142 |
+
"CallToolAction",
|
| 143 |
+
"ListToolsObservation",
|
| 144 |
+
"CallToolObservation",
|
| 145 |
+
"WSMCPMessage",
|
| 146 |
+
"WSMCPResponse",
|
| 147 |
+
"RESERVED_TOOL_NAMES",
|
| 148 |
+
"MCPEnvironment",
|
| 149 |
+
# JSON-RPC types
|
| 150 |
+
"JsonRpcErrorCode",
|
| 151 |
+
"JsonRpcError",
|
| 152 |
+
"JsonRpcRequest",
|
| 153 |
+
"JsonRpcResponse",
|
| 154 |
+
"McpMethod",
|
| 155 |
+
]
|
src/core/env_server/_utils.py
ADDED
|
@@ -0,0 +1,10 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""Shared env server implementation helpers."""
|
| 2 |
+
|
| 3 |
+
from __future__ import annotations
|
| 4 |
+
|
| 5 |
+
from typing import Any
|
| 6 |
+
|
| 7 |
+
|
| 8 |
+
def overrides_method(method: Any, base_method: Any) -> bool:
|
| 9 |
+
"""Return whether a bound method differs from the base implementation."""
|
| 10 |
+
return getattr(method, "__func__", method) is not base_method
|
src/core/env_server/base_transforms.py
ADDED
|
@@ -0,0 +1,25 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# SPDX-License-Identifier: BSD-3-Clause
|
| 2 |
+
|
| 3 |
+
"""Base transform implementations for composing environment-specific transforms."""
|
| 4 |
+
|
| 5 |
+
from .interfaces import Transform
|
| 6 |
+
from .types import Observation
|
| 7 |
+
|
| 8 |
+
|
| 9 |
+
class CompositeTransform(Transform):
|
| 10 |
+
"""Combines multiple transforms into a single transform."""
|
| 11 |
+
|
| 12 |
+
def __init__(self, transforms: list[Transform]):
|
| 13 |
+
self.transforms = transforms
|
| 14 |
+
|
| 15 |
+
def __call__(self, observation: Observation) -> Observation:
|
| 16 |
+
for transform in self.transforms:
|
| 17 |
+
observation = transform(observation)
|
| 18 |
+
return observation
|
| 19 |
+
|
| 20 |
+
|
| 21 |
+
class NullTransform(Transform):
|
| 22 |
+
"""Default transform that passes through unchanged."""
|
| 23 |
+
|
| 24 |
+
def __call__(self, observation: Observation) -> Observation:
|
| 25 |
+
return observation
|
src/core/env_server/exceptions.py
ADDED
|
@@ -0,0 +1,101 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# SPDX-License-Identifier: BSD-3-Clause
|
| 2 |
+
|
| 3 |
+
"""Custom exceptions for environment server operations."""
|
| 4 |
+
|
| 5 |
+
from typing import Optional
|
| 6 |
+
|
| 7 |
+
|
| 8 |
+
class OpenEnvError(Exception):
|
| 9 |
+
"""Base exception for all OpenEnv errors."""
|
| 10 |
+
|
| 11 |
+
pass
|
| 12 |
+
|
| 13 |
+
|
| 14 |
+
class ConcurrencyConfigurationError(OpenEnvError):
|
| 15 |
+
"""
|
| 16 |
+
Raised when an environment is misconfigured for concurrent sessions.
|
| 17 |
+
|
| 18 |
+
This error is raised during server startup when max_concurrent_envs > 1
|
| 19 |
+
is specified for an environment that is not marked as SUPPORTS_CONCURRENT_SESSIONS.
|
| 20 |
+
"""
|
| 21 |
+
|
| 22 |
+
def __init__(
|
| 23 |
+
self,
|
| 24 |
+
environment_name: str,
|
| 25 |
+
max_concurrent_envs: int,
|
| 26 |
+
message: Optional[str] = None,
|
| 27 |
+
):
|
| 28 |
+
self.environment_name = environment_name
|
| 29 |
+
self.max_concurrent_envs = max_concurrent_envs
|
| 30 |
+
|
| 31 |
+
if message is None:
|
| 32 |
+
message = (
|
| 33 |
+
f"Environment '{environment_name}' is not marked as SUPPORTS_CONCURRENT_SESSIONS. "
|
| 34 |
+
f"Cannot run with max_concurrent_envs={max_concurrent_envs}. "
|
| 35 |
+
f"Either set max_concurrent_envs=1 or ensure the environment "
|
| 36 |
+
f"properly isolates session state and set SUPPORTS_CONCURRENT_SESSIONS=True."
|
| 37 |
+
)
|
| 38 |
+
|
| 39 |
+
super().__init__(message)
|
| 40 |
+
|
| 41 |
+
|
| 42 |
+
class SessionCapacityError(OpenEnvError):
|
| 43 |
+
"""
|
| 44 |
+
Raised when the server cannot accept new sessions due to capacity limits.
|
| 45 |
+
|
| 46 |
+
This error is raised when a new WebSocket connection is attempted but
|
| 47 |
+
the server has already reached max_concurrent_envs active sessions.
|
| 48 |
+
"""
|
| 49 |
+
|
| 50 |
+
def __init__(
|
| 51 |
+
self,
|
| 52 |
+
active_sessions: int,
|
| 53 |
+
max_sessions: int,
|
| 54 |
+
message: Optional[str] = None,
|
| 55 |
+
):
|
| 56 |
+
self.active_sessions = active_sessions
|
| 57 |
+
self.max_sessions = max_sessions
|
| 58 |
+
|
| 59 |
+
if message is None:
|
| 60 |
+
message = (
|
| 61 |
+
f"Server at capacity: {active_sessions}/{max_sessions} sessions active. "
|
| 62 |
+
f"Cannot accept new connections."
|
| 63 |
+
)
|
| 64 |
+
|
| 65 |
+
super().__init__(message)
|
| 66 |
+
|
| 67 |
+
|
| 68 |
+
class SessionNotFoundError(OpenEnvError):
|
| 69 |
+
"""Raised when attempting to access a session that does not exist."""
|
| 70 |
+
|
| 71 |
+
def __init__(self, session_id: str, message: Optional[str] = None):
|
| 72 |
+
self.session_id = session_id
|
| 73 |
+
|
| 74 |
+
if message is None:
|
| 75 |
+
message = f"Session '{session_id}' not found."
|
| 76 |
+
|
| 77 |
+
super().__init__(message)
|
| 78 |
+
|
| 79 |
+
|
| 80 |
+
class SessionCreationError(OpenEnvError):
|
| 81 |
+
"""Raised when a session cannot be created."""
|
| 82 |
+
|
| 83 |
+
def __init__(self, reason: str, message: Optional[str] = None):
|
| 84 |
+
self.reason = reason
|
| 85 |
+
|
| 86 |
+
if message is None:
|
| 87 |
+
message = f"Failed to create session: {reason}"
|
| 88 |
+
|
| 89 |
+
super().__init__(message)
|
| 90 |
+
|
| 91 |
+
|
| 92 |
+
class EnvironmentFactoryError(OpenEnvError):
|
| 93 |
+
"""Raised when the environment factory fails to create an instance."""
|
| 94 |
+
|
| 95 |
+
def __init__(self, factory_name: str, message: Optional[str] = None):
|
| 96 |
+
self.factory_name = factory_name
|
| 97 |
+
|
| 98 |
+
if message is None:
|
| 99 |
+
message = f"Environment factory '{factory_name}' failed to create instance."
|
| 100 |
+
|
| 101 |
+
super().__init__(message)
|
src/core/env_server/gradio_theme.py
ADDED
|
@@ -0,0 +1,142 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# SPDX-License-Identifier: BSD-3-Clause
|
| 2 |
+
|
| 3 |
+
"""Theme for the OpenEnv Gradio UI: white, hairline borders, one green accent."""
|
| 4 |
+
|
| 5 |
+
from __future__ import annotations
|
| 6 |
+
|
| 7 |
+
import gradio as gr
|
| 8 |
+
|
| 9 |
+
_GREEN_HUE = gr.themes.Color(
|
| 10 |
+
c50="#e6f4ea",
|
| 11 |
+
c100="#ceead6",
|
| 12 |
+
c200="#a8dab5",
|
| 13 |
+
c300="#6fcc8b",
|
| 14 |
+
c400="#3fb950",
|
| 15 |
+
c500="#238636",
|
| 16 |
+
c600="#1a7f37",
|
| 17 |
+
c700="#116329",
|
| 18 |
+
c800="#0a4620",
|
| 19 |
+
c900="#033a16",
|
| 20 |
+
c950="#04200d",
|
| 21 |
+
)
|
| 22 |
+
|
| 23 |
+
_NEUTRAL_HUE = gr.themes.Color(
|
| 24 |
+
c50="#fafafa",
|
| 25 |
+
c100="#f4f4f5",
|
| 26 |
+
c200="#e4e4e7",
|
| 27 |
+
c300="#d4d4d8",
|
| 28 |
+
c400="#a1a1aa",
|
| 29 |
+
c500="#71717a",
|
| 30 |
+
c600="#52525b",
|
| 31 |
+
c700="#3f3f46",
|
| 32 |
+
c800="#27272a",
|
| 33 |
+
c900="#18181b",
|
| 34 |
+
c950="#0a0a0a",
|
| 35 |
+
)
|
| 36 |
+
|
| 37 |
+
OPENENV_GRADIO_THEME = gr.themes.Base(
|
| 38 |
+
primary_hue=_GREEN_HUE,
|
| 39 |
+
secondary_hue=_NEUTRAL_HUE,
|
| 40 |
+
neutral_hue=_NEUTRAL_HUE,
|
| 41 |
+
font=(gr.themes.GoogleFont("Geist"), "system-ui", "sans-serif"),
|
| 42 |
+
font_mono=(gr.themes.GoogleFont("Geist Mono"), "ui-monospace", "monospace"),
|
| 43 |
+
radius_size=gr.themes.sizes.radius_md,
|
| 44 |
+
).set(
|
| 45 |
+
body_background_fill="#ffffff",
|
| 46 |
+
body_text_color="#0a0a0a",
|
| 47 |
+
background_fill_primary="#ffffff",
|
| 48 |
+
background_fill_secondary="#fafafa",
|
| 49 |
+
block_background_fill="#ffffff",
|
| 50 |
+
block_border_width="0px",
|
| 51 |
+
block_label_text_color="#52525b",
|
| 52 |
+
block_label_background_fill="transparent",
|
| 53 |
+
border_color_primary="#ecedef",
|
| 54 |
+
input_background_fill="#ffffff",
|
| 55 |
+
input_border_color="#e4e4e7",
|
| 56 |
+
input_border_color_focus="#1a7f37",
|
| 57 |
+
button_primary_background_fill="#0a0a0a",
|
| 58 |
+
button_primary_background_fill_hover="#27272a",
|
| 59 |
+
button_primary_text_color="#ffffff",
|
| 60 |
+
button_primary_border_color="#0a0a0a",
|
| 61 |
+
button_secondary_background_fill="#ffffff",
|
| 62 |
+
button_secondary_background_fill_hover="#fafafa",
|
| 63 |
+
button_secondary_text_color="#0a0a0a",
|
| 64 |
+
button_secondary_border_color="#e4e4e7",
|
| 65 |
+
button_border_width="1px",
|
| 66 |
+
checkbox_background_color_selected="#1a7f37",
|
| 67 |
+
checkbox_border_color_selected="#1a7f37",
|
| 68 |
+
color_accent="#1a7f37",
|
| 69 |
+
color_accent_soft="#f0f7f2",
|
| 70 |
+
color_accent_soft_dark="#0d2416",
|
| 71 |
+
body_text_color_subdued="#52525b",
|
| 72 |
+
body_text_color_subdued_dark="#a1a1aa",
|
| 73 |
+
body_background_fill_dark="#0a0a0a",
|
| 74 |
+
body_text_color_dark="#fafafa",
|
| 75 |
+
background_fill_primary_dark="#0a0a0a",
|
| 76 |
+
background_fill_secondary_dark="#141414",
|
| 77 |
+
block_background_fill_dark="#0a0a0a",
|
| 78 |
+
block_label_text_color_dark="#a1a1aa",
|
| 79 |
+
border_color_primary_dark="#27272a",
|
| 80 |
+
input_background_fill_dark="#0a0a0a",
|
| 81 |
+
input_border_color_dark="#3f3f46",
|
| 82 |
+
button_primary_background_fill_dark="#fafafa",
|
| 83 |
+
button_primary_background_fill_hover_dark="#e4e4e7",
|
| 84 |
+
button_primary_text_color_dark="#0a0a0a",
|
| 85 |
+
button_primary_border_color_dark="#fafafa",
|
| 86 |
+
button_secondary_background_fill_dark="#0a0a0a",
|
| 87 |
+
button_secondary_background_fill_hover_dark="#18181b",
|
| 88 |
+
button_secondary_text_color_dark="#fafafa",
|
| 89 |
+
button_secondary_border_color_dark="#3f3f46",
|
| 90 |
+
)
|
| 91 |
+
|
| 92 |
+
OPENENV_GRADIO_CSS = """
|
| 93 |
+
.gradio-container { max-width: 1120px !important; margin: 0 auto !important; }
|
| 94 |
+
.oe-header { padding: 24px 0 12px; }
|
| 95 |
+
.oe-header h1 { margin: 6px 0 4px; font-size: 40px; line-height: 1.1; font-weight: 700; letter-spacing: -.03em; }
|
| 96 |
+
.oe-header p, .oe-step p, .oe-episode small, .oe-muted { color: var(--body-text-color-subdued); }
|
| 97 |
+
.oe-header p { margin: 0; font-size: 16px; max-width: 640px; }
|
| 98 |
+
.oe-eyebrow { font-family: var(--font-mono); font-size: 13px; font-weight: 500; color: var(--color-accent); }
|
| 99 |
+
.oe-card { border: 1px solid var(--border-color-primary) !important; border-radius: 12px !important; padding: 18px 20px !important; background: var(--background-fill-primary) !important; gap: 12px !important; }
|
| 100 |
+
.oe-row { align-items: center !important; }
|
| 101 |
+
.oe-step h2, .oe-episode h2 { margin: 0; font-size: 17px; font-weight: 600; }
|
| 102 |
+
.oe-step p { margin: 2px 0 0; font-size: 14px; }
|
| 103 |
+
.oe-step code, .oe-episode code, .oe-result code { font-family: var(--font-mono); background: none; padding: 0; }
|
| 104 |
+
.oe-run { flex: none !important; align-self: flex-start !important; width: auto !important; min-width: 140px !important; padding: 0 22px !important; margin-left: var(--block-padding, 10px) !important; }
|
| 105 |
+
.oe-card textarea, .oe-card input[type=number], .oe-card input[type=password], .oe-card input[type=text] { border: 1px solid var(--input-border-color) !important; border-radius: 8px !important; padding: 10px 12px !important; font-family: var(--font-mono) !important; }
|
| 106 |
+
.oe-tools .wrap, .oe-actions .wrap { gap: 8px; }
|
| 107 |
+
.oe-tools .wrap { display: grid !important; grid-template-columns: repeat(auto-fit, minmax(220px, 1fr)); }
|
| 108 |
+
.oe-actions .wrap { display: flex !important; flex-wrap: wrap; }
|
| 109 |
+
.oe-tools .wrap label, .oe-actions .wrap label { border: 1px solid var(--input-border-color) !important; border-radius: 10px !important; background: var(--background-fill-primary) !important; font-family: var(--font-mono); cursor: pointer; }
|
| 110 |
+
.oe-tools .wrap label { padding: 12px 14px !important; font-size: 13px; }
|
| 111 |
+
.oe-actions .wrap label { padding: 10px 16px !important; font-size: 14px; }
|
| 112 |
+
.oe-tools .wrap label.selected { border: 1.5px solid var(--color-accent) !important; background: var(--color-accent-soft) !important; }
|
| 113 |
+
.oe-actions input[type=radio] { display: none; }
|
| 114 |
+
.oe-result { border: 1px solid var(--border-color-primary); border-radius: 10px; overflow: hidden; }
|
| 115 |
+
.oe-output, .oe-fields { background: var(--background-fill-secondary); border-bottom: 1px solid var(--border-color-primary); }
|
| 116 |
+
.oe-output { margin: 0; padding: 14px 16px; font-family: var(--font-mono); font-size: 15px; white-space: pre-wrap; word-break: break-word; }
|
| 117 |
+
.oe-fields { padding: 10px 16px; }
|
| 118 |
+
.oe-fields div { display: grid; grid-template-columns: minmax(96px, 28%) minmax(0, 1fr); gap: 12px; padding: 3px 0; font-size: 13px; }
|
| 119 |
+
.oe-fields code { white-space: pre-wrap; word-break: break-word; }
|
| 120 |
+
.oe-fields span, .oe-stats span, .oe-episode li > span { font-family: var(--font-mono); color: var(--body-text-color-subdued); }
|
| 121 |
+
.oe-fields span { overflow: hidden; text-overflow: ellipsis; white-space: nowrap; }
|
| 122 |
+
.oe-stats { display: grid; grid-template-columns: repeat(3, minmax(0, 1fr)); }
|
| 123 |
+
.oe-stats div { padding: 8px 16px; border-right: 1px solid var(--border-color-primary); }
|
| 124 |
+
.oe-stats div:last-child { border-right: none; }
|
| 125 |
+
.oe-stats span { display: block; font-size: 12px; }
|
| 126 |
+
.oe-stats code { display: block; overflow: hidden; text-overflow: ellipsis; white-space: nowrap; }
|
| 127 |
+
.oe-error { color: #b42318; }
|
| 128 |
+
.dark .oe-error { color: #ff8a80; }
|
| 129 |
+
.oe-episode ol { list-style: none; margin: 8px 0 0; padding: 0; }
|
| 130 |
+
.oe-episode li { display: flex; gap: 12px; padding: 8px 0; border-top: 1px solid var(--border-color-primary); }
|
| 131 |
+
.oe-episode li > span { flex: none; width: 20px; font-size: 12px; }
|
| 132 |
+
.oe-episode code { display: block; font-size: 13px; word-break: break-word; }
|
| 133 |
+
.oe-episode small { display: block; font-size: 13px; word-break: break-word; }
|
| 134 |
+
.oe-muted { margin: 8px 0 0; font-size: 14px; }
|
| 135 |
+
.oe-code pre code { white-space: pre-wrap !important; word-break: break-word; }
|
| 136 |
+
.oe-json .code_wrap, .oe-json .cm-editor { max-height: 360px; overflow: auto; }
|
| 137 |
+
.oe-code pre { background: var(--background-fill-secondary) !important; border: 1px solid var(--border-color-primary); border-radius: 8px; }
|
| 138 |
+
.oe-raw { border: none !important; }
|
| 139 |
+
.oe-visual { flex: none !important; width: auto !important; min-width: 0 !important; }
|
| 140 |
+
.dark { --color-accent: #3fb950; }
|
| 141 |
+
footer { display: none !important; }
|
| 142 |
+
"""
|