> ## Documentation Index
> Fetch the complete documentation index at: https://docs.coreweave.com/llms.txt
> Use this file to discover all available pages before exploring further.

# Configure an optimization engine

Choose an optimization engine for Weights & Biases Sweeps. Use the built-in Weights & Biases engine for random, grid, or Bayesian search, or use Optuna or Ax for their search algorithms and customization options.

The engine, scheduler, and sweep agents work together:

| Component | Role |
| - | - |
| Optimization engine | Suggests hyperparameter configurations and, depending on the search strategy, uses run results to guide the search. |
| Sweep scheduler | Requests configurations from the engine, submits them to Weights & Biases, and passes run results back to the engine. |
| Sweep agent | Retrieves assigned configurations from Weights & Biases and executes your training code. |

This page describes how to configure an engine for a scheduler running in your own environment. For complete YAML examples, see [Sweep scheduler examples](https://docs.coreweave.com/products/wandb/sweeps/define-sweep-configuration.md#sweep-scheduler-examples).

## Choose an engine

Choose an engine based on the search methods and customization options you need:

| Engine | Description |
| - | - |
| Weights & Biases (built-in) | Random, grid, or Bayesian search configured through sweep configuration keys. |
| [Optuna](https://optuna.org/) | An open source hyperparameter optimization framework with configurable samplers and pruners. |
| [Ax](https://ax.dev/) | An open source platform for adaptive experimentation with configurable generation and stopping strategies. |

Use Optuna or Ax to:

* [Define multiple objectives](#define-multiple-objectives) to optimize more than one metric.
* [Define a custom sweep-level stopping function](#stop-the-sweep) to control when the scheduler stops proposing new runs.

With Optuna, you can also [define a conditional search space](#define-a-conditional-search-space) to select parameters based on earlier choices in a trial.

## Install dependencies

Install Weights & Biases Python SDK version 0.31.0 or later with the extra dependencies for your engine. Run the installation command in the environment where you run the scheduler.

> [tab] Weights & Biases
>
> Use Python 3.10 or later. Install the SDK with the `sweeps` extra to include `sweeps>=0.2.0` in the scheduler environment:
>
> ```bash theme={"system"}
> pip install "wandb[sweeps]>=0.31.0"
> ```

> [tab] Optuna
>
> Use Python 3.10 or later. Install the SDK with the `optuna` extra to include `optuna>=4.0.0` in the scheduler environment:
>
> ```bash theme={"system"}
> pip install "wandb[optuna]>=0.31.0"
> ```
>
> To use the OptunaHub Terminator example on this page, also install `optunahub`.
>
> ```bash theme={"system"}
> pip install optunahub
> ```

> [tab] Ax
>
> Use Python 3.11 or later. Install the SDK with the `ax` extra to include `ax-platform>=1.3.1` in the scheduler environment:
>
> ```bash theme={"system"}
> pip install "wandb[ax]>=0.31.0"
> ```

## Configure optimization

Configure the search strategy, objectives, and search space for your engine. The built-in Weights & Biases engine uses the sweep configuration. For Optuna and Ax, you can provide a Python function that configures an Optuna study or Ax client.

### Configure the search strategy

Set `scheduler.engine` to `wandb`, `optuna`, or `ax`.

For Optuna and Ax, custom optimizer functions are optional. You can switch engines by changing `scheduler.engine` and use the search space and objectives defined in the sweep configuration, including multiple objectives defined with `metrics`.

To use a custom optimizer function, set `scheduler.source` to your Python filename and `scheduler.optimizer` to the function name. For Optuna, you can also define a custom search space with `scheduler.search_space`.

> [tab] Weights & Biases
>
> Configure the built-in Weights & Biases engine using a sweep configuration.
>
> Set `method` to `random`, `grid`, or `bayes`. The engine uses `method`, `metric`, and `parameters` to select configurations for sweep runs.
>
> For example, to configure Bayesian search using the built-in Weights & Biases engine:
>
> ```yaml title="sweep.yaml" theme={"system"}
> program: train.py
> method: bayes
> metric:
>   name: validation_loss
>   goal: minimize
> parameters:
>   learning_rate:
>     min: 0.0001
>     max: 0.01
>   batch_size:
>     values: [16, 32, 64]
> scheduler:
>   engine: wandb
> ```
>
> For details about each search method, see [Search methods](https://docs.coreweave.com/products/wandb/sweeps/sweep-config-keys.md#method).

> [tab] Optuna
>
> To customize the optimizer, define a function that accepts one `SweepInfo` argument and returns an Optuna `Study`.
>
> For example, suppose you have the following sweep configuration:
>
> ```yaml title="sweep.yaml" theme={"system"}
> program: train.py
>
> # Required by the W&B sweep schema, but the Optuna scheduler
> # controls the search strategy.
> method: random
>
> metric:
>   name: val_loss
>   goal: minimize
>
> parameters:
>   n_estimators:
>     min: 20
>     max: 200
>   max_depth:
>     min: 2
>     max: 20
>   min_samples_split:
>     min: 0.01
>     max: 0.5
>
> # Stop after at most 100 runs, even if the terminator does not
> # detect convergence first.
> run_cap: 100
>
> scheduler:
>   engine: optuna
>   source: optimizer.py
>   optimizer: create_optimizer
> ```
>
> In `optimizer.py`, create a study with a TPE sampler and return it with the Terminator’s stopping function:
>
> ```python title="optimizer.py" theme={"system"}
> import optuna
> import optunahub
>
> terminator_module = optunahub.load_module("callbacks/terminator")
>
> EMMREvaluator = terminator_module.EMMREvaluator
> MedianErrorEvaluator = terminator_module.MedianErrorEvaluator
> Terminator = terminator_module.Terminator
>
>
> def create_optimizer(sweep_info):
>     study = optuna.create_study(
>         direction="minimize",
>         sampler=optuna.samplers.TPESampler(seed=42),
>     )
>
>     improvement_evaluator = EMMREvaluator()
>     error_evaluator = MedianErrorEvaluator(improvement_evaluator)
>
>     terminator = Terminator(
>         improvement_evaluator=improvement_evaluator,
>         error_evaluator=error_evaluator,
>         min_n_trials=20,
>     )
>
>     return study, terminator.should_terminate
> ```
>
> The previous example configures two parts of the Optuna integration:
>
> * A `TPESampler`, which suggests hyperparameter configurations.
>
> * An [OptunaHub Terminator](https://hub.optuna.org/callbacks/terminator/), which stops the sweep when further improvement is unlikely.
>
> > [!NOTE]
> >
> > Don't call `study.optimize()` in your optimizer function. The scheduler uses Optuna's `ask()` and `tell()` methods to request trials and report results. This integration doesn't expose the `Study.optimize()` callback interface.

> [tab] Ax
>
> To customize the optimizer, define a function that accepts one `SweepInfo` argument and returns an Ax `Client`.
>
> Import Ax helper functions to configure the client's search space and objectives from the sweep configuration:
>
> | Helper function | Description |
> | - | - |
> | `sweep_config_to_search_space(config)` | Converts `parameters` into the parameter configurations accepted by `Client.configure_experiment(parameters=...)`. |
> | `configure_sweep_objective(client, config)` | Configures optimization from `metric` or `metrics`. It preserves metric names containing characters such as hyphens, which Ax's expression parser can misinterpret. We recommend this helper over `Client.configure_optimization()` for these names. |
>
> Mismatched objective names or directions fail validation when the scheduler starts.
>
> For example, the following function creates an Ax client and configures it from `sweep_info.config`:
>
> ```python title="optimizer.py" theme={"system"}
> from ax.api.client import Client
> from wandb.sdk.sweeps.scheduler.ax import (
>     configure_sweep_objective,
>     sweep_config_to_search_space,
> )
>
> def create_optimizer(sweep_info):
>     client = Client()
>     client.configure_experiment(
>         parameters=sweep_config_to_search_space(sweep_info.config)
>     )
>     configure_sweep_objective(client, sweep_info.config)
>     return client
> ```
>
> In `sweep.yaml`, set `scheduler.source` to `optimizer.py` and `scheduler.optimizer` to `create_optimizer`:
>
> ```yaml title="sweep.yaml" theme={"system"}
> program: train.py
> method: custom
> metric:
>   name: validation_loss
>   goal: minimize
> parameters:
>   learning_rate:
>     min: 0.0001
>     max: 0.01
>   batch_size:
>     values: [16, 32, 64]
> scheduler:
>   engine: ax
>   source: optimizer.py
>   optimizer: create_optimizer
> ```

For individual-run and sweep-level stopping options, see [Configure stopping](#configure-stopping).

### Define multiple objectives

If you use Optuna or Ax, you can use [`metrics`](https://docs.coreweave.com/products/wandb/sweeps/sweep-config-keys#metrics) (plural) to define multiple objectives.

When you create a sweep with `wandb sweep`, Weights & Biases stores the sweep configuration. When the scheduler starts, it retrieves that configuration and passes it to your custom optimizer function as part of a `SweepInfo` object.

Use `sweep_info.config` to access the stored sweep configuration.

> [!NOTE]
>
> The built-in Weights & Biases engine supports only a single objective, defined with `metric`.

You don't need a custom optimizer function to use multiple objectives. The following examples show how to configure the objectives when you provide one.

> [tab] Optuna
>
> For a multi-objective sweep with Optuna, define each objective in the `metrics` list. Use the `sweep_directions()` helper from `wandb.sdk.sweeps.scheduler.optuna` to convert the goals to Optuna study directions in the same order.
>
> For example, the following sweep minimizes validation loss and maximizes accuracy:
>
> ```yaml title="sweep.yaml" theme={"system"}
> name: optuna-multi-objective
> program: train.py
>
> scheduler:
>   engine: optuna
>   source: optimizer.py
>   optimizer: create_optimizer
>
> metrics:
>   - name: validation_loss
>     goal: minimize
>   - name: accuracy
>     goal: maximize
>
> parameters:
>   learning_rate:
>     min: 0.0001
>     max: 0.01
>   batch_size:
>     values: [16, 32, 64]
> ```
>
> The following demo script reads the suggested `learning_rate` and `batch_size` values, calculates two synthetic objective values, and logs them to the run. The final logged values appear in the run summary, where the scheduler can retrieve them.
>
> ```python title="train.py" theme={"system"}
> import math
> import wandb
>
> with wandb.init() as run:
>     learning_rate = run.config["learning_rate"]
>     batch_size = run.config["batch_size"]
>
>     # Replace these calculations with model training and evaluation.
>     validation_loss = (
>         abs(math.log10(learning_rate) + 3)
>         + 16 / batch_size
>     )
>     accuracy = min(
>         1.0,
>         0.8 - abs(math.log10(learning_rate) + 2.5) * 0.05 + batch_size / 1000,
>     )
>
>     run.log(
>         {
>             "validation_loss": validation_loss,
>             "accuracy": accuracy,
>         }
>     )
> ```
>
> Within your `optimizer.py` file, pass the stored sweep configuration from `sweep_info.config` to `sweep_directions()` when you create the Optuna study:
>
> ```python title="optimizer.py" theme={"system"}
> import optuna
> from wandb.sdk.sweeps.scheduler.optuna import sweep_directions
>
> def create_optimizer(sweep_info):
>     return optuna.create_study(
>         directions=sweep_directions(sweep_info.config),
>     )
> ```
>
> The resulting study minimizes `validation_loss` and maximizes `accuracy`.

> [tab] Ax
>
> For a multi-objective sweep with Ax, use the `configure_sweep_objective()` helper from `wandb.sdk.sweeps.scheduler.ax` to configure the client with the name and optimization direction of each entry in `metrics`.
>
> For example, the following sweep minimizes validation loss and maximizes accuracy:
>
> ```yaml title="sweep.yaml" theme={"system"}
> name: ax-multi-objective
> program: train.py
>
> scheduler:
>   engine: ax
>   source: optimizer.py
>   optimizer: create_optimizer
>
> metrics:
>   - name: validation_loss
>     goal: minimize
>   - name: accuracy
>     goal: maximize
>
> parameters:
>   learning_rate:
>     min: 0.0001
>     max: 0.01
>   batch_size:
>     values: [16, 32, 64]
> ```
>
> The following demo script reads the suggested `learning_rate` and `batch_size` values, calculates two synthetic objective values, and logs them to the run. The final logged values appear in the run summary, where the scheduler can retrieve them.
>
> ```python title="train.py" theme={"system"}
> import math
> import wandb
>
> with wandb.init() as run:
>     learning_rate = run.config["learning_rate"]
>     batch_size = run.config["batch_size"]
>
>     # Replace these calculations with model training and evaluation.
>     validation_loss = (
>         abs(math.log10(learning_rate) + 3)
>         + 16 / batch_size
>     )
>     accuracy = min(
>         1.0,
>         0.8 - abs(math.log10(learning_rate) + 2.5) * 0.05 + batch_size / 1000,
>     )
>
>     run.log(
>         {
>             "validation_loss": validation_loss,
>             "accuracy": accuracy,
>         }
>     )
> ```
>
> In `optimizer.py`, create an Ax client and configure its search space and objectives from the stored sweep configuration. Use `sweep_config_to_search_space()` to convert `parameters` into Ax parameter definitions, and `configure_sweep_objective()` to configure the objectives defined in `metrics`.
>
> ```python title="optimizer.py" theme={"system"}
> from ax.api.client import Client
> from wandb.sdk.sweeps.scheduler.ax import (
>     configure_sweep_objective,
>     sweep_config_to_search_space,
> )
>
>
> def create_optimizer(sweep_info):
>     client = Client()
>     client.configure_experiment(
>         parameters=sweep_config_to_search_space(sweep_info.config)
>     )
>     configure_sweep_objective(client, sweep_info.config)
>     return client
> ```
>
> In this example, `configure_sweep_objective()` configures Ax to minimize `validation_loss` and maximize `accuracy`.

### Define a conditional search space

> [!NOTE]
>
> Only the Optuna engine supports `scheduler.search_space`.

For Optuna, use `scheduler.search_space` to define conditional parameters in Python. For example, you can define different hyperparameters for different model types or vary the number of parameters based on earlier choices.

Set `scheduler.search_space` to the name of a function in the file specified by `scheduler.source`. The scheduler passes an Optuna trial to the function and uses the returned dictionary as the run configuration.

> [!TIP]
>
> For a fixed search space, define `parameters` in the sweep configuration.

When you set `scheduler.search_space`, the scheduler ignores the top-level `parameters` definitions. Define every parameter you want to pass to a run in the Python function.

For example, suppose you have the following sweep configuration:

```yaml title="sweep.yaml" theme={"system"}
metric:
  name: validation_loss
  goal: minimize

parameters: {}

scheduler:
  engine: optuna
  source: search.py
  search_space: define_search_space
```

For each trial, the scheduler calls `define_search_space(trial)` and uses the returned dictionary as the run configuration:

```python title="search.py" theme={"system"}
def define_search_space(trial):
    model = trial.suggest_categorical("model", ["linear", "transformer"])

    config = {
        "model": model,
        "learning_rate": trial.suggest_float(
            "learning_rate", 1e-5, 1e-2, log=True
        ),
    }

    if model == "transformer":
        config["num_layers"] = trial.suggest_int("num_layers", 2, 12)
        config["num_heads"] = trial.suggest_categorical(
            "num_heads", [4, 8, 16]
        )
    else:
        config["regularization"] = trial.suggest_float(
            "regularization", 0.0, 1.0
        )

    return config
```

Every run receives `model` and `learning_rate`. Runs with `model="transformer"` also receive `num_layers` and `num_heads`. Runs with `model="linear"` receive `regularization` instead.

## Configure stopping

Stopping can apply to an individual run or to the overall sweep. Individual-run early stopping ends runs that are unlikely to produce a useful result. Sweep-level stopping prevents the scheduler from proposing new runs after a condition is met.

### Stop individual runs

To evaluate a run’s progress, log intermediate values for the objective metric throughout training. Each engine provides its own mechanism for deciding whether to stop a run.

> [!NOTE]
>
> The top-level `early_terminate` key applies only to the built-in Weights & Biases engine. For Optuna or Ax, configure the engine’s native early-stopping mechanism.

> [tab] Weights & Biases
>
> Use [Hyperband](https://arxiv.org/abs/1603.06560) for hyperparameter optimization with early stopping based on intermediate results.
>
> Set `early_terminate.type` to `hyperband` and use `min_iter` to set the minimum number of iterations.
>
> For example, the following sweep configuration uses Hyperband for early stopping with a minimum of 3 iterations:
>
> ```yaml title="sweep.yaml" theme={"system"}
> program: train.py
> method: bayes
> metric:
>   name: validation_loss
>   goal: minimize
> parameters:
>   learning_rate:
>     min: 0.0001
>     max: 0.01
>   batch_size:
>     values: [16, 32, 64]
> early_terminate:
>   type: hyperband
>   min_iter: 3
> scheduler:
>   engine: wandb
> ```
>
> See [`early_terminate`](https://docs.coreweave.com/products/wandb/sweeps/sweep-config-keys#early_terminate) for configuration options.

> [tab] Optuna
>
> In Optuna, a [pruner](https://optuna.readthedocs.io/en/stable/reference/pruners.html) decides whether to stop a trial early based on intermediate results.
>
> For example, suppose you have the following sweep configuration and want to apply early stopping with Optuna:
>
> ```yaml title="sweep.yaml" theme={"system"}
> name: optuna-pruning-demo
> program: train.py
> method: grid
> run_cap: 2
>
> metric:
>   name: loss
>   goal: minimize
>
> parameters:
>   loss:
>     values: [0.2, 0.9]
>
> scheduler:
>   engine: optuna
>   source: optimizer.py
>   optimizer: make_study
> ```
>
> In `optimizer.py`, create a study with a grid sampler and a threshold pruner. The pruner recommends stopping trials whose intermediate loss exceeds 0.5:
>
> ```python title="optimizer.py" theme={"system"}
> import optuna
>
> def make_study(sweep_info):
>     """Create a study with a grid sampler and a threshold pruner."""
>     return optuna.create_study(
>         direction="minimize",
>         sampler=optuna.samplers.GridSampler(
>             {"loss": [0.2, 0.9]},
>         ),
>         # Stop any run whose intermediate loss exceeds 0.5.
>         pruner=optuna.pruners.ThresholdPruner(
>             upper=0.5,
>             n_warmup_steps=0,
>         ),
>     )
> ```
>
> In `train.py`, log the configured loss at each step so the scheduler can pass intermediate results to the pruner:
>
> ```python title="train.py" theme={"system"}
> import time
> import wandb
>
> with wandb.init() as run:
>     loss = run.config["loss"]
>
>     for step in range(30):
>         run.log({"loss": loss}, step=step)
>         time.sleep(1)
> ```
>
> > [!NOTE]
> >
> > Optuna pruning is not available for multi-objective studies.

> [tab] Ax
>
> To enable individual-run early stopping with Ax, call `client.set_early_stopping_strategy()` to configure an early-stopping strategy, then return the client.
>
> For example, suppose you have the following sweep configuration and want to apply early stopping with Ax:
>
> ```yaml title="sweep.yaml" theme={"system"}
> name: ax-run-stopping-demo
> program: train.py
>
> # Required by the sweep schema; Ax controls parameter generation.
> method: random
> run_cap: 2
> metric:
>   name: loss
>   goal: minimize
> parameters:
>   difficulty:
>     values: [0.1, 1.0]
> scheduler:
>   engine: ax
>   source: optimizer.py
>   optimizer: make_client
> ```
>
> In `optimizer.py`, configure the client with the search space, objective, and a percentile early-stopping strategy. Because the objective minimizes `loss`, the strategy uses higher loss values to identify underperforming runs.
>
> ```python title="optimizer.py" theme={"system"}
> from ax.api.client import Client
> from ax.early_stopping.strategies import PercentileEarlyStoppingStrategy
>
> from wandb.sdk.sweeps.scheduler.ax import (
>     configure_sweep_objective,
>     sweep_config_to_search_space,
> )
>
> def make_client(sweep_info):
>     client = Client()
>     client.configure_experiment(
>         parameters=sweep_config_to_search_space(sweep_info.config),
>     )
>     configure_sweep_objective(client, sweep_info.config)
>
>     client.set_early_stopping_strategy(
>         PercentileEarlyStoppingStrategy(
>             percentile_threshold=50.0,
>             min_progression=2,
>             min_curves=2,
>             seconds_between_polls=0,
>             normalize_progressions=False,
>         )
>     )
>
>     return client
> ```
>
> In `train.py`, calculate a synthetic loss from `difficulty` and log it at each step. These intermediate values let Ax compare the runs’ progress:
>
> ```python title="train.py" theme={"system"}
> import time
> import wandb
>
> with wandb.init() as run:
>     difficulty = run.config["difficulty"]
>
>     for step in range(30):
>         # The run with difficulty=1.0 is consistently worse.
>         loss = difficulty + 1.0 / (step + 1)
>
>         run.log(
>             {"loss": loss},
>             step=step,
>         )
>
>         # Keep the runs alive long enough for the scheduler to poll them.
>         time.sleep(1)
> ```
>
> Run two sweep agents concurrently so the scheduler can collect intermediate results from both runs. When Ax recommends stopping a run, the scheduler requests that Weights & Biases stop it.

### Stop the sweep

Sweep-level stopping controls when the scheduler stops proposing new runs. Use `run_cap` to set a fixed maximum number of runs for any engine. For Optuna or Ax, you can also return a stopping function that evaluates the current state of the optimization.

> [!NOTE]
>
> In the following examples, the scheduler controls the search strategy, so it ignores `method: random`. However, the Weights & Biases sweep schema requires `method`.

The scheduler might already have runs queued or running when a stopping condition is met. Those runs can continue after the scheduler stops proposing new ones.

> [tab] Optuna
>
> To enable sweep-level stopping, the optimization function you specify in `scheduler.optimizer` must return an Optuna study and a stopping function as a tuple:
>
> ```python theme={"system"}
> (study, terminator_callback)
> ```
>
> After processing results, the scheduler passes the current study to the stopping function. The function returns `True` to stop the sweep or `False` to continue proposing new runs.
>
> For example, the following sweep configuration searches for `x`, limits the sweep to 20 runs, and instructs the Optuna scheduler to call `make_study()` from the `optimizer.py` file:
>
> ```yaml title="sweep.yaml" theme={"system"}
> name: optuna-sweep-stopping-demo
> program: train.py
> method: random
> run_cap: 20
>
> metric:
>   name: loss
>   goal: minimize
>
> parameters:
>   x:
>     min: 0
>     max: 10
>
> scheduler:
>   engine: optuna
>   source: optimizer.py
>   optimizer: make_study
> ```
>
> In this example, `run_cap` provides a fallback limit if the stopping function doesn't end the sweep first.
>
> Next, define the Optuna study and stopping function. `make_study()` configures a seeded random sampler and returns the study with `stop_after_three_completed_trials()`:
>
> ```python title="optimizer.py" theme={"system"}
> import optuna
>
>
> def stop_after_three_completed_trials(study: optuna.Study) -> bool:
>     """Stop the sweep after three completed trials."""
>     completed = sum(
>         trial.state == optuna.trial.TrialState.COMPLETE
>         for trial in study.trials
>     )
>
>     return completed >= 3
>
>
> def make_study(sweep_info):
>     study = optuna.create_study(
>         direction="minimize",
>         sampler=optuna.samplers.RandomSampler(seed=42),
>     )
>
>     return study, stop_after_three_completed_trials
> ```
>
> The stopping function counts only successfully completed trials. It returns `True` after the third completed trial, which tells the scheduler to stop proposing configurations. Failed, pruned, and running trials don't contribute to the count.
>
> Finally, define the objective that each sweep run evaluates. The training script retrieves the suggested value of `x`, calculates its loss, and logs that loss to Weights & Biases:
>
> ```python title="train.py" theme={"system"}
> import wandb
>
> with wandb.init() as run:
>     x = run.config["x"]
>     loss = (x - 3) ** 2
>
>     run.log({"loss": loss})
> ```
>
> After each run finishes, the scheduler passes its logged loss to the Optuna study. Once three trials complete, the stopping function returns `True` and the scheduler stops proposing new runs. Runs that are already queued or running can still finish.

> [tab] Ax
>
> To enable sweep-level stopping, return a tuple containing the Ax client and a stopping function from the function specified by `scheduler.optimizer`:
>
> ```python theme={"system"}
> (client, terminator)
> ```
>
> The following configuration limits the sweep to 20 runs. The stopping function can end it earlier if the objective stops improving:
>
> ```yaml title="sweep.yaml" theme={"system"}
> name: ax-sweep-stopping-example
> program: train.py
> method: random
> run_cap: 20
> metric:
>   name: loss
>   goal: minimize
> parameters:
>   x:
>     min: 0.0
>     max: 1.0
> scheduler:
>   engine: ax
>   source: optimizer.py
>   optimizer: make_client
> ```
>
> In `optimizer.py`, create an Ax client and a stopping function. This example uses `ImprovementGlobalStoppingStrategy` to compare recent improvement with the spread of observed objective values:
>
> ```python title="optimizer.py" theme={"system"}
> from ax.api.client import Client
> from ax.global_stopping.strategies.improvement import (
>     ImprovementGlobalStoppingStrategy,
> )
>
> from wandb.sdk.sweeps.scheduler.ax import (
>     configure_sweep_objective,
>     sweep_config_to_search_space,
> )
>
>
> def make_client(sweep_info):
>     """Create the Ax client and sweep-level stopping callback."""
>
>     client = Client()
>
>     # Convert the parameters in sweep.yaml into an Ax search space.
>     client.configure_experiment(
>         parameters=sweep_config_to_search_space(sweep_info.config),
>     )
>
>     # Configure "loss" as the objective Ax should minimize.
>     configure_sweep_objective(client, sweep_info.config)
>
>     client.configure_generation_strategy(
>         method="random_search",
>         initialization_random_seed=7,
>         initialize_with_center=True,
>     )
>
>     stopping_strategy = ImprovementGlobalStoppingStrategy(
>         min_trials=6,
>         window_size=3,
>         improvement_bar=0.01,
>         inactive_when_pending_trials=True,
>     )
>
>     def should_terminate_sweep(client):
>         """Return True when Ax recommends ending the whole sweep."""
>         should_stop, reason = stopping_strategy.should_stop_optimization(
>             experiment=client._experiment,
>         )
>
>         if should_stop:
>             print(f"Ax is stopping the sweep: {reason}")
>
>         return should_stop
>
>     return client, should_terminate_sweep
> ```
>
> The following training script evaluates a synthetic objective whose minimum occurs at `x=0.5`. The generation strategy evaluates that point first, so subsequent trials cannot improve the best result:
>
> ```python title="train.py" theme={"system"}
> import time
> import wandb
>
> with wandb.init() as run:
>     x = run.config["x"]
>
>     for step in range(5):
>         loss = (x - 0.5) ** 2 + 1.0 / (step + 1)
>
>         run.log(
>             {"loss": loss},
>             step=step,
>         )
>
>         time.sleep(0.5)
> ```
>
> This stopping strategy uses only the final objective value from each completed run. The script logs intermediate values to simulate training; this example does not configure individual-run early stopping.

## Handle scheduler restarts

The sweep scheduler operates on an existing sweep, so you can stop and restart the scheduler without creating a new sweep. Each time the scheduler starts, it reconstructs the optimization history from the sweep’s existing run data and then continues the optimization loop. It does not restore the engine from a saved optimizer checkpoint.

When the scheduler starts or restarts, it calls your custom optimizer function to create a new engine instance. Return a fresh Optuna study or Ax client each time the function runs. Don't load a separate optimizer checkpoint or reuse an engine instance from a previous scheduler process.

If you use a `scheduler.source` file, the scheduler loads its current contents each time it starts. You can update the file between starts to fix bugs or modify the custom `optimizer` or `search_space` function. These changes can affect subsequent suggestions, so preserve the version of the `source` file used for each scheduler session if you need to reproduce the optimization.