chore: import upstream snapshot with attribution

doc/source/train/doc_code/asynchronous_validation.py

@@ -0,0 +1,311 @@
+# __validation_fn_simple_start__
+import os
+import torch
+
+import ray.train
+import ray.data
+
+# Define Ray Data validation dataset outside validation function because it is not json serializable
+validation_dataset = ...
+
+
+def validation_fn(checkpoint: ray.train.Checkpoint) -> dict:
+    # Load the checkpoint
+    model = ...
+    with checkpoint.as_directory() as checkpoint_dir:
+        model_state_dict = torch.load(os.path.join(checkpoint_dir, "model.pt"))
+        model.load_state_dict(model_state_dict)
+    model.eval()
+
+    # Perform validation on the data
+    total_accuracy = 0
+    with torch.no_grad():
+        for batch in validation_dataset.iter_torch_batches(batch_size=128):
+            images, labels = batch["image"], batch["label"]
+            outputs = model(images)
+            total_accuracy += (outputs.argmax(1) == labels).sum().item()
+    return {"score": total_accuracy / len(validation_dataset)}
+# __validation_fn_simple_end__
+
+# __validation_fn_torch_trainer_start__
+import torchmetrics
+from torch.nn import CrossEntropyLoss
+
+import ray.train.torch
+from ray.data import ExecutionOptions
+
+
+def eval_only_train_fn(config_dict: dict) -> dict:
+    # Load the checkpoint
+    model = ...
+    with config_dict["checkpoint"].as_directory() as checkpoint_dir:
+        model_state_dict = torch.load(os.path.join(checkpoint_dir, "model.pt"))
+        model.load_state_dict(model_state_dict)
+    model.cuda().eval()
+
+    # Set up metrics and data loaders
+    criterion = CrossEntropyLoss()
+    mean_valid_loss = torchmetrics.MeanMetric().cuda()
+    test_data_shard = ray.train.get_dataset_shard("validation")
+    test_dataloader = test_data_shard.iter_torch_batches(batch_size=128)
+
+    # Compute metric and return it directly from the train function
+    with torch.no_grad():
+        for batch in test_dataloader:
+            images, labels = batch["image"], batch["label"]
+            outputs = model(images)
+            loss = criterion(outputs, labels)
+            mean_valid_loss(loss)
+    return {"score": mean_valid_loss.compute().item()}
+
+
+def validation_fn(checkpoint: ray.train.Checkpoint, train_run_name: str, epoch: int) -> dict:
+    trainer = ray.train.torch.TorchTrainer(
+        eval_only_train_fn,
+        train_loop_config={"checkpoint": checkpoint},
+        scaling_config=ray.train.ScalingConfig(
+            num_workers=2, use_gpu=True, accelerator_type="A10G"
+        ),
+        # Give unique name to validation run so it does not attempt to load placeholder checkpoint.
+        # Also allows you to better associate training runs with validation runs.
+        run_config=ray.train.RunConfig(
+            name=f"{train_run_name}_validation_epoch_{epoch}"
+        ),
+        # Use weaker GPUs for validation
+        datasets={"validation": validation_dataset},
+        # Pin to the "validation" subcluster so it doesn't compete with
+        # training. See https://docs.ray.io/en/latest/data/concurrent-dataset-execution.html.
+        dataset_config=ray.train.DataConfig(
+            execution_options={
+                "validation": ExecutionOptions(
+                    label_selector={"ray-subcluster": "validation"}
+                ),
+            },
+        ),
+    )
+    result = trainer.fit()
+    # return_value holds the value returned by train function of worker 0
+    return result.return_value
+# __validation_fn_torch_trainer_end__
+
+# __validation_fn_map_batches_start__
+import ray.data
+
+
+class Predictor:
+    def __init__(self, checkpoint: ray.train.Checkpoint):
+        self.model = ...
+        with checkpoint.as_directory() as checkpoint_dir:
+            model_state_dict = torch.load(os.path.join(checkpoint_dir, "model.pt"))
+            self.model.load_state_dict(model_state_dict)
+        self.model.cuda().eval()
+
+    def __call__(self, batch: dict) -> dict:
+        image = torch.as_tensor(batch["image"], dtype=torch.float32, device="cuda")
+        label = torch.as_tensor(batch["label"], dtype=torch.float32, device="cuda")
+        pred = self.model(image)
+        return {"res": (pred.argmax(1) == label).cpu().numpy()}
+
+
+# Construct ``validation_dataset`` under a DataContext copy pinned to the
+# "validation" subcluster. ``Dataset.context`` is a deep copy of the
+# current context taken at construction, so the selector is baked in and
+# every downstream operator (including the ``map_batches`` below) inherits
+# it — no in-function mutation needed. See
+# https://docs.ray.io/en/latest/data/concurrent-dataset-execution.html.
+ctx = ray.data.DataContext.get_current().copy()
+ctx.execution_options.label_selector = {"ray-subcluster": "validation"}
+with ray.data.DataContext.current(ctx):
+    validation_dataset = ray.data.read_parquet(...)
+
+
+def validation_fn(checkpoint: ray.train.Checkpoint) -> dict:
+    # Set name to avoid confusion; default name is "Dataset"
+    validation_dataset.set_name("validation")
+    eval_res = validation_dataset.map_batches(
+        Predictor,
+        batch_size=128,
+        num_gpus=1,
+        fn_constructor_kwargs={"checkpoint": checkpoint},
+        concurrency=2,
+    )
+    mean = eval_res.mean(["res"])
+    return {
+        "score": mean,
+    }
+# __validation_fn_map_batches_end__
+
+# __validation_fn_report_start__
+import tempfile
+
+from ray.data import ExecutionOptions
+from ray.train import ValidationConfig, ValidationTaskConfig
+
+
+def train_func(config: dict) -> None:
+    ...
+    epochs = ...
+    model = ...
+    rank = ray.train.get_context().get_world_rank()
+    for epoch in epochs:
+        ...  # training step
+        if rank == 0:
+            training_metrics = {"loss": ..., "epoch": epoch}
+            local_checkpoint_dir = tempfile.mkdtemp()
+            torch.save(
+                model.module.state_dict(),
+                os.path.join(local_checkpoint_dir, "model.pt"),
+            )
+            ray.train.report(
+                training_metrics,
+                checkpoint=ray.train.Checkpoint.from_directory(local_checkpoint_dir),
+                checkpoint_upload_mode=ray.train.CheckpointUploadMode.ASYNC,
+                validation=ValidationTaskConfig(fn_kwargs={
+                    "train_run_name": ray.train.get_context().get_experiment_name(),
+                    "epoch": epoch,
+                }),
+            )
+        else:
+            ray.train.report({}, None)
+
+
+def run_trainer() -> ray.train.Result:
+    # 1) Construction-time tasks (parquet schema inference, file listing)
+    # read the current DataContext. Pin them to "training" with a copy of
+    # the DataContext applied via the DataContext.current() context
+    # manager — scoped to the `with` block so it doesn't leak. See
+    # https://docs.ray.io/en/latest/data/concurrent-dataset-execution.html.
+    ctx = ray.data.DataContext.get_current().copy()
+    ctx.execution_options.label_selector = {"ray-subcluster": "training"}
+    with ray.data.DataContext.current(ctx):
+        train_dataset = ray.data.read_parquet(...)
+
+    trainer = ray.train.torch.TorchTrainer(
+        train_func,
+        validation_config=ValidationConfig(fn=validation_fn),
+        # Pass training dataset in datasets arg to split it across training workers
+        datasets={"train": train_dataset},
+        # 2) DataConfig.execution_options REPLACES ds.context.execution_options
+        # wholesale at training start, dropping anything not re-specified
+        # (including label_selector). Restate the selector here so per-worker
+        # ingest stays pinned to "training".
+        dataset_config=ray.train.DataConfig(
+            datasets_to_split=["train"],
+            execution_options={
+                "train": ExecutionOptions(
+                    label_selector={"ray-subcluster": "training"}
+                ),
+            },
+        ),
+        scaling_config=ray.train.ScalingConfig(
+            num_workers=2,
+            use_gpu=True,
+            # Use powerful GPUs for training
+            accelerator_type="A100",
+        ),
+    )
+    return trainer.fit()
+# __validation_fn_report_end__
+
+# __exp_tracking_same_run_wandb_start__
+import wandb
+import ray.train
+from ray.train import ValidationConfig, ValidationTaskConfig
+
+
+entity = "my_entity"
+project = "my_project"
+num_epochs = ...
+
+
+def validation_fn(checkpoint: ray.train.Checkpoint, wandb_run_id: str, val_step: int) -> dict:
+    wandb.init(
+        entity=entity,
+        project=project,
+        settings=wandb.Settings(mode="shared", x_primary=False),
+        id=wandb_run_id,
+    )
+    score = ...
+    wandb.log({"validation/loss": score, "val_step": val_step})
+    wandb.finish()  # flush the metrics
+    return {"validation/loss": score}
+
+
+def train_func():
+    if ray.train.get_context().get_world_rank() == 0:
+        run = wandb.init(
+            entity=entity,
+            project=project,
+            settings=wandb.Settings(mode="shared", x_primary=True,)
+        )
+        wandb.define_metric("val_step", hidden=True)
+        wandb.define_metric("train_step", hidden=True)
+        wandb.define_metric("validation/loss", step_metric="val_step")
+        wandb.define_metric("train/loss", step_metric="train_step")
+
+    for epoch in range(num_epochs):
+        loss = ...
+        if ray.train.get_context().get_world_rank() == 0:
+            wandb.log({"train/loss": loss, "train_step": epoch})
+            checkpoint = ...
+            ray.train.report(
+                {"train/loss": loss},
+                checkpoint=checkpoint,
+                validation=ValidationTaskConfig(
+                    fn_kwargs={"wandb_run_id": run.id, "val_step": epoch}
+                ),
+            )
+        else:
+            ray.train.report({}, None)
+
+    if ray.train.get_context().get_world_rank() == 0:
+        wandb.finish()
+
+
+# __exp_tracking_same_run_wandb_end__
+
+# __exp_tracking_same_run_mlflow_start__
+import mlflow
+from mlflow.tracking import MlflowClient
+import ray.train
+from ray.train import ValidationConfig, ValidationTaskConfig
+
+
+tracking_uri = "my_uri"
+experiment_name = "my_experiment"
+num_epochs = ...
+
+def validation_fn(
+    checkpoint: ray.train.Checkpoint, mlflow_run_id: str, val_step: int
+) -> dict:
+    client = MlflowClient(tracking_uri=tracking_uri)
+    score = ...
+    client.log_metric(mlflow_run_id, "val_score", score, step=val_step)
+    return {"val_score": score}
+
+
+def train_func():
+    if ray.train.get_context().get_world_rank() == 0:
+        client = MlflowClient(tracking_uri=tracking_uri)
+        experiment = client.get_experiment_by_name(experiment_name)
+        run = client.create_run(experiment_id=experiment.experiment_id)
+
+    for epoch in range(num_epochs):
+        loss = ...
+        if ray.train.get_context().get_world_rank() == 0:
+            client.log_metric(run.info.run_id, "train_loss", loss, step=epoch)
+            checkpoint = ...
+            ray.train.report(
+                {"train_loss": loss},
+                checkpoint=checkpoint,
+                validation=ValidationTaskConfig(
+                    fn_kwargs={"mlflow_run_id": run.info.run_id, "val_step": epoch}
+                ),
+            )
+        else:
+            ray.train.report({}, None)
+
+    if ray.train.get_context().get_world_rank() == 0:
+        client.set_terminated(run.info.run_id)
+
+# __exp_tracking_same_run_mlflow_end__

