openenv/coding_env
21
1# Copyright (c) Meta Platforms, Inc. and affiliates.2# All rights reserved.3#4# This source code is licensed under the BSD-style license found in the5# LICENSE file in the root directory of this source tree.6 7"""8Environment Auto-Discovery System9==================================10 11This module provides automatic discovery of OpenEnv environments by:121. Discovering installed openenv-* packages using importlib.metadata132. Loading manifests (openenv.yaml) from package resources143. Caching results for performance154. Supporting HuggingFace Hub downloads16 17This enables AutoEnv to work without coupling to src/envs/ directory.18"""19 20import importlib21import importlib.metadata22import importlib.resources23import json24import logging25import re26import tempfile27from dataclasses import asdict, dataclass28from pathlib import Path29from typing import Any, Dict, Optional, Type30 31import yaml32 33logger = logging.getLogger(__name__)34 35 36@dataclass37class EnvironmentInfo:38 """39 Rich information about a discovered environment.40 41 Attributes:42 env_key: Environment key (e.g., "echo", "coding")43 name: Full environment name (e.g., "echo_env")44 package_name: Package name (e.g., "openenv-echo_env")45 version: Version string46 description: Human-readable description47 client_module_path: Full module path to client (e.g., "echo_env.client")48 client_class_name: Client class name (e.g., "EchoEnv")49 action_class_name: Action class name (e.g., "EchoAction")50 observation_class_name: Observation class name (e.g., "EchoObservation")51 default_image: Default Docker image name (e.g., "echo-env:latest")52 spec_version: OpenEnv spec version (from openenv.yaml)53 manifest: Original manifest data54 """55 56 env_key: str57 name: str58 package_name: str59 version: str60 description: str61 client_module_path: str62 client_class_name: str63 action_class_name: str64 observation_class_name: str65 default_image: str66 spec_version: Optional[int] = None67 manifest: Optional[Dict[str, Any]] = None68 69 def get_client_class(self) -> Type:70 """71 Dynamically import and return the client class.72 73 Returns:74 Client class (e.g., EchoEnv)75 76 Raises:77 ImportError: If module or class cannot be imported78 """79 try:80 module = importlib.import_module(self.client_module_path)81 return getattr(module, self.client_class_name)82 except ImportError as e:83 raise ImportError(84 f"Failed to import {self.client_class_name} from {self.client_module_path}: {e}\n"85 f"Make sure the package '{self.package_name}' is installed: "86 f"pip install {self.package_name}"87 ) from e88 except AttributeError as e:89 raise ImportError(90 f"Class {self.client_class_name} not found in {self.client_module_path}: {e}"91 ) from e92 93 def get_action_class(self) -> Type:94 """95 Dynamically import and return the action class.96 97 Returns:98 Action class (e.g., EchoAction)99 100 Raises:101 ImportError: If module or class cannot be imported102 """103 try:104 module = importlib.import_module(self.client_module_path)105 return getattr(module, self.action_class_name)106 except ImportError as e:107 raise ImportError(108 f"Failed to import {self.action_class_name} from {self.client_module_path}: {e}\n"109 f"Make sure the package '{self.package_name}' is installed: "110 f"pip install {self.package_name}"111 ) from e112 except AttributeError as e:113 raise ImportError(114 f"Class {self.action_class_name} not found in {self.client_module_path}: {e}"115 ) from e116 117 def get_observation_class(self) -> Type:118 """119 Dynamically import and return the observation class.120 121 Returns:122 Observation class (e.g., EchoObservation)123 124 Raises:125 ImportError: If module or class cannot be imported126 """127 try:128 module = importlib.import_module(self.client_module_path)129 return getattr(module, self.observation_class_name)130 except ImportError as e:131 raise ImportError(132 f"Failed to import {self.observation_class_name} from {self.client_module_path}: {e}\n"133 f"Make sure the package '{self.package_name}' is installed: "134 f"pip install {self.package_name}"135 ) from e136 except AttributeError as e:137 raise ImportError(138 f"Class {self.observation_class_name} not found in {self.client_module_path}: {e}"139 ) from e140 141 142def _normalize_env_name(name: str) -> str:143 """144 Normalize environment name to standard format.145 146 Args:147 name: Input name (e.g., "echo", "echo-env", "echo_env")148 149 Returns:150 Normalized name (e.g., "echo_env")151 152 Examples:153 >>> _normalize_env_name("echo")154 'echo_env'155 >>> _normalize_env_name("echo-env")156 'echo_env'157 >>> _normalize_env_name("echo_env")158 'echo_env'159 """160 # Remove common suffixes161 name = re.sub(r"[-_]env$", "", name)162 # Convert hyphens to underscores163 name = name.replace("-", "_")164 # Add _env suffix if not present165 if not name.endswith("_env"):166 name = f"{name}_env"167 return name168 169 170def _is_hub_url(name: str) -> bool:171 """172 Check if name is a HuggingFace Hub URL or repo ID.173 174 Args:175 name: Input name176 177 Returns:178 True if it looks like a Hub URL179 180 Examples:181 >>> _is_hub_url("meta-pytorch/echo_env")182 True183 >>> _is_hub_url("https://huggingface.co/meta-pytorch/echo_env")184 True185 >>> _is_hub_url("echo")186 False187 """188 # Contains org/repo pattern or huggingface.co domain189 return "/" in name or "huggingface.co" in name190 191 192def _infer_class_name(env_name: str, class_type: str) -> str:193 """194 Infer class name from environment name using simple conventions.195 196 Args:197 env_name: Environment name (e.g., "echo_env")198 class_type: Type of class ("client", "action", "observation")199 200 Returns:201 Inferred class name202 203 Examples:204 >>> _infer_class_name("echo_env", "client")205 'EchoEnv'206 >>> _infer_class_name("echo_env", "action")207 'EchoAction'208 """209 # Remove _env suffix for base name210 base_name = env_name.replace("_env", "")211 212 # Convert to PascalCase213 pascal_name = "".join(word.capitalize() for word in base_name.split("_"))214 215 # Add suffix based on type216 if class_type == "client":217 return f"{pascal_name}Env"218 elif class_type == "action":219 return f"{pascal_name}Action"220 elif class_type == "observation":221 return f"{pascal_name}Observation"222 else:223 raise ValueError(f"Unknown class type: {class_type}")224 225 226def _load_manifest_from_package(227 package_name: str, module_name: str228) -> Optional[Dict[str, Any]]:229 """230 Load openenv.yaml manifest from an installed package.231 232 Args:233 package_name: Package name (e.g., "openenv-echo_env")234 module_name: Module name (e.g., "echo_env")235 236 Returns:237 Parsed manifest dictionary, or None if not found238 239 """240 try:241 # Try to read openenv.yaml from package242 if hasattr(importlib.resources, "files"):243 # Python 3.9+244 package_files = importlib.resources.files(module_name)245 if (package_files / "openenv.yaml").is_file():246 manifest_text = (package_files / "openenv.yaml").read_text()247 return yaml.safe_load(manifest_text)248 else:249 # Python 3.7-3.8 fallback250 with importlib.resources.open_text(module_name, "openenv.yaml") as f:251 return yaml.safe_load(f)252 except (FileNotFoundError, ModuleNotFoundError, AttributeError):253 logger.debug(f"No openenv.yaml found in {module_name}")254 return None255 except Exception as e:256 logger.warning(f"Failed to load openenv.yaml from {module_name}: {e}")257 return None258 259 260def _create_env_info_from_package(261 package_name: str, module_name: str, version: str262) -> Optional[EnvironmentInfo]:263 """264 Create EnvironmentInfo from an installed package.265 266 Args:267 package_name: Package name (e.g., "openenv-echo_env")268 module_name: Module name (e.g., "echo_env")269 version: Package version270 271 Returns:272 EnvironmentInfo instance, or None if invalid273 """274 # Load manifest275 manifest = _load_manifest_from_package(package_name, module_name)276 277 # Get environment name278 if manifest and "name" in manifest:279 env_name = manifest["name"]280 else:281 # Infer from module name282 env_name = module_name283 284 # Normalize to ensure _env suffix285 if not env_name.endswith("_env"):286 env_name = f"{env_name}_env"287 288 # Determine env_key (e.g., "echo_env" → "echo")289 env_key = env_name.replace("_env", "") if env_name.endswith("_env") else env_name290 291 # Get description292 description = (293 manifest.get("description", f"{env_name} environment")294 if manifest295 else f"{env_name} environment"296 )297 298 # Get spec version299 spec_version = manifest.get("spec_version") if manifest else None300 301 # Determine class names302 # Check if manifest has custom class names (custom format)303 if manifest and "action" in manifest and "observation" in manifest:304 # Custom format (like coding_env)305 client_class_name = _infer_class_name(env_name, "client")306 action_class_name = manifest.get(307 "action", _infer_class_name(env_name, "action")308 )309 observation_class_name = manifest.get(310 "observation", _infer_class_name(env_name, "observation")311 )312 else:313 # Use conventions314 client_class_name = _infer_class_name(env_name, "client")315 action_class_name = _infer_class_name(env_name, "action")316 observation_class_name = _infer_class_name(env_name, "observation")317 318 # Module path is just module_name.client319 client_module_path = f"{module_name}.client"320 321 # Determine default Docker image name322 image_name = env_name.replace("_", "-")323 default_image = f"{image_name}:latest"324 325 return EnvironmentInfo(326 env_key=env_key,327 name=env_name,328 package_name=package_name,329 version=version,330 description=description,331 client_module_path=client_module_path,332 client_class_name=client_class_name,333 action_class_name=action_class_name,334 observation_class_name=observation_class_name,335 default_image=default_image,336 spec_version=spec_version,337 manifest=manifest,338 )339 340 341class EnvironmentDiscovery:342 """343 Auto-discovery system for OpenEnv environments using installed packages.344 345 This class discovers installed openenv-* packages and loads their metadata.346 """347 348 def __init__(self):349 """Initialize discovery system."""350 self._cache: Optional[Dict[str, EnvironmentInfo]] = None351 self._cache_file = Path(tempfile.gettempdir()) / "openenv_discovery_cache.json"352 353 def _discover_installed_packages(self) -> Dict[str, EnvironmentInfo]:354 """355 Discover all installed openenv-* packages.356 357 Returns:358 Dictionary mapping env_key to EnvironmentInfo359 """360 environments = {}361 362 # Invalidate import caches to ensure we pick up newly installed packages363 importlib.invalidate_caches()364 365 # Get all installed packages366 try:367 distributions = importlib.metadata.distributions()368 except Exception as e:369 logger.warning(f"Failed to get installed packages: {e}")370 return environments371 372 # Filter for openenv-* packages (exclude openenv-core)373 for dist in distributions:374 package_name = dist.metadata["Name"]375 376 if not package_name.startswith("openenv-"):377 continue378 379 if package_name == "openenv-core":380 continue381 382 # Get module name (e.g., "openenv-echo_env" → "echo_env")383 module_name = package_name.replace("openenv-", "").replace("-", "_")384 385 # Get version386 version = dist.version387 388 try:389 # Create environment info390 env_info = _create_env_info_from_package(391 package_name, module_name, version392 )393 394 if env_info:395 environments[env_info.env_key] = env_info396 logger.debug(397 f"Discovered environment: {env_info.env_key} ({package_name})"398 )399 400 except Exception as e:401 logger.warning(f"Failed to load environment from {package_name}: {e}")402 continue403 404 return environments405 406 def _load_cache(self) -> Optional[Dict[str, EnvironmentInfo]]:407 """408 Load cached discovery results.409 410 Returns:411 Dictionary of env_key -> EnvironmentInfo, or None if cache invalid412 """413 if not self._cache_file.exists():414 return None415 416 try:417 with open(self._cache_file, "r") as f:418 cache_data = json.load(f)419 420 # Reconstruct EnvironmentInfo objects421 cache = {}422 for env_key, env_data in cache_data.items():423 cache[env_key] = EnvironmentInfo(**env_data)424 425 return cache426 except Exception as e:427 logger.warning(f"Failed to load discovery cache: {e}")428 return None429 430 def _save_cache(self, environments: Dict[str, EnvironmentInfo]) -> None:431 """432 Save discovery results to cache.433 434 Args:435 environments: Dictionary of env_key -> EnvironmentInfo436 """437 try:438 cache_data = {}439 for env_key, env_info in environments.items():440 cache_data[env_key] = asdict(env_info)441 442 with open(self._cache_file, "w") as f:443 json.dump(cache_data, f, indent=2)444 445 except Exception as e:446 logger.warning(f"Failed to save discovery cache: {e}")447 448 def discover(self, use_cache: bool = True) -> Dict[str, EnvironmentInfo]:449 """450 Discover all installed OpenEnv environments.451 452 Args:453 use_cache: If True, try to load from cache first454 455 Returns:456 Dictionary mapping env_key to EnvironmentInfo457 458 Examples:459 >>> discovery = EnvironmentDiscovery()460 >>> envs = discovery.discover()461 >>> print(envs.keys())462 dict_keys(['echo', 'coding', ...])463 """464 # Try to load from memory cache first465 if use_cache and self._cache is not None:466 return self._cache467 468 # Try to load from file cache469 if use_cache:470 cached = self._load_cache()471 if cached is not None:472 self._cache = cached473 return self._cache474 475 # Discover from installed packages476 environments = self._discover_installed_packages()477 478 # Save to cache479 self._save_cache(environments)480 self._cache = environments481 482 return environments483 484 def get_environment(self, env_key: str) -> Optional[EnvironmentInfo]:485 """486 Get information about a specific environment.487 488 Args:489 env_key: Environment key (e.g., "echo", "coding")490 491 Returns:492 EnvironmentInfo if found, None otherwise493 494 Examples:495 >>> discovery = EnvironmentDiscovery()496 >>> env = discovery.get_environment("echo")497 >>> print(env.client_class_name)498 'EchoEnv'499 """500 environments = self.discover()501 return environments.get(env_key)502 503 def get_environment_by_name(self, name: str) -> Optional[EnvironmentInfo]:504 """505 Get environment info by flexible name matching.506 507 Args:508 name: Environment name (e.g., "echo", "echo-env", "echo_env")509 510 Returns:511 EnvironmentInfo if found, None otherwise512 """513 # Normalize name to env_key514 normalized = _normalize_env_name(name)515 env_key = normalized.replace("_env", "")516 517 return self.get_environment(env_key)518 519 def list_environments(self) -> None:520 """521 Print a formatted list of all discovered environments.522 523 Examples:524 >>> discovery = EnvironmentDiscovery()525 >>> discovery.list_environments()526 Available OpenEnv Environments:527 ----------------------------------------------------------------------528 echo : Echo Environment (v0.1.0) - openenv-echo_env529 coding : Coding Environment (v0.1.0) - openenv-coding_env530 ...531 """532 environments = self.discover()533 534 print("Available OpenEnv Environments:")535 print("-" * 70)536 537 if not environments:538 print(" No OpenEnv environments found.")539 print(" Install environments with: pip install openenv-<env-name>")540 else:541 for env_key in sorted(environments.keys()):542 env = environments[env_key]543 print(f" {env_key:<15}: {env.description} (v{env.version})")544 print(f" Package: {env.package_name}")545 546 print("-" * 70)547 print(f"Total: {len(environments)} environments")548 549 def clear_cache(self) -> None:550 """Clear the discovery cache."""551 if self._cache_file.exists():552 self._cache_file.unlink()553 self._cache = None554 555 556# Global discovery instance557_global_discovery: Optional[EnvironmentDiscovery] = None558 559 560def get_discovery() -> EnvironmentDiscovery:561 """562 Get or create the global discovery instance.563 564 Returns:565 Global EnvironmentDiscovery instance566 567 Examples:568 >>> discovery = get_discovery()569 >>> envs = discovery.discover()570 """571 global _global_discovery572 573 if _global_discovery is None:574 _global_discovery = EnvironmentDiscovery()575 576 return _global_discovery577 578 579def reset_discovery() -> None:580 """Reset the global discovery instance (useful for testing)."""581 global _global_discovery582 if _global_discovery is not None:583 _global_discovery.clear_cache()584 _global_discovery = None585