chore: import upstream snapshot with attribution

doc/source/tune/doc_code/faq.py

@@ -0,0 +1,352 @@
+# flake8: noqa
+
+# __reproducible_start__
+import numpy as np
+from ray import tune
+
+
+def train_func(config):
+    # Set seed for trainable random result.
+    # If you remove this line, you will get different results
+    # each time you run the trial, even if the configuration
+    # is the same.
+    np.random.seed(config["seed"])
+    random_result = np.random.uniform(0, 100, size=1).item()
+    tune.report({"result": random_result})
+
+
+# Set seed for Ray Tune's random search.
+# If you remove this line, you will get different configurations
+# each time you run the script.
+np.random.seed(1234)
+tuner = tune.Tuner(
+    train_func,
+    tune_config=tune.TuneConfig(
+        num_samples=10,
+        search_alg=tune.search.BasicVariantGenerator(),
+    ),
+    param_space={"seed": tune.randint(0, 1000)},
+)
+tuner.fit()
+# __reproducible_end__
+
+# __basic_config_start__
+config = {"a": {"x": tune.uniform(0, 10)}, "b": tune.choice([1, 2, 3])}
+# __basic_config_end__
+
+# __conditional_spaces_start__
+config = {
+    "a": tune.randint(5, 10),
+    "b": tune.sample_from(lambda config: np.random.randint(0, config["a"])),
+}
+# __conditional_spaces_end__
+
+
+# __iter_start__
+def _iter():
+    for a in range(5, 10):
+        for b in range(a):
+            yield a, b
+
+
+config = {
+    "ab": tune.grid_search(list(_iter())),
+}
+# __iter_end__
+
+
+def train_func(config):
+    random_result = np.random.uniform(0, 100, size=1).item()
+    tune.report({"result": random_result})
+
+
+train_fn = train_func
+MOCK = True
+# Note we put this check here to make sure at least the syntax of
+# the code is correct. Some of these snippets simply can't be run on the nose.
+
+if not MOCK:
+    # __resources_start__
+    tuner = tune.Tuner(
+        tune.with_resources(
+            train_fn, resources={"cpu": 2, "gpu": 0.5, "custom_resources": {"hdd": 80}}
+        ),
+    )
+    tuner.fit()
+    # __resources_end__
+
+    # __resources_pgf_start__
+    tuner = tune.Tuner(
+        tune.with_resources(
+            train_fn,
+            resources=tune.PlacementGroupFactory(
+                [
+                    {"CPU": 2, "GPU": 0.5, "hdd": 80},
+                    {"CPU": 1},
+                    {"CPU": 1},
+                ],
+                strategy="PACK",
+            ),
+        )
+    )
+    tuner.fit()
+    # __resources_pgf_end__
+
+    # __resources_lambda_start__
+    tuner = tune.Tuner(
+        tune.with_resources(
+            train_fn,
+            resources=lambda config: {"GPU": 1} if config["use_gpu"] else {"GPU": 0},
+        ),
+        param_space={
+            "use_gpu": True,
+        },
+    )
+    tuner.fit()
+    # __resources_lambda_end__
+
+    metric = None
+
+    # __modin_start__
+    def train_fn(config):
+        # some Modin operations here
+        # import modin.pandas as pd
+        tune.report({"metric": metric})
+
+    tuner = tune.Tuner(
+        tune.with_resources(
+            train_fn,
+            resources=tune.PlacementGroupFactory(
+                [
+                    {"CPU": 1},  # this bundle will be used by the trainable itself
+                    {"CPU": 1},  # this bundle will be used by Modin
+                ],
+                strategy="PACK",
+            ),
+        )
+    )
+    tuner.fit()
+# __modin_end__
+
+# __huge_data_start__
+from ray import tune
+import numpy as np
+
+
+def train_func(config, num_epochs=5, data=None):
+    for i in range(num_epochs):
+        for sample in data:
+            # ... train on sample
+            pass
+
+
+# Some huge dataset
+data = np.random.random(size=100000000)
+
+tuner = tune.Tuner(tune.with_parameters(train_func, num_epochs=5, data=data))
+tuner.fit()
+# __huge_data_end__
+
+
+# __seeded_1_start__
+import random
+
+random.seed(1234)
+output = [random.randint(0, 100) for _ in range(10)]
+
+# The output will always be the same.
+assert output == [99, 56, 14, 0, 11, 74, 4, 85, 88, 10]
+# __seeded_1_end__
+
+
+# __seeded_2_start__
+# This should suffice to initialize the RNGs for most Python-based libraries
+import random
+import numpy as np
+
+random.seed(1234)
+np.random.seed(5678)
+# __seeded_2_end__
+
+
+# __torch_tf_seeds_start__
+import torch
+
+torch.manual_seed(0)
+
+import tensorflow as tf
+
+tf.random.set_seed(0)
+# __torch_tf_seeds_end__
+
+# __torch_seed_example_start__
+import random
+import numpy as np
+from ray import tune
+
+
+def trainable(config):
+    # config["seed"] is set deterministically, but differs between training runs
+    random.seed(config["seed"])
+    np.random.seed(config["seed"])
+    # torch.manual_seed(config["seed"])
+    # ... training code
+
+
+config = {
+    "seed": tune.randint(0, 10000),
+    # ...
+}
+
+if __name__ == "__main__":
+    # Set seed for the search algorithms/schedulers
+    random.seed(1234)
+    np.random.seed(1234)
+    # Don't forget to check if the search alg has a `seed` parameter
+    tuner = tune.Tuner(trainable, param_space=config)
+    tuner.fit()
+# __torch_seed_example_end__
+
+# __large_data_start__
+from ray import tune
+import numpy as np
+
+
+def f(config, data=None):
+    pass
+    # use data
+
+
+data = np.random.random(size=100000000)
+
+tuner = tune.Tuner(tune.with_parameters(f, data=data))
+tuner.fit()
+# __large_data_end__
+
+
+import ray
+
+ray.shutdown()
+
+# __grid_search_start__
+parameters = {
+    "qux": tune.sample_from(lambda spec: 2 + 2),
+    "bar": tune.grid_search([True, False]),
+    "foo": tune.grid_search([1, 2, 3]),
+    "baz": "asd",  # a constant value
+}
+
+tuner = tune.Tuner(train_fn, param_space=parameters)
+tuner.fit()
+# __grid_search_end__
+
+# __grid_search_2_start__
+# num_samples=10 repeats the 3x3 grid search 10 times, for a total of 90 trials
+tuner = tune.Tuner(
+    train_fn,
+    run_config=tune.RunConfig(name="my_trainable"),
+    param_space={
+        "alpha": tune.uniform(100, 200),
+        "beta": tune.sample_from(lambda config: config["alpha"] * np.random.normal()),
+        "nn_layers": [
+            tune.grid_search([16, 64, 256]),
+            tune.grid_search([16, 64, 256]),
+        ],
+    },
+    tune_config=tune.TuneConfig(num_samples=10),
+)
+# __grid_search_2_end__
+
+if not MOCK:
+    import os
+    from pathlib import Path
+
+    # __no_chdir_start__
+    def train_func(config):
+        # Read from relative paths
+        print(open("./read.txt").read())
+
+        # The working directory shouldn't have changed from the original
+        # NOTE: The `TUNE_ORIG_WORKING_DIR` environment variable is deprecated.
+        assert os.getcwd() == os.environ["TUNE_ORIG_WORKING_DIR"]
+
+        # Write to the Tune trial directory, not the shared working dir
+        tune_trial_dir = Path(ray.tune.get_context().get_trial_dir())
+        with open(tune_trial_dir / "write.txt", "w") as f:
+            f.write("trial saved artifact")
+
+    os.environ["RAY_CHDIR_TO_TRIAL_DIR"] = "0"
+    tuner = tune.Tuner(train_func)
+    tuner.fit()
+    # __no_chdir_end__
+
+
+# __iter_experimentation_initial_start__
+import os
+import tempfile
+
+import torch
+
+from ray import tune
+from ray.tune import Checkpoint
+import random
+
+
+def trainable(config):
+    for epoch in range(1, config["num_epochs"]):
+        # Do some training...
+
+        with tempfile.TemporaryDirectory() as tempdir:
+            torch.save(
+                {"model_state_dict": {"x": 1}}, os.path.join(tempdir, "model.pt")
+            )
+            tune.report(
+                {"score": random.random()},
+                checkpoint=Checkpoint.from_directory(tempdir),
+            )
+
+
+tuner = tune.Tuner(
+    trainable,
+    param_space={"num_epochs": 10, "hyperparam": tune.grid_search([1, 2, 3])},
+    tune_config=tune.TuneConfig(metric="score", mode="max"),
+)
+result_grid = tuner.fit()
+
+best_result = result_grid.get_best_result()
+best_checkpoint = best_result.checkpoint
+# __iter_experimentation_initial_end__
+
+
+# __iter_experimentation_resume_start__
+import ray
+
+
+def trainable(config):
+    # Add logic to handle the initial checkpoint.
+    checkpoint: Checkpoint = config["start_from_checkpoint"]
+    with checkpoint.as_directory() as checkpoint_dir:
+        model_state_dict = torch.load(os.path.join(checkpoint_dir, "model.pt"))
+
+    # Initialize a model from the checkpoint...
+    # model = ...
+    # model.load_state_dict(model_state_dict)
+
+    for epoch in range(1, config["num_epochs"]):
+        # Do some more training...
+        ...
+
+        tune.report({"score": random.random()})
+
+
+new_tuner = tune.Tuner(
+    trainable,
+    param_space={
+        "num_epochs": 10,
+        "hyperparam": tune.grid_search([4, 5, 6]),
+        "start_from_checkpoint": best_checkpoint,
+    },
+    tune_config=tune.TuneConfig(metric="score", mode="max"),
+)
+result_grid = new_tuner.fit()
+# __iter_experimentation_resume_end__

