Knowledge Base + Multi-Task Warm-starting

.pre-commit-config.yaml

@@ -1,5 +1,9 @@
 ---
 repos:
+  - repo: /~https://github.com/pre-commit/pre-commit-hooks
+    rev: v4.3.0  # Use the ref you want to point at
+    hooks:
+      - id: check-merge-conflict
   - repo: /~https://github.com/python/black
     rev: 22.6.0
     hooks:

src/orion/algo/base.py

@@ -18,8 +18,10 @@
 import inspect
 import logging
 from abc import abstractmethod
+from typing import Any
 
 from orion.algo.registry import Registry
+from orion.algo.space import Space
 from orion.core.utils import GenericFactory
 from orion.core.worker.trial import Trial
 
@@ -101,7 +103,9 @@ def observe(self, points, results):
     requires_shape = None
     requires_dist = None
 
-    def __init__(self, space, **kwargs):
+    max_trials: int | None = None
+
+    def __init__(self, space: Space, **kwargs):
         log.debug(
             "Creating Algorithm object of %s type with parameters:\n%s",
             type(self).__name__,
@@ -142,18 +146,19 @@ def state_dict(self):
         """Return a state dict that can be used to reset the state of the algorithm."""
         return {"registry": self.registry.state_dict}
 
-    def set_state(self, state_dict):
+    def set_state(self, state_dict: dict):
         """Reset the state of the algorithm based on the given state_dict
 
         :param state_dict: Dictionary representing state of an algorithm
         """
         self.registry.set_state(state_dict["registry"])
 
-    def get_id(self, trial, ignore_fidelity=False, ignore_parent=False):
+    def get_id(
+        self, trial: Trial, ignore_fidelity: bool = False, ignore_parent: bool = False
+    ) -> str:
         """Return unique hash for a trials based on params
 
-        The trial is assumed to be in the transformed space if the algorithm is working in a
-        transformed space.
+        The trial is assumed to be in the optimization space of the algorithm.
 
         Parameters
         ----------
@@ -170,13 +175,12 @@ def get_id(self, trial, ignore_fidelity=False, ignore_parent=False):
         return trial.compute_trial_hash(
             trial,
             ignore_fidelity=ignore_fidelity,
-            ignore_experiment=True,
             ignore_lie=True,
             ignore_parent=ignore_parent,
         )
 
     @property
-    def fidelity_index(self):
+    def fidelity_index(self) -> str | None:
         """Returns the name of the first fidelity dimension if there is one, otherwise `None`."""
         fidelity_dims = [dim for dim in self.space.values() if dim.type == "fidelity"]
         if fidelity_dims:
@@ -209,7 +213,7 @@ def suggest(self, num: int) -> list[Trial]:
         has suggested/observed, and for the auto-generated unit-tests to pass.
         """
 
-    def observe(self, trials):
+    def observe(self, trials: list[Trial]) -> None:
         """Observe the `results` of the evaluation of the `trials` in the
         process defined in user's script.
 
@@ -223,7 +227,7 @@ def observe(self, trials):
             if not self.has_observed(trial):
                 self.register(trial)
 
-    def register(self, trial):
+    def register(self, trial: Trial) -> None:
         """Save the trial as one suggested or observed by the algorithm.
 
         Parameters
@@ -234,16 +238,16 @@ def register(self, trial):
         self.registry.register(trial)
 
     @property
-    def n_suggested(self):
+    def n_suggested(self) -> int:
         """Number of trials suggested by the algorithm"""
         return len(self.registry)
 
     @property
-    def n_observed(self):
+    def n_observed(self) -> int:
         """Number of completed trials observed by the algorithm."""
         return sum(self.has_observed(trial) for trial in self.registry)
 
-    def has_suggested(self, trial):
+    def has_suggested(self, trial: Trial) -> bool:
         """Whether the algorithm has suggested a given point.
 
         Parameters
@@ -259,7 +263,7 @@ def has_suggested(self, trial):
         """
         return self.registry.has_suggested(trial)
 
-    def has_observed(self, trial):
+    def has_observed(self, trial: Trial) -> bool:
         """Whether the algorithm has observed a given point objective.
 
         This only counts observed completed trials.
@@ -312,15 +316,11 @@ def has_completed_max_trials(self) -> bool:
         """Returns True if the algorithm has a `max_trials` attribute, and has completed more trials
         than its value.
         """
-        if not hasattr(self, "max_trials"):
-            return False
-        max_trials = getattr(self, "max_trials")
-        if max_trials is None:
+        if self.max_trials is None:
             return False
 
         fidelity_index = self.fidelity_index
         max_fidelity_value = None
-
         # When a fidelity dimension is present, we only count trials that have the maximum value.
         if fidelity_index is not None:
             _, max_fidelity_value = self.space[fidelity_index].interval()
@@ -333,10 +333,11 @@ def _is_completed(trial: Trial) -> bool:
                 and trial.params[fidelity_index] >= max_fidelity_value
             )
 
-        return sum(map(_is_completed, self.registry)) >= max_trials
+        return sum(map(_is_completed, self.registry)) >= self.max_trials
 
-    def score(self, trial):  # pylint:disable=no-self-use,unused-argument
-        """Allow algorithm to evaluate `point` based on a prediction about
+    # pylint:disable=no-self-use,unused-argument
+    def score(self, trial: Trial) -> float:
+        """Allow algorithm to evaluate `trial` based on a prediction about
         this parameter set's performance.
 
         By default, return the same score any parameter (no preference).
@@ -353,7 +354,8 @@ def score(self, trial):  # pylint:disable=no-self-use,unused-argument
         """
         return 0
 
-    def judge(self, trial, measurements):  # pylint:disable=no-self-use,unused-argument
+    # pylint:disable=no-self-use,unused-argument
+    def judge(self, trial: Trial, measurements: Any) -> dict | None:
         """Inform an algorithm about online `measurements` of a running trial.
 
         This method is to be used as a callback in a client-server communication
@@ -381,7 +383,7 @@ def judge(self, trial, measurements):  # pylint:disable=no-self-use,unused-argum
         """
         return None
 
-    def should_suspend(self, trial):
+    def should_suspend(self, trial: Trial) -> bool:
         """Allow algorithm to decide whether a particular running trial is still
         worth to complete its evaluation, based on information provided by the
         `judge` method.
@@ -390,10 +392,12 @@ def should_suspend(self, trial):
         return False
 
     @property
-    def configuration(self):
+    def configuration(self) -> dict[str, Any]:
         """Return tunable elements of this algorithm in a dictionary form
         appropriate for saving.
 
+        By default, returns a dictionary containing the attributes of `self` which are also
+        constructor arguments.
         """
         dict_form = dict()
         for attrname in self._param_names:
@@ -404,14 +408,13 @@ def configuration(self):
         return {self.__class__.__name__.lower(): dict_form}
 
     @property
-    def space(self):
+    def space(self) -> Space:
         """Domain of problem associated with this algorithm's instance."""
         return self._space
 
-    @space.setter
-    def space(self, space):
-        """Set space."""
-        self._space = space
+    @property
+    def unwrapped(self):
+        return self
 
 
 algo_factory = GenericFactory(BaseAlgorithm)

src/orion/algo/gridsearch.py

@@ -133,8 +133,11 @@ def __init__(
             if not isinstance(n_values, dict)
             else n_values
         )
+        max_trials = 10_000 if self.max_trials is None else self.max_trials
         self.grid = self.build_grid(
-            self.space, n_values_dict, getattr(self, "max_trials", 10000)
+            self.space,
+            n_values_dict,
+            max_trials=max_trials,
         )
         self.index = 0
 

src/orion/algo/pbt/pb2.py

@@ -17,6 +17,7 @@
 from orion.algo.pbt.pbt import PBT
 from orion.core.utils.flatten import flatten
 from orion.core.utils.random_state import RandomState, control_randomness
+from orion.core.worker.transformer import ReshapedSpace, TransformedSpace
 from orion.core.worker.trial import Trial
 
 logger = logging.getLogger(__name__)
@@ -172,6 +173,7 @@ def _generate_offspring(self, trial):
             ]
 
             new_trial = trial_to_branch.branch(params=new_params)
