Use context manager for per-operation websocket connection in DockerExecutor by albertvillanova · Pull Request #1750 · huggingface/smolagents (original) (raw)

Use context manager for per-operation websocket connection in DockerExecutor.

Refactor DockerExecutor to use per-operation websocket connections like ModalExecutor.

The current DockerExecutor maintains a persistent websocket connection that is created during initialization and stored as an instance variable. This approach has several drawbacks:

• Requires manual cleanup logic in the cleanup() method
• More error-prone due to connection state management
• Inconsistent with ModalExecutor pattern
• Connection can become stale over time

This PR refactors DockerExecutor to follow the same pattern as ModalExecutor by creating websocket connections per operation using context managers:

1. Removed persistent websocket connection: No longer store self.ws during initialization
2. Per-operation connections: Create websocket connection in run_code_raise_errors() method
3. Context manager usage: Use contextlib.closing() to ensure automatic resource cleanup
4. Simplified cleanup: Removed websocket cleanup from cleanup() method