147 lines
5.5 KiB
Python
147 lines
5.5 KiB
Python
"""Configuration utilities for loading YAML configs and merging with CLI args."""
|
|
|
|
import os
|
|
import sys
|
|
from typing import Dict, Any, Type, TypeVar
|
|
import yaml
|
|
from dataclasses import fields, is_dataclass
|
|
|
|
from experiment_logger.unified_logger import get_logger
|
|
|
|
log = get_logger()
|
|
|
|
T = TypeVar("T")
|
|
|
|
|
|
def load_yaml_config(config_path: str) -> Dict[str, Any]:
|
|
"""Load configuration from YAML file."""
|
|
if not os.path.exists(config_path):
|
|
raise FileNotFoundError(f"Config file not found: {config_path}")
|
|
|
|
with open(config_path, "r") as f:
|
|
config = yaml.safe_load(f)
|
|
|
|
if config is None:
|
|
return {}
|
|
|
|
log.info(f"Loaded configuration from: {config_path}")
|
|
return config
|
|
|
|
|
|
def save_yaml_config(config: Dict[str, Any], config_path: str):
|
|
"""Save configuration to YAML file."""
|
|
os.makedirs(os.path.dirname(config_path), exist_ok=True)
|
|
|
|
with open(config_path, "w") as f:
|
|
yaml.dump(config, f, default_flow_style=False, indent=2, sort_keys=False)
|
|
|
|
log.info(f"Saved configuration to: {config_path}")
|
|
|
|
|
|
def dataclass_from_dict(cls: Type[T], config_dict: Dict[str, Any]) -> T:
|
|
"""Create dataclass instance from dictionary, handling type conversions."""
|
|
if not is_dataclass(cls):
|
|
raise ValueError(f"{cls} is not a dataclass")
|
|
|
|
# Get field names and types
|
|
field_map = {f.name: f for f in fields(cls)} # type: ignore
|
|
|
|
# Filter config to only include valid fields
|
|
filtered_config: Dict[str, Any] = {}
|
|
for key, value in config_dict.items():
|
|
if key in field_map:
|
|
field = field_map[key]
|
|
# Handle type conversion if needed
|
|
try:
|
|
# Handle None values and optional types
|
|
if value is None:
|
|
filtered_config[key] = None
|
|
elif hasattr(field.type, "__origin__") and field.type.__origin__ is type(None):
|
|
# Optional type (Union[X, None])
|
|
filtered_config[key] = value
|
|
else:
|
|
# Try to convert to the expected type
|
|
if field.type is bool and isinstance(value, str):
|
|
filtered_config[key] = value.lower() in ("true", "1", "yes", "on")
|
|
else:
|
|
filtered_config[key] = field.type(value) if value is not None else None # type: ignore
|
|
except (ValueError, TypeError) as e:
|
|
log.warning(f"Could not convert {key}={value} to {field.type}: {e}")
|
|
filtered_config[key] = value
|
|
else:
|
|
log.warning(f"Unknown configuration parameter: {key}")
|
|
|
|
return cls(**filtered_config)
|
|
|
|
|
|
def merge_config_with_cli(config_class: Type[T], config_file: str | None = None) -> T:
|
|
"""Merge YAML config with CLI arguments, with CLI taking precedence.
|
|
|
|
Args:
|
|
config_class: Dataclass type to create
|
|
config_file: Path to YAML config file (optional)
|
|
|
|
Returns:
|
|
Instance of config_class with merged configuration
|
|
"""
|
|
# Parse CLI args first to get the default/CLI values
|
|
import tyro
|
|
|
|
# Check if --config is in sys.argv and extract it
|
|
extracted_config_file = config_file
|
|
if "--config" in sys.argv:
|
|
config_idx = sys.argv.index("--config")
|
|
if config_idx + 1 < len(sys.argv):
|
|
extracted_config_file = sys.argv[config_idx + 1]
|
|
# Remove from sys.argv so tyro doesn't see it
|
|
sys.argv.pop(config_idx) # Remove --config
|
|
sys.argv.pop(config_idx) # Remove config file path
|
|
|
|
# Load YAML config if available
|
|
yaml_config = {}
|
|
if extracted_config_file and os.path.exists(extracted_config_file):
|
|
yaml_config = load_yaml_config(extracted_config_file)
|
|
log.info(f"Merging YAML config from {extracted_config_file} with CLI args")
|
|
elif extracted_config_file:
|
|
log.warning(f"Config file not found: {extracted_config_file}, using CLI args only")
|
|
|
|
# Create default instance to know what the defaults are
|
|
default_instance = config_class()
|
|
default_dict = {f.name: getattr(default_instance, f.name) for f in fields(config_class)} # type: ignore
|
|
|
|
# Parse CLI args
|
|
cli_instance = tyro.cli(config_class)
|
|
cli_dict = {f.name: getattr(cli_instance, f.name) for f in fields(config_class)} # type: ignore
|
|
|
|
# Merge configs: YAML as base, CLI overrides non-default values
|
|
final_config = {}
|
|
|
|
for field in fields(config_class): # type: ignore
|
|
field_name = field.name
|
|
default_value = default_dict[field_name]
|
|
yaml_value = yaml_config.get(field_name, default_value)
|
|
cli_value = cli_dict[field_name]
|
|
|
|
# Use CLI value if it's different from default, otherwise use YAML value
|
|
if cli_value != default_value:
|
|
final_config[field_name] = cli_value
|
|
if yaml_value != default_value and yaml_value != cli_value:
|
|
log.info(f"CLI override: {field_name}={cli_value} (YAML had {yaml_value})")
|
|
else:
|
|
final_config[field_name] = yaml_value
|
|
if yaml_value != default_value:
|
|
log.info(f"YAML config: {field_name}={yaml_value}")
|
|
|
|
return config_class(**final_config)
|
|
|
|
|
|
def print_config(config: Any, title: str = "Configuration"):
|
|
"""Pretty print configuration."""
|
|
log.info(f"{title}:")
|
|
if is_dataclass(config):
|
|
for field in fields(config):
|
|
value = getattr(config, field.name)
|
|
log.info(f" {field.name}: {value}")
|
|
else:
|
|
for key, value in vars(config).items():
|
|
log.info(f" {key}: {value}")
|