doc/source/tune/doc_code/fault_tolerance.py

@@ -0,0 +1,162 @@
+# flake8: noqa
+
+# __ft_initial_run_start__
+import json
+import os
+import tempfile
+
+from ray import tune
+
+
+def trainable(config):
+    # Checkpoint loading
+    checkpoint = tune.get_checkpoint()
+    start = 1
+    if checkpoint:
+        with checkpoint.as_directory() as checkpoint_dir:
+            with open(os.path.join(checkpoint_dir, "checkpoint.json"), "r") as f:
+                state = json.load(f)
+        start = state["epoch"] + 1
+
+    for epoch in range(start, config["num_epochs"]):
+        # Do some training...
+
+        # Checkpoint saving
+        with tempfile.TemporaryDirectory() as temp_checkpoint_dir:
+            with open(os.path.join(temp_checkpoint_dir, "checkpoint.json"), "w") as f:
+                json.dump({"epoch": epoch}, f)
+            tune.report(
+                {"epoch": epoch},
+                checkpoint=tune.Checkpoint.from_directory(temp_checkpoint_dir),
+            )
+
+
+tuner = tune.Tuner(
+    trainable,
+    param_space={"num_epochs": 10},
+    run_config=tune.RunConfig(
+        storage_path=os.path.expanduser("~/ray_results"),
+        name="tune_fault_tolerance_guide",
+    ),
+)
+result_grid = tuner.fit()
+# __ft_initial_run_end__
+
+assert not result_grid.errors
+
+# __ft_restored_run_start__
+tuner = tune.Tuner.restore(
+    os.path.expanduser("~/ray_results/tune_fault_tolerance_guide"),
+    trainable=trainable,
+    resume_errored=True,
+)
+tuner.fit()
+# __ft_restored_run_end__
+
+# __ft_restore_options_start__
+tuner = tune.Tuner.restore(
+    os.path.expanduser("~/ray_results/tune_fault_tolerance_guide"),
+    trainable=trainable,
+    resume_errored=True,
+    restart_errored=False,
+    resume_unfinished=True,
+)
+# __ft_restore_options_end__
+
+# __ft_restore_multiplexing_start__
+import os
+from ray import tune
+
+storage_path = os.path.expanduser("~/ray_results")
+exp_name = "tune_fault_tolerance_guide"
+path = os.path.join(storage_path, exp_name)
+
+if tune.Tuner.can_restore(path):
+    tuner = tune.Tuner.restore(path, trainable=trainable, resume_errored=True)
+else:
+    tuner = tune.Tuner(
+        trainable,
+        param_space={"num_epochs": 10},
+        run_config=tune.RunConfig(storage_path=storage_path, name=exp_name),
+    )
+tuner.fit()
+# __ft_restore_multiplexing_end__
+
+
+# Run the multiplexed logic again to make sure it goes through the restore branch.
+if tune.Tuner.can_restore(path):
+    tuner = tune.Tuner.restore(path, trainable=trainable, resume_errored=True)
+else:
+    tuner = tune.Tuner(
+        trainable,
+        param_space={"num_epochs": 10},
+        run_config=tune.RunConfig(storage_path=storage_path, name=exp_name),
+    )
+assert tuner.get_results()
+
+
+# __ft_restore_objrefs_initial_start__
+import ray
+from ray import tune
+
+
+class LargeModel:
+    def __init__(self, model_id):
+        self.model_id = model_id
+        # Load weights based on the `model_id`...
+
+
+def train_fn(config):
+    # Retrieve the model from the object store.
+    model = ray.get(config["model_ref"])
+    print(model.model_id)
+
+
+# These models may be large, so `ray.put` them in the Ray Object Store
+# to share the models between trials.
+model_refs = [ray.put(LargeModel(1)), ray.put(LargeModel(2))]
+
+tuner = tune.Tuner(
+    train_fn,
+    # Tune over the object references!
+    param_space={"model_ref": tune.grid_search(model_refs)},
+    run_config=tune.RunConfig(
+        storage_path=os.path.expanduser("~/ray_results"), name="restore_object_refs"
+    ),
+)
+tuner.fit()
+# __ft_restore_objrefs_initial_end__
+
+if ray.is_initialized():
+    ray.shutdown()
+
+# __ft_restore_objrefs_restored_start__
+# Re-create the objects and put them in the object store.
+param_space = {
+    "model_ref": tune.grid_search([ray.put(LargeModel(1)), ray.put(LargeModel(2))])
+}
+
+tuner = tune.Tuner.restore(
+    os.path.expanduser("~/ray_results/restore_object_refs"),
+    trainable=train_fn,
+    # Re-specify the `param_space` to update the object references.
+    param_space=param_space,
+    resume_errored=True,
+)
+tuner.fit()
+# __ft_restore_objrefs_restored_end__
+
+# __ft_trial_failure_start__
+from ray import tune
+
+tuner = tune.Tuner(
+    trainable,
+    param_space={"num_epochs": 10},
+    run_config=tune.RunConfig(
+        storage_path=os.path.expanduser("~/ray_results"),
+        name="trial_fault_tolerance",
+        failure_config=tune.FailureConfig(max_failures=3),
+    ),
+)
+tuner.fit()
+# __ft_trial_failure_end__

