diff --git a/.gitignore b/.gitignore index fed309b0..25342208 100644 --- a/.gitignore +++ b/.gitignore @@ -109,3 +109,4 @@ outputs/ .uv/ *.backup*/ +src/envs/.discovery_cache.json diff --git a/examples/auto_env_example.py b/examples/auto_env_example.py new file mode 100755 index 00000000..956461f8 --- /dev/null +++ b/examples/auto_env_example.py @@ -0,0 +1,317 @@ +#!/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 + +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_name("coding-env") + print("✓ Environment created!") + print() + + # Get the Action class automatically + print("Getting Action class using AutoAction...") + CodeAction = AutoAction.from_name("coding-env") + 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 environment key""" + print("=" * 70) + print("Example 2: Alternative Syntax") + print("=" * 70) + print() + + # You can also use just the environment key + print("Getting Action class by environment name...") + CodeAction = AutoAction.from_name("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" Observation Class: {info['observation_class']}") + print(f" Module: {info['module']}") + print(f" Version: {info['version']}") + print(f" Spec Version: {info['spec_version']}") + 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_name("nonexistent-env") + 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_name("cooding-env") + except ValueError as e: + print(f"✓ Got helpful suggestion: {e}") + print() + + # Try deprecated julia environment + print("Trying deprecated 'julia' environment...") + try: + env = AutoEnv.from_name("julia-env") + 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_name('dipg-env')") + print() + print(" # Correct usage:") + print(" env = AutoEnv.from_name(") + print(" 'dipg-env',") + 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_name(") + print(" 'finrl-env',") + 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 + # Use the simplified name format + env_image_name = f"{env_name}-env" if not env_name.endswith("-env") else env_name + env = AutoEnv.from_name(env_image_name, wait_timeout=60.0) + print("✓ Environment created!") + + # Get action class + ActionClass = AutoAction.from_name(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/envs/__init__.py b/src/envs/__init__.py new file mode 100644 index 00000000..7a583800 --- /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 +environment names. + +Example: + >>> from envs import AutoEnv, AutoAction + >>> + >>> # Automatically detect and create environment from name + >>> client = AutoEnv.from_name("coding-env") + >>> + >>> # Get the corresponding Action class + >>> CodeAction = AutoAction.from_name("coding-env") + >>> + >>> # 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/_discovery.py b/src/envs/_discovery.py new file mode 100644 index 00000000..3536b125 --- /dev/null +++ b/src/envs/_discovery.py @@ -0,0 +1,565 @@ +# 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 Auto-Discovery System +================================== + +This module provides automatic discovery of OpenEnv environments by: +1. Discovering installed openenv-* packages using importlib.metadata +2. Loading manifests (openenv.yaml) from package resources +3. Caching results for performance +4. Supporting HuggingFace Hub downloads + +This enables AutoEnv to work without coupling to src/envs/ directory. +""" + +import importlib +import importlib.metadata +import importlib.resources +import json +import logging +import re +import tempfile +from dataclasses import dataclass, asdict +from pathlib import Path +from typing import Dict, List, Optional, Type, Any + +import yaml + +logger = logging.getLogger(__name__) + + +@dataclass +class EnvironmentInfo: + """ + Rich information about a discovered environment. + + Attributes: + env_key: Environment key (e.g., "echo", "coding") + name: Full environment name (e.g., "echo_env") + package_name: Package name (e.g., "openenv-echo_env") + version: Version string + description: Human-readable description + client_module_path: Full module path to client (e.g., "echo_env.client") + client_class_name: Client class name (e.g., "EchoEnv") + action_class_name: Action class name (e.g., "EchoAction") + observation_class_name: Observation class name (e.g., "EchoObservation") + default_image: Default Docker image name (e.g., "echo-env:latest") + spec_version: OpenEnv spec version (from openenv.yaml) + manifest: Original manifest data + """ + + env_key: str + name: str + package_name: str + version: str + description: str + client_module_path: str + client_class_name: str + action_class_name: str + observation_class_name: str + default_image: str + spec_version: Optional[int] = None + manifest: Optional[Dict[str, Any]] = None + + def get_client_class(self) -> Type: + """ + Dynamically import and return the client class. + + Returns: + Client class (e.g., EchoEnv) + + Raises: + ImportError: If module or class cannot be imported + """ + try: + module = importlib.import_module(self.client_module_path) + return getattr(module, self.client_class_name) + except ImportError as e: + raise ImportError( + f"Failed to import {self.client_class_name} from {self.client_module_path}: {e}\n" + f"Make sure the package '{self.package_name}' is installed: " + f"pip install {self.package_name}" + ) from e + except AttributeError as e: + raise ImportError( + f"Class {self.client_class_name} not found in {self.client_module_path}: {e}" + ) from e + + def get_action_class(self) -> Type: + """ + Dynamically import and return the action class. + + Returns: + Action class (e.g., EchoAction) + + Raises: + ImportError: If module or class cannot be imported + """ + try: + module = importlib.import_module(self.client_module_path) + return getattr(module, self.action_class_name) + except ImportError as e: + raise ImportError( + f"Failed to import {self.action_class_name} from {self.client_module_path}: {e}\n" + f"Make sure the package '{self.package_name}' is installed: " + f"pip install {self.package_name}" + ) from e + except AttributeError as e: + raise ImportError( + f"Class {self.action_class_name} not found in {self.client_module_path}: {e}" + ) from e + + def get_observation_class(self) -> Type: + """ + Dynamically import and return the observation class. + + Returns: + Observation class (e.g., EchoObservation) + + Raises: + ImportError: If module or class cannot be imported + """ + try: + module = importlib.import_module(self.client_module_path) + return getattr(module, self.observation_class_name) + except ImportError as e: + raise ImportError( + f"Failed to import {self.observation_class_name} from {self.client_module_path}: {e}\n" + f"Make sure the package '{self.package_name}' is installed: " + f"pip install {self.package_name}" + ) from e + except AttributeError as e: + raise ImportError( + f"Class {self.observation_class_name} not found in {self.client_module_path}: {e}" + ) from e + + +def _normalize_env_name(name: str) -> str: + """ + Normalize environment name to standard format. + + Args: + name: Input name (e.g., "echo", "echo-env", "echo_env") + + Returns: + Normalized name (e.g., "echo_env") + + Examples: + >>> _normalize_env_name("echo") + 'echo_env' + >>> _normalize_env_name("echo-env") + 'echo_env' + >>> _normalize_env_name("echo_env") + 'echo_env' + """ + # Remove common suffixes + name = re.sub(r"[-_]env$", "", name) + # Convert hyphens to underscores + name = name.replace("-", "_") + # Add _env suffix if not present + if not name.endswith("_env"): + name = f"{name}_env" + return name + + +def _is_hub_url(name: str) -> bool: + """ + Check if name is a HuggingFace Hub URL or repo ID. + + Args: + name: Input name + + Returns: + True if it looks like a Hub URL + + Examples: + >>> _is_hub_url("meta-pytorch/echo-env") + True + >>> _is_hub_url("https://huggingface.co/meta-pytorch/echo-env") + True + >>> _is_hub_url("echo") + False + """ + # Contains org/repo pattern or huggingface.co domain + return "/" in name or "huggingface.co" in name + + +def _infer_class_name(env_name: str, class_type: str) -> str: + """ + Infer class name from environment name using simple conventions. + + Args: + env_name: Environment name (e.g., "echo_env") + class_type: Type of class ("client", "action", "observation") + + Returns: + Inferred class name + + Examples: + >>> _infer_class_name("echo_env", "client") + 'EchoEnv' + >>> _infer_class_name("echo_env", "action") + 'EchoAction' + """ + # Remove _env suffix for base name + base_name = env_name.replace("_env", "") + + # Convert to PascalCase + pascal_name = "".join(word.capitalize() for word in base_name.split("_")) + + # Add suffix based on type + if class_type == "client": + return f"{pascal_name}Env" + elif class_type == "action": + return f"{pascal_name}Action" + elif class_type == "observation": + return f"{pascal_name}Observation" + else: + raise ValueError(f"Unknown class type: {class_type}") + + +def _load_manifest_from_package(package_name: str, module_name: str) -> Optional[Dict[str, Any]]: + """ + Load openenv.yaml manifest from an installed package. + + Args: + package_name: Package name (e.g., "openenv-echo_env") + module_name: Module name (e.g., "echo_env") + + Returns: + Parsed manifest dictionary, or None if not found + + """ + try: + # Try to read openenv.yaml from package + if hasattr(importlib.resources, 'files'): + # Python 3.9+ + package_files = importlib.resources.files(module_name) + if (package_files / "openenv.yaml").is_file(): + manifest_text = (package_files / "openenv.yaml").read_text() + return yaml.safe_load(manifest_text) + else: + # Python 3.7-3.8 fallback + with importlib.resources.open_text(module_name, "openenv.yaml") as f: + return yaml.safe_load(f) + except (FileNotFoundError, ModuleNotFoundError, AttributeError): + logger.debug(f"No openenv.yaml found in {module_name}") + return None + except Exception as e: + logger.warning(f"Failed to load openenv.yaml from {module_name}: {e}") + return None + + +def _create_env_info_from_package(package_name: str, module_name: str, version: str) -> Optional[EnvironmentInfo]: + """ + Create EnvironmentInfo from an installed package. + + Args: + package_name: Package name (e.g., "openenv-echo_env") + module_name: Module name (e.g., "echo_env") + version: Package version + + Returns: + EnvironmentInfo instance, or None if invalid + """ + # Load manifest + manifest = _load_manifest_from_package(package_name, module_name) + + # Get environment name + if manifest and "name" in manifest: + env_name = manifest["name"] + else: + # Infer from module name + env_name = module_name + + # Normalize to ensure _env suffix + if not env_name.endswith("_env"): + env_name = f"{env_name}_env" + + # Determine env_key (e.g., "echo_env" → "echo") + env_key = env_name.replace("_env", "") if env_name.endswith("_env") else env_name + + # Get description + description = manifest.get("description", f"{env_name} environment") if manifest else f"{env_name} environment" + + # Get spec version + spec_version = manifest.get("spec_version") if manifest else None + + # Determine class names + # Check if manifest has custom class names (custom format) + if manifest and "action" in manifest and "observation" in manifest: + # Custom format (like coding_env) + client_class_name = _infer_class_name(env_name, "client") + action_class_name = manifest.get("action", _infer_class_name(env_name, "action")) + observation_class_name = manifest.get("observation", _infer_class_name(env_name, "observation")) + else: + # Use conventions + client_class_name = _infer_class_name(env_name, "client") + action_class_name = _infer_class_name(env_name, "action") + observation_class_name = _infer_class_name(env_name, "observation") + + # Module path is just module_name.client + client_module_path = f"{module_name}.client" + + # Determine default Docker image name + image_name = env_name.replace("_", "-") + default_image = f"{image_name}:latest" + + return EnvironmentInfo( + env_key=env_key, + name=env_name, + package_name=package_name, + version=version, + description=description, + client_module_path=client_module_path, + client_class_name=client_class_name, + action_class_name=action_class_name, + observation_class_name=observation_class_name, + default_image=default_image, + spec_version=spec_version, + manifest=manifest, + ) + + +class EnvironmentDiscovery: + """ + Auto-discovery system for OpenEnv environments using installed packages. + + This class discovers installed openenv-* packages and loads their metadata. + """ + + def __init__(self): + """Initialize discovery system.""" + self._cache: Optional[Dict[str, EnvironmentInfo]] = None + self._cache_file = Path(tempfile.gettempdir()) / "openenv_discovery_cache.json" + + def _discover_installed_packages(self) -> Dict[str, EnvironmentInfo]: + """ + Discover all installed openenv-* packages. + + Returns: + Dictionary mapping env_key to EnvironmentInfo + """ + environments = {} + + # Get all installed packages + try: + distributions = importlib.metadata.distributions() + except Exception as e: + logger.warning(f"Failed to get installed packages: {e}") + return environments + + # Filter for openenv-* packages (exclude openenv-core) + for dist in distributions: + package_name = dist.metadata["Name"] + + if not package_name.startswith("openenv-"): + continue + + if package_name == "openenv-core": + continue + + # Get module name (e.g., "openenv-echo_env" → "echo_env") + module_name = package_name.replace("openenv-", "").replace("-", "_") + + # Get version + version = dist.version + + try: + # Create environment info + env_info = _create_env_info_from_package(package_name, module_name, version) + + if env_info: + environments[env_info.env_key] = env_info + logger.debug(f"Discovered environment: {env_info.env_key} ({package_name})") + + except Exception as e: + logger.warning(f"Failed to load environment from {package_name}: {e}") + continue + + return environments + + def _load_cache(self) -> Optional[Dict[str, EnvironmentInfo]]: + """ + Load cached discovery results. + + Returns: + Dictionary of env_key -> EnvironmentInfo, or None if cache invalid + """ + if not self._cache_file.exists(): + return None + + try: + with open(self._cache_file, "r") as f: + cache_data = json.load(f) + + # Reconstruct EnvironmentInfo objects + cache = {} + for env_key, env_data in cache_data.items(): + cache[env_key] = EnvironmentInfo(**env_data) + + return cache + except Exception as e: + logger.warning(f"Failed to load discovery cache: {e}") + return None + + def _save_cache(self, environments: Dict[str, EnvironmentInfo]) -> None: + """ + Save discovery results to cache. + + Args: + environments: Dictionary of env_key -> EnvironmentInfo + """ + try: + cache_data = {} + for env_key, env_info in environments.items(): + cache_data[env_key] = asdict(env_info) + + with open(self._cache_file, "w") as f: + json.dump(cache_data, f, indent=2) + + except Exception as e: + logger.warning(f"Failed to save discovery cache: {e}") + + def discover(self, use_cache: bool = True) -> Dict[str, EnvironmentInfo]: + """ + Discover all installed OpenEnv environments. + + Args: + use_cache: If True, try to load from cache first + + Returns: + Dictionary mapping env_key to EnvironmentInfo + + Examples: + >>> discovery = EnvironmentDiscovery() + >>> envs = discovery.discover() + >>> print(envs.keys()) + dict_keys(['echo', 'coding', ...]) + """ + # Try to load from memory cache first + if use_cache and self._cache is not None: + return self._cache + + # Try to load from file cache + if use_cache: + cached = self._load_cache() + if cached is not None: + self._cache = cached + return self._cache + + # Discover from installed packages + environments = self._discover_installed_packages() + + # Save to cache + self._save_cache(environments) + self._cache = environments + + return environments + + def get_environment(self, env_key: str) -> Optional[EnvironmentInfo]: + """ + Get information about a specific environment. + + Args: + env_key: Environment key (e.g., "echo", "coding") + + Returns: + EnvironmentInfo if found, None otherwise + + Examples: + >>> discovery = EnvironmentDiscovery() + >>> env = discovery.get_environment("echo") + >>> print(env.client_class_name) + 'EchoEnv' + """ + environments = self.discover() + return environments.get(env_key) + + def get_environment_by_name(self, name: str) -> Optional[EnvironmentInfo]: + """ + Get environment info by flexible name matching. + + Args: + name: Environment name (e.g., "echo", "echo-env", "echo_env") + + Returns: + EnvironmentInfo if found, None otherwise + """ + # Normalize name to env_key + normalized = _normalize_env_name(name) + env_key = normalized.replace("_env", "") + + return self.get_environment(env_key) + + def list_environments(self) -> None: + """ + Print a formatted list of all discovered environments. + + Examples: + >>> discovery = EnvironmentDiscovery() + >>> discovery.list_environments() + Available OpenEnv Environments: + ---------------------------------------------------------------------- + echo : Echo Environment (v0.1.0) - openenv-echo_env + coding : Coding Environment (v0.1.0) - openenv-coding_env + ... + """ + environments = self.discover() + + print("Available OpenEnv Environments:") + print("-" * 70) + + if not environments: + print(" No OpenEnv environments found.") + print(" Install environments with: pip install openenv-") + else: + for env_key in sorted(environments.keys()): + env = environments[env_key] + print(f" {env_key:<15}: {env.description} (v{env.version})") + print(f" Package: {env.package_name}") + + print("-" * 70) + print(f"Total: {len(environments)} environments") + + def clear_cache(self) -> None: + """Clear the discovery cache.""" + if self._cache_file.exists(): + self._cache_file.unlink() + self._cache = None + + +# Global discovery instance +_global_discovery: Optional[EnvironmentDiscovery] = None + + +def get_discovery() -> EnvironmentDiscovery: + """ + Get or create the global discovery instance. + + Returns: + Global EnvironmentDiscovery instance + + Examples: + >>> discovery = get_discovery() + >>> envs = discovery.discover() + """ + global _global_discovery + + if _global_discovery is None: + _global_discovery = EnvironmentDiscovery() + + return _global_discovery + + +def reset_discovery() -> None: + """Reset the global discovery instance (useful for testing).""" + global _global_discovery + if _global_discovery is not None: + _global_discovery.clear_cache() + _global_discovery = None diff --git a/src/envs/auto_action.py b/src/envs/auto_action.py new file mode 100644 index 00000000..f21689b5 --- /dev/null +++ b/src/envs/auto_action.py @@ -0,0 +1,264 @@ +# 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 from installed packages or HuggingFace Hub. + +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_name("coding") + >>> action = CodeAction(code="print('Hello!')") + >>> + >>> # From HuggingFace Hub + >>> CodeAction = AutoAction.from_name("meta-pytorch/coding-env") + >>> + >>> # Use with AutoEnv + >>> env = AutoEnv.from_name("coding-env") + >>> result = env.step(action) +""" + +from __future__ import annotations + +import logging +from typing import Type, Dict, Any + +from ._discovery import get_discovery, _is_hub_url +from .auto_env import AutoEnv + +logger = logging.getLogger(__name__) + + +class AutoAction: + """ + AutoAction automatically retrieves the correct Action class based on + environment names or HuggingFace Hub repositories. + + 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 and + return the class (not an instance) for you to instantiate. + + Example: + >>> # From installed package + >>> CodeAction = AutoAction.from_name("coding") + >>> action = CodeAction(code="print('test')") + >>> + >>> # From HuggingFace Hub + >>> CodeAction = AutoAction.from_name("meta-pytorch/coding-env") + >>> action = CodeAction(code="print('test')") + >>> + >>> # Use with AutoEnv for a complete workflow + >>> env = AutoEnv.from_name("coding-env") + >>> ActionClass = AutoAction.from_name("coding-env") + >>> action = ActionClass(code="print('Hello, AutoAction!')") + >>> result = env.step(action) + + Note: + AutoAction is not meant to be instantiated directly. Use the class + method from_name() 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_name() instead." + ) + + @classmethod + def from_name(cls, name: str) -> Type: + """ + Get the Action class from environment name or HuggingFace Hub repository. + + This method automatically: + 1. Checks if the name is a HuggingFace Hub URL/repo ID + 2. If Hub: downloads and installs the environment package + 3. If local: looks up the installed openenv-* package + 4. Imports and returns the Action class + + Args: + name: Environment name or HuggingFace Hub repo ID + Examples: + - "coding" / "coding-env" / "coding_env" + - "meta-pytorch/coding-env" (Hub repo ID) + - "https://huggingface.co/meta-pytorch/coding-env" (Hub URL) + + Returns: + Action class (not an instance!) + + Raises: + ValueError: If environment not found + ImportError: If environment package is not installed + + Examples: + >>> # From installed package + >>> CodeAction = AutoAction.from_name("coding-env") + >>> action = CodeAction(code="print('Hello!')") + >>> + >>> # From HuggingFace Hub + >>> CodeAction = AutoAction.from_name("meta-pytorch/coding-env") + >>> action = CodeAction(code="print('Hello!')") + >>> + >>> # Different name formats + >>> EchoAction = AutoAction.from_name("echo") + >>> EchoAction = AutoAction.from_name("echo-env") + >>> EchoAction = AutoAction.from_name("echo_env") + """ + # Check if it's a HuggingFace Hub URL or repo ID + if _is_hub_url(name): + # Download from Hub and install (reuse AutoEnv logic) + env_path = AutoEnv._download_from_hub(name) + package_name = AutoEnv._install_from_path(env_path) + + # Clear discovery cache to pick up the newly installed package + get_discovery().clear_cache() + + # Extract environment name from package name + # "openenv-coding_env" -> "coding_env" + env_name = package_name.replace("openenv-", "").replace("-", "_") + else: + env_name = name + + # Get environment info from discovery + discovery = get_discovery() + env_info = discovery.get_environment_by_name(env_name) + + if not env_info: + # Environment not found - provide helpful error message + available_envs = discovery.discover() + + if not available_envs: + raise ValueError( + f"No OpenEnv environments found.\n" + f"Install an environment with: pip install openenv-\n" + f"Or specify a HuggingFace Hub repository: AutoAction.from_name('org/repo')" + ) + + # Try to suggest similar environment names + from difflib import get_close_matches + + env_keys = list(available_envs.keys()) + suggestions = get_close_matches(env_name, env_keys, n=3, cutoff=0.6) + + error_msg = f"Unknown environment '{env_name}'.\n" + if suggestions: + error_msg += f"Did you mean: {', '.join(suggestions)}?\n" + error_msg += f"Available environments: {', '.join(sorted(env_keys))}" + + raise ValueError(error_msg) + + # Get the action class + try: + action_class = env_info.get_action_class() + return action_class + except ImportError as e: + raise ImportError( + f"Failed to import action class for '{env_name}'.\n" + f"Package '{env_info.package_name}' appears to be installed but the module cannot be imported.\n" + f"Try reinstalling: pip install --force-reinstall {env_info.package_name}\n" + f"Original error: {e}" + ) from e + + @classmethod + def from_env(cls, env_name: str) -> Type: + """ + Get the Action class from environment name. + + This is an alias for from_name() for backward compatibility and clarity. + + Args: + env_name: Environment name (e.g., "coding", "echo") + + Returns: + Action class (not an instance!) + + Examples: + >>> CodeAction = AutoAction.from_env("coding") + >>> action = CodeAction(code="print('Hello!')") + """ + return cls.from_name(env_name) + + @classmethod + def get_action_info(cls, name: str) -> Dict[str, Any]: + """ + Get detailed information about an action class. + + Args: + name: Environment name + + Returns: + Dictionary with action class metadata + + Raises: + ValueError: If environment not found + + Examples: + >>> info = AutoAction.get_action_info("coding") + >>> print(info['action_class']) + 'CodingAction' + >>> print(info['module']) + 'coding_env.client' + """ + discovery = get_discovery() + env_info = discovery.get_environment_by_name(name) + + if not env_info: + raise ValueError(f"Unknown environment: {name}") + + return { + "env_key": env_info.env_key, + "env_name": env_info.name, + "package": env_info.package_name, + "action_class": env_info.action_class_name, + "observation_class": env_info.observation_class_name, + "module": env_info.client_module_path, + } + + @classmethod + def list_actions(cls) -> None: + """ + Print a formatted list of all available action classes. + + This discovers all installed openenv-* packages and displays + their action class information in a user-friendly format. + + Examples: + >>> AutoAction.list_actions() + Available Action Classes: + ---------------------------------------------------------------------- + echo : EchoAction (from openenv-echo-env) + coding : CodingAction (from openenv-coding_env) + ---------------------------------------------------------------------- + Total: 2 action classes + """ + discovery = get_discovery() + environments = discovery.discover() + + print("Available Action Classes:") + print("-" * 70) + + if not environments: + print(" No OpenEnv environments found.") + print(" Install environments with: pip install openenv-") + else: + for env_key in sorted(environments.keys()): + env = environments[env_key] + print(f" {env_key:<15}: {env.action_class_name}") + print(f" Package: {env.package_name}") + + print("-" * 70) + print(f"Total: {len(environments)} action classes") diff --git a/src/envs/auto_env.py b/src/envs/auto_env.py new file mode 100644 index 00000000..52ed4d64 --- /dev/null +++ b/src/envs/auto_env.py @@ -0,0 +1,412 @@ +# 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 from installed packages or +HuggingFace Hub. + +This module simplifies environment creation by automatically detecting the +environment type from the name and instantiating the appropriate client class. + +Example: + >>> from envs import AutoEnv, AutoAction + >>> + >>> # From installed package + >>> env = AutoEnv.from_name("coding-env") + >>> + >>> # From HuggingFace Hub + >>> env = AutoEnv.from_name("meta-pytorch/coding-env") + >>> + >>> # With configuration + >>> env = AutoEnv.from_name("coding", env_vars={"DEBUG": "1"}) +""" + +from __future__ import annotations + +import importlib +import logging +import os +import subprocess +import tempfile +from pathlib import Path +from typing import Any, Optional, TYPE_CHECKING, Dict + +from ._discovery import get_discovery, _is_hub_url, _normalize_env_name + +if TYPE_CHECKING: + from core.containers.runtime import ContainerProvider + from core.http_env_client import HTTPEnvClient + +logger = logging.getLogger(__name__) + + +class AutoEnv: + """ + AutoEnv automatically selects and instantiates the correct environment client + based on environment names or HuggingFace Hub repositories. + + 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: + 1. Check if name is a HuggingFace Hub URL/repo ID + 2. If Hub: download and install the environment package + 3. If local: look up the installed openenv-* package + 4. Import and instantiate the client class + + Example: + >>> # From installed package + >>> env = AutoEnv.from_name("coding-env") + >>> + >>> # From HuggingFace Hub + >>> env = AutoEnv.from_name("meta-pytorch/coding-env") + >>> + >>> # List available environments + >>> AutoEnv.list_environments() + + Note: + AutoEnv is not meant to be instantiated directly. Use the class method + from_name() 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_name() instead." + ) + + @classmethod + def _download_from_hub( + cls, repo_id: str, cache_dir: Optional[Path] = None + ) -> Path: + """ + Download environment from HuggingFace Hub. + + Args: + repo_id: HuggingFace repo ID (e.g., "meta-pytorch/coding-env") + cache_dir: Optional cache directory + + Returns: + Path to downloaded environment directory + + Raises: + ImportError: If huggingface_hub is not installed + ValueError: If download fails + """ + try: + from huggingface_hub import snapshot_download + except ImportError: + raise ImportError( + "HuggingFace Hub support requires huggingface_hub package.\n" + "Install it with: pip install huggingface_hub" + ) + + # Clean up repo_id if it's a full URL + if "huggingface.co" in repo_id: + # Extract org/repo from URL + # https://huggingface.co/meta-pytorch/coding-env -> meta-pytorch/coding-env + parts = repo_id.split("/") + if len(parts) >= 2: + repo_id = f"{parts[-2]}/{parts[-1]}" + + logger.info(f"Downloading environment from HuggingFace Hub: {repo_id}") + + try: + # Download to cache + env_path = snapshot_download( + repo_id=repo_id, + cache_dir=cache_dir or Path(tempfile.gettempdir()) / "openenv_hub_cache", + repo_type="space", # OpenEnv environments are published as Spaces + ) + return Path(env_path) + except Exception as e: + raise ValueError( + f"Failed to download environment from HuggingFace Hub: {repo_id}\n" + f"Error: {e}\n" + f"Make sure the repository exists and is accessible." + ) from e + + @classmethod + def _install_from_path(cls, env_path: Path) -> str: + """ + Install environment package from a local path. + + Args: + env_path: Path to environment directory containing pyproject.toml + + Returns: + Package name that was installed + + Raises: + ValueError: If installation fails + """ + if not (env_path / "pyproject.toml").exists(): + raise ValueError( + f"Environment directory does not contain pyproject.toml: {env_path}" + ) + + logger.info(f"Installing environment from: {env_path}") + + try: + # Install in editable mode + subprocess.run( + ["pip", "install", "-e", str(env_path)], + check=True, + capture_output=True, + text=True, + ) + + # Read package name from pyproject.toml + import toml + + with open(env_path / "pyproject.toml", "r") as f: + pyproject = toml.load(f) + + package_name = pyproject.get("project", {}).get("name") + if not package_name: + raise ValueError("Could not determine package name from pyproject.toml") + + logger.info(f"Successfully installed: {package_name}") + return package_name + + except subprocess.CalledProcessError as e: + raise ValueError( + f"Failed to install environment package from {env_path}\n" + f"Error: {e.stderr}" + ) from e + except Exception as e: + raise ValueError(f"Failed to install environment package: {e}") from e + + @classmethod + def from_name( + cls, + name: str, + base_url: Optional[str] = None, + docker_image: Optional[str] = None, + container_provider: Optional[ContainerProvider] = None, + wait_timeout: float = 30.0, + env_vars: Optional[Dict[str, str]] = None, + **kwargs: Any, + ) -> HTTPEnvClient: + """ + Create an environment client from a name or HuggingFace Hub repository. + + This method automatically: + 1. Checks if the name is a HuggingFace Hub URL/repo ID + 2. If Hub: downloads and installs the environment package + 3. If local: looks up the installed openenv-* package + 4. Imports the client class and instantiates it + + Args: + name: Environment name or HuggingFace Hub repo ID + Examples: + - "coding" / "coding-env" / "coding_env" + - "meta-pytorch/coding-env" (Hub repo ID) + - "https://huggingface.co/meta-pytorch/coding-env" (Hub URL) + base_url: Optional base URL for HTTP connection + docker_image: Optional Docker image name (overrides default) + container_provider: Optional container provider + wait_timeout: Timeout for container startup (seconds) + env_vars: Optional environment variables for the container + **kwargs: Additional arguments passed to the client class + + Returns: + Instance of the environment client class + + Raises: + ValueError: If environment not found or cannot be loaded + ImportError: If environment package is not installed + + Examples: + >>> # From installed package + >>> env = AutoEnv.from_name("coding-env") + >>> + >>> # From HuggingFace Hub + >>> env = AutoEnv.from_name("meta-pytorch/coding-env") + >>> + >>> # With custom Docker image + >>> env = AutoEnv.from_name("coding", docker_image="my-coding-env:v2") + >>> + >>> # With environment variables + >>> env = AutoEnv.from_name( + ... "dipg", + ... env_vars={"DIPG_DATASET_PATH": "/data/dipg"} + ... ) + """ + # Check if it's a HuggingFace Hub URL or repo ID + if _is_hub_url(name): + # Download from Hub and install + env_path = cls._download_from_hub(name) + package_name = cls._install_from_path(env_path) + + # Clear discovery cache to pick up the newly installed package + get_discovery().clear_cache() + + # Extract environment name from package name + # "openenv-coding_env" -> "coding_env" + env_name = package_name.replace("openenv-", "").replace("-", "_") + else: + env_name = name + + # Get environment info from discovery + discovery = get_discovery() + env_info = discovery.get_environment_by_name(env_name) + + if not env_info: + # Environment not found - provide helpful error message + available_envs = discovery.discover() + + if not available_envs: + raise ValueError( + f"No OpenEnv environments found.\n" + f"Install an environment with: pip install openenv-\n" + f"Or specify a HuggingFace Hub repository: AutoEnv.from_name('org/repo')" + ) + + # Try to suggest similar environment names + from difflib import get_close_matches + + env_keys = list(available_envs.keys()) + suggestions = get_close_matches(env_name, env_keys, n=3, cutoff=0.6) + + error_msg = f"Unknown environment '{env_name}'.\n" + if suggestions: + error_msg += f"Did you mean: {', '.join(suggestions)}?\n" + error_msg += f"Available environments: {', '.join(sorted(env_keys))}" + + raise ValueError(error_msg) + + # Get the client class + try: + client_class = env_info.get_client_class() + except ImportError as e: + raise ImportError( + f"Failed to import environment client for '{env_name}'.\n" + f"Package '{env_info.package_name}' appears to be installed but the module cannot be imported.\n" + f"Try reinstalling: pip install --force-reinstall {env_info.package_name}\n" + f"Original error: {e}" + ) from e + + # Determine Docker image to use + if docker_image is None: + docker_image = env_info.default_image + + # Create client instance + try: + if base_url: + # Connect to existing server at URL + return client_class(base_url=base_url, **kwargs) + else: + # Start new Docker container + return client_class.from_docker_image( + image=docker_image, + container_provider=container_provider, + wait_timeout=wait_timeout, + env_vars=env_vars or {}, + **kwargs, + ) + except Exception as e: + raise ValueError( + f"Failed to create environment client for '{env_name}'.\n" + f"Client class: {client_class.__name__}\n" + f"Docker image: {docker_image}\n" + f"Error: {e}" + ) from e + + @classmethod + def get_env_class(cls, name: str): + """ + Get the environment client class without instantiating it. + + Args: + name: Environment name + + Returns: + The environment client class + + Raises: + ValueError: If environment not found + + Examples: + >>> CodingEnv = AutoEnv.get_env_class("coding") + >>> # Now you can instantiate it yourself + >>> env = CodingEnv(base_url="http://localhost:8000") + """ + discovery = get_discovery() + env_info = discovery.get_environment_by_name(name) + + if not env_info: + raise ValueError(f"Unknown environment: {name}") + + return env_info.get_client_class() + + @classmethod + def get_env_info(cls, name: str) -> Dict[str, Any]: + """ + Get detailed information about an environment. + + Args: + name: Environment name + + Returns: + Dictionary with environment metadata + + Raises: + ValueError: If environment not found + + Examples: + >>> info = AutoEnv.get_env_info("coding") + >>> print(info['description']) + 'Coding environment for OpenEnv' + >>> print(info['default_image']) + 'coding-env:latest' + """ + discovery = get_discovery() + env_info = discovery.get_environment_by_name(name) + + if not env_info: + raise ValueError(f"Unknown environment: {name}") + + return { + "env_key": env_info.env_key, + "name": env_info.name, + "package": env_info.package_name, + "version": env_info.version, + "description": env_info.description, + "env_class": env_info.client_class_name, + "action_class": env_info.action_class_name, + "observation_class": env_info.observation_class_name, + "module": env_info.client_module_path, + "default_image": env_info.default_image, + "spec_version": env_info.spec_version, + } + + @classmethod + def list_environments(cls) -> None: + """ + Print a formatted list of all available environments. + + This discovers all installed openenv-* packages and displays + their metadata in a user-friendly format. + + Examples: + >>> AutoEnv.list_environments() + Available OpenEnv Environments: + ---------------------------------------------------------------------- + echo : Echo Environment (v0.1.0) + Package: openenv-echo-env + coding : Coding Environment (v0.1.0) + Package: openenv-coding_env + ---------------------------------------------------------------------- + Total: 2 environments + """ + discovery = get_discovery() + discovery.list_environments() diff --git a/src/envs/coding_env/openenv.yaml b/src/envs/coding_env/openenv.yaml index ba42db55..b5e919b3 100644 --- a/src/envs/coding_env/openenv.yaml +++ b/src/envs/coding_env/openenv.yaml @@ -1,5 +1,5 @@ name: coding_env version: "0.1.0" description: "Coding environment for OpenEnv" -action: CodingAction -observation: CodingObservation +action: CodeAction +observation: CodeObservation diff --git a/src/envs/echo_env/pyproject.toml b/src/envs/echo_env/pyproject.toml index a337f8fa..b7f6a07d 100644 --- a/src/envs/echo_env/pyproject.toml +++ b/src/envs/echo_env/pyproject.toml @@ -35,7 +35,8 @@ dev = [ server = "echo_env.server.app:main" [tool.setuptools] -package-dir = {"" = "."} +packages = ["echo_env", "echo_env.server"] +package-dir = { "echo_env" = ".", "echo_env.server" = "server" } -[tool.setuptools.packages.find] -where = ["."] +[tool.setuptools.package-data] +echo_env = ["**/*.yaml", "**/*.yml"] diff --git a/tests/envs/test_auto_integration.py b/tests/envs/test_auto_integration.py new file mode 100644 index 00000000..7a9dde74 --- /dev/null +++ b/tests/envs/test_auto_integration.py @@ -0,0 +1,281 @@ +# 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. + +""" +Integration tests for AutoEnv and AutoAction +============================================= + +Tests the full integration of package-based discovery with AutoEnv/AutoAction. + +These tests use the actual installed packages (echo_env, coding_env) to verify +the complete flow works end-to-end. +""" + +import pytest +from envs import AutoEnv, AutoAction +from envs._discovery import reset_discovery + + +class TestAutoEnvIntegration: + """Test AutoEnv integration with package discovery.""" + + def setup_method(self): + """Reset discovery before each test to ensure clean state.""" + reset_discovery() + + def test_auto_env_get_env_class(self): + """Test getting environment class by name.""" + # Test with echo environment (should work if echo_env package is installed) + try: + EchoEnv = AutoEnv.get_env_class("echo") + assert EchoEnv.__name__ == "EchoEnv" + assert "echo_env.client" in EchoEnv.__module__ + except (ValueError, ImportError) as e: + # If package not installed or can't be imported, skip test + pytest.skip(f"echo_env package not properly installed: {e}") + + def test_auto_env_get_env_class_flexible_naming(self): + """Test flexible name matching.""" + try: + # All these should work + EchoEnv1 = AutoEnv.get_env_class("echo") + EchoEnv2 = AutoEnv.get_env_class("echo-env") + EchoEnv3 = AutoEnv.get_env_class("echo_env") + + # Should all return the same class + assert EchoEnv1 is EchoEnv2 + assert EchoEnv2 is EchoEnv3 + except (ValueError, ImportError): + pytest.skip("echo_env package not properly installed") + + def test_auto_env_get_env_info(self): + """Test getting environment info.""" + try: + info = AutoEnv.get_env_info("echo") + assert info["name"] == "echo_env" + assert info["env_class"] == "EchoEnv" + assert info["action_class"] == "EchoAction" + assert "description" in info + assert "default_image" in info + assert "package" in info + assert info["package"].startswith("openenv-") + except (ValueError, ImportError): + pytest.skip("echo_env package not properly installed") + + def test_auto_env_list_environments(self, capsys): + """Test listing all environments.""" + AutoEnv.list_environments() + captured = capsys.readouterr() + + assert "Available OpenEnv Environments" in captured.out + # Should show at least the pattern, even if no envs installed + assert "Total:" in captured.out + + def test_auto_env_unknown_environment(self): + """Test error handling for unknown environment.""" + with pytest.raises(ValueError) as exc_info: + AutoEnv.get_env_class("nonexistent-environment") + + assert "Unknown environment" in str(exc_info.value) + + def test_auto_env_get_env_info_unknown(self): + """Test getting info for unknown environment.""" + with pytest.raises(ValueError) as exc_info: + AutoEnv.get_env_info("nonexistent") + + assert "Unknown environment" in str(exc_info.value) + + +class TestAutoActionIntegration: + """Test AutoAction integration with package discovery.""" + + def setup_method(self): + """Reset discovery before each test.""" + reset_discovery() + + def test_auto_action_from_name_simple(self): + """Test getting action class from simple name.""" + try: + EchoAction = AutoAction.from_name("echo") + assert EchoAction.__name__ == "EchoAction" + assert "echo_env" in EchoAction.__module__ + except (ValueError, ImportError): + pytest.skip("echo_env package not properly installed") + + def test_auto_action_from_name_flexible(self): + """Test getting action class with different name formats.""" + try: + # All these should work + Action1 = AutoAction.from_name("echo") + Action2 = AutoAction.from_name("echo-env") + Action3 = AutoAction.from_name("echo_env") + + # Should all return the same class + assert Action1 is Action2 + assert Action2 is Action3 + except (ValueError, ImportError): + pytest.skip("echo_env package not properly installed") + + def test_auto_action_from_env(self): + """Test from_env() alias method.""" + try: + Action1 = AutoAction.from_name("echo") + Action2 = AutoAction.from_env("echo") + + # Should return the same class + assert Action1 is Action2 + except (ValueError, ImportError): + pytest.skip("echo_env package not properly installed") + + def test_auto_action_coding_env(self): + """Test with coding_env if installed.""" + try: + CodeAction = AutoAction.from_name("coding") + assert CodeAction.__name__ == "CodeAction" + assert "coding_env" in CodeAction.__module__ + except ValueError: + pytest.skip("coding_env package not installed") + + def test_auto_action_get_action_info(self): + """Test getting action info.""" + try: + info = AutoAction.get_action_info("echo") + assert info["action_class"] == "EchoAction" + assert info["env_name"] == "echo_env" + assert "package" in info + except (ValueError, ImportError): + pytest.skip("echo_env package not properly installed") + + def test_auto_action_list_actions(self, capsys): + """Test listing all action classes.""" + AutoAction.list_actions() + captured = capsys.readouterr() + + assert "Available Action Classes" in captured.out + assert "Total:" in captured.out + + def test_auto_action_unknown_environment(self): + """Test error handling for unknown environment.""" + with pytest.raises(ValueError) as exc_info: + AutoAction.from_name("nonexistent-environment") + + assert "Unknown environment" in str(exc_info.value) + + +class TestAutoEnvAutoActionTogether: + """Test using AutoEnv and AutoAction together.""" + + def setup_method(self): + """Reset discovery before each test.""" + reset_discovery() + + def test_auto_env_and_action_together(self): + """Test getting both environment and action class.""" + try: + # Get environment class + EchoEnv = AutoEnv.get_env_class("echo") + assert EchoEnv.__name__ == "EchoEnv" + + # Get action class + EchoAction = AutoAction.from_name("echo") + assert EchoAction.__name__ == "EchoAction" + + # Verify they're related + info = AutoEnv.get_env_info("echo") + assert info["action_class"] == "EchoAction" + except (ValueError, ImportError): + pytest.skip("echo_env package not properly installed") + + def test_multiple_environments(self): + """Test working with multiple environments.""" + try: + # Try echo + EchoAction = AutoAction.from_name("echo") + assert EchoAction is not None + + # Try coding (if installed) + try: + CodeAction = AutoAction.from_name("coding") + assert CodeAction is not None + # Should be different classes + assert EchoAction is not CodeAction + except ValueError: + # coding_env not installed, that's ok + pass + + except (ValueError, ImportError): + pytest.skip("No environment packages properly installed") + + def test_action_creation(self): + """Test creating action instances.""" + try: + EchoAction = AutoAction.from_name("echo") + + # Create an action instance + action = EchoAction(message="Hello, World!") + + # Verify it's the right type + assert isinstance(action, EchoAction) + assert hasattr(action, "message") + except (ValueError, ImportError): + pytest.skip("echo_env package not properly installed") + + +class TestDiscoveryPerformance: + """Test discovery caching and performance.""" + + def setup_method(self): + """Reset discovery before each test.""" + reset_discovery() + + def test_discovery_uses_cache(self): + """Test that discovery uses cache on subsequent calls.""" + from envs._discovery import get_discovery + + discovery = get_discovery() + + # First call - should discover + envs1 = discovery.discover(use_cache=False) + + # Second call with cache - should be fast + envs2 = discovery.discover(use_cache=True) + + # Should return the same data (from cache) + assert envs1.keys() == envs2.keys() + + def test_cache_invalidation(self): + """Test that cache can be cleared.""" + from envs._discovery import get_discovery + + discovery = get_discovery() + + # Discover and cache + discovery.discover() + + # Clear cache + discovery.clear_cache() + + # Should rediscover + envs = discovery.discover(use_cache=False) + assert envs is not None + + +class TestHubDetection: + """Test HuggingFace Hub URL detection.""" + + def test_hub_url_detection(self): + """Test that Hub URLs are detected correctly.""" + from envs._discovery import _is_hub_url + + # Hub URLs + assert _is_hub_url("meta-pytorch/coding-env") + assert _is_hub_url("org/repo") + assert _is_hub_url("https://huggingface.co/meta-pytorch/coding-env") + + # Local names + assert not _is_hub_url("coding") + assert not _is_hub_url("coding-env") + assert not _is_hub_url("echo_env") diff --git a/tests/envs/test_discovery.py b/tests/envs/test_discovery.py new file mode 100644 index 00000000..12f2ad9e --- /dev/null +++ b/tests/envs/test_discovery.py @@ -0,0 +1,382 @@ +# 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. + +""" +Unit tests for Package-Based Environment Discovery +=================================================== + +Tests cover: +1. Package discovery using importlib.metadata +2. Manifest loading from package resources +3. Class name inference +4. Cache management +5. Helper functions (_normalize_env_name, _is_hub_url, etc.) +""" + +import pytest +from unittest.mock import Mock, patch, MagicMock +from pathlib import Path + +from envs._discovery import ( + EnvironmentDiscovery, + EnvironmentInfo, + get_discovery, + reset_discovery, + _normalize_env_name, + _is_hub_url, + _infer_class_name, + _create_env_info_from_package, +) + + +class TestEnvironmentInfo: + """Test EnvironmentInfo dataclass and methods.""" + + def test_environment_info_creation(self): + """Test creating EnvironmentInfo instance.""" + env_info = EnvironmentInfo( + env_key="echo", + name="echo_env", + package_name="openenv-echo-env", + version="0.1.0", + description="Echo environment", + client_module_path="echo_env.client", + client_class_name="EchoEnv", + action_class_name="EchoAction", + observation_class_name="EchoObservation", + default_image="echo-env:latest" + ) + + assert env_info.env_key == "echo" + assert env_info.name == "echo_env" + assert env_info.package_name == "openenv-echo-env" + assert env_info.client_class_name == "EchoEnv" + assert env_info.default_image == "echo-env:latest" + + +class TestHelperFunctions: + """Test helper functions.""" + + def test_normalize_env_name_simple(self): + """Test normalizing simple names.""" + assert _normalize_env_name("echo") == "echo_env" + assert _normalize_env_name("coding") == "coding_env" + + def test_normalize_env_name_with_suffix(self): + """Test normalizing names with -env suffix.""" + assert _normalize_env_name("echo-env") == "echo_env" + assert _normalize_env_name("coding-env") == "coding_env" + + def test_normalize_env_name_with_underscore(self): + """Test normalizing names with _env suffix.""" + assert _normalize_env_name("echo_env") == "echo_env" + assert _normalize_env_name("coding_env") == "coding_env" + + def test_is_hub_url_with_slash(self): + """Test Hub URL detection with org/repo pattern.""" + assert _is_hub_url("meta-pytorch/coding-env") + assert _is_hub_url("myorg/myenv") + + def test_is_hub_url_with_domain(self): + """Test Hub URL detection with full URL.""" + assert _is_hub_url("https://huggingface.co/meta-pytorch/coding-env") + assert _is_hub_url("huggingface.co/spaces/myenv") + + def test_is_hub_url_local(self): + """Test that local names are not detected as Hub URLs.""" + assert not _is_hub_url("echo") + assert not _is_hub_url("coding-env") + assert not _is_hub_url("echo_env") + + def test_infer_class_name_client(self): + """Test inferring client class names.""" + assert _infer_class_name("echo_env", "client") == "EchoEnv" + assert _infer_class_name("coding_env", "client") == "CodingEnv" + assert _infer_class_name("browser_gym_env", "client") == "BrowserGymEnv" + + def test_infer_class_name_action(self): + """Test inferring action class names.""" + assert _infer_class_name("echo_env", "action") == "EchoAction" + assert _infer_class_name("coding_env", "action") == "CodingAction" + + def test_infer_class_name_observation(self): + """Test inferring observation class names.""" + assert _infer_class_name("echo_env", "observation") == "EchoObservation" + assert _infer_class_name("coding_env", "observation") == "CodingObservation" + + +class TestCreateEnvInfoFromPackage: + """Test creating EnvironmentInfo from package data.""" + + @patch('envs._discovery._load_manifest_from_package') + def test_create_env_info_with_manifest(self, mock_load_manifest): + """Test creating env info when manifest exists.""" + # Mock manifest data + mock_load_manifest.return_value = { + "name": "echo_env", + "version": "0.1.0", + "description": "Echo environment for OpenEnv", + "spec_version": 1, + } + + env_info = _create_env_info_from_package( + package_name="openenv-echo-env", + module_name="echo_env", + version="0.1.0" + ) + + assert env_info is not None + assert env_info.env_key == "echo" + assert env_info.name == "echo_env" + assert env_info.package_name == "openenv-echo-env" + assert env_info.version == "0.1.0" + assert env_info.client_class_name == "EchoEnv" + assert env_info.action_class_name == "EchoAction" + + @patch('envs._discovery._load_manifest_from_package') + def test_create_env_info_with_custom_class_names(self, mock_load_manifest): + """Test creating env info with custom class names from manifest.""" + # Mock manifest with custom class names + mock_load_manifest.return_value = { + "name": "coding_env", + "version": "0.1.0", + "description": "Coding environment", + "action": "CodeAction", # Custom name + "observation": "CodeObservation", # Custom name + } + + env_info = _create_env_info_from_package( + package_name="openenv-coding_env", + module_name="coding_env", + version="0.1.0" + ) + + assert env_info.action_class_name == "CodeAction" + assert env_info.observation_class_name == "CodeObservation" + + @patch('envs._discovery._load_manifest_from_package') + def test_create_env_info_without_manifest(self, mock_load_manifest): + """Test creating env info when no manifest exists (uses conventions).""" + mock_load_manifest.return_value = None + + env_info = _create_env_info_from_package( + package_name="openenv-test-env", + module_name="test_env", + version="1.0.0" + ) + + assert env_info is not None + assert env_info.env_key == "test" + assert env_info.name == "test_env" + assert env_info.client_class_name == "TestEnv" + assert env_info.action_class_name == "TestAction" + + +class TestEnvironmentDiscovery: + """Test EnvironmentDiscovery class.""" + + @patch('importlib.metadata.distributions') + @patch('envs._discovery._create_env_info_from_package') + def test_discover_installed_packages(self, mock_create_info, mock_distributions): + """Test discovering installed packages.""" + # Mock distribution objects + mock_dist1 = Mock() + mock_dist1.metadata = {"Name": "openenv-echo-env"} + mock_dist1.version = "0.1.0" + + mock_dist2 = Mock() + mock_dist2.metadata = {"Name": "openenv-coding_env"} + mock_dist2.version = "0.2.0" + + mock_dist3 = Mock() + mock_dist3.metadata = {"Name": "openenv-core"} # Should be filtered out + mock_dist3.version = "1.0.0" + + mock_distributions.return_value = [mock_dist1, mock_dist2, mock_dist3] + + # Mock env info creation + def create_info_side_effect(package_name, module_name, version): + return EnvironmentInfo( + env_key=module_name.replace("_env", ""), + name=f"{module_name}", + package_name=package_name, + version=version, + description=f"{module_name} environment", + client_module_path=f"{module_name}.client", + client_class_name=f"{module_name.replace('_env', '').capitalize()}Env", + action_class_name=f"{module_name.replace('_env', '').capitalize()}Action", + observation_class_name=f"{module_name.replace('_env', '').capitalize()}Observation", + default_image=f"{module_name.replace('_', '-')}:latest" + ) + + mock_create_info.side_effect = create_info_side_effect + + discovery = EnvironmentDiscovery() + envs = discovery._discover_installed_packages() + + # Should discover 2 environments (not openenv-core) + assert len(envs) == 2 + assert "echo" in envs + assert "coding" in envs + + def test_get_environment(self): + """Test getting a specific environment.""" + discovery = EnvironmentDiscovery() + + # Mock the discover method + with patch.object(discovery, 'discover') as mock_discover: + mock_discover.return_value = { + "echo": EnvironmentInfo( + env_key="echo", + name="echo_env", + package_name="openenv-echo-env", + version="0.1.0", + description="Echo", + client_module_path="echo_env.client", + client_class_name="EchoEnv", + action_class_name="EchoAction", + observation_class_name="EchoObservation", + default_image="echo-env:latest" + ) + } + + env = discovery.get_environment("echo") + assert env is not None + assert env.env_key == "echo" + + def test_get_environment_not_found(self): + """Test getting a non-existent environment.""" + discovery = EnvironmentDiscovery() + + with patch.object(discovery, 'discover') as mock_discover: + mock_discover.return_value = {} + + env = discovery.get_environment("nonexistent") + assert env is None + + def test_get_environment_by_name_flexible(self): + """Test getting environment with flexible name matching.""" + discovery = EnvironmentDiscovery() + + mock_env = EnvironmentInfo( + env_key="echo", + name="echo_env", + package_name="openenv-echo-env", + version="0.1.0", + description="Echo", + client_module_path="echo_env.client", + client_class_name="EchoEnv", + action_class_name="EchoAction", + observation_class_name="EchoObservation", + default_image="echo-env:latest" + ) + + with patch.object(discovery, 'discover') as mock_discover: + mock_discover.return_value = {"echo": mock_env} + + # All these should work + assert discovery.get_environment_by_name("echo") is not None + assert discovery.get_environment_by_name("echo-env") is not None + assert discovery.get_environment_by_name("echo_env") is not None + + def test_cache_management(self): + """Test cache loading and saving.""" + discovery = EnvironmentDiscovery() + + # Create mock environment + mock_env = EnvironmentInfo( + env_key="test", + name="test_env", + package_name="openenv-test", + version="1.0.0", + description="Test", + client_module_path="test_env.client", + client_class_name="TestEnv", + action_class_name="TestAction", + observation_class_name="TestObservation", + default_image="test-env:latest" + ) + + envs = {"test": mock_env} + + # Test saving cache + discovery._save_cache(envs) + assert discovery._cache_file.exists() + + # Test loading cache + loaded = discovery._load_cache() + assert loaded is not None + assert "test" in loaded + + # Clean up + discovery.clear_cache() + assert not discovery._cache_file.exists() + + +class TestGlobalDiscovery: + """Test global discovery instance management.""" + + def test_get_discovery_singleton(self): + """Test that get_discovery returns singleton.""" + reset_discovery() + + discovery1 = get_discovery() + discovery2 = get_discovery() + + assert discovery1 is discovery2 + + def test_reset_discovery(self): + """Test resetting global discovery instance.""" + discovery1 = get_discovery() + + reset_discovery() + + discovery2 = get_discovery() + + # Should be different instances after reset + assert discovery1 is not discovery2 + + +class TestListEnvironments: + """Test list_environments output.""" + + def test_list_environments_with_envs(self, capsys): + """Test listing when environments are found.""" + discovery = EnvironmentDiscovery() + + mock_envs = { + "echo": EnvironmentInfo( + env_key="echo", + name="echo_env", + package_name="openenv-echo-env", + version="0.1.0", + description="Echo environment", + client_module_path="echo_env.client", + client_class_name="EchoEnv", + action_class_name="EchoAction", + observation_class_name="EchoObservation", + default_image="echo-env:latest" + ) + } + + with patch.object(discovery, 'discover', return_value=mock_envs): + discovery.list_environments() + + captured = capsys.readouterr() + assert "Available OpenEnv Environments" in captured.out + assert "echo" in captured.out + assert "Total: 1 environments" in captured.out + + def test_list_environments_empty(self, capsys): + """Test listing when no environments are found.""" + discovery = EnvironmentDiscovery() + + with patch.object(discovery, 'discover', return_value={}): + discovery.list_environments() + + captured = capsys.readouterr() + assert "No OpenEnv environments found" in captured.out + assert "pip install openenv-" in captured.out