mlos_bench.schedulers.base_scheduler

Base class for the optimization loop scheduling policies.

Classes

Scheduler

Base class for the optimization loop scheduling policies.

Module Contents

class mlos_bench.schedulers.base_scheduler.Scheduler(*, config: dict[str, Any], global_config: dict[str, Any], trial_runners: collections.abc.Iterable[mlos_bench.schedulers.trial_runner.TrialRunner], optimizer: mlos_bench.optimizers.base_optimizer.Optimizer, storage: mlos_bench.storage.base_storage.Storage, root_env_config: str)[source]

Bases: contextlib.AbstractContextManager

Base class for the optimization loop scheduling policies.

Create a new instance of the scheduler. The constructor of this and the derived classes is called by the persistence service after reading the class JSON configuration. Other objects like the TrialRunner(s) and their Environment(s) and Optimizer are provided by the Launcher.

Parameters:

config (dict) – The configuration for the Scheduler.
global_config (dict) – The global configuration for the Experiment.
trial_runner (Iterable[TrialRunner]) – The set of TrialRunner(s) (and associated Environment(s)) to benchmark/optimize.
optimizer (Optimizer) – The Optimizer to use.
storage (Storage) – The Storage to use.
root_env_config (str) – Path to the root Environment configuration.
trial_runners (collections.abc.Iterable[mlos_bench.schedulers.trial_runner.TrialRunner])

__enter__() → Scheduler[source]

Enter the scheduler’s context.

Return type:: Scheduler

__exit__(ex_type: type[BaseException] | None, ex_val: BaseException | None, ex_tb: types.TracebackType | None) → Literal[False][source]

Exit the context of the scheduler.

Parameters:

ex_type (type[BaseException] | None)
ex_val (BaseException | None)
ex_tb (types.TracebackType | None)

Return type:

Literal[False]

__repr__() → str[source]

Produce a human-readable version of the Scheduler (mostly for logging).

Returns:: string – A human-readable version of the Scheduler.
Return type:: str

add_new_optimizer_suggestions() → bool[source]

Optimizer part of the loop.

Load the results of the executed trials into the Optimizer, suggest new configurations, and add them to the queue.

Returns:: The return value indicates whether the optimization process should continue to get suggestions from the Optimizer or not. See Also: not_done().
Return type:: bool

add_trial_to_queue(tunables: mlos_bench.tunables.tunable_groups.TunableGroups, ts_start: datetime.datetime | None = None) → None[source]

Add a configuration to the queue of trials 1 or more times.

(e.g., according to the trial_config_repeat_count)

Parameters:

tunables (TunableGroups) – The tunable configuration to add to the queue.
ts_start (datetime.datetime | None) – Optional timestamp to use to start the trial.

Return type:

None

Notes

Alternative scheduling policies may prefer to expand repeats over time as well as space, or adjust the number of repeats (budget) of a given trial based on whether initial results are promising.

assign_trial_runners(trials: collections.abc.Iterable[mlos_bench.storage.base_storage.Storage.Trial]) → None[source]

Assigns a TrialRunner to each Trial in the batch.

The base class implements a simple round-robin scheduling algorithm for each Trial in sequence.

Subclasses can override this method to implement a more sophisticated policy. For instance:

def assign_trial_runners(
    self,
    trials: Iterable[Storage.Trial],
) -> TrialRunner:
    trial_runners_map = {}
    # Implement a more sophisticated policy here.
    # For example, to assign the Trial to the TrialRunner with the least
    # number of running Trials.
    # Or assign the Trial to the TrialRunner that hasn't executed this
    # TunableValues Config yet.
    for (trial, trial_runner) in trial_runners_map:
        # Call the base class method to assign the TrialRunner in the Trial's metadata.
        trial.set_trial_runner(trial_runner)
    ...

Notes

Subclasses are not required to assign a TrialRunner to the Trial (e.g., if the Trial should be deferred to a later time).

