Include component defaults in Optuna search space

src/lenskit/cli/__init__.py

@@ -43,7 +43,14 @@ def main():
 
 @click.group("lenskit", invoke_without_command=True)
 @click.help_option("-h", "--help")
-@click.option("-v", "--verbose", "verbosity", count=True, help="Enable verbose logging output.")
+@click.option(
+    "-v",
+    "--verbose",
+    "verbosity",
+    count=True,
+    envvar="LK_VERBOSE",
+    help="Enable verbose logging output.",
+)
 @click.option("--log-file", type=Path, help="Save log messages to FILE")
 @click.option(
     "--log-file-verbose",

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__)
 
@@ -58,7 +58,6 @@ def run(self) -> OptunaTuneResults:
             pruner=CompositePruner(),
             direction=StudyDirection.MINIMIZE if self.mode == "min" else StudyDirection.MAXIMIZE,
         )
-        # TODO: add parallelism support
 
         self.log = _log.bind(pipeline=self.pipeline.meta.name, dataset=self.data.name)
         self.log.info(
@@ -83,6 +82,7 @@ def run(self) -> OptunaTuneResults:
         return OptunaTuneResults(spec=self.spec, study=study, iterative=self.iterative, task=task)
 
     def _run_study(self, study: Study):
+        self._enqueue_defaults(study)
         npts = self.spec.search.num_search_points()
         task = Task.current()
         with item_progress("Search trials", total=npts) as pb:
@@ -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)
@@ -252,11 +253,36 @@ def _run_simple_trial(self, study: Study, trial: Trial, config):
 
         study.tell(trial, results[self.metric], state=TrialState.COMPLETE)
 
+    def _enqueue_defaults(self, study: Study):
+        """
+        Enqueue the component's default hyperparameters into the study.
+        """
+        self.log.debug(
+            "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[self.spec.component_name], comp.dump_config())
+        self.log.info("enqueueing default point", config=config)
+        study.enqueue_trial(
+            config,
+            user_attrs={
+                "memo": "initial default configuration",
+                "component": self.spec.component_name,
+            },
+        )
+
     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
 
 
@@ -269,28 +295,97 @@ 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],
+    *,
+    out: dict[str, JsonValue] | None = None,
+    prefix: str = "",
+) -> dict[str, JsonValue]:
+    # 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):
+            _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:
+            # FIXME this is an ugly hack for single / multiple configs
+            out[k] = config
+    return out
+
+
 @dataclass
 class OptunaTuneResults(TuneResults):
     spec: TuningSpec