doc/source/tune/doc_code/keras_hyperopt.py

@@ -0,0 +1,35 @@
+# flake8: noqa
+
+accuracy = 42
+
+# __keras_hyperopt_start__
+from ray import tune
+from ray.tune.search.hyperopt import HyperOptSearch
+import keras
+
+
+def objective(config):  # <1>
+    model = keras.models.Sequential()
+    model.add(keras.layers.Dense(784, activation=config["activation"]))
+    model.add(keras.layers.Dense(10, activation="softmax"))
+
+    model.compile(loss="binary_crossentropy", optimizer="adam", metrics=["accuracy"])
+    # model.fit(...)
+    # loss, accuracy = model.evaluate(...)
+    return {"accuracy": accuracy}
+
+
+search_space = {"activation": tune.choice(["relu", "tanh"])}  # <2>
+algo = HyperOptSearch()
+
+tuner = tune.Tuner(  # <3>
+    objective,
+    tune_config=tune.TuneConfig(
+        metric="accuracy",
+        mode="max",
+        search_alg=algo,
+    ),
+    param_space=search_space,
+)
+results = tuner.fit()
+# __keras_hyperopt_end__

doc/source/tune/doc_code/key_concepts.py

@@ -0,0 +1,162 @@
+# flake8: noqa
+
+# __function_api_start__
+from ray import tune
+
+
+def objective(x, a, b):  # Define an objective function.
+    return a * (x**2) + b
+
+
+def trainable(config):  # Pass a "config" dictionary into your trainable.
+
+    for x in range(20):  # "Train" for 20 iterations and compute intermediate scores.
+        score = objective(x, config["a"], config["b"])
+
+        tune.report({"score": score})  # Send the score to Tune.
+
+
+# __function_api_end__
+
+
+# __class_api_start__
+from ray import tune
+
+
+def objective(x, a, b):
+    return a * (x**2) + b
+
+
+class Trainable(tune.Trainable):
+    def setup(self, config):
+        # config (dict): A dict of hyperparameters
+        self.x = 0
+        self.a = config["a"]
+        self.b = config["b"]
+
+    def step(self):  # This is called iteratively.
+        score = objective(self.x, self.a, self.b)
+        self.x += 1
+        return {"score": score}
+
+    # __class_api_end__
+
+    # TODO: this example does not work as advertised. Errors out.
+    def save_checkpoint(self, checkpoint_dir):
+        pass
+
+    def load_checkpoint(self, checkpoint_dir):
+        pass
+
+
+# __run_tunable_start__
+# Pass in a Trainable class or function, along with a search space "config".
+tuner = tune.Tuner(trainable, param_space={"a": 2, "b": 4})
+tuner.fit()
+# __run_tunable_end__
+
+# __run_tunable_samples_start__
+tuner = tune.Tuner(
+    trainable, param_space={"a": 2, "b": 4}, tune_config=tune.TuneConfig(num_samples=10)
+)
+tuner.fit()
+# __run_tunable_samples_end__
+
+# __search_space_start__
+space = {"a": tune.uniform(0, 1), "b": tune.uniform(0, 1)}
+tuner = tune.Tuner(
+    trainable, param_space=space, tune_config=tune.TuneConfig(num_samples=10)
+)
+tuner.fit()
+# __search_space_end__
+
+# __config_start__
+config = {
+    "uniform": tune.uniform(-5, -1),  # Uniform float between -5 and -1
+    "quniform": tune.quniform(3.2, 5.4, 0.2),  # Round to multiples of 0.2
+    "loguniform": tune.loguniform(1e-4, 1e-1),  # Uniform float in log space
+    "qloguniform": tune.qloguniform(1e-4, 1e-1, 5e-5),  # Round to multiples of 0.00005
+    "randn": tune.randn(10, 2),  # Normal distribution with mean 10 and sd 2
+    "qrandn": tune.qrandn(10, 2, 0.2),  # Round to multiples of 0.2
+    "randint": tune.randint(-9, 15),  # Random integer between -9 and 15
+    "qrandint": tune.qrandint(-21, 12, 3),  # Round to multiples of 3 (includes 12)
+    "lograndint": tune.lograndint(1, 10),  # Random integer in log space
+    "qlograndint": tune.qlograndint(1, 10, 2),  # Round to multiples of 2
+    "choice": tune.choice(["a", "b", "c"]),  # Choose one of these options uniformly
+    "func": tune.sample_from(
+        lambda config: config["uniform"] * 0.01
+    ),  # Depends on other value
+    "grid": tune.grid_search([32, 64, 128]),  # Search over all these values
+}
+# __config_end__
+
+# __bayes_start__
+from ray.tune.search.bayesopt import BayesOptSearch
+
+# Define the search space
+search_space = {"a": tune.uniform(0, 1), "b": tune.uniform(0, 20)}
+
+algo = BayesOptSearch(random_search_steps=4)
+
+tuner = tune.Tuner(
+    trainable,
+    tune_config=tune.TuneConfig(
+        metric="score",
+        mode="min",
+        search_alg=algo,
+    ),
+    run_config=tune.RunConfig(stop={"training_iteration": 20}),
+    param_space=search_space,
+)
+tuner.fit()
+# __bayes_end__
+
+# __hyperband_start__
+from ray.tune.schedulers import HyperBandScheduler
+
+# Create HyperBand scheduler and minimize the score
+hyperband = HyperBandScheduler(metric="score", mode="max")
+
+config = {"a": tune.uniform(0, 1), "b": tune.uniform(0, 1)}
+
+tuner = tune.Tuner(
+    trainable,
+    tune_config=tune.TuneConfig(
+        num_samples=20,
+        scheduler=hyperband,
+    ),
+    param_space=config,
+)
+tuner.fit()
+# __hyperband_end__
+
+# __analysis_start__
+tuner = tune.Tuner(
+    trainable,
+    tune_config=tune.TuneConfig(
+        metric="score",
+        mode="min",
+        search_alg=BayesOptSearch(random_search_steps=4),
+    ),
+    run_config=tune.RunConfig(
+        stop={"training_iteration": 20},
+    ),
+    param_space=config,
+)
+results = tuner.fit()
+
+best_result = results.get_best_result()  # Get best result object
+best_config = best_result.config  # Get best trial's hyperparameters
+best_logdir = best_result.path  # Get best trial's result directory
+best_checkpoint = best_result.checkpoint  # Get best trial's best checkpoint
+best_metrics = best_result.metrics  # Get best trial's last results
+best_result_df = best_result.metrics_dataframe  # Get best result as pandas dataframe
+# __analysis_end__
+
+# __results_start__
+# Get a dataframe with the last results for each trial
+df_results = results.get_dataframe()
+
+# Get a dataframe of results for a specific score or mode
+df = results.get_dataframe(filter_metric="score", filter_mode="max")
+# __results_end__

