feat: Updated BaseConfig class for non primitive fields

src/oumi/core/configs/base_config.py

@@ -13,9 +13,11 @@
 # limitations under the License.
 
 import dataclasses
+import inspect
 import logging
 import re
 from collections.abc import Iterator
+from enum import Enum
 from io import StringIO
 from pathlib import Path
 from typing import Any, Optional, TypeVar, Union, cast
@@ -28,6 +30,74 @@
 
 _CLI_IGNORED_PREFIXES = ["--local-rank"]
 
+# Set of primitive types that OmegaConf can handle directly
+_PRIMITIVE_TYPES = {str, int, float, bool, type(None), bytes, Path, Enum}
+
+
+def _is_primitive_type(value: Any) -> bool:
+    """Check if a value is of a primitive type that OmegaConf can handle."""
+    return (
+        isinstance(value, (str, int, float, bool, bytes))
+        or value is None
+        or isinstance(value, Path)
+        or isinstance(value, Enum)
+    )
+
+
+def _handle_non_primitives(config: Any, removed_paths, path: str = "") -> Any:
+    """Recursively process config object to handle non-primitive values.
+
+    Args:
+        config: The config object to process
+        removed_paths: Set to track paths of removed non-primitive values
+        path: The current path in the config (for logging)
+
+    Returns:
+        The processed config with non-primitive values removed
+    """
+    if isinstance(config, list):
+        return [
+            _handle_non_primitives(item, removed_paths, f"{path}[{i}]")
+            for i, item in enumerate(config)
+        ]
+
+    if isinstance(config, dict):
+        result = {}
+        for key, value in config.items():
+            current_path = f"{path}.{key}" if path else key
+            if _is_primitive_type(value):
+                result[key] = value
+            else:
+                # Recursively process nested dictionaries and other non-primitive values
+                processed_value = _handle_non_primitives(
+                    value, removed_paths, current_path
+                )
+                if processed_value is not None:
+                    result[key] = processed_value
+                else:
+                    removed_paths.add(current_path)
+                    result[key] = None
+        return result
+
+    if _is_primitive_type(config):
+        return config
+
+    # Try to convert functions to their source code
+    if callable(config):
+        try:
+            # Lambda functions and built-in functions can't have source extracted
+            source = inspect.getsource(config)
+            # Only return source if we successfully got it
+            return source
+        except (TypeError, OSError):
+            # Can't get source for lambdas, built-ins, or C extensions
+            removed_paths.add(path)
+            return None
+
+    # For any other type, remove it and track the path
+    removed_paths.add(path)
+    return None
+
 
 def _filter_ignored_args(arg_list: list[str]) -> list[str]:
     """Filters out ignored CLI arguments."""
@@ -57,8 +127,31 @@ def _read_config_without_interpolation(config_path: str) -> str:
 @dataclasses.dataclass
 class BaseConfig:
     def to_yaml(self, config_path: Union[str, Path, StringIO]) -> None:
-        """Saves the configuration to a YAML file."""
-        OmegaConf.save(config=self, f=config_path)
+        """Saves the configuration to a YAML file.
+
+        Non-primitive values are removed and warnings are logged.
+
+        Args:
+            config_path: Path to save the config to
+        """
+        # Convert the dataclass to an OmegaConf structure first
+        omega_config = OmegaConf.structured(self)
+        config_dict = OmegaConf.to_container(omega_config, resolve=True)
+        removed_paths = set()
+        processed_config = _handle_non_primitives(
+            config_dict, removed_paths=removed_paths
+        )
+
+        # Log warnings for removed values
+        if removed_paths:
+            logger = logging.getLogger(__name__)
+            logger.warning(
+                "The following non-primitive values were removed from the config "
+                "as they cannot be saved to YAML:\n"
+                + "\n".join(f"- {path}" for path in sorted(removed_paths))
+            )
+
+        OmegaConf.save(config=processed_config, f=config_path)
 
     @classmethod
     def from_yaml(
@@ -182,7 +275,9 @@ def print_config(self, logger: Optional[logging.Logger] = None) -> None:
         if logger is None:
             logger = logging.getLogger(__name__)
 
-        config_yaml = OmegaConf.to_yaml(self, resolve=True)
+        # Convert the dataclass to an OmegaConf structure first
+        omega_config = OmegaConf.structured(self)
+        config_yaml = OmegaConf.to_yaml(omega_config, resolve=True)
         logger.info(f"Configuration:\n{config_yaml}")
 
     def finalize_and_validate(self) -> None: