[ENHANCEMENT] Support for Concurrent Environments

[rycerzes]

Currently, each environment instance requires its own Docker container. This creates significant resource overhead for parallel workloads:

• Running N parallel episodes requires N separate Docker containers
• Each container must be managed individually
• LocalDockerProvider creates one container per client instance

For RL training with dozens or hundreds of parallel rollouts, this becomes resource-intensive and impractical.

Current Architecture

Server side: Environment classes maintain state in instance variables (self._state, self._executor). One container = one environment instance = one episode session.

Client side: HTTPEnvClient.from_docker_image() starts a dedicated container via LocalDockerProvider.

Example:

# Each creates a separate container
env1 = CodingEnv.from_docker_image("coding-env:latest")  # Container 1
env2 = CodingEnv.from_docker_image("coding-env:latest")  # Container 2
env3 = CodingEnv.from_docker_image("coding-env:latest")  # Container 3

Use Case

I was trying to create concurrent environments for the poke-env wrapper. The underlying poke-env library supports concurrency natively, but the only way to support it via the OpenEnv wrapper is by running and managing multiple containers via Docker Swarm, Docker Compose, or similar container orchestration tools. This adds significant complexity compared to the native library.

Impact

• Resource usage: Each container has overhead (memory, CPU, startup time)
• Complexity: Managing dozens of containers requires orchestration tooling
• Developer experience: Creating parallel environments is cumbersome compared to native libraries

Possible Solutions

1. Multi-session per container: One server handles multiple isolated sessions via session IDs in the HTTP API
2. Vectorized environments: Single container runs N environments with batch API like VectorEnv
3. Something else: Open to ideas

[rycerzes]

CC: @zkwentz @burtenshaw @init27 @Darktex @pankit-eng - Would appreciate your thoughts on this

[guipenedo]

One possible implementation which wouldn't require servers to keep track of ids would be to have the clients keep track of state.
There would be a single type of request to the server (step) which instead of just taking action would take (action, current state) (state would be None at the start), and the server would return (action result, new state). An example would be a wordle game where the action result is the colors of the current attempt, while state would hold be the number of attempts so far.
This could be hidden behind abstractions client side, where the client could start a new episode (episode = env.episode()) which essentially just keeps track of state. Actions are then called on the episode (episode.step(action)). This would also allow to easily cache requests client side (as state being a step argument would make the calls to the env essentially stateless: episode.set_state(state)).
The server could then be a simple "action processor" with an async queue and a max number of requests to serve simultaneously (for example, the python sandbox env can run 5 executions simultaneously while the remaining wait in the queue)

[burtenshaw]

Thanks for reporting this @rycerzes . I've had the same experience with RL workloads. 'Multi-session per container' makes sense and is straightforward to implement. We should add this!

That said, it won't completely resolve the problem. What about a expanding the parallel container runtime providers. For example, we could add a provider for docker swarm like what we have for kubernetes, say SwarmProvider. IMO, this would be the most in keeping with the library's approach and require the least work for env builders.

[burtenshaw]

@guipenedo welcome to the OpenEnv repo! IMO, your proposal adds quite a lot of overhead to the env builder who is responsible for the server and client code, which goes against the repos RFCs. Ideally, we could solve this without requiring too much work on the env side.

[taha-yassine]

What about a expanding the parallel container runtime providers. For example, we could add a provider for docker swarm like what we have for kubernetes, say SwarmProvider. IMO, this would be the most in keeping with the library's approach and require the least work for env builders.

I think Docker Swarm would help with scaling and management, but imo it would not really fix the core issue. Swarm makes it easier to run lots of containers, but each environment would still need its own full container with its own memory, CPU use, and startup cost.

Also, since support for Kubernetes is planned, I feel like Swarm would not add much beyond being a simpler alternative for people who prefer it.

I still think a multi-session per container setup would actually reduce the overhead, since multiple environments could run inside one container. So Swarm might improve convenience, but it probably would not remove the fundamental inefficiency.

[burtenshaw]

Completely agree @taha-yassine. There are two sides to the OP problem. I will distil our two comments into a single solution:

Container Concurrency for sessions

We need multiple sessions so that the container is not locked from reset <> done. Instead, it will only be locked during single requests like step. Which for some tasks is significant.