doc/source/train/doc_code/checkpoints.py

@@ -0,0 +1,569 @@
+# flake8: noqa
+# isort: skip_file
+
+# __pytorch_save_start__
+import os
+import tempfile
+
+import numpy as np
+import torch
+import torch.nn as nn
+from torch.optim import Adam
+
+import ray.train.torch
+from ray import train
+from ray.train import Checkpoint, ScalingConfig
+from ray.train.torch import TorchTrainer
+
+
+def train_func(config):
+    n = 100
+    # create a toy dataset
+    # data   : X - dim = (n, 4)
+    # target : Y - dim = (n, 1)
+    X = torch.Tensor(np.random.normal(0, 1, size=(n, 4)))
+    Y = torch.Tensor(np.random.uniform(0, 1, size=(n, 1)))
+    # toy neural network : 1-layer
+    # Wrap the model in DDP
+    model = ray.train.torch.prepare_model(nn.Linear(4, 1))
+    criterion = nn.MSELoss()
+
+    optimizer = Adam(model.parameters(), lr=3e-4)
+    for epoch in range(config["num_epochs"]):
+        y = model.forward(X)
+        loss = criterion(y, Y)
+        optimizer.zero_grad()
+        loss.backward()
+        optimizer.step()
+
+        metrics = {"loss": loss.item()}
+
+        with tempfile.TemporaryDirectory() as temp_checkpoint_dir:
+            checkpoint = None
+
+            should_checkpoint = epoch % config.get("checkpoint_freq", 1) == 0
+            # In standard DDP training, where the model is the same across all ranks,
+            # only the global rank 0 worker needs to save and report the checkpoint
+            if train.get_context().get_world_rank() == 0 and should_checkpoint:
+                torch.save(
+                    model.module.state_dict(),  # NOTE: Unwrap the model.
+                    os.path.join(temp_checkpoint_dir, "model.pt"),
+                )
+                checkpoint = Checkpoint.from_directory(temp_checkpoint_dir)
+
+            train.report(metrics, checkpoint=checkpoint)
+
+
+trainer = TorchTrainer(
+    train_func,
+    train_loop_config={"num_epochs": 5},
+    scaling_config=ScalingConfig(num_workers=2),
+)
+result = trainer.fit()
+# __pytorch_save_end__
+
+# __pytorch_restore_start__
+import os
+import tempfile
+
+import numpy as np
+import torch
+import torch.nn as nn
+from torch.optim import Adam
+
+import ray.train.torch
+from ray import train
+from ray.train import Checkpoint, ScalingConfig
+from ray.train.torch import TorchTrainer
+
+
+def train_func(config):
+    n = 100
+    # create a toy dataset
+    # data   : X - dim = (n, 4)
+    # target : Y - dim = (n, 1)
+    X = torch.Tensor(np.random.normal(0, 1, size=(n, 4)))
+    Y = torch.Tensor(np.random.uniform(0, 1, size=(n, 1)))
+    # toy neural network : 1-layer
+    model = nn.Linear(4, 1)
+    optimizer = Adam(model.parameters(), lr=3e-4)
+    criterion = nn.MSELoss()
+
+    # Wrap the model in DDP and move it to GPU.
+    model = ray.train.torch.prepare_model(model)
+
+    # ====== Resume training state from the checkpoint. ======
+    start_epoch = 0
+    checkpoint = train.get_checkpoint()
+    if checkpoint:
+        with checkpoint.as_directory() as checkpoint_dir:
+            model_state_dict = torch.load(
+                os.path.join(checkpoint_dir, "model.pt"),
+                # map_location=...,  # Load onto a different device if needed.
+            )
+            model.module.load_state_dict(model_state_dict)
+            optimizer.load_state_dict(
+                torch.load(os.path.join(checkpoint_dir, "optimizer.pt"))
+            )
+            start_epoch = (
+                torch.load(os.path.join(checkpoint_dir, "extra_state.pt"))["epoch"] + 1
+            )
+    # ========================================================
+
+    for epoch in range(start_epoch, config["num_epochs"]):
+        y = model.forward(X)
+        loss = criterion(y, Y)
+        optimizer.zero_grad()
+        loss.backward()
+        optimizer.step()
+
+        metrics = {"loss": loss.item()}
+
+        with tempfile.TemporaryDirectory() as temp_checkpoint_dir:
+            checkpoint = None
+
+            should_checkpoint = epoch % config.get("checkpoint_freq", 1) == 0
+            # In standard DDP training, where the model is the same across all ranks,
+            # only the global rank 0 worker needs to save and report the checkpoint
+            if train.get_context().get_world_rank() == 0 and should_checkpoint:
+                # === Make sure to save all state needed for resuming training ===
+                torch.save(
+                    model.module.state_dict(),  # NOTE: Unwrap the model.
+                    os.path.join(temp_checkpoint_dir, "model.pt"),
+                )
+                torch.save(
+                    optimizer.state_dict(),
+                    os.path.join(temp_checkpoint_dir, "optimizer.pt"),
+                )
+                torch.save(
+                    {"epoch": epoch},
+                    os.path.join(temp_checkpoint_dir, "extra_state.pt"),
+                )
+                # ================================================================
+                checkpoint = Checkpoint.from_directory(temp_checkpoint_dir)
+
+            train.report(metrics, checkpoint=checkpoint)
+
+        if epoch == 1:
+            raise RuntimeError("Intentional error to showcase restoration!")
+
+
+trainer = TorchTrainer(
+    train_func,
+    train_loop_config={"num_epochs": 5},
+    scaling_config=ScalingConfig(num_workers=2),
+    run_config=train.RunConfig(failure_config=train.FailureConfig(max_failures=1)),
+)
+result = trainer.fit()
+# __pytorch_restore_end__
+
+# __checkpoint_from_single_worker_start__
+import tempfile
+
+from ray import train
+
+
+def train_fn(config):
+    ...
+
+    metrics = {...}
+    with tempfile.TemporaryDirectory() as temp_checkpoint_dir:
+        checkpoint = None
+
+        # Only the global rank 0 worker saves and reports the checkpoint
+        if train.get_context().get_world_rank() == 0:
+            ...  # Save checkpoint to temp_checkpoint_dir
+
+            checkpoint = Checkpoint.from_directory(temp_checkpoint_dir)
+
+        train.report(metrics, checkpoint=checkpoint)
+# __checkpoint_from_single_worker_end__
+
+
+# __lightning_save_example_start__
+import lightning.pytorch as pl
+
+from ray import train
+from ray.train.lightning import RayTrainReportCallback
+from ray.train.torch import TorchTrainer
+
+
+class MyLightningModule(pl.LightningModule):
+    # ...
+
+    def on_validation_epoch_end(self):
+        ...
+        mean_acc = calculate_accuracy()
+        self.log("mean_accuracy", mean_acc, sync_dist=True)
+
+
+def train_func():
+    ...
+    model = MyLightningModule(...)
+    datamodule = MyLightningDataModule(...)
+
+    trainer = pl.Trainer(
+        # ...
+        callbacks=[RayTrainReportCallback()]
+    )
+    trainer.fit(model, datamodule=datamodule)
+
+
+ray_trainer = TorchTrainer(
+    train_func,
+    scaling_config=train.ScalingConfig(num_workers=2),
+    run_config=train.RunConfig(
+        checkpoint_config=train.CheckpointConfig(
+            num_to_keep=2,
+            checkpoint_score_attribute="mean_accuracy",
+            checkpoint_score_order="max",
+        ),
+    ),
+)
+# __lightning_save_example_end__
+
+
+# __lightning_custom_save_example_start__
+import os
+from tempfile import TemporaryDirectory
+
+from lightning.pytorch.callbacks import Callback
+
+import ray
+import ray.train
+from ray.train import Checkpoint
+
+
+class CustomRayTrainReportCallback(Callback):
+    def on_train_epoch_end(self, trainer, pl_module):
+        should_checkpoint = trainer.current_epoch % 3 == 0
+
+        with TemporaryDirectory() as tmpdir:
+            # Fetch metrics from `self.log(..)` in the LightningModule
+            metrics = trainer.callback_metrics
+            metrics = {k: v.item() for k, v in metrics.items()}
+
+            # Add customized metrics
+            metrics["epoch"] = trainer.current_epoch
+            metrics["custom_metric"] = 123
+
+            checkpoint = None
+            global_rank = ray.train.get_context().get_world_rank() == 0
+            if global_rank == 0 and should_checkpoint:
+                # Save model checkpoint file to tmpdir
+                ckpt_path = os.path.join(tmpdir, "ckpt.pt")
+                trainer.save_checkpoint(ckpt_path, weights_only=False)
+
+                checkpoint = Checkpoint.from_directory(tmpdir)
+
+            # Report to train session
+            ray.train.report(metrics=metrics, checkpoint=checkpoint)
+# __lightning_custom_save_example_end__
+
+# __lightning_restore_example_start__
+import os
+
+from ray import train
+from ray.train import Checkpoint
+from ray.train.torch import TorchTrainer
+from ray.train.lightning import RayTrainReportCallback
+
+
+def train_func():
+    model = MyLightningModule(...)
+    datamodule = MyLightningDataModule(...)
+    trainer = pl.Trainer(..., callbacks=[RayTrainReportCallback()])
+
+    checkpoint = train.get_checkpoint()
+    if checkpoint:
+        with checkpoint.as_directory() as ckpt_dir:
+            ckpt_path = os.path.join(ckpt_dir, RayTrainReportCallback.CHECKPOINT_NAME)
+            trainer.fit(model, datamodule=datamodule, ckpt_path=ckpt_path)
+    else:
+        trainer.fit(model, datamodule=datamodule)
+
+
+ray_trainer = TorchTrainer(
+    train_func,
+    scaling_config=train.ScalingConfig(num_workers=2),
+    run_config=train.RunConfig(
+        checkpoint_config=train.CheckpointConfig(num_to_keep=2),
+    ),
+)
+# __lightning_restore_example_end__
+
+
+# __transformers_save_example_start__
+from transformers import TrainingArguments
+
+from ray import train
+from ray.train.huggingface.transformers import RayTrainReportCallback, prepare_trainer
+from ray.train.torch import TorchTrainer
+
+
+def train_func(config):
+    ...
+
+    # Configure logging, saving, evaluation strategies as usual.
+    args = TrainingArguments(
+        ...,
+        eval_strategy="epoch",
+        save_strategy="epoch",
+        logging_strategy="step",
+    )
+
+    trainer = transformers.Trainer(args, ...)
+
+    # Add a report callback to transformers Trainer
+    # =============================================
+    trainer.add_callback(RayTrainReportCallback())
+    trainer = prepare_trainer(trainer)
+
+    trainer.train()
+
+
+ray_trainer = TorchTrainer(
+    train_func,
+    run_config=train.RunConfig(
+        checkpoint_config=train.CheckpointConfig(
+            num_to_keep=3,
+            checkpoint_score_attribute="eval_loss",  # The monitoring metric
+            checkpoint_score_order="min",
+        )
+    ),
+)
+# __transformers_save_example_end__
+
+
+# __transformers_custom_save_example_start__
+from ray import train
+
+from transformers.trainer_callback import TrainerCallback
+
+
+class MyTrainReportCallback(TrainerCallback):
+    def __init__(self):
+        super().__init__()
+        self.metrics = {}
+
+    def on_log(self, args, state, control, model=None, logs=None, **kwargs):
+        """Log is called on evaluation step and logging step."""
+        self.metrics.update(logs)
+
+    def on_save(self, args, state, control, **kwargs):
+        """Event called after a checkpoint save."""
+
+        checkpoint = None
+        if train.get_context().get_world_rank() == 0:
+            # Build a Ray Train Checkpoint from the latest checkpoint
+            checkpoint_path = transformers.trainer.get_last_checkpoint(args.output_dir)
+            checkpoint = Checkpoint.from_directory(checkpoint_path)
+
+        # Report to Ray Train with up-to-date metrics
+        ray.train.report(metrics=self.metrics, checkpoint=checkpoint)
+
+        # Clear the metrics buffer
+        self.metrics = {}
+# __transformers_custom_save_example_end__
+
+
+# __distributed_checkpointing_start__
+from ray import train
+from ray.train import Checkpoint
+from ray.train.torch import TorchTrainer
+
+
+def train_func(config):
+    ...
+
+    with tempfile.TemporaryDirectory() as temp_checkpoint_dir:
+        rank = train.get_context().get_world_rank()
+        torch.save(
+            ...,
+            os.path.join(temp_checkpoint_dir, f"model-rank={rank}.pt"),
+        )
+        checkpoint = Checkpoint.from_directory(temp_checkpoint_dir)
+
+        train.report(metrics, checkpoint=checkpoint)
+
+
+trainer = TorchTrainer(
+    train_func,
+    scaling_config=train.ScalingConfig(num_workers=2),
+    run_config=train.RunConfig(storage_path="s3://bucket/"),
+)
+# The checkpoint in cloud storage will contain: model-rank=0.pt, model-rank=1.pt
+# __distributed_checkpointing_end__
+
+# __inspect_checkpoint_example_start__
+from pathlib import Path
+
+from ray.train import Checkpoint
+
+# For demonstration, create a locally available directory with a `model.pt` file.
+example_checkpoint_dir = Path("/tmp/test-checkpoint")
+example_checkpoint_dir.mkdir()
+example_checkpoint_dir.joinpath("model.pt").touch()
+
+# Create the checkpoint, which is a reference to the directory.
+checkpoint = Checkpoint.from_directory(example_checkpoint_dir)
+
+# Inspect the checkpoint's contents with either `as_directory` or `to_directory`:
+with checkpoint.as_directory() as checkpoint_dir:
+    assert Path(checkpoint_dir).joinpath("model.pt").exists()
+
+checkpoint_dir = checkpoint.to_directory()
+assert Path(checkpoint_dir).joinpath("model.pt").exists()
+# __inspect_checkpoint_example_end__
+
+# __inspect_transformers_checkpoint_example_start__
+# After training finished
+checkpoint = result.checkpoint
+with checkpoint.as_directory() as checkpoint_dir:
+    hf_checkpoint_path = f"{checkpoint_dir}/checkpoint/"
+# __inspect_transformers_checkpoint_example_end__
+
+# __inspect_lightning_checkpoint_example_start__
+# After training finished
+checkpoint = result.checkpoint
+with checkpoint.as_directory() as checkpoint_dir:
+    lightning_checkpoint_path = f"{checkpoint_dir}/checkpoint.ckpt"
+# __inspect_lightning_checkpoint_example_end__
+
+# __checkpoint_upload_mode_sync_start__
+def train_fn(config):
+    ...
+    metrics = {...}
+    with tempfile.TemporaryDirectory() as tmpdir:
+        ...  # Save checkpoint to tmpdir
+        checkpoint = Checkpoint.from_directory(tmpdir)
+        train.report(
+            metrics,
+            checkpoint=checkpoint,
+            checkpoint_upload_mode=train.CheckpointUploadMode.SYNC,
+        )
+# __checkpoint_upload_mode_sync_end__
+
+# __checkpoint_upload_mode_async_start__
+def train_fn(config):
+    ...
+    metrics = {...}
+    tmpdir = tempfile.mkdtemp()
+    ...  # Save checkpoint to tmpdir
+    checkpoint = Checkpoint.from_directory(tmpdir)
+    train.report(
+        metrics,
+        checkpoint=checkpoint,
+        checkpoint_upload_mode=train.CheckpointUploadMode.ASYNC,
+    )
+# __checkpoint_upload_mode_async_end__
+
+# __checkpoint_upload_mode_no_upload_start__
+from s3torchconnector.dcp import S3StorageWriter
+from torch.distributed.checkpoint.state_dict_saver import save
+from torch.distributed.checkpoint.state_dict import get_state_dict
+
+
+def train_fn(config):
+    ...
+    for epoch in range(config["num_epochs"]):
+        # Directly upload checkpoint to s3 with Torch
+        model, optimizer = ...
+        storage_context = ray.train.get_context().get_storage()
+        checkpoint_path = (
+            f"s3://{storage_context.build_checkpoint_path_from_name(str(epoch))}"
+        )
+        storage_writer = S3StorageWriter(region="us-west-2", path=checkpoint_path)
+        model_dict, opt_dict = get_state_dict(model=model, optimizers=optimizer)
+        save(
+            {"model": model_dict, "opt": opt_dict},
+            storage_writer=storage_writer,
+        )
+
+        # Report that checkpoint to Ray Train
+        metrics = {...}
+        checkpoint = Checkpoint(checkpoint_path)
+        train.report(
+            metrics,
+            checkpoint=checkpoint,
+            checkpoint_upload_mode=train.CheckpointUploadMode.NO_UPLOAD,
+        )
+# __checkpoint_upload_mode_no_upload_end__
+
+
+# __checkpoint_upload_fn_start__
+from torch.distributed.checkpoint.state_dict_saver import async_save
+from s3torchconnector.dcp import S3StorageWriter
+from torch.distributed.checkpoint.state_dict import get_state_dict
+
+from ray import train
+from ray.train import Checkpoint
+
+
+def train_fn(config):
+    ...
+    for epoch in config["num_epochs"]:
+        # Start async checkpoint upload to s3 with Torch
+        model, optimizer = ...
+        storage_context = train.get_context().get_storage()
+        checkpoint_path = (
+            f"s3://{storage_context.build_checkpoint_path_from_name(str(epoch))}"
+        )
+        storage_writer = S3StorageWriter(region="us-west-2", path=checkpoint_path)
+        model_dict, opt_dict = get_state_dict(model=model, optimizers=optimizer)
+        ckpt_ref = async_save(
+            {"model": model_dict, "opt": opt_dict},
+            storage_writer=storage_writer,
+        )
+
+        def wait_async_save(checkpoint, checkpoint_dir_name):
+            # This function waits for checkpoint to be finalized before returning it as is
+            ckpt_ref.result()
+            return checkpoint
+
+        # Ray Train kicks off a thread that waits for the async checkpoint upload to complete
+        # before reporting the checkpoint
+        metrics = {...}
+        checkpoint = Checkpoint(checkpoint_path)
+        train.report(
+            metrics=metrics,
+            checkpoint=checkpoint,
+            checkpoint_upload_mode=train.CheckpointUploadMode.ASYNC,
+            checkpoint_upload_fn=wait_async_save,
+            # As uploading into the experiment directory then don't delete the checkpoint after upload is complete
+            delete_local_checkpoint_after_upload=False,
+        )
+
+trainer = TorchTrainer(
+   train_fn,
+   train_loop_config={"num_epochs": 3},
+   scaling_config=train.ScalingConfig(num_workers=2, use_gpu=True),
+   # we need a cpu backend for async_save and a gpu backend for training
+   torch_config=train.torch.TorchConfig(backend="cpu:gloo,cuda:nccl"),
+   run_config=train.RunConfig(storage_path="s3://bucket/")
+)
+# __checkpoint_upload_fn_end__
+
+# __get_all_reported_checkpoints_example_start__
+import ray.train
+from ray.train import CheckpointConsistencyMode
+
+def train_fn():
+    for epoch in range(2):
+        metrics = {"train/loss": 0.1}
+        checkpoint = ...
+        ray.train.report(
+            metrics,
+            checkpoint=checkpoint,
+            validation=...,
+        )
+
+    # Get committed checkpoints which may still have ongoing validations.
+    committed_checkpoints = ray.train.get_all_reported_checkpoints(
+        consistency_mode=CheckpointConsistencyMode.COMMITTED)
+
+    # Wait for all pending validations to finish to access reported checkpoints
+    # with validation metrics attached.
+    validated_checkpoints = ray.train.get_all_reported_checkpoints()
+    ...
+# __get_all_reported_checkpoints_example_end__

doc/source/train/doc_code/collate_utils.py

@@ -0,0 +1,118 @@
+from dataclasses import dataclass
+from typing import Dict, List, Tuple, Union
+import torch
+from ray import cloudpickle as pickle
+import pyarrow as pa
+
+# (dtype, shape, offset)
+FEATURE_TYPE = Tuple[torch.dtype, torch.Size, int]
+TORCH_BYTE_ELEMENT_TYPE = torch.uint8
+
+def _create_binary_array_from_buffer(buffer: bytes) -> pa.BinaryArray:
+    """Zero-copy create a binary array from a buffer."""
+    data_buffer = pa.py_buffer(buffer)
+    return pa.Array.from_buffers(
+        pa.binary(),
+        1,
+        [
+            None,
+            pa.array([0, data_buffer.size], type=pa.int32()).buffers()[1],
+            data_buffer,
+        ],
+    )
+
+@dataclass
+class _Metadata:
+    features: Dict[str, List[FEATURE_TYPE]]
+    total_buffer_size: int
+
+@dataclass
+class _TensorBatch:
+    """Internal class for serializing/deserializing tensor batches."""
+    buffer: torch.Tensor
+    metadata: _Metadata
+
+    @classmethod
+    def from_batch(cls, batch: Dict[str, Union[List[torch.Tensor], torch.Tensor]]) -> '_TensorBatch':
+        """Serialize a batch of tensors into a single buffer."""
+        features: Dict[str, List[FEATURE_TYPE]] = {}
+        flattened_binary_tensors = []
+        total_buffer_size = 0
+        
+        for name, tensors in batch.items():
+            features[name] = []
+            if not isinstance(tensors, list):
+                tensors = [tensors]
+            for tensor in tensors:
+                flattened_tensor = tensor.flatten().contiguous().view(TORCH_BYTE_ELEMENT_TYPE)
+                flattened_binary_tensors.append(flattened_tensor)
+                features[name].append((tensor.dtype, tensor.shape, total_buffer_size))
+                total_buffer_size += flattened_tensor.shape[0]
+        
+        buffer = torch.empty(total_buffer_size, dtype=TORCH_BYTE_ELEMENT_TYPE)
+        cur_offset = 0
+        for flattened_tensor in flattened_binary_tensors:
+            buffer[cur_offset:cur_offset + flattened_tensor.shape[0]] = flattened_tensor
+            cur_offset += flattened_tensor.shape[0]
+        
+        return _TensorBatch(
+            buffer=buffer,
+            metadata=_Metadata(
+                features=features,
+                total_buffer_size=total_buffer_size,
+            ),
+        )
+
+    def to_table(self) -> pa.Table:
+        """Convert to a single-row PyArrow table."""
+        buffer_array = _create_binary_array_from_buffer(self.buffer.numpy().data)
+        metadata_array = _create_binary_array_from_buffer(pickle.dumps(self.metadata))
+        return pa.Table.from_arrays(
+            arrays=[buffer_array, metadata_array],
+            names=["_buffer", "_metadata"],
+        )
+
+    @classmethod
+    def from_table(cls, table: pa.Table) -> '_TensorBatch':
+        """Deserialize from a single-row PyArrow table."""
+        return _TensorBatch(
+            buffer=torch.frombuffer(
+                table["_buffer"].chunks[0].buffers()[2],
+                dtype=TORCH_BYTE_ELEMENT_TYPE
+            ),
+            metadata=pickle.loads(table["_metadata"].chunks[0].buffers()[2]),
+        )
+
+    def to_batch(self, pin_memory: bool = False) -> Dict[str, List[torch.Tensor]]:
+        """Deserialize back to a batch of tensors."""
+        batch = {}
+        storage_buffer = self.buffer.untyped_storage()
+        offsets = []
+        for name, features in self.metadata.features.items():
+            for _, _, offset in features:
+                offsets.append(offset)
+        offsets.append(self.metadata.total_buffer_size)
+        
+        offset_id = 0
+        for name, features in self.metadata.features.items():
+            batch[name] = []
+            for dtype, shape, _ in features:
+                # Create a zero-copy view of the byte slice.
+                byte_slice = self.buffer[offsets[offset_id]:offsets[offset_id + 1]]
+                tensor = torch.frombuffer(
+                    byte_slice.numpy().data, dtype=dtype
+                ).view(shape)
+                if pin_memory:
+                    tensor = tensor.pin_memory()
+                batch[name].append(tensor)
+                offset_id += 1
+        return batch
+
+# Helper functions for use in your code
+def serialize_tensors_to_table(batch: Dict[str, Union[List[torch.Tensor], torch.Tensor]]) -> pa.Table:
+    """Serialize a batch of tensors to a PyArrow table."""
+    return _TensorBatch.from_batch(batch).to_table()
+
+def deserialize_table_to_tensors(table: pa.Table, pin_memory: bool = False) -> Dict[str, List[torch.Tensor]]:
+    """Deserialize a PyArrow table back to tensors."""
+    return _TensorBatch.from_table(table).to_batch(pin_memory=pin_memory)

doc/source/train/doc_code/data_ingest_torch_new.py

@@ -0,0 +1,147 @@
+# flake8: noqa
+# isort: skip_file
+
+# __basic__
+import ray
+from ray import train
+from ray.train import ScalingConfig
+from ray.train.torch import TorchTrainer
+
+import numpy as np
+from typing import Dict
+
+# Load the data.
+train_ds = ray.data.read_parquet("s3://anonymous@ray-example-data/iris.parquet")
+## Uncomment to randomize the block order each epoch.
+# train_ds = train_ds.randomize_block_order()
+
+
+# Define a preprocessing function.
+def normalize_length(batch: Dict[str, np.ndarray]) -> Dict[str, np.ndarray]:
+    new_col = batch["sepal.length"] / np.max(batch["sepal.length"])
+    batch["normalized.sepal.length"] = new_col
+    del batch["sepal.length"]
+    return batch
+
+
+# Preprocess your data any way you want. This will be re-run each epoch.
+# You can use Ray Data preprocessors here as well,
+# e.g., preprocessor.fit_transform(train_ds)
+train_ds = train_ds.map_batches(normalize_length)
+
+
+def train_loop_per_worker():
+    # Get an iterator to the dataset we passed in below.
+    it = train.get_dataset_shard("train")
+
+    # Train for 10 epochs over the data. We'll use a shuffle buffer size
+    # of 10k elements, and prefetch up to 10 batches of size 128 each.
+    for _ in range(10):
+        for batch in it.iter_batches(
+            local_shuffle_buffer_size=10000, batch_size=128, prefetch_batches=10
+        ):
+            print("Do some training on batch", batch)
+
+
+my_trainer = TorchTrainer(
+    train_loop_per_worker,
+    scaling_config=ScalingConfig(num_workers=2),
+    datasets={"train": train_ds},
+)
+my_trainer.fit()
+# __basic_end__
+
+# __custom_split__
+dataset_a = ray.data.read_text(
+    "s3://anonymous@ray-example-data/sms_spam_collection_subset.txt"
+)
+dataset_b = ray.data.read_csv("s3://anonymous@ray-example-data/dow_jones.csv")
+
+my_trainer = TorchTrainer(
+    train_loop_per_worker,
+    scaling_config=ScalingConfig(num_workers=2),
+    datasets={"a": dataset_a, "b": dataset_b},
+    dataset_config=ray.train.DataConfig(
+        datasets_to_split=["a"],
+    ),
+)
+# __custom_split_end__
+
+
+def augment_data(batch):
+    return batch
+
+
+# __materialized__
+# Load the data.
+train_ds = ray.data.read_parquet("s3://anonymous@ray-example-data/iris.parquet")
+
+# Preprocess the data. Transformations that are made to the materialize call below
+# will only be run once.
+train_ds = train_ds.map_batches(normalize_length)
+
+# Materialize the dataset in object store memory.
+train_ds = train_ds.materialize()
+
+# Add per-epoch preprocessing. Transformations that you want to run per-epoch, such
+# as data augmentation, should go after the materialize call.
+train_ds = train_ds.map_batches(augment_data)
+# __materialized_end__
+
+# __options__
+from ray.train import DataConfig
+
+options = DataConfig.default_ingest_options()
+options.resource_limits = options.resource_limits.copy(object_store_memory=10e9)
+
+
+my_trainer = TorchTrainer(
+    train_loop_per_worker,
+    scaling_config=ScalingConfig(num_workers=2),
+    dataset_config=ray.train.DataConfig(
+        execution_options=options,
+    ),
+)
+# __options_end__
+
+# __custom__
+# Note that this example class is doing the same thing as the basic DataConfig
+# impl included with Ray Train.
+from typing import Optional, Dict, List
+
+from ray.data import Dataset, DataIterator, NodeIdStr
+from ray.actor import ActorHandle
+
+
+class MyCustomDataConfig(DataConfig):
+    def configure(
+        self,
+        datasets: Dict[str, Dataset],
+        world_size: int,
+        worker_handles: Optional[List[ActorHandle]],
+        worker_node_ids: Optional[List[NodeIdStr]],
+        **kwargs,
+    ) -> List[Dict[str, DataIterator]]:
+        assert len(datasets) == 1, "This example only handles the simple case"
+
+        # Configure Ray Data for ingest.
+        ctx = ray.data.DataContext.get_current()
+        ctx.execution_options = DataConfig.default_ingest_options()
+
+        # Split the stream into shards.
+        iterator_shards = datasets["train"].streaming_split(
+            world_size, equal=True, locality_hints=worker_node_ids
+        )
+
+        # Return the assigned iterators for each worker.
+        return [{"train": it} for it in iterator_shards]
+
+
+my_trainer = TorchTrainer(
+    train_loop_per_worker,
+    scaling_config=ScalingConfig(num_workers=2),
+    datasets={"train": train_ds},
+    dataset_config=MyCustomDataConfig(),
+)
+my_trainer.fit()
+# __custom_end__

doc/source/train/doc_code/dl_guide.py

@@ -0,0 +1,110 @@
+# flake8: noqa
+
+# TODO: [V2] Deprecated doc code to delete.
+import os
+
+os.environ["RAY_TRAIN_V2_ENABLED"] = "0"
+
+MOCK = True
+
+# __ft_initial_run_start__
+import os
+import tempfile
+from typing import Dict, Optional
+
+import torch
+
+import ray
+from ray import train
+from ray.train import Checkpoint
+from ray.train.torch import TorchTrainer
+
+
+def get_datasets() -> Dict[str, ray.data.Dataset]:
+    return {"train": ray.data.from_items([{"x": i, "y": 2 * i} for i in range(10)])}
+
+
+def train_loop_per_worker(config: dict):
+    from torchvision.models import resnet18
+
+    model = resnet18()
+
+    # Checkpoint loading
+    checkpoint: Optional[Checkpoint] = train.get_checkpoint()
+    if checkpoint:
+        with checkpoint.as_directory() as checkpoint_dir:
+            model_state_dict = torch.load(os.path.join(checkpoint_dir, "model.pt"))
+            model.load_state_dict(model_state_dict)
+
+    model = train.torch.prepare_model(model)
+
+    train_ds = train.get_dataset_shard("train")
+
+    for epoch in range(5):
+        # Do some training...
+
+        # Checkpoint saving
+        with tempfile.TemporaryDirectory() as tmpdir:
+            torch.save(model.module.state_dict(), os.path.join(tmpdir, "model.pt"))
+            train.report({"epoch": epoch}, checkpoint=Checkpoint.from_directory(tmpdir))
+
+
+trainer = TorchTrainer(
+    train_loop_per_worker=train_loop_per_worker,
+    datasets=get_datasets(),
+    scaling_config=train.ScalingConfig(num_workers=2),
+    run_config=train.RunConfig(
+        name="dl_trainer_restore", storage_path=os.path.expanduser("~/ray_results")
+    ),
+)
+result = trainer.fit()
+# __ft_initial_run_end__
+
+# __ft_restored_run_start__
+from ray.train.torch import TorchTrainer
+
+restored_trainer = TorchTrainer.restore(
+    path=os.path.expanduser("~/ray_results/dl_trainer_restore"),
+    datasets=get_datasets(),
+)
+# __ft_restored_run_end__
+
+
+if not MOCK:
+    # __ft_restore_from_cloud_initial_start__
+    original_trainer = TorchTrainer(
+        # ...
+        run_config=train.RunConfig(
+            # Configure cloud storage
+            storage_path="s3://results-bucket",
+            name="dl_trainer_restore",
+        ),
+    )
+    result = trainer.fit()
+    # __ft_restore_from_cloud_initial_end__
+
+    # __ft_restore_from_cloud_restored_start__
+    restored_trainer = TorchTrainer.restore(
+        "s3://results-bucket/dl_trainer_restore",
+        datasets=get_datasets(),
+    )
+    # __ft_restore_from_cloud_restored_end__
+
+
+# __ft_autoresume_start__
+experiment_path = os.path.expanduser("~/ray_results/dl_restore_autoresume")
+if TorchTrainer.can_restore(experiment_path):
+    trainer = TorchTrainer.restore(experiment_path, datasets=get_datasets())
+    result = trainer.fit()
+else:
+    trainer = TorchTrainer(
+        train_loop_per_worker=train_loop_per_worker,
+        datasets=get_datasets(),
+        scaling_config=train.ScalingConfig(num_workers=2),
+        run_config=train.RunConfig(
+            storage_path=os.path.expanduser("~/ray_results"),
+            name="dl_restore_autoresume",
+        ),
+    )
+result = trainer.fit()
+# __ft_autoresume_end__

doc/source/train/doc_code/fault_tolerance.py

@@ -0,0 +1,107 @@
+import os
+
+os.environ["RAY_TRAIN_V2_ENABLED"] = "1"
+
+# __failure_config_start__
+import ray.train
+
+# Tries to recover a run up to this many times.
+failure_config = ray.train.FailureConfig(max_failures=2)
+
+# No limit on the number of retries.
+failure_config = ray.train.FailureConfig(max_failures=-1)
+# __failure_config_end__
+
+# __worker_fault_tolerance_start__
+import tempfile
+import uuid
+
+import ray.train
+import ray.train.torch
+
+
+def train_fn_per_worker(train_loop_config: dict):
+    # [1] Train worker restoration logic.
+    checkpoint = ray.train.get_checkpoint()
+    if checkpoint:
+        with checkpoint.as_directory() as temp_checkpoint_dir:
+            # model.load_state_dict(torch.load(...))
+            ...
+
+    # [2] Checkpoint saving and reporting logic.
+    with tempfile.TemporaryDirectory() as temp_checkpoint_dir:
+        # torch.save(...)
+        ray.train.report(
+            {"loss": 0.1},
+            checkpoint=ray.train.Checkpoint.from_directory(temp_checkpoint_dir),
+        )
+
+
+trainer = ray.train.torch.TorchTrainer(
+    train_fn_per_worker,
+    scaling_config=ray.train.ScalingConfig(num_workers=4),
+    run_config=ray.train.RunConfig(
+        # (If multi-node, configure S3 / NFS as the storage path.)
+        # storage_path="s3://...",
+        name=f"train_run-{uuid.uuid4().hex}",
+        # [3] Enable worker-level fault tolerance to gracefully handle
+        # Train worker failures.
+        failure_config=ray.train.FailureConfig(max_failures=3),
+    ),
+)
+trainer.fit()
+# __worker_fault_tolerance_end__
+
+# Avoid running the code below so that the argument parser is not used.
+__name__ = "__dummy__"
+
+# __job_driver_fault_tolerance_start__
+# entrypoint.py
+
+import argparse
+import tempfile
+import uuid
+
+import ray.train
+import ray.train.torch
+
+
+def train_fn_per_worker(train_loop_config: dict):
+    # [1] Train worker restoration logic.
+    checkpoint = ray.train.get_checkpoint()
+    if checkpoint:
+        with checkpoint.as_directory() as temp_checkpoint_dir:
+            # model.load_state_dict(torch.load(...))
+            ...
+
+    # [2] Checkpoint saving and reporting logic.
+    with tempfile.TemporaryDirectory() as temp_checkpoint_dir:
+        # torch.save(...)
+        ray.train.report(
+            {"loss": 0.1},
+            checkpoint=ray.train.Checkpoint.from_directory(temp_checkpoint_dir),
+        )
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--storage_path", type=str, required=True)
+    parser.add_argument("--run_name", type=str, required=True)
+    args = parser.parse_args()
+
+    trainer = ray.train.torch.TorchTrainer(
+        train_fn_per_worker,
+        scaling_config=ray.train.ScalingConfig(num_workers=4),
+        run_config=ray.train.RunConfig(
+            # [3] Enable worker-level fault tolerance to gracefully handle
+            # Train worker failures.
+            failure_config=ray.train.FailureConfig(max_failures=3),
+            # [4] (Recommendation) The (storage_path, name) pair should be
+            # determined by the job submitter and passed in as arguments
+            # to the entrypoint script.
+            storage_path=args.storage_path,
+            name=args.run_name,
+        ),
+    )
+    trainer.fit()
+# __job_driver_fault_tolerance_end__

doc/source/train/doc_code/hvd_trainer.py

@@ -0,0 +1,80 @@
+# TODO: [V2] Deprecated doc code to delete.
+import os
+
+os.environ["RAY_TRAIN_V2_ENABLED"] = "0"
+
+import tempfile
+
+import horovod.torch as hvd
+import ray
+from ray import train
+from ray.train import Checkpoint, ScalingConfig
+import ray.train.torch  # Need this to use `train.torch.get_device()`
+from ray.train.horovod import HorovodTrainer
+import torch
+import torch.nn as nn
+
+# If using GPUs, set this to True.
+use_gpu = False
+
+
+input_size = 1
+layer_size = 15
+output_size = 1
+num_epochs = 3
+
+
+class NeuralNetwork(nn.Module):
+    def __init__(self):
+        super(NeuralNetwork, self).__init__()
+        self.layer1 = nn.Linear(input_size, layer_size)
+        self.relu = nn.ReLU()
+        self.layer2 = nn.Linear(layer_size, output_size)
+
+    def forward(self, input):
+        return self.layer2(self.relu(self.layer1(input)))
+
+
+def train_loop_per_worker():
+    hvd.init()
+    dataset_shard = train.get_dataset_shard("train")
+    model = NeuralNetwork()
+    device = train.torch.get_device()
+    model.to(device)
+    loss_fn = nn.MSELoss()
+    lr_scaler = 1
+    optimizer = torch.optim.SGD(model.parameters(), lr=0.1 * lr_scaler)
+    # Horovod: wrap optimizer with DistributedOptimizer.
+    optimizer = hvd.DistributedOptimizer(
+        optimizer,
+        named_parameters=model.named_parameters(),
+        op=hvd.Average,
+    )
+    for epoch in range(num_epochs):
+        model.train()
+        for batch in dataset_shard.iter_torch_batches(
+            batch_size=32, dtypes=torch.float
+        ):
+            inputs, labels = torch.unsqueeze(batch["x"], 1), batch["y"]
+            outputs = model(inputs)
+            loss = loss_fn(outputs, labels)
+            optimizer.zero_grad()
+            loss.backward()
+            optimizer.step()
+            print(f"epoch: {epoch}, loss: {loss.item()}")
+
+        with tempfile.TemporaryDirectory() as tmpdir:
+            torch.save(model.state_dict(), os.path.join(tmpdir, "model.pt"))
+            train.report(
+                {"loss": loss.item()}, checkpoint=Checkpoint.from_directory(tmpdir)
+            )
+
+
+train_dataset = ray.data.from_items([{"x": x, "y": x + 1} for x in range(32)])
+scaling_config = ScalingConfig(num_workers=3, use_gpu=use_gpu)
+trainer = HorovodTrainer(
+    train_loop_per_worker=train_loop_per_worker,
+    scaling_config=scaling_config,
+    datasets={"train": train_dataset},
+)
+result = trainer.fit()

doc/source/train/doc_code/key_concepts.py

@@ -0,0 +1,136 @@
+# flake8: noqa
+# isort: skip_file
+
+from pathlib import Path
+import tempfile
+
+import ray.train
+from ray.train.v2.api.data_parallel_trainer import DataParallelTrainer
+
+
+def train_fn(config):
+    for i in range(3):
+        with tempfile.TemporaryDirectory() as temp_checkpoint_dir:
+            Path(temp_checkpoint_dir).joinpath("model.pt").touch()
+            ray.train.report(
+                {"loss": i},
+                checkpoint=ray.train.Checkpoint.from_directory(temp_checkpoint_dir),
+            )
+
+    return {"total loss": 3}
+
+
+trainer = DataParallelTrainer(
+    train_fn, scaling_config=ray.train.ScalingConfig(num_workers=2)
+)
+
+
+# __run_config_start__
+import os
+
+from ray.train import RunConfig
+
+run_config = RunConfig(
+    # Name of the training run (directory name).
+    name="my_train_run",
+    # The experiment results will be saved to: storage_path/name
+    storage_path=os.path.expanduser("~/ray_results"),
+    # storage_path="s3://my_bucket/tune_results",
+)
+# __run_config_end__
+
+# __checkpoint_config_start__
+from ray.train import RunConfig, CheckpointConfig
+
+# Example 1: Only keep the 2 *most recent* checkpoints and delete the others.
+run_config = RunConfig(checkpoint_config=CheckpointConfig(num_to_keep=2))
+
+
+# Example 2: Only keep the 2 *best* checkpoints and delete the others.
+run_config = RunConfig(
+    checkpoint_config=CheckpointConfig(
+        num_to_keep=2,
+        # *Best* checkpoints are determined by these params:
+        checkpoint_score_attribute="mean_accuracy",
+        checkpoint_score_order="max",
+    ),
+    # This will store checkpoints on S3.
+    storage_path="s3://remote-bucket/location",
+)
+# __checkpoint_config_end__
+
+
+# __result_metrics_start__
+result = trainer.fit()
+
+print("Observed metrics:", result.metrics)
+# __result_metrics_end__
+
+
+# __result_dataframe_start__
+df = result.metrics_dataframe
+print("Minimum loss", min(df["loss"]))
+# __result_dataframe_end__
+
+
+# __result_return_value_start__
+print("Returned data", result.return_value)
+# __result_return_value_end__
+
+
+# __result_checkpoint_start__
+print("Last checkpoint:", result.checkpoint)
+
+with result.checkpoint.as_directory() as tmpdir:
+    # Load model from directory
+    ...
+# __result_checkpoint_end__
+
+# __result_best_checkpoint_start__
+# Print available checkpoints
+for checkpoint, metrics in result.best_checkpoints:
+    print("Loss", metrics["loss"], "checkpoint", checkpoint)
+
+# Get checkpoint with minimal loss
+best_checkpoint = min(
+    result.best_checkpoints, key=lambda checkpoint: checkpoint[1]["loss"]
+)[0]
+
+with best_checkpoint.as_directory() as tmpdir:
+    # Load model from directory
+    ...
+# __result_best_checkpoint_end__
+
+import pyarrow
+
+# __result_path_start__
+result_path: str = result.path
+result_filesystem: pyarrow.fs.FileSystem = result.filesystem
+
+print(f"Results location (fs, path) = ({result_filesystem}, {result_path})")
+# __result_path_end__
+
+
+# __result_restore_start__
+from ray.train import Result
+
+restored_result = Result.from_path(result_path)
+print("Restored loss", restored_result.metrics["loss"])
+# __result_restore_end__
+
+
+def error_train_fn(config):
+    raise RuntimeError("Simulated training error")
+
+
+trainer = DataParallelTrainer(
+    error_train_fn, scaling_config=ray.train.ScalingConfig(num_workers=1)
+)
+
+# __result_error_start__
+try:
+    result = trainer.fit()
+except ray.train.TrainingFailedError as e:
+    if isinstance(e, ray.train.WorkerGroupError):
+        print(e.worker_failures)
+# __result_error_end__