+            assert isinstance(self.space, (TransformedSpace, ReshapedSpace))
             new_trial = self.space.transform(self.space.reverse(new_trial))
 
             logger.debug("Attempt %s - Creating new trial %s", attempts, new_trial)

src/orion/algo/pbt/pbt.py

@@ -162,6 +162,7 @@ def __init__(
         explore: dict | None = None,
         fork_timeout: int = 60,
     ):
+        super().__init__(space)
         if exploit is None:
             exploit = {
                 "of_type": "PipelineExploit",
@@ -217,7 +218,6 @@ def __init__(
 
         self.lineages = Lineages()
 
-        super().__init__(space)
         self.seed = seed
         self.population_size = population_size
         self.generations = generations
@@ -318,14 +318,15 @@ def suggest(self, num: int) -> list[Trial]:
             A list of trials representing values suggested by the algorithm.
 
         """
-
-        # Sample points until num is met, or population_size
-        num_random_samples = min(max(self.population_size - self._num_root, 0), num)
+        assert num > 0
         logger.debug(
             "PBT has %s pending or completed trials at root, %s broken trials.",
             self._num_root,
             len(self.lineages) - self._num_root,
         )
+
+        # Sample points until num is met, or population_size
+        num_random_samples = min(max(self.population_size - self._num_root, 0), num)
         logger.debug("Sampling %s new trials", num_random_samples)
         trials = self._sample(num_random_samples)
         logger.debug("Sampled %s new trials", len(trials))

src/orion/algo/random.py

@@ -5,9 +5,14 @@
 Draw and deliver samples from prior defined in problem's domain.
 
 """
+from __future__ import annotations
+
+from typing import Sequence
+
 import numpy
 
 from orion.algo.base import BaseAlgorithm
+from orion.algo.space import Space
 
 
 class Random(BaseAlgorithm):
@@ -23,10 +28,12 @@ class Random(BaseAlgorithm):
 
     """
 
-    def __init__(self, space, seed=None):
-        super().__init__(space, seed=seed)
+    def __init__(self, space: Space, seed: int | Sequence[int] | None = None):
+        super().__init__(space)
+        self.seed = seed
+        self.seed_rng(seed)
 
-    def seed_rng(self, seed):
+    def seed_rng(self, seed: int | Sequence[int] | None):
         """Seed the state of the random number generator.
 
         :param seed: Integer seed for the random number generator.

src/orion/algo/registry.py

@@ -4,18 +4,27 @@
 import copy
 from collections import defaultdict
 from logging import getLogger as get_logger
-from typing import Any, Container, Iterator, Mapping
+from typing import Any, Container, Iterable, Iterator, Mapping
 
 from orion.core.worker.trial import Trial, TrialCM
 
 logger = get_logger(__name__)
 
 
 class Registry(Container[Trial]):
-    """In-memory container for the trials that the algorithm suggests/observes/etc."""
+    """In-memory container for the trials that the algorithm suggests/observes/etc.
 
-    def __init__(self):
+    This behaves a bit like a managed dictionary, but the "keys" are trials ids, which
+    (at the time of writing) can vary depending on how we chose to compute them.
+    """
+
+    def __init__(self, trials: Iterable[Trial] = ()):
         self._trials: dict[str, Trial] = {}
+        for trial in trials:
+            self.register(trial)
+
+    def __repr__(self) -> str:
+        return f"{type(self).__qualname__}({list(iter(self))})"
 
     def __contains__(self, trial_or_id: str | Trial | Any) -> bool:
         if isinstance(trial_or_id, TrialCM):
@@ -94,7 +103,7 @@ def get_existing(self, trial: Trial) -> Trial:
 class RegistryMapping(Mapping[Trial, "list[Trial]"]):
     """A map between the original and transformed registries.
 
-    This object is used in the `SpaceTransformAlgoWrapper` to check if a trial in the original space
+    This object is used in the `SpaceTransform` to check if a trial in the original space
     has equivalent trials in the transformed space.
 
     The goal is to make it so the algorithms don't have to care about the transforms/etc.
@@ -123,10 +132,15 @@ def set_state(self, statedict: dict):
         self._mapping = copy.deepcopy(statedict["_mapping"])
 
     def __iter__(self) -> Iterator[Trial]:
+        """Iterate over the trials in the original registry."""
         for trial_id in self._mapping:
             yield self.original_registry[trial_id]
 
     def __len__(self) -> int:
+        """Give the number of trials in the mapping.
+
+        This should be the same as the number of trials in the original registry.
+        """
         return len(self._mapping)
 
     def __contains__(self, trial: Trial):
@@ -160,11 +174,14 @@ def register(self, original_trial: Trial, transformed_trial: Trial) -> str:
         self._mapping[original_trial_id].add(transformed_trial_id)
         return original_trial_id
 
+    def __repr__(self) -> str:
+        return f"{type(self).__qualname__}({list((trial, self.get_trials(trial)) for trial in self)})"
+
 
 def _get_id(trial: Trial) -> str:
     """Returns the unique identifier to be used to store the trial.
 
     Only to be used internally in this module. This ignores the `experiment`
     attribute of the trial.
     """
-    return Trial.compute_trial_hash(trial, ignore_experiment=True)
+    return Trial.compute_trial_hash(trial)

src/orion/benchmark/benchmark_client.py

@@ -76,9 +76,9 @@ def get_or_create_benchmark(
     )
 
     if input_configure and input_benchmark.configuration != benchmark.configuration:
-        logger.warn(
-            "Benchmark with same name is found but has different configuration, "
-            "which will be used for this creation.\n{}".format(benchmark.configuration)
+        logger.warning(
+            f"Benchmark with same name is found but has different configuration, "
+            f"which will be used for this creation.\n{benchmark.configuration}"
         )
 
     if benchmark_id is None: