From 254a52d37a49bf8ed53bf28a0eab02f14119f706 Mon Sep 17 00:00:00 2001
From: Eric Alt <ericallenalt@gmail.com>
Date: Fri, 9 Jan 2026 16:34:23 -0800
Subject: [PATCH 01/31] Create resolve mlflow config

---
 simplexity/structured_configs/mlflow.py | 36 +++++++++++++++++++++++++
 1 file changed, 36 insertions(+)

diff --git a/simplexity/structured_configs/mlflow.py b/simplexity/structured_configs/mlflow.py
index f3b88403..e478cabf 100644
--- a/simplexity/structured_configs/mlflow.py
+++ b/simplexity/structured_configs/mlflow.py
@@ -11,10 +11,12 @@
 
 from dataclasses import dataclass
 
+from mlflow import MlflowClient
 from omegaconf import DictConfig
 
 from simplexity.structured_configs.validation import validate_bool, validate_nonempty_str, validate_uri
 from simplexity.utils.config_utils import dynamic_resolve
+from simplexity.utils.mlflow_utils import get_experiment, get_run, resolve_registry_uri
 
 
 @dataclass
@@ -53,6 +55,40 @@ def validate_mlflow_config(cfg: DictConfig) -> None:
     validate_uri(registry_uri, "MLFlowConfig.registry_uri", is_none_allowed=True)
 
 
+@dynamic_resolve
+def resolve_mlflow_config(cfg: DictConfig) -> None:
+    """Resolve a MLFlowConfig."""
+    # Resolve registry uri
+    tracking_uri: str | None = cfg.get("tracking_uri")
+    registry_uri: str | None = cfg.get("registry_uri")
+    downgrade_unity_catalog: bool = cfg.get("downgrade_unity_catalog", True)
+    cfg.registry_uri = resolve_registry_uri(
+        registry_uri=registry_uri,
+        tracking_uri=tracking_uri,
+        downgrade_unity_catalog=downgrade_unity_catalog,
+    )
+
+    client = MlflowClient(tracking_uri=tracking_uri, registry_uri=cfg.registry_uri)
+
+    # Resolve experiment id and name
+    experiment_id: str | None = cfg.get("experiment_id")
+    experiment_name: str | None = cfg.get("experiment_name")
+    experiment = get_experiment(experiment_id=experiment_id, experiment_name=experiment_name, client=client)
+    if experiment is None:
+        raise ValueError(f"Experiment not found for id: {experiment_id} and name: {experiment_name}")
+    cfg.experiment_id = experiment.experiment_id
+    cfg.experiment_name = experiment.name
+
+    # Resolve run id and name
+    run_id: str | None = cfg.get("run_id")
+    run_name: str | None = cfg.get("run_name")
+    run = get_run(run_id=run_id, run_name=run_name, experiment_id=cfg.experiment_id, client=client)
+    if run is None:
+        raise ValueError(f"Run not found for id: {run_id} and name: {run_name}")
+    cfg.run_id = run.info.run_id
+    cfg.run_name = run.info.run_name
+
+
 @dynamic_resolve
 def update_mlflow_config(cfg: DictConfig, updated_cfg: DictConfig) -> None:
     """Update a MLFlowConfig with the updated configuration."""

From c082130abcd40e1ca5bf92517b0ed4f9fcaa18d5 Mon Sep 17 00:00:00 2001
From: Eric Alt <ericallenalt@gmail.com>
Date: Fri, 9 Jan 2026 16:34:48 -0800
Subject: [PATCH 02/31] Create mlflow defaults

---
 simplexity/run_management/run_management.py   |   2 +
 .../structured_configs/mlflow_defaults.py     | 182 ++++++
 tests/run_management/test_mlflow_defaults.py  | 516 ++++++++++++++++++
 3 files changed, 700 insertions(+)
 create mode 100644 simplexity/structured_configs/mlflow_defaults.py
 create mode 100644 tests/run_management/test_mlflow_defaults.py

diff --git a/simplexity/run_management/run_management.py b/simplexity/run_management/run_management.py
index ed75a730..6f1c7764 100644
--- a/simplexity/run_management/run_management.py
+++ b/simplexity/run_management/run_management.py
@@ -64,6 +64,7 @@
     validate_metric_tracker_config,
 )
 from simplexity.structured_configs.mlflow import update_mlflow_config
+from simplexity.structured_configs.mlflow_defaults import load_mlflow_defaults
 from simplexity.structured_configs.optimizer import (
     is_optimizer_target,
     is_pytorch_optimizer_config,
@@ -633,6 +634,7 @@ def decorator(fn: Callable[..., Any]) -> Callable[..., Any]:
         def wrapper(*args: Any, **kwargs: Any) -> Any:
             try:
                 cfg = get_config(args, kwargs)
+                cfg = load_mlflow_defaults(cfg)
                 validate_base_config(cfg)
                 resolve_base_config(cfg, strict=strict)
                 with _setup_device(cfg), _setup_mlflow(cfg):
diff --git a/simplexity/structured_configs/mlflow_defaults.py b/simplexity/structured_configs/mlflow_defaults.py
new file mode 100644
index 00000000..9b04ab78
--- /dev/null
+++ b/simplexity/structured_configs/mlflow_defaults.py
@@ -0,0 +1,182 @@
+"""Config utilities."""
+
+# pylint: disable=all
+# Temporarily disable all pylint checkers during AST traversal to prevent crash.
+# The imports checker crashes when resolving simplexity package imports due to a bug
+# in pylint/astroid: https://github.com/pylint-dev/pylint/issues/10185
+# pylint: enable=all
+# Re-enable all pylint checkers for the checking phase. This allows other checks
+# (code quality, style, undefined names, etc.) to run normally while bypassing
+# the problematic imports checker that would crash during AST traversal.
+
+import re
+import tempfile
+from typing import Any, NamedTuple, cast
+
+from mlflow import MlflowClient
+from omegaconf import DictConfig, ListConfig, OmegaConf
+
+from simplexity.exceptions import ConfigValidationError
+from simplexity.logger import SIMPLEXITY_LOGGER
+from simplexity.structured_configs.mlflow import resolve_mlflow_config, validate_mlflow_config
+from simplexity.utils.config_utils import dynamic_resolve
+
+# TARGET(@PACKAGE)?
+# [optional|override]? TARGET(@PACKAGE)? : OPTION
+FLAGS_STR = r"(?P<flags>(?:(?:optional|override)\s+)*)"
+MLFLOW_CONFIG_ENTRY_STR = r"(?P<target>[\w\.]+)(?:@(?P<package>[\w\.]+))?"
+OPTION_STR = r"(?P<option>.*?)"
+MLFLOW_DEFAULT_ITEM_PATTERN = re.compile(rf"^{FLAGS_STR}(?:{MLFLOW_CONFIG_ENTRY_STR})?(?:\s*:\s*{OPTION_STR})?$")
+
+
+class _ParsedEntry(NamedTuple):
+    """Parsed MLflow default item."""
+
+    optional: bool
+    target: str
+    package: str
+    artifact_path: str | None
+    select_path: str | None
+
+
+def _parse_option(option: str) -> tuple[str, str | None]:
+    """Parse artifact path and select path from option string."""
+    if "#" in option:
+        artifact_part, select_part = option.split("#", 1)
+        artifact_path = artifact_part.strip() or "config"
+        select_path = select_part.strip() or None
+        return artifact_path, select_path
+
+    if "/" in option:
+        artifact_path = option.strip()
+        select_path = None
+        return artifact_path, select_path
+
+    artifact_path = "config"
+    select_path = option.strip() or None
+    return artifact_path, select_path
+
+
+def _parse_entry(item: str) -> _ParsedEntry:
+    """Parse a single MLflow default item."""
+    match = MLFLOW_DEFAULT_ITEM_PATTERN.match(item)
+    if not match:
+        raise ValueError(
+            f"Invalid MLflow default entry: {item}. Must be in format '[optional|override]* TARGET(@PACKAGE)?: OPTION'"
+        )
+
+    groups = match.groupdict()
+    flags_str = groups.get("flags", "")
+    optional = "optional" in flags_str
+
+    target = groups["target"]
+    if not target:
+        raise ValueError(
+            f"Invalid MLflow default entry: {item}. Must be in format "
+            "'TARGET(@PACKAGE)? | [optional|override]? TARGET(@PACKAGE)?: OPTION | _self_'"
+        )
+
+    package = groups.get("package") or "."
+    option = groups.get("option")
+
+    if option == "null":
+        return _ParsedEntry(optional, target, package, None, None)
+
+    artifact_path, select_path = _parse_option(option or "")
+    return _ParsedEntry(optional, target, package, artifact_path, select_path)
+
+
+def _get_target_config(cfg: DictConfig, parsed_entry: _ParsedEntry) -> Any | None:
+    if parsed_entry.artifact_path is None:
+        SIMPLEXITY_LOGGER.warning("Config is mandatory but OPTION is null for entry: %s", parsed_entry.target)
+        return None
+
+    target_node: DictConfig | None = OmegaConf.select(cfg, parsed_entry.target)
+    if target_node is None:
+        SIMPLEXITY_LOGGER.warning("Target node '%s' not found in config", parsed_entry.target)
+        return None
+
+    try:
+        validate_mlflow_config(target_node)
+    except ConfigValidationError as e:
+        SIMPLEXITY_LOGGER.warning("Error validating MLflow config: %s", e)
+        return None
+
+    try:
+        resolve_mlflow_config(target_node)
+    except ValueError as e:
+        SIMPLEXITY_LOGGER.warning("Error resolving MLflow config: %s", e)
+        return None
+
+    tracking_uri: str | None = target_node.get("tracking_uri")
+    run_id: str = target_node.get("run_id")
+    client = MlflowClient(tracking_uri=tracking_uri)
+
+    with tempfile.TemporaryDirectory() as tmp_dir:
+        try:
+            local_path = client.download_artifacts(run_id=run_id, path=parsed_entry.artifact_path, dst_path=tmp_dir)
+        except Exception as e:
+            SIMPLEXITY_LOGGER.warning("Failed to download artifact from MLflow '%s': %s", parsed_entry.target, e)
+            return None
+
+        try:
+            loaded_config = OmegaConf.load(local_path)
+        except Exception as e:
+            SIMPLEXITY_LOGGER.warning("Failed to load MLflow default '%s': %s", parsed_entry.target, e)
+            return None
+
+    if parsed_entry.select_path is None:
+        return loaded_config
+
+    selected_config = OmegaConf.select(loaded_config, parsed_entry.select_path)
+    if selected_config is None:
+        SIMPLEXITY_LOGGER.warning(
+            "Selected path '%s' not found in artifact '%s'", parsed_entry.select_path, parsed_entry.artifact_path
+        )
+        return None
+
+    return selected_config
+
+
+def _process_entry(cfg: DictConfig, accumulator: DictConfig, item: str) -> DictConfig:
+    """Process a single MLflow default item."""
+    if item == "_self_":
+        return cast(DictConfig, OmegaConf.merge(accumulator, cfg))
+
+    parsed_entry = _parse_entry(item)
+
+    loaded_config = _get_target_config(cfg, parsed_entry)
+
+    if loaded_config is None:
+        if parsed_entry.optional:
+            return accumulator
+        raise ValueError(f"Target config not found for entry: {item}")
+
+    if parsed_entry.package == ".":
+        if not isinstance(loaded_config, DictConfig):
+            if parsed_entry.optional:
+                return accumulator
+            raise ValueError(f"Target config not found for entry: {item}")
+        return cast(DictConfig, OmegaConf.merge(accumulator, loaded_config))
+
+    package_conf = OmegaConf.create()
+    OmegaConf.update(package_conf, parsed_entry.package, loaded_config)
+    return cast(DictConfig, OmegaConf.merge(accumulator, package_conf))
+
+
+@dynamic_resolve
+def load_mlflow_defaults(cfg: DictConfig) -> DictConfig:
+    """Load defaults from MLflow runs."""
+    mlflow_defaults: ListConfig | None = cfg.get("mlflow_defaults")
+    if mlflow_defaults is None:
+        return cfg
+
+    if "_self_" not in mlflow_defaults:
+        mlflow_defaults.append("_self_")
+
+    accumulator = OmegaConf.create()
+
+    for item in mlflow_defaults:
+        accumulator = _process_entry(cfg, accumulator, item)
+
+    return accumulator
diff --git a/tests/run_management/test_mlflow_defaults.py b/tests/run_management/test_mlflow_defaults.py
new file mode 100644
index 00000000..5b83911f
--- /dev/null
+++ b/tests/run_management/test_mlflow_defaults.py
@@ -0,0 +1,516 @@
+"""Test MLflow defaults."""
+
+from pathlib import Path
+from unittest.mock import ANY, MagicMock
+
+import pytest
+from mlflow.client import MlflowClient
+from omegaconf import DictConfig, OmegaConf
+
+from simplexity.exceptions import ConfigValidationError
+from simplexity.structured_configs.mlflow_defaults import load_mlflow_defaults
+
+
+@pytest.fixture(autouse=True)
+def _patch_mlflow_helpers(mocker):
+    """Automatically patch common MLflow helper functions for all tests."""
+    mocker.patch("simplexity.structured_configs.mlflow_defaults.validate_mlflow_config")
+    mocker.patch("simplexity.structured_configs.mlflow_defaults.resolve_mlflow_config")
+
+
+@pytest.fixture
+def base_cfg() -> DictConfig:
+    """Set up the base config."""
+    return OmegaConf.create(
+        {
+            "previous_run": {
+                "tracking_uri": "databricks",
+                "run_id": "test_run_id",
+            },
+            "other_section": {"foo": "bar"},
+        }
+    )
+
+
+@pytest.fixture
+def mock_download(mocker) -> MagicMock:
+    """Create a mock MLflow client instance."""
+    mock_client_instance = MagicMock(spec=MlflowClient)
+    mocker.patch(
+        "simplexity.structured_configs.mlflow_defaults.MlflowClient",
+        return_value=mock_client_instance,
+    )
+    return mock_client_instance.download_artifacts
+
+
+def test_load_default_config_at_root(base_cfg: DictConfig, mock_download: MagicMock, tmp_path: Path):
+    """Test loading basic config."""
+    artifact_path = tmp_path / "config.yaml"
+    artifact_path.write_text("loaded_key: loaded_value\n")
+    mock_download.return_value = str(artifact_path)
+
+    cfg = OmegaConf.merge(
+        base_cfg,
+        {
+            "mlflow_defaults": ["previous_run"],
+        },
+    )
+
+    loaded_cfg: DictConfig = load_mlflow_defaults(cfg)
+
+    assert loaded_cfg.get("loaded_key") == "loaded_value"
+    assert OmegaConf.select(loaded_cfg, "other_section.foo") == "bar"
+
+
+def test_load_default_config_at_package(base_cfg: DictConfig, mock_download: MagicMock, tmp_path: Path):
+    """Test loading at package."""
+    artifact_path = tmp_path / "config.yaml"
+    artifact_path.write_text("loaded_key: loaded_value\n")
+    mock_download.return_value = str(artifact_path)
+
+    cfg = OmegaConf.merge(
+        base_cfg,
+        {
+            "mlflow_defaults": ["previous_run@dest"],
+        },
+    )
+
+    loaded_cfg: DictConfig = load_mlflow_defaults(cfg)
+    assert OmegaConf.select(loaded_cfg, "dest.loaded_key") == "loaded_value"
+    assert OmegaConf.select(loaded_cfg, "other_section.foo") == "bar"
+    assert mock_download.call_args.kwargs["path"] == "config"
+
+
+def test_load_nondefault_config(base_cfg: DictConfig, mock_download: MagicMock, tmp_path: Path):
+    """Test loading basic config."""
+    artifact_path = tmp_path / "nondefault_config.yaml"
+    artifact_path.write_text("loaded_key: loaded_value\n")
+    mock_download.return_value = str(artifact_path)
+
+    cfg = OmegaConf.merge(
+        base_cfg,
+        {
+            "mlflow_defaults": ["previous_run: nondefault_config#"],
+        },
+    )
+
+    loaded_cfg: DictConfig = load_mlflow_defaults(cfg)
+
+    assert loaded_cfg.get("loaded_key") == "loaded_value"
+    assert OmegaConf.select(loaded_cfg, "other_section.foo") == "bar"
+    assert mock_download.call_args.kwargs["path"] == "nondefault_config"
+
+
+def test_load_subconfig(base_cfg: DictConfig, mock_download: MagicMock, tmp_path: Path):
+    """Test loading subconfig."""
+    artifact_path = tmp_path / "config.yaml"
+    artifact_path.write_text("sub:\n  nested: value\n")
+    mock_download.return_value = str(artifact_path)
+
+    cfg = OmegaConf.merge(
+        base_cfg,
+        {
+            "mlflow_defaults": ["previous_run: sub"],
+        },
+    )
+
+    loaded_cfg: DictConfig = load_mlflow_defaults(cfg)
+    assert loaded_cfg.get("nested") == "value"
+    assert mock_download.call_args.kwargs["path"] == "config"
+
+
+def test_load_subconfig_select(base_cfg: DictConfig, mock_download: MagicMock, tmp_path: Path):
+    """Test loading subconfig select."""
+    artifact_path = tmp_path / "nondefault_config.yaml"
+    artifact_path.write_text("sub:\n  nested: value\n")
+    mock_download.return_value = str(artifact_path)
+
+    cfg = OmegaConf.merge(
+        base_cfg,
+        {
+            "mlflow_defaults": ["previous_run@dest: nondefault_config#sub.nested"],
+        },
+    )
+
+    loaded_cfg: DictConfig = load_mlflow_defaults(cfg)
+    assert loaded_cfg.get("dest") == "value"
+
+
+def test_implicit_config_select(base_cfg: DictConfig, mock_download: MagicMock, tmp_path: Path):
+    """Test implicit artifact="config" when option has no / or #."""
+    artifact_path = tmp_path / "config.yaml"
+    artifact_path.write_text("sub: value\n")
+    mock_download.return_value = str(artifact_path)
+
+    cfg = OmegaConf.merge(
+        base_cfg,
+        {
+            "mlflow_defaults": ["previous_run@dest: sub"],
+        },
+    )
+
+    loaded_cfg: DictConfig = load_mlflow_defaults(cfg)
+    assert loaded_cfg.get("dest") == "value"
+    mock_download.assert_called_with(run_id="test_run_id", path="config", dst_path=ANY)
+
+
+def test_implicit_artifact_select(base_cfg: DictConfig, mock_download: MagicMock, tmp_path: Path):
+    """Test implicit select=root when option has /."""
+    artifact_path = tmp_path / "artifact.yaml"
+    artifact_path.write_text("key: value\n")
+    mock_download.return_value = str(artifact_path)
+
+    cfg = OmegaConf.merge(
+        base_cfg,
+        {
+            "mlflow_defaults": ["previous_run@dest: path/to/artifact"],
+        },
+    )
+    loaded_cfg: DictConfig = load_mlflow_defaults(cfg)
+    assert OmegaConf.select(loaded_cfg, "dest.key") == "value"
+    mock_download.assert_called_with(run_id="test_run_id", path="path/to/artifact", dst_path=ANY)
+
+
+def test_override_flag(base_cfg: DictConfig, mock_download: MagicMock, tmp_path: Path):
+    """Test override flag."""
+    artifact_path = tmp_path / "nondefault_config.yaml"
+    artifact_path.write_text("foo: new_bar\n")
+    mock_download.return_value = str(artifact_path)
+
+    # Place _self_ first so the loaded config overrides the base config
+    cfg = OmegaConf.merge(
+        base_cfg,
+        {
+            "mlflow_defaults": ["_self_", "override previous_run@other_section: nondefault_config#"],
+        },
+    )
+    loaded_cfg: DictConfig = load_mlflow_defaults(cfg)
+    assert OmegaConf.select(loaded_cfg, "other_section.foo") == "new_bar"
+
+
+def test_optional_flag_missing_artifact(base_cfg: DictConfig):
+    """Test optional flag with missing artifact."""
+
+    cfg = OmegaConf.merge(
+        base_cfg,
+        {
+            "mlflow_defaults": ["optional previous_run: null"],
+        },
+    )
+    # Should not raise
+    loaded_cfg: DictConfig = load_mlflow_defaults(cfg)
+    assert loaded_cfg == cfg
+
+
+def test_null_option_error(base_cfg: DictConfig):
+    """Test null option raises error."""
+    cfg = OmegaConf.merge(
+        base_cfg,
+        {
+            "mlflow_defaults": ["previous_run: null"],
+        },
+    )
+    # null option returns just the target config in new logic unless specific logic handles it differently
+    # Actually user logic for option="null" returns _ParsedEntry with artifact_path=None
+    # And _get_target_config handles artifact_path=None by warning and returning None
+    # So "mlflow: null" should be valid but yield nothing in accumulator if not trapped earlier?
+    # Re-reading logic in Step 245:
+    # _parse_entry returns artifact_path=None if option="null"
+    # _get_target_config checks if parsed_entry.artifact_path is None: logs Warning, returns None.
+    # _process_entry sees loaded_config is None.
+    # It raises ValueError("Target config not found for entry...") if not optional.
+    # So it SHOULD raise ValueError.
+
+    with pytest.raises(ValueError, match="Target config not found for entry"):
+        load_mlflow_defaults(cfg)
+
+
+def test_null_option_optional(base_cfg: DictConfig):
+    """Test null option with optional flag."""
+    cfg = OmegaConf.merge(
+        base_cfg,
+        {
+            "mlflow_defaults": ["optional previous_run: null"],
+        },
+    )
+    loaded_cfg: DictConfig = load_mlflow_defaults(cfg)
+    assert OmegaConf.select(loaded_cfg, "other_section") == cfg.other_section
+
+
+def test_composition_order(base_cfg: DictConfig, mock_download: MagicMock, tmp_path: Path):
+    """Test _self_ placement in composition order."""
+    artifact_path = tmp_path / "nondefault_config.yaml"
+    artifact_path.write_text("foo: mlflow_value\n")
+    mock_download.return_value = str(artifact_path)
+
+    # Case 1: _self_ (default/original) before mlflow -> mlflow overrides
+    cfg1 = OmegaConf.merge(base_cfg, {"mlflow_defaults": ["_self_", "previous_run@other_section: nondefault_config#"]})
+    loaded1: DictConfig = load_mlflow_defaults(cfg1)
+    assert OmegaConf.select(loaded1, "other_section.foo") == "mlflow_value"
+
+    # Case 2: mlflow before _self_ -> original overrides
+    cfg2 = OmegaConf.merge(base_cfg, {"mlflow_defaults": ["previous_run@other_section: nondefault_config#", "_self_"]})
+    loaded2: DictConfig = load_mlflow_defaults(cfg2)
+    assert OmegaConf.select(loaded2, "other_section.foo") == "bar"  # Original value
+
+
+def test_config_entry_syntax(base_cfg: DictConfig, mock_download: MagicMock, tmp_path: Path):
+    """Test CONFIG syntax: TARGET (no option)."""
+    artifact_path = tmp_path / "nondefault_config.yaml"
+    artifact_path.write_text("key: val\n")
+    mock_download.return_value = str(artifact_path)
+
+    cfg = OmegaConf.merge(
+        base_cfg,
+        {
+            "mlflow_defaults": ["previous_run"],
+        },
+    )
+    loaded_cfg: DictConfig = load_mlflow_defaults(cfg)
+    assert OmegaConf.select(loaded_cfg, "key") == "val"
+    # Implicitly loads "config" artifact
+    mock_download.assert_called_with(run_id="test_run_id", path="config", dst_path=ANY)
+
+
+def test_missing_target_error(base_cfg: DictConfig):
+    """Test missing target raises error."""
+    cfg = OmegaConf.merge(
+        base_cfg,
+        {
+            "mlflow_defaults": ["optional : option"],
+        },
+    )
+    with pytest.raises(ValueError, match="Invalid MLflow default entry"):
+        load_mlflow_defaults(cfg)
+
+
+def test_invalid_entry_format(base_cfg: DictConfig):
+    """Test invalid entry format raises error (line 64)."""
+    cfg = OmegaConf.merge(
+        base_cfg,
+        {
+            "mlflow_defaults": ["invalid@format@too@many@ats"],
+        },
+    )
+    with pytest.raises(ValueError, match="Invalid MLflow default entry"):
+        load_mlflow_defaults(cfg)
+
+
+def test_missing_target_node(base_cfg: DictConfig):
+    """Test missing target node in config (lines 96-97)."""
+    cfg = OmegaConf.merge(
+        base_cfg,
+        {
+            "mlflow_defaults": ["nonexistent: config"],
+        },
+    )
+    # Should raise ValueError because target node not found and not optional
+    with pytest.raises(ValueError, match="Target config not found for entry"):
+        load_mlflow_defaults(cfg)
+
+
+def test_missing_target_node_optional(base_cfg: DictConfig):
+    """Test missing target node with optional flag."""
+    cfg = OmegaConf.merge(
+        base_cfg,
+        {
+            "mlflow_defaults": ["optional nonexistent: config"],
+        },
+    )
+    # Should not raise, just return accumulator
+    loaded_cfg: DictConfig = load_mlflow_defaults(cfg)
+    assert OmegaConf.select(loaded_cfg, "other_section") == cfg.other_section
+
+
+def test_validate_mlflow_config_error(base_cfg: DictConfig, mocker):
+    """Test ConfigValidationError from validate_mlflow_config (lines 101-103)."""
+    # Make validate_mlflow_config raise ConfigValidationError
+    mock_validate = mocker.patch("simplexity.structured_configs.mlflow_defaults.validate_mlflow_config")
+    mock_validate.side_effect = ConfigValidationError("Invalid config")
+
+    cfg = OmegaConf.merge(
+        base_cfg,
+        {
+            "mlflow_defaults": ["previous_run: config"],
+        },
+    )
+    # Should raise ValueError because validation failed and not optional
+    with pytest.raises(ValueError, match="Target config not found for entry"):
+        load_mlflow_defaults(cfg)
+
+
+def test_validate_mlflow_config_error_optional(base_cfg: DictConfig, mocker):
+    """Test ConfigValidationError with optional flag."""
+    # Make validate_mlflow_config raise ConfigValidationError
+    mock_validate = mocker.patch("simplexity.structured_configs.mlflow_defaults.validate_mlflow_config")
+    mock_validate.side_effect = ConfigValidationError("Invalid config")
+
+    cfg = OmegaConf.merge(
+        base_cfg,
+        {
+            "mlflow_defaults": ["optional previous_run: config"],
+        },
+    )
+    # Should not raise, just return accumulator
+    loaded_cfg: DictConfig = load_mlflow_defaults(cfg)
+    assert OmegaConf.select(loaded_cfg, "other_section") == cfg.other_section
+
+
+def test_resolve_mlflow_config_error(base_cfg: DictConfig, mocker):
+    """Test ValueError from resolve_mlflow_config (lines 107-109)."""
+    # Make resolve_mlflow_config raise ValueError
+    mock_resolve = mocker.patch("simplexity.structured_configs.mlflow_defaults.resolve_mlflow_config")
+    mock_resolve.side_effect = ValueError("Resolution failed")
+
+    cfg = OmegaConf.merge(
+        base_cfg,
+        {
+            "mlflow_defaults": ["previous_run: config"],
+        },
+    )
+    # Should raise ValueError because resolution failed and not optional
+    with pytest.raises(ValueError, match="Target config not found for entry"):
+        load_mlflow_defaults(cfg)
+
+
+def test_resolve_mlflow_config_error_optional(base_cfg: DictConfig, mocker):
+    """Test ValueError from resolve_mlflow_config with optional flag."""
+    # Make resolve_mlflow_config raise ValueError
+    mock_resolve = mocker.patch("simplexity.structured_configs.mlflow_defaults.resolve_mlflow_config")
+    mock_resolve.side_effect = ValueError("Resolution failed")
+
+    cfg = OmegaConf.merge(
+        base_cfg,
+        {
+            "mlflow_defaults": ["optional mlflow: config"],
+        },
+    )
+    # Should not raise, just return accumulator
+    loaded_cfg: DictConfig = load_mlflow_defaults(cfg)
+    assert OmegaConf.select(loaded_cfg, "other_section") == cfg.other_section
+
+
+def test_omegaconf_load_error(base_cfg: DictConfig, mock_download: MagicMock, tmp_path: Path):
+    """Test Exception from OmegaConf.load (lines 124-126)."""
+    # Create a file with invalid YAML content
+    artifact_path = tmp_path / "invalid.yaml"
+    artifact_path.write_text("invalid: yaml: content: [unclosed")
+    mock_download.return_value = str(artifact_path)
+
+    cfg = OmegaConf.merge(
+        base_cfg,
+        {
+            "mlflow_defaults": ["mlflow: config"],
+        },
+    )
+    # Should raise ValueError because load failed and not optional
+    with pytest.raises(ValueError, match="Target config not found for entry"):
+        load_mlflow_defaults(cfg)
+
+
+def test_omegaconf_load_error_optional(base_cfg: DictConfig, mock_download: MagicMock, tmp_path: Path):
+    """Test Exception from OmegaConf.load with optional flag."""
+    # Create a file with invalid YAML content
+    artifact_path = tmp_path / "invalid.yaml"
+    artifact_path.write_text("invalid: yaml: content: [unclosed")
+    mock_download.return_value = str(artifact_path)
+
+    cfg = OmegaConf.merge(
+        base_cfg,
+        {
+            "mlflow_defaults": ["optional previous_run: config"],
+        },
+    )
+    # Should not raise, just return accumulator
+    loaded_cfg: DictConfig = load_mlflow_defaults(cfg)
+    assert OmegaConf.select(loaded_cfg, "other_section") == cfg.other_section
+
+
+def test_select_path_not_found(base_cfg: DictConfig, mock_download: MagicMock, tmp_path: Path):
+    """Test select path not found in artifact (lines 133-136)."""
+    artifact_path = tmp_path / "nondefault_config.yaml"
+    artifact_path.write_text("key: value\n")
+    mock_download.return_value = str(artifact_path)
+
+    cfg = OmegaConf.merge(
+        base_cfg,
+        {
+            "mlflow_defaults": ["previous_run@dest: nondefault_config#nonexistent.path"],
+        },
+    )
+    # Should raise ValueError because select path not found and not optional
+    with pytest.raises(ValueError, match="Target config not found for entry"):
+        load_mlflow_defaults(cfg)
+
+
+def test_select_path_not_found_optional(base_cfg: DictConfig, mock_download: MagicMock, tmp_path: Path):
+    """Test select path not found with optional flag."""
+    artifact_path = tmp_path / "nondefault_config.yaml"
+    artifact_path.write_text("key: value\n")
+    mock_download.return_value = str(artifact_path)
+
+    cfg = OmegaConf.merge(
+        base_cfg,
+        {
+            "mlflow_defaults": ["optional previous_run@dest: nondefault_config#nonexistent.path"],
+        },
+    )
+    # Should not raise, just return accumulator
+    loaded_cfg: DictConfig = load_mlflow_defaults(cfg)
+    assert OmegaConf.select(loaded_cfg, "other_section") == cfg.other_section
+
+
+def test_non_dictconfig_error(base_cfg: DictConfig, mock_download: MagicMock, tmp_path: Path):
+    """Test non-DictConfig error when not optional (lines 157-159)."""
+    # Create artifact with list (non-dict) root value
+    # Use "nondefault_config#" to avoid select_path, so we get the full ListConfig
+    artifact_path = tmp_path / "nondefault_config.yaml"
+    artifact_path.write_text("- item1\n- item2\n")
+    mock_download.return_value = str(artifact_path)
+
+    cfg = OmegaConf.merge(
+        base_cfg,
+        {
+            "mlflow_defaults": ["previous_run: nondefault_config#"],
+        },
+    )
+    # Should raise ValueError because loaded_config is ListConfig (not DictConfig) and not optional
+    with pytest.raises(ValueError, match="Target config not found for entry"):
+        load_mlflow_defaults(cfg)
+
+
+def test_non_dictconfig_optional(base_cfg: DictConfig, mock_download: MagicMock, tmp_path: Path):
+    """Test non-DictConfig with optional flag."""
+    # Create artifact with list (non-dict) root value
+    # Use "nondefault_config#" to avoid select_path, so we get the full ListConfig
+    artifact_path = tmp_path / "nondefault_config.yaml"
+    artifact_path.write_text("- item1\n- item2\n")
+    mock_download.return_value = str(artifact_path)
+
+    cfg = OmegaConf.merge(
+        base_cfg,
+        {
+            "mlflow_defaults": ["optional previous_run: nondefault_config#"],
+        },
+    )
+    # Should not raise, just return accumulator
+    loaded_cfg: DictConfig = load_mlflow_defaults(cfg)
+    assert OmegaConf.select(loaded_cfg, "other_section") == cfg.other_section
+
+
+def test_no_mlflow_defaults_key():
+    """Test early return when mlflow_defaults key is missing (line 172)."""
+    # Config without mlflow_defaults key
+    cfg = OmegaConf.create(
+        {
+            "previous_run": {
+                "tracking_uri": "databricks",
+                "run_id": "test_run_id",
+            },
+            "other_section": {"foo": "bar"},
+        }
+    )
+    # Should return cfg unchanged
+    loaded_cfg: DictConfig = load_mlflow_defaults(cfg)
+    assert loaded_cfg == cfg
+    assert OmegaConf.select(loaded_cfg, "other_section.foo") == "bar"

From 27b814c860ddd28489fd458ccc2b8a2dff866ebc Mon Sep 17 00:00:00 2001
From: Eric Alt <ericallenalt@gmail.com>
Date: Thu, 15 Jan 2026 12:36:20 -0800
Subject: [PATCH 03/31] Update defaults e2e test

---
 .../full_default_copy.yaml                    |   4 +
 .../setup/group_a/subconfig_a1.yaml           |   3 +
 .../setup/group_a/subconfig_a2.yaml           |   3 +
 .../setup/group_b/subconfig_b1.yaml           |   3 +
 .../setup/group_b/subconfig_b2.yaml           |   3 +
 .../setup/prev_config_1.yaml                  |   8 +
 .../setup/prev_config_2.yaml                  |   8 +
 .../setup/special_1.yaml                      |  11 ++
 .../setup/special_2.yaml                      |  11 ++
 tests/end_to_end/test_mlflow_defaults.py      | 145 ++++++++++++++++++
 10 files changed, 199 insertions(+)
 create mode 100644 tests/end_to_end/mlflow_defaults_configs/full_default_copy.yaml
 create mode 100644 tests/end_to_end/mlflow_defaults_configs/setup/group_a/subconfig_a1.yaml
 create mode 100644 tests/end_to_end/mlflow_defaults_configs/setup/group_a/subconfig_a2.yaml
 create mode 100644 tests/end_to_end/mlflow_defaults_configs/setup/group_b/subconfig_b1.yaml
 create mode 100644 tests/end_to_end/mlflow_defaults_configs/setup/group_b/subconfig_b2.yaml
 create mode 100644 tests/end_to_end/mlflow_defaults_configs/setup/prev_config_1.yaml
 create mode 100644 tests/end_to_end/mlflow_defaults_configs/setup/prev_config_2.yaml
 create mode 100644 tests/end_to_end/mlflow_defaults_configs/setup/special_1.yaml
 create mode 100644 tests/end_to_end/mlflow_defaults_configs/setup/special_2.yaml
 create mode 100644 tests/end_to_end/test_mlflow_defaults.py

diff --git a/tests/end_to_end/mlflow_defaults_configs/full_default_copy.yaml b/tests/end_to_end/mlflow_defaults_configs/full_default_copy.yaml
new file mode 100644
index 00000000..a1fe6e49
--- /dev/null
+++ b/tests/end_to_end/mlflow_defaults_configs/full_default_copy.yaml
@@ -0,0 +1,4 @@
+defaults:
+  - mlflow@load_source: prev_run_1
+mlflow_defaults:
+  - load_source
\ No newline at end of file
diff --git a/tests/end_to_end/mlflow_defaults_configs/setup/group_a/subconfig_a1.yaml b/tests/end_to_end/mlflow_defaults_configs/setup/group_a/subconfig_a1.yaml
new file mode 100644
index 00000000..53a0ed55
--- /dev/null
+++ b/tests/end_to_end/mlflow_defaults_configs/setup/group_a/subconfig_a1.yaml
@@ -0,0 +1,3 @@
+a_key: 1
+a_dict:
+  a_subkey: 10
\ No newline at end of file
diff --git a/tests/end_to_end/mlflow_defaults_configs/setup/group_a/subconfig_a2.yaml b/tests/end_to_end/mlflow_defaults_configs/setup/group_a/subconfig_a2.yaml
new file mode 100644
index 00000000..ba32fc12
--- /dev/null
+++ b/tests/end_to_end/mlflow_defaults_configs/setup/group_a/subconfig_a2.yaml
@@ -0,0 +1,3 @@
+a_key: 2
+a_dict:
+  a_subkey: 20
\ No newline at end of file
diff --git a/tests/end_to_end/mlflow_defaults_configs/setup/group_b/subconfig_b1.yaml b/tests/end_to_end/mlflow_defaults_configs/setup/group_b/subconfig_b1.yaml
new file mode 100644
index 00000000..a89b185c
--- /dev/null
+++ b/tests/end_to_end/mlflow_defaults_configs/setup/group_b/subconfig_b1.yaml
@@ -0,0 +1,3 @@
+b_key: 1
+b_dict:
+  b_subkey: 10
\ No newline at end of file
diff --git a/tests/end_to_end/mlflow_defaults_configs/setup/group_b/subconfig_b2.yaml b/tests/end_to_end/mlflow_defaults_configs/setup/group_b/subconfig_b2.yaml
new file mode 100644
index 00000000..9e407359
--- /dev/null
+++ b/tests/end_to_end/mlflow_defaults_configs/setup/group_b/subconfig_b2.yaml
@@ -0,0 +1,3 @@
+b_key: 2
+b_dict:
+  b_subkey: 20
\ No newline at end of file
diff --git a/tests/end_to_end/mlflow_defaults_configs/setup/prev_config_1.yaml b/tests/end_to_end/mlflow_defaults_configs/setup/prev_config_1.yaml
new file mode 100644
index 00000000..375b3a59
--- /dev/null
+++ b/tests/end_to_end/mlflow_defaults_configs/setup/prev_config_1.yaml
@@ -0,0 +1,8 @@
+defaults:
+  - _self_
+  - group_a: subconfig_a1
+  - group_b: subconfig_b1
+
+key: 1
+dict:
+  subkey: 10
diff --git a/tests/end_to_end/mlflow_defaults_configs/setup/prev_config_2.yaml b/tests/end_to_end/mlflow_defaults_configs/setup/prev_config_2.yaml
new file mode 100644
index 00000000..d35c26bd
--- /dev/null
+++ b/tests/end_to_end/mlflow_defaults_configs/setup/prev_config_2.yaml
@@ -0,0 +1,8 @@
+defaults:
+  - _self_
+  - group_a: subconfig_a2
+  - group_b: subconfig_b2
+
+key: 2
+dict:
+  subkey: 20
\ No newline at end of file
diff --git a/tests/end_to_end/mlflow_defaults_configs/setup/special_1.yaml b/tests/end_to_end/mlflow_defaults_configs/setup/special_1.yaml
new file mode 100644
index 00000000..7b78e00b
--- /dev/null
+++ b/tests/end_to_end/mlflow_defaults_configs/setup/special_1.yaml
@@ -0,0 +1,11 @@
+key: 3
+dict:
+  subkey: 30
+group_a:
+  a_key: 3
+  a_dict:
+    a_subkey: 30
+group_b:
+  b_key: 3
+  b_dict:
+    b_subkey: 30
diff --git a/tests/end_to_end/mlflow_defaults_configs/setup/special_2.yaml b/tests/end_to_end/mlflow_defaults_configs/setup/special_2.yaml
new file mode 100644
index 00000000..a2d9aabc
--- /dev/null
+++ b/tests/end_to_end/mlflow_defaults_configs/setup/special_2.yaml
@@ -0,0 +1,11 @@
+key: 4
+dict:
+  subkey: 40
+group_a:
+  a_key: 4
+  a_dict:
+    a_subkey: 40
+group_b:
+  b_key: 4
+  b_dict:
+    b_subkey: 40
diff --git a/tests/end_to_end/test_mlflow_defaults.py b/tests/end_to_end/test_mlflow_defaults.py
new file mode 100644
index 00000000..a770835f
--- /dev/null
+++ b/tests/end_to_end/test_mlflow_defaults.py
@@ -0,0 +1,145 @@
+"""End-to-end mlflow defaults tests."""
+
+# pylint: disable=all
+# Temporarily disable all pylint checkers during AST traversal to prevent crash.
+# The imports checker crashes when resolving simplexity package imports due to a bug
+# in pylint/astroid: https://github.com/pylint-dev/pylint/issues/10185
+# pylint: enable=all
+# Re-enable all pylint checkers for the checking phase. This allows other checks
+# (code quality, style, undefined names, etc.) to run normally while bypassing
+# the problematic imports checker that would crash during AST traversal.
+
+import shutil
+from pathlib import Path
+
+import mlflow
+import pytest
+from hydra import compose, initialize_config_dir
+from omegaconf import DictConfig
+
+from simplexity.logging.mlflow_logger import MLFlowLogger
+from simplexity.structured_configs.mlflow_defaults import load_mlflow_defaults
+from simplexity.utils.mlflow_utils import get_experiment, get_run
+
+CONFIGS_SRC = str(Path(__file__).parent / "mlflow_defaults_configs" / "setup")
+
+EXPERIMENT_NAME = "test_mlflow_defaults"
+
+MLFLOW_CONFIG = """
+mlflow:
+  experiment_name: {experiment_name}
+  run_name: {run_name}
+  tracking_uri: {tracking_uri}
+"""
+
+
+@pytest.fixture(scope="session")
+def setup_dir(tmp_path: Path) -> Path:
+    """Setup function."""
+    config_path = str(tmp_path / "configs")
+    shutil.copytree(CONFIGS_SRC, config_path)
+    tracking_uri = f"sqlite:///{tmp_path.resolve()}/mlflow.db"
+    mlflow.set_tracking_uri(tracking_uri)
+
+    # for config_file in Path(config_path).rglob("*.yaml"):
+    #     lines = config_file.read_text().splitlines()
+    #     new_lines = []
+    #     for line in lines:
+    #         if "tracking_uri:" in line:
+    #             key, _ = line.split(":", 1)
+    #             new_lines.append(f"{key}: {tracking_uri}")
+    #         else:
+    #             new_lines.append(line)
+    #     config_file.write_text("\n".join(new_lines) + "\n")
+
+    def log_to_mlflow(cfg: DictConfig | None, config_names: list[str], run_name: str) -> tuple[str, str]:
+        """Save config."""
+        experiment = get_experiment(experiment_name=EXPERIMENT_NAME)
+        assert experiment is not None
+        experiment_id = experiment.experiment_id
+        run = get_run(run_name=run_name, experiment_id=experiment_id)
+        assert run is not None
+        run_id = run.info.run_id
+        with mlflow.start_run(
+            run_id=run_id,
+            experiment_id=experiment_id,
+            run_name=run_name,
+            log_system_metrics=False,
+        ):
+            logger = MLFlowLogger(tracking_uri=tracking_uri)
+            if cfg is not None:
+                logger.log_config(cfg, resolve=False)
+            for config_name in config_names:
+                logger.log_artifact(local_path=f"config_path/{config_name}", artifact_path="subdir/special.yaml")
+        return experiment_id, run_id
+
+    def previous_run(config_name: str | None, config_names: list[str], run_name: str) -> None:
+        """Previous run."""
+        if config_name is None:
+            cfg = None
+        else:
+            with initialize_config_dir(config_dir=config_path):
+                cfg = compose(config_name=config_name)
+        experiment_id, run_id = log_to_mlflow(cfg, config_names, run_name)
+        output_path = tmp_path / "configs" / "mlflow" / f"{run_name}.yaml"
+        output_path.parent.mkdir(parents=True, exist_ok=True)
+        with open(output_path, "w", encoding="utf-8") as f:
+            f.write(
+                MLFLOW_CONFIG.format(
+                    experiment_id=experiment_id,
+                    run_id=run_id,
+                    tracking_uri=tracking_uri,
+                )
+            )
+
+    previous_run(config_name="prev_config_1", config_names=["special_1.yaml"], run_name="prev_run_1")
+    previous_run(config_name="prev_config_2", config_names=["special_2.yaml"], run_name="prev_run_2")
+
+    return tmp_path
+
+
+def test_setup(setup_dir: Path, tmp_path: Path) -> None:
+    """Test setup."""
+    tracking_uri = f"sqlite:///{setup_dir.resolve()}/mlflow.db"
+    client = mlflow.MlflowClient(tracking_uri=tracking_uri)
+    experiments = client.get_experiment_by_name(name=EXPERIMENT_NAME)
+    assert experiments is not None
+    experiment_id = experiments.experiment_id
+
+    for run_num in (1, 2):
+        run = client.search_runs(
+            experiment_ids=[experiment_id], filter_string=f"attributes.run_name = 'prev_run_{run_num}'"
+        )
+        assert len(run) == 1
+        run_id = run[0].info.run_id
+
+        config_path = client.download_artifacts(
+            run_id=run_id, path="config.yaml", dst_path=str(tmp_path / f"config_{run_num}.yaml")
+        )
+        expected_config_path = setup_dir / "configs" / f"prev_run_{run_num}.yaml"
+        with open(config_path, encoding="utf-8") as f, open(expected_config_path, encoding="utf-8") as expected_f:
+            assert f.read() == expected_f.read()
+
+        special_config_path = client.download_artifacts(
+            run_id=run_id, path="subdir/special.yaml", dst_path=str(tmp_path / f"special_{run_num}.yaml")
+        )
+        expected_special_config_path = setup_dir / "configs" / f"special_{run_num}.yaml"
+        with (
+            open(special_config_path, encoding="utf-8") as f,
+            open(expected_special_config_path, encoding="utf-8") as expected_f,
+        ):
+            assert f.read() == expected_f.read()
+
+
+@pytest.mark.parametrize(
+    "test_case",
+    ["full_default_copy"],
+)
+def test_mlflow_defaults(setup_dir: Path, test_case: str) -> None:
+    """Test mlflow defaults."""
+    tracking_uri = f"sqlite:///{setup_dir.resolve()}/mlflow.db"
+    with initialize_config_dir(config_dir=str(setup_dir / "configs")):
+        cfg = compose(config_name=test_case, overrides=[f"mlflow.tracking_uri={tracking_uri}"])
+        expected = compose(config_name=f"{test_case}_expected", overrides=[f"mlflow.tracking_uri={tracking_uri}"])
+    actual = load_mlflow_defaults(cfg)
+    assert actual == expected

From 60b74d17ebf7d899eeac65d4910ead21574b3b5d Mon Sep 17 00:00:00 2001
From: Eric Alt <ericallenalt@gmail.com>
Date: Thu, 15 Jan 2026 14:04:36 -0800
Subject: [PATCH 04/31] Fix name clash

---
 ...lts.py => test_load_config_from_mlflow.py} |  13 +-
 .../test_load_mlflow_defaults.py              | 516 ++++++++++++++++++
 2 files changed, 523 insertions(+), 6 deletions(-)
 rename tests/end_to_end/{test_mlflow_defaults.py => test_load_config_from_mlflow.py} (92%)
 create mode 100644 tests/run_management/test_load_mlflow_defaults.py