doc/source/train/doc_code/lightgbm_quickstart.py

@@ -0,0 +1,128 @@
+# flake8: noqa
+# isort: skip_file
+
+# __lightgbm_start__
+import pandas as pd
+import lightgbm as lgb
+
+# 1. Load your data as a `lightgbm.Dataset`.
+train_df = pd.read_csv("s3://ray-example-data/iris/train/1.csv")
+eval_df = pd.read_csv("s3://ray-example-data/iris/val/1.csv")
+
+train_X = train_df.drop("target", axis=1)
+train_y = train_df["target"]
+eval_X = eval_df.drop("target", axis=1)
+eval_y = eval_df["target"]
+
+train_set = lgb.Dataset(train_X, label=train_y)
+eval_set = lgb.Dataset(eval_X, label=eval_y)
+
+# 2. Define your LightGBM model training parameters.
+params = {
+    "objective": "multiclass",
+    "num_class": 3,
+    "metric": ["multi_logloss", "multi_error"],
+    "verbosity": -1,
+    "boosting_type": "gbdt",
+    "num_leaves": 31,
+    "learning_rate": 0.05,
+    "feature_fraction": 0.9,
+    "bagging_fraction": 0.8,
+    "bagging_freq": 5,
+}
+
+# 3. Do non-distributed training.
+model = lgb.train(
+    params,
+    train_set,
+    valid_sets=[eval_set],
+    valid_names=["eval"],
+    num_boost_round=100,
+)
+# __lightgbm_end__
+
+
+# __lightgbm_ray_start__
+import lightgbm as lgb
+
+import ray.train
+from ray.train.lightgbm import LightGBMTrainer, RayTrainReportCallback
+
+# 1. Load your data as a Ray Data Dataset.
+train_dataset = ray.data.read_csv("s3://anonymous@ray-example-data/iris/train")
+eval_dataset = ray.data.read_csv("s3://anonymous@ray-example-data/iris/val")
+
+
+def train_func():
+    # 2. Load your data shard as a `lightgbm.Dataset`.
+
+    # Get dataset shards for this worker
+    train_shard = ray.train.get_dataset_shard("train")
+    eval_shard = ray.train.get_dataset_shard("eval")
+
+    # Convert shards to PyArrow tables. LightGBM (>=4.2.0) supports PyArrow
+    # natively, which avoids a round-trip through pandas.
+    import pyarrow as pa
+
+    train_table = pa.concat_tables(
+        train_shard.iter_batches(batch_format="pyarrow", batch_size=None)
+    )
+    eval_table = pa.concat_tables(
+        eval_shard.iter_batches(batch_format="pyarrow", batch_size=None)
+    )
+
+    train_X = train_table.drop(["target"])
+    train_y = train_table.column("target")
+    eval_X = eval_table.drop(["target"])
+    eval_y = eval_table.column("target")
+
+    train_set = lgb.Dataset(train_X, label=train_y)
+    eval_set = lgb.Dataset(eval_X, label=eval_y)
+
+    # 3. Define your LightGBM model training parameters.
+    params = {
+        "objective": "multiclass",
+        "num_class": 3,
+        "metric": ["multi_logloss", "multi_error"],
+        "verbosity": -1,
+        "boosting_type": "gbdt",
+        "num_leaves": 31,
+        "learning_rate": 0.05,
+        "feature_fraction": 0.9,
+        "bagging_fraction": 0.8,
+        "bagging_freq": 5,
+        # Adding the lines below are the only changes needed
+        # for your `lgb.train` call!
+        "tree_learner": "data_parallel",
+        "pre_partition": True,
+        **ray.train.lightgbm.get_network_params(),
+    }
+
+    # 4. Do distributed data-parallel training.
+    # Ray Train sets up the necessary coordinator processes and
+    # environment variables for your workers to communicate with each other.
+    model = lgb.train(
+        params,
+        train_set,
+        valid_sets=[eval_set],
+        valid_names=["eval"],
+        num_boost_round=100,
+        # Optional: Use the `RayTrainReportCallback` to save and report checkpoints.
+        callbacks=[RayTrainReportCallback()],
+    )
+
+
+# 5. Configure scaling and resource requirements.
+scaling_config = ray.train.ScalingConfig(num_workers=2, resources_per_worker={"CPU": 2})
+
+# 6. Launch distributed training job.
+trainer = LightGBMTrainer(
+    train_func,
+    scaling_config=scaling_config,
+    datasets={"train": train_dataset, "eval": eval_dataset},
+)
+result = trainer.fit()
+
+# 7. Load the trained model.
+model = RayTrainReportCallback.get_model(result.checkpoint)
+# __lightgbm_ray_end__

doc/source/train/doc_code/metric_logging.py

@@ -0,0 +1,138 @@
+# flake8: noqa
+# isort: skip_file
+
+import os
+
+os.environ["RAY_TRAIN_V2_ENABLED"] = "1"
+
+
+# __torchmetrics_start__
+
+# First, pip install torchmetrics
+# This code is tested with torchmetrics==0.7.3 and torch==1.12.1
+
+import os
+import tempfile
+
+import ray.train.torch
+from ray import train
+from ray.train import ScalingConfig
+from ray.train.torch import TorchTrainer
+
+import torch
+import torch.nn as nn
+import torchmetrics
+from torch.optim import Adam
+import numpy as np
+
+
+def train_func(config):
+    n = 100
+    # create a toy dataset
+    X = torch.Tensor(np.random.normal(0, 1, size=(n, 4)))
+    X_valid = torch.Tensor(np.random.normal(0, 1, size=(n, 4)))
+    Y = torch.Tensor(np.random.uniform(0, 1, size=(n, 1)))
+    Y_valid = torch.Tensor(np.random.uniform(0, 1, size=(n, 1)))
+    # toy neural network : 1-layer
+    # wrap the model in DDP
+    model = ray.train.torch.prepare_model(nn.Linear(4, 1))
+    criterion = nn.MSELoss()
+
+    mape = torchmetrics.MeanAbsolutePercentageError()
+    # for averaging loss
+    mean_valid_loss = torchmetrics.MeanMetric()
+
+    optimizer = Adam(model.parameters(), lr=3e-4)
+    for epoch in range(config["num_epochs"]):
+        model.train()
+        y = model.forward(X)
+
+        # compute loss
+        loss = criterion(y, Y)
+
+        # back-propagate loss
+        optimizer.zero_grad()
+        loss.backward()
+        optimizer.step()
+
+        # evaluate
+        model.eval()
+        with torch.no_grad():
+            pred = model(X_valid)
+            valid_loss = criterion(pred, Y_valid)
+            # save loss in aggregator
+            mean_valid_loss(valid_loss)
+            mape(pred, Y_valid)
+
+        # collect all metrics
+        # use .item() to obtain a value that can be reported
+        valid_loss = valid_loss.item()
+        mape_collected = mape.compute().item()
+        mean_valid_loss_collected = mean_valid_loss.compute().item()
+
+        with tempfile.TemporaryDirectory() as temp_checkpoint_dir:
+            torch.save(
+                model.state_dict(), os.path.join(temp_checkpoint_dir, "model.pt")
+            )
+
+            train.report(
+                {
+                    "mape_collected": mape_collected,
+                    "valid_loss": valid_loss,
+                    "mean_valid_loss_collected": mean_valid_loss_collected,
+                },
+                checkpoint=train.Checkpoint.from_directory(temp_checkpoint_dir),
+            )
+
+        # reset for next epoch
+        mape.reset()
+        mean_valid_loss.reset()
+
+
+trainer = TorchTrainer(
+    train_func,
+    train_loop_config={"num_epochs": 5},
+    scaling_config=ScalingConfig(num_workers=2),
+)
+result = trainer.fit()
+print(result.metrics["valid_loss"], result.metrics["mean_valid_loss_collected"])
+# 0.5109779238700867 0.5512474775314331
+
+# __torchmetrics_end__
+
+# __report_callback_start__
+import os
+
+assert os.environ["RAY_TRAIN_V2_ENABLED"] == "1"
+
+from typing import Any, Dict, List, Optional
+
+import ray.train
+import ray.train.torch
+
+
+def train_fn_per_worker(config):
+    # Free-floating metrics can be accessed from the callback below.
+    ray.train.report({"rank": ray.train.get_context().get_world_rank()})
+
+
+class CustomMetricsCallback(ray.train.UserCallback):
+    def after_report(
+        self,
+        run_context,
+        metrics: List[Dict[str, Any]],
+        checkpoint: Optional[ray.train.Checkpoint],
+    ):
+        rank_0_metrics = metrics[0]
+        print(rank_0_metrics)
+        # Ex: Write metrics to a file...
+
+
+trainer = ray.train.torch.TorchTrainer(
+    train_fn_per_worker,
+    scaling_config=ray.train.ScalingConfig(num_workers=2),
+    run_config=ray.train.RunConfig(callbacks=[CustomMetricsCallback()]),
+)
+trainer.fit()
+
+# __report_callback_end__

doc/source/train/doc_code/random_text_generator.py

@@ -0,0 +1,58 @@
+import random
+import string
+import ray
+
+def random_text(length: int) -> str:
+    """Generate random text of specified length."""
+    if length <= 0:
+        return ""
+
+    if length <= 3:
+        return "".join(random.choices(string.ascii_lowercase, k=length))
+
+    words = []
+    current_length = 0
+
+    while current_length < length:
+        remaining = length - current_length
+        
+        if remaining <= 4:
+            word_length = remaining
+            word = "".join(random.choices(string.ascii_lowercase, k=word_length))
+            words.append(word)
+            break
+        else:
+            max_word_length = min(10, remaining - 1)
+            if max_word_length >= 3:
+                word_length = random.randint(3, max_word_length)
+            else:
+                word_length = remaining
+            word = "".join(random.choices(string.ascii_lowercase, k=word_length))
+            words.append(word)
+            current_length += len(word) + 1
+
+    text = " ".join(words)
+    return text[:length]
+
+def random_label() -> int:
+    """Pick a random label."""
+    labels = [0, 1, 2, 3, 4, 5, 6, 7]
+    return random.choice(labels)
+
+def create_mock_ray_text_dataset(dataset_size: int = 96, min_len: int = 5, max_len: int = 100):
+    """Create a mock Ray dataset with random text and labels."""
+    numbers = random.choices(range(min_len, max_len + 1), k=dataset_size)
+    ray_dataset = ray.data.from_items(numbers)
+
+    def map_to_text_and_label(item):
+        length = item['item']
+        text = random_text(length)
+        label = random_label()
+        return {
+            "length": length,
+            "text": text,
+            "label": label
+        }
+
+    text_dataset = ray_dataset.map(map_to_text_and_label)
+    return text_dataset