doc/source/tune/doc_code/pytorch_optuna.py

@@ -0,0 +1,116 @@
+# flake8: noqa
+
+import os
+from filelock import FileLock
+import torch.nn as nn
+import torch.nn.functional as F
+from torchvision import datasets, transforms
+
+EPOCH_SIZE = 512
+TEST_SIZE = 256
+
+
+def train_epoch(model, optimizer, train_loader, device=None):
+    device = device or torch.device("cpu")
+    model.train()
+    for batch_idx, (data, target) in enumerate(train_loader):
+        if batch_idx * len(data) > EPOCH_SIZE:
+            return
+        data, target = data.to(device), target.to(device)
+        optimizer.zero_grad()
+        output = model(data)
+        loss = F.nll_loss(output, target)
+        loss.backward()
+        optimizer.step()
+
+
+def test(model, data_loader, device=None):
+    device = device or torch.device("cpu")
+    model.eval()
+    correct = 0
+    total = 0
+    with torch.no_grad():
+        for batch_idx, (data, target) in enumerate(data_loader):
+            if batch_idx * len(data) > TEST_SIZE:
+                break
+            data, target = data.to(device), target.to(device)
+            outputs = model(data)
+            _, predicted = torch.max(outputs.data, 1)
+            total += target.size(0)
+            correct += (predicted == target).sum().item()
+
+    return correct / total
+
+
+def load_data():
+    mnist_transforms = transforms.Compose(
+        [transforms.ToTensor(), transforms.Normalize((0.1307,), (0.3081,))]
+    )
+    with FileLock(os.path.expanduser("~/data.lock")):
+        train_loader = torch.utils.data.DataLoader(
+            datasets.MNIST(
+                "~/data", train=True, download=True, transform=mnist_transforms
+            ),
+            batch_size=64,
+            shuffle=True,
+        )
+        test_loader = torch.utils.data.DataLoader(
+            datasets.MNIST(
+                "~/data", train=False, download=True, transform=mnist_transforms
+            ),
+            batch_size=64,
+            shuffle=True,
+        )
+    return train_loader, test_loader
+
+
+class ConvNet(nn.Module):
+    def __init__(self):
+        super(ConvNet, self).__init__()
+        self.conv1 = nn.Conv2d(1, 3, kernel_size=3)
+        self.fc = nn.Linear(192, 10)
+
+    def forward(self, x):
+        x = F.relu(F.max_pool2d(self.conv1(x), 3))
+        x = x.view(-1, 192)
+        x = self.fc(x)
+        return F.log_softmax(x, dim=1)
+
+
+# __pytorch_optuna_start__
+import torch
+from ray import tune
+from ray.tune.search.optuna import OptunaSearch
+
+
+def objective(config):  # <1>
+    train_loader, test_loader = load_data()  # Load some data
+    model = ConvNet().to("cpu")  # Create a PyTorch conv net
+    optimizer = torch.optim.SGD(  # Tune the optimizer
+        model.parameters(), lr=config["lr"], momentum=config["momentum"]
+    )
+
+    while True:
+        train_epoch(model, optimizer, train_loader)  # Train the model
+        acc = test(model, test_loader)  # Compute test accuracy
+        tune.report({"mean_accuracy": acc})  # Report to Tune
+
+
+search_space = {"lr": tune.loguniform(1e-4, 1e-2), "momentum": tune.uniform(0.1, 0.9)}
+algo = OptunaSearch()  # <2>
+
+tuner = tune.Tuner(  # <3>
+    objective,
+    tune_config=tune.TuneConfig(
+        metric="mean_accuracy",
+        mode="max",
+        search_alg=algo,
+    ),
+    run_config=tune.RunConfig(
+        stop={"training_iteration": 5},
+    ),
+    param_space=search_space,
+)
+results = tuner.fit()
+print("Best config is:", results.get_best_result().config)
+# __pytorch_optuna_end__

