copybara-service · Mar 9, 2023 · Mar 9, 2023
diff --git a/vizier/_src/algorithms/designers/grid.py b/vizier/_src/algorithms/designers/grid.py
@@ -81,10 +81,10 @@ def __init__(
   def from_problem(
       cls,
       problem: pyvizier.ProblemStatement,
-      shuffle_seed: Optional[int] = None,
+      seed: Optional[int] = None,
   ):
     """For wrapping via `PartiallySerializableDesignerPolicy`."""
-    return GridSearchDesigner(problem.search_space, shuffle_seed)
+    return GridSearchDesigner(problem.search_space, seed)
 
   def update(
       self,

diff --git a/vizier/_src/algorithms/designers/quasi_random.py b/vizier/_src/algorithms/designers/quasi_random.py
@@ -296,8 +296,11 @@ def create_input_converter(pc):
     self._output_specs = tuple(self._converter.output_specs.values())
 
   @classmethod
-  def from_problem(cls, problem: vz.ProblemStatement):
+  def from_problem(
+      cls, problem: vz.ProblemStatement, seed: Optional[int] = None
+  ):
     """For wrapping via `PartiallySerializableDesignerPolicy`."""
+    del seed
     return QuasiRandomDesigner(problem.search_space)
 
   def load(self, metadata: vz.Metadata) -> None:

diff --git a/vizier/_src/algorithms/evolution/nsga2.py b/vizier/_src/algorithms/evolution/nsga2.py
@@ -197,7 +197,8 @@ def create_nsga2(
     *,
     ranking_fn: Callable[[np.ndarray], np.ndarray] = _pareto_rank,
     eviction_limit: Optional[int] = None,
-    metadata_namespace: str = 'nsga2'
+    metadata_namespace: str = 'nsga2',
+    seed: Optional[int] = None
 ) -> templates.CanonicalEvolutionDesigner[Population, Offspring]:
   """Creates NSGA2 Designer.
 
@@ -213,11 +214,12 @@ def create_nsga2(
       you want to improve performance, your own implementation can be injected.
     eviction_limit: Evict a gene that has been alive for this many generations.
     metadata_namespace: Metadata namespace to use.
+    seed: Random seed.
 
   Returns:
     NSGA2 Designer.
   """
-
+  del seed
   algorithm = templates.CanonicalEvolutionDesigner(
       numpy_populations.PopulationConverter(
           problem.search_space,

diff --git a/vizier/_src/algorithms/policies/designer_policy.py b/vizier/_src/algorithms/policies/designer_policy.py
@@ -14,11 +14,16 @@
 
 from __future__ import annotations
 
-"""Wrappers for Designer into Policy."""
+"""Wrappers for Designer into Policy.
+
+This file should contain the core logic to run Designers
+efficiently for testing, benchmarking, or production purposes. Typical
+usage is via PolicySuggester.from_designer_factory.
+"""
 
 import abc
 import json
-from typing import Callable, Generic, Sequence, Type, TypeVar
+from typing import Generic, Optional, Sequence, Type, TypeVar
 
 from absl import logging
 from vizier import algorithms as vza
@@ -109,10 +114,11 @@ def __init__(
       self,
       problem_statement: vz.ProblemStatement,
       supporter: pythia.PolicySupporter,
-      designer_factory: Callable[[vz.ProblemStatement], _T],
+      designer_factory: vza.DesignerFactory[_T],
       *,
       ns_root: str = 'designer_policy_v0',
       verbose: int = 0,
+      seed: Optional[int] = None,
   ):
     """Init.
 
@@ -123,6 +129,7 @@ def __init__(
         (Partially)SerializableDesigner designers.
       ns_root: Root of the namespace where policy state is stored.
       verbose: Logging verbosity.
+      seed: Random seed to be passed to the designer factory.
     """
     self._supporter = supporter
     self._designer_factory = designer_factory
@@ -131,6 +138,7 @@ def __init__(
     self._problem_statement = problem_statement
     self._verbose = verbose
     self._designer = None
+    self._seed = seed
 
   def suggest(self, request: pythia.SuggestRequest) -> pythia.SuggestDecision:
     """Perform a suggest operation.
@@ -269,7 +277,9 @@ def _initialize_designer(self,
       logging.log_if(logging.INFO, 'Failed to decode state. %s',
                      self._verbose >= 1, e)
       # Restart the state of the policy.
-      self._designer = self._designer_factory(problem_statement)
+      self._designer = self._designer_factory(
+          problem_statement, seed=self._seed
+      )
       self._incorporated_completed_trial_ids: set[int] = set()
 
   def dump(self) -> vz.Metadata:
@@ -326,6 +336,20 @@ def _get_trials(
     return new_trials
 
 
+# TODO: Rewrite base to separate incorporated trial logic from
+# serialization logic. For now, it simply disables the serialization.
+class InRamDesignerPolicy(_SerializableDesignerPolicyBase[vza.Designer]):
+  """Wraps a Designer in RAM for efficient updates()."""
+
+  # Override restore() to simply return the designer.
+  def _restore_designer(self, designer_metadata: vz.Metadata) -> vza.Designer:
+    return self._designer
+
+  # Override dump() since this is in ram.
+  def dump(self) -> vz.Metadata:
+    return vz.Metadata()
+
+
 class PartiallySerializableDesignerPolicy(
     _SerializableDesignerPolicyBase[vza.PartiallySerializableDesigner]):
   """Wraps a PartiallySerializableDesigner."""

diff --git a/vizier/_src/algorithms/policies/designer_policy_test.py b/vizier/_src/algorithms/policies/designer_policy_test.py
@@ -130,7 +130,7 @@ def test_restore_fully_serializable_designer(self):
     policy = dp.SerializableDesignerPolicy(
         problem_statement=vz.ProblemStatement(),
         supporter=runner,
-        designer_factory=lambda _: _FakeSerializableDesigner(),
+        designer_factory=lambda _, **kwargs: _FakeSerializableDesigner(),
         designer_cls=_FakeSerializableDesigner,
         ns_root='test',
         verbose=2,
@@ -141,7 +141,7 @@ def test_restore_fully_serializable_designer(self):
     restored_policy = dp.SerializableDesignerPolicy(
         problem_statement=vz.ProblemStatement(metadata=metadata),
         supporter=runner,
-        designer_factory=lambda _: _FakeSerializableDesigner(),
+        designer_factory=lambda _, **kwargs: _FakeSerializableDesigner(),
         designer_cls=_FakeSerializableDesigner,
         ns_root='test',
         verbose=2,
@@ -158,7 +158,7 @@ def test_restore_partially_serializable_designer(self):
     policy = dp.PartiallySerializableDesignerPolicy(
         vz.ProblemStatement(),
         runner,
-        lambda _: _FakeSerializableDesigner(),
+        lambda _, **kwargs: _FakeSerializableDesigner(),
         ns_root='test',
         verbose=2,
     )
@@ -168,7 +168,7 @@ def test_restore_partially_serializable_designer(self):
     restored_policy = dp.PartiallySerializableDesignerPolicy(
         vz.ProblemStatement(),
         runner,
-        lambda _: _FakeSerializableDesigner(),
+        lambda _, **kwargs: _FakeSerializableDesigner(),
         ns_root='test',
         verbose=2,
     )
@@ -185,7 +185,7 @@ def test_restore_partially_serializable_designer(self):
   def test_update_stateless_designer(self):
     runner = _create_runner()
     designer = _FakeDesigner()
-    policy = dp.DesignerPolicy(runner, lambda _: designer)
+    policy = dp.DesignerPolicy(runner, lambda _, **kwargs: designer)
     runner.SuggestTrials(policy, 5)
     self.assertLen(
         designer.last_completed,
@@ -215,7 +215,7 @@ def test_update_stateful_designer(self):
     policy = dp.PartiallySerializableDesignerPolicy(
         vz.ProblemStatement(),
         runner,
-        lambda _: designer,
+        lambda _, **kwargs: designer,
         ns_root='test',
         verbose=2,
     )

diff --git a/vizier/_src/algorithms/testing/test_runners.py b/vizier/_src/algorithms/testing/test_runners.py
@@ -17,13 +17,15 @@
 """Test runners for algorithms."""
 from typing import Any, Callable, Collection, Optional, Sequence
 
+
 from absl import logging
 import attr
 import numpy as np
 from vizier import algorithms as vza
 from vizier import benchmarks
 from vizier import pythia
 from vizier import pyvizier as vz
+from vizier._src.algorithms.policies.designer_policy import InRamDesignerPolicy
 
 
 @attr.define
@@ -78,11 +80,9 @@ class RandomMetricsRunner:
   verbose: int = attr.field(default=0, kw_only=True)
   validate_parameters: bool = attr.field(default=False, kw_only=True)
 
-  def _run(self,
-           algorithm: benchmarks.AlgorithmSuggester) -> Collection[vz.Trial]:
+  def _run(self, algorithm: benchmarks.PolicySuggester) -> Collection[vz.Trial]:
     """Implementation of run methods."""
     rng = np.random.RandomState(self.seed)
-    all_trials = []
     for it in range(self.iters):
       suggestions = algorithm.suggest(self.batch_size)
       if not suggestions:
@@ -98,16 +98,12 @@ def _run(self,
         for mi in self.problem.metric_information:
           measurement.metrics[mi.name] = rng.uniform(
               mi.min_value_or(lambda: -10.), mi.max_value_or(lambda: 10.))
-        trials.append(
-            suggestion.to_trial(len(all_trials) + 1).complete(measurement))
+        trials.append(suggestion.complete(measurement))
       if self.verbose:
         logging.info('At iteration %s, trials suggested and evaluated:\n%s', it,
                      trials)
-      algorithm.post_completion_callback(
-          completed=vza.CompletedTrials(trials), all_active=vza.ActiveTrials()
-      )
-      all_trials.extend(trials)
-    return all_trials
+      algorithm.supporter.AddTrials(trials)
+    return algorithm.supporter.GetTrials()
 
   def run_policy(self, policy_factory: Callable[[pythia.PolicySupporter],
                                                 pythia.Policy]):
@@ -118,8 +114,13 @@ def run_policy(self, policy_factory: Callable[[pythia.PolicySupporter],
 
   def run_designer(self, designer: vza.Designer):
     """Run the specified $designer."""
-    runner = pythia.InRamPolicySupporter(self.problem)
-    return self._run(benchmarks.DesignerSuggester(designer, runner))
+    supporter = pythia.InRamPolicySupporter(self.problem)
+    policy = InRamDesignerPolicy(
+        problem_statement=self.problem,
+        supporter=supporter,
+        designer_factory=lambda _, **kwargs: designer,
+    )
+    return self._run(benchmarks.PolicySuggester(policy, supporter))
 
 
 def run_with_random_metrics(

diff --git a/vizier/_src/benchmarks/runners/algorithm_suggester.py b/vizier/_src/benchmarks/runners/algorithm_suggester.py