Pythia Policies and Hosting Designers

This documentation will allow a developer to:

• Understand the basic structure of a Pythia Policy.

• Host Designers in the service.

Installation and reference imports

!pip install google-vizier

from typing import Optional, Sequence

from vizier import pythia
from vizier import algorithms
from vizier.service import pyvizier as vz
from vizier._src.algorithms.policies import designer_policy
from vizier._src.algorithms.evolution import nsga2

Pythia Policies

The Pythia Service maps algorithm names to Policy objects. All algorithms which need to be hosted on the server must eventually be wrapped into a Policy.

Every Policy is injected with a PolicySupporter, which is a client used for fetching data from the datastore. This design choice serves two core purposes:

1. The Policy is effectively stateless, and thus can be deleted and recovered at any time (e.g. due to a server preemption or failure).

2. Consequently, this avoids needing to save an explicit and potentially complicated algorithm state. Instead, the “algorithm state” can be recovered purely from the entire study containing (metadata, study_config, trials).

We show the Policy abstract class explicitly below. Exact class entrypoint can be found here.

class Policy(abc.ABC):
  """Interface for Pythia Policy subclasses."""

  @abc.abstractmethod
  def suggest(self, request: SuggestRequest) -> SuggestDecision:
    """Compute suggestions that Vizier will eventually hand to the user."""

  @abc.abstractmethod
  def early_stop(self, request: EarlyStopRequest) -> EarlyStopDecisions:
    """Decide which Trials Vizier should stop."""

  @property
  def should_be_cached(self) -> bool:
    """Returns True if it's safe & worthwhile to cache this Policy in RAM."""
    return False

Fundamental Rule of Service Pythia Policies

For algorithms used in the Pythia Service, the fundamental rule is to assume that a Pythia policy class instance will only call once per user interaction:

• __init__

• suggest()

and be immediately deleted afterwards. Thus a typical policy will use a stateless_algorithm and roughly look like:

class TypicalPolicy(Policy):

  def __init__(self, policy_supporter: PolicySupporter):
    self._policy_supporter = policy_supporter

  def suggest(self, request: SuggestRequest) -> SuggestDecision:
    all_completed = policy_supporter.GetTrials(status_matches=COMPLETED)
    all_active = policy_supporter.GetTrials(status_matches=ACTIVE)
    suggestions = stateless_algorithm(all_completed, all_active)
    return SuggestDecision(suggestions)

Example Pythia Policy

Here, we write a toy policy, where we only act on CATEGORICAL parameters for simplicity. The make_parameters function will simply for-loop over every category and then cycle back.

def make_parameters(
    search_space: vz.SearchSpace, index: int
) -> vz.ParameterDict:
  parameter_dict = vz.ParameterDict()
  for parameter_config in search_space.parameters:
    if parameter_config.type != vz.ParamterType.CATEGORICAL:
      raise ValueError("This function only supports CATEGORICAL parameters.")
    feasible_values = parameter_config.feasible_values
    parameter_dict[parameter_config.name] = vz.ParameterValue(
        value=feasible_values[index % len(feasible_values)]
    )
  return parameter_dict

To collect the index from the database, we will use the PolicySupporter to obtain the maximum trial ID based on completed and active trials.

def get_next_index(policy_supporter: pythia.PolicySupporter):
  """Returns current trial index."""
  completed = policy_supporter.GetTrials(status_matches=vz.TrialStatus.COMPLETED)
  active = policy_supporter.GetTrials(status_matches=vz.TrialStatus.ACTIVE)
  trial_ids = [t.id for t in completed + active]

  if trial_ids:
    return max(trial_ids)
  return 0

We can now put it all together into our Pythia Policy.

class MyPolicy(pythia.Policy):

  def __init__(self, policy_supporter: pythia.PolicySupporter):
    self._policy_supporter = policy_supporter

  def suggest(self, request: pythia.SuggestRequest) -> pythia.SuggestDecision:
    """Gets number of Trials to propose, and produces Trials."""
    suggest_decision_list = []
    for _ in range(request.count):
      index = get_next_index(self._policy_supporter)
      parameters = make_parameters(request.study_config.search_space, index)
      suggest_decision_list.append(vz.TrialSuggestion(parameters=parameters))
    return pythia.SuggestDecision(
        suggestions=suggest_decision_list, metadata=vz.MetadataDelta()
    )