doc/source/tune/doc_code/stopping.py

@@ -0,0 +1,169 @@
+# flake8: noqa
+
+# fmt: off
+# __stopping_example_trainable_start__
+from ray import tune
+import time
+
+def my_trainable(config):
+    i = 1
+    while True:
+        # Do some training...
+        time.sleep(1)
+
+        # Report some metrics for demonstration...
+        tune.report({"mean_accuracy": min(i / 10, 1.0)})
+        i += 1
+# __stopping_example_trainable_end__
+# fmt: on
+
+
+def my_trainable(config):
+    # NOTE: This re-defines the training loop with the sleep removed for faster testing.
+    i = 1
+    # Training won't finish unless one of the stopping criteria is met!
+    while True:
+        # Do some training, and report some metrics for demonstration...
+        tune.report({"mean_accuracy": min(i / 10, 1.0)})
+        i += 1
+
+
+# __stopping_dict_start__
+from ray import tune
+
+tuner = tune.Tuner(
+    my_trainable,
+    run_config=tune.RunConfig(stop={"training_iteration": 10, "mean_accuracy": 0.8}),
+)
+result_grid = tuner.fit()
+# __stopping_dict_end__
+
+final_iter = result_grid[0].metrics["training_iteration"]
+assert final_iter == 8, final_iter
+
+# __stopping_fn_start__
+from ray import tune
+
+
+def stop_fn(trial_id: str, result: dict) -> bool:
+    return result["mean_accuracy"] >= 0.8 or result["training_iteration"] >= 10
+
+
+tuner = tune.Tuner(my_trainable, run_config=tune.RunConfig(stop=stop_fn))
+result_grid = tuner.fit()
+# __stopping_fn_end__
+
+final_iter = result_grid[0].metrics["training_iteration"]
+assert final_iter == 8, final_iter
+
+# __stopping_cls_start__
+from ray import tune
+from ray.tune import Stopper
+
+
+class CustomStopper(Stopper):
+    def __init__(self):
+        self.should_stop = False
+
+    def __call__(self, trial_id: str, result: dict) -> bool:
+        if not self.should_stop and result["mean_accuracy"] >= 0.8:
+            self.should_stop = True
+        return self.should_stop
+
+    def stop_all(self) -> bool:
+        """Returns whether to stop trials and prevent new ones from starting."""
+        return self.should_stop
+
+
+stopper = CustomStopper()
+tuner = tune.Tuner(
+    my_trainable,
+    run_config=tune.RunConfig(stop=stopper),
+    tune_config=tune.TuneConfig(num_samples=2),
+)
+result_grid = tuner.fit()
+# __stopping_cls_end__
+
+for result in result_grid:
+    final_iter = result.metrics.get("training_iteration", 0)
+    assert final_iter <= 8, final_iter
+
+# __stopping_on_trial_error_start__
+from ray import tune
+import time
+
+
+def my_failing_trainable(config):
+    if config["should_fail"]:
+        raise RuntimeError("Failing (on purpose)!")
+    # Do some training...
+    time.sleep(10)
+    tune.report({"mean_accuracy": 0.9})
+
+
+tuner = tune.Tuner(
+    my_failing_trainable,
+    param_space={"should_fail": tune.grid_search([True, False])},
+    run_config=tune.RunConfig(failure_config=tune.FailureConfig(fail_fast=True)),
+)
+result_grid = tuner.fit()
+# __stopping_on_trial_error_end__
+
+for result in result_grid:
+    # Should never get to report
+    final_iter = result.metrics.get("training_iteration")
+    assert not final_iter, final_iter
+
+# __early_stopping_start__
+from ray import tune
+from ray.tune.schedulers import AsyncHyperBandScheduler
+
+
+scheduler = AsyncHyperBandScheduler(time_attr="training_iteration")
+
+tuner = tune.Tuner(
+    my_trainable,
+    run_config=tune.RunConfig(stop={"training_iteration": 10}),
+    tune_config=tune.TuneConfig(
+        scheduler=scheduler, num_samples=2, metric="mean_accuracy", mode="max"
+    ),
+)
+result_grid = tuner.fit()
+# __early_stopping_end__
+
+
+def my_trainable(config):
+    # NOTE: Introduce the sleep again for the time-based unit-tests.
+    i = 1
+    while True:
+        time.sleep(1)
+        # Do some training, and report some metrics for demonstration...
+        tune.report({"mean_accuracy": min(i / 10, 1.0)})
+        i += 1
+
+
+# __stopping_trials_by_time_start__
+from ray import tune
+
+tuner = tune.Tuner(
+    my_trainable,
+    # Stop a trial after it's run for more than 5 seconds.
+    run_config=tune.RunConfig(stop={"time_total_s": 5}),
+)
+result_grid = tuner.fit()
+# __stopping_trials_by_time_end__
+
+# Should only get ~5 reports
+assert result_grid[0].metrics["training_iteration"] < 8
+
+
+# __stopping_experiment_by_time_start__
+from ray import tune
+
+# Stop the entire experiment after ANY trial has run for more than 5 seconds.
+tuner = tune.Tuner(my_trainable, tune_config=tune.TuneConfig(time_budget_s=5.0))
+result_grid = tuner.fit()
+# __stopping_experiment_by_time_end__
+
+# Should only get ~5 reports
+assert result_grid[0].metrics["training_iteration"] < 8

