From 895df9912162983127e2a26cdf0238ec9b6ac968 Mon Sep 17 00:00:00 2001 From: pankit-eng Date: Wed, 8 Oct 2025 13:29:55 -0400 Subject: [PATCH] basic skeleton for creating a client from docker image'd server --- src/__init__.py | 38 --- src/core/__init__.py | 8 +- src/core/base.py | 49 ---- src/core/base_env_client.py | 82 ------ src/core/docker/__init__.py | 11 - src/core/docker/docker_executor.py | 234 ------------------ src/core/http_env_client.py | 136 ++++++++++ src/core/local_docker.py | 124 ++++++++++ src/envs/__init__.py | 11 - .../coding_env/client/coding_env_client.py | 2 +- src/envs/coding_env/client/example_usage.py | 10 +- 11 files changed, 274 insertions(+), 431 deletions(-) delete mode 100644 src/core/base.py delete mode 100644 src/core/base_env_client.py delete mode 100644 src/core/docker/__init__.py delete mode 100644 src/core/docker/docker_executor.py create mode 100644 src/core/http_env_client.py create mode 100644 src/core/local_docker.py delete mode 100644 src/envs/__init__.py diff --git a/src/__init__.py b/src/__init__.py index 5493527247..47e552876c 100644 --- a/src/__init__.py +++ b/src/__init__.py @@ -5,41 +5,3 @@ # LICENSE file in the root directory of this source tree. """EnvTorch: Standardized agentic execution environments.""" - -# Core interfaces and types -from .core.env import ( - Environment, Transform, Tool, ToolRegistry, - Action, CodeAction, Observation, CodeObservation, - State, CodeState, ExecutionResult, - CompositeTransform, NullTransform, - CodeExecutionEnvironment -) - -# Docker execution -from .core.docker import DockerExecutor - -# Environment implementations -from .envs import CodingEnv - -__version__ = "0.1.0" - -__all__ = [ - # Core interfaces - "Environment", "Transform", "Tool", "ToolRegistry", - - # Types - "Action", "CodeAction", "Observation", "CodeObservation", - "State", "CodeState", "ExecutionResult", - - # Base transforms - "CompositeTransform", "NullTransform", - - # Base environment implementation - "CodeExecutionEnvironment", - - # Execution engines - "DockerExecutor", - - # Concrete environment implementations - "CodingEnv", -] diff --git a/src/core/__init__.py b/src/core/__init__.py index 841c9416e6..022ecbcca7 100644 --- a/src/core/__init__.py +++ b/src/core/__init__.py @@ -8,6 +8,12 @@ # Re-export main components from submodules for convenience from .env_server import * -from .docker import * +from .http_env_client import HTTPEnvClient +from .types import StepResult # Note: MCP module doesn't export anything yet + +__all__ = [ + "HTTPEnvClient", + "StepResult", +] diff --git a/src/core/base.py b/src/core/base.py deleted file mode 100644 index c9f684ebdb..0000000000 --- a/src/core/base.py +++ /dev/null @@ -1,49 +0,0 @@ -# Abstract base classes for EnvTorch -from __future__ import annotations - -from abc import ABC, abstractmethod -from typing import Generic, TypeVar - -from .types import StepResult - -# Generic type variables -ActT = TypeVar("ActT") # Type for the action sent to the environment -ObsT = TypeVar("ObsT") # Type for the observation returned by the environment - - -class BaseEnv(ABC, Generic[ActT, ObsT]): - """ - Abstract base class for all environments. - - Each environment must implement: - - reset(): to initialize or reinitialize environment state - - step(action): to execute an action and return results - - close(): to release resources (containers, sessions, etc.) - """ - - @abstractmethod - def reset(self) -> ObsT: - """ - Resets the environment to its initial state. - - Returns: - ObsT: The initial observation after resetting. - """ - raise NotImplementedError - - @abstractmethod - def step(self, action: ActT) -> StepResult[ObsT]: - """ - Performs one logical step in the environment. - - Args: - action (ActT): The action to perform in the environment. - - Returns: - StepResult[ObsT]: The resulting observation, reward, done flag, and info. - """ - raise NotImplementedError - - def close(self) -> None: - """Release resources (containers, sessions, etc.).""" - pass diff --git a/src/core/base_env_client.py b/src/core/base_env_client.py deleted file mode 100644 index f2cd6b0b25..0000000000 --- a/src/core/base_env_client.py +++ /dev/null @@ -1,82 +0,0 @@ -""" -core/runner_env.py -Minimal HTTP-based environment client. -- Talks to a single env worker exposing: POST /reset, POST /step - -Future hooks (commented below) for: -- episode_id, seed on reset -- request_id on step -- custom headers (auth/trace) -""" - -from __future__ import annotations - -from abc import abstractmethod -from typing import Any, Dict, Generic, Optional, TypeVar - -import requests - -from .base import BaseEnv -from .types import StepResult - -ActT = TypeVar("ActT") -ObsT = TypeVar("ObsT") - - -class HTTPEnvClient(BaseEnv[ActT, ObsT], Generic[ActT, ObsT]): - def __init__( - self, - base_url: str, - request_timeout_s: float = 15.0, - default_headers: Optional[Dict[str, str]] = None, - ): - self._base = base_url.rstrip("/") - self._timeout = float(request_timeout_s) - self._http = requests.Session() - self._headers = default_headers or {} - - @abstractmethod - def _step_payload(self, action: ActT) -> dict: - """Convert an Action object to the JSON body expected by the env server.""" - raise NotImplementedError - - @abstractmethod - def _parse_result(self, payload: dict) -> StepResult[ObsT]: - """Convert a JSON response from the env server to StepResult[ObsT].""" - raise NotImplementedError - - # ---------- BaseEnv ---------- - def reset(self) -> ObsT: - body: Dict[str, Any] = {} - # TODO: later: - # body["seed"] = seed - # body["episode_id"] = episode_id - r = self._http.post( - f"{self._base}/reset", - json=body, - headers=self._headers, - timeout=self._timeout, - ) - r.raise_for_status() - return self._parse_result(r.json()).observation - - def step(self, action: ActT) -> StepResult[ObsT]: - body: Dict[str, Any] = { - "action": self._step_payload(action), - "timeout_s": int(self._timeout), - } - # TODO: later: - # body["request_id"] = str(uuid.uuid4()) - # body["episode_id"] = current_episode_id - r = self._http.post( - f"{self._base}/step", - json=body, - headers=self._headers, - timeout=self._timeout, - ) - r.raise_for_status() - return self._parse_result(r.json()) - - def close(self) -> None: - # nothing to close; higher-level libraries own lifecycles of the endpoints - pass diff --git a/src/core/docker/__init__.py b/src/core/docker/__init__.py deleted file mode 100644 index 0864cb9162..0000000000 --- a/src/core/docker/__init__.py +++ /dev/null @@ -1,11 +0,0 @@ -# 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. - -"""Docker-based code execution.""" - -from .docker_executor import DockerExecutor - -__all__ = ["DockerExecutor"] \ No newline at end of file diff --git a/src/core/docker/docker_executor.py b/src/core/docker/docker_executor.py deleted file mode 100644 index fe06460b0a..0000000000 --- a/src/core/docker/docker_executor.py +++ /dev/null @@ -1,234 +0,0 @@ -# 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. - -import json -import subprocess -import time -from typing import Any - -from ..env.types import ExecutionResult - - -class DockerExecutor: - """Simple Docker-based Python code executor with persistent session.""" - - def __init__(self, image: str = "python:3.11-slim", timeout_seconds: int = 30): - self.image = image - self.timeout_seconds = timeout_seconds - self.container_id: str | None = None - self._process: subprocess.Popen | None = None - - def start_session(self) -> None: - """Start new Docker container with persistent Python session.""" - if self.container_id: - self.stop_session() - - # Run interactive Python in container - cmd = [ - "docker", "run", "--rm", "-i", - "--memory=512m", - "--cpus=1.0", - "--network=host", # For MCP integration later - self.image, - "python", "-u", "-c", - "import sys; [exec(input()) for _ in iter(int, 1)]" - ] - - try: - self._process = subprocess.Popen( - cmd, - stdin=subprocess.PIPE, - stdout=subprocess.PIPE, - stderr=subprocess.PIPE, - text=True, - bufsize=0 - ) - - # Get container ID for potential cleanup - # This is a bit hacky but works for the simple approach - time.sleep(0.1) # Give container time to start - - except Exception as e: - raise RuntimeError(f"Failed to start Docker container: {e}") - - def execute_code(self, code: str) -> ExecutionResult: - """Send code to running container and get results.""" - if not self._process: - raise RuntimeError("No active session. Call start_session() first.") - - start_time = time.time() - - try: - # Wrap user code to capture output and results - wrapped_code = self._wrap_code_for_execution(code) - - # Send code to container - self._process.stdin.write(wrapped_code + "\n") - self._process.stdin.flush() - - # Read result with timeout - stdout_lines = [] - stderr_lines = [] - - # Simple timeout mechanism - not perfect but works for MVP - end_time = start_time + self.timeout_seconds - result_captured = False - - while time.time() < end_time and not result_captured: - if self._process.poll() is not None: - # Process died - break - - # Try to read a line with short timeout - try: - # This is simplified - in production we'd use select/threading - self._process.stdout.settimeout(0.1) - line = self._process.stdout.readline() - if line: - if line.startswith("__ENVTORCH_RESULT__"): - result_captured = True - result_json = line[len("__ENVTORCH_RESULT__"):].strip() - break - else: - stdout_lines.append(line.rstrip()) - except: - time.sleep(0.01) - - if not result_captured: - return ExecutionResult.from_exception( - TimeoutError(f"Code execution timed out after {self.timeout_seconds}s"), - stdout="\n".join(stdout_lines), - stderr="\n".join(stderr_lines) - ) - - # Parse result - try: - result_data = json.loads(result_json) - execution_time_ms = (time.time() - start_time) * 1000 - - if result_data.get("success", True): - return ExecutionResult.from_success( - return_value=result_data.get("return_value"), - stdout=result_data.get("stdout", ""), - stderr=result_data.get("stderr", ""), - execution_time_ms=execution_time_ms - ) - else: - # Reconstruct exception from data - exc_type = result_data.get("exception_type", "Exception") - exc_message = result_data.get("exception_message", "") - - # Create a generic exception for now - exc = Exception(f"{exc_type}: {exc_message}") - result = ExecutionResult.from_exception( - exc, - stdout=result_data.get("stdout", ""), - stderr=result_data.get("stderr", "") - ) - result.traceback_str = result_data.get("traceback", "") - result.execution_time_ms = execution_time_ms - return result - - except json.JSONDecodeError as e: - return ExecutionResult.from_exception( - RuntimeError(f"Failed to parse execution result: {e}"), - stdout="\n".join(stdout_lines) - ) - - except Exception as e: - return ExecutionResult.from_exception(e) - - def get_variable_dump(self) -> dict[str, Any]: - """Get all variables for render() - send globals() inspection command.""" - if not self._process: - raise RuntimeError("No active session. Call start_session() first.") - - dump_code = ''' -import json -result = {} -for name, value in globals().items(): - if not name.startswith('_') and name not in ['json']: - try: - # Try to get a readable representation - if hasattr(value, '__dict__') and not callable(value): - result[name] = f"<{type(value).__name__}: {str(value)[:100]}>" - else: - result[name] = repr(value)[:200] # Limit length - except: - result[name] = f"<{type(value).__name__} object>" -print("__ENVTORCH_DUMP__" + json.dumps(result)) -''' - - # Execute the dump code - exec_result = self.execute_code(dump_code) - if not exec_result.success: - return {"error": "Failed to dump variables", "details": exec_result.exception_message} - - # Extract dump from stdout - for line in exec_result.stdout.split('\n'): - if line.startswith("__ENVTORCH_DUMP__"): - try: - return json.loads(line[len("__ENVTORCH_DUMP__"):]) - except json.JSONDecodeError: - pass - - return {"error": "No variable dump found in output"} - - def stop_session(self) -> None: - """Kill the container process.""" - if self._process: - try: - self._process.terminate() - self._process.wait(timeout=5) - except subprocess.TimeoutExpired: - self._process.kill() - self._process.wait() - finally: - self._process = None - self.container_id = None - - def _wrap_code_for_execution(self, code: str) -> str: - """Wrap user code to capture results and exceptions.""" - return f''' -import sys -import json -import traceback -from io import StringIO - -# Capture stdout/stderr -old_stdout = sys.stdout -old_stderr = sys.stderr -stdout_capture = StringIO() -stderr_capture = StringIO() -sys.stdout = stdout_capture -sys.stderr = stderr_capture - -result = {{"success": True, "return_value": None, "stdout": "", "stderr": ""}} - -try: - # Execute user code - exec_result = None - exec("""{code}""") - result["return_value"] = exec_result if 'exec_result' in locals() else None -except Exception as e: - result["success"] = False - result["exception_type"] = e.__class__.__name__ - result["exception_message"] = str(e) - result["traceback"] = traceback.format_exc() -finally: - # Restore stdout/stderr and capture output - sys.stdout = old_stdout - sys.stderr = old_stderr - result["stdout"] = stdout_capture.getvalue() - result["stderr"] = stderr_capture.getvalue() - -# Send result back -print("__ENVTORCH_RESULT__" + json.dumps(result, default=str)) -''' - - def __del__(self): - """Cleanup on destruction.""" - self.stop_session() \ No newline at end of file diff --git a/src/core/http_env_client.py b/src/core/http_env_client.py new file mode 100644 index 0000000000..11eef854f6 --- /dev/null +++ b/src/core/http_env_client.py @@ -0,0 +1,136 @@ +""" +core/runner_env.py +Minimal HTTP-based environment client. +- Talks to a single env worker exposing: POST /reset, POST /step + +Future hooks (commented below) for: +- episode_id, seed on reset +- request_id on step +- custom headers (auth/trace) +""" + +from __future__ import annotations + +from abc import ABC, abstractmethod +from typing import Any, Dict, Generic, Optional, Type, TypeVar + +import requests +from . import local_docker +from .types import StepResult + +ActT = TypeVar("ActT") +ObsT = TypeVar("ObsT") +EnvClientT = TypeVar("EnvClientT", bound="HTTPEnvClient") + + +class HTTPEnvClient(ABC, Generic[ActT, ObsT]): + def __init__( + self, + base_url: str, + request_timeout_s: float = 15.0, + default_headers: Optional[Dict[str, str]] = None, + ): + self._base = base_url.rstrip("/") + self._timeout = float(request_timeout_s) + self._http = requests.Session() + self._headers = default_headers or {} + + @classmethod + def from_docker_image(cls: Type[EnvClientT], image: str) -> EnvClientT: + """ + Create an environment client by spinning up a Docker container locally. + + This is a development utility that: + 1. Starts a Docker container from the specified image + 2. Waits for the server to be ready + 3. Creates and returns a client instance connected to the container + + Note: The container lifecycle management is left to the user or higher-level + orchestration. The container will keep running until manually stopped. + + Args: + image: Docker image name to run (e.g., "coding-env:latest") + + Returns: + An instance of the client class connected to the running container + + Example: + >>> from envs.coding_env.client import CodingEnv + >>> from envs.coding_env.models import CodeAction + >>> + >>> # Create environment from image + >>> env = CodingEnv.from_docker_image("coding-env:latest") + >>> + >>> # Use the environment + >>> result = env.reset() + >>> print(result.observation) + >>> + >>> step_result = env.step(CodeAction(code="print('hello')")) + >>> print(step_result.observation.stdout) + >>> + >>> # Cleanup (optional) + >>> env.close() + """ + + + # 1. Find available port + port = local_docker._find_available_port() + + # 2. Generate unique container name + container_name = local_docker._generate_container_name(image) + + # 3. Start Docker container + # container = local_docker._start_container(image=image, name=container_name, port=port) + + # 4. Wait for server to be ready + # local_docker._wait_for_server_ready(f"http://localhost:{port}") + + # 5. Create and return client instance + base_url = f"http://localhost:{port}" + return cls(base_url=base_url) + + @abstractmethod + def _step_payload(self, action: ActT) -> dict: + """Convert an Action object to the JSON body expected by the env server.""" + raise NotImplementedError + + @abstractmethod + def _parse_result(self, payload: dict) -> StepResult[ObsT]: + """Convert a JSON response from the env server to StepResult[ObsT].""" + raise NotImplementedError + + # ---------- Environment Server Interface Methods ---------- + def reset(self) -> StepResult[ObsT]: + body: Dict[str, Any] = {} + # TODO: later: + # body["seed"] = seed + # body["episode_id"] = episode_id + r = self._http.post( + f"{self._base}/reset", + json=body, + headers=self._headers, + timeout=self._timeout, + ) + r.raise_for_status() + return self._parse_result(r.json()) + + def step(self, action: ActT) -> StepResult[ObsT]: + body: Dict[str, Any] = { + "action": self._step_payload(action), + "timeout_s": int(self._timeout), + } + # TODO: later: + # body["request_id"] = str(uuid.uuid4()) + # body["episode_id"] = current_episode_id + r = self._http.post( + f"{self._base}/step", + json=body, + headers=self._headers, + timeout=self._timeout, + ) + r.raise_for_status() + return self._parse_result(r.json()) + + def close(self) -> None: + # nothing to close; higher-level libraries own lifecycles of the endpoints + pass diff --git a/src/core/local_docker.py b/src/core/local_docker.py new file mode 100644 index 0000000000..611cedba54 --- /dev/null +++ b/src/core/local_docker.py @@ -0,0 +1,124 @@ +# 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. + +""" +Docker utilities for running environment servers locally. + +This module provides low-level Docker helpers for container management, +used by HTTPEnvClient.from_docker_image() and subclasses. +""" + +from __future__ import annotations + +from typing import Any + + +def _find_available_port() -> int: + """ + Find an available port on localhost. + + Returns: + An available port number + """ + # TODO: Implement port finding logic + # Can use socket to find an available port + import socket + + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s: + s.bind(("", 0)) + s.listen(1) + port = s.getsockname()[1] + return port + + +def _generate_container_name(image: str) -> str: + """ + Generate a unique container name based on image name and timestamp. + + Args: + image: Docker image name + + Returns: + A unique container name + """ + # TODO: Implement container name generation + # Extract image name without tag/registry + # Add timestamp or UUID for uniqueness + import time + + clean_image = image.split("/")[-1].split(":")[0] + timestamp = int(time.time() * 1000) + return f"{clean_image}-{timestamp}" + + +def _start_container(image: str, name: str, port: int) -> Any: + """ + Start a Docker container from the specified image. + + Args: + image: Docker image name + name: Container name + port: Port to expose + + Returns: + Container handle/reference for cleanup + """ + # TODO: Implement using docker SDK or subprocess + # docker_client = docker.from_env() + # container = docker_client.containers.run( + # image, + # name=name, + # ports={"8000/tcp": port}, # Assuming server runs on 8000 internally + # detach=True, + # remove=True, # Auto-remove when stopped + # ) + # return container + raise NotImplementedError("Docker container start logic not yet implemented") + + +def _wait_for_server_ready(base_url: str) -> None: + """ + Wait for the environment server to be ready to accept requests. + + Args: + base_url: Base URL of the server + + Raises: + TimeoutError: If server doesn't become ready within timeout + """ + # TODO: Implement health check polling + # import time + # import requests + # + # timeout_s = 30.0 + # start_time = time.time() + # while time.time() - start_time < timeout_s: + # try: + # response = requests.get(f"{base_url}/health", timeout=1.0) + # if response.status_code == 200: + # return + # except requests.RequestException: + # pass + # time.sleep(0.5) + # raise TimeoutError(f"Server at {base_url} did not become ready within {timeout_s}s") + raise NotImplementedError("Server health check logic not yet implemented") + + +def _stop_container(container: Any) -> None: + """ + Stop and remove a Docker container. + + Args: + container: Container handle from _start_container + """ + # TODO: Implement container cleanup + # try: + # container.stop(timeout=5) + # container.remove() + # except Exception as e: + # # Log error but don't fail + # print(f"Warning: Failed to clean up container: {e}") + raise NotImplementedError("Docker container cleanup logic not yet implemented") diff --git a/src/envs/__init__.py b/src/envs/__init__.py deleted file mode 100644 index ebcfab16f8..0000000000 --- a/src/envs/__init__.py +++ /dev/null @@ -1,11 +0,0 @@ -# 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. - -"""This module encapsulates all environment implementations.""" - -from .coding import CodingEnv - -__all__ = ["CodingEnv"] diff --git a/src/envs/coding_env/client/coding_env_client.py b/src/envs/coding_env/client/coding_env_client.py index 513f95eb02..8386428949 100644 --- a/src/envs/coding_env/client/coding_env_client.py +++ b/src/envs/coding_env/client/coding_env_client.py @@ -15,7 +15,7 @@ from typing import Optional -from core.base_env_client import HTTPEnvClient +from core.http_env_client import HTTPEnvClient from core.types import StepResult from ..models import CodeAction, CodeObservation diff --git a/src/envs/coding_env/client/example_usage.py b/src/envs/coding_env/client/example_usage.py index d99bb96da4..de2c3d51b0 100644 --- a/src/envs/coding_env/client/example_usage.py +++ b/src/envs/coding_env/client/example_usage.py @@ -2,7 +2,9 @@ from .coding_env_client import CodingEnv -env = CodingEnv(base_url="http://localhost:8080") -obs0 = env.reset() -result = env.step(CodeAction(code="print('hi')")) -print(result.observation.stdout.strip(), result.reward) +env = CodingEnv.from_docker_image(image="coding_env:latest") +reset_result = env.reset() +print(f"Initial observation: {reset_result.observation}") + +step_result = env.step(CodeAction(code="print('hi')")) +print(step_result.observation.stdout.strip(), step_result.reward)