File size: 4,624 Bytes
4729fab
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
# 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 for benchmarking infrastructure and concurrency.
Actions specify how many seconds to wait, allowing testing of parallel execution.
"""

import asyncio
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 models import BenchmarkAction, BenchmarkObservation


def _get_host_url() -> str:
    """Get the host URL for this server."""
    hostname = socket.gethostname()
    port = os.environ.get("PORT", "8000")
    # Try to get the actual IP if possible
    try:
        ip = socket.gethostbyname(hostname)
    except socket.gaierror:
        ip = "127.0.0.1"
    return f"http://{ip}:{port}"


class BenchmarkEnvironment(Environment):
    """
    A benchmark environment for testing concurrency and infrastructure.

    Actions specify a number of seconds to wait (sleep), which is useful for
    testing parallel execution and concurrency limits. The environment returns
    identity information (host_url, pid, session_hash) to help verify which
    server instance handled the request.

    Example:
        >>> env = BenchmarkEnvironment()
        >>> obs = env.reset()
        >>> print(obs.host_url)  # "http://192.168.1.1:8000"
        >>> print(obs.pid)  # 12345
        >>> print(obs.session_hash)  # "a1b2c3d4..."
        >>>
        >>> obs = env.step(BenchmarkAction(wait_seconds=2.0))
        >>> print(obs.waited_seconds)  # 2.0
    """

    def __init__(self):
        """Initialize the benchmark environment."""
        self._state = State(episode_id=str(uuid4()), step_count=0)
        self._session_hash = hashlib.sha256(
            f"{uuid4()}-{time.time()}-{os.getpid()}".encode()
        ).hexdigest()[:16]
        self._pid = os.getpid()
        self._host_url = _get_host_url()

    def _make_observation(
        self, waited_seconds: float = 0.0, done: bool = False, reward: float = 0.0
    ) -> BenchmarkObservation:
        """Create an observation with current server identity."""
        return BenchmarkObservation(
            host_url=self._host_url,
            pid=self._pid,
            session_hash=self._session_hash,
            waited_seconds=waited_seconds,
            timestamp=time.time(),
            done=done,
            reward=reward,
        )

    def reset(self) -> BenchmarkObservation:
        """
        Reset the environment.

        Returns:
            BenchmarkObservation with server identity info
        """
        self._state = State(episode_id=str(uuid4()), step_count=0)

        return self._make_observation(waited_seconds=0.0, done=False, reward=0.0)

    def step(self, action: BenchmarkAction) -> BenchmarkObservation:  # type: ignore[override]
        """
        Execute a step by waiting for the specified seconds.

        Args:
            action: BenchmarkAction containing wait_seconds

        Returns:
            BenchmarkObservation with server identity and timing info
        """
        self._state.step_count += 1

        wait_time = max(0.0, action.wait_seconds)

        # Synchronous sleep - for async version, use step_async
        if wait_time > 0:
            time.sleep(wait_time)

        # Reward based on wait time (inverse - faster is better)
        reward = 1.0 / (1.0 + wait_time)

        return self._make_observation(
            waited_seconds=wait_time,
            done=False,
            reward=reward,
        )

    async def step_async(self, action: BenchmarkAction) -> BenchmarkObservation:
        """
        Async version of step - uses asyncio.sleep for better concurrency.

        Args:
            action: BenchmarkAction containing wait_seconds

        Returns:
            BenchmarkObservation with server identity and timing info
        """
        self._state.step_count += 1

        wait_time = max(0.0, action.wait_seconds)

        if wait_time > 0:
            await asyncio.sleep(wait_time)

        reward = 1.0 / (1.0 + wait_time)

        return self._make_observation(
            waited_seconds=wait_time,
            done=False,
            reward=reward,
        )

    @property
    def state(self) -> State:
        """
        Get the current environment state.

        Returns:
            Current State with episode_id and step_count
        """
        return self._state