doc/source/tune/doc_code/trainable.py

@@ -0,0 +1,82 @@
+# flake8: noqa
+
+# fmt: off
+# __example_objective_start__
+def objective(x, a, b):
+    return a * (x ** 0.5) + b
+# __example_objective_end__
+# fmt: on
+
+# __function_api_report_intermediate_metrics_start__
+from ray import tune
+
+
+def trainable(config: dict):
+    intermediate_score = 0
+    for x in range(20):
+        intermediate_score = objective(x, config["a"], config["b"])
+        tune.report({"score": intermediate_score})  # This sends the score to Tune.
+
+
+tuner = tune.Tuner(trainable, param_space={"a": 2, "b": 4})
+results = tuner.fit()
+# __function_api_report_intermediate_metrics_end__
+
+# __function_api_report_final_metrics_start__
+from ray import tune
+
+
+def trainable(config: dict):
+    final_score = 0
+    for x in range(20):
+        final_score = objective(x, config["a"], config["b"])
+
+    tune.report({"score": final_score})  # This sends the score to Tune.
+
+
+tuner = tune.Tuner(trainable, param_space={"a": 2, "b": 4})
+results = tuner.fit()
+# __function_api_report_final_metrics_end__
+
+# fmt: off
+# __function_api_return_final_metrics_start__
+def trainable(config: dict):
+    final_score = 0
+    for x in range(20):
+        final_score = objective(x, config["a"], config["b"])
+
+    return {"score": final_score}  # This sends the score to Tune.
+# __function_api_return_final_metrics_end__
+# fmt: on
+
+# __class_api_example_start__
+from ray import tune
+
+
+class Trainable(tune.Trainable):
+    def setup(self, config: dict):
+        # config (dict): A dict of hyperparameters
+        self.x = 0
+        self.a = config["a"]
+        self.b = config["b"]
+
+    def step(self):  # This is called iteratively.
+        score = objective(self.x, self.a, self.b)
+        self.x += 1
+        return {"score": score}
+
+
+tuner = tune.Tuner(
+    Trainable,
+    run_config=tune.RunConfig(
+        # Train for 20 steps
+        stop={"training_iteration": 20},
+        checkpoint_config=tune.CheckpointConfig(
+            # We haven't implemented checkpointing yet. See below!
+            checkpoint_at_end=False
+        ),
+    ),
+    param_space={"a": 2, "b": 4},
+)
+results = tuner.fit()
+# __class_api_example_end__

