# Copyright 2021 Amazon.com, Inc. or its affiliates. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License").
# You may not use this file except in compliance with the License.
# A copy of the License is located at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# or in the "license" file accompanying this file. This file is distributed
# on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either
# express or implied. See the License for the specific language governing
# permissions and limitations under the License.
from typing import Dict
import json
import time
from syne_tune.tuner_callback import TunerCallback
from syne_tune import Tuner
from syne_tune.backend.trial_status import Trial
from syne_tune.optimizer.schedulers import FIFOScheduler
from syne_tune.optimizer.schedulers.searchers import ModelBasedSearcher
[docs]
class StoreSearcherStatesCallback(TunerCallback):
"""
Stores list of searcher states alongside a tuning run. The list
is extended by a new state whenever the ``TuningJobState`` has changed
compared to the last recently added one.
This callback is useful to create meaningful unit tests, by sampling
a given searcher alongside a realistic experiment.
Works only for ``ModelBasedSearcher`` searchers. For other searchers, nothing
is stored.
"""
def __init__(self):
super().__init__()
self._states = []
self._num_observations = None
self._start_time = time.time()
self._searcher = None
[docs]
def on_tuning_start(self, tuner: Tuner):
scheduler = tuner.scheduler
if isinstance(scheduler, FIFOScheduler):
searcher = scheduler.searcher
if isinstance(searcher, ModelBasedSearcher):
self._searcher = searcher
[docs]
def on_trial_result(self, trial: Trial, status: str, result: Dict, decision: str):
if self._searcher is not None:
state = self._searcher.state_transformer.state
num_observations = state.num_observed_cases()
if (
self._num_observations is None
or num_observations != self._num_observations
):
searcher_state = self._searcher.get_state()
searcher_state["elapsed_time"] = time.time() - self._start_time
searcher_state["num_observations"] = num_observations
searcher_state["num_configs"] = len(state.candidate_evaluations)
self._states.append(searcher_state)
self._num_observations = num_observations
@property
def states(self):
return self._states
[docs]
def searcher_state_as_code(self, pos: int, add_info: bool = False):
assert 0 <= pos < len(self._states)
searcher_state = self._states[pos]
lines = []
if add_info:
lines.append(f"# elapsed_time = {searcher_state['elapsed_time']}")
lines.append(f"# num_observations = {searcher_state['num_observations']}")
lines.append(f"# num_configs = {searcher_state['num_configs']}")
model_params = searcher_state["model_params"]
lines.append(f"_model_params = '{json.dumps(model_params)}'")
# lines.append("model_params = json.loads(_model_params)")
state = searcher_state["state"]
lines.append(f"_state = '{json.dumps(state)}'")
# lines.append("state = decode_state(enc_state=json.loads(_state), hp_ranges=hp_ranges)")
return "\n".join(lines)