[Deprecated] Writing Algorithms and Pythia Policies

This documentation will allow a developer to:

• Understand the basic structure of a Pythia policy.

• Use the Designer API for simplfying algorithm design.

Installation and reference imports

!pip install google-vizier

from typing import Optional, Sequence

from vizier import pythia
from vizier import pyvizier as vz
from vizier import algorithms
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(
    study_config: vz.StudyConfig, index: int
) -> vz.ParameterDict:
  parameter_dict = vz.ParameterDict()
  for parameter_config in study_config.search_space.parameters:
    if parameter_config.type != vz.ParamterType.CATEGORICAL:
      raise ValueError("This function only supports CATEGORICAL parameters.")
    feasible_values = parameter_config.feasible_values
    categorical_size = len(feasible_values)
    parameter_dict[parameter_config.name] = vz.ParameterValue(
        value=feasible_values[index % categorical_size]
    )
  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_and_active = policy_supporter.GetTrials(
      status_matches=vz.TrialStatus.COMPLETED
  ) + policy_supporter.GetTrials(status_matches=vz.TrialStatus.ACTIVE)
  trial_ids = [t.id for t in completed_and_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, index)
      suggest_decision_list.append(vz.TrialSuggestion(parameters=parameters))
    return pythia.SuggestDecision(
        suggestions=suggest_decision_list, metadata=vz.MetadataDelta()
    )

Designers

While Pythia policies are the proper interface for hosting algorithms on the server, we also provide the Designer API to simplify algorithm development. A Policy can use PolicySupporter to actively fetch trials from the datastore. Unlike a Policy, a Designer cannot actively fetch trials, but it can be updated with previously COMPLETED or ACTIVE trials.

We display the Designer class below. The source of truth for Designer can be found here.

class Designer(...):
  """Suggestion algorithm for sequential usage."""

  @abc.abstractmethod
  def update(self, completed: CompletedTrials, all_active: ActiveTrials) -> None:
    """Updates recently completed and ALL active trials into the designer's state."""

  @abc.abstractmethod
  def suggest(self, count: Optional[int] = None) -> Sequence[vz.TrialSuggestion]:
    """Make new suggestions."""

Every time update() is called, the Designer will get any newly COMPLETED trials since the last update() call, and will get all ACTIVE trials at the current moment in time.

Note that trials which may have been provided as ACTIVE in previous update() calls, can be provided as COMPLETED in subsequent update() calls.

To implement our same algorithm above in a Designer, only the update() and suggest() methods need to be implemented using our previous make_parameters function. The designer class can now store completed trials inside its self._completed_trials attribute.

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, current_index + i)
        for i in range(count)
    ]

Wrapping a Designer to a Pythia Policy

Note that in the above implementation of MyDesigner, the entire algorithm (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(pythia.Policy):
  """Wraps a Designer into a Pythia Policy."""

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

  def suggest(self, request: pythia.SuggestRequest) -> pythia.SuggestDecision:
    completed = self._supporter.GetTrials(status_matches=vz.TrialStatus.COMPLETED)
    active = self._supporter.GetTrials(status_matches=vz.TrialStatus.ACTIVE)
    designer.update(vza.CompletedTrials(completed), vza.ActiveTrials(active))
    return pythia.SuggestDecision(
        designer.suggest(request.count), metadata=vz.MetadataDelta())

Below is the actual act of wrapping:

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

Serializing algorithm 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(algorithms.Designer):

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

  def suggest(
      self, count: Optional[int] = None) -> Sequence[vz.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.

• Our designers folder contains examples of designers.

• Our evolution folder contains examples of creating evolutionary designers, such as NSGA2.

• Our designer testing routine contains up-to-date examples on interacting with designers.