doc/source/tune/doc_code/trial_checkpoint.py

@@ -0,0 +1,188 @@
+# flake8: noqa
+
+# __class_api_checkpointing_start__
+import os
+import torch
+from torch import nn
+
+from ray import tune
+
+
+class MyTrainableClass(tune.Trainable):
+    def setup(self, config):
+        self.model = nn.Sequential(
+            nn.Linear(config.get("input_size", 32), 32), nn.ReLU(), nn.Linear(32, 10)
+        )
+
+    def step(self):
+        return {}
+
+    def save_checkpoint(self, tmp_checkpoint_dir):
+        checkpoint_path = os.path.join(tmp_checkpoint_dir, "model.pth")
+        torch.save(self.model.state_dict(), checkpoint_path)
+        return tmp_checkpoint_dir
+
+    def load_checkpoint(self, tmp_checkpoint_dir):
+        checkpoint_path = os.path.join(tmp_checkpoint_dir, "model.pth")
+        self.model.load_state_dict(torch.load(checkpoint_path))
+
+
+tuner = tune.Tuner(
+    MyTrainableClass,
+    param_space={"input_size": 64},
+    run_config=tune.RunConfig(
+        stop={"training_iteration": 2},
+        checkpoint_config=tune.CheckpointConfig(checkpoint_frequency=2),
+    ),
+)
+tuner.fit()
+# __class_api_checkpointing_end__
+
+# __class_api_manual_checkpointing_start__
+import random
+
+
+# to be implemented by user.
+def detect_instance_preemption():
+    choice = random.randint(1, 100)
+    # simulating a 1% chance of preemption.
+    return choice <= 1
+
+
+def train_func(self):
+    # training code
+    result = {"mean_accuracy": "my_accuracy"}
+    if detect_instance_preemption():
+        result.update(should_checkpoint=True)
+    return result
+
+
+# __class_api_manual_checkpointing_end__
+
+# __class_api_periodic_checkpointing_start__
+
+tuner = tune.Tuner(
+    MyTrainableClass,
+    run_config=tune.RunConfig(
+        stop={"training_iteration": 2},
+        checkpoint_config=tune.CheckpointConfig(checkpoint_frequency=10),
+    ),
+)
+tuner.fit()
+
+# __class_api_periodic_checkpointing_end__
+
+
+# __class_api_end_checkpointing_start__
+tuner = tune.Tuner(
+    MyTrainableClass,
+    run_config=tune.RunConfig(
+        stop={"training_iteration": 2},
+        checkpoint_config=tune.CheckpointConfig(
+            checkpoint_frequency=10, checkpoint_at_end=True
+        ),
+    ),
+)
+tuner.fit()
+
+# __class_api_end_checkpointing_end__
+
+
+class MyModel:
+    def state_dict(self) -> dict:
+        return {}
+
+    def load_state_dict(self, state_dict):
+        pass
+
+
+# __function_api_checkpointing_from_dir_start__
+import os
+import tempfile
+
+from ray import tune
+from ray.tune import Checkpoint
+
+
+def train_func(config):
+    start = 1
+    my_model = MyModel()
+
+    checkpoint = tune.get_checkpoint()
+    if checkpoint:
+        with checkpoint.as_directory() as checkpoint_dir:
+            checkpoint_dict = torch.load(os.path.join(checkpoint_dir, "checkpoint.pt"))
+            start = checkpoint_dict["epoch"] + 1
+            my_model.load_state_dict(checkpoint_dict["model_state"])
+
+    for epoch in range(start, config["epochs"] + 1):
+        # Model training here
+        # ...
+
+        metrics = {"metric": 1}
+        with tempfile.TemporaryDirectory() as tempdir:
+            torch.save(
+                {"epoch": epoch, "model_state": my_model.state_dict()},
+                os.path.join(tempdir, "checkpoint.pt"),
+            )
+            tune.report(metrics=metrics, checkpoint=Checkpoint.from_directory(tempdir))
+
+
+tuner = tune.Tuner(train_func, param_space={"epochs": 5})
+result_grid = tuner.fit()
+# __function_api_checkpointing_from_dir_end__
+
+assert not result_grid.errors
+
+# __function_api_checkpointing_periodic_start__
+NUM_EPOCHS = 12
+# checkpoint every three epochs.
+CHECKPOINT_FREQ = 3
+
+
+def train_func(config):
+    for epoch in range(1, config["epochs"] + 1):
+        # Model training here
+        # ...
+
+        # Report metrics and save a checkpoint
+        metrics = {"metric": "my_metric"}
+        if epoch % CHECKPOINT_FREQ == 0:
+            with tempfile.TemporaryDirectory() as tempdir:
+                # Save a checkpoint in tempdir.
+                tune.report(metrics, checkpoint=Checkpoint.from_directory(tempdir))
+        else:
+            tune.report(metrics)
+
+
+tuner = tune.Tuner(train_func, param_space={"epochs": NUM_EPOCHS})
+result_grid = tuner.fit()
+# __function_api_checkpointing_periodic_end__
+
+assert not result_grid.errors
+assert len(result_grid[0].best_checkpoints) == NUM_EPOCHS // CHECKPOINT_FREQ
+
+# __callback_api_checkpointing_start__
+from ray import tune
+from ray.tune.experiment import Trial
+from ray.tune.result import SHOULD_CHECKPOINT, TRAINING_ITERATION
+
+
+class CheckpointByStepsTaken(tune.Callback):
+    def __init__(self, iterations_per_checkpoint: int):
+        self.steps_per_checkpoint = iterations_per_checkpoint
+        self._trials_last_checkpoint = {}
+
+    def on_trial_result(
+        self, iteration: int, trials: list[Trial], trial: Trial, result: dict, **info
+    ):
+        current_iteration = result[TRAINING_ITERATION]
+        if (
+            current_iteration - self._trials_last_checkpoint.get(trial, -1)
+            >= self.steps_per_checkpoint
+        ):
+            result[SHOULD_CHECKPOINT] = True
+            self._trials_last_checkpoint[trial] = current_iteration
+
+
+# __callback_api_checkpointing_end__

