Spaces:
Sleeping
Sleeping
| # Copyright (c) Meta Platforms, Inc. and affiliates. | |
| # All rights reserved. | |
| # | |
| # This source code is licensed under the BSD-style license found in the | |
| # LICENSE file in the root directory of this source tree. | |
| """ | |
| Benchmark Environment Implementation. | |
| A test environment designed for benchmarking server concurrency. | |
| Each session gets its own environment instance, enabling true concurrent execution. | |
| """ | |
| import hashlib | |
| import os | |
| import socket | |
| import time | |
| from uuid import uuid4 | |
| from openenv.core.env_server.interfaces import Environment | |
| from openenv.core.env_server.types import State | |
| from benchmark.models import BenchmarkAction, BenchmarkObservation | |
| class BenchmarkEnvironment(Environment): | |
| """ | |
| A benchmark environment for testing server concurrency. | |
| This environment is designed for testing concurrent WebSocket sessions. | |
| Each instance maintains its own state and can run independently. | |
| Key features: | |
| - Sleeps for a configurable duration on each step | |
| - Returns process ID and session hash to verify concurrency | |
| - Thread-safe and supports multiple concurrent sessions | |
| Example: | |
| >>> env = BenchmarkEnvironment() | |
| >>> obs = env.reset() | |
| >>> print(obs.session_hash) # Unique session identifier | |
| >>> | |
| >>> obs = env.step(BenchmarkAction(wait_seconds=1.0)) | |
| >>> print(obs.waited_seconds) # ~1.0 | |
| >>> print(obs.pid) # Process ID | |
| """ | |
| # Enable concurrent WebSocket sessions | |
| SUPPORTS_CONCURRENT_SESSIONS = True | |
| def __init__(self): | |
| """Initialize the benchmark environment.""" | |
| self._session_id = str(uuid4()) | |
| self._session_hash = hashlib.sha256( | |
| self._session_id.encode() | |
| ).hexdigest()[:12] | |
| self._state = State(episode_id=self._session_id, step_count=0) | |
| self._pid = os.getpid() | |
| self._host_url = self._get_host_url() | |
| def _get_host_url(self) -> str: | |
| """Get the host URL for debugging.""" | |
| try: | |
| hostname = socket.gethostname() | |
| return f"{hostname}:{os.getenv('PORT', '8000')}" | |
| except Exception: | |
| return "unknown" | |
| def reset(self) -> BenchmarkObservation: | |
| """ | |
| Reset the environment. | |
| Returns: | |
| BenchmarkObservation with session info | |
| """ | |
| self._state = State(episode_id=self._session_id, step_count=0) | |
| return BenchmarkObservation( | |
| waited_seconds=0.0, | |
| pid=self._pid, | |
| session_hash=self._session_hash, | |
| host_url=self._host_url, | |
| step_count=0, | |
| ) | |
| def step(self, action: BenchmarkAction) -> BenchmarkObservation: | |
| """ | |
| Execute a step - wait for the specified duration. | |
| Args: | |
| action: BenchmarkAction containing wait_seconds | |
| Returns: | |
| BenchmarkObservation with timing and concurrency info | |
| """ | |
| self._state.step_count += 1 | |
| wait_seconds = action.wait_seconds | |
| # Actually sleep to test concurrency | |
| if wait_seconds > 0: | |
| time.sleep(wait_seconds) | |
| return BenchmarkObservation( | |
| waited_seconds=wait_seconds, | |
| pid=self._pid, | |
| session_hash=self._session_hash, | |
| host_url=self._host_url, | |
| step_count=self._state.step_count, | |
| ) | |
| def state(self) -> State: | |
| """ | |
| Get the current environment state. | |
| Returns: | |
| Current State with episode_id and step_count | |
| """ | |
| return self._state | |