doc/source/train/doc_code/tf_starter.py

@@ -0,0 +1,84 @@
+# flake8: noqa
+# isort: skip_file
+
+# __tf_train_start__
+import os
+os.environ["TF_USE_LEGACY_KERAS"] = "1"
+
+import ray
+import tensorflow as tf
+
+from ray import train
+from ray.train import ScalingConfig
+from ray.train.tensorflow import TensorflowTrainer
+from ray.train.tensorflow.keras import ReportCheckpointCallback
+
+
+# If using GPUs, set this to True.
+use_gpu = False
+
+a = 5
+b = 10
+size = 100
+
+
+def build_model() -> tf.keras.Model:
+    model = tf.keras.Sequential(
+        [
+            tf.keras.layers.InputLayer(input_shape=()),
+            # Add feature dimension, expanding (batch_size,) to (batch_size, 1).
+            tf.keras.layers.Flatten(),
+            tf.keras.layers.Dense(10),
+            tf.keras.layers.Dense(1),
+        ]
+    )
+    return model
+
+
+def train_func(config: dict):
+    import os
+    os.environ["TF_USE_LEGACY_KERAS"] = "1"
+    import tensorflow as tf
+
+    batch_size = config.get("batch_size", 64)
+    epochs = config.get("epochs", 3)
+
+    strategy = tf.distribute.MultiWorkerMirroredStrategy()
+    with strategy.scope():
+        # Model building/compiling need to be within `strategy.scope()`.
+        multi_worker_model = build_model()
+        multi_worker_model.compile(
+            optimizer=tf.keras.optimizers.SGD(learning_rate=config.get("lr", 1e-3)),
+            loss="mean_squared_error",
+            metrics=["mean_squared_error"],
+        )
+
+    dataset = train.get_dataset_shard("train")
+
+    results = []
+    for _ in range(epochs):
+        tf_dataset = dataset.to_tf(
+            feature_columns="x", label_columns="y", batch_size=batch_size
+        )
+        history = multi_worker_model.fit(
+            tf_dataset, callbacks=[ReportCheckpointCallback()]
+        )
+        results.append(history.history)
+    return results
+
+
+config = {"lr": 1e-3, "batch_size": 32, "epochs": 4}
+
+train_dataset = ray.data.from_items(
+    [{"x": x / 200, "y": 2 * x / 200} for x in range(200)]
+)
+scaling_config = ScalingConfig(num_workers=2, use_gpu=use_gpu)
+trainer = TensorflowTrainer(
+    train_loop_per_worker=train_func,
+    train_loop_config=config,
+    scaling_config=scaling_config,
+    datasets={"train": train_dataset},
+)
+result = trainer.fit()
+print(result.metrics)
+# __tf_train_end__

doc/source/train/doc_code/train_tune_interop.py

@@ -0,0 +1,204 @@
+# flake8: noqa
+# isort: skip_file
+
+import os
+
+os.environ["RAY_TRAIN_V2_ENABLED"] = "1"
+
+# __quickstart_start__
+import random
+import tempfile
+import uuid
+
+import ray.train
+import ray.train.torch
+import ray.tune
+from ray.tune.integration.ray_train import TuneReportCallback
+
+
+# [1] Define your Ray Train worker code.
+def train_fn_per_worker(train_loop_config: dict):
+    # Unpack train worker hyperparameters.
+    # Train feeds in the `train_loop_config` defined below.
+    lr = train_loop_config["lr"]
+
+    # training code here...
+    print(
+        ray.train.get_context().get_world_size(),
+        ray.train.get_context().get_world_rank(),
+        train_loop_config,
+    )
+    # model = ray.train.torch.prepare_model(...)  # Wrap model in DDP.
+    with tempfile.TemporaryDirectory() as temp_checkpoint_dir:
+        ray.train.report(
+            {"loss": random.random()},
+            checkpoint=ray.train.Checkpoint.from_directory(temp_checkpoint_dir),
+        )
+
+
+# [2] Define a function that launches the Ray Train run.
+def train_driver_fn(config: dict):
+    # Unpack run-level hyperparameters.
+    # Tune feeds in hyperparameters defined in the `param_space` below.
+    num_workers = config["num_workers"]
+
+    trainer = ray.train.torch.TorchTrainer(
+        train_fn_per_worker,
+        train_loop_config=config["train_loop_config"],
+        scaling_config=ray.train.ScalingConfig(
+            num_workers=num_workers,
+            # Uncomment to use GPUs.
+            # use_gpu=True,
+        ),
+        run_config=ray.train.RunConfig(
+            # [3] Assign unique names to each run.
+            # Recommendation: use the trial id as part of the run name.
+            name=f"train-trial_id={ray.tune.get_context().get_trial_id()}",
+            # [4] (Optional) Pass in a `TuneReportCallback` to propagate
+            # reported results to the Tuner.
+            callbacks=[TuneReportCallback()],
+            # (If multi-node, configure S3 / NFS as the storage path.)
+            # storage_path="s3://...",
+        ),
+    )
+    trainer.fit()
+
+
+# Launch a single Train run.
+# Note that you can only create a TuneReportCallback in a Ray Tune session.
+# train_driver_fn({"num_workers": 4, "train_loop_config": {"lr": 1e-3}})
+
+
+# Launch a sweep of hyperparameters with Ray Tune.
+tuner = ray.tune.Tuner(
+    train_driver_fn,
+    param_space={
+        "num_workers": ray.tune.choice([2, 4]),
+        "train_loop_config": {
+            "lr": ray.tune.grid_search([1e-3, 3e-4]),
+            "batch_size": ray.tune.grid_search([32, 64]),
+        },
+    },
+    run_config=ray.tune.RunConfig(
+        name=f"tune_train_example-{uuid.uuid4().hex[:6]}",
+        # (If multi-node, configure S3 / NFS as the storage path.)
+        # storage_path="s3://...",
+    ),
+    # [5] (Optional) Set the maximum number of concurrent trials
+    # in order to prevent too many Train driver processes from
+    # being launched at once.
+    tune_config=ray.tune.TuneConfig(max_concurrent_trials=2),
+)
+results = tuner.fit()
+
+print(results.get_best_result(metric="loss", mode="min"))
+# __quickstart_end__
+
+# __max_concurrent_trials_start__
+# For a fixed size cluster, calculate this based on the limiting resource (ex: GPUs).
+total_cluster_gpus = 8
+num_gpu_workers_per_trial = 4
+max_concurrent_trials = total_cluster_gpus // num_gpu_workers_per_trial
+
+
+def train_driver_fn(config: dict):
+    trainer = ray.train.torch.TorchTrainer(
+        train_fn_per_worker,
+        scaling_config=ray.train.ScalingConfig(
+            num_workers=num_gpu_workers_per_trial, use_gpu=True
+        ),
+    )
+    trainer.fit()
+
+
+tuner = ray.tune.Tuner(
+    train_driver_fn,
+    tune_config=ray.tune.TuneConfig(max_concurrent_trials=max_concurrent_trials),
+)
+# __max_concurrent_trials_end__
+
+# __trainable_resources_start__
+# Cluster setup:
+# head_node:
+#     resources:
+#         CPU: 16.0
+# worker_node_cpu:
+#     resources:
+#         CPU: 32.0
+#         TRAIN_DRIVER_RESOURCE: 1.0
+# worker_node_gpu:
+#     resources:
+#         GPU: 4.0
+
+import ray.tune
+
+
+def train_driver_fn(config):
+    # trainer = TorchTrainer(...)
+    ...
+
+
+tuner = ray.tune.Tuner(
+    ray.tune.with_resources(
+        train_driver_fn,
+        # Note: 0.01 is an arbitrary value to schedule the actor
+        # onto the `worker_node_cpu` node type.
+        {"TRAIN_DRIVER_RESOURCE": 0.01},
+    ),
+)
+# __trainable_resources_end__
+
+
+# __fault_tolerance_start__
+import tempfile
+
+import ray.tune
+import ray.train
+import ray.train.torch
+
+
+def train_fn_per_worker(train_loop_config: dict):
+    # [1] Train worker restoration logic.
+    checkpoint = ray.train.get_checkpoint()
+    if checkpoint:
+        with checkpoint.as_directory() as temp_checkpoint_dir:
+            # model.load_state_dict(torch.load(...))
+            ...
+
+    with tempfile.TemporaryDirectory() as temp_checkpoint_dir:
+        # torch.save(...)
+        ray.train.report(
+            {"loss": 0.1},
+            checkpoint=ray.train.Checkpoint.from_directory(temp_checkpoint_dir),
+        )
+
+
+def train_fn_driver(config: dict):
+    trainer = ray.train.torch.TorchTrainer(
+        train_fn_per_worker,
+        run_config=ray.train.RunConfig(
+            # [2] Train driver restoration is automatic, as long as
+            # the (storage_path, name) remains the same across trial restarts.
+            # The easiest way to do this is to attach the trial ID in the name.
+            # **Do not include any timestamps or random values in the name.**
+            name=f"train-trial_id={ray.tune.get_context().get_trial_id()}",
+            # [3] Enable worker-level fault tolerance to gracefully handle
+            # Train worker failures.
+            failure_config=ray.train.FailureConfig(max_failures=3),
+            # (If multi-node, configure S3 / NFS as the storage path.)
+            # storage_path="s3://...",
+        ),
+    )
+    trainer.fit()
+
+
+tuner = ray.tune.Tuner(
+    train_fn_driver,
+    run_config=ray.tune.RunConfig(
+        # [4] Enable trial-level fault tolerance to gracefully handle
+        # Train driver process failures.
+        failure_config=ray.tune.FailureConfig(max_failures=3)
+    ),
+)
+tuner.fit()
+# __fault_tolerance_end__

doc/source/train/doc_code/tuner.py

