tuning: fix pow2 in default setup

Author: mdekstrand

src/lenskit/tuning/_optuna.py

@@ -26,7 +26,7 @@
 from typing_extensions import override
 
 from lenskit.data import unflatten_dict
-from lenskit.logging import Task, get_logger, item_progress
+from lenskit.logging import Task, get_logger, item_progress, trace
 from lenskit.logging.tasks import add_context_task
 from lenskit.parallel import NestedPool
 from lenskit.pipeline import Pipeline, PipelineBuilder
@@ -38,7 +38,7 @@
 from ._base import BasePipelineTuner, TuneResults
 from ._measure import measure_pipeline
 from ._stopping import PlateauStopRule
-from .spec import SearchSpace, TuningSpec
+from .spec import SearchParam, SearchSpace, TuningSpec
 
 _log = get_logger(__name__)
 
@@ -126,6 +126,7 @@ def _run_trial(self, study: Study, *, nested_pool: bool = False):
             tags=["tune", "trial"],
             reset_hwm=True,
         ):
+            self.log.debug("beginning search", config=config)
             try:
                 if self.iterative:
                     self._run_iter_trial(study, trial, config)
@@ -260,13 +261,15 @@ def _enqueue_defaults(self, study: Study):
             "instantiating pipeline to extract defaults", component=self.spec.component_name
         )
         pipe = Pipeline.from_config(self.pipeline)
+        assert self.spec.component_name is not None
         comp = pipe.component(self.spec.component_name)
         assert comp is not None
         if not isinstance(comp, Component):
             self.log.warn("component is not pipeline", component=self.spec.component_name)
             return
 
-        config = _extract_defaults(self.spec.space, comp.dump_config())
+        config = _extract_defaults(self.spec.space[self.spec.component_name], comp.dump_config())
+        self.log.info("enqueueing default point", config=config)
         study.enqueue_trial(
             config,
             user_attrs={
@@ -279,7 +282,7 @@ def _ask_config(self, trial: Trial):
         # we have exactly one
         for space in self.spec.space.values():
             cfg = _ask_space(trial, space)
-            self.log.info("obtained search point", cfg=cfg)
+            self.log.info("sampled search point", config=cfg)
             return cfg
 
 
@@ -292,37 +295,94 @@ def _ask_space(trial: Trial, space: SearchSpace, *, prefix: str = ""):
             assert isinstance(spec.min, int)
             assert isinstance(spec.max, int)
             out[name] = trial.suggest_int(prefix + name, spec.min, spec.max)
+            trace(
+                _log, "sampled int", name=prefix + name, min=spec.min, max=spec.max, value=out[name]
+            )
         elif spec.type == "int" and spec.scale == "log":
             assert isinstance(spec.min, int)
             assert isinstance(spec.max, int)
             out[name] = trial.suggest_int(prefix + name, spec.min, spec.max, log=True)
+            trace(
+                _log,
+                "sampled int/log2",
+                name=prefix + name,
+                min=spec.min,
+                max=spec.max,
+                value=out[name],
+            )
         elif spec.type == "int" and spec.scale == "pow2":
             min = int(math.log2(spec.min))
             max = int(math.log2(spec.max))
-            out[name] = 2 ** trial.suggest_int(prefix + name, min, max)
+            p = trial.suggest_int(prefix + name, min, max)
+            out[name] = 2**p
+            trace(
+                _log,
+                "sampled int/pow2",
+                name=prefix + name,
+                min=min,
+                max=max,
+                power=p,
+                value=out[name],
+            )
         elif spec.type == "float" and spec.scale == "uniform":
             out[name] = trial.suggest_float(prefix + name, spec.min, spec.max)
+            trace(
+                _log,
+                "sampled float",
+                name=prefix + name,
+                min=spec.min,
+                max=spec.max,
+                value=out[name],
+            )
         elif spec.type == "float" and spec.scale == "log":
             out[name] = trial.suggest_float(prefix + name, spec.min, spec.max, log=True)
+            trace(
+                _log,
+                "sampled float/log",
+                name=prefix + name,
+                min=spec.min,
+                max=spec.max,
+                value=out[name],
+            )
         elif spec.type == "bool":
             out[name] = trial.suggest_categorical(prefix + name, [False, True])
+            trace(_log, "sampled bool", name=prefix + name, value=out[name])
         elif spec.type == "choice":
             out[name] = trial.suggest_categorical(prefix + name, spec.choices)
+            trace(_log, "sampled choice", name=prefix + name, value=out[name])
         else:  # pragma: nocover
             raise ValueError(f"unsupported configuration {space}")
 
     return out
 
 
 def _extract_defaults(
-    space: dict[str, SearchSpace], config: dict[str, JsonValue]
+    space: dict[str, SearchSpace],
+    config: dict[str, JsonValue],
+    *,
+    out: dict[str, JsonValue] | None = None,
+    prefix: str = "",
 ) -> dict[str, JsonValue]:
-    out = {}
+    # FIXME this whole function is ugly but seems to work
+    if out is None:
+        out = {}
     for k, v in space.items():
         if isinstance(v, dict):
-            out[k] = _extract_defaults(v, config[k])  # type: ignore
+            _extract_defaults(v, config[k], out=out, prefix=f"{prefix}{k}.")  # type: ignore
+        elif isinstance(config, dict):
+            assert isinstance(v, SearchParam)
+            try:
+                out[k] = config[k]
+            except KeyError as e:
+                if k.endswith("_exp"):
+                    out[k] = int(math.log2(config[k[:-4]]))  # type: ignore
+                else:
+                    raise e
+            if v.scale == "pow2":
+                out[k] = int(math.log2(out[k]))  # type: ignore
         else:
-            out[k] = config[k]
+            # FIXME this is an ugly hack for single / multiple configs
+            out[k] = config
     return out