Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
27 changes: 26 additions & 1 deletion flaml/tune/searcher/search_thread.py
Original file line number Diff line number Diff line change
Expand Up @@ -25,6 +25,31 @@
logger = logging.getLogger(__name__)


def _recursive_dict_update(target: Dict, source: Dict) -> None:
"""Recursively update target dictionary with source dictionary.

Unlike dict.update(), this function merges nested dictionaries instead of
replacing them entirely. This is crucial for configurations with nested
structures (e.g., XGBoost params).

Args:
target: The dictionary to be updated (modified in place).
source: The dictionary containing values to merge into target.

Example:
>>> target = {'params': {'eta': 0.1, 'max_depth': 3}}
>>> source = {'params': {'verbosity': 0}}
>>> _recursive_dict_update(target, source)
>>> target
{'params': {'eta': 0.1, 'max_depth': 3, 'verbosity': 0}}
"""
for key, value in source.items():
if isinstance(value, dict) and key in target and isinstance(target[key], dict):
_recursive_dict_update(target[key], value)
else:
target[key] = value


class SearchThread:
"""Class of global or local search thread."""

Expand Down Expand Up @@ -65,7 +90,7 @@ def suggest(self, trial_id: str) -> Optional[Dict]:
try:
config = self._search_alg.suggest(trial_id)
if isinstance(self._search_alg._space, dict):
config.update(self._const)
_recursive_dict_update(config, self._const)
else:
# define by run
config, self.space = unflatten_hierarchical(config, self._space)
Expand Down
99 changes: 99 additions & 0 deletions test/tune/test_search_thread.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,99 @@
"""Tests for SearchThread nested dictionary update fix."""

import pytest

from flaml.tune.searcher.search_thread import _recursive_dict_update


def test_recursive_dict_update_simple():
"""Test simple non-nested dictionary update."""
target = {"a": 1, "b": 2}
source = {"c": 3}
_recursive_dict_update(target, source)
assert target == {"a": 1, "b": 2, "c": 3}


def test_recursive_dict_update_override():
"""Test that source values override target values for non-dict values."""
target = {"a": 1, "b": 2}
source = {"b": 3}
_recursive_dict_update(target, source)
assert target == {"a": 1, "b": 3}


def test_recursive_dict_update_nested():
"""Test nested dictionary merge (the main use case for XGBoost params)."""
target = {
"num_boost_round": 10,
"params": {
"max_depth": 12,
"eta": 0.020168455186106736,
"min_child_weight": 1.4504723523894132,
"scale_pos_weight": 3.794258636185337,
"gamma": 0.4985070123025904,
},
}
source = {
"params": {
"verbosity": 3,
"booster": "gbtree",
"eval_metric": "auc",
"tree_method": "hist",
"objective": "binary:logistic",
}
}
_recursive_dict_update(target, source)

# Check that sampled params are preserved
assert target["params"]["max_depth"] == 12
assert target["params"]["eta"] == 0.020168455186106736
assert target["params"]["min_child_weight"] == 1.4504723523894132
assert target["params"]["scale_pos_weight"] == 3.794258636185337
assert target["params"]["gamma"] == 0.4985070123025904

# Check that const params are added
assert target["params"]["verbosity"] == 3
assert target["params"]["booster"] == "gbtree"
assert target["params"]["eval_metric"] == "auc"
assert target["params"]["tree_method"] == "hist"
assert target["params"]["objective"] == "binary:logistic"

# Check top-level param is preserved
assert target["num_boost_round"] == 10


def test_recursive_dict_update_deeply_nested():
"""Test deeply nested dictionary merge."""
target = {"a": {"b": {"c": 1, "d": 2}}}
source = {"a": {"b": {"e": 3}}}
_recursive_dict_update(target, source)
assert target == {"a": {"b": {"c": 1, "d": 2, "e": 3}}}


def test_recursive_dict_update_mixed_types():
"""Test that non-dict values in source replace dict values in target."""
target = {"a": {"b": 1}}
source = {"a": 2}
_recursive_dict_update(target, source)
assert target == {"a": 2}


def test_recursive_dict_update_empty_dicts():
"""Test with empty dictionaries."""
target = {}
source = {"a": 1}
_recursive_dict_update(target, source)
assert target == {"a": 1}

target = {"a": 1}
source = {}
_recursive_dict_update(target, source)
assert target == {"a": 1}


def test_recursive_dict_update_none_values():
"""Test that None values are properly handled."""
target = {"a": 1, "b": None}
source = {"b": 2, "c": None}
_recursive_dict_update(target, source)
assert target == {"a": 1, "b": 2, "c": None}
Loading