diff --git a/examples/auto_env_example.py b/examples/auto_env_example.py new file mode 100755 index 00000000..690e5277 --- /dev/null +++ b/examples/auto_env_example.py @@ -0,0 +1,320 @@ +#!/usr/bin/env python3 +# 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. + +""" +Comprehensive AutoEnv and AutoAction Example +============================================= + +This example demonstrates how to use the AutoEnv and AutoAction classes +to automatically select and use environments without manual imports. + +The AutoEnv/AutoAction API follows the HuggingFace pattern, making it easy +to work with different environments using a consistent interface. + +Run this example with: + python examples/auto_env_example.py + +Or test a specific environment: + python examples/auto_env_example.py --env coding +""" + +import sys +import argparse +from pathlib import Path + +# Add src to path +sys.path.insert(0, str(Path(__file__).parent.parent / "src")) + +from envs import AutoEnv, AutoAction + + +def example_basic_usage(): + """Example 1: Basic usage with AutoEnv and AutoAction""" + print("=" * 70) + print("Example 1: Basic Usage") + print("=" * 70) + print() + + # Instead of: + # from envs.coding_env import CodingEnv, CodeAction + # client = CodingEnv.from_docker_image("coding-env:latest") + + # You can now do: + print("Creating environment using AutoEnv...") + client = AutoEnv.from_docker_image("coding-env:latest") + print("✓ Environment created!") + print() + + # Get the Action class automatically + print("Getting Action class using AutoAction...") + CodeAction = AutoAction.from_image("coding-env:latest") + print(f"✓ Got Action class: {CodeAction.__name__}") + print() + + # Use them together + print("Testing the environment:") + result = client.reset() + print(f" Reset: exit_code={result.observation.exit_code}") + + action = CodeAction(code="print('Hello from AutoEnv!')") + step_result = client.step(action) + print(f" Step result: {step_result.observation.stdout.strip()}") + + client.close() + print("✓ Environment closed") + print() + + +def example_alternative_syntax(): + """Example 2: Alternative syntax using from_env()""" + print("=" * 70) + print("Example 2: Alternative Syntax") + print("=" * 70) + print() + + # You can also use environment names directly + print("Getting Action class by environment name...") + CodeAction = AutoAction.from_env("coding") + print(f"✓ Got Action class: {CodeAction.__name__}") + print() + + # Create instance + action = CodeAction(code="x = 5 + 3\nprint(f'Result: {x}')") + print(f"Created action: {action}") + print() + + +def example_list_environments(): + """Example 3: List all available environments""" + print("=" * 70) + print("Example 3: List Available Environments") + print("=" * 70) + print() + + # List all available environments + AutoEnv.list_environments() + print() + + +def example_list_actions(): + """Example 4: List all available action classes""" + print("=" * 70) + print("Example 4: List Available Action Classes") + print("=" * 70) + print() + + # List all available action classes + AutoAction.list_actions() + print() + + +def example_environment_info(): + """Example 5: Get detailed environment information""" + print("=" * 70) + print("Example 5: Environment Information") + print("=" * 70) + print() + + # Get detailed info about a specific environment + env_name = "coding" + print(f"Information about '{env_name}' environment:") + print("-" * 70) + + info = AutoEnv.get_env_info(env_name) + print(f" Description: {info['description']}") + print(f" Docker Image: {info['default_image']}") + print(f" Environment Class: {info['env_class']}") + print(f" Action Class: {info['action_class']}") + print(f" Special Requirements: {info['special_requirements'] or 'None'}") + print() + + print(" Supported Features:") + for feature in info["supported_features"]: + print(f" - {feature}") + print() + + +def example_error_handling(): + """Example 6: Error handling with helpful messages""" + print("=" * 70) + print("Example 6: Error Handling") + print("=" * 70) + print() + + # Try an unknown environment + print("Trying unknown environment 'nonexistent'...") + try: + env = AutoEnv.from_docker_image("nonexistent-env:latest") + except ValueError as e: + print(f"✓ Got expected error: {e}") + print() + + # Try a typo - should suggest similar names + print("Trying typo 'cooding' (should suggest 'coding')...") + try: + env = AutoEnv.from_docker_image("cooding-env:latest") + except ValueError as e: + print(f"✓ Got helpful suggestion: {e}") + print() + + # Try deprecated julia environment + print("Trying deprecated 'julia' environment...") + try: + env = AutoEnv.from_docker_image("julia-env:latest") + except ValueError as e: + print(f"✓ Got deprecation notice: {e}") + print() + + +def example_special_requirements(): + """Example 7: Environments with special requirements""" + print("=" * 70) + print("Example 7: Special Requirements") + print("=" * 70) + print() + + # DIPG environment requires dataset path + print("DIPG environment requires DIPG_DATASET_PATH:") + print() + print(" # This would show a warning:") + print(" # env = AutoEnv.from_docker_image('dipg-env:latest')") + print() + print(" # Correct usage:") + print(" env = AutoEnv.from_docker_image(") + print(" 'dipg-env:latest',") + print(" env_vars={'DIPG_DATASET_PATH': '/data/dipg'}") + print(" )") + print() + + # FinRL environment has optional config + print("FinRL environment accepts optional config:") + print() + print(" env = AutoEnv.from_docker_image(") + print(" 'finrl-env:latest',") + print(" env_vars={'FINRL_CONFIG_PATH': '/config.json'}") + print(" )") + print() + + +def test_specific_environment(env_name: str): + """Test a specific environment by name""" + print("=" * 70) + print(f"Testing {env_name} Environment") + print("=" * 70) + print() + + try: + # Get environment info + info = AutoEnv.get_env_info(env_name) + image = info["default_image"] + + print(f"Creating {env_name} environment...") + print(f" Docker image: {image}") + print() + + # Create environment with extended timeout for slow containers + env = AutoEnv.from_docker_image(image, wait_timeout=60.0) + print("✓ Environment created!") + + # Get action class + ActionClass = AutoAction.from_env(env_name) + print(f"✓ Action class: {ActionClass.__name__}") + print() + + # Test reset + print("Testing reset()...") + result = env.reset() + print(f"✓ Reset successful") + print() + + # Get state + state = env.state() + print(f"State: episode_id={state.episode_id}, step_count={state.step_count}") + print() + + # Close + env.close() + print("✓ Environment closed") + print() + + print("=" * 70) + print(f"✓ {env_name} environment test passed!") + print("=" * 70) + + return True + + except Exception as e: + print(f"\n❌ Error testing {env_name}: {e}") + import traceback + + traceback.print_exc() + return False + + +def main(): + """Main function to run examples""" + parser = argparse.ArgumentParser( + description="AutoEnv and AutoAction Examples", + formatter_class=argparse.RawDescriptionHelpFormatter, + ) + parser.add_argument( + "--env", + type=str, + help="Test a specific environment (e.g., coding, echo, git)", + ) + parser.add_argument( + "--all-examples", + action="store_true", + help="Run all examples (without Docker)", + ) + + args = parser.parse_args() + + if args.env: + # Test specific environment + success = test_specific_environment(args.env) + sys.exit(0 if success else 1) + + elif args.all_examples: + # Run all examples (no Docker needed) + example_basic_usage() # This requires Docker + # Skip Docker examples, run info-only examples + example_alternative_syntax() + example_list_environments() + example_list_actions() + example_environment_info() + example_error_handling() + example_special_requirements() + + else: + # Show usage info and examples that don't need Docker + print("AutoEnv and AutoAction Examples") + print("=" * 70) + print() + print("This demonstrates the HuggingFace-style API for OpenEnv.") + print() + print("Usage:") + print(" python examples/auto_env_example.py --all-examples") + print(" python examples/auto_env_example.py --env coding") + print() + print("Running info examples (no Docker required)...") + print() + + example_list_environments() + example_list_actions() + example_environment_info() + example_error_handling() + example_special_requirements() + + print() + print("To test with actual Docker environments:") + print(" python examples/auto_env_example.py --env coding") + print() + + +if __name__ == "__main__": + main() diff --git a/src/core/containers/runtime/providers.py b/src/core/containers/runtime/providers.py index a8022ddc..3b9703d5 100644 --- a/src/core/containers/runtime/providers.py +++ b/src/core/containers/runtime/providers.py @@ -118,7 +118,11 @@ def __init__(self): capture_output=True, timeout=5, ) - except (subprocess.CalledProcessError, FileNotFoundError, subprocess.TimeoutExpired): + except ( + subprocess.CalledProcessError, + FileNotFoundError, + subprocess.TimeoutExpired, + ): raise RuntimeError( "Docker is not available. Please install Docker Desktop or Docker Engine." ) @@ -138,26 +142,44 @@ def start_container( port: Port to expose (if None, finds available port) env_vars: Environment variables for the container **kwargs: Additional Docker run options + - memory_gb: Memory limit in GB (default: 4GB) + - command_override: List of command args to override container CMD Returns: Base URL to connect to the container """ import subprocess import time + import logging + + logger = logging.getLogger(__name__) # Find available port if not specified if port is None: port = self._find_available_port() + # Use default memory limit if not specified + memory_gb = kwargs.get("memory_gb", 16) + # Generate container name self._container_name = self._generate_container_name(image) # Build docker run command + # Use host networking for better performance and consistency with podman + # NOTE: Do NOT use --rm initially - if container fails to start, we need logs cmd = [ - "docker", "run", + "docker", + "run", "-d", # Detached - "--name", self._container_name, - "-p", f"{port}:8000", # Map port + "--name", + self._container_name, + "--network", + "host", # Use host network + "--memory", + f"{memory_gb}g", # Limit container memory + "--memory-swap", + f"{memory_gb}g", # Prevent swap usage (set equal to --memory) + "--oom-kill-disable=false", # Allow OOM killer (exit gracefully) ] # Add environment variables @@ -165,13 +187,24 @@ def start_container( for key, value in env_vars.items(): cmd.extend(["-e", f"{key}={value}"]) + # Pass custom port via environment variable instead of overriding command + # This allows the container to use its proper entrypoint/CMD + if port != 8000: + cmd.extend(["-e", f"PORT={port}"]) + # Add image cmd.append(image) + # Add command override if provided (explicit override by user) + if "command_override" in kwargs: + cmd.extend(kwargs["command_override"]) + # Run container try: + logger.debug(f"Starting container with command: {' '.join(cmd)}") result = subprocess.run(cmd, capture_output=True, text=True, check=True) self._container_id = result.stdout.strip() + logger.debug(f"Container started with ID: {self._container_id}") except subprocess.CalledProcessError as e: error_msg = f"Failed to start Docker container.\nCommand: {' '.join(cmd)}\nExit code: {e.returncode}\nStderr: {e.stderr}\nStdout: {e.stdout}" raise RuntimeError(error_msg) from e @@ -192,24 +225,47 @@ def stop_container(self) -> None: import subprocess try: - # Stop container - subprocess.run( - ["docker", "stop", self._container_id], - capture_output=True, - check=True, - timeout=10, - ) + # Try graceful stop first (with longer timeout) + print(f"Stopping container {self._container_id[:12]}...") + try: + subprocess.run( + ["docker", "stop", "-t", "5", self._container_id], + capture_output=True, + timeout=30, + ) + except subprocess.TimeoutExpired: + # If graceful stop times out, force kill + print(f"Graceful stop timed out, forcing kill...") + subprocess.run( + ["docker", "kill", self._container_id], + capture_output=True, + timeout=10, + ) # Remove container + print(f"Removing container {self._container_id[:12]}...") subprocess.run( - ["docker", "rm", self._container_id], + ["docker", "rm", "-f", self._container_id], capture_output=True, - check=True, - timeout=10, + timeout=15, ) - except subprocess.CalledProcessError: - # Container might already be stopped/removed - pass + + print(f"✓ Container cleaned up successfully") + + except subprocess.TimeoutExpired: + # Last resort: force remove + print(f"Remove timed out, trying force remove...") + try: + subprocess.run( + ["docker", "rm", "-f", self._container_id], + capture_output=True, + timeout=10, + ) + except Exception: + pass + except Exception as e: + # Log but don't fail - container might already be gone + print(f"Note: Cleanup had issues (container may already be removed): {e}") finally: self._container_id = None self._container_name = None @@ -241,8 +297,28 @@ def wait_for_ready(self, base_url: str, timeout_s: float = 30.0) -> None: time.sleep(0.5) + # Get container logs for debugging + logs_snippet = "" + if self._container_id: + try: + import subprocess + + result = subprocess.run( + ["docker", "logs", "--tail", "20", self._container_id], + capture_output=True, + text=True, + timeout=5, + ) + if result.stdout or result.stderr: + logs_snippet = "\n\nContainer logs (last 20 lines):\n" + logs_snippet += result.stdout + result.stderr + except Exception: + pass + raise TimeoutError( - f"Container at {base_url} did not become ready within {timeout_s}s" + f"Container at {base_url} did not become ready within {timeout_s}s. " + f"The container is still running and will be cleaned up automatically. " + f"Try increasing wait_timeout (e.g., wait_timeout=60.0 or higher).{logs_snippet}" ) def _find_available_port(self) -> int: @@ -290,4 +366,5 @@ class KubernetesProvider(ContainerProvider): >>> # Pod running in k8s, accessible via service or port-forward >>> provider.stop_container() """ + pass diff --git a/src/core/http_env_client.py b/src/core/http_env_client.py index 16bbfa5d..f8e815b9 100644 --- a/src/core/http_env_client.py +++ b/src/core/http_env_client.py @@ -46,6 +46,7 @@ def from_docker_image( cls: Type[EnvClientT], image: str, provider: Optional["ContainerProvider"] = None, + wait_timeout: float = 30.0, **kwargs: Any, ) -> EnvClientT: """ @@ -62,6 +63,7 @@ def from_docker_image( Args: image: Docker image name to run (e.g., "echo-env:latest") provider: Container provider to use (defaults to LocalDockerProvider) + wait_timeout: Maximum time (in seconds) to wait for container to be ready (default: 30.0) **kwargs: Additional arguments to pass to provider.start_container() (e.g., env_vars, port) @@ -81,6 +83,12 @@ def from_docker_image( ... env_vars={"MY_VAR": "value"} ... ) >>> + >>> # Create with custom wait timeout (useful for slow containers) + >>> env = CodingEnv.from_docker_image( + ... "coding-env:latest", + ... wait_timeout=60.0 # Wait up to 60 seconds + ... ) + >>> >>> # Use the environment >>> result = env.reset() >>> print(result.observation) @@ -99,28 +107,41 @@ def from_docker_image( # 1. Start container with optional kwargs (e.g., env_vars, port) base_url = provider.start_container(image, **kwargs) - # 2. Wait for server to be ready - provider.wait_for_ready(base_url) + # 2. Wait for server to be ready with custom timeout + try: + provider.wait_for_ready(base_url, timeout_s=wait_timeout) + except TimeoutError: + # Cleanup: stop and remove the container if it didn't become ready + print( + f"Container failed to become ready within {wait_timeout}s. Cleaning up..." + ) + provider.stop_container() + raise # 3. Create and return client instance with provider reference return cls(base_url=base_url, provider=provider) @classmethod - def from_hub(cls: Type[EnvClientT], repo_id: str, provider: Optional["ContainerProvider"] = None, **kwargs: Any) -> EnvClientT: + def from_hub( + cls: Type[EnvClientT], + repo_id: str, + provider: Optional["ContainerProvider"] = None, + **kwargs: Any, + ) -> EnvClientT: """ Create an environment client by pulling from a Hugging Face model hub. """ - + if provider is None: provider = LocalDockerProvider() - + if "tag" in kwargs: tag = kwargs["tag"] else: tag = "latest" - + base_url = f"registry.hf.space/{repo_id.replace('/', '-')}:{tag}" - + return cls.from_docker_image(image=base_url, provider=provider) @abstractmethod diff --git a/src/envs/__init__.py b/src/envs/__init__.py new file mode 100644 index 00000000..293453b0 --- /dev/null +++ b/src/envs/__init__.py @@ -0,0 +1,62 @@ +# 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. + +""" +OpenEnv Environments +==================== + +This package contains all environment implementations for OpenEnv. + +Each environment provides: +- An environment client class (e.g., CodingEnv, AtariEnv) +- Action and Observation data classes +- Server implementations for the HTTP API + +Auto Classes +------------ +The AutoEnv and AutoAction classes provide a HuggingFace-style API for +automatically selecting the correct environment and action types based on +Docker image names. + +Example: + >>> from envs import AutoEnv, AutoAction + >>> + >>> # Automatically detect and create environment from image + >>> client = AutoEnv.from_docker_image("coding-env:latest") + >>> + >>> # Get the corresponding Action class + >>> CodeAction = AutoAction.from_image("coding-env:latest") + >>> + >>> # Use them together + >>> result = client.reset() + >>> action = CodeAction(code="print('Hello, AutoEnv!')") + >>> step_result = client.step(action) + >>> client.close() + +Direct Imports +-------------- +You can also import specific environment classes directly: + + >>> from envs.coding_env import CodingEnv, CodeAction + >>> from envs.echo_env import EchoEnv, EchoAction + >>> from envs.git_env import GitEnv, GitAction + >>> # ... etc + +List Available Environments +--------------------------- +To see all available environments: + + >>> AutoEnv.list_environments() + >>> AutoAction.list_actions() +""" + +from .auto_env import AutoEnv +from .auto_action import AutoAction + +__all__ = [ + "AutoEnv", + "AutoAction", +] diff --git a/src/envs/_registry.py b/src/envs/_registry.py new file mode 100644 index 00000000..dc4d7c0f --- /dev/null +++ b/src/envs/_registry.py @@ -0,0 +1,241 @@ +# 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. + +""" +Environment Registry for AutoEnv and AutoAction +================================================ + +This module provides a centralized registry mapping environment names to +their corresponding client classes, action classes, and default Docker +image names. + +The registry enables the AutoEnv and AutoAction classes to automatically +instantiate the correct environment and action types based on Docker +image names. +""" + +from typing import Any, Dict + +# Registry structure: +# env_key: (module_path, env_class_name, action_class_name, +# default_image, special_notes) +ENV_REGISTRY: Dict[str, Dict[str, Any]] = { + "atari": { + "module": "envs.atari_env", + "env_class": "AtariEnv", + "action_class": "AtariAction", + "default_image": "atari-env:latest", + "description": "Atari 2600 games environment (100+ games)", + "special_requirements": None, + "supported_features": [ + "Multiple games (100+)", + "RGB/grayscale/RAM observations", + "Configurable action spaces (minimal/full)", + "Frame skipping and sticky actions", + ], + }, + "browsergym": { + "module": "envs.browsergym_env", + "env_class": "BrowserGymEnv", + "action_class": "BrowserGymAction", + "default_image": "browsergym-env:latest", + "description": "Web browsing environment with multiple benchmarks", + "special_requirements": "WebArena tasks require backend setup with env vars", + "supported_features": [ + "MiniWoB/WebArena/VisualWebArena benchmarks", + "Natural language actions", + "Multi-modal observations (text/visual)", + ], + }, + "chat": { + "module": "envs.chat_env", + "env_class": "ChatEnv", + "action_class": "ChatAction", + "default_image": "chat-env:latest", + "description": "Chat environment with tokenization support", + "special_requirements": None, + "supported_features": [ + "PyTorch tensor handling", + "Hugging Face chat format", + "Optional tokenization with TOKENIZER_NAME env var", + ], + }, + "coding": { + "module": "envs.coding_env", + "env_class": "CodingEnv", + "action_class": "CodeAction", + "default_image": "coding-env:latest", + "description": "Python code execution environment", + "special_requirements": None, + "supported_features": [ + "Python code execution", + "Persistent execution context", + "stdout/stderr/exit_code capture", + ], + }, + "connect4": { + "module": "envs.connect4_env", + "env_class": "Connect4Env", + "action_class": "Connect4Action", + "default_image": "connect4-env:latest", + "description": "Connect Four board game environment", + "special_requirements": None, + "supported_features": [ + "Two-player game (6x7 grid)", + "Legal actions masking", + "Turn tracking", + ], + }, + "dipg": { + "module": "envs.dipg_safety_env", + "env_class": "DIPGSafetyEnv", + "action_class": "DIPGAction", + "default_image": "dipg-env:latest", + "description": "DIPG safety-critical medical decision environment", + "special_requirements": "Requires DIPG_DATASET_PATH env var pointing to dataset", + "supported_features": [ + "Safety-critical medical domain", + "LLM response scoring", + "Conflict/abstention rewards", + ], + }, + "echo": { + "module": "envs.echo_env", + "env_class": "EchoEnv", + "action_class": "EchoAction", + "default_image": "echo-env:latest", + "description": "Simple echo test environment", + "special_requirements": None, + "supported_features": [ + "Message echoing", + "Basic HTTP server testing", + ], + }, + "finrl": { + "module": "envs.finrl_env", + "env_class": "FinRLEnv", + "action_class": "FinRLAction", + "default_image": "finrl-env:latest", + "description": "Financial trading environment", + "special_requirements": "Optional FINRL_CONFIG_PATH env var for custom configuration", + "supported_features": [ + "Stock trading simulation", + "Technical indicators", + "Custom configuration support", + ], + }, + "git": { + "module": "envs.git_env", + "env_class": "GitEnv", + "action_class": "GitAction", + "default_image": "git-env:latest", + "description": "Git repository management with Gitea integration", + "special_requirements": None, + "supported_features": [ + "Repository cloning", + "Git command execution", + "Gitea server integration", + ], + }, + "openspiel": { + "module": "envs.openspiel_env", + "env_class": "OpenSpielEnv", + "action_class": "OpenSpielAction", + "default_image": "openspiel-env:latest", + "description": "OpenSpiel game environment (multiple games)", + "special_requirements": None, + "supported_features": [ + "6 supported games (catch/tic-tac-toe/kuhn_poker/cliff_walking/2048/blackjack)", + "Single and multi-player support", + "Optional opponent policies", + ], + }, + "sumo_rl": { + "module": "envs.sumo_rl_env", + "env_class": "SumoRLEnv", + "action_class": "SumoAction", + "default_image": "sumo-rl-env:latest", + "description": "SUMO traffic signal control environment", + "special_requirements": "Custom network files can be provided via volume mounts", + "supported_features": [ + "Traffic signal control", + "SUMO simulator integration", + "Multiple reward functions", + "Phase-based actions with configurable timings", + ], + }, + "textarena": { + "module": "envs.textarena_env", + "env_class": "TextArenaEnv", + "action_class": "TextArenaAction", + "default_image": "textarena-env:latest", + "description": "Text-based game environment (word games, reasoning tasks)", + "special_requirements": None, + "supported_features": [ + "Word and reasoning games", + "Multi-agent support", + "Environment configuration via kwargs", + ], + }, +} + +# Deprecated or removed environments +DEPRECATED_ENVS: Dict[str, str] = { + "julia": "julia_env has been removed from this version of OpenEnv. " + "The Julia environment is no longer maintained.", +} + + +def get_env_info(env_key: str) -> Dict[str, Any]: + """ + Get environment information from registry. + + Args: + env_key: Environment key (e.g., "coding", "atari") + + Returns: + Dictionary with environment information + + Raises: + ValueError: If environment key is not found in registry + """ + env_key = env_key.lower() + + # Check if deprecated + if env_key in DEPRECATED_ENVS: + raise ValueError(DEPRECATED_ENVS[env_key]) + + # Get from registry + if env_key not in ENV_REGISTRY: + # Try to suggest similar environment names + from difflib import get_close_matches + + suggestions = get_close_matches(env_key, ENV_REGISTRY.keys(), n=3, cutoff=0.6) + suggestion_str = "" + if suggestions: + suggestion_str = f" Did you mean: {', '.join(suggestions)}?" + + raise ValueError( + f"Unknown environment '{env_key}'. " + f"Supported environments: {', '.join(sorted(ENV_REGISTRY.keys()))}.{suggestion_str}" + ) + + return ENV_REGISTRY[env_key] + + +def list_available_environments() -> Dict[str, str]: + """ + List all available environments with their descriptions. + + Returns: + Dictionary mapping environment keys to descriptions + """ + return {key: info["description"] for key, info in ENV_REGISTRY.items()} + + +def get_all_env_keys() -> list[str]: + """Get list of all registered environment keys.""" + return sorted(ENV_REGISTRY.keys()) diff --git a/src/envs/auto_action.py b/src/envs/auto_action.py new file mode 100644 index 00000000..4d5cb3e9 --- /dev/null +++ b/src/envs/auto_action.py @@ -0,0 +1,322 @@ +# 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. + +""" +AutoAction - Automatic Action Class Selection +============================================== + +AutoAction provides a HuggingFace-style API for automatically retrieving the +correct Action class based on environment names or Docker image names. + +This module simplifies working with environment actions by automatically +detecting and returning the appropriate Action class without requiring +manual imports. + +Example: + >>> from envs import AutoEnv, AutoAction + >>> + >>> # Get Action class from environment name + >>> CodeAction = AutoAction.from_env("coding") + >>> + >>> # Or get Action class from Docker image + >>> CodeAction = AutoAction.from_image("coding-env:latest") + >>> + >>> # Use the Action class + >>> action = CodeAction(code="print('Hello!')") + >>> + >>> # Use with AutoEnv + >>> env = AutoEnv.from_docker_image("coding-env:latest") + >>> result = env.step(action) +""" + +from __future__ import annotations + +import importlib +import re +from typing import Type + +from ._registry import get_env_info + + +class AutoAction: + """ + AutoAction automatically retrieves the correct Action class based on + environment names or Docker image names. + + This class follows the HuggingFace AutoModel pattern, making it easy to + get the right Action class without needing to know which module to import. + + The class provides factory methods that look up the Action class in the + registry and return the class (not an instance) for you to instantiate. + + Example: + >>> # Get Action class from environment name + >>> CodeAction = AutoAction.from_env("coding") + >>> action = CodeAction(code="print('test')") + >>> + >>> # Get Action class from Docker image name + >>> CodeAction = AutoAction.from_image("coding-env:latest") + >>> action = CodeAction(code="print('test')") + >>> + >>> # Use with AutoEnv for a complete workflow + >>> env = AutoEnv.from_docker_image("coding-env:latest") + >>> ActionClass = AutoAction.from_image("coding-env:latest") + >>> action = ActionClass(code="print('Hello, AutoAction!')") + >>> result = env.step(action) + + Note: + AutoAction is not meant to be instantiated directly. Use the class + methods like from_env() or from_image() instead. + """ + + def __init__(self): + """AutoAction should not be instantiated directly. Use class methods instead.""" + raise TypeError( + "AutoAction is a factory class and should not be instantiated directly. " + "Use AutoAction.from_env() or AutoAction.from_image() instead." + ) + + @classmethod + def _parse_env_name_from_image(cls, image: str) -> str: + """ + Extract environment name from Docker image string. + + This method uses the same parsing logic as AutoEnv to ensure consistency. + + Supports various image name formats: + - "coding-env:latest" -> "coding" + - "ghcr.io/openenv/coding-env:v1.0" -> "coding" + - "registry.hf.space/org-name-coding-env:latest" -> "coding" + + Args: + image: Docker image name + + Returns: + Environment key (e.g., "coding", "atari") + + Raises: + ValueError: If image name format is invalid + """ + # Remove registry prefix if present + image_without_registry = re.sub(r"^[a-z0-9._-]+\.[a-z]+/", "", image, flags=re.IGNORECASE) + + # Remove organization/path prefix if present + image_without_org = image_without_registry.split("/")[-1] + + # Remove tag if present + image_without_tag = image_without_org.split(":")[0] + + # Extract environment name + # Pattern: "{env-name}-env" -> "{env-name}" + # Also support HF format: "org-name-{env-name}-env" -> "{env-name}" + if image_without_tag.endswith("-env"): + # Remove the "-env" suffix + base_name = image_without_tag[:-4] + + # For HF format like "org-name-coding-env", we need the last part before "-env" + # Split by hyphen and look for known environment names from the end + parts = base_name.split("-") + + # Try to find a match from the registry starting from the end + # This handles cases like "openenv-coding" -> "coding" + for i in range(len(parts)): + potential_env = "-".join(parts[i:]).replace("-", "_") + if potential_env in ["sumo_rl"]: # Special case for underscore envs + return potential_env.lower() + + # Check if it could be a valid env name (simple word) + if i == len(parts) - 1 or len(parts[i:]) == 1: + # Last part or single word - likely the env name + env_name = parts[-1] + return env_name.lower() + + # If we got here, just use the base name + env_name = base_name + else: + # No "-env" suffix, use as-is + env_name = image_without_tag + + # Clean up: keep underscores + env_name = env_name.replace("_", "_") + + # Validate it looks like an environment name + if not re.match(r"^[a-z0-9_]+$", env_name, re.IGNORECASE): + raise ValueError( + f"Invalid Docker image name format: '{image}'. " + f"Expected format: '{{env-name}}-env:{{tag}}' or '{{registry}}/{{org}}/{{env-name}}-env:{{tag}}'" + ) + + return env_name.lower() + + @classmethod + def _get_action_class(cls, env_key: str) -> Type: + """ + Dynamically import and return the Action class for an environment. + + Args: + env_key: Environment key from registry (e.g., "coding", "atari") + + Returns: + Action class type (not an instance) + + Raises: + ImportError: If module or class cannot be imported + ValueError: If environment not found in registry + """ + env_info = get_env_info(env_key) + + module_path = env_info["module"] + action_class_name = env_info["action_class"] + + try: + # Dynamically import the module + module = importlib.import_module(module_path) + + # Get the Action class from the module + action_class = getattr(module, action_class_name) + + return action_class + + except ImportError as e: + raise ImportError( + f"Failed to import environment module '{module_path}': {e}. " + f"Make sure the environment package is installed." + ) from e + except AttributeError as e: + raise ImportError( + f"Failed to find Action class '{action_class_name}' in module '{module_path}': {e}" + ) from e + + @classmethod + def from_env(cls, env_name: str) -> Type: + """ + Get the Action class for a specific environment by name. + + This method takes an environment name (key in the registry) and returns + the corresponding Action class. + + Args: + env_name: Environment name (e.g., "coding", "atari", "echo") + + Returns: + The Action class for the specified environment (not an instance) + + Raises: + ValueError: If environment name is not found in registry + ImportError: If Action class module cannot be imported + + Examples: + >>> # Get CodeAction class + >>> CodeAction = AutoAction.from_env("coding") + >>> action = CodeAction(code="print('Hello!')") + >>> + >>> # Get AtariAction class + >>> AtariAction = AutoAction.from_env("atari") + >>> action = AtariAction(action=0) # Fire button + >>> + >>> # Get EchoAction class + >>> EchoAction = AutoAction.from_env("echo") + >>> action = EchoAction(message="Hello!") + """ + env_key = env_name.lower() + return cls._get_action_class(env_key) + + @classmethod + def from_image(cls, image: str) -> Type: + """ + Get the Action class for an environment by parsing its Docker image name. + + This method takes a Docker image name, extracts the environment type, + and returns the corresponding Action class. + + Args: + image: Docker image name (e.g., "coding-env:latest") + + Returns: + The Action class for the environment (not an instance) + + Raises: + ValueError: If image name cannot be parsed or environment not found + ImportError: If Action class module cannot be imported + + Examples: + >>> # Get CodeAction from image name + >>> CodeAction = AutoAction.from_image("coding-env:latest") + >>> action = CodeAction(code="print('Hello!')") + >>> + >>> # With full registry path + >>> CodeAction = AutoAction.from_image("ghcr.io/openenv/coding-env:v1.0") + >>> action = CodeAction(code="x = 5 + 3") + >>> + >>> # From Hugging Face Hub format + >>> CodeAction = AutoAction.from_image("registry.hf.space/openenv-coding-env:latest") + >>> action = CodeAction(code="import math") + """ + env_key = cls._parse_env_name_from_image(image) + return cls._get_action_class(env_key) + + @classmethod + def get_action_info(cls, env_name: str) -> dict: + """ + Get information about the Action class for an environment. + + This is a convenience method to get details about what fields the + Action class expects without having to instantiate it. + + Args: + env_name: Environment name (e.g., "coding", "atari") + + Returns: + Dictionary with Action class information including module and class name + + Example: + >>> info = AutoAction.get_action_info("coding") + >>> print(info["action_class"]) # "CodeAction" + >>> print(info["module"]) # "envs.coding_env" + """ + env_key = env_name.lower() + env_info = get_env_info(env_key) + + return { + "action_class": env_info["action_class"], + "module": env_info["module"], + "env_class": env_info["env_class"], + "description": env_info["description"], + } + + @classmethod + def list_actions(cls) -> None: + """ + Print a list of all available Action classes. + + This is a convenience method for discovering what Action classes are available. + + Example: + >>> AutoAction.list_actions() + Available Action Classes: + ------------------------- + coding : CodeAction (Python code execution environment) + atari : AtariAction (Atari 2600 games environment (100+ games)) + echo : EchoAction (Simple echo test environment) + ... + """ + from ._registry import ENV_REGISTRY + + print("Available Action Classes:") + print("-" * 70) + + for env_key in sorted(ENV_REGISTRY.keys()): + info = ENV_REGISTRY[env_key] + action_class = info["action_class"] + description = info["description"] + print(f" {env_key:<15}: {action_class:<20} ({description})") + + print("-" * 70) + print(f"Total: {len(ENV_REGISTRY)} Action classes") + print("\nUsage:") + print(" ActionClass = AutoAction.from_env('env-name')") + print(" # or") + print(" ActionClass = AutoAction.from_image('env-name-env:latest')") diff --git a/src/envs/auto_env.py b/src/envs/auto_env.py new file mode 100644 index 00000000..77132782 --- /dev/null +++ b/src/envs/auto_env.py @@ -0,0 +1,415 @@ +# 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. + +""" +AutoEnv - Automatic Environment Selection +========================================== + +AutoEnv provides a HuggingFace-style API for automatically selecting and +instantiating the correct environment client based on Docker image names. + +This module simplifies environment creation by automatically detecting the +environment type from the Docker image name and instantiating the appropriate +client class. + +Example: + >>> from envs import AutoEnv, AutoAction + >>> + >>> # Automatically detect and create the right environment + >>> client = AutoEnv.from_docker_image("coding-env:latest") + >>> + >>> # Get the corresponding Action class + >>> CodeAction = AutoAction.from_image("coding-env:latest") + >>> + >>> # Use them together + >>> result = client.reset() + >>> action = CodeAction(code="print('Hello, AutoEnv!')") + >>> step_result = client.step(action) + >>> client.close() +""" + +from __future__ import annotations + +import importlib +import re +from typing import Any, Optional, TYPE_CHECKING + +from ._registry import get_env_info, list_available_environments + +if TYPE_CHECKING: + from core.containers.runtime import ContainerProvider + from core.http_env_client import HTTPEnvClient + + +class AutoEnv: + """ + AutoEnv automatically selects and instantiates the correct environment client + based on Docker image names. + + This class follows the HuggingFace AutoModel pattern, making it easy to work + with different environments without needing to import specific client classes. + + The class provides factory methods that parse Docker image names, look up the + corresponding environment in the registry, and return an instance of the + appropriate client class. + + Example: + >>> # Simple usage - just specify the image + >>> env = AutoEnv.from_docker_image("coding-env:latest") + >>> + >>> # With custom configuration + >>> env = AutoEnv.from_docker_image( + ... "dipg-env:latest", + ... env_vars={"DIPG_DATASET_PATH": "/data/dipg"} + ... ) + >>> + >>> # From Hugging Face Hub + >>> env = AutoEnv.from_hub("openenv/coding-env", tag="v1.0") + >>> + >>> # List available environments + >>> AutoEnv.list_environments() + + Note: + AutoEnv is not meant to be instantiated directly. Use the class methods + like from_docker_image() or from_hub() instead. + """ + + def __init__(self): + """AutoEnv should not be instantiated directly. Use class methods instead.""" + raise TypeError( + "AutoEnv is a factory class and should not be instantiated directly. " + "Use AutoEnv.from_docker_image() or AutoEnv.from_hub() instead." + ) + + @classmethod + def _parse_env_name_from_image(cls, image: str) -> str: + """ + Extract environment name from Docker image string. + + Supports various image name formats: + - "coding-env:latest" -> "coding" + - "ghcr.io/openenv/coding-env:v1.0" -> "coding" + - "registry.hf.space/org-name-coding-env:latest" -> "coding" + + Args: + image: Docker image name + + Returns: + Environment key (e.g., "coding", "atari") + + Raises: + ValueError: If image name format is invalid + """ + # Remove registry prefix if present + # Examples: "ghcr.io/openenv/coding-env:latest", "registry.hf.space/..." + image_without_registry = re.sub( + r"^[a-z0-9._-]+\.[a-z]+/", "", image, flags=re.IGNORECASE + ) + + # Remove organization/path prefix if present + # Example: "openenv/coding-env:latest" -> "coding-env:latest" + image_without_org = image_without_registry.split("/")[-1] + + # Remove tag if present + # Example: "coding-env:latest" -> "coding-env" + image_without_tag = image_without_org.split(":")[0] + + # Extract environment name + # Pattern: "{env-name}-env" -> "{env-name}" + # Also support HF format: "org-name-{env-name}-env" -> "{env-name}" + # First try to match the "-env" suffix pattern + if image_without_tag.endswith("-env"): + # Remove the "-env" suffix + base_name = image_without_tag[:-4] + + # For HF format like "org-name-coding-env", we need the last part before "-env" + # Split by hyphen and look for known environment names from the end + parts = base_name.split("-") + + # Try to find a match from the registry starting from the end + # This handles cases like "openenv-coding" -> "coding" + for i in range(len(parts)): + potential_env = "-".join(parts[i:]).replace("-", "_") + if potential_env in ["sumo_rl"]: # Special case for underscore envs + return potential_env.lower() + + # Check if it could be a valid env name (simple word) + if i == len(parts) - 1 or len(parts[i:]) == 1: + # Last part or single word - likely the env name + env_name = parts[-1] + return env_name.lower() + + # If we got here, just use the base name + env_name = base_name + else: + # No "-env" suffix, use as-is + env_name = image_without_tag + + # Clean up: convert underscores as needed + env_name = env_name.replace("_", "_") # Keep underscores + + # Validate it looks like an environment name + if not re.match(r"^[a-z0-9_]+$", env_name, re.IGNORECASE): + raise ValueError( + f"Invalid Docker image name format: '{image}'. " + f"Expected format: '{{env-name}}-env:{{tag}}' or '{{registry}}/{{org}}/{{env-name}}-env:{{tag}}'" + ) + + return env_name.lower() + + @classmethod + def _get_env_class(cls, env_key: str) -> type: + """ + Dynamically import and return the environment class. + + Args: + env_key: Environment key from registry + + Returns: + Environment class type + + Raises: + ImportError: If module or class cannot be imported + """ + env_info = get_env_info(env_key) + + module_path = env_info["module"] + class_name = env_info["env_class"] + + try: + # Dynamically import the module + module = importlib.import_module(module_path) + + # Get the class from the module + env_class = getattr(module, class_name) + + return env_class + + except ImportError as e: + raise ImportError( + f"Failed to import environment module '{module_path}': {e}. " + f"Make sure the environment package is installed." + ) from e + except AttributeError as e: + raise ImportError( + f"Failed to find class '{class_name}' in module '{module_path}': {e}" + ) from e + + @classmethod + def from_docker_image( + cls, + image: str, + provider: Optional["ContainerProvider"] = None, + wait_timeout: float = 30.0, + **kwargs: Any, + ) -> "HTTPEnvClient": + """ + Create an environment client from a Docker image, automatically detecting + the environment type. + + This method: + 1. Parses the Docker image name to identify the environment type + 2. Looks up the environment in the registry + 3. Dynamically imports the appropriate client class + 4. Calls that class's from_docker_image() method + 5. Returns the instantiated client + + Args: + image: Docker image name (e.g., "coding-env:latest") + provider: Optional container provider (defaults to LocalDockerProvider) + wait_timeout: Maximum time (in seconds) to wait for container to be ready (default: 30.0) + Increase this for slow-starting containers or low-resource environments + **kwargs: Additional arguments passed to provider.start_container() + Common kwargs: + - env_vars: Dict of environment variables + - port: Port to expose + - volumes: Volume mounts + + Returns: + An instance of the appropriate environment client class + + Raises: + ValueError: If image name cannot be parsed or environment not found + ImportError: If environment module cannot be imported + TimeoutError: If container doesn't become ready within wait_timeout + + Examples: + >>> # Simple usage + >>> env = AutoEnv.from_docker_image("coding-env:latest") + >>> result = env.reset() + >>> env.close() + >>> + >>> # With custom timeout (useful for slow containers) + >>> env = AutoEnv.from_docker_image( + ... "coding-env:latest", + ... wait_timeout=60.0 # Wait up to 60 seconds + ... ) + >>> + >>> # With environment variables (for DIPG environment) + >>> env = AutoEnv.from_docker_image( + ... "dipg-env:latest", + ... wait_timeout=60.0, + ... env_vars={"DIPG_DATASET_PATH": "/data/dipg"} + ... ) + >>> + >>> # With custom provider + >>> from core.containers.runtime import LocalDockerProvider + >>> provider = LocalDockerProvider() + >>> env = AutoEnv.from_docker_image( + ... "coding-env:latest", + ... provider=provider, + ... wait_timeout=45.0 + ... ) + """ + # Parse environment name from image + env_key = cls._parse_env_name_from_image(image) + + # Get environment class + env_class = cls._get_env_class(env_key) + + # Get environment info for special requirements + env_info = get_env_info(env_key) + + # Warn about special requirements if not provided + special_req = env_info.get("special_requirements") + if special_req and "env_vars" not in kwargs: + import warnings + + warnings.warn( + f"Environment '{env_key}' has special requirements: {special_req}. " + f"You may need to provide appropriate env_vars.", + UserWarning, + ) + + # Create and return instance using the class's from_docker_image method + return env_class.from_docker_image( + image=image, provider=provider, wait_timeout=wait_timeout, **kwargs + ) + + @classmethod + def from_hub( + cls, + repo_id: str, + provider: Optional["ContainerProvider"] = None, + **kwargs: Any, + ) -> "HTTPEnvClient": + """ + Create an environment client from Hugging Face Hub. + + This is a convenience method that constructs the appropriate Docker image + name from a Hugging Face repository ID and calls from_docker_image(). + + Args: + repo_id: Hugging Face repository ID (e.g., "openenv/coding-env") + provider: Optional container provider (defaults to LocalDockerProvider) + **kwargs: Additional arguments, including: + - tag: Docker image tag (default: "latest") + - env_vars: Dict of environment variables + - Other provider kwargs + + Returns: + An instance of the appropriate environment client class + + Example: + >>> # Pull from Hugging Face Hub + >>> env = AutoEnv.from_hub("openenv/coding-env") + >>> + >>> # With specific version + >>> env = AutoEnv.from_hub("openenv/coding-env", tag="v1.0") + """ + # Extract tag if provided + tag = kwargs.pop("tag", "latest") + + # Construct image name for HF registry + image = f"registry.hf.space/{repo_id.replace('/', '-')}:{tag}" + + # Use from_docker_image with the constructed image name + return cls.from_docker_image(image=image, provider=provider, **kwargs) + + @classmethod + def list_environments(cls) -> None: + """ + Print a list of all available environments with descriptions. + + This is a convenience method for discovering what environments are available. + + Example: + >>> AutoEnv.list_environments() + Available Environments: + ---------------------- + atari : Atari 2600 games environment (100+ games) + browsergym : Web browsing environment with multiple benchmarks + chat : Chat environment with tokenization support + ... + """ + envs = list_available_environments() + + print("Available Environments:") + print("-" * 60) + + for env_key in sorted(envs.keys()): + description = envs[env_key] + print(f" {env_key:<15}: {description}") + + print("-" * 60) + print(f"Total: {len(envs)} environments") + print("\nUsage:") + print(" env = AutoEnv.from_docker_image('{env-name}-env:latest')") + + @classmethod + def from_name(cls, env_name: str) -> type: + """ + Get the environment class for a specific environment by name. + + This method takes an environment name (key in the registry) and returns + the corresponding environment class (not an instance). + + Args: + env_name: Environment name (e.g., "coding", "atari", "echo") + + Returns: + The environment class for the specified environment (not an instance) + + Raises: + ValueError: If environment name is not found in registry + ImportError: If environment class module cannot be imported + + Examples: + >>> # Get CodingEnv class + >>> CodingEnv = AutoEnv.from_name("coding") + >>> + >>> # Get AtariEnv class + >>> AtariEnv = AutoEnv.from_name("atari") + >>> + >>> # Get EchoEnv class + >>> EchoEnv = AutoEnv.from_name("echo") + """ + env_key = env_name.lower() + return cls._get_env_class(env_key) + + @classmethod + def get_env_info(cls, env_key: str) -> dict: + """ + Get detailed information about a specific environment. + + Args: + env_key: Environment key (e.g., "coding", "atari") + + Returns: + Dictionary with environment information including: + - description + - special_requirements + - supported_features + - default_image + + Example: + >>> info = AutoEnv.get_env_info("coding") + >>> print(info["description"]) + >>> print(info["special_requirements"]) + >>> for feature in info["supported_features"]: + ... print(f" - {feature}") + """ + return get_env_info(env_key) diff --git a/src/envs/echo_env/models.py b/src/envs/echo_env/models.py index d73134ba..083c3989 100644 --- a/src/envs/echo_env/models.py +++ b/src/envs/echo_env/models.py @@ -27,4 +27,4 @@ class EchoObservation(Observation): """Observation from the Echo environment - the echoed message.""" echoed_message: str - message_length: int = 0 \ No newline at end of file + message_length: int = 0