Container Parallelism

Providers will simplify parallelism on various setups. Docker swarm is just convenience for single machine users. IMO, we could also add providers beyond this and Kubernetes.

For my own workloads, I've been using uvicorn server.app:app --workers 8 to get around this within the pending UVProvider.

An asynchronous client can help free the client too, but we can solve that separately.

[rycerzes]

Thanks everyone for the great ideas!
I think we're converging on a solid approach. @burtenshaw's two-pronged solution makes a lot of sense.

Container Concurrency (Priority 1)

Multi-session per container seems like the most impactful change. This would address the core resource inefficiency I'm experiencing with RL workloads. The key being that containers would only be locked during individual step requests rather than entire episodes.

Container Parallelism (Priority 2)

Enhanced orchestration providers would be valuable for different deployment scenarios, though as @taha-yassine correctly points out, this doesn't solve the fundamental overhead issue. #199 is for docker swarm provider but we will need more such providers for deployment.

Questions for Implementation

1. Session Management: How should we handle session IDs? HTTP headers, URL paths, or request body parameters?
2. Backward Compatibility: Can we implement this in a way that doesn't break existing single-session usage?
3. Error Isolation: How do we ensure that one session's errors don't affect others in the same container?

@guipenedo's stateless approach is interesting but I agree with @burtenshaw that it might add too much complexity for environment builders. The multi-session approach seems more aligned with OpenEnv's design philosophy.

Would be happy to contribute to the implementation once we settle on the API design.

[burtenshaw]

Would be happy to contribute to the implementation once we settle on the API design.

🎉 Legend! Thanks.

Session Management: How should we handle session IDs? HTTP headers, URL paths, or request body parameters?

IMO, we should keep this extremely simple for now with a session_id as uuid type that defaults to new uuid. Most users won't need to interact much with this.

I know @Wauplin had some further feedback on this.

Backward Compatibility: Can we implement this in a way that doesn't break existing single-session usage?

Yes. IMO the above approach achieves that. That said, this is a beta library and we should not let backward compatibility prevent core functionality implementation.

Error Isolation: How do we ensure that one session's errors don't affect others in the same container?

We won't in this PR. This is the use case for concurrency and existing error handling.

@rycerzes Really excited about this. Feel free to open an early draft PR and @pankit-eng @zkwentz and I can input to design, if you need.

[Wauplin]

Hey there, thanks for the ping @burtenshaw! 🤗 I went through the issue and agree with the current limitation, even if multiple containers/workers can be run in parallel.

I think the only way to correctly handle several envs on a single server is to redesign the API to use websockets instead of normal routes. I.e. have "1 client == 1 websocket connection == 1 running env". I don't see another way to properly assign an env to a client. Using a session_id header would work but be quickly be limited as the server wouldn't be able to know if a running environment can be garbage collected or not.

In practice here is how I would do it:

• create_app currently takes an Environment as input, which unfortunately put a hard constraint "1 server == 1 running env". Let's pass a env_factory instead
• create_app should also take as input a max_concurrent_envs: Optional[int] = None argument defaulting to either 1 (or an env variable ?). Not all envs are taking the same resources so we do want the user launching the server to control this.
  • (lower prio) let's accept MAX_CONCURRENT_ENVS environment variable as well? Make it easily configurable if running a docker container e.g. -e MAX_CONCURRENT_ENVS=12.
  • (lower prio) let's also accept -1 if we don't want to set a maximum number of envs?
• we replace the /reset , /step , /state routes by a /ws one which keeps the env running as long as the connection is alive. Client can send messages like {"type": "reset"} or {"type": "step", "data": {...}} to control the env. Server responds either with a state or an error message.
• when env is done (env.state.done is True), let's shut down the connection. Client has to open a new connection to start a new env.
• when calling /ws , if max number of envs is reached then user gets a proper error "Maximum concurrent environments reached"

This would be a hard breaking change in the API, requiring to update how the HTTPEnvClient and web interface are implemented. Luckily, it should not impact much the usability of the library since most changes would be handled internally. The client class could still have abstractions like client.reset, client.step, etc. The main user-facing breaking change would be the need to correctly close sessions if not in use anymore.

Let me know what you think. Happy to help on implementation details or code reviews if necessary :)