@@ -0,0 +1,245 @@
+# flake8: noqa
+# isort: skip_file
+
+# TODO: [V2] Deprecated doc code to delete.
+import os
+
+os.environ["RAY_TRAIN_V2_ENABLED"] = "0"
+
+# __basic_start__
+import ray
+import ray.tune
+import ray.train
+from ray.tune import Tuner
+from ray.train.xgboost import XGBoostTrainer
+
+dataset = ray.data.read_csv("s3://anonymous@air-example-data/breast_cancer.csv")
+
+trainer = XGBoostTrainer(
+    label_column="target",
+    params={
+        "objective": "binary:logistic",
+        "eval_metric": ["logloss", "error"],
+        "max_depth": 4,
+    },
+    datasets={"train": dataset},
+    scaling_config=ray.train.ScalingConfig(num_workers=2),
+)
+
+# Create Tuner
+tuner = Tuner(
+    trainer,
+    # Add some parameters to tune
+    param_space={"params": {"max_depth": ray.tune.choice([4, 5, 6])}},
+    # Specify tuning behavior
+    tune_config=ray.tune.TuneConfig(metric="train-logloss", mode="min", num_samples=2),
+)
+# Run tuning job
+tuner.fit()
+# __basic_end__
+
+# __xgboost_start__
+import ray.data
+import ray.train
+import ray.tune
+from ray.tune import Tuner
+from ray.train.xgboost import XGBoostTrainer
+
+dataset = ray.data.read_csv("s3://anonymous@air-example-data/breast_cancer.csv")
+
+# Create an XGBoost trainer
+trainer = XGBoostTrainer(
+    label_column="target",
+    params={
+        "objective": "binary:logistic",
+        "eval_metric": ["logloss", "error"],
+        "max_depth": 4,
+    },
+    num_boost_round=10,
+    datasets={"train": dataset},
+)
+
+param_space = {
+    # Tune parameters directly passed into the XGBoostTrainer
+    "num_boost_round": ray.tune.randint(5, 20),
+    # `params` will be merged with the `params` defined in the above XGBoostTrainer
+    "params": {
+        "min_child_weight": ray.tune.uniform(0.8, 1.0),
+        # Below will overwrite the XGBoostTrainer setting
+        "max_depth": ray.tune.randint(1, 5),
+    },
+    # Tune the number of distributed workers
+    "scaling_config": ray.train.ScalingConfig(num_workers=ray.tune.grid_search([1, 2])),
+}
+
+tuner = Tuner(
+    trainable=trainer,
+    run_config=ray.tune.RunConfig(name="test_tuner_xgboost"),
+    param_space=param_space,
+    tune_config=ray.tune.TuneConfig(
+        mode="min", metric="train-logloss", num_samples=2, max_concurrent_trials=2
+    ),
+)
+result_grid = tuner.fit()
+# __xgboost_end__
+
+# __torch_start__
+import os
+
+import ray.train
+import ray.tune
+from ray.tune import Tuner
+from ray.train.examples.pytorch.torch_linear_example import (
+    train_func as linear_train_func,
+)
+from ray.train.torch import TorchTrainer
+
+trainer = TorchTrainer(
+    train_loop_per_worker=linear_train_func,
+    train_loop_config={"lr": 1e-2, "batch_size": 4, "epochs": 10},
+    scaling_config=ray.train.ScalingConfig(num_workers=1, use_gpu=False),
+)
+
+param_space = {
+    # The params will be merged with the ones defined in the TorchTrainer
+    "train_loop_config": {
+        # This is a parameter that hasn't been set in the TorchTrainer
+        "hidden_size": ray.tune.randint(1, 4),
+        # This will overwrite whatever was set when TorchTrainer was instantiated
+        "batch_size": ray.tune.choice([4, 8]),
+    },
+    # Tune the number of distributed workers
+    "scaling_config": ray.train.ScalingConfig(num_workers=ray.tune.grid_search([1, 2])),
+}
+
+tuner = Tuner(
+    trainable=trainer,
+    run_config=ray.tune.RunConfig(
+        name="test_tuner", storage_path=os.path.expanduser("~/ray_results")
+    ),
+    param_space=param_space,
+    tune_config=ray.tune.TuneConfig(
+        mode="min", metric="loss", num_samples=2, max_concurrent_trials=2
+    ),
+)
+result_grid = tuner.fit()
+# __torch_end__
+
+
+# __tune_dataset_start__
+import ray.data
+import ray.tune
+from ray.data.preprocessors import StandardScaler
+
+
+def get_dataset():
+    ds1 = ray.data.read_csv("s3://anonymous@air-example-data/breast_cancer.csv")
+    prep_v1 = StandardScaler(["worst radius", "worst area"])
+    ds1 = prep_v1.fit_transform(ds1)
+    return ds1
+
+
+def get_another_dataset():
+    ds2 = ray.data.read_csv(
+        "s3://anonymous@air-example-data/breast_cancer_with_categorical.csv"
+    )
+    prep_v2 = StandardScaler(["worst concavity", "worst smoothness"])
+    ds2 = prep_v2.fit_transform(ds2)
+    return ds2
+
+
+dataset_1 = get_dataset()
+dataset_2 = get_another_dataset()
+
+tuner = ray.tune.Tuner(
+    trainer,
+    param_space={
+        "datasets": {
+            "train": ray.tune.grid_search([dataset_1, dataset_2]),
+        }
+        # Your other parameters go here
+    },
+)
+# __tune_dataset_end__
+
+# __tune_optimization_start__
+from ray.tune.search.bayesopt import BayesOptSearch
+from ray.tune.schedulers import HyperBandScheduler
+from ray.tune import TuneConfig
+
+config = TuneConfig(
+    # ...
+    search_alg=BayesOptSearch(),
+    scheduler=HyperBandScheduler(),
+)
+# __tune_optimization_end__
+
+# __result_grid_inspection_start__
+from ray.tune import Tuner, TuneConfig
+
+tuner = Tuner(
+    trainable=trainer,
+    param_space=param_space,
+    tune_config=TuneConfig(mode="min", metric="loss", num_samples=5),
+)
+result_grid = tuner.fit()
+
+num_results = len(result_grid)
+
+# Check if there have been errors
+if result_grid.errors:
+    print("At least one trial failed.")
+
+# Get the best result
+best_result = result_grid.get_best_result()
+
+# And the best checkpoint
+best_checkpoint = best_result.checkpoint
+
+# And the best metrics
+best_metric = best_result.metrics
+
+# Or a dataframe for further analysis
+results_df = result_grid.get_dataframe()
+print("Shortest training time:", results_df["time_total_s"].min())
+
+# Iterate over results
+for result in result_grid:
+    if result.error:
+        print("The trial had an error:", result.error)
+        continue
+
+    print("The trial finished successfully with the metrics:", result.metrics["loss"])
+# __result_grid_inspection_end__
+
+# __run_config_start__
+import ray.tune
+
+run_config = ray.tune.RunConfig(
+    name="MyExperiment",
+    storage_path="s3://...",
+    checkpoint_config=ray.tune.CheckpointConfig(checkpoint_frequency=2),
+)
+# __run_config_end__
+
+# __tune_config_start__
+from ray.tune import TuneConfig
+from ray.tune.search.bayesopt import BayesOptSearch
+
+tune_config = TuneConfig(
+    metric="loss",
+    mode="min",
+    max_concurrent_trials=10,
+    num_samples=100,
+    search_alg=BayesOptSearch(),
+)
+# __tune_config_end__
+
+# __tune_restore_start__
+tuner = Tuner.restore(
+    path=os.path.expanduser("~/ray_results/test_tuner"),
+    trainable=trainer,
+    restart_errored=True,
+)
+tuner.fit()
+# __tune_restore_end__

doc/source/train/doc_code/xgboost_quickstart.py

@@ -0,0 +1,113 @@
+# flake8: noqa
+# isort: skip_file
+
+# __xgboost_start__
+import pandas as pd
+import xgboost
+
+# 1. Load your data as an `xgboost.DMatrix`.
+train_df = pd.read_csv("s3://ray-example-data/iris/train/1.csv")
+eval_df = pd.read_csv("s3://ray-example-data/iris/val/1.csv")
+
+train_X = train_df.drop("target", axis=1)
+train_y = train_df["target"]
+eval_X = eval_df.drop("target", axis=1)
+eval_y = eval_df["target"]
+
+dtrain = xgboost.DMatrix(train_X, label=train_y)
+deval = xgboost.DMatrix(eval_X, label=eval_y)
+
+# 2. Define your xgboost model training parameters.
+params = {
+    "tree_method": "approx",
+    "objective": "reg:squarederror",
+    "eta": 1e-4,
+    "subsample": 0.5,
+    "max_depth": 2,
+}
+
+# 3. Do non-distributed training.
+bst = xgboost.train(
+    params,
+    dtrain=dtrain,
+    evals=[(deval, "validation")],
+    num_boost_round=10,
+)
+# __xgboost_end__
+
+
+# __xgboost_ray_start__
+import xgboost
+
+import ray.train
+from ray.train.xgboost import XGBoostTrainer, RayTrainReportCallback
+
+# 1. Load your data as a Ray Data Dataset.
+train_dataset = ray.data.read_csv("s3://anonymous@ray-example-data/iris/train")
+eval_dataset = ray.data.read_csv("s3://anonymous@ray-example-data/iris/val")
+
+
+def train_func():
+    # 2. Load your data shard as an `xgboost.DMatrix`.
+
+    # Get dataset shards for this worker
+    train_shard = ray.train.get_dataset_shard("train")
+    eval_shard = ray.train.get_dataset_shard("eval")
+
+    # Convert shards to pandas DataFrames
+    train_df = train_shard.materialize().to_pandas()
+    eval_df = eval_shard.materialize().to_pandas()
+
+    train_X = train_df.drop("target", axis=1)
+    train_y = train_df["target"]
+    eval_X = eval_df.drop("target", axis=1)
+    eval_y = eval_df["target"]
+
+    dtrain = xgboost.DMatrix(train_X, label=train_y)
+    deval = xgboost.DMatrix(eval_X, label=eval_y)
+
+    # 3. Define your xgboost model training parameters.
+    params = {
+        "tree_method": "approx",
+        "objective": "reg:squarederror",
+        "eta": 1e-4,
+        "subsample": 0.5,
+        "max_depth": 2,
+    }
+
+    # 4. Do distributed data-parallel training.
+    # Ray Train sets up the necessary coordinator processes and
+    # environment variables for your workers to communicate with each other.
+    bst = xgboost.train(
+        params,
+        dtrain=dtrain,
+        evals=[(deval, "validation")],
+        num_boost_round=10,
+        # Optional: Use the `RayTrainReportCallback` to save and report checkpoints.
+        callbacks=[RayTrainReportCallback()],
+    )
+
+
+# 5. Configure scaling and resource requirements.
+scaling_config = ray.train.ScalingConfig(num_workers=2, resources_per_worker={"CPU": 2})
+
+# 6. Launch distributed training job.
+trainer = XGBoostTrainer(
+    train_func,
+    scaling_config=scaling_config,
+    datasets={"train": train_dataset, "eval": eval_dataset},
+    # If running in a multi-node cluster, this is where you
+    # should configure the run's persistent storage that is accessible
+    # across all worker nodes.
+    # run_config=ray.train.RunConfig(storage_path="s3://..."),
+)
+result = trainer.fit()
+
+# 7. Load the trained model
+import os
+
+with result.checkpoint.as_directory() as checkpoint_dir:
+    model_path = os.path.join(checkpoint_dir, RayTrainReportCallback.CHECKPOINT_NAME)
+    model = xgboost.Booster()
+    model.load_model(model_path)
+# __xgboost_ray_end__