Wrapping Designers as Pythia Policies

Consider if your algorithm code fits in the simpler Designer abstraction, which avoids needing to deal with distributed systems logic.

For example, the same exact behavior above can be re-written as a Designer:

class MyDesigner(algorithms.Designer):

  def __init__(self, study_config: vz.StudyConfig):
    self._study_config = study_config
    self._completed_trials = []
    self._active_trials = []

  def update(
      self,
      completed: algorithms.CompletedTrials,
      all_active: algorithms.ActiveTrials,
  ) -> None:
    self._completed_trials.extend(completed.trials)
    self._active_trials = all_active.trials

  def suggest(
      self, count: Optional[int] = None
  ) -> Sequence[vz.TrialSuggestion]:
    if count is None:
      return []
    trial_ids = [t.id for t in self._completed_trials + self._active_trials]
    current_index = max(trial_ids)
    return [
        make_parameters(self._study_config.search_space, current_index + i)
        for i in range(count)
    ]

The entire designer (if deleted or preempted) can conveniently be recovered in just a single call of update() after __init__.

Thus we may immediately wrap MyDesigner into a Pythia Policy with the following Pythia suggest() implementation:

• Create the designer temporarily.

• Update the temporary designer with all previously completed trials and active trials.

• Obtain suggestions from the temporary designer.

This is done conveniently with the DesignerPolicy wrapper (code):

class DesignerPolicy(Policy):
  """Wraps a Designer into a Pythia Policy."""

  def __init__(self, supporter: PolicySupporter, designer_factory: Factory[Designer]):
    self._supporter = supporter
    self._designer_factory = designer_factory

  def suggest(self, request: SuggestRequest) -> SuggestDecision:
    completed = self._supporter.GetTrials(status_matches=COMPLETED)
    active = self._supporter.GetTrials(status_matches=ACTIVE)
    designer = self._designer_factory(...)
    designer.update(CompletedTrials(completed), ActiveTrials(active))
    return SuggestDecision(designer.suggest(request.count))

Below is the actual act of wrapping:

designer_factory = lambda study_config: MyDesigner(study_config)
supporter: PolicySupporter = ... # Assume PolicySupporter was created.
pythia_policy = DesignerPolicy(supporter, designer_factory)

Serializing Designer States

The above method can gradually become slower as the number of completed trials in the study increases.

Thus we may consider storing a compressed representation of the algorithm state instead. Examples include:

• The coordinate position in a grid search algorithm.

• The population for evolutionary algorithms such as NSGA2.

• Directory location for stored neural network weights.

As a simple example, consider the case if our designer stores a _counter of all suggestions it has made:

class CounterDesigner(Designer):

  def __init__(self, ...):
    ...
    self._counter = 0

  def suggest(self, count: Optional[int] = None) -> Sequence[TrialSuggestion]:
    ...
    self._counter += len(suggestions)
    return suggestions

Vizier offers two Designer subclasses, both of which will use the Metadata primitive to store algorithm state data:

• SerializableDesigner will use additional recover/dump methods and should be used if the entire algorithm state can be easily serialized and can be saved and restored in full.

• PartiallySerializableDesigner will use additional load/dump methods and be used if the algorithm has subcomponents that are not easily serializable. State recovery will be handled by calling the Designer’s __init__ (with same arguments as before) and then load.

They can also be converted into Pythia Policies using SerializableDesignerPolicy and PartiallySerializableDesignerPolicy respectively.

Below is an example modifying our CounterDesigner into CounterSerialDesigner and CounterPartialDesigner respectively:

class CounterSerialDesigner(algorithms.SerializableDesigner):

  def __init__(self, counter: int):
    self._counter = counter

  @classmethod
  def recover(cls, metadata: vz.Metadata) -> CounterSerialDesigner:
    return cls(metadata['counter'])

  def dump(self) -> vz.Metadata:
    metadata = vz.Metadata()
    metadata['counter'] = str(self._counter)
    return metadata


class CounterPartialDesigner(algorithms.PartiallySerializableDesigner):

  def load(self, metadata: vz.Metadata) -> None:
    self._counter = int(metadata['counter'])

  def dump(self) -> vz.Metadata:
    metadata = vz.Metadata()
    metadata['counter'] = str(self._counter)
    return metadata

Additional References

• Our policies folder contains examples of Pythia policies.