doc/source/tune/doc_code/tune.py

@@ -0,0 +1,60 @@
+# flake8: noqa
+
+# fmt: off
+# __step1_begin__
+from ray import tune
+import ray
+import os
+
+NUM_MODELS = 100
+
+def train_model(config):
+    score = config["model_id"]
+
+    # Import model libraries, etc...
+    # Load data and train model code here...
+
+    # Return final stats. You can also return intermediate progress
+    # using ray.tune.report() if needed.
+    # To return your model, you could write it to storage and return its
+    # URI in this dict, or return it as a Tune Checkpoint:
+    # https://docs.ray.io/en/latest/tune/tutorials/tune-checkpoints.html
+    return {"score": score, "other_data": ...}
+# __step1_end__
+
+# __step2_begin__
+# Define trial parameters as a single grid sweep.
+trial_space = {
+    # This is an example parameter. You could replace it with filesystem paths,
+    # model types, or even full nested Python dicts of model configurations, etc.,
+    # that enumerate the set of trials to run.
+    "model_id": tune.grid_search([
+        "model_{}".format(i)
+        for i in range(NUM_MODELS)
+    ])
+}
+# __step2_end__
+
+# __step3_begin__
+# Can customize resources per trial, here we set 1 CPU each.
+train_model = tune.with_resources(train_model, {"cpu": 1})
+# __step3_end__
+
+# __step4_begin__
+# Start a Tune run and print the best result.
+tuner = tune.Tuner(train_model, param_space=trial_space)
+results = tuner.fit()
+
+# Access individual results.
+print(results[0])
+print(results[1])
+print(results[2])
+# __step4_end__
+
+# __tasks_begin__
+remote_train = ray.remote(train_model)
+futures = [remote_train.remote({"model_id": i}) for i in range(NUM_MODELS)]
+print("Submitting tasks...")
+results = ray.get(futures)
+print("Trial results", results)
+# __tasks_end__