Parameters:: trials (Iterable[Storage.Trial]) – The trial to assign a TrialRunner to.
Return type:: None

get_best_observation() → tuple[dict[str, float] | None, mlos_bench.tunables.tunable_groups.TunableGroups | None][source]

Get the best observation from the optimizer.

Return type:: tuple[dict[str, float] | None, mlos_bench.tunables.tunable_groups.TunableGroups | None]

get_trial_runner(trial: mlos_bench.storage.base_storage.Storage.Trial) → mlos_bench.schedulers.trial_runner.TrialRunner[source]

Gets the TrialRunner associated with the given Trial.

Parameters:: trial (Storage.Trial) – The trial to get the associated TrialRunner for.
Return type:: TrialRunner

load_tunable_config(config_id: int) → mlos_bench.tunables.tunable_groups.TunableGroups[source]

Load the existing tunable configuration from the storage.

Parameters:: config_id (int)
Return type:: mlos_bench.tunables.tunable_groups.TunableGroups

not_done() → bool[source]

Check the stopping conditions.

By default, stop when the Optimizer converges or the limit of max_trials is reached.

Return type:: bool

run_schedule(running: bool = False) → None[source]

Runs the current schedule of trials.

Check for Trial instances with :py:attr:.Status.PENDING` and an assigned trial_runner_id in the queue and run them with run_trial().

Subclasses can override this method to implement a more sophisticated scheduling policy.

Parameters:: running (bool) – If True, run the trials that are already in a “running” state (e.g., to resume them). If False (default), run the trials that are pending.
Return type:: None

abstractmethod run_trial(trial: mlos_bench.storage.base_storage.Storage.Trial) → None[source]

Set up and run a single trial.

Save the results in the storage.

Parameters:: trial (mlos_bench.storage.base_storage.Storage.Trial)
Return type:: None

start() → None[source]

Start the scheduling loop.

Return type:: None

teardown() → None[source]

Tear down the TrialRunners/Environment(s).

Call it after the completion of the Scheduler.start() in the Scheduler context.

Return type:: None

property environments: collections.abc.Iterable[mlos_bench.environments.base_environment.Environment][source]

Gets the Environment from the TrialRunners.

Return type:: collections.abc.Iterable[mlos_bench.environments.base_environment.Environment]

property experiment: mlos_bench.storage.base_storage.Storage.Experiment | None[source]

Gets the Experiment Storage.

Return type:: mlos_bench.storage.base_storage.Storage.Experiment | None

global_config[source]

property max_trials: int[source]

Gets the maximum number of trials to run for a given experiment, or -1 for no limit.

Return type:: int

property optimizer: mlos_bench.optimizers.base_optimizer.Optimizer[source]

Gets the Optimizer.

Return type:: mlos_bench.optimizers.base_optimizer.Optimizer

property ran_trials: list[mlos_bench.storage.base_storage.Storage.Trial][source]

Get the list of trials that were run.

Return type:: list[mlos_bench.storage.base_storage.Storage.Trial]

property root_environment: mlos_bench.environments.base_environment.Environment[source]

Gets the root (prototypical) Environment from the first TrialRunner.

Notes

All TrialRunners have the same Environment config and are made unique by their use of the unique trial_runner_id assigned to each TrialRunner’s Environment’s global_config.

Return type:: mlos_bench.environments.base_environment.Environment

property storage: mlos_bench.storage.base_storage.Storage[source]

Gets the Storage.

Return type:: mlos_bench.storage.base_storage.Storage

property trial_config_repeat_count: int[source]

Gets the number of trials to run for a given config.

Return type:: int

property trial_count: int[source]

Gets the current number of trials run for the experiment.

Return type:: int

property trial_runners: dict[int, mlos_bench.schedulers.trial_runner.TrialRunner][source]

Gets the set of Trial Runners.

Return type:: dict[int, mlos_bench.schedulers.trial_runner.TrialRunner]