Note: with this format, it would also be easy to add parameters to the reset method as requested in #183 (to provide hyper-parameter about an env like a seed or a game name in atari env).

[burtenshaw]

Thanks for detailed suggestion @Wauplin ! It would be very cool to take this on, and I like that the changes are in openenv-core and server side. A few follow ups:

The main user-facing breaking change would be the need to correctly close sessions if not in use anymore.

In practice, this would mean that legacy implementation would leave sessions open. IMO, this is ok because that essentially happens anyway with the current implementation.

Let me know what you think. Happy to help on implementation details or code reviews if necessary :)

That would be super useful. Thanks

@rycerzes What do you think? For context, @Wauplin is the maintainer of the hf library.

[Wauplin]

Since a snippet is sometimes more explicit than an explanation, this is how I imagine the core logic could be implemented. Disclaimer: this is half-AI-generated, half human-(thoroughly)-reviewed. For a proper PR, it'd still require many changes to adapt server and client creation, and session lifecycle :)

@app.websocket("/ws")
async def websocket_endpoint(websocket: WebSocket):
    # Check if we can create a new environment
    if (
        self.max_concurrent_envs is not None
        and len(self._active_envs) >= self.max_concurrent_envs
    ):
        await websocket.close(
            code=1013,
            reason="Maximum concurrent environments reached. Please try again later.",
        )
        return

    # Accept the connection and create a new environment instance
    await websocket.accept()
    env = self.env_factory()
    self._active_envs[websocket] = env

    try:
        while True:
            # Receive message from client
            message = await websocket.receive_text()
            data = json.loads(message)

            msg_type = data.get("type")
            msg_data = data.get("data", {})

            if msg_type == "reset":
                observation = env.reset()  # TODO: support reset parameters? https://github.com/meta-pytorch/OpenEnv/issues/183
                await websocket.send_json(
                    {
                        "type": "observation",
                        "data": self._serialize_observation(observation),
                    }
                )

            elif msg_type == "step":
                action = self._deserialize_action(msg_data)
                observation = env.step(action)
                await websocket.send_json(
                    {
                        "type": "observation",
                        "data": self._serialize_observation(observation),
                    }
                )

            elif msg_type == "state":
                await websocket.send_json(
                    {"type": "state", "data": asdict(env.state)}
                )

            else:
                await websocket.send_json(
                    {
                        "type": "error",
                        "data": {
                            "message": f"Unknown message type: {msg_type}"
                        },
                    }
                )

    except WebSocketDisconnect:
        # Connection closed => cleanup environment without raising
        pass
    except Exception as e:
        # Send error and close connection
        try:
            await websocket.send_json(
                {
                    "type": "error",
                    "data": {"message": str(e)},
                }
            )
        except Exception:
            pass
    finally:
        # Cleanup
        self._active_envs.pop(websocket, None)

In practice, this would mean that legacy implementation would leave sessions open. IMO, this is ok because that essentially happens anyway with the current implementation.

Agree!

For context, @Wauplin is the maintainer of the hf library.

Forgot to present myself but yes indeed! 🤗 I'm external to this project but happy to be pinged from time to time :)

[rycerzes]

Thanks for the detailed insight and code snippet, @Wauplin! The WebSocket-based approach is really interesting.

I think this is the right direction. This approach would provide true session isolation where each connection equals one environment instance, which naturally handles session lifecycle. The automatic cleanup when connections drop is elegant, and the max_concurrent_envs parameter gives us fine-grained control over resource usage. This also opens the door for real-time features and makes the API more extensible with the message-based protocol.

While this is a significant breaking change, the abstraction at the client level should minimize user-facing impact.

A few implementation questions:

• Do we maintain the interface while using WebSockets under the hood to make migration smoother?
• For RL workloads that create/destroy many episodes, should we implement connection pooling to avoid WebSocket handshake overhead?
• And should we support both HTTP and WebSocket modes during a transition period?

I'm excited to work on this! I propose starting with a proof of concept following @Wauplin's snippet, then building a client compatibility layer to ensure existing HTTPEnvClient usage patterns still work, and finally testing with my poke-env use case to validate the approach.

@burtenshaw @Wauplin - Should I open a draft PR to start experimenting with this approach? I can begin with the server-side WebSocket implementation and we can iterate on the client abstractions.