[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:
The
Policy
is effectively stateless, and thus can be deleted and recovered at any time (e.g. due to a server preemption or failure).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 additionalrecover
/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 additionalload
/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 thenload
.
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.