diff --git a/tests/end_to_end/test_mlflow_defaults.py b/tests/end_to_end/test_load_config_from_mlflow.py
similarity index 92%
rename from tests/end_to_end/test_mlflow_defaults.py
rename to tests/end_to_end/test_load_config_from_mlflow.py
index a770835f..573418f6 100644
--- a/tests/end_to_end/test_mlflow_defaults.py
+++ b/tests/end_to_end/test_load_config_from_mlflow.py
@@ -34,10 +34,11 @@
 
 
 @pytest.fixture(scope="session")
-def setup_dir(tmp_path: Path) -> Path:
+def setup_dir(tmp_path_factory: pytest.TempPathFactory) -> Path:
     """Setup function."""
-    config_path = str(tmp_path / "configs")
-    shutil.copytree(CONFIGS_SRC, config_path)
+    tmp_path = tmp_path_factory.mktemp("mlflow_defaults")
+    config_path = tmp_path / "configs"
+    shutil.copytree(CONFIGS_SRC, str(config_path))
     tracking_uri = f"sqlite:///{tmp_path.resolve()}/mlflow.db"
     mlflow.set_tracking_uri(tracking_uri)
 
@@ -70,7 +71,7 @@ def log_to_mlflow(cfg: DictConfig | None, config_names: list[str], run_name: str
             if cfg is not None:
                 logger.log_config(cfg, resolve=False)
             for config_name in config_names:
-                logger.log_artifact(local_path=f"config_path/{config_name}", artifact_path="subdir/special.yaml")
+                logger.log_artifact(local_path=str(config_path / config_name), artifact_path="subdir/special.yaml")
         return experiment_id, run_id
 
     def previous_run(config_name: str | None, config_names: list[str], run_name: str) -> None:
@@ -78,10 +79,10 @@ def previous_run(config_name: str | None, config_names: list[str], run_name: str
         if config_name is None:
             cfg = None
         else:
-            with initialize_config_dir(config_dir=config_path):
+            with initialize_config_dir(config_dir=str(config_path)):
                 cfg = compose(config_name=config_name)
         experiment_id, run_id = log_to_mlflow(cfg, config_names, run_name)
-        output_path = tmp_path / "configs" / "mlflow" / f"{run_name}.yaml"
+        output_path = config_path / "mlflow" / f"{run_name}.yaml"
         output_path.parent.mkdir(parents=True, exist_ok=True)
         with open(output_path, "w", encoding="utf-8") as f:
             f.write(
diff --git a/tests/run_management/test_load_mlflow_defaults.py b/tests/run_management/test_load_mlflow_defaults.py
new file mode 100644
index 00000000..5b83911f
--- /dev/null
+++ b/tests/run_management/test_load_mlflow_defaults.py
@@ -0,0 +1,516 @@
+"""Test MLflow defaults."""
+
+from pathlib import Path
+from unittest.mock import ANY, MagicMock
+
+import pytest
+from mlflow.client import MlflowClient
+from omegaconf import DictConfig, OmegaConf
+
+from simplexity.exceptions import ConfigValidationError
+from simplexity.structured_configs.mlflow_defaults import load_mlflow_defaults
+
+
+@pytest.fixture(autouse=True)
+def _patch_mlflow_helpers(mocker):
+    """Automatically patch common MLflow helper functions for all tests."""
+    mocker.patch("simplexity.structured_configs.mlflow_defaults.validate_mlflow_config")
+    mocker.patch("simplexity.structured_configs.mlflow_defaults.resolve_mlflow_config")
+
+
+@pytest.fixture
+def base_cfg() -> DictConfig:
+    """Set up the base config."""
+    return OmegaConf.create(
+        {
+            "previous_run": {
+                "tracking_uri": "databricks",
+                "run_id": "test_run_id",
+            },
+            "other_section": {"foo": "bar"},
+        }
+    )
+
+
+@pytest.fixture
+def mock_download(mocker) -> MagicMock:
+    """Create a mock MLflow client instance."""
+    mock_client_instance = MagicMock(spec=MlflowClient)
+    mocker.patch(
+        "simplexity.structured_configs.mlflow_defaults.MlflowClient",
+        return_value=mock_client_instance,
+    )
+    return mock_client_instance.download_artifacts
+
+
+def test_load_default_config_at_root(base_cfg: DictConfig, mock_download: MagicMock, tmp_path: Path):
+    """Test loading basic config."""
+    artifact_path = tmp_path / "config.yaml"
+    artifact_path.write_text("loaded_key: loaded_value\n")
+    mock_download.return_value = str(artifact_path)
+
+    cfg = OmegaConf.merge(
+        base_cfg,
+        {
+            "mlflow_defaults": ["previous_run"],
+        },
+    )
+
+    loaded_cfg: DictConfig = load_mlflow_defaults(cfg)
+
+    assert loaded_cfg.get("loaded_key") == "loaded_value"
+    assert OmegaConf.select(loaded_cfg, "other_section.foo") == "bar"
+
+
+def test_load_default_config_at_package(base_cfg: DictConfig, mock_download: MagicMock, tmp_path: Path):
+    """Test loading at package."""
+    artifact_path = tmp_path / "config.yaml"
+    artifact_path.write_text("loaded_key: loaded_value\n")
+    mock_download.return_value = str(artifact_path)
+
+    cfg = OmegaConf.merge(
+        base_cfg,
+        {
+            "mlflow_defaults": ["previous_run@dest"],
+        },
+    )
+
+    loaded_cfg: DictConfig = load_mlflow_defaults(cfg)
+    assert OmegaConf.select(loaded_cfg, "dest.loaded_key") == "loaded_value"
+    assert OmegaConf.select(loaded_cfg, "other_section.foo") == "bar"
+    assert mock_download.call_args.kwargs["path"] == "config"
+
+
+def test_load_nondefault_config(base_cfg: DictConfig, mock_download: MagicMock, tmp_path: Path):
+    """Test loading basic config."""
+    artifact_path = tmp_path / "nondefault_config.yaml"
+    artifact_path.write_text("loaded_key: loaded_value\n")
+    mock_download.return_value = str(artifact_path)
+
+    cfg = OmegaConf.merge(
+        base_cfg,
+        {
+            "mlflow_defaults": ["previous_run: nondefault_config#"],
+        },
+    )
+
+    loaded_cfg: DictConfig = load_mlflow_defaults(cfg)
+
+    assert loaded_cfg.get("loaded_key") == "loaded_value"
+    assert OmegaConf.select(loaded_cfg, "other_section.foo") == "bar"
+    assert mock_download.call_args.kwargs["path"] == "nondefault_config"
+
+
+def test_load_subconfig(base_cfg: DictConfig, mock_download: MagicMock, tmp_path: Path):
+    """Test loading subconfig."""
+    artifact_path = tmp_path / "config.yaml"
+    artifact_path.write_text("sub:\n  nested: value\n")
+    mock_download.return_value = str(artifact_path)
+
+    cfg = OmegaConf.merge(
+        base_cfg,
+        {
+            "mlflow_defaults": ["previous_run: sub"],
+        },
+    )
+
+    loaded_cfg: DictConfig = load_mlflow_defaults(cfg)
+    assert loaded_cfg.get("nested") == "value"
+    assert mock_download.call_args.kwargs["path"] == "config"
+
+
+def test_load_subconfig_select(base_cfg: DictConfig, mock_download: MagicMock, tmp_path: Path):
+    """Test loading subconfig select."""
+    artifact_path = tmp_path / "nondefault_config.yaml"
+    artifact_path.write_text("sub:\n  nested: value\n")
+    mock_download.return_value = str(artifact_path)
+
+    cfg = OmegaConf.merge(
+        base_cfg,
+        {
+            "mlflow_defaults": ["previous_run@dest: nondefault_config#sub.nested"],
+        },
+    )
+
+    loaded_cfg: DictConfig = load_mlflow_defaults(cfg)
+    assert loaded_cfg.get("dest") == "value"
+
+
+def test_implicit_config_select(base_cfg: DictConfig, mock_download: MagicMock, tmp_path: Path):
+    """Test implicit artifact="config" when option has no / or #."""
+    artifact_path = tmp_path / "config.yaml"
+    artifact_path.write_text("sub: value\n")
+    mock_download.return_value = str(artifact_path)
+
+    cfg = OmegaConf.merge(
+        base_cfg,
+        {
+            "mlflow_defaults": ["previous_run@dest: sub"],
+        },
+    )
+
+    loaded_cfg: DictConfig = load_mlflow_defaults(cfg)
+    assert loaded_cfg.get("dest") == "value"
+    mock_download.assert_called_with(run_id="test_run_id", path="config", dst_path=ANY)
+
+
+def test_implicit_artifact_select(base_cfg: DictConfig, mock_download: MagicMock, tmp_path: Path):
+    """Test implicit select=root when option has /."""
+    artifact_path = tmp_path / "artifact.yaml"
+    artifact_path.write_text("key: value\n")
+    mock_download.return_value = str(artifact_path)
+
+    cfg = OmegaConf.merge(
+        base_cfg,
+        {
+            "mlflow_defaults": ["previous_run@dest: path/to/artifact"],
+        },
+    )
+    loaded_cfg: DictConfig = load_mlflow_defaults(cfg)
+    assert OmegaConf.select(loaded_cfg, "dest.key") == "value"
+    mock_download.assert_called_with(run_id="test_run_id", path="path/to/artifact", dst_path=ANY)
+
+
+def test_override_flag(base_cfg: DictConfig, mock_download: MagicMock, tmp_path: Path):
+    """Test override flag."""
+    artifact_path = tmp_path / "nondefault_config.yaml"
+    artifact_path.write_text("foo: new_bar\n")
+    mock_download.return_value = str(artifact_path)
+
+    # Place _self_ first so the loaded config overrides the base config
+    cfg = OmegaConf.merge(
+        base_cfg,
+        {
+            "mlflow_defaults": ["_self_", "override previous_run@other_section: nondefault_config#"],
+        },
+    )
+    loaded_cfg: DictConfig = load_mlflow_defaults(cfg)
+    assert OmegaConf.select(loaded_cfg, "other_section.foo") == "new_bar"
+
+
+def test_optional_flag_missing_artifact(base_cfg: DictConfig):
+    """Test optional flag with missing artifact."""
+
+    cfg = OmegaConf.merge(
+        base_cfg,
+        {
+            "mlflow_defaults": ["optional previous_run: null"],
+        },
+    )
+    # Should not raise
+    loaded_cfg: DictConfig = load_mlflow_defaults(cfg)
+    assert loaded_cfg == cfg
+
+
+def test_null_option_error(base_cfg: DictConfig):
+    """Test null option raises error."""
+    cfg = OmegaConf.merge(
+        base_cfg,
+        {
+            "mlflow_defaults": ["previous_run: null"],
+        },
+    )
+    # null option returns just the target config in new logic unless specific logic handles it differently
+    # Actually user logic for option="null" returns _ParsedEntry with artifact_path=None
+    # And _get_target_config handles artifact_path=None by warning and returning None
+    # So "mlflow: null" should be valid but yield nothing in accumulator if not trapped earlier?
+    # Re-reading logic in Step 245:
+    # _parse_entry returns artifact_path=None if option="null"
+    # _get_target_config checks if parsed_entry.artifact_path is None: logs Warning, returns None.
+    # _process_entry sees loaded_config is None.
+    # It raises ValueError("Target config not found for entry...") if not optional.
+    # So it SHOULD raise ValueError.
+
+    with pytest.raises(ValueError, match="Target config not found for entry"):
+        load_mlflow_defaults(cfg)
+
+
+def test_null_option_optional(base_cfg: DictConfig):
+    """Test null option with optional flag."""
+    cfg = OmegaConf.merge(
+        base_cfg,
+        {
+            "mlflow_defaults": ["optional previous_run: null"],
+        },
+    )
+    loaded_cfg: DictConfig = load_mlflow_defaults(cfg)
+    assert OmegaConf.select(loaded_cfg, "other_section") == cfg.other_section
+
+
+def test_composition_order(base_cfg: DictConfig, mock_download: MagicMock, tmp_path: Path):
+    """Test _self_ placement in composition order."""
+    artifact_path = tmp_path / "nondefault_config.yaml"
+    artifact_path.write_text("foo: mlflow_value\n")
+    mock_download.return_value = str(artifact_path)
+
+    # Case 1: _self_ (default/original) before mlflow -> mlflow overrides
+    cfg1 = OmegaConf.merge(base_cfg, {"mlflow_defaults": ["_self_", "previous_run@other_section: nondefault_config#"]})
+    loaded1: DictConfig = load_mlflow_defaults(cfg1)
+    assert OmegaConf.select(loaded1, "other_section.foo") == "mlflow_value"
+
+    # Case 2: mlflow before _self_ -> original overrides
+    cfg2 = OmegaConf.merge(base_cfg, {"mlflow_defaults": ["previous_run@other_section: nondefault_config#", "_self_"]})
+    loaded2: DictConfig = load_mlflow_defaults(cfg2)
+    assert OmegaConf.select(loaded2, "other_section.foo") == "bar"  # Original value
+
+
+def test_config_entry_syntax(base_cfg: DictConfig, mock_download: MagicMock, tmp_path: Path):
+    """Test CONFIG syntax: TARGET (no option)."""
+    artifact_path = tmp_path / "nondefault_config.yaml"
+    artifact_path.write_text("key: val\n")
+    mock_download.return_value = str(artifact_path)
+
+    cfg = OmegaConf.merge(
+        base_cfg,
+        {
+            "mlflow_defaults": ["previous_run"],
+        },
+    )
+    loaded_cfg: DictConfig = load_mlflow_defaults(cfg)
+    assert OmegaConf.select(loaded_cfg, "key") == "val"
+    # Implicitly loads "config" artifact
+    mock_download.assert_called_with(run_id="test_run_id", path="config", dst_path=ANY)
+
+
+def test_missing_target_error(base_cfg: DictConfig):
+    """Test missing target raises error."""
+    cfg = OmegaConf.merge(
+        base_cfg,
+        {
+            "mlflow_defaults": ["optional : option"],
+        },
+    )
+    with pytest.raises(ValueError, match="Invalid MLflow default entry"):
+        load_mlflow_defaults(cfg)
+
+
+def test_invalid_entry_format(base_cfg: DictConfig):
+    """Test invalid entry format raises error (line 64)."""
+    cfg = OmegaConf.merge(
+        base_cfg,
+        {
+            "mlflow_defaults": ["invalid@format@too@many@ats"],
+        },
+    )
+    with pytest.raises(ValueError, match="Invalid MLflow default entry"):
+        load_mlflow_defaults(cfg)
+
+
+def test_missing_target_node(base_cfg: DictConfig):
+    """Test missing target node in config (lines 96-97)."""
+    cfg = OmegaConf.merge(
+        base_cfg,
+        {
+            "mlflow_defaults": ["nonexistent: config"],
+        },
+    )
+    # Should raise ValueError because target node not found and not optional
+    with pytest.raises(ValueError, match="Target config not found for entry"):
+        load_mlflow_defaults(cfg)
+
+
+def test_missing_target_node_optional(base_cfg: DictConfig):
+    """Test missing target node with optional flag."""
+    cfg = OmegaConf.merge(
+        base_cfg,
+        {
+            "mlflow_defaults": ["optional nonexistent: config"],
+        },
+    )
+    # Should not raise, just return accumulator
+    loaded_cfg: DictConfig = load_mlflow_defaults(cfg)
+    assert OmegaConf.select(loaded_cfg, "other_section") == cfg.other_section
+
+
+def test_validate_mlflow_config_error(base_cfg: DictConfig, mocker):
+    """Test ConfigValidationError from validate_mlflow_config (lines 101-103)."""
+    # Make validate_mlflow_config raise ConfigValidationError
+    mock_validate = mocker.patch("simplexity.structured_configs.mlflow_defaults.validate_mlflow_config")
+    mock_validate.side_effect = ConfigValidationError("Invalid config")
+
+    cfg = OmegaConf.merge(
+        base_cfg,
+        {
+            "mlflow_defaults": ["previous_run: config"],
+        },
+    )
+    # Should raise ValueError because validation failed and not optional
+    with pytest.raises(ValueError, match="Target config not found for entry"):
+        load_mlflow_defaults(cfg)
+
+
+def test_validate_mlflow_config_error_optional(base_cfg: DictConfig, mocker):
+    """Test ConfigValidationError with optional flag."""
+    # Make validate_mlflow_config raise ConfigValidationError
+    mock_validate = mocker.patch("simplexity.structured_configs.mlflow_defaults.validate_mlflow_config")
+    mock_validate.side_effect = ConfigValidationError("Invalid config")
+
+    cfg = OmegaConf.merge(
+        base_cfg,
+        {
+            "mlflow_defaults": ["optional previous_run: config"],
+        },
+    )
+    # Should not raise, just return accumulator
+    loaded_cfg: DictConfig = load_mlflow_defaults(cfg)
+    assert OmegaConf.select(loaded_cfg, "other_section") == cfg.other_section
+
+
+def test_resolve_mlflow_config_error(base_cfg: DictConfig, mocker):
+    """Test ValueError from resolve_mlflow_config (lines 107-109)."""
+    # Make resolve_mlflow_config raise ValueError
+    mock_resolve = mocker.patch("simplexity.structured_configs.mlflow_defaults.resolve_mlflow_config")
+    mock_resolve.side_effect = ValueError("Resolution failed")
+
+    cfg = OmegaConf.merge(
+        base_cfg,
+        {
+            "mlflow_defaults": ["previous_run: config"],
+        },
+    )
+    # Should raise ValueError because resolution failed and not optional
+    with pytest.raises(ValueError, match="Target config not found for entry"):
+        load_mlflow_defaults(cfg)
+
+
+def test_resolve_mlflow_config_error_optional(base_cfg: DictConfig, mocker):
+    """Test ValueError from resolve_mlflow_config with optional flag."""
+    # Make resolve_mlflow_config raise ValueError
+    mock_resolve = mocker.patch("simplexity.structured_configs.mlflow_defaults.resolve_mlflow_config")
+    mock_resolve.side_effect = ValueError("Resolution failed")
+
+    cfg = OmegaConf.merge(
+        base_cfg,
+        {
+            "mlflow_defaults": ["optional mlflow: config"],
+        },
+    )
+    # Should not raise, just return accumulator
+    loaded_cfg: DictConfig = load_mlflow_defaults(cfg)
+    assert OmegaConf.select(loaded_cfg, "other_section") == cfg.other_section
+
+
+def test_omegaconf_load_error(base_cfg: DictConfig, mock_download: MagicMock, tmp_path: Path):
+    """Test Exception from OmegaConf.load (lines 124-126)."""
+    # Create a file with invalid YAML content
+    artifact_path = tmp_path / "invalid.yaml"
+    artifact_path.write_text("invalid: yaml: content: [unclosed")
+    mock_download.return_value = str(artifact_path)
+
+    cfg = OmegaConf.merge(
+        base_cfg,
+        {
+            "mlflow_defaults": ["mlflow: config"],
+        },
+    )
+    # Should raise ValueError because load failed and not optional
+    with pytest.raises(ValueError, match="Target config not found for entry"):
+        load_mlflow_defaults(cfg)
+
+
+def test_omegaconf_load_error_optional(base_cfg: DictConfig, mock_download: MagicMock, tmp_path: Path):
+    """Test Exception from OmegaConf.load with optional flag."""
+    # Create a file with invalid YAML content
+    artifact_path = tmp_path / "invalid.yaml"
+    artifact_path.write_text("invalid: yaml: content: [unclosed")
+    mock_download.return_value = str(artifact_path)
+
+    cfg = OmegaConf.merge(
+        base_cfg,
+        {
+            "mlflow_defaults": ["optional previous_run: config"],
+        },
+    )
+    # Should not raise, just return accumulator
+    loaded_cfg: DictConfig = load_mlflow_defaults(cfg)
+    assert OmegaConf.select(loaded_cfg, "other_section") == cfg.other_section
+
+
+def test_select_path_not_found(base_cfg: DictConfig, mock_download: MagicMock, tmp_path: Path):
+    """Test select path not found in artifact (lines 133-136)."""
+    artifact_path = tmp_path / "nondefault_config.yaml"
+    artifact_path.write_text("key: value\n")
+    mock_download.return_value = str(artifact_path)
+
+    cfg = OmegaConf.merge(
+        base_cfg,
+        {
+            "mlflow_defaults": ["previous_run@dest: nondefault_config#nonexistent.path"],
+        },
+    )
+    # Should raise ValueError because select path not found and not optional
+    with pytest.raises(ValueError, match="Target config not found for entry"):
+        load_mlflow_defaults(cfg)
+
+
+def test_select_path_not_found_optional(base_cfg: DictConfig, mock_download: MagicMock, tmp_path: Path):
+    """Test select path not found with optional flag."""
+    artifact_path = tmp_path / "nondefault_config.yaml"
+    artifact_path.write_text("key: value\n")
+    mock_download.return_value = str(artifact_path)
+
+    cfg = OmegaConf.merge(
+        base_cfg,
+        {
+            "mlflow_defaults": ["optional previous_run@dest: nondefault_config#nonexistent.path"],
+        },
+    )
+    # Should not raise, just return accumulator
+    loaded_cfg: DictConfig = load_mlflow_defaults(cfg)
+    assert OmegaConf.select(loaded_cfg, "other_section") == cfg.other_section
+
+
+def test_non_dictconfig_error(base_cfg: DictConfig, mock_download: MagicMock, tmp_path: Path):
+    """Test non-DictConfig error when not optional (lines 157-159)."""
+    # Create artifact with list (non-dict) root value
+    # Use "nondefault_config#" to avoid select_path, so we get the full ListConfig
+    artifact_path = tmp_path / "nondefault_config.yaml"
+    artifact_path.write_text("- item1\n- item2\n")
+    mock_download.return_value = str(artifact_path)
+
+    cfg = OmegaConf.merge(
+        base_cfg,
+        {
+            "mlflow_defaults": ["previous_run: nondefault_config#"],
+        },
+    )
+    # Should raise ValueError because loaded_config is ListConfig (not DictConfig) and not optional
+    with pytest.raises(ValueError, match="Target config not found for entry"):
+        load_mlflow_defaults(cfg)
+
+
+def test_non_dictconfig_optional(base_cfg: DictConfig, mock_download: MagicMock, tmp_path: Path):
+    """Test non-DictConfig with optional flag."""
+    # Create artifact with list (non-dict) root value
+    # Use "nondefault_config#" to avoid select_path, so we get the full ListConfig
+    artifact_path = tmp_path / "nondefault_config.yaml"
+    artifact_path.write_text("- item1\n- item2\n")
+    mock_download.return_value = str(artifact_path)
+
+    cfg = OmegaConf.merge(
+        base_cfg,
+        {
+            "mlflow_defaults": ["optional previous_run: nondefault_config#"],
+        },
+    )
+    # Should not raise, just return accumulator
+    loaded_cfg: DictConfig = load_mlflow_defaults(cfg)
+    assert OmegaConf.select(loaded_cfg, "other_section") == cfg.other_section
+
+
+def test_no_mlflow_defaults_key():
+    """Test early return when mlflow_defaults key is missing (line 172)."""
+    # Config without mlflow_defaults key
+    cfg = OmegaConf.create(
+        {
+            "previous_run": {
+                "tracking_uri": "databricks",
+                "run_id": "test_run_id",
+            },
+            "other_section": {"foo": "bar"},
+        }
+    )
+    # Should return cfg unchanged
+    loaded_cfg: DictConfig = load_mlflow_defaults(cfg)
+    assert loaded_cfg == cfg
+    assert OmegaConf.select(loaded_cfg, "other_section.foo") == "bar"

From 86d71856b6b90507ae75797c4e4ac3d3987d2b8f Mon Sep 17 00:00:00 2001
From: Eric Alt <ericallenalt@gmail.com>
Date: Thu, 15 Jan 2026 15:20:14 -0800
Subject: [PATCH 05/31] Fix test setup

---
 .../test_load_config_from_mlflow.py           | 29 +++++++++++--------
 1 file changed, 17 insertions(+), 12 deletions(-)

diff --git a/tests/end_to_end/test_load_config_from_mlflow.py b/tests/end_to_end/test_load_config_from_mlflow.py
index 573418f6..3c305f91 100644
--- a/tests/end_to_end/test_load_config_from_mlflow.py
+++ b/tests/end_to_end/test_load_config_from_mlflow.py
@@ -15,20 +15,20 @@
 import mlflow
 import pytest
 from hydra import compose, initialize_config_dir
-from omegaconf import DictConfig
+from omegaconf import DictConfig, OmegaConf
 
 from simplexity.logging.mlflow_logger import MLFlowLogger
 from simplexity.structured_configs.mlflow_defaults import load_mlflow_defaults
 from simplexity.utils.mlflow_utils import get_experiment, get_run
 
-CONFIGS_SRC = str(Path(__file__).parent / "mlflow_defaults_configs" / "setup")
+CONFIG_DIR = str(Path(__file__).parent / "mlflow_defaults_configs")
 
 EXPERIMENT_NAME = "test_mlflow_defaults"
 
 MLFLOW_CONFIG = """
 mlflow:
-  experiment_name: {experiment_name}
-  run_name: {run_name}
+  experiment_id: {experiment_id}
+  run_id: {run_id}
   tracking_uri: {tracking_uri}
 """
 
@@ -38,7 +38,7 @@ def setup_dir(tmp_path_factory: pytest.TempPathFactory) -> Path:
     """Setup function."""
     tmp_path = tmp_path_factory.mktemp("mlflow_defaults")
     config_path = tmp_path / "configs"
-    shutil.copytree(CONFIGS_SRC, str(config_path))
+    shutil.copytree(f"{CONFIG_DIR}/setup", str(config_path))
     tracking_uri = f"sqlite:///{tmp_path.resolve()}/mlflow.db"
     mlflow.set_tracking_uri(tracking_uri)
 
@@ -71,7 +71,7 @@ def log_to_mlflow(cfg: DictConfig | None, config_names: list[str], run_name: str
             if cfg is not None:
                 logger.log_config(cfg, resolve=False)
             for config_name in config_names:
-                logger.log_artifact(local_path=str(config_path / config_name), artifact_path="subdir/special.yaml")
+                logger.log_artifact(local_path=str(config_path / config_name), artifact_path="subdir")
         return experiment_id, run_id
 
     def previous_run(config_name: str | None, config_names: list[str], run_name: str) -> None:
@@ -117,13 +117,18 @@ def test_setup(setup_dir: Path, tmp_path: Path) -> None:
         config_path = client.download_artifacts(
             run_id=run_id, path="config.yaml", dst_path=str(tmp_path / f"config_{run_num}.yaml")
         )
-        expected_config_path = setup_dir / "configs" / f"prev_run_{run_num}.yaml"
-        with open(config_path, encoding="utf-8") as f, open(expected_config_path, encoding="utf-8") as expected_f:
-            assert f.read() == expected_f.read()
 
-        special_config_path = client.download_artifacts(
-            run_id=run_id, path="subdir/special.yaml", dst_path=str(tmp_path / f"special_{run_num}.yaml")
+        actual = OmegaConf.load(config_path)
+        with initialize_config_dir(config_dir=str(setup_dir / "configs")):
+            expected = compose(config_name=f"prev_config_{run_num}.yaml")
+        assert actual == expected
+
+        downloaded_dir = client.download_artifacts(
+            run_id=run_id,
+            path="subdir",
+            dst_path=str(tmp_path),
         )
+        special_config_path = Path(downloaded_dir) / f"special_{run_num}.yaml"
         expected_special_config_path = setup_dir / "configs" / f"special_{run_num}.yaml"
         with (
             open(special_config_path, encoding="utf-8") as f,
@@ -139,7 +144,7 @@ def test_setup(setup_dir: Path, tmp_path: Path) -> None:
 def test_mlflow_defaults(setup_dir: Path, test_case: str) -> None:
     """Test mlflow defaults."""
     tracking_uri = f"sqlite:///{setup_dir.resolve()}/mlflow.db"
-    with initialize_config_dir(config_dir=str(setup_dir / "configs")):
+    with initialize_config_dir(config_dir=CONFIG_DIR):
         cfg = compose(config_name=test_case, overrides=[f"mlflow.tracking_uri={tracking_uri}"])
         expected = compose(config_name=f"{test_case}_expected", overrides=[f"mlflow.tracking_uri={tracking_uri}"])
     actual = load_mlflow_defaults(cfg)

From 8d106d6f15bdc489e775ece155080f5553324080 Mon Sep 17 00:00:00 2001
From: Eric Alt <ericallenalt@gmail.com>
Date: Thu, 15 Jan 2026 15:37:42 -0800
Subject: [PATCH 06/31] Create more configs

---
 .../full_default_copy_expected.yaml                   | 11 +++++++++++
 .../mlflow_defaults_configs/mlflow/prev_run_1.yaml    |  3 +++
 .../mlflow_defaults_configs/mlflow/prev_run_2.yaml    |  3 +++
 3 files changed, 17 insertions(+)
 create mode 100644 tests/end_to_end/mlflow_defaults_configs/full_default_copy_expected.yaml
 create mode 100644 tests/end_to_end/mlflow_defaults_configs/mlflow/prev_run_1.yaml
 create mode 100644 tests/end_to_end/mlflow_defaults_configs/mlflow/prev_run_2.yaml

diff --git a/tests/end_to_end/mlflow_defaults_configs/full_default_copy_expected.yaml b/tests/end_to_end/mlflow_defaults_configs/full_default_copy_expected.yaml
new file mode 100644
index 00000000..5b0c86f8
--- /dev/null
+++ b/tests/end_to_end/mlflow_defaults_configs/full_default_copy_expected.yaml
@@ -0,0 +1,11 @@
+key: 1
+dict:
+  subkey: 10
+group_a:
+  a_key: 1
+  a_dict:
+    a_subkey: 10
+group_b:
+  b_key: 1
+  b_dict:
+    b_subkey: 10
diff --git a/tests/end_to_end/mlflow_defaults_configs/mlflow/prev_run_1.yaml b/tests/end_to_end/mlflow_defaults_configs/mlflow/prev_run_1.yaml
new file mode 100644
index 00000000..27c1b0c8
--- /dev/null
+++ b/tests/end_to_end/mlflow_defaults_configs/mlflow/prev_run_1.yaml
@@ -0,0 +1,3 @@
+experiment_name: test_mlflow_defaults
+run_name: prev_run_1
+tracking_uri: placeholder
\ No newline at end of file
diff --git a/tests/end_to_end/mlflow_defaults_configs/mlflow/prev_run_2.yaml b/tests/end_to_end/mlflow_defaults_configs/mlflow/prev_run_2.yaml
new file mode 100644
index 00000000..d0a49224
--- /dev/null
+++ b/tests/end_to_end/mlflow_defaults_configs/mlflow/prev_run_2.yaml
@@ -0,0 +1,3 @@
+experiment_name: test_mlflow_defaults
+run_name: prev_run_2
+tracking_uri: placeholder
\ No newline at end of file

From 8b3a28fa4c27efd2306ca79d1e6967d6d8851496 Mon Sep 17 00:00:00 2001
From: Eric Alt <ericallenalt@gmail.com>
Date: Thu, 15 Jan 2026 15:46:00 -0800
Subject: [PATCH 07/31] Test fixes

---
 .../full_default_copy.yaml                    |  2 +-
 .../test_load_config_from_mlflow.py           | 37 ++-----------------
 2 files changed, 5 insertions(+), 34 deletions(-)

diff --git a/tests/end_to_end/mlflow_defaults_configs/full_default_copy.yaml b/tests/end_to_end/mlflow_defaults_configs/full_default_copy.yaml
index a1fe6e49..dce6d0fc 100644
--- a/tests/end_to_end/mlflow_defaults_configs/full_default_copy.yaml
+++ b/tests/end_to_end/mlflow_defaults_configs/full_default_copy.yaml
@@ -1,4 +1,4 @@
 defaults:
   - mlflow@load_source: prev_run_1
 mlflow_defaults:
-  - load_source
\ No newline at end of file
+  - load_source: config
\ No newline at end of file
diff --git a/tests/end_to_end/test_load_config_from_mlflow.py b/tests/end_to_end/test_load_config_from_mlflow.py
index 3c305f91..f6cb0d9e 100644
--- a/tests/end_to_end/test_load_config_from_mlflow.py
+++ b/tests/end_to_end/test_load_config_from_mlflow.py
@@ -25,13 +25,6 @@
 
 EXPERIMENT_NAME = "test_mlflow_defaults"
 
-MLFLOW_CONFIG = """
-mlflow:
-  experiment_id: {experiment_id}
-  run_id: {run_id}
-  tracking_uri: {tracking_uri}
-"""
-
 
 @pytest.fixture(scope="session")
 def setup_dir(tmp_path_factory: pytest.TempPathFactory) -> Path:
@@ -42,18 +35,7 @@ def setup_dir(tmp_path_factory: pytest.TempPathFactory) -> Path:
     tracking_uri = f"sqlite:///{tmp_path.resolve()}/mlflow.db"
     mlflow.set_tracking_uri(tracking_uri)
 
-    # for config_file in Path(config_path).rglob("*.yaml"):
-    #     lines = config_file.read_text().splitlines()
-    #     new_lines = []
-    #     for line in lines:
-    #         if "tracking_uri:" in line:
-    #             key, _ = line.split(":", 1)
-    #             new_lines.append(f"{key}: {tracking_uri}")
-    #         else:
-    #             new_lines.append(line)
-    #     config_file.write_text("\n".join(new_lines) + "\n")
-
-    def log_to_mlflow(cfg: DictConfig | None, config_names: list[str], run_name: str) -> tuple[str, str]:
+    def log_to_mlflow(cfg: DictConfig | None, config_names: list[str], run_name: str) -> None:
         """Save config."""
         experiment = get_experiment(experiment_name=EXPERIMENT_NAME)
         assert experiment is not None
@@ -72,7 +54,6 @@ def log_to_mlflow(cfg: DictConfig | None, config_names: list[str], run_name: str
                 logger.log_config(cfg, resolve=False)
             for config_name in config_names:
                 logger.log_artifact(local_path=str(config_path / config_name), artifact_path="subdir")
-        return experiment_id, run_id
 
     def previous_run(config_name: str | None, config_names: list[str], run_name: str) -> None:
         """Previous run."""
@@ -81,17 +62,7 @@ def previous_run(config_name: str | None, config_names: list[str], run_name: str
         else:
             with initialize_config_dir(config_dir=str(config_path)):
                 cfg = compose(config_name=config_name)
-        experiment_id, run_id = log_to_mlflow(cfg, config_names, run_name)
-        output_path = config_path / "mlflow" / f"{run_name}.yaml"
-        output_path.parent.mkdir(parents=True, exist_ok=True)
-        with open(output_path, "w", encoding="utf-8") as f:
-            f.write(
-                MLFLOW_CONFIG.format(
-                    experiment_id=experiment_id,
-                    run_id=run_id,
-                    tracking_uri=tracking_uri,
-                )
-            )
+        log_to_mlflow(cfg, config_names, run_name)
 
     previous_run(config_name="prev_config_1", config_names=["special_1.yaml"], run_name="prev_run_1")
     previous_run(config_name="prev_config_2", config_names=["special_2.yaml"], run_name="prev_run_2")
@@ -145,7 +116,7 @@ def test_mlflow_defaults(setup_dir: Path, test_case: str) -> None:
     """Test mlflow defaults."""
     tracking_uri = f"sqlite:///{setup_dir.resolve()}/mlflow.db"
     with initialize_config_dir(config_dir=CONFIG_DIR):
-        cfg = compose(config_name=test_case, overrides=[f"mlflow.tracking_uri={tracking_uri}"])
-        expected = compose(config_name=f"{test_case}_expected", overrides=[f"mlflow.tracking_uri={tracking_uri}"])
+        cfg = compose(config_name=test_case, overrides=[f"load_source.tracking_uri={tracking_uri}"])
+        expected = compose(config_name=f"{test_case}_expected")
     actual = load_mlflow_defaults(cfg)
     assert actual == expected

From ceaf397b49944392f3528939f331ee35a9d1eac7 Mon Sep 17 00:00:00 2001
From: Eric Alt <ericallenalt@gmail.com>
Date: Thu, 15 Jan 2026 16:27:38 -0800
Subject: [PATCH 08/31] Fix test

---
 .../structured_configs/mlflow_defaults.py     | 31 +++++++++++++++++--
 .../full_default_copy_expected.yaml           |  7 +++++
 .../test_load_config_from_mlflow.py           | 15 ++++++++-
 3 files changed, 50 insertions(+), 3 deletions(-)

diff --git a/simplexity/structured_configs/mlflow_defaults.py b/simplexity/structured_configs/mlflow_defaults.py
index 9b04ab78..3c1725b5 100644
--- a/simplexity/structured_configs/mlflow_defaults.py
+++ b/simplexity/structured_configs/mlflow_defaults.py
@@ -110,11 +110,14 @@ def _get_target_config(cfg: DictConfig, parsed_entry: _ParsedEntry) -> Any | Non
 
     tracking_uri: str | None = target_node.get("tracking_uri")
     run_id: str = target_node.get("run_id")
+    artifact_path = parsed_entry.artifact_path
+    if not artifact_path.endswith((".yaml", ".yml")):
+        artifact_path = f"{artifact_path}.yaml"
     client = MlflowClient(tracking_uri=tracking_uri)
 
     with tempfile.TemporaryDirectory() as tmp_dir:
         try:
-            local_path = client.download_artifacts(run_id=run_id, path=parsed_entry.artifact_path, dst_path=tmp_dir)
+            local_path = client.download_artifacts(run_id=run_id, path=artifact_path, dst_path=tmp_dir)
         except Exception as e:
             SIMPLEXITY_LOGGER.warning("Failed to download artifact from MLflow '%s': %s", parsed_entry.target, e)
             return None
@@ -138,6 +141,28 @@ def _get_target_config(cfg: DictConfig, parsed_entry: _ParsedEntry) -> Any | Non
     return selected_config
 
 
+def _normalize_item(item: str | DictConfig) -> str:
+    """Normalize an item to a string format.
+
+    If item is a DictConfig with a single key-value pair:
+    - If value is "config" (the default artifact), treat as just the key (no option)
+    - Otherwise, convert to "key: value" format
+    Otherwise, convert to string.
+    """
+    if isinstance(item, DictConfig):
+        keys = list(item.keys())
+        if len(keys) == 1:
+            key = keys[0]
+            value = item[key]
+            # If value is "config" (default artifact), treat as simple CONFIG entry
+            if value == "config":
+                return key
+            return f"{key}: {value}"
+        # Multiple keys - convert entire dict to string representation
+        return str(item)
+    return str(item)
+
+
 def _process_entry(cfg: DictConfig, accumulator: DictConfig, item: str) -> DictConfig:
     """Process a single MLflow default item."""
     if item == "_self_":
@@ -177,6 +202,8 @@ def load_mlflow_defaults(cfg: DictConfig) -> DictConfig:
     accumulator = OmegaConf.create()
 
     for item in mlflow_defaults:
-        accumulator = _process_entry(cfg, accumulator, item)
+        normalized_item = _normalize_item(item)
+        accumulator = _process_entry(cfg, accumulator, normalized_item)
 
+    del accumulator["mlflow_defaults"]
     return accumulator
diff --git a/tests/end_to_end/mlflow_defaults_configs/full_default_copy_expected.yaml b/tests/end_to_end/mlflow_defaults_configs/full_default_copy_expected.yaml
index 5b0c86f8..610debcc 100644
--- a/tests/end_to_end/mlflow_defaults_configs/full_default_copy_expected.yaml
+++ b/tests/end_to_end/mlflow_defaults_configs/full_default_copy_expected.yaml
@@ -9,3 +9,10 @@ group_b:
   b_key: 1
   b_dict:
     b_subkey: 10
+load_source:
+  experiment_name: test_mlflow_defaults
+  run_name: prev_run_1
+  tracking_uri: sqlite:///mlflow.db  # placeholder
+  registry_uri: null
+  experiment_id: "1"  # placeholder
+  run_id: "00000000000000000000000000000000"  # placeholder
diff --git a/tests/end_to_end/test_load_config_from_mlflow.py b/tests/end_to_end/test_load_config_from_mlflow.py
index f6cb0d9e..c6743aa3 100644
--- a/tests/end_to_end/test_load_config_from_mlflow.py
+++ b/tests/end_to_end/test_load_config_from_mlflow.py
@@ -117,6 +117,19 @@ def test_mlflow_defaults(setup_dir: Path, test_case: str) -> None:
     tracking_uri = f"sqlite:///{setup_dir.resolve()}/mlflow.db"
     with initialize_config_dir(config_dir=CONFIG_DIR):
         cfg = compose(config_name=test_case, overrides=[f"load_source.tracking_uri={tracking_uri}"])
-        expected = compose(config_name=f"{test_case}_expected")
     actual = load_mlflow_defaults(cfg)
+    experiment_id = OmegaConf.select(actual, "load_source.experiment_id")
+    assert experiment_id is not None
+    run_id = OmegaConf.select(actual, "load_source.run_id")
+    assert run_id is not None
+    with initialize_config_dir(config_dir=CONFIG_DIR):
+        expected = compose(
+            config_name=f"{test_case}_expected",
+            overrides=[
+                f"load_source.tracking_uri={tracking_uri}",
+                f'load_source.experiment_id="{experiment_id}"',
+                f'load_source.run_id="{run_id}"',
+            ],
+        )
+
     assert actual == expected

From 9230bfe9e3e89b0be38a220ea3188eeb6edede3f Mon Sep 17 00:00:00 2001
From: Eric Alt <ericallenalt@gmail.com>
Date: Thu, 15 Jan 2026 16:28:27 -0800
Subject: [PATCH 09/31] Fix test

---
 tests/end_to_end/mlflow_defaults_configs/full_default_copy.yaml | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tests/end_to_end/mlflow_defaults_configs/full_default_copy.yaml b/tests/end_to_end/mlflow_defaults_configs/full_default_copy.yaml
index dce6d0fc..a1fe6e49 100644
--- a/tests/end_to_end/mlflow_defaults_configs/full_default_copy.yaml
+++ b/tests/end_to_end/mlflow_defaults_configs/full_default_copy.yaml
@@ -1,4 +1,4 @@
 defaults:
   - mlflow@load_source: prev_run_1
 mlflow_defaults:
-  - load_source: config
\ No newline at end of file
+  - load_source
\ No newline at end of file

From c6008437ba39a5369e9e9b3048a5cd9b6e47b5ce Mon Sep 17 00:00:00 2001
From: Eric Alt <ericallenalt@gmail.com>
Date: Thu, 15 Jan 2026 16:45:47 -0800
Subject: [PATCH 10/31] Create new tests

---
 .../composition_order_self_first.yaml         |  7 +++++++
 ...composition_order_self_first_expected.yaml | 19 +++++++++++++++++++
 .../composition_order_self_last.yaml          |  7 +++++++
 .../composition_order_self_last_expected.yaml |  9 +++++++++
 .../implicit_artifact_select.yaml             |  4 ++++
 .../implicit_artifact_select_expected.yaml    | 19 +++++++++++++++++++
 .../implicit_config_select.yaml               |  4 ++++
 .../implicit_config_select_expected.yaml      | 11 +++++++++++
 .../load_default_config_at_package.yaml       |  4 ++++
 ...ad_default_config_at_package_expected.yaml | 19 +++++++++++++++++++
 .../load_nondefault_config.yaml               |  4 ++++
 .../load_nondefault_config_expected.yaml      | 18 ++++++++++++++++++
 .../load_subconfig.yaml                       |  4 ++++
 .../load_subconfig_expected.yaml              | 11 +++++++++++
 .../load_subconfig_select.yaml                |  4 ++++
 .../load_subconfig_select_expected.yaml       |  8 ++++++++
 .../override_flag.yaml                        |  7 +++++++
 .../override_flag_expected.yaml               | 19 +++++++++++++++++++
 .../test_load_config_from_mlflow.py           | 13 ++++++++++++-
 19 files changed, 190 insertions(+), 1 deletion(-)
 create mode 100644 tests/end_to_end/mlflow_defaults_configs/composition_order_self_first.yaml
 create mode 100644 tests/end_to_end/mlflow_defaults_configs/composition_order_self_first_expected.yaml
 create mode 100644 tests/end_to_end/mlflow_defaults_configs/composition_order_self_last.yaml
 create mode 100644 tests/end_to_end/mlflow_defaults_configs/composition_order_self_last_expected.yaml
 create mode 100644 tests/end_to_end/mlflow_defaults_configs/implicit_artifact_select.yaml
 create mode 100644 tests/end_to_end/mlflow_defaults_configs/implicit_artifact_select_expected.yaml
 create mode 100644 tests/end_to_end/mlflow_defaults_configs/implicit_config_select.yaml
 create mode 100644 tests/end_to_end/mlflow_defaults_configs/implicit_config_select_expected.yaml
 create mode 100644 tests/end_to_end/mlflow_defaults_configs/load_default_config_at_package.yaml
 create mode 100644 tests/end_to_end/mlflow_defaults_configs/load_default_config_at_package_expected.yaml
 create mode 100644 tests/end_to_end/mlflow_defaults_configs/load_nondefault_config.yaml
 create mode 100644 tests/end_to_end/mlflow_defaults_configs/load_nondefault_config_expected.yaml
 create mode 100644 tests/end_to_end/mlflow_defaults_configs/load_subconfig.yaml
 create mode 100644 tests/end_to_end/mlflow_defaults_configs/load_subconfig_expected.yaml
 create mode 100644 tests/end_to_end/mlflow_defaults_configs/load_subconfig_select.yaml
 create mode 100644 tests/end_to_end/mlflow_defaults_configs/load_subconfig_select_expected.yaml
 create mode 100644 tests/end_to_end/mlflow_defaults_configs/override_flag.yaml
 create mode 100644 tests/end_to_end/mlflow_defaults_configs/override_flag_expected.yaml

diff --git a/tests/end_to_end/mlflow_defaults_configs/composition_order_self_first.yaml b/tests/end_to_end/mlflow_defaults_configs/composition_order_self_first.yaml
new file mode 100644
index 00000000..92e56a7a
--- /dev/null
+++ b/tests/end_to_end/mlflow_defaults_configs/composition_order_self_first.yaml
@@ -0,0 +1,7 @@
+defaults:
+  - mlflow@load_source: prev_run_1
+other_section:
+  foo: bar
+mlflow_defaults:
+  - _self_
+  - load_source@other_section: subdir/special_1#
diff --git a/tests/end_to_end/mlflow_defaults_configs/composition_order_self_first_expected.yaml b/tests/end_to_end/mlflow_defaults_configs/composition_order_self_first_expected.yaml
new file mode 100644
index 00000000..8d4d28d3
--- /dev/null
+++ b/tests/end_to_end/mlflow_defaults_configs/composition_order_self_first_expected.yaml
@@ -0,0 +1,19 @@
+other_section:
+  key: 3
+  dict:
+    subkey: 30
+  group_a:
+    a_key: 3
+    a_dict:
+      a_subkey: 30
+  group_b:
+    b_key: 3
+    b_dict:
+      b_subkey: 30
+load_source:
+  experiment_name: test_mlflow_defaults
+  run_name: prev_run_1
+  tracking_uri: sqlite:///mlflow.db  # placeholder
+  registry_uri: null
+  experiment_id: "1"  # placeholder
+  run_id: "00000000000000000000000000000000"  # placeholder
diff --git a/tests/end_to_end/mlflow_defaults_configs/composition_order_self_last.yaml b/tests/end_to_end/mlflow_defaults_configs/composition_order_self_last.yaml
new file mode 100644
index 00000000..95bb779b
--- /dev/null
+++ b/tests/end_to_end/mlflow_defaults_configs/composition_order_self_last.yaml
@@ -0,0 +1,7 @@
+defaults:
+  - mlflow@load_source: prev_run_1
+other_section:
+  foo: bar
+mlflow_defaults:
+  - load_source@other_section: subdir/special_1#
+  - _self_
diff --git a/tests/end_to_end/mlflow_defaults_configs/composition_order_self_last_expected.yaml b/tests/end_to_end/mlflow_defaults_configs/composition_order_self_last_expected.yaml
new file mode 100644
index 00000000..08545844
--- /dev/null
+++ b/tests/end_to_end/mlflow_defaults_configs/composition_order_self_last_expected.yaml
@@ -0,0 +1,9 @@
+other_section:
+  foo: bar
+load_source:
+  experiment_name: test_mlflow_defaults
+  run_name: prev_run_1
+  tracking_uri: sqlite:///mlflow.db  # placeholder
+  registry_uri: null
+  experiment_id: "1"  # placeholder
+  run_id: "00000000000000000000000000000000"  # placeholder
diff --git a/tests/end_to_end/mlflow_defaults_configs/implicit_artifact_select.yaml b/tests/end_to_end/mlflow_defaults_configs/implicit_artifact_select.yaml
new file mode 100644
index 00000000..e8451c42
--- /dev/null
+++ b/tests/end_to_end/mlflow_defaults_configs/implicit_artifact_select.yaml
@@ -0,0 +1,4 @@
+defaults:
+  - mlflow@load_source: prev_run_1
+mlflow_defaults:
+  - load_source@dest: subdir/special_1
diff --git a/tests/end_to_end/mlflow_defaults_configs/implicit_artifact_select_expected.yaml b/tests/end_to_end/mlflow_defaults_configs/implicit_artifact_select_expected.yaml
new file mode 100644
index 00000000..61848826
--- /dev/null
+++ b/tests/end_to_end/mlflow_defaults_configs/implicit_artifact_select_expected.yaml
@@ -0,0 +1,19 @@
+dest:
+  key: 3
+  dict:
+    subkey: 30
+  group_a:
+    a_key: 3
+    a_dict:
+      a_subkey: 30
+  group_b:
+    b_key: 3
+    b_dict:
+      b_subkey: 30
+load_source:
+  experiment_name: test_mlflow_defaults
+  run_name: prev_run_1
+  tracking_uri: sqlite:///mlflow.db  # placeholder
+  registry_uri: null
+  experiment_id: "1"  # placeholder
+  run_id: "00000000000000000000000000000000"  # placeholder
diff --git a/tests/end_to_end/mlflow_defaults_configs/implicit_config_select.yaml b/tests/end_to_end/mlflow_defaults_configs/implicit_config_select.yaml
new file mode 100644
index 00000000..e79d0038
--- /dev/null
+++ b/tests/end_to_end/mlflow_defaults_configs/implicit_config_select.yaml
@@ -0,0 +1,4 @@
+defaults:
+  - mlflow@load_source: prev_run_1
+mlflow_defaults:
+  - load_source@dest: group_a
diff --git a/tests/end_to_end/mlflow_defaults_configs/implicit_config_select_expected.yaml b/tests/end_to_end/mlflow_defaults_configs/implicit_config_select_expected.yaml
new file mode 100644
index 00000000..dbf6f5f3
--- /dev/null
+++ b/tests/end_to_end/mlflow_defaults_configs/implicit_config_select_expected.yaml
@@ -0,0 +1,11 @@
+dest:
+  a_key: 1
+  a_dict:
+    a_subkey: 10
+load_source:
+  experiment_name: test_mlflow_defaults
+  run_name: prev_run_1
+  tracking_uri: sqlite:///mlflow.db  # placeholder
+  registry_uri: null
+  experiment_id: "1"  # placeholder
+  run_id: "00000000000000000000000000000000"  # placeholder
diff --git a/tests/end_to_end/mlflow_defaults_configs/load_default_config_at_package.yaml b/tests/end_to_end/mlflow_defaults_configs/load_default_config_at_package.yaml
new file mode 100644
index 00000000..7fdc40f9
--- /dev/null
+++ b/tests/end_to_end/mlflow_defaults_configs/load_default_config_at_package.yaml
@@ -0,0 +1,4 @@
+defaults:
+  - mlflow@load_source: prev_run_1
+mlflow_defaults:
+  - load_source@dest
diff --git a/tests/end_to_end/mlflow_defaults_configs/load_default_config_at_package_expected.yaml b/tests/end_to_end/mlflow_defaults_configs/load_default_config_at_package_expected.yaml
new file mode 100644
index 00000000..c741220c
--- /dev/null
+++ b/tests/end_to_end/mlflow_defaults_configs/load_default_config_at_package_expected.yaml
@@ -0,0 +1,19 @@
+dest:
+  key: 1
+  dict:
+    subkey: 10
+  group_a:
+    a_key: 1
+    a_dict:
+      a_subkey: 10
+  group_b:
+    b_key: 1
+    b_dict:
+      b_subkey: 10
+load_source:
+  experiment_name: test_mlflow_defaults
+  run_name: prev_run_1
+  tracking_uri: sqlite:///mlflow.db  # placeholder
+  registry_uri: null
+  experiment_id: "1"  # placeholder
+  run_id: "00000000000000000000000000000000"  # placeholder
diff --git a/tests/end_to_end/mlflow_defaults_configs/load_nondefault_config.yaml b/tests/end_to_end/mlflow_defaults_configs/load_nondefault_config.yaml
new file mode 100644
index 00000000..82756626
--- /dev/null
+++ b/tests/end_to_end/mlflow_defaults_configs/load_nondefault_config.yaml
@@ -0,0 +1,4 @@
+defaults:
+  - mlflow@load_source: prev_run_1
+mlflow_defaults:
+  - load_source: subdir/special_1#
diff --git a/tests/end_to_end/mlflow_defaults_configs/load_nondefault_config_expected.yaml b/tests/end_to_end/mlflow_defaults_configs/load_nondefault_config_expected.yaml
new file mode 100644
index 00000000..f75a3b3a
--- /dev/null
+++ b/tests/end_to_end/mlflow_defaults_configs/load_nondefault_config_expected.yaml
@@ -0,0 +1,18 @@
+key: 3
+dict:
+  subkey: 30
+group_a:
+  a_key: 3
+  a_dict:
+    a_subkey: 30
+group_b:
+  b_key: 3
+  b_dict:
+    b_subkey: 30
+load_source:
+  experiment_name: test_mlflow_defaults
+  run_name: prev_run_1
+  tracking_uri: sqlite:///mlflow.db  # placeholder
+  registry_uri: null
+  experiment_id: "1"  # placeholder
+  run_id: "00000000000000000000000000000000"  # placeholder
diff --git a/tests/end_to_end/mlflow_defaults_configs/load_subconfig.yaml b/tests/end_to_end/mlflow_defaults_configs/load_subconfig.yaml
new file mode 100644
index 00000000..88e2b479
--- /dev/null
+++ b/tests/end_to_end/mlflow_defaults_configs/load_subconfig.yaml
@@ -0,0 +1,4 @@
+defaults:
+  - mlflow@load_source: prev_run_1
+mlflow_defaults:
+  - load_source: group_a
diff --git a/tests/end_to_end/mlflow_defaults_configs/load_subconfig_expected.yaml b/tests/end_to_end/mlflow_defaults_configs/load_subconfig_expected.yaml
new file mode 100644
index 00000000..877bcb86
--- /dev/null
+++ b/tests/end_to_end/mlflow_defaults_configs/load_subconfig_expected.yaml
@@ -0,0 +1,11 @@
+group_a:
+  a_key: 1
+  a_dict:
+    a_subkey: 10
+load_source:
+  experiment_name: test_mlflow_defaults
+  run_name: prev_run_1
+  tracking_uri: sqlite:///mlflow.db  # placeholder
+  registry_uri: null
+  experiment_id: "1"  # placeholder
+  run_id: "00000000000000000000000000000000"  # placeholder
diff --git a/tests/end_to_end/mlflow_defaults_configs/load_subconfig_select.yaml b/tests/end_to_end/mlflow_defaults_configs/load_subconfig_select.yaml
new file mode 100644
index 00000000..ddfa0b29
--- /dev/null
+++ b/tests/end_to_end/mlflow_defaults_configs/load_subconfig_select.yaml
@@ -0,0 +1,4 @@
+defaults:
+  - mlflow@load_source: prev_run_1
+mlflow_defaults:
+  - load_source@dest: subdir/special_1#group_a.a_key
diff --git a/tests/end_to_end/mlflow_defaults_configs/load_subconfig_select_expected.yaml b/tests/end_to_end/mlflow_defaults_configs/load_subconfig_select_expected.yaml
new file mode 100644
index 00000000..e219585b
--- /dev/null
+++ b/tests/end_to_end/mlflow_defaults_configs/load_subconfig_select_expected.yaml
@@ -0,0 +1,8 @@
+dest: 3
+load_source:
+  experiment_name: test_mlflow_defaults
+  run_name: prev_run_1
+  tracking_uri: sqlite:///mlflow.db  # placeholder
+  registry_uri: null
+  experiment_id: "1"  # placeholder
+  run_id: "00000000000000000000000000000000"  # placeholder
diff --git a/tests/end_to_end/mlflow_defaults_configs/override_flag.yaml b/tests/end_to_end/mlflow_defaults_configs/override_flag.yaml
new file mode 100644
index 00000000..a6d13a5e
--- /dev/null
+++ b/tests/end_to_end/mlflow_defaults_configs/override_flag.yaml
@@ -0,0 +1,7 @@
+defaults:
+  - mlflow@load_source: prev_run_1
+other_section:
+  foo: bar
+mlflow_defaults:
+  - _self_
+  - override load_source@other_section: subdir/special_1#
diff --git a/tests/end_to_end/mlflow_defaults_configs/override_flag_expected.yaml b/tests/end_to_end/mlflow_defaults_configs/override_flag_expected.yaml
new file mode 100644
index 00000000..8d4d28d3
--- /dev/null
+++ b/tests/end_to_end/mlflow_defaults_configs/override_flag_expected.yaml
@@ -0,0 +1,19 @@
+other_section:
+  key: 3
+  dict:
+    subkey: 30
+  group_a:
+    a_key: 3
+    a_dict:
+      a_subkey: 30
+  group_b:
+    b_key: 3
+    b_dict:
+      b_subkey: 30
+load_source:
+  experiment_name: test_mlflow_defaults
+  run_name: prev_run_1
+  tracking_uri: sqlite:///mlflow.db  # placeholder
+  registry_uri: null
+  experiment_id: "1"  # placeholder
+  run_id: "00000000000000000000000000000000"  # placeholder
diff --git a/tests/end_to_end/test_load_config_from_mlflow.py b/tests/end_to_end/test_load_config_from_mlflow.py
index c6743aa3..b3b7088d 100644
--- a/tests/end_to_end/test_load_config_from_mlflow.py
+++ b/tests/end_to_end/test_load_config_from_mlflow.py
@@ -110,7 +110,18 @@ def test_setup(setup_dir: Path, tmp_path: Path) -> None:
 
 @pytest.mark.parametrize(
     "test_case",
-    ["full_default_copy"],
+    [
+        "full_default_copy",
+        "load_default_config_at_package",
+        "load_nondefault_config",
+        "load_subconfig",
+        "load_subconfig_select",
+        "implicit_config_select",
+        "implicit_artifact_select",
+        "override_flag",
+        "composition_order_self_first",
+        "composition_order_self_last",
+    ],
 )
 def test_mlflow_defaults(setup_dir: Path, test_case: str) -> None:
     """Test mlflow defaults."""

From 025a9a67d1a6f56ce523c3d4d2c20ed1d05e744e Mon Sep 17 00:00:00 2001
From: Eric Alt <ericallenalt@gmail.com>
Date: Fri, 16 Jan 2026 12:00:47 -0800
Subject: [PATCH 11/31] Add tests

---
 .../multiple_runs.yaml                        |   6 +
 .../multiple_runs_expected.yaml               |  38 +++++
 .../test_load_config_from_mlflow.py           |  58 ++++---
 tests/run_management/test_mlflow_defaults.py  | 159 ++++++++++++++++++
 4 files changed, 236 insertions(+), 25 deletions(-)
 create mode 100644 tests/end_to_end/mlflow_defaults_configs/multiple_runs.yaml
 create mode 100644 tests/end_to_end/mlflow_defaults_configs/multiple_runs_expected.yaml

diff --git a/tests/end_to_end/mlflow_defaults_configs/multiple_runs.yaml b/tests/end_to_end/mlflow_defaults_configs/multiple_runs.yaml
new file mode 100644
index 00000000..9bb9b463
--- /dev/null
+++ b/tests/end_to_end/mlflow_defaults_configs/multiple_runs.yaml
@@ -0,0 +1,6 @@
+defaults:
+  - mlflow@load_source_1: prev_run_1
+  - mlflow@load_source_2: prev_run_2
+mlflow_defaults:
+  - load_source_1@old_model_1
+  - load_source_2@old_model_2
diff --git a/tests/end_to_end/mlflow_defaults_configs/multiple_runs_expected.yaml b/tests/end_to_end/mlflow_defaults_configs/multiple_runs_expected.yaml
new file mode 100644
index 00000000..65fcb9a0
--- /dev/null
+++ b/tests/end_to_end/mlflow_defaults_configs/multiple_runs_expected.yaml
@@ -0,0 +1,38 @@
+old_model_1:
+  key: 1
+  dict:
+    subkey: 10
+  group_a:
+    a_key: 1
+    a_dict:
+      a_subkey: 10
+  group_b:
+    b_key: 1
+    b_dict:
+      b_subkey: 10
+old_model_2:
+  key: 2
+  dict:
+    subkey: 20
+  group_a:
+    a_key: 2
+    a_dict:
+      a_subkey: 20
+  group_b:
+    b_key: 2
+    b_dict:
+      b_subkey: 20
+load_source_1:
+  experiment_name: test_mlflow_defaults
+  run_name: prev_run_1
+  tracking_uri: sqlite:///mlflow.db  # placeholder
+  registry_uri: null
+  experiment_id: "1"  # placeholder
+  run_id: "00000000000000000000000000000000"  # placeholder
+load_source_2:
+  experiment_name: test_mlflow_defaults
+  run_name: prev_run_2
+  tracking_uri: sqlite:///mlflow.db  # placeholder
+  registry_uri: null
+  experiment_id: "1"  # placeholder
+  run_id: "00000000000000000000000000000000"  # placeholder
diff --git a/tests/end_to_end/test_load_config_from_mlflow.py b/tests/end_to_end/test_load_config_from_mlflow.py
index b3b7088d..8440e553 100644
--- a/tests/end_to_end/test_load_config_from_mlflow.py
+++ b/tests/end_to_end/test_load_config_from_mlflow.py
@@ -109,38 +109,46 @@ def test_setup(setup_dir: Path, tmp_path: Path) -> None:
 
 
 @pytest.mark.parametrize(
-    "test_case",
+    ("test_case", "load_sources"),
     [
-        "full_default_copy",
-        "load_default_config_at_package",
-        "load_nondefault_config",
-        "load_subconfig",
-        "load_subconfig_select",
-        "implicit_config_select",
-        "implicit_artifact_select",
-        "override_flag",
-        "composition_order_self_first",
-        "composition_order_self_last",
+        ("full_default_copy", ["load_source"]),
+        ("load_default_config_at_package", ["load_source"]),
+        ("load_nondefault_config", ["load_source"]),
+        ("load_subconfig", ["load_source"]),
+        ("load_subconfig_select", ["load_source"]),
+        ("implicit_config_select", ["load_source"]),
+        ("implicit_artifact_select", ["load_source"]),
+        ("override_flag", ["load_source"]),
+        ("composition_order_self_first", ["load_source"]),
+        ("composition_order_self_last", ["load_source"]),
+        ("multiple_runs", ["load_source_1", "load_source_2"]),
     ],
 )
-def test_mlflow_defaults(setup_dir: Path, test_case: str) -> None:
+def test_mlflow_defaults(setup_dir: Path, test_case: str, load_sources: list[str]) -> None:
     """Test mlflow defaults."""
     tracking_uri = f"sqlite:///{setup_dir.resolve()}/mlflow.db"
     with initialize_config_dir(config_dir=CONFIG_DIR):
-        cfg = compose(config_name=test_case, overrides=[f"load_source.tracking_uri={tracking_uri}"])
+        # Generate overrides for all load sources
+        overrides = [f"{load_source}.tracking_uri={tracking_uri}" for load_source in load_sources]
+        cfg = compose(config_name=test_case, overrides=overrides)
     actual = load_mlflow_defaults(cfg)
-    experiment_id = OmegaConf.select(actual, "load_source.experiment_id")
-    assert experiment_id is not None
-    run_id = OmegaConf.select(actual, "load_source.run_id")
-    assert run_id is not None
-    with initialize_config_dir(config_dir=CONFIG_DIR):
-        expected = compose(
-            config_name=f"{test_case}_expected",
-            overrides=[
-                f"load_source.tracking_uri={tracking_uri}",
-                f'load_source.experiment_id="{experiment_id}"',
-                f'load_source.run_id="{run_id}"',
-            ],
+
+    # Extract experiment_id and run_id for all load sources
+    expected_overrides = []
+    for load_source in load_sources:
+        experiment_id = OmegaConf.select(actual, f"{load_source}.experiment_id")
+        run_id = OmegaConf.select(actual, f"{load_source}.run_id")
+        assert experiment_id is not None
+        assert run_id is not None
+        expected_overrides.extend(
+            [
+                f"{load_source}.tracking_uri={tracking_uri}",
+                f'{load_source}.experiment_id="{experiment_id}"',
+                f'{load_source}.run_id="{run_id}"',
+            ]
         )
 
+    with initialize_config_dir(config_dir=CONFIG_DIR):
+        expected = compose(config_name=f"{test_case}_expected", overrides=expected_overrides)
+
     assert actual == expected
diff --git a/tests/run_management/test_mlflow_defaults.py b/tests/run_management/test_mlflow_defaults.py
index 5b83911f..698792b1 100644
--- a/tests/run_management/test_mlflow_defaults.py
+++ b/tests/run_management/test_mlflow_defaults.py
@@ -514,3 +514,162 @@ def test_no_mlflow_defaults_key():
     loaded_cfg: DictConfig = load_mlflow_defaults(cfg)
     assert loaded_cfg == cfg
     assert OmegaConf.select(loaded_cfg, "other_section.foo") == "bar"
+
+
+def test_multiple_entries_different_runs(base_cfg: DictConfig, mock_download: MagicMock, tmp_path: Path):
+    """Test loading configs from multiple different runs."""
+    # Setup two different runs with different configs
+    artifact_path_1 = tmp_path / "config_run1.yaml"
+    artifact_path_1.write_text("run1_key: run1_value\nrun1_section:\n  nested: run1_nested\n")
+    artifact_path_2 = tmp_path / "config_run2.yaml"
+    artifact_path_2.write_text("run2_key: run2_value\nrun2_section:\n  nested: run2_nested\n")
+
+    # Mock download_artifacts to return different paths for different run_ids
+    def download_side_effect(run_id, path, dst_path):
+        if run_id == "test_run_id_1":
+            return str(artifact_path_1)
+        elif run_id == "test_run_id_2":
+            return str(artifact_path_2)
+        return str(artifact_path_1)
+
+    mock_download.side_effect = download_side_effect
+
+    cfg = OmegaConf.merge(
+        base_cfg,
+        {
+            "run1": {
+                "tracking_uri": "databricks",
+                "run_id": "test_run_id_1",
+            },
+            "run2": {
+                "tracking_uri": "databricks",
+                "run_id": "test_run_id_2",
+            },
+            "mlflow_defaults": ["run1@model1", "run2@model2"],
+        },
+    )
+
+    loaded_cfg: DictConfig = load_mlflow_defaults(cfg)
+    assert OmegaConf.select(loaded_cfg, "model1.run1_key") == "run1_value"
+    assert OmegaConf.select(loaded_cfg, "model1.run1_section.nested") == "run1_nested"
+    assert OmegaConf.select(loaded_cfg, "model2.run2_key") == "run2_value"
+    assert OmegaConf.select(loaded_cfg, "model2.run2_section.nested") == "run2_nested"
+    assert OmegaConf.select(loaded_cfg, "other_section.foo") == "bar"
+
+
+def test_multiple_entries_same_run(base_cfg: DictConfig, mock_download: MagicMock, tmp_path: Path):
+    """Test loading multiple artifacts from the same run."""
+    # Setup multiple artifacts from the same run
+    artifact_path_1 = tmp_path / "config.yaml"
+    artifact_path_1.write_text("artifact1_key: artifact1_value\n")
+    artifact_path_2 = tmp_path / "other_artifact.yaml"
+    artifact_path_2.write_text("artifact2_key: artifact2_value\n")
+
+    # Mock download_artifacts to return different paths based on artifact path
+    def download_side_effect(run_id, path, dst_path):
+        if path == "config":
+            return str(artifact_path_1)
+        elif path == "other_artifact":
+            return str(artifact_path_2)
+        return str(artifact_path_1)
+
+    mock_download.side_effect = download_side_effect
+
+    cfg = OmegaConf.merge(
+        base_cfg,
+        {
+            "mlflow_defaults": [
+                "previous_run@section1: config",
+                "previous_run@section2: other_artifact",
+            ],
+        },
+    )
+
+    loaded_cfg: DictConfig = load_mlflow_defaults(cfg)
+    assert OmegaConf.select(loaded_cfg, "section1.artifact1_key") == "artifact1_value"
+    assert OmegaConf.select(loaded_cfg, "section2.artifact2_key") == "artifact2_value"
+    assert OmegaConf.select(loaded_cfg, "other_section.foo") == "bar"
+
+
+def test_multiple_entries_shared_keys_last_wins(base_cfg: DictConfig, mock_download: MagicMock, tmp_path: Path):
+    """Test that last entry wins when keys conflict."""
+    # Setup entries that load configs with overlapping keys
+    artifact_path_1 = tmp_path / "config1.yaml"
+    artifact_path_1.write_text("shared_key: first_value\nunique1: value1\n")
+    artifact_path_2 = tmp_path / "config2.yaml"
+    artifact_path_2.write_text("shared_key: second_value\nunique2: value2\n")
+
+    def download_side_effect(run_id, path, dst_path):
+        if path == "config1":
+            return str(artifact_path_1)
+        elif path == "config2":
+            return str(artifact_path_2)
+        return str(artifact_path_1)
+
+    mock_download.side_effect = download_side_effect
+
+    cfg = OmegaConf.merge(
+        base_cfg,
+        {
+            "run1": {
+                "tracking_uri": "databricks",
+                "run_id": "test_run_id_1",
+            },
+            "run2": {
+                "tracking_uri": "databricks",
+                "run_id": "test_run_id_2",
+            },
+            "mlflow_defaults": [
+                "run1: config1",
+                "run2: config2",
+            ],
+        },
+    )
+
+    loaded_cfg: DictConfig = load_mlflow_defaults(cfg)
+    # Last entry should win for shared_key
+    assert OmegaConf.select(loaded_cfg, "shared_key") == "second_value"
+    # Both unique keys should be present
+    assert OmegaConf.select(loaded_cfg, "unique1") == "value1"
+    assert OmegaConf.select(loaded_cfg, "unique2") == "value2"
+    assert OmegaConf.select(loaded_cfg, "other_section.foo") == "bar"
+
+
+def test_option_trailing_hash(base_cfg: DictConfig, mock_download: MagicMock, tmp_path: Path):
+    """Test option with trailing # (artifact path only)."""
+    # Test: "custom#" should load artifact "custom" at root
+    artifact_path = tmp_path / "custom.yaml"
+    artifact_path.write_text("key: value\n")
+    mock_download.return_value = str(artifact_path)
+
+    cfg = OmegaConf.merge(
+        base_cfg,
+        {
+            "mlflow_defaults": ["previous_run: custom#"],
+        },
+    )
+
+    loaded_cfg: DictConfig = load_mlflow_defaults(cfg)
+    assert OmegaConf.select(loaded_cfg, "key") == "value"
+    assert mock_download.call_args.kwargs["path"] == "custom"
+
+
+def test_option_hash_empty_select(base_cfg: DictConfig, mock_download: MagicMock, tmp_path: Path):
+    """Test option with # but empty select path."""
+    # Test: "artifact#" should load artifact "artifact" at root
+    artifact_path = tmp_path / "artifact.yaml"
+    artifact_path.write_text("root_key: root_value\nnested:\n  nested_key: nested_value\n")
+    mock_download.return_value = str(artifact_path)
+
+    cfg = OmegaConf.merge(
+        base_cfg,
+        {
+            "mlflow_defaults": ["previous_run: artifact#"],
+        },
+    )
+
+    loaded_cfg: DictConfig = load_mlflow_defaults(cfg)
+    # Should load entire artifact at root
+    assert OmegaConf.select(loaded_cfg, "root_key") == "root_value"
+    assert OmegaConf.select(loaded_cfg, "nested.nested_key") == "nested_value"
+    assert mock_download.call_args.kwargs["path"] == "artifact"

From 452eed5a40f5fdbc4d7450dd144530296b9005ca Mon Sep 17 00:00:00 2001
From: Eric Alt <ericallenalt@gmail.com>
Date: Fri, 16 Jan 2026 14:08:46 -0800
Subject: [PATCH 12/31] Consolidate and fix tests

---
 .../test_load_mlflow_defaults.py              | 179 ++++-
 tests/run_management/test_mlflow_defaults.py  | 675 ------------------
 2 files changed, 169 insertions(+), 685 deletions(-)
 delete mode 100644 tests/run_management/test_mlflow_defaults.py

diff --git a/tests/run_management/test_load_mlflow_defaults.py b/tests/run_management/test_load_mlflow_defaults.py
index 5b83911f..f6e87998 100644
--- a/tests/run_management/test_load_mlflow_defaults.py
+++ b/tests/run_management/test_load_mlflow_defaults.py
@@ -78,7 +78,7 @@ def test_load_default_config_at_package(base_cfg: DictConfig, mock_download: Mag
     loaded_cfg: DictConfig = load_mlflow_defaults(cfg)
     assert OmegaConf.select(loaded_cfg, "dest.loaded_key") == "loaded_value"
     assert OmegaConf.select(loaded_cfg, "other_section.foo") == "bar"
-    assert mock_download.call_args.kwargs["path"] == "config"
+    assert mock_download.call_args.kwargs["path"] == "config.yaml"
 
 
 def test_load_nondefault_config(base_cfg: DictConfig, mock_download: MagicMock, tmp_path: Path):
@@ -98,25 +98,25 @@ def test_load_nondefault_config(base_cfg: DictConfig, mock_download: MagicMock,
 
     assert loaded_cfg.get("loaded_key") == "loaded_value"
     assert OmegaConf.select(loaded_cfg, "other_section.foo") == "bar"
-    assert mock_download.call_args.kwargs["path"] == "nondefault_config"
+    assert mock_download.call_args.kwargs["path"] == "nondefault_config.yaml"
 
 
 def test_load_subconfig(base_cfg: DictConfig, mock_download: MagicMock, tmp_path: Path):
     """Test loading subconfig."""
     artifact_path = tmp_path / "config.yaml"
-    artifact_path.write_text("sub:\n  nested: value\n")
+    artifact_path.write_text("root:\n  sub:\n    nested: value\n")
     mock_download.return_value = str(artifact_path)
 
     cfg = OmegaConf.merge(
         base_cfg,
         {
-            "mlflow_defaults": ["previous_run: sub"],
+            "mlflow_defaults": ["previous_run: root.sub"],
         },
     )
 
     loaded_cfg: DictConfig = load_mlflow_defaults(cfg)
-    assert loaded_cfg.get("nested") == "value"
-    assert mock_download.call_args.kwargs["path"] == "config"
+    assert OmegaConf.select(loaded_cfg, "sub.nested") == "value"
+    assert mock_download.call_args.kwargs["path"] == "config.yaml"
 
 
 def test_load_subconfig_select(base_cfg: DictConfig, mock_download: MagicMock, tmp_path: Path):
@@ -151,7 +151,7 @@ def test_implicit_config_select(base_cfg: DictConfig, mock_download: MagicMock,
 
     loaded_cfg: DictConfig = load_mlflow_defaults(cfg)
     assert loaded_cfg.get("dest") == "value"
-    mock_download.assert_called_with(run_id="test_run_id", path="config", dst_path=ANY)
+    mock_download.assert_called_with(run_id="test_run_id", path="config.yaml", dst_path=ANY)
 
 
 def test_implicit_artifact_select(base_cfg: DictConfig, mock_download: MagicMock, tmp_path: Path):
@@ -168,7 +168,7 @@ def test_implicit_artifact_select(base_cfg: DictConfig, mock_download: MagicMock
     )
     loaded_cfg: DictConfig = load_mlflow_defaults(cfg)
     assert OmegaConf.select(loaded_cfg, "dest.key") == "value"
-    mock_download.assert_called_with(run_id="test_run_id", path="path/to/artifact", dst_path=ANY)
+    mock_download.assert_called_with(run_id="test_run_id", path="path/to/artifact.yaml", dst_path=ANY)
 
 
 def test_override_flag(base_cfg: DictConfig, mock_download: MagicMock, tmp_path: Path):
@@ -199,7 +199,7 @@ def test_optional_flag_missing_artifact(base_cfg: DictConfig):
     )
     # Should not raise
     loaded_cfg: DictConfig = load_mlflow_defaults(cfg)
-    assert loaded_cfg == cfg
+    assert loaded_cfg == base_cfg
 
 
 def test_null_option_error(base_cfg: DictConfig):
@@ -269,7 +269,7 @@ def test_config_entry_syntax(base_cfg: DictConfig, mock_download: MagicMock, tmp
     loaded_cfg: DictConfig = load_mlflow_defaults(cfg)
     assert OmegaConf.select(loaded_cfg, "key") == "val"
     # Implicitly loads "config" artifact
-    mock_download.assert_called_with(run_id="test_run_id", path="config", dst_path=ANY)
+    mock_download.assert_called_with(run_id="test_run_id", path="config.yaml", dst_path=ANY)
 
 
 def test_missing_target_error(base_cfg: DictConfig):
@@ -514,3 +514,162 @@ def test_no_mlflow_defaults_key():
     loaded_cfg: DictConfig = load_mlflow_defaults(cfg)
     assert loaded_cfg == cfg
     assert OmegaConf.select(loaded_cfg, "other_section.foo") == "bar"
+
+
+def test_multiple_entries_different_runs(base_cfg: DictConfig, mock_download: MagicMock, tmp_path: Path):
+    """Test loading configs from multiple different runs."""
+    # Setup two different runs with different configs
+    artifact_path_1 = tmp_path / "config_run1.yaml"
+    run_1 = DictConfig({"run1_key": "run1_value", "run1_section": {"nested": "run1_nested"}})
+    artifact_path_1.write_text(OmegaConf.to_yaml(run_1))
+    artifact_path_2 = tmp_path / "config_run2.yaml"
+    run_2 = DictConfig({"run2_key": "run2_value", "run2_section": {"nested": "run2_nested"}})
+    artifact_path_2.write_text(OmegaConf.to_yaml(run_2))
+
+    # Mock download_artifacts to return different paths for different run_ids
+    def download_side_effect(run_id, _path, _dst_path):
+        if run_id == "test_run_id_1":
+            return str(artifact_path_1)
+        if run_id == "test_run_id_2":
+            return str(artifact_path_2)
+        return ""
+
+    mock_download.side_effect = download_side_effect
+
+    cfg = OmegaConf.merge(
+        base_cfg,
+        {
+            "run1": {
+                "tracking_uri": "databricks",
+                "run_id": "test_run_id_1",
+            },
+            "run2": {
+                "tracking_uri": "databricks",
+                "run_id": "test_run_id_2",
+            },
+            "mlflow_defaults": ["run1@model1", "run2@model2"],
+        },
+    )
+
+    loaded_cfg: DictConfig = load_mlflow_defaults(cfg)
+    assert OmegaConf.select(loaded_cfg, "model1") == run_1
+    assert OmegaConf.select(loaded_cfg, "model2") == run_2
+    assert OmegaConf.select(loaded_cfg, "other_section.foo") == "bar"
+
+
+def test_multiple_entries_same_run(base_cfg: DictConfig, mock_download: MagicMock, tmp_path: Path):
+    """Test loading multiple artifacts from the same run."""
+    # Setup multiple artifacts from the same run
+    artifact_path_1 = tmp_path / "config.yaml"
+    artifact_path_1.write_text("artifact1_key: artifact1_value\n")
+    artifact_path_2 = tmp_path / "other_artifact.yaml"
+    artifact_path_2.write_text("artifact2_key: artifact2_value\n")
+
+    # Mock download_artifacts to return different paths based on artifact path
+    def download_side_effect(_run_id, path, _dst_path):
+        if path == "config":
+            return str(artifact_path_1)
+        elif path == "other_artifact":
+            return str(artifact_path_2)
+        return str(artifact_path_1)
+
+    mock_download.side_effect = download_side_effect
+
+    cfg = OmegaConf.merge(
+        base_cfg,
+        {
+            "mlflow_defaults": [
+                "previous_run@section1: config",
+                "previous_run@section2: other_artifact",
+            ],
+        },
+    )
+
+    loaded_cfg: DictConfig = load_mlflow_defaults(cfg)
+    assert OmegaConf.select(loaded_cfg, "section1.artifact1_key") == "artifact1_value"
+    assert OmegaConf.select(loaded_cfg, "section2.artifact2_key") == "artifact2_value"
+    assert OmegaConf.select(loaded_cfg, "other_section.foo") == "bar"
+
+
+def test_multiple_entries_shared_keys_last_wins(base_cfg: DictConfig, mock_download: MagicMock, tmp_path: Path):
+    """Test that last entry wins when keys conflict."""
+    # Setup entries that load configs with overlapping keys
+    artifact_path_1 = tmp_path / "config1.yaml"
+    artifact_path_1.write_text("shared_key: first_value\nunique1: value1\n")
+    artifact_path_2 = tmp_path / "config2.yaml"
+    artifact_path_2.write_text("shared_key: second_value\nunique2: value2\n")
+
+    def download_side_effect(_run_id, path, _dst_path):
+        if path == "config1":
+            return str(artifact_path_1)
+        elif path == "config2":
+            return str(artifact_path_2)
+        return str(artifact_path_1)
+
+    mock_download.side_effect = download_side_effect
+
+    cfg = OmegaConf.merge(
+        base_cfg,
+        {
+            "run1": {
+                "tracking_uri": "databricks",
+                "run_id": "test_run_id_1",
+            },
+            "run2": {
+                "tracking_uri": "databricks",
+                "run_id": "test_run_id_2",
+            },
+            "mlflow_defaults": [
+                "run1: config1",
+                "run2: config2",
+            ],
+        },
+    )
+
+    loaded_cfg: DictConfig = load_mlflow_defaults(cfg)
+    # Last entry should win for shared_key
+    assert OmegaConf.select(loaded_cfg, "shared_key") == "second_value"
+    # Both unique keys should be present
+    assert OmegaConf.select(loaded_cfg, "unique1") == "value1"
+    assert OmegaConf.select(loaded_cfg, "unique2") == "value2"
+    assert OmegaConf.select(loaded_cfg, "other_section.foo") == "bar"
+
+
+def test_option_trailing_hash(base_cfg: DictConfig, mock_download: MagicMock, tmp_path: Path):
+    """Test option with trailing # (artifact path only)."""
+    # Test: "custom#" should load artifact "custom" at root
+    artifact_path = tmp_path / "custom.yaml"
+    artifact_path.write_text("key: value\n")
+    mock_download.return_value = str(artifact_path)
+
+    cfg = OmegaConf.merge(
+        base_cfg,
+        {
+            "mlflow_defaults": ["previous_run: custom#"],
+        },
+    )
+
+    loaded_cfg: DictConfig = load_mlflow_defaults(cfg)
+    assert OmegaConf.select(loaded_cfg, "key") == "value"
+    assert mock_download.call_args.kwargs["path"] == "custom"
+
+
+def test_option_hash_empty_select(base_cfg: DictConfig, mock_download: MagicMock, tmp_path: Path):
+    """Test option with # but empty select path."""
+    # Test: "artifact#" should load artifact "artifact" at root
+    artifact_path = tmp_path / "artifact.yaml"
+    artifact_path.write_text("root_key: root_value\nnested:\n  nested_key: nested_value\n")
+    mock_download.return_value = str(artifact_path)
+
+    cfg = OmegaConf.merge(
+        base_cfg,
+        {
+            "mlflow_defaults": ["previous_run: artifact#"],
+        },
+    )
+
+    loaded_cfg: DictConfig = load_mlflow_defaults(cfg)
+    # Should load entire artifact at root
+    assert OmegaConf.select(loaded_cfg, "root_key") == "root_value"
+    assert OmegaConf.select(loaded_cfg, "nested.nested_key") == "nested_value"
+    assert mock_download.call_args.kwargs["path"] == "artifact"
diff --git a/tests/run_management/test_mlflow_defaults.py b/tests/run_management/test_mlflow_defaults.py
deleted file mode 100644
index 698792b1..00000000
--- a/tests/run_management/test_mlflow_defaults.py
+++ /dev/null
@@ -1,675 +0,0 @@
-"""Test MLflow defaults."""
-
-from pathlib import Path
-from unittest.mock import ANY, MagicMock
-
-import pytest
-from mlflow.client import MlflowClient
-from omegaconf import DictConfig, OmegaConf
-
-from simplexity.exceptions import ConfigValidationError
-from simplexity.structured_configs.mlflow_defaults import load_mlflow_defaults
-
-
-@pytest.fixture(autouse=True)
-def _patch_mlflow_helpers(mocker):
-    """Automatically patch common MLflow helper functions for all tests."""
-    mocker.patch("simplexity.structured_configs.mlflow_defaults.validate_mlflow_config")
-    mocker.patch("simplexity.structured_configs.mlflow_defaults.resolve_mlflow_config")
-
-
-@pytest.fixture
-def base_cfg() -> DictConfig:
-    """Set up the base config."""
-    return OmegaConf.create(
-        {
-            "previous_run": {
-                "tracking_uri": "databricks",
-                "run_id": "test_run_id",
-            },
-            "other_section": {"foo": "bar"},
-        }
-    )
-
-
-@pytest.fixture
-def mock_download(mocker) -> MagicMock:
-    """Create a mock MLflow client instance."""
-    mock_client_instance = MagicMock(spec=MlflowClient)
-    mocker.patch(
-        "simplexity.structured_configs.mlflow_defaults.MlflowClient",
-        return_value=mock_client_instance,
-    )
-    return mock_client_instance.download_artifacts
-
-
-def test_load_default_config_at_root(base_cfg: DictConfig, mock_download: MagicMock, tmp_path: Path):
-    """Test loading basic config."""
-    artifact_path = tmp_path / "config.yaml"
-    artifact_path.write_text("loaded_key: loaded_value\n")
-    mock_download.return_value = str(artifact_path)
-
-    cfg = OmegaConf.merge(
-        base_cfg,
-        {
-            "mlflow_defaults": ["previous_run"],
-        },
-    )
-
-    loaded_cfg: DictConfig = load_mlflow_defaults(cfg)
-
-    assert loaded_cfg.get("loaded_key") == "loaded_value"
-    assert OmegaConf.select(loaded_cfg, "other_section.foo") == "bar"
-
-
-def test_load_default_config_at_package(base_cfg: DictConfig, mock_download: MagicMock, tmp_path: Path):
-    """Test loading at package."""
-    artifact_path = tmp_path / "config.yaml"
-    artifact_path.write_text("loaded_key: loaded_value\n")
-    mock_download.return_value = str(artifact_path)
-
-    cfg = OmegaConf.merge(
-        base_cfg,
-        {
-            "mlflow_defaults": ["previous_run@dest"],
-        },
-    )
-
-    loaded_cfg: DictConfig = load_mlflow_defaults(cfg)
-    assert OmegaConf.select(loaded_cfg, "dest.loaded_key") == "loaded_value"
-    assert OmegaConf.select(loaded_cfg, "other_section.foo") == "bar"
-    assert mock_download.call_args.kwargs["path"] == "config"
-
-
-def test_load_nondefault_config(base_cfg: DictConfig, mock_download: MagicMock, tmp_path: Path):
-    """Test loading basic config."""
-    artifact_path = tmp_path / "nondefault_config.yaml"
-    artifact_path.write_text("loaded_key: loaded_value\n")
-    mock_download.return_value = str(artifact_path)
-
-    cfg = OmegaConf.merge(
-        base_cfg,
-        {
-            "mlflow_defaults": ["previous_run: nondefault_config#"],
-        },
-    )
-
-    loaded_cfg: DictConfig = load_mlflow_defaults(cfg)
-
-    assert loaded_cfg.get("loaded_key") == "loaded_value"
-    assert OmegaConf.select(loaded_cfg, "other_section.foo") == "bar"
-    assert mock_download.call_args.kwargs["path"] == "nondefault_config"
-
-
-def test_load_subconfig(base_cfg: DictConfig, mock_download: MagicMock, tmp_path: Path):
-    """Test loading subconfig."""
-    artifact_path = tmp_path / "config.yaml"
-    artifact_path.write_text("sub:\n  nested: value\n")
-    mock_download.return_value = str(artifact_path)
-
-    cfg = OmegaConf.merge(
-        base_cfg,
-        {
-            "mlflow_defaults": ["previous_run: sub"],
-        },
-    )
-
-    loaded_cfg: DictConfig = load_mlflow_defaults(cfg)
-    assert loaded_cfg.get("nested") == "value"
-    assert mock_download.call_args.kwargs["path"] == "config"
-
-
-def test_load_subconfig_select(base_cfg: DictConfig, mock_download: MagicMock, tmp_path: Path):
-    """Test loading subconfig select."""
-    artifact_path = tmp_path / "nondefault_config.yaml"
-    artifact_path.write_text("sub:\n  nested: value\n")
-    mock_download.return_value = str(artifact_path)
-
-    cfg = OmegaConf.merge(
-        base_cfg,
-        {
-            "mlflow_defaults": ["previous_run@dest: nondefault_config#sub.nested"],
-        },
-    )
-
-    loaded_cfg: DictConfig = load_mlflow_defaults(cfg)
-    assert loaded_cfg.get("dest") == "value"
-
-
-def test_implicit_config_select(base_cfg: DictConfig, mock_download: MagicMock, tmp_path: Path):
-    """Test implicit artifact="config" when option has no / or #."""
-    artifact_path = tmp_path / "config.yaml"
-    artifact_path.write_text("sub: value\n")
-    mock_download.return_value = str(artifact_path)
-
-    cfg = OmegaConf.merge(
-        base_cfg,
-        {
-            "mlflow_defaults": ["previous_run@dest: sub"],
-        },
-    )
-
-    loaded_cfg: DictConfig = load_mlflow_defaults(cfg)
-    assert loaded_cfg.get("dest") == "value"
-    mock_download.assert_called_with(run_id="test_run_id", path="config", dst_path=ANY)
-
-
-def test_implicit_artifact_select(base_cfg: DictConfig, mock_download: MagicMock, tmp_path: Path):
-    """Test implicit select=root when option has /."""
-    artifact_path = tmp_path / "artifact.yaml"
-    artifact_path.write_text("key: value\n")
-    mock_download.return_value = str(artifact_path)
-
-    cfg = OmegaConf.merge(
-        base_cfg,
-        {
-            "mlflow_defaults": ["previous_run@dest: path/to/artifact"],
-        },
-    )
-    loaded_cfg: DictConfig = load_mlflow_defaults(cfg)
-    assert OmegaConf.select(loaded_cfg, "dest.key") == "value"
-    mock_download.assert_called_with(run_id="test_run_id", path="path/to/artifact", dst_path=ANY)
-
-
-def test_override_flag(base_cfg: DictConfig, mock_download: MagicMock, tmp_path: Path):
-    """Test override flag."""
-    artifact_path = tmp_path / "nondefault_config.yaml"
-    artifact_path.write_text("foo: new_bar\n")
-    mock_download.return_value = str(artifact_path)
-
-    # Place _self_ first so the loaded config overrides the base config
-    cfg = OmegaConf.merge(
-        base_cfg,
-        {
-            "mlflow_defaults": ["_self_", "override previous_run@other_section: nondefault_config#"],
-        },
-    )
-    loaded_cfg: DictConfig = load_mlflow_defaults(cfg)
-    assert OmegaConf.select(loaded_cfg, "other_section.foo") == "new_bar"
-
-
-def test_optional_flag_missing_artifact(base_cfg: DictConfig):
-    """Test optional flag with missing artifact."""
-
-    cfg = OmegaConf.merge(
-        base_cfg,
-        {
-            "mlflow_defaults": ["optional previous_run: null"],
-        },
-    )
-    # Should not raise
-    loaded_cfg: DictConfig = load_mlflow_defaults(cfg)
-    assert loaded_cfg == cfg
-
-
-def test_null_option_error(base_cfg: DictConfig):
-    """Test null option raises error."""
-    cfg = OmegaConf.merge(
-        base_cfg,
-        {
-            "mlflow_defaults": ["previous_run: null"],
-        },
-    )
-    # null option returns just the target config in new logic unless specific logic handles it differently
-    # Actually user logic for option="null" returns _ParsedEntry with artifact_path=None
-    # And _get_target_config handles artifact_path=None by warning and returning None
-    # So "mlflow: null" should be valid but yield nothing in accumulator if not trapped earlier?
-    # Re-reading logic in Step 245:
-    # _parse_entry returns artifact_path=None if option="null"
-    # _get_target_config checks if parsed_entry.artifact_path is None: logs Warning, returns None.
-    # _process_entry sees loaded_config is None.
-    # It raises ValueError("Target config not found for entry...") if not optional.
-    # So it SHOULD raise ValueError.
-
-    with pytest.raises(ValueError, match="Target config not found for entry"):
-        load_mlflow_defaults(cfg)
-
-
-def test_null_option_optional(base_cfg: DictConfig):
-    """Test null option with optional flag."""
-    cfg = OmegaConf.merge(
-        base_cfg,
-        {
-            "mlflow_defaults": ["optional previous_run: null"],
-        },
-    )
-    loaded_cfg: DictConfig = load_mlflow_defaults(cfg)
-    assert OmegaConf.select(loaded_cfg, "other_section") == cfg.other_section
-
-
-def test_composition_order(base_cfg: DictConfig, mock_download: MagicMock, tmp_path: Path):
-    """Test _self_ placement in composition order."""
-    artifact_path = tmp_path / "nondefault_config.yaml"
-    artifact_path.write_text("foo: mlflow_value\n")
-    mock_download.return_value = str(artifact_path)
-
-    # Case 1: _self_ (default/original) before mlflow -> mlflow overrides
-    cfg1 = OmegaConf.merge(base_cfg, {"mlflow_defaults": ["_self_", "previous_run@other_section: nondefault_config#"]})
-    loaded1: DictConfig = load_mlflow_defaults(cfg1)
-    assert OmegaConf.select(loaded1, "other_section.foo") == "mlflow_value"
-
-    # Case 2: mlflow before _self_ -> original overrides
-    cfg2 = OmegaConf.merge(base_cfg, {"mlflow_defaults": ["previous_run@other_section: nondefault_config#", "_self_"]})
-    loaded2: DictConfig = load_mlflow_defaults(cfg2)
-    assert OmegaConf.select(loaded2, "other_section.foo") == "bar"  # Original value
-
-
-def test_config_entry_syntax(base_cfg: DictConfig, mock_download: MagicMock, tmp_path: Path):
-    """Test CONFIG syntax: TARGET (no option)."""
-    artifact_path = tmp_path / "nondefault_config.yaml"
-    artifact_path.write_text("key: val\n")
-    mock_download.return_value = str(artifact_path)
-
-    cfg = OmegaConf.merge(
-        base_cfg,
-        {
-            "mlflow_defaults": ["previous_run"],
-        },
-    )
-    loaded_cfg: DictConfig = load_mlflow_defaults(cfg)
-    assert OmegaConf.select(loaded_cfg, "key") == "val"
-    # Implicitly loads "config" artifact
-    mock_download.assert_called_with(run_id="test_run_id", path="config", dst_path=ANY)
-
-
-def test_missing_target_error(base_cfg: DictConfig):
-    """Test missing target raises error."""
-    cfg = OmegaConf.merge(
-        base_cfg,
-        {
-            "mlflow_defaults": ["optional : option"],
-        },
-    )
-    with pytest.raises(ValueError, match="Invalid MLflow default entry"):
-        load_mlflow_defaults(cfg)
-
-
-def test_invalid_entry_format(base_cfg: DictConfig):
-    """Test invalid entry format raises error (line 64)."""
-    cfg = OmegaConf.merge(
-        base_cfg,
-        {
-            "mlflow_defaults": ["invalid@format@too@many@ats"],
-        },
-    )
-    with pytest.raises(ValueError, match="Invalid MLflow default entry"):
-        load_mlflow_defaults(cfg)
-
-
-def test_missing_target_node(base_cfg: DictConfig):
-    """Test missing target node in config (lines 96-97)."""
-    cfg = OmegaConf.merge(
-        base_cfg,
-        {
-            "mlflow_defaults": ["nonexistent: config"],
-        },
-    )
-    # Should raise ValueError because target node not found and not optional
-    with pytest.raises(ValueError, match="Target config not found for entry"):
-        load_mlflow_defaults(cfg)
-
-
-def test_missing_target_node_optional(base_cfg: DictConfig):
-    """Test missing target node with optional flag."""
-    cfg = OmegaConf.merge(
-        base_cfg,
-        {
-            "mlflow_defaults": ["optional nonexistent: config"],
-        },
-    )
-    # Should not raise, just return accumulator
-    loaded_cfg: DictConfig = load_mlflow_defaults(cfg)
-    assert OmegaConf.select(loaded_cfg, "other_section") == cfg.other_section
-
-
-def test_validate_mlflow_config_error(base_cfg: DictConfig, mocker):
-    """Test ConfigValidationError from validate_mlflow_config (lines 101-103)."""
-    # Make validate_mlflow_config raise ConfigValidationError
-    mock_validate = mocker.patch("simplexity.structured_configs.mlflow_defaults.validate_mlflow_config")
-    mock_validate.side_effect = ConfigValidationError("Invalid config")
-
-    cfg = OmegaConf.merge(
-        base_cfg,
-        {
-            "mlflow_defaults": ["previous_run: config"],
-        },
-    )
-    # Should raise ValueError because validation failed and not optional
-    with pytest.raises(ValueError, match="Target config not found for entry"):
-        load_mlflow_defaults(cfg)
-
-
-def test_validate_mlflow_config_error_optional(base_cfg: DictConfig, mocker):
-    """Test ConfigValidationError with optional flag."""
-    # Make validate_mlflow_config raise ConfigValidationError
-    mock_validate = mocker.patch("simplexity.structured_configs.mlflow_defaults.validate_mlflow_config")
-    mock_validate.side_effect = ConfigValidationError("Invalid config")
-
-    cfg = OmegaConf.merge(
-        base_cfg,
-        {
-            "mlflow_defaults": ["optional previous_run: config"],
-        },
-    )
-    # Should not raise, just return accumulator
-    loaded_cfg: DictConfig = load_mlflow_defaults(cfg)
-    assert OmegaConf.select(loaded_cfg, "other_section") == cfg.other_section
-
-
-def test_resolve_mlflow_config_error(base_cfg: DictConfig, mocker):
-    """Test ValueError from resolve_mlflow_config (lines 107-109)."""
-    # Make resolve_mlflow_config raise ValueError
-    mock_resolve = mocker.patch("simplexity.structured_configs.mlflow_defaults.resolve_mlflow_config")
-    mock_resolve.side_effect = ValueError("Resolution failed")
-
-    cfg = OmegaConf.merge(
-        base_cfg,
-        {
-            "mlflow_defaults": ["previous_run: config"],
-        },
-    )
-    # Should raise ValueError because resolution failed and not optional
-    with pytest.raises(ValueError, match="Target config not found for entry"):
-        load_mlflow_defaults(cfg)
-
-
-def test_resolve_mlflow_config_error_optional(base_cfg: DictConfig, mocker):
-    """Test ValueError from resolve_mlflow_config with optional flag."""
-    # Make resolve_mlflow_config raise ValueError
-    mock_resolve = mocker.patch("simplexity.structured_configs.mlflow_defaults.resolve_mlflow_config")
-    mock_resolve.side_effect = ValueError("Resolution failed")
-
-    cfg = OmegaConf.merge(
-        base_cfg,
-        {
-            "mlflow_defaults": ["optional mlflow: config"],
-        },
-    )
-    # Should not raise, just return accumulator
-    loaded_cfg: DictConfig = load_mlflow_defaults(cfg)
-    assert OmegaConf.select(loaded_cfg, "other_section") == cfg.other_section
-
-
-def test_omegaconf_load_error(base_cfg: DictConfig, mock_download: MagicMock, tmp_path: Path):
-    """Test Exception from OmegaConf.load (lines 124-126)."""
-    # Create a file with invalid YAML content
-    artifact_path = tmp_path / "invalid.yaml"
-    artifact_path.write_text("invalid: yaml: content: [unclosed")
-    mock_download.return_value = str(artifact_path)
-
-    cfg = OmegaConf.merge(
-        base_cfg,
-        {
-            "mlflow_defaults": ["mlflow: config"],
-        },
-    )
-    # Should raise ValueError because load failed and not optional
-    with pytest.raises(ValueError, match="Target config not found for entry"):
-        load_mlflow_defaults(cfg)
-
-
-def test_omegaconf_load_error_optional(base_cfg: DictConfig, mock_download: MagicMock, tmp_path: Path):
-    """Test Exception from OmegaConf.load with optional flag."""
-    # Create a file with invalid YAML content
-    artifact_path = tmp_path / "invalid.yaml"
-    artifact_path.write_text("invalid: yaml: content: [unclosed")
-    mock_download.return_value = str(artifact_path)
-
-    cfg = OmegaConf.merge(
-        base_cfg,
-        {
-            "mlflow_defaults": ["optional previous_run: config"],
-        },
-    )
-    # Should not raise, just return accumulator
-    loaded_cfg: DictConfig = load_mlflow_defaults(cfg)
-    assert OmegaConf.select(loaded_cfg, "other_section") == cfg.other_section
-
-
-def test_select_path_not_found(base_cfg: DictConfig, mock_download: MagicMock, tmp_path: Path):
-    """Test select path not found in artifact (lines 133-136)."""
-    artifact_path = tmp_path / "nondefault_config.yaml"
-    artifact_path.write_text("key: value\n")
-    mock_download.return_value = str(artifact_path)
-
-    cfg = OmegaConf.merge(
-        base_cfg,
-        {
-            "mlflow_defaults": ["previous_run@dest: nondefault_config#nonexistent.path"],
-        },
-    )
-    # Should raise ValueError because select path not found and not optional
-    with pytest.raises(ValueError, match="Target config not found for entry"):
-        load_mlflow_defaults(cfg)
-
-
-def test_select_path_not_found_optional(base_cfg: DictConfig, mock_download: MagicMock, tmp_path: Path):
-    """Test select path not found with optional flag."""
-    artifact_path = tmp_path / "nondefault_config.yaml"
-    artifact_path.write_text("key: value\n")
-    mock_download.return_value = str(artifact_path)
-
-    cfg = OmegaConf.merge(
-        base_cfg,
-        {
-            "mlflow_defaults": ["optional previous_run@dest: nondefault_config#nonexistent.path"],
-        },
-    )
-    # Should not raise, just return accumulator
-    loaded_cfg: DictConfig = load_mlflow_defaults(cfg)
-    assert OmegaConf.select(loaded_cfg, "other_section") == cfg.other_section
-
-
-def test_non_dictconfig_error(base_cfg: DictConfig, mock_download: MagicMock, tmp_path: Path):
-    """Test non-DictConfig error when not optional (lines 157-159)."""
-    # Create artifact with list (non-dict) root value
-    # Use "nondefault_config#" to avoid select_path, so we get the full ListConfig
-    artifact_path = tmp_path / "nondefault_config.yaml"
-    artifact_path.write_text("- item1\n- item2\n")
-    mock_download.return_value = str(artifact_path)
-
-    cfg = OmegaConf.merge(
-        base_cfg,
-        {
-            "mlflow_defaults": ["previous_run: nondefault_config#"],
-        },
-    )
-    # Should raise ValueError because loaded_config is ListConfig (not DictConfig) and not optional
-    with pytest.raises(ValueError, match="Target config not found for entry"):
-        load_mlflow_defaults(cfg)
-
-
-def test_non_dictconfig_optional(base_cfg: DictConfig, mock_download: MagicMock, tmp_path: Path):
-    """Test non-DictConfig with optional flag."""
-    # Create artifact with list (non-dict) root value
-    # Use "nondefault_config#" to avoid select_path, so we get the full ListConfig
-    artifact_path = tmp_path / "nondefault_config.yaml"
-    artifact_path.write_text("- item1\n- item2\n")
-    mock_download.return_value = str(artifact_path)
-
-    cfg = OmegaConf.merge(
-        base_cfg,
-        {
-            "mlflow_defaults": ["optional previous_run: nondefault_config#"],
-        },
-    )
-    # Should not raise, just return accumulator
-    loaded_cfg: DictConfig = load_mlflow_defaults(cfg)
-    assert OmegaConf.select(loaded_cfg, "other_section") == cfg.other_section
-
-
-def test_no_mlflow_defaults_key():
-    """Test early return when mlflow_defaults key is missing (line 172)."""
-    # Config without mlflow_defaults key
-    cfg = OmegaConf.create(
-        {
-            "previous_run": {
-                "tracking_uri": "databricks",
-                "run_id": "test_run_id",
-            },
-            "other_section": {"foo": "bar"},
-        }
-    )
-    # Should return cfg unchanged
-    loaded_cfg: DictConfig = load_mlflow_defaults(cfg)
-    assert loaded_cfg == cfg
-    assert OmegaConf.select(loaded_cfg, "other_section.foo") == "bar"
-
-
-def test_multiple_entries_different_runs(base_cfg: DictConfig, mock_download: MagicMock, tmp_path: Path):
-    """Test loading configs from multiple different runs."""
-    # Setup two different runs with different configs
-    artifact_path_1 = tmp_path / "config_run1.yaml"
-    artifact_path_1.write_text("run1_key: run1_value\nrun1_section:\n  nested: run1_nested\n")
-    artifact_path_2 = tmp_path / "config_run2.yaml"
-    artifact_path_2.write_text("run2_key: run2_value\nrun2_section:\n  nested: run2_nested\n")
-
-    # Mock download_artifacts to return different paths for different run_ids
-    def download_side_effect(run_id, path, dst_path):
-        if run_id == "test_run_id_1":
-            return str(artifact_path_1)
-        elif run_id == "test_run_id_2":
-            return str(artifact_path_2)
-        return str(artifact_path_1)
-
-    mock_download.side_effect = download_side_effect
-
-    cfg = OmegaConf.merge(
-        base_cfg,
-        {
-            "run1": {
-                "tracking_uri": "databricks",
-                "run_id": "test_run_id_1",
-            },
-            "run2": {
-                "tracking_uri": "databricks",
-                "run_id": "test_run_id_2",
-            },
-            "mlflow_defaults": ["run1@model1", "run2@model2"],
-        },
-    )
-
-    loaded_cfg: DictConfig = load_mlflow_defaults(cfg)
-    assert OmegaConf.select(loaded_cfg, "model1.run1_key") == "run1_value"
-    assert OmegaConf.select(loaded_cfg, "model1.run1_section.nested") == "run1_nested"
-    assert OmegaConf.select(loaded_cfg, "model2.run2_key") == "run2_value"
-    assert OmegaConf.select(loaded_cfg, "model2.run2_section.nested") == "run2_nested"
-    assert OmegaConf.select(loaded_cfg, "other_section.foo") == "bar"
-
-
-def test_multiple_entries_same_run(base_cfg: DictConfig, mock_download: MagicMock, tmp_path: Path):
-    """Test loading multiple artifacts from the same run."""
-    # Setup multiple artifacts from the same run
-    artifact_path_1 = tmp_path / "config.yaml"
-    artifact_path_1.write_text("artifact1_key: artifact1_value\n")
-    artifact_path_2 = tmp_path / "other_artifact.yaml"
-    artifact_path_2.write_text("artifact2_key: artifact2_value\n")
-
-    # Mock download_artifacts to return different paths based on artifact path
-    def download_side_effect(run_id, path, dst_path):
-        if path == "config":
-            return str(artifact_path_1)
-        elif path == "other_artifact":
-            return str(artifact_path_2)
-        return str(artifact_path_1)
-
-    mock_download.side_effect = download_side_effect
-
-    cfg = OmegaConf.merge(
-        base_cfg,
-        {
-            "mlflow_defaults": [
-                "previous_run@section1: config",
-                "previous_run@section2: other_artifact",
-            ],
-        },
-    )
-
-    loaded_cfg: DictConfig = load_mlflow_defaults(cfg)
-    assert OmegaConf.select(loaded_cfg, "section1.artifact1_key") == "artifact1_value"
-    assert OmegaConf.select(loaded_cfg, "section2.artifact2_key") == "artifact2_value"
-    assert OmegaConf.select(loaded_cfg, "other_section.foo") == "bar"
-
-
-def test_multiple_entries_shared_keys_last_wins(base_cfg: DictConfig, mock_download: MagicMock, tmp_path: Path):
-    """Test that last entry wins when keys conflict."""
-    # Setup entries that load configs with overlapping keys
-    artifact_path_1 = tmp_path / "config1.yaml"
-    artifact_path_1.write_text("shared_key: first_value\nunique1: value1\n")
-    artifact_path_2 = tmp_path / "config2.yaml"
-    artifact_path_2.write_text("shared_key: second_value\nunique2: value2\n")
-
-    def download_side_effect(run_id, path, dst_path):
-        if path == "config1":
-            return str(artifact_path_1)
-        elif path == "config2":
-            return str(artifact_path_2)
-        return str(artifact_path_1)
-
-    mock_download.side_effect = download_side_effect
-
-    cfg = OmegaConf.merge(
-        base_cfg,
-        {
-            "run1": {
-                "tracking_uri": "databricks",
-                "run_id": "test_run_id_1",
-            },
-            "run2": {
-                "tracking_uri": "databricks",
-                "run_id": "test_run_id_2",
-            },
-            "mlflow_defaults": [
-                "run1: config1",
-                "run2: config2",
-            ],
-        },
-    )
-
-    loaded_cfg: DictConfig = load_mlflow_defaults(cfg)
-    # Last entry should win for shared_key
-    assert OmegaConf.select(loaded_cfg, "shared_key") == "second_value"
-    # Both unique keys should be present
-    assert OmegaConf.select(loaded_cfg, "unique1") == "value1"
-    assert OmegaConf.select(loaded_cfg, "unique2") == "value2"
-    assert OmegaConf.select(loaded_cfg, "other_section.foo") == "bar"
-
-
-def test_option_trailing_hash(base_cfg: DictConfig, mock_download: MagicMock, tmp_path: Path):
-    """Test option with trailing # (artifact path only)."""
-    # Test: "custom#" should load artifact "custom" at root
-    artifact_path = tmp_path / "custom.yaml"
-    artifact_path.write_text("key: value\n")
-    mock_download.return_value = str(artifact_path)
-
-    cfg = OmegaConf.merge(
-        base_cfg,
-        {
-            "mlflow_defaults": ["previous_run: custom#"],
-        },
-    )
-
-    loaded_cfg: DictConfig = load_mlflow_defaults(cfg)
-    assert OmegaConf.select(loaded_cfg, "key") == "value"
-    assert mock_download.call_args.kwargs["path"] == "custom"
-
-
-def test_option_hash_empty_select(base_cfg: DictConfig, mock_download: MagicMock, tmp_path: Path):
-    """Test option with # but empty select path."""
-    # Test: "artifact#" should load artifact "artifact" at root
-    artifact_path = tmp_path / "artifact.yaml"
-    artifact_path.write_text("root_key: root_value\nnested:\n  nested_key: nested_value\n")
-    mock_download.return_value = str(artifact_path)
-
-    cfg = OmegaConf.merge(
-        base_cfg,
-        {
-            "mlflow_defaults": ["previous_run: artifact#"],
-        },
-    )
-
-    loaded_cfg: DictConfig = load_mlflow_defaults(cfg)
-    # Should load entire artifact at root
-    assert OmegaConf.select(loaded_cfg, "root_key") == "root_value"
-    assert OmegaConf.select(loaded_cfg, "nested.nested_key") == "nested_value"
-    assert mock_download.call_args.kwargs["path"] == "artifact"

From b11dc2542fe23556a516c0a987cd77b23c90000a Mon Sep 17 00:00:00 2001
From: Eric Alt <ericallenalt@gmail.com>
Date: Fri, 16 Jan 2026 14:17:41 -0800
Subject: [PATCH 13/31] Fix tests

---
 .../test_load_mlflow_defaults.py              | 28 ++++++++++---------
 1 file changed, 15 insertions(+), 13 deletions(-)

diff --git a/tests/run_management/test_load_mlflow_defaults.py b/tests/run_management/test_load_mlflow_defaults.py
index f6e87998..653f51c6 100644
--- a/tests/run_management/test_load_mlflow_defaults.py
+++ b/tests/run_management/test_load_mlflow_defaults.py
@@ -527,7 +527,7 @@ def test_multiple_entries_different_runs(base_cfg: DictConfig, mock_download: Ma
     artifact_path_2.write_text(OmegaConf.to_yaml(run_2))
 
     # Mock download_artifacts to return different paths for different run_ids
-    def download_side_effect(run_id, _path, _dst_path):
+    def download_side_effect(run_id=None, _path=None, _dst_path=None, **_kwargs):
         if run_id == "test_run_id_1":
             return str(artifact_path_1)
         if run_id == "test_run_id_2":
@@ -566,10 +566,11 @@ def test_multiple_entries_same_run(base_cfg: DictConfig, mock_download: MagicMoc
     artifact_path_2.write_text("artifact2_key: artifact2_value\n")
 
     # Mock download_artifacts to return different paths based on artifact path
-    def download_side_effect(_run_id, path, _dst_path):
-        if path == "config":
+    def download_side_effect(_run_id=None, path=None, _dst_path=None, **_kwargs):
+        # Path might be "config" or "config.yaml" depending on parsing
+        if path in ("config", "config.yaml"):
             return str(artifact_path_1)
-        elif path == "other_artifact":
+        elif path in ("other_artifact", "other_artifact.yaml"):
             return str(artifact_path_2)
         return str(artifact_path_1)
 
@@ -579,8 +580,8 @@ def download_side_effect(_run_id, path, _dst_path):
         base_cfg,
         {
             "mlflow_defaults": [
-                "previous_run@section1: config",
-                "previous_run@section2: other_artifact",
+                "previous_run@section1: config#",
+                "previous_run@section2: other_artifact#",
             ],
         },
     )
@@ -599,10 +600,11 @@ def test_multiple_entries_shared_keys_last_wins(base_cfg: DictConfig, mock_downl
     artifact_path_2 = tmp_path / "config2.yaml"
     artifact_path_2.write_text("shared_key: second_value\nunique2: value2\n")
 
-    def download_side_effect(_run_id, path, _dst_path):
-        if path == "config1":
+    def download_side_effect(_run_id=None, path=None, _dst_path=None, **_kwargs):
+        # Path might be "config1" or "config1.yaml" depending on parsing
+        if path in ("config1", "config1.yaml"):
             return str(artifact_path_1)
-        elif path == "config2":
+        elif path in ("config2", "config2.yaml"):
             return str(artifact_path_2)
         return str(artifact_path_1)
 
@@ -620,8 +622,8 @@ def download_side_effect(_run_id, path, _dst_path):
                 "run_id": "test_run_id_2",
             },
             "mlflow_defaults": [
-                "run1: config1",
-                "run2: config2",
+                "run1: config1#",
+                "run2: config2#",
             ],
         },
     )
@@ -651,7 +653,7 @@ def test_option_trailing_hash(base_cfg: DictConfig, mock_download: MagicMock, tm
 
     loaded_cfg: DictConfig = load_mlflow_defaults(cfg)
     assert OmegaConf.select(loaded_cfg, "key") == "value"
-    assert mock_download.call_args.kwargs["path"] == "custom"
+    assert mock_download.call_args.kwargs["path"] == "custom.yaml"
 
 
 def test_option_hash_empty_select(base_cfg: DictConfig, mock_download: MagicMock, tmp_path: Path):
@@ -672,4 +674,4 @@ def test_option_hash_empty_select(base_cfg: DictConfig, mock_download: MagicMock
     # Should load entire artifact at root
     assert OmegaConf.select(loaded_cfg, "root_key") == "root_value"
     assert OmegaConf.select(loaded_cfg, "nested.nested_key") == "nested_value"
-    assert mock_download.call_args.kwargs["path"] == "artifact"
+    assert mock_download.call_args.kwargs["path"] == "artifact.yaml"

From 00026e0944689fd23f1d4b552531be161c6278ac Mon Sep 17 00:00:00 2001
From: Eric Alt <ericallenalt@gmail.com>
Date: Fri, 16 Jan 2026 15:21:44 -0800
Subject: [PATCH 14/31] Fix some stuff

---
 LOAD_SUBCONFIGS.md                            | 355 ++++++++++++++++++
 .../structured_configs/mlflow_defaults.py     |  57 ++-
 ...composition_order_self_first_expected.yaml |   1 +
 .../composition_order_self_last_expected.yaml |  11 +
 .../override_flag_expected.yaml               |   1 +
 .../test_load_mlflow_defaults.py              |   4 +-
 6 files changed, 424 insertions(+), 5 deletions(-)
 create mode 100644 LOAD_SUBCONFIGS.md

diff --git a/LOAD_SUBCONFIGS.md b/LOAD_SUBCONFIGS.md
new file mode 100644
index 00000000..03b2af68
--- /dev/null
+++ b/LOAD_SUBCONFIGS.md
@@ -0,0 +1,355 @@
+# Loading Subconfigs from MLflow runs
+
+## Goal:
+
+The goal is to be able to compose a config taking parts/entire configs that have been logged as MLflow artifacts from previous runs
+
+## Design:
+
+The proposed design is patterned after the Hydra [defaults list](https://hydra.cc/docs/advanced/defaults_list/):
+
+```
+mlflow_defaults:
+  (- MLFLOW_ENTRY)*
+
+MLFLOW_ENTRY      : CONFIG | OPTION_CONFIG | _self_
+CONFIG            : TARGET(@PACKAGE)?
+OPTION_CONFIG     : [optional|override]? TARGET(@PACKAGE)?: OPTION
+OPTION            : ARTIFACT_PATH | SELECT_PATH | ARTIFACT_PATH#SELECT_PATH | null
+```
+
+- `TARGET`: OmegaConf dot path within current config pointing to the [MLFlowConfig](https://github.com/Astera-org/simplexity/blob/657aff777cd7d7a6edfadd1f846bbf62028caeec/simplexity/structured_configs/mlflow.py#L21) node
+- `PACKAGE`: OmegaConf dot path where to place merged content (default: ".", i.e., root)
+- `ARTIFACT_PATH`: path within the MLflow run’s artifact dir to the source YAML without system extension (default: "config").
+  - If the path contains `/`, it is treated as an `ARTIFACT_PATH`.
+  - If the path does not contain `/`, it is ambiguous and defaults to `SELECT_PATH`. To force it to be an `ARTIFACT_PATH`, append `#` (e.g., `my_artifact#`).
+- `SELECT_PATH`: OmegaConf dot path within the source YAML to the subconfig to import (default: root).
+- `_self_`, `optional`, and `override` function the same as their Hydra equivalents.
+  - `_self_` determines the composition order. If missing, it is appended to the end.
+  - Standard merge semantics apply: **whatever comes last wins**, whether it's `_self_` or an MLflow entry.
+    - If `_self_` is first: MLflow entries merged after `_self_` will override the original config.
+    - If `_self_` is last: Original config merged after MLflow entries will override MLflow content.
+  - `optional` suppresses errors if the artifact or selected subconfig is missing.
+  - `override` is a semantic indicator for intention but functionally both override and standard merge result in overwrites.
+
+### Merge Semantics for Overlapping Keys
+
+When multiple entries in `mlflow_defaults` (or `_self_`) have overlapping keys, the merge behavior follows Hydra's deep merge semantics:
+
+1. **Conflicting keys (same key path)**: The last entry wins. The value from the later entry completely replaces the value from earlier entries.
+
+   ```yaml
+   # Entry 1: other_section.foo = "first"
+   # Entry 2: other_section.foo = "second"
+   # Result: other_section.foo = "second" (Entry 2 wins)
+   ```
+
+2. **Non-conflicting keys at the same level**: Both keys are preserved. Nested dictionaries are merged deeply.
+
+   ```yaml
+   # Entry 1: other_section.key = 3, other_section.dict.subkey = 30
+   # Entry 2: other_section.foo = "bar"
+   # Result: other_section.key = 3, other_section.dict.subkey = 30, other_section.foo = "bar"
+   # (All keys preserved, nested dicts merged)
+   ```
+
+3. **Partially overlapping nested structures**: Deep merge preserves non-conflicting nested keys.
+   ```yaml
+   # Entry 1: other_section.dict.key1 = "value1", other_section.dict.key2 = "value2"
+   # Entry 2: other_section.dict.key2 = "new_value2", other_section.dict.key3 = "value3"
+   # Result:
+   #   other_section.dict.key1 = "value1" (preserved from Entry 1)
+   #   other_section.dict.key2 = "new_value2" (Entry 2 wins for conflicting key)
+   #   other_section.dict.key3 = "value3" (added from Entry 2)
+   ```
+
+This matches Hydra's defaults list behavior: entries are merged in order using `OmegaConf.merge()`, which performs deep merging of nested dictionaries.
+
+### OPTION Parsing Logic
+
+To resolve ambiguity in `OPTION`:
+
+1. If `OPTION` contains `#`:
+   - Split into `ARTIFACT_PART` and `SELECT_PART`.
+   - `ARTIFACT_PATH` = `ARTIFACT_PART` (if empty, defaults to "config").
+   - `SELECT_PATH` = `SELECT_PART` (if empty, defaults to root).
+2. Else (no `#`):
+   - If `OPTION` contains `/`:
+     - `ARTIFACT_PATH` = `OPTION`.
+     - `SELECT_PATH` = root.
+   - Else:
+     - `ARTIFACT_PATH` = "config".
+     - `SELECT_PATH` = `OPTION`.
+
+## Examples:
+
+### load entire `config.yaml` and merge at root
+
+```yaml
+defaults:
+  - mlflow@load_source: previous_run
+mlflow_defaults:
+  - load_source
+```
+
+### load `configs/model.yaml` and merge as old_model
+
+```yaml
+defaults:
+  - mlflow@load_source: previous_run
+mlflow_defaults:
+  # 'configs/model' contains '/', so it is treated as ARTIFACT_PATH
+  - load_source@old_model: configs/model
+```
+
+### load `persistence` subconfig from `config.yaml` if it exists
+
+```yaml
+defaults:
+  - mlflow@load_source: previous_run
+mlflow_defaults:
+  # 'persistence' has no '/' or '#', so it is treated as SELECT_PATH from default artifact 'config'
+  - optional load_source@persistence: persistence
+```
+
+### override `train.optimizer` with subconfig from` train.yaml`
+
+```yaml
+defaults:
+  - mlflow@load_source: previous_run
+  - train: smoke
+mlflow_defaults:
+  # usage of '#' explicit defines ARTIFACT_PATH (train) and SELECT_PATH (optimizer)
+  - override load_source@train.optimizer: train#optimizer
+```
+
+### load `model` and `generative_process` from previous run
+
+```yaml
+defaults:
+  - mlflow@load_source: previous_run
+mlflow_defaults:
+  - load_source@model: model
+  - load_source@generative_process: generative_process
+```
+
+### load models from multiple runs
+
+```yaml
+defaults:
+  - model@current_model: transformer
+  - mlflow@load_source_1: previous_run_1
+  - mlflow@load_source_2: previous_run_2
+mlflow_defaults:
+  - load_source_1@old_model_1: model
+  - load_source_2@old_model_2: model
+```
+
+### Load a top-level artifact `custom.yaml` at root
+
+```yaml
+defaults:
+  - mlflow@load_source: previous_run
+mlflow_defaults:
+  # Use '#' to indicate 'custom' is the artifact, and select root
+  - load_source@custom_section: custom#
+```
+
+### Merge behavior with overlapping keys
+
+Deep merge semantics apply consistently everywhere, matching Hydra's defaults list behavior:
+
+1. **Loading at root (no package specified)**: Deep merge - multiple entries merge together, preserving non-conflicting keys.
+
+2. **Loading at a specific package path (e.g., `@other_section`)**: Deep merge - content merges with any existing content at that package path, preserving non-conflicting keys.
+
+   ```yaml
+   defaults:
+     - mlflow@load_source: previous_run
+   other_section:
+     foo: bar
+   mlflow_defaults:
+     - _self_
+     - load_source@other_section: config#
+   # Result: other_section contains both foo: bar (from _self_)
+   # and content from config (merged together)
+   ```
+
+3. **Multiple entries at the same package path**: Deep merge preserves non-conflicting keys:
+   ```yaml
+   mlflow_defaults:
+     - load_source_1@other_section: config1# # Loads at other_section
+     - load_source_2@other_section: config2# # Merges with config1's other_section
+   ```
+
+**Example of multiple entries with overlapping keys:**
+
+```yaml
+defaults:
+  - mlflow@load_source_1: previous_run_1
+  - mlflow@load_source_2: previous_run_2
+mlflow_defaults:
+  # Load config from run 1
+  - load_source_1@other_section: config1#
+  # Load config from run 2 (will merge with run 1's other_section)
+  - load_source_2@other_section: config2#
+```
+
+If `config1.yaml` contains:
+
+```yaml
+other_section:
+  key: 3
+  dict:
+    subkey: 30
+```
+
+And `config2.yaml` contains:
+
+```yaml
+other_section:
+  foo: bar
+  dict:
+    other_subkey: 40
+```
+
+The final result will be:
+
+```yaml
+other_section:
+  key: 3 # Preserved from config1 (non-conflicting)
+  foo: bar # Added from config2 (non-conflicting)
+  dict:
+    subkey: 30 # Preserved from config1 (non-conflicting nested key)
+    other_subkey: 40 # Added from config2 (non-conflicting nested key)
+```
+
+If both configs had `other_section.foo`, the value from `config2` would win (last entry wins for conflicting keys).
+
+**Note**: The `override` flag is currently a semantic indicator and doesn't change merge behavior. In the future, it may be used to specify replacement instead of merge.
+
+## Implementation:
+
+The proposed implementation is as a stand alone function that dynamically resolves the config:
+
+```python
+from mlflow.client import MlflowClient
+from omegaconf import DictConfig, OmegaConf
+
+from simplexity.utils.config_utils import dynamic_resolve
+
+def load_mlflow_defaults(cfg: DictConfig) -> DictConfig
+    # 1. Parse 'mlflow_defaults' list.
+    #    - If '_self_' is missing, append it to the end.
+
+    # 2. Iterate through items:
+    #    - If item is '_self_':
+    #         Merge original cfg (passed in argument) into the accumulator.
+    #    - Else (MLFLOW_ENTRY):
+    #         Parse TARGET, PACKAGE, OPTION.
+    #         If OPTION is "null":
+    #             If "optional" flag is set: continue (ignore)
+    #             Else: throw Error ("Config is mandatory but OPTION is null")
+    #
+    #         Resolve ARTIFACT_PATH and SELECT_PATH from OPTION.
+    #
+    #         Resolve 'tracking_uri', 'run_id' from cfg[TARGET].
+    #         Instantiate MlflowClient.
+    #
+    #         Download artifact (with caching/tempfile).
+    #         Load artifact as DictConfig.
+    #
+    #         Select subconfig if SELECT_PATH is set.
+    #
+    #         Merge into accumulator at PACKAGE.
+
+    # 3. Return accumulator.
+    ...
+```
+
+This could then be included in the `managed_run` decorator
+
+```python
+...
+cfg = get_config(args, kwargs)
+cfg = load_mlflow_defaults(cfg) # <- load subconfigs here
+validate_base_config(cfg)
+resolve_base_config(cfg, strict=strict)
+...
+```
+
+Or used as by itself, such as in a notebook:
+
+```python
+# %%
+cfg = DictConfig({
+    "load_source": DictConfig({
+        "experiment_id": "9828318895773678"
+        "run_id": "93c47bf390aef1273573b9dd53de2d3a"
+        "tracking_uri": "databricks"
+    }),
+    "mlflow_defaults": "load_source",
+})
+
+cfg = load_mlflow_defaults(cfg)
+```
+
+## Testing
+
+### Test Cases
+
+**Entries**
+
+- Single `ENTRY`
+- Multiple `ENTRY` items for different runs
+- Multiple `ENTRY` items for the same run
+- Multiple `ENTRY` items with shared keys but differing values
+  - Resulting value comes from last `ENTRY` in the list with that key
+- Explicit `_self_` omitted
+  - included implicitly at the end
+- Explicit `_self_` before other `ENTRY` item(s)
+
+**Packages**
+
+- Explicit `PACKAGE` omitted
+  - Load at config root by default
+- Load at specified `PACKAGE` path
+
+**Artifact Paths**
+
+- Explicit `ARTIFACT_PATH` omitted
+  - `config.yaml` loaded by default
+- Load config specified by `ARTIFACT_PATH`
+
+**Select Paths**
+
+- Explicit `SELECT_PATH` omitted
+  - Entire `TARGET` config loaded by default
+- Load subconfig specified by `SELECT_PATH`
+- Load a single key specified by `SELECT_PATH`
+
+**Options**
+
+- Explicit `optional` omitted
+  - Throws exception if `OPTION` is `null`
+  - Throws exception if there is an issue loading given `OPTION`
+- Explicit `optional` specified
+  - Accepts `null` value for `OPTION` effectively skipping that `ENTRY`
+  - Loads valid `OPTION` normally
+  - Skips that `ENTRY` if there is an issue loading given `OPTION`
+- Explicit `override` omitted
+  - No effect
+- Explicit `override` specified
+  - No effect
+
+### Unit Tests
+
+- There should be individual test functions for each of the [Test Cases]
+- Each case should be tested in relative isolation to the extent that that is feasible (use simplest defaults for components not under test)
+- Calls to `MlflowClient.download_artifacts` should be mocked with the side effect of saving a yaml file with content needed by the test in a temp directory to avoid this dependency not under test
+
+### Integration Tests
+
+- Should aim to broad coverage of the possibilites in [Test Cases]
+- Should test realistic workflows (such as corresponding to the structures of [Examples]) with individual tests typically covering multiple non-default elements of [Test Cases]
+- Should avoid any mocking (use local MLflow with `sqlite://` in a temp directory as the `tracking_uri`)
diff --git a/simplexity/structured_configs/mlflow_defaults.py b/simplexity/structured_configs/mlflow_defaults.py
index 3c1725b5..2580fb96 100644
--- a/simplexity/structured_configs/mlflow_defaults.py
+++ b/simplexity/structured_configs/mlflow_defaults.py
@@ -76,9 +76,19 @@ def _parse_entry(item: str) -> _ParsedEntry:
             "'TARGET(@PACKAGE)? | [optional|override]? TARGET(@PACKAGE)?: OPTION | _self_'"
         )
 
-    package = groups.get("package") or "."
+    package = groups.get("package")
     option = groups.get("option")
 
+    # Special case: TARGET: VALUE (no @) where VALUE doesn't contain / or #
+    # Treat VALUE as both PACKAGE and SELECT_PATH
+    if package is None and option is not None and "/" not in option and "#" not in option:
+        package = option.rsplit(".", 1)[1] if "." in option else option
+        artifact_path = "config"
+        select_path = option
+        return _ParsedEntry(optional, target, package, artifact_path, select_path)
+
+    package = package or "."
+
     if option == "null":
         return _ParsedEntry(optional, target, package, None, None)
 
@@ -156,17 +166,53 @@ def _normalize_item(item: str | DictConfig) -> str:
             value = item[key]
             # If value is "config" (default artifact), treat as simple CONFIG entry
             if value == "config":
-                return key
+                return str(key)
             return f"{key}: {value}"
         # Multiple keys - convert entire dict to string representation
         return str(item)
     return str(item)
 
 
+def _resolve_load_source(accumulator: DictConfig) -> None:
+    """Resolve load_source if it exists in the accumulator.
+
+    This is needed when _self_ is processed and load_source is merged unresolved.
+    If load_source is not a valid MLflow config or can't be resolved, it is silently skipped.
+
+    Args:
+        accumulator: The accumulator DictConfig that may contain load_source.
+    """
+    load_source = accumulator.get("load_source")
+    if load_source is not None and isinstance(load_source, DictConfig):
+        try:
+            validate_mlflow_config(load_source)
+            resolve_mlflow_config(load_source)
+        except (ConfigValidationError, ValueError):
+            # If load_source is not a valid MLflow config or can't be resolved, skip
+            pass
+
+
 def _process_entry(cfg: DictConfig, accumulator: DictConfig, item: str) -> DictConfig:
     """Process a single MLflow default item."""
     if item == "_self_":
-        return cast(DictConfig, OmegaConf.merge(accumulator, cfg))
+        # Standard merge semantics: last entry wins (matches Hydra behavior)
+        # OmegaConf.merge does deep merge (nested dicts get merged), which matches
+        # how Hydra composes configs with defaults list
+        result = cast(DictConfig, OmegaConf.merge(accumulator, cfg))
+        # Resolve load_source if it exists, as it may be merged unresolved
+        _resolve_load_source(result)
+        return result
+        # # When merging _self_, we need to allow cfg to override accumulator
+        # # Use merge but ensure top-level keys from cfg replace those in accumulator
+        # result = OmegaConf.create()
+        # # Copy all keys from accumulator first
+        # for key in accumulator.keys():
+        #     result[key] = accumulator[key]
+        # # Then copy keys from cfg, which will override accumulator keys
+        # for key in cfg.keys():
+        #     if key != "mlflow_defaults":  # Skip mlflow_defaults as it's metadata
+        #         result[key] = cfg[key]
+        # return cast(DictConfig, result)
 
     parsed_entry = _parse_entry(item)
 
@@ -184,6 +230,9 @@ def _process_entry(cfg: DictConfig, accumulator: DictConfig, item: str) -> DictC
             raise ValueError(f"Target config not found for entry: {item}")
         return cast(DictConfig, OmegaConf.merge(accumulator, loaded_config))
 
+    # When merging at a package (not root), use deep merge semantics (matches Hydra behavior)
+    # This ensures that MLflow content merges with any existing content at that package path,
+    # preserving non-conflicting keys and deeply merging nested dictionaries
     package_conf = OmegaConf.create()
     OmegaConf.update(package_conf, parsed_entry.package, loaded_config)
     return cast(DictConfig, OmegaConf.merge(accumulator, package_conf))
@@ -205,5 +254,7 @@ def load_mlflow_defaults(cfg: DictConfig) -> DictConfig:
         normalized_item = _normalize_item(item)
         accumulator = _process_entry(cfg, accumulator, normalized_item)
 
+    # if "mlflow_defaults" in accumulator:
     del accumulator["mlflow_defaults"]
+
     return accumulator
diff --git a/tests/end_to_end/mlflow_defaults_configs/composition_order_self_first_expected.yaml b/tests/end_to_end/mlflow_defaults_configs/composition_order_self_first_expected.yaml
index 8d4d28d3..e07228c3 100644
--- a/tests/end_to_end/mlflow_defaults_configs/composition_order_self_first_expected.yaml
+++ b/tests/end_to_end/mlflow_defaults_configs/composition_order_self_first_expected.yaml
@@ -1,4 +1,5 @@
 other_section:
+  foo: bar
   key: 3
   dict:
     subkey: 30
diff --git a/tests/end_to_end/mlflow_defaults_configs/composition_order_self_last_expected.yaml b/tests/end_to_end/mlflow_defaults_configs/composition_order_self_last_expected.yaml
index 08545844..6127550f 100644
--- a/tests/end_to_end/mlflow_defaults_configs/composition_order_self_last_expected.yaml
+++ b/tests/end_to_end/mlflow_defaults_configs/composition_order_self_last_expected.yaml
@@ -1,4 +1,15 @@
 other_section:
+  key: 3
+  dict:
+    subkey: 30
+  group_a:
+    a_key: 3
+    a_dict:
+      a_subkey: 30
+  group_b:
+    b_key: 3
+    b_dict:
+      b_subkey: 30
   foo: bar
 load_source:
   experiment_name: test_mlflow_defaults
diff --git a/tests/end_to_end/mlflow_defaults_configs/override_flag_expected.yaml b/tests/end_to_end/mlflow_defaults_configs/override_flag_expected.yaml
index 8d4d28d3..e07228c3 100644
--- a/tests/end_to_end/mlflow_defaults_configs/override_flag_expected.yaml
+++ b/tests/end_to_end/mlflow_defaults_configs/override_flag_expected.yaml
@@ -1,4 +1,5 @@
 other_section:
+  foo: bar
   key: 3
   dict:
     subkey: 30
diff --git a/tests/run_management/test_load_mlflow_defaults.py b/tests/run_management/test_load_mlflow_defaults.py
index 653f51c6..796b3dc8 100644
--- a/tests/run_management/test_load_mlflow_defaults.py
+++ b/tests/run_management/test_load_mlflow_defaults.py
@@ -248,10 +248,10 @@ def test_composition_order(base_cfg: DictConfig, mock_download: MagicMock, tmp_p
     loaded1: DictConfig = load_mlflow_defaults(cfg1)
     assert OmegaConf.select(loaded1, "other_section.foo") == "mlflow_value"
 
-    # Case 2: mlflow before _self_ -> original overrides
+    # Case 2: mlflow before _self_ -> original overrides (last entry wins)
     cfg2 = OmegaConf.merge(base_cfg, {"mlflow_defaults": ["previous_run@other_section: nondefault_config#", "_self_"]})
     loaded2: DictConfig = load_mlflow_defaults(cfg2)
-    assert OmegaConf.select(loaded2, "other_section.foo") == "bar"  # Original value
+    assert OmegaConf.select(loaded2, "other_section.foo") == "bar"  # Original value wins (last entry)
 
 
 def test_config_entry_syntax(base_cfg: DictConfig, mock_download: MagicMock, tmp_path: Path):

From 727983e1df919d8994819880ae378a7ec9f78a3d Mon Sep 17 00:00:00 2001
From: Eric Alt <ericallenalt@gmail.com>
Date: Fri, 16 Jan 2026 15:42:59 -0800
Subject: [PATCH 15/31] Update comments/docstrings

---
 .../structured_configs/mlflow_defaults.py     | 112 +++++++++++++-----
 1 file changed, 85 insertions(+), 27 deletions(-)

diff --git a/simplexity/structured_configs/mlflow_defaults.py b/simplexity/structured_configs/mlflow_defaults.py
index 2580fb96..414ac310 100644
--- a/simplexity/structured_configs/mlflow_defaults.py
+++ b/simplexity/structured_configs/mlflow_defaults.py
@@ -1,4 +1,11 @@
-"""Config utilities."""
+"""Load and merge configs from MLflow runs.
+
+This module implements the mlflow_defaults functionality, allowing configs to be
+composed from parts or entire configs that have been logged as MLflow artifacts
+from previous runs. The design is patterned after Hydra's defaults list.
+
+See LOAD_SUBCONFIGS.md for the full specification.
+"""
 
 # pylint: disable=all
 # Temporarily disable all pylint checkers during AST traversal to prevent crash.
@@ -97,6 +104,16 @@ def _parse_entry(item: str) -> _ParsedEntry:
 
 
 def _get_target_config(cfg: DictConfig, parsed_entry: _ParsedEntry) -> Any | None:
+    """Download and load config from MLflow run.
+
+    Args:
+        cfg: The current config containing MLflow config nodes.
+        parsed_entry: The parsed entry specifying target, package, and paths.
+
+    Returns:
+        The loaded config (or selected subconfig), or None if loading failed
+        and the entry is optional.
+    """
     if parsed_entry.artifact_path is None:
         SIMPLEXITY_LOGGER.warning("Config is mandatory but OPTION is null for entry: %s", parsed_entry.target)
         return None
@@ -173,46 +190,73 @@ def _normalize_item(item: str | DictConfig) -> str:
     return str(item)
 
 
-def _resolve_load_source(accumulator: DictConfig) -> None:
-    """Resolve load_source if it exists in the accumulator.
+def _resolve_mlflow_configs_recursive(cfg: DictConfig) -> None:
+    """Recursively find and resolve all MLflow configs in the config.
 
-    This is needed when _self_ is processed and load_source is merged unresolved.
-    If load_source is not a valid MLflow config or can't be resolved, it is silently skipped.
+    This is needed when _self_ is processed and MLflow configs are merged unresolved.
+    MLflow configs that are not valid or can't be resolved are silently skipped.
 
     Args:
-        accumulator: The accumulator DictConfig that may contain load_source.
+        cfg: The DictConfig to search for MLflow configs.
     """
-    load_source = accumulator.get("load_source")
-    if load_source is not None and isinstance(load_source, DictConfig):
+    if not isinstance(cfg, DictConfig):
+        return
+
+    # Only try to validate/resolve if this looks like an MLflow config
+    # Check for key identifying fields (not just optional fields like registry_uri)
+    mlflow_identifying_fields = {"experiment_id", "experiment_name", "run_id", "run_name"}
+    has_mlflow_fields = any(field in cfg for field in mlflow_identifying_fields)
+
+    if has_mlflow_fields:
+        # Try to validate and resolve the current node as an MLflow config
         try:
-            validate_mlflow_config(load_source)
-            resolve_mlflow_config(load_source)
+            validate_mlflow_config(cfg)
+            # If validation passes, this looks like an MLflow config, try to resolve it
+            resolve_mlflow_config(cfg)
         except (ConfigValidationError, ValueError):
-            # If load_source is not a valid MLflow config or can't be resolved, skip
+            # Not an MLflow config or can't be resolved, continue to check nested configs
             pass
 
+    # Recursively process all nested DictConfig values
+    for key in cfg:
+        try:
+            value = cfg[key]
+        except Exception:
+            # Skip keys that can't be accessed (e.g., missing mandatory values)
+            continue
+
+        if isinstance(value, DictConfig):
+            _resolve_mlflow_configs_recursive(value)
+        elif isinstance(value, ListConfig):
+            # Also check ListConfig items that might be DictConfigs
+            for item in value:
+                if isinstance(item, DictConfig):
+                    _resolve_mlflow_configs_recursive(item)
+
 
 def _process_entry(cfg: DictConfig, accumulator: DictConfig, item: str) -> DictConfig:
-    """Process a single MLflow default item."""
+    """Process a single MLflow default item.
+
+    Handles both "_self_" entries (merging the original config) and MLflow entries
+    (downloading and merging configs from MLflow runs). Uses deep merge semantics
+    matching Hydra's defaults list behavior.
+
+    Args:
+        cfg: The original config being processed.
+        accumulator: The accumulated merged config so far.
+        item: The entry to process (either "_self_" or an MLflow entry string).
+
+    Returns:
+        The updated accumulator with the entry merged in.
+    """
     if item == "_self_":
         # Standard merge semantics: last entry wins (matches Hydra behavior)
         # OmegaConf.merge does deep merge (nested dicts get merged), which matches
         # how Hydra composes configs with defaults list
         result = cast(DictConfig, OmegaConf.merge(accumulator, cfg))
-        # Resolve load_source if it exists, as it may be merged unresolved
-        _resolve_load_source(result)
+        # Resolve any MLflow configs that may have been merged unresolved
+        _resolve_mlflow_configs_recursive(result)
         return result
-        # # When merging _self_, we need to allow cfg to override accumulator
-        # # Use merge but ensure top-level keys from cfg replace those in accumulator
-        # result = OmegaConf.create()
-        # # Copy all keys from accumulator first
-        # for key in accumulator.keys():
-        #     result[key] = accumulator[key]
-        # # Then copy keys from cfg, which will override accumulator keys
-        # for key in cfg.keys():
-        #     if key != "mlflow_defaults":  # Skip mlflow_defaults as it's metadata
-        #         result[key] = cfg[key]
-        # return cast(DictConfig, result)
 
     parsed_entry = _parse_entry(item)
 
@@ -240,7 +284,22 @@ def _process_entry(cfg: DictConfig, accumulator: DictConfig, item: str) -> DictC
 
 @dynamic_resolve
 def load_mlflow_defaults(cfg: DictConfig) -> DictConfig:
-    """Load defaults from MLflow runs."""
+    """Load and merge configs from MLflow runs based on mlflow_defaults list.
+
+    Processes the mlflow_defaults list in the config, downloading artifacts from
+    MLflow runs and merging them according to the specified package paths and
+    select paths. If "_self_" is not present, it is appended to the end.
+
+    Args:
+        cfg: The config containing mlflow_defaults list and MLflow config nodes.
+
+    Returns:
+        A new config with MLflow defaults merged in. The mlflow_defaults key is
+        removed from the result.
+
+    See Also:
+        LOAD_SUBCONFIGS.md for the full specification and examples.
+    """
     mlflow_defaults: ListConfig | None = cfg.get("mlflow_defaults")
     if mlflow_defaults is None:
         return cfg
@@ -254,7 +313,6 @@ def load_mlflow_defaults(cfg: DictConfig) -> DictConfig:
         normalized_item = _normalize_item(item)
         accumulator = _process_entry(cfg, accumulator, normalized_item)
 
-    # if "mlflow_defaults" in accumulator:
     del accumulator["mlflow_defaults"]
 
     return accumulator

From 6d65c0b2cbf63d0264939e6c3e99a863624733fc Mon Sep 17 00:00:00 2001
From: Eric Alt <ericallenalt@gmail.com>
Date: Fri, 16 Jan 2026 15:48:40 -0800
Subject: [PATCH 16/31] Fix lint issues

---
 tests/run_management/test_load_mlflow_defaults.py | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

diff --git a/tests/run_management/test_load_mlflow_defaults.py b/tests/run_management/test_load_mlflow_defaults.py
index 796b3dc8..7b0db2e4 100644
--- a/tests/run_management/test_load_mlflow_defaults.py
+++ b/tests/run_management/test_load_mlflow_defaults.py
@@ -570,9 +570,9 @@ def download_side_effect(_run_id=None, path=None, _dst_path=None, **_kwargs):
         # Path might be "config" or "config.yaml" depending on parsing
         if path in ("config", "config.yaml"):
             return str(artifact_path_1)
-        elif path in ("other_artifact", "other_artifact.yaml"):
+        if path in ("other_artifact", "other_artifact.yaml"):
             return str(artifact_path_2)
-        return str(artifact_path_1)
+        return ""
 
     mock_download.side_effect = download_side_effect
 
@@ -604,9 +604,9 @@ def download_side_effect(_run_id=None, path=None, _dst_path=None, **_kwargs):
         # Path might be "config1" or "config1.yaml" depending on parsing
         if path in ("config1", "config1.yaml"):
             return str(artifact_path_1)
-        elif path in ("config2", "config2.yaml"):
+        if path in ("config2", "config2.yaml"):
             return str(artifact_path_2)
-        return str(artifact_path_1)
+        return ""
 
     mock_download.side_effect = download_side_effect
 

From addcf13deda2624a268894e6614c470a1a3a42d3 Mon Sep 17 00:00:00 2001
From: Eric Alt <ericallenalt@gmail.com>
Date: Fri, 16 Jan 2026 15:50:09 -0800
Subject: [PATCH 17/31] Move spec to docs dir

---
 LOAD_SUBCONFIGS.md => docs/LOAD_SUBCONFIGS.md | 0
 1 file changed, 0 insertions(+), 0 deletions(-)
 rename LOAD_SUBCONFIGS.md => docs/LOAD_SUBCONFIGS.md (100%)

diff --git a/LOAD_SUBCONFIGS.md b/docs/LOAD_SUBCONFIGS.md
similarity index 100%
rename from LOAD_SUBCONFIGS.md
rename to docs/LOAD_SUBCONFIGS.md

From d47a271bb7af03ec87ea6a8100cd1ab820c56a62 Mon Sep 17 00:00:00 2001
From: Eric Alt <ericallenalt@gmail.com>
Date: Fri, 16 Jan 2026 16:20:04 -0800
Subject: [PATCH 18/31] Implement override flag

---
 docs/LOAD_SUBCONFIGS.md                       |  35 ++++-
 .../structured_configs/mlflow_defaults.py     |  39 ++++-
 .../override_flag_expected.yaml               |   1 -
 .../test_load_mlflow_defaults.py              | 138 +++++++++++++++++-
 4 files changed, 198 insertions(+), 15 deletions(-)

diff --git a/docs/LOAD_SUBCONFIGS.md b/docs/LOAD_SUBCONFIGS.md
index 03b2af68..b2bed420 100644
--- a/docs/LOAD_SUBCONFIGS.md
+++ b/docs/LOAD_SUBCONFIGS.md
@@ -25,12 +25,13 @@ OPTION            : ARTIFACT_PATH | SELECT_PATH | ARTIFACT_PATH#SELECT_PATH | nu
   - If the path does not contain `/`, it is ambiguous and defaults to `SELECT_PATH`. To force it to be an `ARTIFACT_PATH`, append `#` (e.g., `my_artifact#`).
 - `SELECT_PATH`: OmegaConf dot path within the source YAML to the subconfig to import (default: root).
 - `_self_`, `optional`, and `override` function the same as their Hydra equivalents.
-  - `_self_` determines the composition order. If missing, it is appended to the end.
+  - `_self_` determines the composition order. If missing, it is always appended to the end. If `_self_` is explicitly included, it is processed in the specified order.
   - Standard merge semantics apply: **whatever comes last wins**, whether it's `_self_` or an MLflow entry.
     - If `_self_` is first: MLflow entries merged after `_self_` will override the original config.
     - If `_self_` is last: Original config merged after MLflow entries will override MLflow content.
+    - To override everything from `_self_`, explicitly include `_self_` first, then use `override` at root.
   - `optional` suppresses errors if the artifact or selected subconfig is missing.
-  - `override` is a semantic indicator for intention but functionally both override and standard merge result in overwrites.
+  - `override` causes the loaded config to completely replace the value at the package path instead of merging. When override is used, any existing content at that path is completely replaced with the loaded content (no deep merge).
 
 ### Merge Semantics for Overlapping Keys
 
@@ -120,9 +121,12 @@ defaults:
   - train: smoke
 mlflow_defaults:
   # usage of '#' explicit defines ARTIFACT_PATH (train) and SELECT_PATH (optimizer)
+  # override flag causes complete replacement of train.optimizer (no merge)
   - override load_source@train.optimizer: train#optimizer
 ```
 
+**Note**: With `override`, the entire `train.optimizer` section is replaced with the content from the artifact. Any existing keys in `train.optimizer` that are not in the artifact will be removed.
+
 ### load `model` and `generative_process` from previous run
 
 ```yaml
@@ -226,7 +230,32 @@ other_section:
 
 If both configs had `other_section.foo`, the value from `config2` would win (last entry wins for conflicting keys).
 
-**Note**: The `override` flag is currently a semantic indicator and doesn't change merge behavior. In the future, it may be used to specify replacement instead of merge.
+**Note**: The `override` flag changes merge behavior: when specified, the loaded config completely replaces the value at the package path instead of merging. This means all existing keys at that path are removed and replaced with the loaded content.
+
+### Override flag behavior with `_self_`
+
+When `override` is used, the order of entries still matters. If `_self_` is explicitly included in the list, it will be processed in the order specified:
+
+```yaml
+mlflow_defaults:
+  - load_source_1@other_section: config1# # Merges at other_section
+  - override load_source_2@other_section: config2# # Replaces other_section completely
+  - _self_ # Merges original config, so original values can override the override
+```
+
+In this example:
+
+1. `config1` is merged into `other_section` (preserving existing keys, merging nested dicts)
+2. `config2` completely replaces `other_section` (removing all keys from step 1)
+3. `_self_` merges the original config, so any keys in the original `other_section` will override/replace what was loaded from `config2`
+
+**Important**: `_self_` is always auto-appended if not explicitly included, regardless of override usage. If you want to override everything from `_self_`, explicitly include `_self_` first, then use `override` at root:
+
+```yaml
+mlflow_defaults:
+  - _self_ # Explicitly include first
+  - override load_source # Then override at root, replacing everything
+```
 
 ## Implementation:
 
diff --git a/simplexity/structured_configs/mlflow_defaults.py b/simplexity/structured_configs/mlflow_defaults.py
index 414ac310..4c15bcdd 100644
--- a/simplexity/structured_configs/mlflow_defaults.py
+++ b/simplexity/structured_configs/mlflow_defaults.py
@@ -40,6 +40,7 @@ class _ParsedEntry(NamedTuple):
     """Parsed MLflow default item."""
 
     optional: bool
+    override: bool
     target: str
     package: str
     artifact_path: str | None
@@ -75,6 +76,7 @@ def _parse_entry(item: str) -> _ParsedEntry:
     groups = match.groupdict()
     flags_str = groups.get("flags", "")
     optional = "optional" in flags_str
+    override = "override" in flags_str
 
     target = groups["target"]
     if not target:
@@ -92,15 +94,15 @@ def _parse_entry(item: str) -> _ParsedEntry:
         package = option.rsplit(".", 1)[1] if "." in option else option
         artifact_path = "config"
         select_path = option
-        return _ParsedEntry(optional, target, package, artifact_path, select_path)
+        return _ParsedEntry(optional, override, target, package, artifact_path, select_path)
 
     package = package or "."
 
     if option == "null":
-        return _ParsedEntry(optional, target, package, None, None)
+        return _ParsedEntry(optional, override, target, package, None, None)
 
     artifact_path, select_path = _parse_option(option or "")
-    return _ParsedEntry(optional, target, package, artifact_path, select_path)
+    return _ParsedEntry(optional, override, target, package, artifact_path, select_path)
 
 
 def _get_target_config(cfg: DictConfig, parsed_entry: _ParsedEntry) -> Any | None:
@@ -239,7 +241,10 @@ def _process_entry(cfg: DictConfig, accumulator: DictConfig, item: str) -> DictC
 
     Handles both "_self_" entries (merging the original config) and MLflow entries
     (downloading and merging configs from MLflow runs). Uses deep merge semantics
-    matching Hydra's defaults list behavior.
+    matching Hydra's defaults list behavior, unless the override flag is specified.
+
+    When override is True, the loaded config completely replaces the value at the
+    package path instead of merging with existing content.
 
     Args:
         cfg: The original config being processed.
@@ -247,7 +252,7 @@ def _process_entry(cfg: DictConfig, accumulator: DictConfig, item: str) -> DictC
         item: The entry to process (either "_self_" or an MLflow entry string).
 
     Returns:
-        The updated accumulator with the entry merged in.
+        The updated accumulator with the entry merged in or replaced (if override).
     """
     if item == "_self_":
         # Standard merge semantics: last entry wins (matches Hydra behavior)
@@ -272,9 +277,27 @@ def _process_entry(cfg: DictConfig, accumulator: DictConfig, item: str) -> DictC
             if parsed_entry.optional:
                 return accumulator
             raise ValueError(f"Target config not found for entry: {item}")
+        if parsed_entry.override:
+            # Override at root: replace entire accumulator with loaded config
+            return cast(DictConfig, loaded_config)
         return cast(DictConfig, OmegaConf.merge(accumulator, loaded_config))
 
-    # When merging at a package (not root), use deep merge semantics (matches Hydra behavior)
+    # When merging at a package (not root)
+    if parsed_entry.override:
+        # Override: completely replace the value at the package path
+        # Traverse the path and directly assign the value to replace any existing content
+        package_parts = parsed_entry.package.split(".")
+        target = accumulator
+        # Navigate to the parent of the target key
+        for part in package_parts[:-1]:
+            if part not in target:
+                target[part] = OmegaConf.create()
+            target = target[part]
+        # Directly assign the new value, completely replacing any existing value
+        target[package_parts[-1]] = loaded_config
+        return accumulator
+
+    # Standard merge: use deep merge semantics (matches Hydra behavior)
     # This ensures that MLflow content merges with any existing content at that package path,
     # preserving non-conflicting keys and deeply merging nested dictionaries
     package_conf = OmegaConf.create()
@@ -313,6 +336,8 @@ def load_mlflow_defaults(cfg: DictConfig) -> DictConfig:
         normalized_item = _normalize_item(item)
         accumulator = _process_entry(cfg, accumulator, normalized_item)
 
-    del accumulator["mlflow_defaults"]
+    # Remove mlflow_defaults key if it exists (it may not exist if override at root replaced everything)
+    if "mlflow_defaults" in accumulator:
+        del accumulator["mlflow_defaults"]
 
     return accumulator
diff --git a/tests/end_to_end/mlflow_defaults_configs/override_flag_expected.yaml b/tests/end_to_end/mlflow_defaults_configs/override_flag_expected.yaml
index e07228c3..8d4d28d3 100644
--- a/tests/end_to_end/mlflow_defaults_configs/override_flag_expected.yaml
+++ b/tests/end_to_end/mlflow_defaults_configs/override_flag_expected.yaml
@@ -1,5 +1,4 @@
 other_section:
-  foo: bar
   key: 3
   dict:
     subkey: 30
diff --git a/tests/run_management/test_load_mlflow_defaults.py b/tests/run_management/test_load_mlflow_defaults.py
index 7b0db2e4..0410fafd 100644
--- a/tests/run_management/test_load_mlflow_defaults.py
+++ b/tests/run_management/test_load_mlflow_defaults.py
@@ -172,12 +172,13 @@ def test_implicit_artifact_select(base_cfg: DictConfig, mock_download: MagicMock
 
 
 def test_override_flag(base_cfg: DictConfig, mock_download: MagicMock, tmp_path: Path):
-    """Test override flag."""
+    """Test override flag replaces instead of merging."""
+    # Artifact contains different keys than base config to verify replacement
     artifact_path = tmp_path / "nondefault_config.yaml"
-    artifact_path.write_text("foo: new_bar\n")
+    artifact_path.write_text("new_key: new_value\nother_key: other_value\n")
     mock_download.return_value = str(artifact_path)
 
-    # Place _self_ first so the loaded config overrides the base config
+    # Place _self_ first so original config is merged, then override replaces it
     cfg = OmegaConf.merge(
         base_cfg,
         {
@@ -185,7 +186,136 @@ def test_override_flag(base_cfg: DictConfig, mock_download: MagicMock, tmp_path:
         },
     )
     loaded_cfg: DictConfig = load_mlflow_defaults(cfg)
-    assert OmegaConf.select(loaded_cfg, "other_section.foo") == "new_bar"
+
+    # Verify override replaced the entire other_section (no merge)
+    assert OmegaConf.select(loaded_cfg, "other_section.new_key") == "new_value"
+    assert OmegaConf.select(loaded_cfg, "other_section.other_key") == "other_value"
+    # Original foo: bar should be gone (replaced, not merged)
+    assert OmegaConf.select(loaded_cfg, "other_section.foo") is None
+
+
+def test_override_flag_at_root(base_cfg: DictConfig, mock_download: MagicMock, tmp_path: Path):
+    """Test override flag at root level with explicit _self_ first."""
+    artifact_path = tmp_path / "config.yaml"
+    artifact_path.write_text("root_key: root_value\nother_key: other_value\n")
+    mock_download.return_value = str(artifact_path)
+
+    # To override everything, explicitly include _self_ first, then override
+    cfg = OmegaConf.merge(
+        base_cfg,
+        {
+            "mlflow_defaults": ["_self_", "override previous_run"],
+        },
+    )
+    loaded_cfg: DictConfig = load_mlflow_defaults(cfg)
+
+    # Verify override replaced entire config at root (after _self_ was processed)
+    assert OmegaConf.select(loaded_cfg, "root_key") == "root_value"
+    assert OmegaConf.select(loaded_cfg, "other_key") == "other_value"
+    # Original keys should be gone (override replaced everything)
+    assert OmegaConf.select(loaded_cfg, "other_section") is None
+    assert OmegaConf.select(loaded_cfg, "previous_run") is None
+
+
+def test_override_flag_at_root_without_explicit_self(base_cfg: DictConfig, mock_download: MagicMock, tmp_path: Path):
+    """Test override flag at root level without explicit _self_ - _self_ is auto-appended."""
+    artifact_path = tmp_path / "config.yaml"
+    artifact_path.write_text("root_key: root_value\nother_key: other_value\n")
+    mock_download.return_value = str(artifact_path)
+
+    # Without explicit _self_, it will be auto-appended at the end
+    cfg = OmegaConf.merge(
+        base_cfg,
+        {
+            "mlflow_defaults": ["override previous_run"],
+        },
+    )
+    loaded_cfg: DictConfig = load_mlflow_defaults(cfg)
+
+    # Override replaced everything, but _self_ was auto-appended and merged after
+    # So original keys from _self_ should be back (last entry wins for conflicting keys)
+    assert OmegaConf.select(loaded_cfg, "other_section.foo") == "bar"
+    assert OmegaConf.select(loaded_cfg, "previous_run.tracking_uri") == "databricks"
+    # Keys from override that don't conflict with _self_ should still be present
+    assert OmegaConf.select(loaded_cfg, "root_key") == "root_value"
+    assert OmegaConf.select(loaded_cfg, "other_key") == "other_value"
+
+
+def test_override_flag_nested_path(base_cfg: DictConfig, mock_download: MagicMock, tmp_path: Path):
+    """Test override flag with nested package path."""
+    artifact_path = tmp_path / "config.yaml"
+    artifact_path.write_text("nested_key: nested_value\n")
+    mock_download.return_value = str(artifact_path)
+
+    cfg = OmegaConf.merge(
+        base_cfg,
+        {
+            "mlflow_defaults": ["_self_", "override previous_run@nested.path: config#"],
+        },
+    )
+    loaded_cfg: DictConfig = load_mlflow_defaults(cfg)
+
+    # Verify override replaced at nested path
+    assert OmegaConf.select(loaded_cfg, "nested.path.nested_key") == "nested_value"
+    # Original other_section should still exist (not affected)
+    assert OmegaConf.select(loaded_cfg, "other_section.foo") == "bar"
+
+
+def test_override_with_explicit_self(base_cfg: DictConfig, mock_download: MagicMock, tmp_path: Path):
+    """Test that _self_ is processed when explicitly included after override."""
+    artifact_path = tmp_path / "config.yaml"
+    artifact_path.write_text("key1: value1\nkey2: value2\n")
+    mock_download.return_value = str(artifact_path)
+
+    # Test with override followed by explicit _self_
+    cfg = OmegaConf.merge(
+        base_cfg,
+        {
+            "mlflow_defaults": [
+                "override previous_run@other_section: config#",
+                "_self_",  # Explicitly included, should merge after override
+            ],
+        },
+    )
+    loaded_cfg: DictConfig = load_mlflow_defaults(cfg)
+
+    # Override replaced other_section, but _self_ merged after, so original foo should be back
+    assert OmegaConf.select(loaded_cfg, "other_section.foo") == "bar"  # From _self_
+    assert OmegaConf.select(loaded_cfg, "other_section.key1") == "value1"  # From override
+    assert OmegaConf.select(loaded_cfg, "other_section.key2") == "value2"  # From override
+
+
+def test_override_vs_merge_behavior(base_cfg: DictConfig, mock_download: MagicMock, tmp_path: Path):
+    """Test that override replaces while normal entry merges."""
+    artifact_path = tmp_path / "config.yaml"
+    artifact_path.write_text("key1: value1\nkey2: value2\n")
+    mock_download.return_value = str(artifact_path)
+
+    # Test with merge (no override)
+    cfg_merge = OmegaConf.merge(
+        base_cfg,
+        {
+            "mlflow_defaults": ["_self_", "previous_run@other_section: config#"],
+        },
+    )
+    loaded_merge: DictConfig = load_mlflow_defaults(cfg_merge)
+    # With merge, original foo should be preserved
+    assert OmegaConf.select(loaded_merge, "other_section.foo") == "bar"
+    assert OmegaConf.select(loaded_merge, "other_section.key1") == "value1"
+    assert OmegaConf.select(loaded_merge, "other_section.key2") == "value2"
+
+    # Test with override
+    cfg_override = OmegaConf.merge(
+        base_cfg,
+        {
+            "mlflow_defaults": ["_self_", "override previous_run@other_section: config#"],
+        },
+    )
+    loaded_override: DictConfig = load_mlflow_defaults(cfg_override)
+    # With override, original foo should be replaced
+    assert OmegaConf.select(loaded_override, "other_section.foo") is None
+    assert OmegaConf.select(loaded_override, "other_section.key1") == "value1"
+    assert OmegaConf.select(loaded_override, "other_section.key2") == "value2"
 
 
 def test_optional_flag_missing_artifact(base_cfg: DictConfig):

From f886be3c9a86ef7ceb34b2a40f77adf58532b180 Mon Sep 17 00:00:00 2001
From: Eric Alt <ericallenalt@gmail.com>
Date: Fri, 16 Jan 2026 16:40:16 -0800
Subject: [PATCH 19/31] Address PR comments

---
 .../structured_configs/mlflow_defaults.py     | 31 +++++++++++--------
 1 file changed, 18 insertions(+), 13 deletions(-)

diff --git a/simplexity/structured_configs/mlflow_defaults.py b/simplexity/structured_configs/mlflow_defaults.py
index 4c15bcdd..7502eb79 100644
--- a/simplexity/structured_configs/mlflow_defaults.py
+++ b/simplexity/structured_configs/mlflow_defaults.py
@@ -22,12 +22,15 @@
 
 from mlflow import MlflowClient
 from omegaconf import DictConfig, ListConfig, OmegaConf
+from omegaconf.errors import MissingMandatoryValue
 
 from simplexity.exceptions import ConfigValidationError
 from simplexity.logger import SIMPLEXITY_LOGGER
 from simplexity.structured_configs.mlflow import resolve_mlflow_config, validate_mlflow_config
 from simplexity.utils.config_utils import dynamic_resolve
 
+DEFAULT_ARTIFACT_NAME = "config"
+
 # TARGET(@PACKAGE)?
 # [optional|override]? TARGET(@PACKAGE)? : OPTION
 FLAGS_STR = r"(?P<flags>(?:(?:optional|override)\s+)*)"
@@ -51,7 +54,7 @@ def _parse_option(option: str) -> tuple[str, str | None]:
     """Parse artifact path and select path from option string."""
     if "#" in option:
         artifact_part, select_part = option.split("#", 1)
-        artifact_path = artifact_part.strip() or "config"
+        artifact_path = artifact_part.strip() or DEFAULT_ARTIFACT_NAME
         select_path = select_part.strip() or None
         return artifact_path, select_path
 
@@ -60,7 +63,7 @@ def _parse_option(option: str) -> tuple[str, str | None]:
         select_path = None
         return artifact_path, select_path
 
-    artifact_path = "config"
+    artifact_path = DEFAULT_ARTIFACT_NAME
     select_path = option.strip() or None
     return artifact_path, select_path
 
@@ -92,7 +95,7 @@ def _parse_entry(item: str) -> _ParsedEntry:
     # Treat VALUE as both PACKAGE and SELECT_PATH
     if package is None and option is not None and "/" not in option and "#" not in option:
         package = option.rsplit(".", 1)[1] if "." in option else option
-        artifact_path = "config"
+        artifact_path = DEFAULT_ARTIFACT_NAME
         select_path = option
         return _ParsedEntry(optional, override, target, package, artifact_path, select_path)
 
@@ -147,13 +150,13 @@ def _get_target_config(cfg: DictConfig, parsed_entry: _ParsedEntry) -> Any | Non
     with tempfile.TemporaryDirectory() as tmp_dir:
         try:
             local_path = client.download_artifacts(run_id=run_id, path=artifact_path, dst_path=tmp_dir)
-        except Exception as e:
+        except (OSError, FileNotFoundError) as e:
             SIMPLEXITY_LOGGER.warning("Failed to download artifact from MLflow '%s': %s", parsed_entry.target, e)
             return None
 
         try:
             loaded_config = OmegaConf.load(local_path)
-        except Exception as e:
+        except (OSError, FileNotFoundError, ValueError) as e:
             SIMPLEXITY_LOGGER.warning("Failed to load MLflow default '%s': %s", parsed_entry.target, e)
             return None
 
@@ -183,8 +186,8 @@ def _normalize_item(item: str | DictConfig) -> str:
         if len(keys) == 1:
             key = keys[0]
             value = item[key]
-            # If value is "config" (default artifact), treat as simple CONFIG entry
-            if value == "config":
+            # If value is the default artifact name, treat as simple CONFIG entry
+            if value == DEFAULT_ARTIFACT_NAME:
                 return str(key)
             return f"{key}: {value}"
         # Multiple keys - convert entire dict to string representation
@@ -223,7 +226,7 @@ def _resolve_mlflow_configs_recursive(cfg: DictConfig) -> None:
     for key in cfg:
         try:
             value = cfg[key]
-        except Exception:
+        except MissingMandatoryValue:
             # Skip keys that can't be accessed (e.g., missing mandatory values)
             continue
 
@@ -279,7 +282,7 @@ def _process_entry(cfg: DictConfig, accumulator: DictConfig, item: str) -> DictC
             raise ValueError(f"Target config not found for entry: {item}")
         if parsed_entry.override:
             # Override at root: replace entire accumulator with loaded config
-            return cast(DictConfig, loaded_config)
+            return loaded_config
         return cast(DictConfig, OmegaConf.merge(accumulator, loaded_config))
 
     # When merging at a package (not root)
@@ -327,12 +330,14 @@ def load_mlflow_defaults(cfg: DictConfig) -> DictConfig:
     if mlflow_defaults is None:
         return cfg
 
-    if "_self_" not in mlflow_defaults:
-        mlflow_defaults.append("_self_")
+    # Create a copy to avoid mutating the input config
+    mlflow_defaults_copy = cast(ListConfig, OmegaConf.create(list(mlflow_defaults)))
+    if "_self_" not in mlflow_defaults_copy:
+        mlflow_defaults_copy.append("_self_")
 
-    accumulator = OmegaConf.create()
+    accumulator = cast(DictConfig, OmegaConf.create())
 
-    for item in mlflow_defaults:
+    for item in mlflow_defaults_copy:
         normalized_item = _normalize_item(item)
         accumulator = _process_entry(cfg, accumulator, normalized_item)
 

From 409927cf9ce259acbb063e579e15c42f83f3920a Mon Sep 17 00:00:00 2001
From: Eric Alt <ericallenalt@gmail.com>
Date: Fri, 16 Jan 2026 16:43:01 -0800
Subject: [PATCH 20/31] Fix issue

---
 simplexity/structured_configs/mlflow_defaults.py | 11 +++++++----
 1 file changed, 7 insertions(+), 4 deletions(-)

diff --git a/simplexity/structured_configs/mlflow_defaults.py b/simplexity/structured_configs/mlflow_defaults.py
index 7502eb79..c19fb42b 100644
--- a/simplexity/structured_configs/mlflow_defaults.py
+++ b/simplexity/structured_configs/mlflow_defaults.py
@@ -20,6 +20,7 @@
 import tempfile
 from typing import Any, NamedTuple, cast
 
+import yaml
 from mlflow import MlflowClient
 from omegaconf import DictConfig, ListConfig, OmegaConf
 from omegaconf.errors import MissingMandatoryValue
@@ -91,6 +92,11 @@ def _parse_entry(item: str) -> _ParsedEntry:
     package = groups.get("package")
     option = groups.get("option")
 
+    # Check for null keyword first, before special case handling
+    if option == "null":
+        package = package or "."
+        return _ParsedEntry(optional, override, target, package, None, None)
+
     # Special case: TARGET: VALUE (no @) where VALUE doesn't contain / or #
     # Treat VALUE as both PACKAGE and SELECT_PATH
     if package is None and option is not None and "/" not in option and "#" not in option:
@@ -101,9 +107,6 @@ def _parse_entry(item: str) -> _ParsedEntry:
 
     package = package or "."
 
-    if option == "null":
-        return _ParsedEntry(optional, override, target, package, None, None)
-
     artifact_path, select_path = _parse_option(option or "")
     return _ParsedEntry(optional, override, target, package, artifact_path, select_path)
 
@@ -156,7 +159,7 @@ def _get_target_config(cfg: DictConfig, parsed_entry: _ParsedEntry) -> Any | Non
 
         try:
             loaded_config = OmegaConf.load(local_path)
-        except (OSError, FileNotFoundError, ValueError) as e:
+        except (OSError, FileNotFoundError, ValueError, yaml.YAMLError) as e:
             SIMPLEXITY_LOGGER.warning("Failed to load MLflow default '%s': %s", parsed_entry.target, e)
             return None
 

From 8db0e6330d5618b74494fbb80f30ca3eb0307f63 Mon Sep 17 00:00:00 2001
From: Eric Alt <ericallenalt@gmail.com>
Date: Fri, 16 Jan 2026 16:47:07 -0800
Subject: [PATCH 21/31] Add null option to yaml

---
 simplexity/structured_configs/mlflow_defaults.py  |  4 ++++
 tests/run_management/test_load_mlflow_defaults.py | 15 +++++++++++++++
 2 files changed, 19 insertions(+)

diff --git a/simplexity/structured_configs/mlflow_defaults.py b/simplexity/structured_configs/mlflow_defaults.py
index c19fb42b..29811a2b 100644
--- a/simplexity/structured_configs/mlflow_defaults.py
+++ b/simplexity/structured_configs/mlflow_defaults.py
@@ -181,6 +181,7 @@ def _normalize_item(item: str | DictConfig) -> str:
 
     If item is a DictConfig with a single key-value pair:
     - If value is "config" (the default artifact), treat as just the key (no option)
+    - If value is None (Python None, from YAML null), convert to "key: null" string
     - Otherwise, convert to "key: value" format
     Otherwise, convert to string.
     """
@@ -192,6 +193,9 @@ def _normalize_item(item: str | DictConfig) -> str:
             # If value is the default artifact name, treat as simple CONFIG entry
             if value == DEFAULT_ARTIFACT_NAME:
                 return str(key)
+            # If value is None (from YAML null), convert to "null" string
+            if value is None:
+                return f"{key}: null"
             return f"{key}: {value}"
         # Multiple keys - convert entire dict to string representation
         return str(item)
diff --git a/tests/run_management/test_load_mlflow_defaults.py b/tests/run_management/test_load_mlflow_defaults.py
index 0410fafd..ad44511d 100644
--- a/tests/run_management/test_load_mlflow_defaults.py
+++ b/tests/run_management/test_load_mlflow_defaults.py
@@ -367,6 +367,21 @@ def test_null_option_optional(base_cfg: DictConfig):
     assert OmegaConf.select(loaded_cfg, "other_section") == cfg.other_section
 
 
+def test_null_option_dict_syntax_error(base_cfg: DictConfig):
+    """Test null option with YAML dict syntax raises error."""
+    # YAML dict syntax: - previous_run: null becomes DictConfig with None value
+    # This should be treated the same as string format: "previous_run: null"
+    cfg = OmegaConf.merge(
+        base_cfg,
+        {
+            "mlflow_defaults": [OmegaConf.create({"previous_run": None})],
+        },
+    )
+    # Should raise ValueError because null option without optional flag is mandatory
+    with pytest.raises(ValueError, match="Target config not found for entry"):
+        load_mlflow_defaults(cfg)
+
+
 def test_composition_order(base_cfg: DictConfig, mock_download: MagicMock, tmp_path: Path):
     """Test _self_ placement in composition order."""
     artifact_path = tmp_path / "nondefault_config.yaml"

From a73e1ccf0ab3c7c843eb0584720710fa28e41d0e Mon Sep 17 00:00:00 2001
From: Eric Alt <ericallenalt@gmail.com>
Date: Fri, 16 Jan 2026 16:52:17 -0800
Subject: [PATCH 22/31] Add comments

---
 simplexity/structured_configs/mlflow_defaults.py | 3 +++
 1 file changed, 3 insertions(+)

diff --git a/simplexity/structured_configs/mlflow_defaults.py b/simplexity/structured_configs/mlflow_defaults.py
index 29811a2b..164dbfa6 100644
--- a/simplexity/structured_configs/mlflow_defaults.py
+++ b/simplexity/structured_configs/mlflow_defaults.py
@@ -194,8 +194,11 @@ def _normalize_item(item: str | DictConfig) -> str:
             if value == DEFAULT_ARTIFACT_NAME:
                 return str(key)
             # If value is None (from YAML null), convert to "null" string
+            # IMPORTANT: Check for None BEFORE f-string formatting to avoid converting
+            # Python None to the string "None" instead of "null"
             if value is None:
                 return f"{key}: null"
+            # Format the value - this would convert None to "None" if the check above is bypassed
             return f"{key}: {value}"
         # Multiple keys - convert entire dict to string representation
         return str(item)

From 4024a27fbe94deff637f3f46fb96f67083c26cf1 Mon Sep 17 00:00:00 2001
From: Eric Alt <ericallenalt@gmail.com>
Date: Tue, 27 Jan 2026 09:31:13 -0800
Subject: [PATCH 23/31] Fix bug loading mlflow defaults in run management

---
 simplexity/run_management/run_management.py   |  6 ++
 tests/end_to_end/conftest.py                  | 68 +++++++++++++++++++
 .../test_load_config_from_mlflow.py           | 49 +------------
 .../test_managed_run_mlflow_defaults.py       | 47 +++++++++++++
 4 files changed, 122 insertions(+), 48 deletions(-)
 create mode 100644 tests/end_to_end/conftest.py
 create mode 100644 tests/end_to_end/test_managed_run_mlflow_defaults.py

diff --git a/simplexity/run_management/run_management.py b/simplexity/run_management/run_management.py
index 6f1c7764..e8b3eb9f 100644
--- a/simplexity/run_management/run_management.py
+++ b/simplexity/run_management/run_management.py
@@ -635,6 +635,12 @@ def wrapper(*args: Any, **kwargs: Any) -> Any:
             try:
                 cfg = get_config(args, kwargs)
                 cfg = load_mlflow_defaults(cfg)
+                if "cfg" in kwargs:
+                    kwargs["cfg"] = cfg
+                elif args and isinstance(args[0], DictConfig):
+                    args = (cfg, *args[1:])
+                else:
+                    kwargs["cfg"] = cfg
                 validate_base_config(cfg)
                 resolve_base_config(cfg, strict=strict)
                 with _setup_device(cfg), _setup_mlflow(cfg):
diff --git a/tests/end_to_end/conftest.py b/tests/end_to_end/conftest.py
new file mode 100644
index 00000000..1e8d3679
--- /dev/null
+++ b/tests/end_to_end/conftest.py
@@ -0,0 +1,68 @@
+"""End-to-end tests configuration."""
+
+# pylint: disable=all
+# Temporarily disable all pylint checkers during AST traversal to prevent crash.
+# The imports checker crashes when resolving simplexity package imports due to a bug
+# in pylint/astroid: https://github.com/pylint-dev/pylint/issues/10185
+# pylint: enable=all
+# Re-enable all pylint checkers for the checking phase. This allows other checks
+# (code quality, style, undefined names, etc.) to run normally while bypassing
+# the problematic imports checker that would crash during AST traversal.
+
+import shutil
+from pathlib import Path
+
+import mlflow
+import pytest
+from hydra import compose, initialize_config_dir
+from omegaconf import DictConfig
+
+from simplexity.logging.mlflow_logger import MLFlowLogger
+from simplexity.utils.mlflow_utils import get_experiment, get_run
+
+CONFIG_DIR = str(Path(__file__).parent / "mlflow_defaults_configs")
+EXPERIMENT_NAME = "test_mlflow_defaults"
+
+
+@pytest.fixture(scope="session")
+def setup_dir(tmp_path_factory: pytest.TempPathFactory) -> Path:
+    """Setup function."""
+    tmp_path = tmp_path_factory.mktemp("mlflow_defaults")
+    config_path = tmp_path / "configs"
+    shutil.copytree(f"{CONFIG_DIR}/setup", str(config_path))
+    tracking_uri = f"sqlite:///{tmp_path.resolve()}/mlflow.db"
+    mlflow.set_tracking_uri(tracking_uri)
+
+    def log_to_mlflow(cfg: DictConfig | None, config_names: list[str], run_name: str) -> None:
+        """Save config."""
+        experiment = get_experiment(experiment_name=EXPERIMENT_NAME)
+        assert experiment is not None
+        experiment_id = experiment.experiment_id
+        run = get_run(run_name=run_name, experiment_id=experiment_id)
+        assert run is not None
+        run_id = run.info.run_id
+        with mlflow.start_run(
+            run_id=run_id,
+            experiment_id=experiment_id,
+            run_name=run_name,
+            log_system_metrics=False,
+        ):
+            logger = MLFlowLogger(tracking_uri=tracking_uri)
+            if cfg is not None:
+                logger.log_config(cfg, resolve=False)
+            for config_name in config_names:
+                logger.log_artifact(local_path=str(config_path / config_name), artifact_path="subdir")
+
+    def previous_run(config_name: str | None, config_names: list[str], run_name: str) -> None:
+        """Previous run."""
+        if config_name is None:
+            cfg = None
+        else:
+            with initialize_config_dir(config_dir=str(config_path)):
+                cfg = compose(config_name=config_name)
+        log_to_mlflow(cfg, config_names, run_name)
+
+    previous_run(config_name="prev_config_1", config_names=["special_1.yaml"], run_name="prev_run_1")
+    previous_run(config_name="prev_config_2", config_names=["special_2.yaml"], run_name="prev_run_2")
+
+    return tmp_path
diff --git a/tests/end_to_end/test_load_config_from_mlflow.py b/tests/end_to_end/test_load_config_from_mlflow.py
index 8440e553..b7db1fa1 100644
--- a/tests/end_to_end/test_load_config_from_mlflow.py
+++ b/tests/end_to_end/test_load_config_from_mlflow.py
@@ -9,67 +9,20 @@
 # (code quality, style, undefined names, etc.) to run normally while bypassing
 # the problematic imports checker that would crash during AST traversal.
 
-import shutil
 from pathlib import Path
 
 import mlflow
 import pytest
 from hydra import compose, initialize_config_dir
-from omegaconf import DictConfig, OmegaConf
+from omegaconf import OmegaConf
 
-from simplexity.logging.mlflow_logger import MLFlowLogger
 from simplexity.structured_configs.mlflow_defaults import load_mlflow_defaults
-from simplexity.utils.mlflow_utils import get_experiment, get_run
 
 CONFIG_DIR = str(Path(__file__).parent / "mlflow_defaults_configs")
 
 EXPERIMENT_NAME = "test_mlflow_defaults"
 
 
-@pytest.fixture(scope="session")
-def setup_dir(tmp_path_factory: pytest.TempPathFactory) -> Path:
-    """Setup function."""
-    tmp_path = tmp_path_factory.mktemp("mlflow_defaults")
-    config_path = tmp_path / "configs"
-    shutil.copytree(f"{CONFIG_DIR}/setup", str(config_path))
-    tracking_uri = f"sqlite:///{tmp_path.resolve()}/mlflow.db"
-    mlflow.set_tracking_uri(tracking_uri)
-
-    def log_to_mlflow(cfg: DictConfig | None, config_names: list[str], run_name: str) -> None:
-        """Save config."""
-        experiment = get_experiment(experiment_name=EXPERIMENT_NAME)
-        assert experiment is not None
-        experiment_id = experiment.experiment_id
-        run = get_run(run_name=run_name, experiment_id=experiment_id)
-        assert run is not None
-        run_id = run.info.run_id
-        with mlflow.start_run(
-            run_id=run_id,
-            experiment_id=experiment_id,
-            run_name=run_name,
-            log_system_metrics=False,
-        ):
-            logger = MLFlowLogger(tracking_uri=tracking_uri)
-            if cfg is not None:
-                logger.log_config(cfg, resolve=False)
-            for config_name in config_names:
-                logger.log_artifact(local_path=str(config_path / config_name), artifact_path="subdir")
-
-    def previous_run(config_name: str | None, config_names: list[str], run_name: str) -> None:
-        """Previous run."""
-        if config_name is None:
-            cfg = None
-        else:
-            with initialize_config_dir(config_dir=str(config_path)):
-                cfg = compose(config_name=config_name)
-        log_to_mlflow(cfg, config_names, run_name)
-
-    previous_run(config_name="prev_config_1", config_names=["special_1.yaml"], run_name="prev_run_1")
-    previous_run(config_name="prev_config_2", config_names=["special_2.yaml"], run_name="prev_run_2")
-
-    return tmp_path
-
-
 def test_setup(setup_dir: Path, tmp_path: Path) -> None:
     """Test setup."""
     tracking_uri = f"sqlite:///{setup_dir.resolve()}/mlflow.db"
diff --git a/tests/end_to_end/test_managed_run_mlflow_defaults.py b/tests/end_to_end/test_managed_run_mlflow_defaults.py
new file mode 100644
index 00000000..57210bbf
--- /dev/null
+++ b/tests/end_to_end/test_managed_run_mlflow_defaults.py
@@ -0,0 +1,47 @@
+from pathlib import Path
+
+from hydra import compose, initialize_config_dir
+from omegaconf import DictConfig, OmegaConf
+
+import simplexity
+
+CONFIG_DIR = str(Path(__file__).parent / "mlflow_defaults_configs")
+
+
+def test_managed_run_loads_mlflow_defaults(setup_dir: Path) -> None:
+    tracking_uri = f"sqlite:///{setup_dir.resolve()}/mlflow.db"
+    with initialize_config_dir(config_dir=CONFIG_DIR):
+        cfg = compose(
+            config_name="full_default_copy",
+            overrides=[f"load_source.tracking_uri={tracking_uri}"],
+        )
+
+    captured: dict[str, DictConfig] = {}
+
+    @simplexity.managed_run(strict=False)
+    def run(cfg: DictConfig, components: simplexity.Components) -> None:
+        captured["cfg"] = cfg
+
+    run(cfg)
+
+    actual = captured["cfg"]
+    experiment_id = OmegaConf.select(actual, "load_source.experiment_id")
+    run_id = OmegaConf.select(actual, "load_source.run_id")
+    assert experiment_id is not None
+    assert run_id is not None
+
+    expected_overrides = [
+        f"load_source.tracking_uri={tracking_uri}",
+        f'load_source.experiment_id="{experiment_id}"',
+        f'load_source.run_id="{run_id}"',
+    ]
+    with initialize_config_dir(config_dir=CONFIG_DIR):
+        expected = compose(config_name="full_default_copy_expected", overrides=expected_overrides)
+
+    actual_container = OmegaConf.to_container(actual, resolve=True)
+    actual_filtered = OmegaConf.create(actual_container)
+    for key in ("device", "seed", "tags"):
+        if key in actual_filtered:
+            del actual_filtered[key]
+
+    assert actual_filtered == expected

From 9f3b8233889620bbb293279f799ee683034aec68 Mon Sep 17 00:00:00 2001
From: Eric Alt <ericallenalt@gmail.com>
Date: Tue, 27 Jan 2026 09:31:50 -0800
Subject: [PATCH 24/31] Better error handling

---
 simplexity/structured_configs/mlflow_defaults.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/simplexity/structured_configs/mlflow_defaults.py b/simplexity/structured_configs/mlflow_defaults.py
index 164dbfa6..772f3ad4 100644
--- a/simplexity/structured_configs/mlflow_defaults.py
+++ b/simplexity/structured_configs/mlflow_defaults.py
@@ -139,7 +139,7 @@ def _get_target_config(cfg: DictConfig, parsed_entry: _ParsedEntry) -> Any | Non
 
     try:
         resolve_mlflow_config(target_node)
-    except ValueError as e:
+    except (ValueError, RuntimeError) as e:
         SIMPLEXITY_LOGGER.warning("Error resolving MLflow config: %s", e)
         return None
 
@@ -228,7 +228,7 @@ def _resolve_mlflow_configs_recursive(cfg: DictConfig) -> None:
             validate_mlflow_config(cfg)
             # If validation passes, this looks like an MLflow config, try to resolve it
             resolve_mlflow_config(cfg)
-        except (ConfigValidationError, ValueError):
+        except (ConfigValidationError, ValueError, RuntimeError):
             # Not an MLflow config or can't be resolved, continue to check nested configs
             pass
 

From b1bf935e799959bdd2416c2592d86a8718193387 Mon Sep 17 00:00:00 2001
From: Eric Alt <13019253+ealt@users.noreply.github.com>
Date: Thu, 26 Feb 2026 10:59:27 -0800
Subject: [PATCH 25/31] Enhance coverage check workflow: upload coverage
 artifacts, improve reporting, and adjust coverage check to be advisory only

---
 .github/workflows/simplexity.yaml | 75 ++++++++++++++++++++++++-------
 1 file changed, 59 insertions(+), 16 deletions(-)

diff --git a/.github/workflows/simplexity.yaml b/.github/workflows/simplexity.yaml
index f1013c1d..cc950596 100644
--- a/.github/workflows/simplexity.yaml
+++ b/.github/workflows/simplexity.yaml
@@ -77,46 +77,89 @@ jobs:
             exit 1
           fi
 
-      - name: Check coverage for new code (PRs only)
+      - name: Upload coverage artifact for coverage-check job
         if: github.event_name == 'pull_request'
+        uses: actions/upload-artifact@v4
+        with:
+          name: coverage-xml
+          path: coverage.xml
+          retention-days: 1
+
+      - name: Upload coverage reports to Codecov
+        uses: codecov/codecov-action@v5
+        with:
+          token: ${{ secrets.CODECOV_TOKEN }}
+          slug: Astera-org/simplexity
+          verbose: true
+          files: ./coverage.xml
+          fail_ci_if_error: false
+
+  coverage-check:
+    runs-on: ubuntu-latest
+    needs: unit-tests
+    if: github.event_name == 'pull_request'
+    steps:
+      - uses: actions/checkout@v4
+        with:
+          fetch-depth: 0 # Full history needed for diff-cover
+      - uses: astral-sh/setup-uv@v5
+        with:
+          enable-cache: true
+
+      - name: Download coverage artifact
+        uses: actions/download-artifact@v4
+        with:
+          name: coverage-xml
+          path: .
+
+      - name: Check coverage for new code
+        id: coverage_check
+        continue-on-error: true
         run: |
           # Get the base branch (usually main)
           BASE_BRANCH="${{ github.event.pull_request.base.ref }}"
           echo "Checking coverage for new code against base branch: ${BASE_BRANCH}"
           # Run diff-cover to check only new/changed code
-          # Fail if new code coverage is below 80%
+          # This is a "soft" check - warnings only, doesn't block merges
           # Note: diff-cover only needs --extra dev (tests need all extras for integration testing)
           if uv run --extra dev diff-cover coverage.xml \
             --compare-branch=origin/${BASE_BRANCH} \
             --fail-under=80 \
             --markdown-report diff-coverage-report.md; then
-            echo "::notice::New code coverage check passed (80% threshold)"
+            echo "::notice::✅ New code coverage check passed (80% threshold)"
+            echo "coverage_passed=true" >> $GITHUB_OUTPUT
           else
-            echo "::error::New code coverage is below 80% threshold"
-            exit 1
+            echo "::warning::⚠️ New code coverage is below 80% threshold"
+            echo "::warning::This is a warning only and will not block merging, but please consider improving test coverage."
+            echo "coverage_passed=false" >> $GITHUB_OUTPUT
           fi
-        continue-on-error: false
 
       - name: Upload diff coverage report
-        if: github.event_name == 'pull_request' && always()
+        if: always()
         uses: actions/github-script@v7
+        env:
+          COVERAGE_PASSED: ${{ steps.coverage_check.outputs.coverage_passed }}
         with:
           script: |
             const fs = require('fs');
             const reportPath = 'diff-coverage-report.md';
             if (fs.existsSync(reportPath)) {
               const report = fs.readFileSync(reportPath, 'utf8');
+              const coveragePassed = process.env.COVERAGE_PASSED === 'true';
+              const statusEmoji = coveragePassed ? '✅' : '⚠️';
+              const statusText = coveragePassed ? 'PASSED' : 'BELOW THRESHOLD (Warning)';
+              const header = coveragePassed 
+                ? `## ${statusEmoji} Diff Coverage Report - ${statusText}\n\n`
+                : `## ${statusEmoji} Diff Coverage Report - ${statusText}\n\n**Note:** This is a warning only and does not block merging.\n\n`;
               fs.appendFileSync(
                 process.env.GITHUB_STEP_SUMMARY,
-                `## Diff Coverage Report\n\n${report}\n`
+                `${header}${report}\n`
               );
             }
 
-      - name: Upload coverage reports to Codecov
-        uses: codecov/codecov-action@v5
-        with:
-          token: ${{ secrets.CODECOV_TOKEN }}
-          slug: Astera-org/simplexity
-          verbose: true
-          files: ./coverage.xml
-          fail_ci_if_error: false
+      - name: Always succeed (coverage is advisory only)
+        if: always()
+        run: |
+          # This step ensures the job always shows as "success" even if coverage is low
+          # Coverage warnings are still visible in the PR via annotations and step summary
+          echo "Coverage check completed (advisory only - does not block merges)"

From 27afe71890ad522544baf1fb0d67940184f14fe3 Mon Sep 17 00:00:00 2001
From: Eric Alt <13019253+ealt@users.noreply.github.com>
Date: Tue, 3 Mar 2026 17:17:53 -0800
Subject: [PATCH 26/31] Fix pylint issues

---
 tests/end_to_end/test_managed_run_mlflow_defaults.py | 7 +++++--
 1 file changed, 5 insertions(+), 2 deletions(-)

diff --git a/tests/end_to_end/test_managed_run_mlflow_defaults.py b/tests/end_to_end/test_managed_run_mlflow_defaults.py
index 57210bbf..26908356 100644
--- a/tests/end_to_end/test_managed_run_mlflow_defaults.py
+++ b/tests/end_to_end/test_managed_run_mlflow_defaults.py
@@ -1,3 +1,5 @@
+"""End-to-end tests for managed run with MLflow defaults."""
+
 from pathlib import Path
 
 from hydra import compose, initialize_config_dir
@@ -9,6 +11,7 @@
 
 
 def test_managed_run_loads_mlflow_defaults(setup_dir: Path) -> None:
+    """Verify that managed_run correctly loads and merges MLflow default configs."""
     tracking_uri = f"sqlite:///{setup_dir.resolve()}/mlflow.db"
     with initialize_config_dir(config_dir=CONFIG_DIR):
         cfg = compose(
@@ -19,10 +22,10 @@ def test_managed_run_loads_mlflow_defaults(setup_dir: Path) -> None:
     captured: dict[str, DictConfig] = {}
 
     @simplexity.managed_run(strict=False)
-    def run(cfg: DictConfig, components: simplexity.Components) -> None:
+    def run(cfg: DictConfig, _components: simplexity.Components) -> None:
         captured["cfg"] = cfg
 
-    run(cfg)
+    run(cfg)  # pylint: disable=no-value-for-parameter
 
     actual = captured["cfg"]
     experiment_id = OmegaConf.select(actual, "load_source.experiment_id")

From 6719d980562819e7d998547e71da03db0b448aef Mon Sep 17 00:00:00 2001
From: Eric Alt <13019253+ealt@users.noreply.github.com>
Date: Tue, 3 Mar 2026 17:42:52 -0800
Subject: [PATCH 27/31] Add type assertion for actual_filtered in MLflow
 defaults test

---
 tests/end_to_end/test_managed_run_mlflow_defaults.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/tests/end_to_end/test_managed_run_mlflow_defaults.py b/tests/end_to_end/test_managed_run_mlflow_defaults.py
index 26908356..49aea87f 100644
--- a/tests/end_to_end/test_managed_run_mlflow_defaults.py
+++ b/tests/end_to_end/test_managed_run_mlflow_defaults.py
@@ -43,6 +43,7 @@ def run(cfg: DictConfig, _components: simplexity.Components) -> None:
 
     actual_container = OmegaConf.to_container(actual, resolve=True)
     actual_filtered = OmegaConf.create(actual_container)
+    assert isinstance(actual_filtered, DictConfig)
     for key in ("device", "seed", "tags"):
         if key in actual_filtered:
             del actual_filtered[key]

From 622806839c3d990184eda9bb710530c91a1b01fd Mon Sep 17 00:00:00 2001
From: Eric Alt <13019253+ealt@users.noreply.github.com>
Date: Tue, 3 Mar 2026 17:46:07 -0800
Subject: [PATCH 28/31] Update GitHub Actions workflow to use upload-artifact
 action for coverage report

---
 .github/workflows/simplexity.yaml | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/.github/workflows/simplexity.yaml b/.github/workflows/simplexity.yaml
index b878d39e..386d3e8b 100644
--- a/.github/workflows/simplexity.yaml
+++ b/.github/workflows/simplexity.yaml
@@ -108,7 +108,7 @@ jobs:
 
       - name: Upload diff coverage report
         if: github.event_name == 'pull_request' && always()
-        uses: actions/github-script@v7
+        uses: actions/upload-artifact@v4
         with:
           name: coverage-xml
           path: coverage.xml

From 904e910c7ff78051b2a2becdb3ceb98b939d50fb Mon Sep 17 00:00:00 2001
From: Eric Alt <13019253+ealt@users.noreply.github.com>
Date: Tue, 3 Mar 2026 17:46:30 -0800
Subject: [PATCH 29/31] Refactor argument name in managed_run MLflow test for
 clarity

---
 tests/end_to_end/test_managed_run_mlflow_defaults.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tests/end_to_end/test_managed_run_mlflow_defaults.py b/tests/end_to_end/test_managed_run_mlflow_defaults.py
index 49aea87f..d0ed3170 100644
--- a/tests/end_to_end/test_managed_run_mlflow_defaults.py
+++ b/tests/end_to_end/test_managed_run_mlflow_defaults.py
@@ -22,7 +22,7 @@ def test_managed_run_loads_mlflow_defaults(setup_dir: Path) -> None:
     captured: dict[str, DictConfig] = {}
 
     @simplexity.managed_run(strict=False)
-    def run(cfg: DictConfig, _components: simplexity.Components) -> None:
+    def run(cfg: DictConfig, components: simplexity.Components) -> None:  # pylint: disable=unused-argument
         captured["cfg"] = cfg
 
     run(cfg)  # pylint: disable=no-value-for-parameter

From bb305073491e06089600b2220b6936628c2df974 Mon Sep 17 00:00:00 2001
From: Eric Alt <13019253+ealt@users.noreply.github.com>
Date: Tue, 3 Mar 2026 18:01:03 -0800
Subject: [PATCH 30/31] Fix bugbot issues

---
 simplexity/structured_configs/mlflow.py       | 14 +++-
 .../structured_configs/mlflow_defaults.py     | 12 ++-
 .../test_load_mlflow_defaults.py              | 73 +++++++++++++++++++
 3 files changed, 92 insertions(+), 7 deletions(-)

diff --git a/simplexity/structured_configs/mlflow.py b/simplexity/structured_configs/mlflow.py
index e478cabf..0273d262 100644
--- a/simplexity/structured_configs/mlflow.py
+++ b/simplexity/structured_configs/mlflow.py
@@ -56,7 +56,7 @@ def validate_mlflow_config(cfg: DictConfig) -> None:
 
 
 @dynamic_resolve
-def resolve_mlflow_config(cfg: DictConfig) -> None:
+def resolve_mlflow_config(cfg: DictConfig, *, create_if_missing: bool = True) -> None:
     """Resolve a MLFlowConfig."""
     # Resolve registry uri
     tracking_uri: str | None = cfg.get("tracking_uri")
@@ -73,7 +73,9 @@ def resolve_mlflow_config(cfg: DictConfig) -> None:
     # Resolve experiment id and name
     experiment_id: str | None = cfg.get("experiment_id")
     experiment_name: str | None = cfg.get("experiment_name")
-    experiment = get_experiment(experiment_id=experiment_id, experiment_name=experiment_name, client=client)
+    experiment = get_experiment(
+        experiment_id=experiment_id, experiment_name=experiment_name, client=client, create_if_missing=create_if_missing
+    )
     if experiment is None:
         raise ValueError(f"Experiment not found for id: {experiment_id} and name: {experiment_name}")
     cfg.experiment_id = experiment.experiment_id
@@ -82,7 +84,13 @@ def resolve_mlflow_config(cfg: DictConfig) -> None:
     # Resolve run id and name
     run_id: str | None = cfg.get("run_id")
     run_name: str | None = cfg.get("run_name")
-    run = get_run(run_id=run_id, run_name=run_name, experiment_id=cfg.experiment_id, client=client)
+    run = get_run(
+        run_id=run_id,
+        run_name=run_name,
+        experiment_id=cfg.experiment_id,
+        client=client,
+        create_if_missing=create_if_missing,
+    )
     if run is None:
         raise ValueError(f"Run not found for id: {run_id} and name: {run_name}")
     cfg.run_id = run.info.run_id
diff --git a/simplexity/structured_configs/mlflow_defaults.py b/simplexity/structured_configs/mlflow_defaults.py
index 772f3ad4..035b132c 100644
--- a/simplexity/structured_configs/mlflow_defaults.py
+++ b/simplexity/structured_configs/mlflow_defaults.py
@@ -22,6 +22,7 @@
 
 import yaml
 from mlflow import MlflowClient
+from mlflow.exceptions import MlflowException, RestException
 from omegaconf import DictConfig, ListConfig, OmegaConf
 from omegaconf.errors import MissingMandatoryValue
 
@@ -138,7 +139,7 @@ def _get_target_config(cfg: DictConfig, parsed_entry: _ParsedEntry) -> Any | Non
         return None
 
     try:
-        resolve_mlflow_config(target_node)
+        resolve_mlflow_config(target_node, create_if_missing=False)
     except (ValueError, RuntimeError) as e:
         SIMPLEXITY_LOGGER.warning("Error resolving MLflow config: %s", e)
         return None
@@ -153,7 +154,7 @@ def _get_target_config(cfg: DictConfig, parsed_entry: _ParsedEntry) -> Any | Non
     with tempfile.TemporaryDirectory() as tmp_dir:
         try:
             local_path = client.download_artifacts(run_id=run_id, path=artifact_path, dst_path=tmp_dir)
-        except (OSError, FileNotFoundError) as e:
+        except (MlflowException, RestException, OSError, FileNotFoundError) as e:
             SIMPLEXITY_LOGGER.warning("Failed to download artifact from MLflow '%s': %s", parsed_entry.target, e)
             return None
 
@@ -227,7 +228,7 @@ def _resolve_mlflow_configs_recursive(cfg: DictConfig) -> None:
         try:
             validate_mlflow_config(cfg)
             # If validation passes, this looks like an MLflow config, try to resolve it
-            resolve_mlflow_config(cfg)
+            resolve_mlflow_config(cfg, create_if_missing=False)
         except (ConfigValidationError, ValueError, RuntimeError):
             # Not an MLflow config or can't be resolved, continue to check nested configs
             pass
@@ -336,10 +337,13 @@ def load_mlflow_defaults(cfg: DictConfig) -> DictConfig:
     See Also:
         LOAD_SUBCONFIGS.md for the full specification and examples.
     """
-    mlflow_defaults: ListConfig | None = cfg.get("mlflow_defaults")
+    mlflow_defaults: ListConfig | str | None = cfg.get("mlflow_defaults")
     if mlflow_defaults is None:
         return cfg
 
+    if isinstance(mlflow_defaults, str):
+        mlflow_defaults = [mlflow_defaults]
+
     # Create a copy to avoid mutating the input config
     mlflow_defaults_copy = cast(ListConfig, OmegaConf.create(list(mlflow_defaults)))
     if "_self_" not in mlflow_defaults_copy:
diff --git a/tests/run_management/test_load_mlflow_defaults.py b/tests/run_management/test_load_mlflow_defaults.py
index ad44511d..e49b119e 100644
--- a/tests/run_management/test_load_mlflow_defaults.py
+++ b/tests/run_management/test_load_mlflow_defaults.py
@@ -5,6 +5,7 @@
 
 import pytest
 from mlflow.client import MlflowClient
+from mlflow.exceptions import MlflowException
 from omegaconf import DictConfig, OmegaConf
 
 from simplexity.exceptions import ConfigValidationError
@@ -820,3 +821,75 @@ def test_option_hash_empty_select(base_cfg: DictConfig, mock_download: MagicMock
     assert OmegaConf.select(loaded_cfg, "root_key") == "root_value"
     assert OmegaConf.select(loaded_cfg, "nested.nested_key") == "nested_value"
     assert mock_download.call_args.kwargs["path"] == "artifact.yaml"
+
+
+def test_mlflow_exception_caught_by_download(base_cfg: DictConfig, mock_download: MagicMock):
+    """Test that MlflowException from download_artifacts is caught."""
+    mock_download.side_effect = MlflowException("Artifact not found on remote store")
+
+    cfg = OmegaConf.merge(
+        base_cfg,
+        {
+            "mlflow_defaults": ["previous_run"],
+        },
+    )
+    with pytest.raises(ValueError, match="Target config not found for entry"):
+        load_mlflow_defaults(cfg)
+
+
+def test_mlflow_exception_caught_optional(base_cfg: DictConfig, mock_download: MagicMock):
+    """Test that MlflowException from download_artifacts is caught with optional flag."""
+    mock_download.side_effect = MlflowException("Artifact not found on remote store")
+
+    cfg = OmegaConf.merge(
+        base_cfg,
+        {
+            "mlflow_defaults": ["optional previous_run"],
+        },
+    )
+    loaded_cfg: DictConfig = load_mlflow_defaults(cfg)
+    assert OmegaConf.select(loaded_cfg, "other_section.foo") == "bar"
+
+
+def test_string_mlflow_defaults(base_cfg: DictConfig, mock_download: MagicMock, tmp_path: Path):
+    """Test that a string mlflow_defaults value is handled correctly."""
+    artifact_path = tmp_path / "config.yaml"
+    artifact_path.write_text("loaded_key: loaded_value\n")
+    mock_download.return_value = str(artifact_path)
+
+    cfg = OmegaConf.merge(
+        base_cfg,
+        {
+            "mlflow_defaults": "previous_run",
+        },
+    )
+
+    loaded_cfg: DictConfig = load_mlflow_defaults(cfg)
+    assert loaded_cfg.get("loaded_key") == "loaded_value"
+    assert OmegaConf.select(loaded_cfg, "other_section.foo") == "bar"
+
+
+def test_resolve_mlflow_config_called_with_create_if_missing_false(
+    base_cfg: DictConfig,
+    mock_download: MagicMock,
+    tmp_path: Path,
+    mocker,
+):
+    """Test that resolve_mlflow_config is called with create_if_missing=False."""
+    artifact_path = tmp_path / "config.yaml"
+    artifact_path.write_text("loaded_key: loaded_value\n")
+    mock_download.return_value = str(artifact_path)
+
+    mock_resolve = mocker.patch("simplexity.structured_configs.mlflow_defaults.resolve_mlflow_config")
+
+    cfg = OmegaConf.merge(
+        base_cfg,
+        {
+            "mlflow_defaults": ["previous_run"],
+        },
+    )
+
+    load_mlflow_defaults(cfg)
+    mock_resolve.assert_called()
+    for c in mock_resolve.call_args_list:
+        assert c.kwargs.get("create_if_missing") is False

From 9588d404d80eb841f3c1cc5e89367fb6d8c779e7 Mon Sep 17 00:00:00 2001
From: Eric Alt <13019253+ealt@users.noreply.github.com>
Date: Tue, 3 Mar 2026 18:04:20 -0800
Subject: [PATCH 31/31] Refactor load_mlflow_defaults function to simplify
 handling of mlflow_defaults input and improve type consistency

---
 simplexity/structured_configs/mlflow_defaults.py | 7 +++----
 1 file changed, 3 insertions(+), 4 deletions(-)

diff --git a/simplexity/structured_configs/mlflow_defaults.py b/simplexity/structured_configs/mlflow_defaults.py
index 035b132c..bcde76b1 100644
--- a/simplexity/structured_configs/mlflow_defaults.py
+++ b/simplexity/structured_configs/mlflow_defaults.py
@@ -337,15 +337,14 @@ def load_mlflow_defaults(cfg: DictConfig) -> DictConfig:
     See Also:
         LOAD_SUBCONFIGS.md for the full specification and examples.
     """
-    mlflow_defaults: ListConfig | str | None = cfg.get("mlflow_defaults")
+    mlflow_defaults = cfg.get("mlflow_defaults")
     if mlflow_defaults is None:
         return cfg
 
-    if isinstance(mlflow_defaults, str):
-        mlflow_defaults = [mlflow_defaults]
+    mlflow_defaults_list = [mlflow_defaults] if isinstance(mlflow_defaults, str) else list(mlflow_defaults)
 
     # Create a copy to avoid mutating the input config
-    mlflow_defaults_copy = cast(ListConfig, OmegaConf.create(list(mlflow_defaults)))
+    mlflow_defaults_copy = cast(ListConfig, OmegaConf.create(mlflow_defaults_list))
     if "_self_" not in mlflow_defaults_copy:
         mlflow_defaults_copy.append("_self_")