chore: import upstream snapshot with attribution

python/ray/train/tests/conftest.py

@@ -0,0 +1,143 @@
+import logging
+
+import boto3
+import pytest
+
+import ray
+from ray._common.test_utils import simulate_s3_bucket
+from ray.cluster_utils import Cluster
+
+# Trigger pytest hook to automatically zip test cluster logs to archive dir on failure
+from ray.tests.conftest import (
+    propagate_logs,  # noqa
+    pytest_runtest_makereport,  # noqa
+)
+
+
+@pytest.fixture
+def ray_start_4_cpus():
+    address_info = ray.init(num_cpus=4)
+    yield address_info
+    # The code after the yield will run as teardown code.
+    ray.shutdown()
+
+
+@pytest.fixture
+def ray_start_1_cpu_1_gpu():
+    address_info = ray.init(num_cpus=1, num_gpus=1)
+    yield address_info
+    ray.shutdown()
+
+
+@pytest.fixture
+def ray_start_4_cpus_2_gpus():
+    address_info = ray.init(num_cpus=4, num_gpus=2)
+    yield address_info
+    # The code after the yield will run as teardown code.
+    ray.shutdown()
+
+
+@pytest.fixture
+def shutdown_only():
+    yield None
+    ray.shutdown()
+
+
+@pytest.fixture
+def ray_2_node_2_gpu():
+    cluster = Cluster()
+    for _ in range(2):
+        cluster.add_node(num_cpus=4, num_gpus=2)
+
+    ray.init(address=cluster.address)
+
+    yield
+
+    ray.shutdown()
+    cluster.shutdown()
+
+
+@pytest.fixture
+def ray_2_node_2_neuron_cores():
+    cluster = Cluster()
+    for _ in range(2):
+        cluster.add_node(num_cpus=4, resources={"neuron_cores": 2})
+
+    ray.init(address=cluster.address)
+
+    yield
+
+    ray.shutdown()
+    cluster.shutdown()
+
+
+@pytest.fixture
+def ray_start_2_cpus():
+    address_info = ray.init(num_cpus=2)
+    yield address_info
+    # The code after the yield will run as teardown code.
+    ray.shutdown()
+
+
+@pytest.fixture
+def ray_4_node_4_cpu():
+    cluster = Cluster()
+    for _ in range(4):
+        cluster.add_node(num_cpus=4)
+
+    ray.init(address=cluster.address)
+
+    yield
+
+    ray.shutdown()
+    cluster.shutdown()
+
+
+@pytest.fixture
+def ray_2_node_4_gpu():
+    cluster = Cluster()
+    for _ in range(2):
+        cluster.add_node(num_cpus=2, num_gpus=4)
+
+    ray.init(address=cluster.address)
+
+    yield
+
+    ray.shutdown()
+    cluster.shutdown()
+
+
+@pytest.fixture
+def ray_2_node_2_cpu():
+    cluster = Cluster()
+    for _ in range(2):
+        cluster.add_node(num_cpus=2)
+
+    ray.init(address=cluster.address)
+
+    yield
+
+    ray.shutdown()
+    cluster.shutdown()
+
+
+@pytest.fixture
+def mock_s3_bucket_uri():
+    from ray.air._internal.uri_utils import URI
+
+    port = 5002
+    region = "us-west-2"
+    with simulate_s3_bucket(port=port, region=region) as s3_uri:
+        s3 = boto3.client(
+            "s3", region_name=region, endpoint_url=f"http://localhost:{port}"
+        )
+        # Bucket name will be autogenerated/unique per test
+        bucket_name = URI(s3_uri).name
+        s3.create_bucket(
+            Bucket=bucket_name,
+            CreateBucketConfiguration={"LocationConstraint": region},
+        )
+        # Disable server HTTP request logging
+        logging.getLogger("werkzeug").setLevel(logging.WARNING)
+        yield s3_uri
+        logging.getLogger("werkzeug").setLevel(logging.INFO)

python/ray/train/tests/dummy_preprocessor.py

@@ -0,0 +1,29 @@
+import uuid
+
+from ray.data.preprocessor import Preprocessor
+
+
+class DummyPreprocessor(Preprocessor):
+    _is_fittable = False
+
+    def __init__(self, transform=None):
+        self.id = uuid.uuid4()
+
+        if transform is None:
+            self.transform = lambda b: b
+        else:
+            self.transform = transform
+
+    def transform_batch(self, batch):
+        self._batch_transformed = True
+        return self.transform(batch)
+
+    def _transform_pandas(self, df):
+        return df
+
+    @property
+    def has_preprocessed(self):
+        return hasattr(self, "_batch_transformed")
+
+    def __eq__(self, other_preprocessor):
+        return self.id == other_preprocessor.id

python/ray/train/tests/lightning_test_utils.py

@@ -0,0 +1,171 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch.utils.data import DataLoader
+from torchmetrics import Accuracy
+
+from ray import train
+from ray.train.lightning._lightning_utils import import_lightning
+
+pl = import_lightning()
+
+
+class LinearModule(pl.LightningModule):
+    def __init__(self, input_dim, output_dim, strategy="ddp", fail_epoch=-1) -> None:
+        super().__init__()
+        self.linear = nn.Linear(input_dim, output_dim)
+        self.loss = []
+        self.strategy = strategy
+        self.restored = train.get_checkpoint() is not None
+        self.fail_epoch = fail_epoch
+
+    def forward(self, input):
+        if isinstance(input, dict) and len(input) == 1:
+            input = list(input.values())[0]
+        return self.linear(input)
+
+    def training_step(self, batch):
+        if not self.restored and self.fail_epoch == self.current_epoch:
+            raise RuntimeError
+
+        output = self.forward(batch)
+        loss = torch.sum(output)
+        self.log("loss", loss)
+        return loss
+
+    def validation_step(self, val_batch, batch_idx):
+        loss = self.forward(val_batch)
+        self.loss.append(loss)
+        return {"val_loss": loss}
+
+    def test_step(self, batch, batch_idx):
+        loss = self.forward(batch)
+        return {"test_loss": loss}
+
+    def on_validation_epoch_end(self) -> None:
+        avg_loss = torch.stack(self.loss).mean()
+        self.log("val_loss", avg_loss)
+        self.loss.clear()
+
+    def predict_step(self, batch, batch_idx):
+        return self.forward(batch)
+
+    def configure_optimizers(self):
+        if self.strategy == "fsdp":
+            # Feed FSDP wrapped model parameters to optimizer
+            return torch.optim.AdamW(self.trainer.model.parameters(), lr=0.1)
+        else:
+            return torch.optim.AdamW(self.parameters(), lr=0.1)
+
+
+class DoubleLinearModule(pl.LightningModule):
+    def __init__(self, input_dim_1, input_dim_2, output_dim) -> None:
+        super().__init__()
+        self.linear_1 = nn.Linear(input_dim_1, output_dim)
+        self.linear_2 = nn.Linear(input_dim_2, output_dim)
+        self.loss = []
+
+    def forward(self, batch):
+        input_1 = batch["input_1"]
+        input_2 = batch["input_2"]
+        return self.linear_1(input_1) + self.linear_2(input_2)
+
+    def training_step(self, batch):
+        output = self.forward(batch)
+        loss = torch.sum(output)
+        self.log("loss", loss)
+        return loss
+
+    def validation_step(self, val_batch, batch_idx):
+        loss = self.forward(val_batch)
+        self.loss.append(loss)
+        return {"val_loss": loss}
+
+    def on_validation_epoch_end(self) -> None:
+        print("Validation Epoch:", self.current_epoch)
+        avg_loss = torch.stack(self.loss).mean()
+        self.log("val_loss", avg_loss)
+        self.loss.clear()
+
+    def predict_step(self, batch, batch_idx):
+        return self.forward(batch)
+
+    def configure_optimizers(self):
+        return torch.optim.AdamW(self.parameters(), lr=0.1)
+
+
+class DummyDataModule(pl.LightningDataModule):
+    def __init__(self, batch_size: int = 8, dataset_size: int = 256) -> None:
+        super().__init__()
+        self.batch_size = batch_size
+        self.train_data = torch.randn(dataset_size, 32)
+        self.val_data = torch.randn(dataset_size, 32)
+        self.test_data = torch.randn(dataset_size, 32)
+
+    def train_dataloader(self):
+        return DataLoader(self.train_data, batch_size=self.batch_size)
+
+    def val_dataloader(self):
+        return DataLoader(self.val_data, batch_size=self.batch_size)
+
+    def test_dataloader(self):
+        return DataLoader(self.test_data, batch_size=self.batch_size)
+
+
+class LightningMNISTClassifier(pl.LightningModule):
+    def __init__(self, lr: float, layer_1: int, layer_2: int):
+        super(LightningMNISTClassifier, self).__init__()
+
+        self.lr = lr
+        # mnist images are (1, 28, 28) (channels, width, height)
+        self.layer_1 = torch.nn.Linear(28 * 28, layer_1)
+        self.layer_2 = torch.nn.Linear(layer_1, layer_2)
+        self.layer_3 = torch.nn.Linear(layer_2, 10)
+        self.accuracy = Accuracy(task="multiclass", num_classes=10, top_k=1)
+        self.val_acc_list = []
+        self.val_loss_list = []
+
+    def forward(self, x):
+        batch_size, channels, width, height = x.size()
+        x = x.view(batch_size, -1)
+        x = self.layer_1(x)
+        x = torch.relu(x)
+        x = self.layer_2(x)
+        x = torch.relu(x)
+        x = self.layer_3(x)
+        x = torch.log_softmax(x, dim=1)
+        return x
+
+    def configure_optimizers(self):
+        return torch.optim.Adam(self.parameters(), lr=self.lr)
+
+    def training_step(self, train_batch, batch_idx):
+        x, y = train_batch
+        logits = self.forward(x)
+        loss = F.nll_loss(logits, y)
+        acc = self.accuracy(logits, y)
+        self.log("ptl/train_loss", loss)
+        self.log("ptl/train_accuracy", acc)
+        return loss
+
+    def validation_step(self, val_batch, batch_idx):
+        x, y = val_batch
+        logits = self.forward(x)
+        loss = F.nll_loss(logits, y)
+        acc = self.accuracy(logits, y)
+        self.val_acc_list.append(acc)
+        self.val_loss_list.append(loss)
+        return {"val_loss": loss, "val_accuracy": acc}
+
+    def on_validation_epoch_end(self):
+        avg_loss = torch.stack(self.val_loss_list).mean()
+        avg_acc = torch.stack(self.val_acc_list).mean()
+        self.log("ptl/val_loss", avg_loss)
+        self.log("ptl/val_accuracy", avg_acc)
+        self.val_acc_list.clear()
+        self.val_loss_list.clear()
+
+    def predict_step(self, batch, batch_idx, dataloader_idx=None):
+        x = batch
+        logits = self.forward(x)
+        return torch.argmax(logits, dim=-1)

python/ray/train/tests/test_api_migrations.py

@@ -0,0 +1,158 @@
+import sys
+import warnings
+
+import pytest
+
+import ray.train
+import ray.tune
+from ray.train.constants import ENABLE_V2_MIGRATION_WARNINGS_ENV_VAR
+from ray.train.data_parallel_trainer import DataParallelTrainer
+from ray.util.annotations import RayDeprecationWarning
+
+
+@pytest.fixture(autouse=True)
+def enable_v2_migration_deprecation_messages(monkeypatch):
+    monkeypatch.setenv(ENABLE_V2_MIGRATION_WARNINGS_ENV_VAR, "1")
+    yield
+    monkeypatch.delenv(ENABLE_V2_MIGRATION_WARNINGS_ENV_VAR)
+
+
+def test_trainer_restore():
+    with pytest.warns(RayDeprecationWarning, match="restore"):
+        try:
+            DataParallelTrainer.restore("dummy")
+        except Exception:
+            pass
+
+    with pytest.warns(RayDeprecationWarning, match="can_restore"):
+        try:
+            DataParallelTrainer.can_restore("dummy")
+        except Exception:
+            pass
+
+
+def test_trainer_valid_configs(ray_start_4_cpus, tmp_path):
+    with warnings.catch_warnings(record=True) as w:
+        warnings.simplefilter("always")
+        DataParallelTrainer(
+            lambda _: None,
+            scaling_config=ray.train.ScalingConfig(num_workers=1),
+            run_config=ray.train.RunConfig(
+                storage_path=tmp_path,
+                failure_config=ray.train.FailureConfig(max_failures=1),
+            ),
+        ).fit()
+
+    for warning in w:
+        assert not (
+            warning.category == RayDeprecationWarning
+            and "`RunConfig` class should be imported from `ray.tune`"
+            in str(warning.message)
+        )
+
+
+def test_trainer_deprecated_configs():
+    with pytest.warns(RayDeprecationWarning, match="metadata"):
+        DataParallelTrainer(
+            lambda _: None,
+            metadata={"dummy": "dummy"},
+        )
+
+    with pytest.warns(RayDeprecationWarning, match="resume_from_checkpoint"):
+        DataParallelTrainer(
+            lambda _: None,
+            resume_from_checkpoint=ray.train.Checkpoint.from_directory("dummy"),
+        )
+
+    with pytest.warns(RayDeprecationWarning, match="fail_fast"):
+        DataParallelTrainer(
+            lambda _: None,
+            run_config=ray.train.RunConfig(
+                failure_config=ray.train.FailureConfig(fail_fast=True)
+            ),
+        )
+
+    with pytest.warns(RayDeprecationWarning, match="trainer_resources"):
+        DataParallelTrainer(
+            lambda _: None,
+            scaling_config=ray.train.ScalingConfig(trainer_resources={"CPU": 1}),
+        )
+
+    with pytest.warns(RayDeprecationWarning, match="verbose"):
+        DataParallelTrainer(
+            lambda _: None,
+            run_config=ray.train.RunConfig(verbose=True),
+        )
+
+    with pytest.warns(RayDeprecationWarning, match="log_to_file"):
+        DataParallelTrainer(
+            lambda _: None,
+            run_config=ray.train.RunConfig(log_to_file=True),
+        )
+
+    with pytest.warns(RayDeprecationWarning, match="stop"):
+        DataParallelTrainer(
+            lambda _: None,
+            run_config=ray.train.RunConfig(stop={"training_iteration": 1}),
+        )
+
+    with pytest.warns(RayDeprecationWarning, match="callbacks"):
+        DataParallelTrainer(
+            lambda _: None,
+            run_config=ray.train.RunConfig(callbacks=[ray.tune.Callback()]),
+        )
+
+    with pytest.warns(RayDeprecationWarning, match="progress_reporter"):
+        DataParallelTrainer(
+            lambda _: None,
+            run_config=ray.train.RunConfig(
+                progress_reporter=ray.tune.ProgressReporter()
+            ),
+        )
+
+    with pytest.warns(RayDeprecationWarning, match="sync_config"):
+        DataParallelTrainer(
+            lambda _: None,
+            run_config=ray.train.RunConfig(
+                sync_config=ray.train.SyncConfig(sync_artifacts=True)
+            ),
+        )
+
+
+def test_train_context_deprecations(ray_start_4_cpus, tmp_path):
+    def train_fn_per_worker(config):
+        with pytest.warns(RayDeprecationWarning, match="get_trial_dir"):
+            ray.train.get_context().get_trial_dir()
+
+        with pytest.warns(RayDeprecationWarning, match="get_trial_id"):
+            ray.train.get_context().get_trial_id()
+
+        with pytest.warns(RayDeprecationWarning, match="get_trial_name"):
+            ray.train.get_context().get_trial_name()
+
+        with pytest.warns(RayDeprecationWarning, match="get_trial_resources"):
+            ray.train.get_context().get_trial_resources()
+
+    trainer = DataParallelTrainer(
+        train_fn_per_worker,
+        scaling_config=ray.train.ScalingConfig(num_workers=1),
+        run_config=ray.train.RunConfig(storage_path=tmp_path),
+    )
+    trainer.fit()
+
+
+def test_v2_enabled_error(monkeypatch):
+    """Running a V1 Trainer with V2 enabled should raise an error."""
+    from ray.train.v2._internal.constants import V2_ENABLED_ENV_VAR
+
+    monkeypatch.setenv(V2_ENABLED_ENV_VAR, "1")
+
+    with pytest.raises(DeprecationWarning, match="Detected use of a deprecated"):
+        DataParallelTrainer(
+            lambda _: None,
+            scaling_config=ray.train.ScalingConfig(num_workers=1),
+        )
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main(["-v", "-x", __file__]))

python/ray/train/tests/test_backend.py

@@ -0,0 +1,660 @@
+import math
+import os
+import sys
+import tempfile
+import time
+from typing import Set
+from unittest.mock import patch
+
+import pytest
+
+import ray
+from ray import train
+from ray._private.accelerators.neuron import NEURON_RT_VISIBLE_CORES_ENV_VAR
+from ray.air._internal.util import StartTraceback
+
+# Trigger pytest hook to automatically zip test cluster logs to archive dir on failure
+from ray.tests.conftest import pytest_runtest_makereport  # noqa
+from ray.train import DataConfig
+from ray.train._internal.backend_executor import (
+    BackendExecutor,
+    InactiveWorkerGroupError,
+    TrainBackendError,
+    TrainingWorkerError,
+)
+from ray.train._internal.storage import StorageContext
+from ray.train._internal.worker_group import WorkerGroup, WorkerMetadata
+from ray.train.backend import Backend, BackendConfig
+from ray.train.constants import (
+    ENABLE_SHARE_CUDA_VISIBLE_DEVICES_ENV,
+    ENABLE_SHARE_NEURON_CORES_ACCELERATOR_ENV,
+    JAX_DISTRIBUTED_SHUTDOWN_TIMEOUT_S,
+    TORCH_PROCESS_GROUP_SHUTDOWN_TIMEOUT_S,
+    TRAIN_ENABLE_WORKER_SPREAD_ENV,
+)
+from ray.train.torch import TorchConfig
+from ray.train.v2.jax.config import JaxConfig
+from ray.util.placement_group import get_current_placement_group
+from ray.util.scheduling_strategies import PlacementGroupSchedulingStrategy
+from ray.util.state import list_actors
+
+
+@pytest.fixture
+def ray_start_2_cpus():
+    address_info = ray.init(num_cpus=2)
+    yield address_info
+    ray.shutdown()
+
+
+def _start_training(backend_executor: BackendExecutor, fn):
+    storage = StorageContext(
+        storage_path=tempfile.mkdtemp(),
+        experiment_dir_name="exp_name",
+        trial_dir_name="trial_name",
+    )
+    backend_executor.start_training(
+        train_func=fn,
+        datasets={},
+        metadata={},
+        data_config=DataConfig(),
+        storage=storage,
+    )
+
+
+def gen_execute_special(special_f):
+    def execute_async_special(self, f):
+        """Runs f on worker 0, special_f on other workers."""
+        futures = [
+            self.workers[0]
+            .actor._RayTrainWorker__execute.options(name=f.__name__)
+            .remote(f)
+        ]
+        for worker in self.workers[1:]:
+            futures.append(
+                worker.actor._RayTrainWorker__execute.options(
+                    name=special_f.__name__
+                ).remote(special_f)
+            )
+        return futures
+
+    return execute_async_special
+
+
+class TestConfig(BackendConfig):
+    @property
+    def backend_cls(self):
+        return TestBackend
+
+
+class TestBackend(Backend):
+    def on_start(self, worker_group: WorkerGroup, backend_config: TestConfig):
+        pass
+
+    def on_shutdown(self, worker_group: WorkerGroup, backend_config: TestConfig):
+        pass
+
+
+original_add_workers = WorkerGroup.add_workers
+
+
+def mock_add_workers(self, num_workers):
+    original_add_workers(self, num_workers)
+    for i, worker in enumerate(self.workers):
+        metadata = WorkerMetadata(
+            node_id=str(i % 2),
+            node_ip=str(i % 2),
+            hostname=0,
+            resource_ids={"GPU": ["0"]},
+            pid=0,
+        )
+        worker.metadata = metadata
+
+
+def mock_add_workers_to_nodes_with_same_ip(self, num_workers):
+    original_add_workers(self, num_workers)
+    for i, worker in enumerate(self.workers):
+        metadata = WorkerMetadata(
+            node_id=str(i % 2),
+            node_ip=0,
+            hostname=0,
+            resource_ids={"GPU": ["0"]},
+            pid=0,
+        )
+        worker.metadata = metadata
+
+
+def test_start(ray_start_2_cpus):
+    config = TestConfig()
+    e = BackendExecutor(config, num_workers=2)
+    with pytest.raises(InactiveWorkerGroupError):
+        _start_training(e, lambda: 1)
+    e.start()
+    assert len(e.worker_group) == 2
+
+
+def test_initialization_hook(ray_start_2_cpus):
+    config = TestConfig()
+    e = BackendExecutor(config, num_workers=2)
+
+    def init_hook():
+        import os
+
+        os.environ["TEST"] = "1"
+
+    e.start(initialization_hook=init_hook)
+
+    def check():
+        import os
+
+        return os.getenv("TEST", "0")
+
+    _start_training(e, check)
+    assert e.finish_training() == ["1", "1"]
+
+
+def test_shutdown(ray_start_2_cpus):
+    config = TestConfig()
+    e = BackendExecutor(config, num_workers=2)
+    e.start()
+    assert len(e.worker_group) == 2
+    e.shutdown()
+    with pytest.raises(InactiveWorkerGroupError):
+        _start_training(e, lambda: 1)
+
+
+def test_train(ray_start_2_cpus):
+    config = TestConfig()
+    e = BackendExecutor(config, num_workers=2)
+    e.start()
+
+    _start_training(e, lambda: 1)
+    assert e.finish_training() == [1, 1]
+
+
+def test_local_ranks(ray_start_2_cpus):
+    config = TestConfig()
+    e = BackendExecutor(config, num_workers=2)
+    e.start()
+
+    def train_func():
+        return train.get_context().get_local_rank()
+
+    _start_training(e, train_func)
+    assert set(e.finish_training()) == {0, 1}
+
+
+def test_local_ranks_with_same_ip_nodes(ray_2_node_2_cpu):
+    config = TestConfig()
+    e = BackendExecutor(config, num_workers=4)
+    e.start()
+
+    def train_func():
+        return train.get_context().get_local_rank()
+
+    _start_training(e, train_func)
+    assert list(e.finish_training()) == [0, 1, 0, 1]
+
+
+def test_local_world_size(ray_2_node_2_cpu):
+    config = TestConfig()
+    with patch.object(WorkerGroup, "add_workers", mock_add_workers):
+        e = BackendExecutor(config, num_workers=3)
+        e.start()
+
+        def train_func():
+            return train.get_context().get_local_world_size()
+
+        _start_training(e, train_func)
+        assert list(e.finish_training()) == [2, 2, 1]
+
+
+def test_local_world_size_with_same_ip_nodes(ray_2_node_2_cpu):
+    config = TestConfig()
+    with patch.object(
+        WorkerGroup, "add_workers", mock_add_workers_to_nodes_with_same_ip
+    ):
+        e = BackendExecutor(config, num_workers=3)
+        e.start()
+
+        def train_func():
+            return train.get_context().get_local_world_size()
+
+        _start_training(e, train_func)
+        assert list(e.finish_training()) == [2, 2, 1]
+
+
+def test_node_ranks(ray_2_node_2_cpu):
+    config = TestConfig()
+    with patch.object(WorkerGroup, "add_workers", mock_add_workers):
+        e = BackendExecutor(config, num_workers=3)
+        e.start()
+
+        def train_func():
+            return train.get_context().get_node_rank()
+
+        _start_training(e, train_func)
+        assert list(e.finish_training()) == [0, 0, 1]
+
+
+def test_node_ranks_with_same_ip_nodes(ray_2_node_2_cpu):
+    config = TestConfig()
+    with patch.object(
+        WorkerGroup, "add_workers", mock_add_workers_to_nodes_with_same_ip
+    ):
+        e = BackendExecutor(config, num_workers=3)
+        e.start()
+
+        def train_func():
+            return train.get_context().get_node_rank()
+
+        _start_training(e, train_func)
+        assert list(e.finish_training()) == [0, 0, 1]
+
+
+def test_train_failure(ray_start_2_cpus):
+    config = TestConfig()
+    e = BackendExecutor(config, num_workers=2)
+    e.start()
+
+    with pytest.raises(StartTraceback) as exc:
+        e.get_next_results()
+    assert isinstance(exc.value.__cause__, TrainBackendError)
+
+    with pytest.raises(StartTraceback) as exc:
+        e.pause_reporting()
+    assert isinstance(exc.value.__cause__, TrainBackendError)
+
+    with pytest.raises(StartTraceback) as exc:
+        e.finish_training()
+    assert isinstance(exc.value.__cause__, TrainBackendError)
+
+    _start_training(e, lambda: 1)
+
+    with pytest.raises(StartTraceback) as exc:
+        _start_training(e, lambda: 2)
+    assert isinstance(exc.value.__cause__, TrainBackendError)
+
+    assert e.finish_training() == [1, 1]
+
+
+def test_single_worker_user_failure(ray_start_2_cpus):
+    """Tests if training fails immediately if one worker raises an Exception
+    while executing the user training code."""
+    config = TestConfig()
+    e = BackendExecutor(config, num_workers=2)
+    e.start()
+
+    def single_worker_user_failure():
+        if train.get_context().get_world_rank() == 0:
+            raise RuntimeError
+        else:
+            time.sleep(1000000)
+
+    _start_training(e, single_worker_user_failure)
+
+    with pytest.raises(StartTraceback) as exc:
+        e.get_next_results()
+    assert isinstance(exc.value.__cause__, RuntimeError)
+
+
+def test_single_worker_actor_failure(ray_start_2_cpus):
+    """Tests is training fails immediately if one worker actor dies."""
+    config = TestConfig()
+    e = BackendExecutor(config, num_workers=2)
+    e.start()
+
+    def single_worker_actor_failure():
+        if train.get_context().get_world_rank() == 0:
+            # Simulate actor failure
+            os._exit(1)
+        else:
+            time.sleep(1000)
+
+    _start_training(e, single_worker_actor_failure)
+
+    with pytest.raises(TrainingWorkerError):
+        e.get_next_results()
+
+
+@pytest.mark.skipif(
+    sys.version_info >= (3, 12), reason="tensorflow is not supported in python 3.12+"
+)
+def test_tensorflow_start(ray_start_2_cpus):
+    from ray.train.tensorflow import TensorflowConfig
+
+    num_workers = 2
+    tensorflow_config = TensorflowConfig()
+    e = BackendExecutor(tensorflow_config, num_workers=num_workers)
+    e.start()
+
+    def get_tf_config():
+        import json
+        import os
+
+        return json.loads(os.environ["TF_CONFIG"])
+
+    _start_training(e, get_tf_config)
+    results = e.finish_training()
+    assert len(results) == num_workers
+
+    workers = [result["cluster"]["worker"] for result in results]
+    assert all(worker == workers[0] for worker in workers)
+
+    indexes = [result["task"]["index"] for result in results]
+    assert len(set(indexes)) == num_workers
+
+
+@pytest.mark.parametrize("init_method", ["env", "tcp"])
+def test_torch_start_shutdown(ray_start_2_cpus, init_method):
+    torch_config = TorchConfig(backend="gloo", init_method=init_method)
+    e = BackendExecutor(torch_config, num_workers=2)
+    e.start()
+
+    def check_process_group():
+        import torch
+
+        return (
+            torch.distributed.is_initialized()
+            and torch.distributed.get_world_size() == 2
+        )
+
+    _start_training(e, check_process_group)
+    assert all(e.finish_training())
+
+    e._backend.on_shutdown(e.worker_group, e._backend_config)
+
+    _start_training(e, check_process_group)
+    assert not any(e.finish_training())
+
+
+@pytest.mark.parametrize(
+    "init_method, timeout_s", [("env", 5), ("tcp", 5), ("env", 0), ("tcp", 0)]
+)
+def test_torch_process_group_shutdown_timeout(
+    ray_start_2_cpus, monkeypatch, init_method, timeout_s
+):
+    monkeypatch.setenv(TORCH_PROCESS_GROUP_SHUTDOWN_TIMEOUT_S, timeout_s)
+    torch_config = TorchConfig(backend="gloo", init_method=init_method)
+    e = BackendExecutor(torch_config, num_workers=2)
+    e.start()
+
+    _start_training(e, lambda: 1)
+    assert e.finish_training() == [1, 1]
+
+    # Verify that we do not raise an exception even if we time out
+    e._backend.on_shutdown(e.worker_group, e._backend_config)
+
+
+@pytest.mark.parametrize(
+    "worker_results",
+    [
+        (1, [[0]]),
+        (2, [[0, 1]] * 2),
+        (3, [[0]] + [[0, 1]] * 2),
+        (4, [[0, 1]] * 4),
+    ],
+)
+def test_cuda_visible_devices(ray_2_node_2_gpu, worker_results):
+    config = TestConfig()
+
+    def get_resources():
+        cuda_visible_devices = os.environ["CUDA_VISIBLE_DEVICES"]
+        # Sort the cuda visible devices to have exact match with expected result.
+        sorted_devices = [
+            int(device) for device in sorted(cuda_visible_devices.split(","))
+        ]
+        return sorted_devices
+
+    num_workers, expected_results = worker_results
+
+    os.environ[ENABLE_SHARE_CUDA_VISIBLE_DEVICES_ENV] = "1"
+    e = BackendExecutor(
+        config, num_workers=num_workers, resources_per_worker={"GPU": 1}
+    )
+    e.start()
+    _start_training(e, get_resources)
+    results = e.finish_training()
+    results.sort()
+    assert results == expected_results
+
+
+@pytest.mark.parametrize(
+    "worker_results",
+    [
+        (1, [[0]]),
+        (
+            2,
+            [[0]] * 2,
+        ),
+        (3, [[0, 1]] * 3),
+        (4, [[0, 1]] * 4),
+        (5, [[0]] + [[0, 1]] * 4),
+        (6, [[0]] * 2 + [[0, 1]] * 4),
+        (7, [[0, 1]] * 7),
+        (8, [[0, 1]] * 8),
+    ],
+)
+def test_cuda_visible_devices_fractional(ray_2_node_2_gpu, worker_results):
+    config = TestConfig()
+
+    if worker_results[0] != len(worker_results[1]):
+        raise ValueError(
+            "Invalid test parameter. Length of expected result should "
+            "match number of workers."
+        )
+
+    def get_resources():
+        cuda_visible_devices = os.environ["CUDA_VISIBLE_DEVICES"]
+        # Sort the cuda visible devices to have exact match with expected result.
+        sorted_devices = [
+            int(device) for device in sorted(cuda_visible_devices.split(","))
+        ]
+        return sorted_devices
+
+    num_workers, expected_results = worker_results
+
+    os.environ[ENABLE_SHARE_CUDA_VISIBLE_DEVICES_ENV] = "1"
+    e = BackendExecutor(
+        config, num_workers=num_workers, resources_per_worker={"GPU": 0.5}
+    )
+    e.start()
+    _start_training(e, get_resources)
+    results = e.finish_training()
+    results.sort()
+    assert results == expected_results
+
+
+@pytest.mark.parametrize(
+    "worker_results",
+    [
+        (1, [[0, 1]]),
+        (2, [[0, 1, 2, 3]] * 2),
+        (3, [[0, 1]] + [[0, 1, 2, 3]] * 2),
+        (4, [[0, 1, 2, 3]] * 4),
+    ],
+)
+def test_cuda_visible_devices_multiple(ray_2_node_4_gpu, worker_results):
+    config = TestConfig()
+
+    def get_resources():
+        cuda_visible_devices = os.environ["CUDA_VISIBLE_DEVICES"]
+        # Sort the cuda visible devices to have exact match with expected result.
+        sorted_devices = [
+            int(device) for device in sorted(cuda_visible_devices.split(","))
+        ]
+        return sorted_devices
+
+    if worker_results[0] != len(worker_results[1]):
+        raise ValueError(
+            "Invalid test parameter. Length of expected result should "
+            "match number of workers."
+        )
+
+    num_workers, expected_results = worker_results
+
+    os.environ[ENABLE_SHARE_CUDA_VISIBLE_DEVICES_ENV] = "1"
+    e = BackendExecutor(
+        config, num_workers=num_workers, resources_per_worker={"GPU": 2}
+    )
+    e.start()
+    _start_training(e, get_resources)
+    results = e.finish_training()
+    results.sort()
+    assert results == expected_results
+
+
+@pytest.mark.parametrize(
+    "worker_results",
+    [
+        (1, [[0]]),
+        (2, [[0, 1]] * 2),
+        (3, [[0]] + [[0, 1]] * 2),
+        (4, [[0, 1]] * 4),
+    ],
+)
+def test_neuron_core_accelerator_ids(ray_2_node_2_neuron_cores, worker_results):
+    config = TestConfig()
+
+    def get_resources():
+        neuron_resource_ids = os.environ[NEURON_RT_VISIBLE_CORES_ENV_VAR]
+        # Sort the runtime ids to have exact match with expected result.
+        sorted_devices = [
+            int(device) for device in sorted(neuron_resource_ids.split(","))
+        ]
+        return sorted_devices
+
+    num_workers, expected_results = worker_results
+    # sharing enabled by default
+    os.environ.pop(ENABLE_SHARE_NEURON_CORES_ACCELERATOR_ENV, None)
+    e = BackendExecutor(
+        config,
+        num_workers=num_workers,
+        resources_per_worker={"neuron_cores": 1},
+    )
+    e.start()
+    _start_training(e, get_resources)
+    results = e.finish_training()
+    results.sort()
+    assert results == expected_results
+
+
+@pytest.mark.parametrize(
+    "worker_results",
+    [
+        (1, [[0]]),
+        (2, [[0]] + [[1]]),
+        (3, [[0]] * 2 + [[1]]),
+        (4, [[0]] * 2 + [[1]] * 2),
+    ],
+)
+def test_neuron_core_accelerator_ids_sharing_disabled(
+    ray_2_node_2_neuron_cores, worker_results
+):
+    config = TestConfig()
+
+    def get_resources():
+        neuron_resource_ids = os.environ[NEURON_RT_VISIBLE_CORES_ENV_VAR]
+        # Sort the runtime ids to have exact match with expected result.
+        sorted_devices = [
+            int(device) for device in sorted(neuron_resource_ids.split(","))
+        ]
+        return sorted_devices
+
+    num_workers, expected_results = worker_results
+
+    os.environ[ENABLE_SHARE_NEURON_CORES_ACCELERATOR_ENV] = "0"
+    e = BackendExecutor(
+        config,
+        num_workers=num_workers,
+        resources_per_worker={"neuron_cores": 1},
+    )
+    e.start()
+    _start_training(e, get_resources)
+    results = e.finish_training()
+    results.sort()
+    assert results == expected_results
+
+
+def get_node_id_set() -> Set[str]:
+    return {a.node_id for a in list_actors()}
+
+
+@pytest.mark.parametrize("num_workers", [3, 4, 5])
+def test_placement_group_pack(ray_4_node_4_cpu, num_workers):
+    """Tests that workers are packed on nodes."""
+    config = TestConfig()
+    e = BackendExecutor(config, num_workers=num_workers)
+    e.start()
+    node_id_set = get_node_id_set()
+    assert len(node_id_set) == math.ceil(num_workers / 4)
+
+
+@pytest.mark.parametrize("num_workers", [3, 4, 5])
+def test_placement_group_spread(ray_4_node_4_cpu, num_workers):
+    """Tests that workers are spread across nodes."""
+    os.environ[TRAIN_ENABLE_WORKER_SPREAD_ENV] = "1"
+    config = TestConfig()
+    e = BackendExecutor(config, num_workers=num_workers)
+    e.start()
+    node_id_set = get_node_id_set()
+    assert len(node_id_set) == min(num_workers, 4)
+
+
+@pytest.mark.parametrize("placement_group_capture_child_tasks", [True, False])
+def test_placement_group_parent(ray_4_node_4_cpu, placement_group_capture_child_tasks):
+    """Tests that parent placement group will be used."""
+    num_workers = 2
+    bundle = {"CPU": 1}
+    bundles = [bundle.copy() for _ in range(num_workers + 1)]
+    placement_group = ray.util.placement_group(bundles)
+
+    def train_func():
+        return get_current_placement_group().id
+
+    @ray.remote
+    def test():
+        config = TestConfig()
+        e = BackendExecutor(config, num_workers=2)
+        e.start()
+        _start_training(e, train_func)
+        return e.finish_training()
+
+    results_future = test.options(
+        scheduling_strategy=PlacementGroupSchedulingStrategy(
+            placement_group=placement_group,
+            placement_group_capture_child_tasks=placement_group_capture_child_tasks,
+        ),
+    ).remote()
+    results = ray.get(results_future)
+    for worker_result in results:
+        if placement_group_capture_child_tasks:
+            assert worker_result == placement_group.id
+        else:
+            assert worker_result != placement_group.id
+
+
+@pytest.mark.parametrize("timeout_s", [5, 0])
+@pytest.mark.skipif(
+    sys.version_info >= (3, 12),
+    reason="Current jax version is not supported in python 3.12+",
+)
+def test_jax_distributed_shutdown_timeout(ray_start_2_cpus, monkeypatch, timeout_s):
+    """Test that JAX distributed shutdown respects the timeout env var."""
+    monkeypatch.setenv(JAX_DISTRIBUTED_SHUTDOWN_TIMEOUT_S, str(timeout_s))
+    jax_config = JaxConfig(use_tpu=True)
+    e = BackendExecutor(jax_config, num_workers=2)
+    e.start()
+
+    _start_training(e, lambda: 1)
+    assert e.finish_training() == [1, 1]
+
+    # Verify that we do not raise an exception even if we time out
+    e._backend.on_shutdown(e.worker_group, e._backend_config)
+
+
+if __name__ == "__main__":
+    import sys
+
+    import pytest
+
+    sys.exit(pytest.main(["-v", "-x", __file__]))

python/ray/train/tests/test_base_trainer.py

@@ -0,0 +1,197 @@
+import logging
+import tempfile
+
+import numpy as np
+import pytest
+
+import ray
+from ray import train, tune
+from ray.data.context import DataContext
+from ray.train import Checkpoint, ScalingConfig
+from ray.train._internal.session import get_session
+from ray.train.base_trainer import format_datasets_for_repr
+from ray.train.trainer import BaseTrainer
+from ray.util.placement_group import get_current_placement_group
+
+logger = logging.getLogger(__name__)
+
+
+class DummyTrainer(BaseTrainer):
+    _scaling_config_allowed_keys = BaseTrainer._scaling_config_allowed_keys + [
+        "num_workers",
+        "use_gpu",
+        "resources_per_worker",
+        "placement_strategy",
+    ]
+
+    def __init__(self, train_loop, custom_arg=None, **kwargs):
+        self.custom_arg = custom_arg
+        self.train_loop = train_loop
+        super().__init__(**kwargs)
+
+    def training_loop(self) -> None:
+        self.train_loop(self)
+
+
+def test_trainer_fit(ray_start_4_cpus):
+    def training_loop(self):
+        train.report(dict(my_metric=1))
+
+    trainer = DummyTrainer(train_loop=training_loop)
+    result = trainer.fit()
+    assert result.metrics["my_metric"] == 1
+
+
+def test_validate_datasets(ray_start_4_cpus):
+    with pytest.raises(ValueError) as e:
+        DummyTrainer(train_loop=None, datasets=1)
+    assert "`datasets` should be a dict mapping" in str(e.value)
+
+    with pytest.raises(ValueError) as e:
+        DummyTrainer(train_loop=None, datasets={"train": 1})
+    assert "The Dataset under train key is not a `ray.data.Dataset`"
+
+
+def test_resources(ray_start_4_cpus):
+    def check_cpus(self):
+        assert ray.available_resources()["CPU"] == 2
+
+    assert ray.available_resources()["CPU"] == 4
+    trainer = DummyTrainer(
+        check_cpus,
+        scaling_config=ScalingConfig(
+            trainer_resources={"CPU": 2}, resources_per_worker={}
+        ),
+    )
+    trainer.fit()
+
+
+def test_arg_override(ray_start_4_cpus):
+    def check_override(self):
+        assert self.scaling_config.num_workers == 1
+        # Should do deep update.
+        assert not self.custom_arg["outer"]["inner"]
+        assert self.custom_arg["outer"]["fixed"] == 1
+
+        pg = get_current_placement_group()
+        assert len(pg.bundle_specs) == 1  # Merged trainer and worker bundle
+
+    scale_config = ScalingConfig(num_workers=4)
+    trainer = DummyTrainer(
+        check_override,
+        custom_arg={"outer": {"inner": True, "fixed": 1}},
+        scaling_config=scale_config,
+    )
+
+    new_config = {
+        "custom_arg": {"outer": {"inner": False}},
+        "scaling_config": ScalingConfig(num_workers=1),
+    }
+
+    tune.run(trainer.as_trainable(), config=new_config)
+
+
+def test_reserved_cpu_warnings_no_cpu_usage(ray_start_1_cpu_1_gpu):
+    """Ensure there is no divide by zero error if trial requires no CPUs."""
+
+    def train_loop(config):
+        pass
+
+    trainer = DummyTrainer(
+        train_loop,
+        scaling_config=ScalingConfig(
+            num_workers=1, use_gpu=True, trainer_resources={"CPU": 0}
+        ),
+        datasets={"train": ray.data.range(10)},
+    )
+    trainer.fit()
+
+
+def test_setup(ray_start_4_cpus):
+    def check_setup(self):
+        assert self._has_setup
+
+    class DummyTrainerWithSetup(DummyTrainer):
+        def setup(self):
+            self._has_setup = True
+
+    trainer = DummyTrainerWithSetup(check_setup)
+    trainer.fit()
+
+
+def test_repr(ray_start_4_cpus):
+    def training_loop(self):
+        pass
+
+    trainer = DummyTrainer(
+        training_loop,
+        datasets={
+            "train": ray.data.from_items([1, 2, 3]),
+        },
+    )
+
+    representation = repr(trainer)
+
+    assert "DummyTrainer" in representation
+
+
+def test_metadata_propagation(ray_start_4_cpus):
+    class MyTrainer(BaseTrainer):
+        def training_loop(self):
+            assert get_session().metadata == {"a": 1, "b": 1}
+            with tempfile.TemporaryDirectory() as path:
+                checkpoint = Checkpoint.from_directory(path)
+                checkpoint.set_metadata({"b": 2, "c": 3})
+                train.report(dict(my_metric=1), checkpoint=checkpoint)
+
+    trainer = MyTrainer(metadata={"a": 1, "b": 1})
+    result = trainer.fit()
+    meta_out = result.checkpoint.get_metadata()
+    assert meta_out == {"a": 1, "b": 2, "c": 3}, meta_out
+
+
+def test_data_context_propagation(ray_start_4_cpus):
+    ctx = DataContext.get_current()
+    # Fake DataContext attribute to propagate to worker.
+    ctx.foo = "bar"
+
+    def training_loop(self):
+        # Dummy train loop that checks that changes in the driver's
+        # DataContext are propagated to the worker.
+        ctx_worker = DataContext.get_current()
+        assert ctx_worker.foo == "bar"
+
+    trainer = DummyTrainer(
+        train_loop=training_loop,
+        datasets={"train": ray.data.range(10)},
+    )
+    trainer.fit()
+
+
+def test_large_params(ray_start_4_cpus):
+    """Tests that large params are not serialized with the trainer actor
+    and are instead put into the object store separately."""
+    huge_array = np.zeros(shape=int(1e8))
+
+    def training_loop(self):
+        _ = huge_array
+
+    trainer = DummyTrainer(training_loop)
+    trainer.fit()
+
+
+def test_format_datasets_for_repr(ray_start_4_cpus):
+    datasets = {"train": ray.data.range(1), "test": ray.data.range(1)}
+
+    actual_repr = format_datasets_for_repr(datasets)
+
+    assert actual_repr == (
+        "{'train': Dataset(num_rows=1, schema={id: int64}), "
+        "'test': Dataset(num_rows=1, schema={id: int64})}"
+    )
+
+
+if __name__ == "__main__":
+    import sys
+
+    sys.exit(pytest.main(sys.argv[1:] + ["-v", "-x", __file__]))

python/ray/train/tests/test_base_worker_group.py

@@ -0,0 +1,46 @@
+"""Tests for BaseWorkerGroup implementation and usage."""
+
+import pytest
+
+from ray.train._internal.base_worker_group import BaseWorkerGroup
+from ray.train._internal.worker_group import WorkerGroup as V1WorkerGroup
+from ray.train.v2._internal.execution.worker_group.worker_group import (
+    WorkerGroup as V2WorkerGroup,
+)
+
+
+def test_interface_abstract_methods():
+    """Test that BaseWorkerGroup enforces its abstract methods."""
+    # Should not be able to instantiate interface directly
+    with pytest.raises(TypeError):
+        BaseWorkerGroup()
+
+    # Should not be able to create incomplete implementation
+    class IncompleteWorkerGroup(BaseWorkerGroup):
+        def execute(self, func, *args, **kwargs):
+            pass
+
+        # Missing other abstract methods
+
+    with pytest.raises(TypeError):
+        IncompleteWorkerGroup()
+
+
+def test_real_implementations_inherit_interface():
+    """Smoke test that real WorkerGroup implementations inherit from interface."""
+    # Test inheritance
+    assert issubclass(V1WorkerGroup, BaseWorkerGroup)
+    assert issubclass(V2WorkerGroup, BaseWorkerGroup)
+
+    # Test that all abstract methods are implemented
+    # If any abstract methods are missing, __abstractmethods__ will be non-empty
+    assert (
+        len(V1WorkerGroup.__abstractmethods__) == 0
+    ), f"V1 WorkerGroup missing abstract methods: {V1WorkerGroup.__abstractmethods__}"
+    assert (
+        len(V2WorkerGroup.__abstractmethods__) == 0
+    ), f"V2 WorkerGroup missing abstract methods: {V2WorkerGroup.__abstractmethods__}"
+
+
+if __name__ == "__main__":
+    pytest.main([__file__])

python/ray/train/tests/test_checkpoint.py

@@ -0,0 +1,232 @@
+import os
+from pathlib import Path
+
+import pyarrow.fs
+import pytest
+
+import ray
+from ray.train._checkpoint import (
+    _CHECKPOINT_TEMP_DIR_PREFIX,
+    _METADATA_FILE_NAME,
+    Checkpoint,
+    _get_del_lock_path,
+    _list_existing_del_locks,
+)
+from ray.train._internal.storage import _exists_at_fs_path, _upload_to_fs_path
+from ray.train.tests.test_new_persistence import _create_mock_custom_fs
+
+_CHECKPOINT_CONTENT_FILE = "dummy.txt"
+
+
+@pytest.fixture(params=["local", "mock", "custom_fs"])
+def checkpoint(request, tmp_path):
+    """Fixture that sets up a checkpoint on different filesystems."""
+    checkpoint_fs_type = request.param
+
+    checkpoint_path = tmp_path / "ckpt_dir"
+    checkpoint_path.mkdir(exist_ok=True)
+    (checkpoint_path / _CHECKPOINT_CONTENT_FILE).write_text("dummy")
+
+    if checkpoint_fs_type == "local":
+        yield Checkpoint.from_directory(str(checkpoint_path))
+    elif checkpoint_fs_type == "mock":
+        _checkpoint = Checkpoint(path="mock:///mock_bucket/ckpt_dir")
+        _upload_to_fs_path(
+            local_path=str(checkpoint_path),
+            fs=_checkpoint.filesystem,
+            fs_path=_checkpoint.path,
+        )
+        # The "mock://" URI doesn't persist across different instances of
+        # the pyarrow.fs.MockFileSystem, so we must make sure to return
+        # the checkpoint with the same filesystem instance that we uploaded
+        # some mock content to.
+        yield _checkpoint
+    elif checkpoint_fs_type == "custom_fs":
+        custom_storage_fs = _create_mock_custom_fs(tmp_path / "custom_fs")
+        _upload_to_fs_path(
+            local_path=str(checkpoint_path),
+            fs=custom_storage_fs,
+            fs_path="mock_bucket/ckpt_dir",
+        )
+        yield Checkpoint(path="mock_bucket/ckpt_dir", filesystem=custom_storage_fs)
+
+
+def test_to_directory(checkpoint: Checkpoint):
+    checkpoint_path = Path(checkpoint.to_directory())
+
+    assert (checkpoint_path / _CHECKPOINT_CONTENT_FILE).exists()
+    assert _CHECKPOINT_TEMP_DIR_PREFIX in checkpoint_path.name
+
+
+def test_to_directory_with_user_specified_path(checkpoint: Checkpoint, tmp_path):
+    # Test with a string
+    checkpoint_path = Path(checkpoint.to_directory(str(tmp_path / "special_dir")))
+    assert (checkpoint_path / _CHECKPOINT_CONTENT_FILE).exists()
+    assert checkpoint_path.name == "special_dir"
+
+    # Test with a PathLike
+    checkpoint_path = Path(checkpoint.to_directory(tmp_path / "special_dir"))
+    assert (checkpoint_path / _CHECKPOINT_CONTENT_FILE).exists()
+    assert checkpoint_path.name == "special_dir"
+
+
+def test_multiprocess_to_directory(checkpoint: Checkpoint):
+    """Test the case where multiple processes are trying to checkpoint.
+
+    Only one process should download the checkpoint, and the others should
+    wait until it's finished. In the end, a single checkpoint dir should be
+    shared by all processes.
+    """
+    if checkpoint.filesystem.type_name == "mock":
+        pytest.skip("Mock filesystem cannot be pickled for use with Ray.")
+
+    @ray.remote
+    def download_checkpoint(checkpoint: Checkpoint) -> str:
+        return checkpoint.to_directory()
+
+    paths = [ray.get(download_checkpoint.remote(checkpoint)) for _ in range(5)]
+    # Check that all the paths are the same (no duplicates).
+    assert len(set(paths)) == 1
+
+
+def test_as_directory(checkpoint: Checkpoint):
+    with checkpoint.as_directory() as checkpoint_path:
+        checkpoint_path = Path(checkpoint_path)
+        assert (checkpoint_path / _CHECKPOINT_CONTENT_FILE).exists()
+
+        if isinstance(checkpoint.filesystem, pyarrow.fs.LocalFileSystem):
+            # We should have directly returned the local path
+            assert str(checkpoint_path) == checkpoint.path
+        else:
+            # We should have downloaded to a temp dir.
+            assert _CHECKPOINT_TEMP_DIR_PREFIX in checkpoint_path.name
+
+    if isinstance(checkpoint.filesystem, pyarrow.fs.LocalFileSystem):
+        # Checkpoint should not be deleted, if we directly gave the local path.
+        assert (checkpoint_path / _CHECKPOINT_CONTENT_FILE).exists()
+    else:
+        # Should have been deleted, if the checkpoint downloaded to a temp dir.
+        assert not checkpoint_path.exists()
+
+
+def test_multiprocess_as_directory(checkpoint: Checkpoint, monkeypatch):
+    """Tests that deletion lock files are created and deleted correctly,
+    when multiple processes are accessing a checkpoint with `as_directory`"""
+    # If this is pointing to a local path, no lockfiles will be created.
+    is_local_checkpoint = isinstance(checkpoint.filesystem, pyarrow.fs.LocalFileSystem)
+
+    monkeypatch.setattr(os, "getpid", lambda: 11111)
+
+    with checkpoint.as_directory() as checkpoint_dir_1:
+        lock_file_1 = Path(_get_del_lock_path(checkpoint_dir_1))
+
+        # Pretend that the 2nd `as_directory` is called from a different process.
+        monkeypatch.setattr(os, "getpid", lambda: 22222)
+
+        with checkpoint.as_directory() as checkpoint_dir_2:
+            lock_file_2 = Path(_get_del_lock_path(checkpoint_dir_2))
+
+            # Should point both processes to the same canonical temp directory.
+            assert checkpoint_dir_1 == checkpoint_dir_2
+
+            if is_local_checkpoint:
+                # Check that 2 different lock files were created.
+                assert not lock_file_1.exists()
+                assert not lock_file_2.exists()
+            else:
+                # Check that 2 different lock files were created.
+                assert lock_file_1.exists()
+                assert lock_file_2.exists()
+
+                assert len(_list_existing_del_locks(checkpoint_dir_1)) == 2
+
+        if not is_local_checkpoint:
+            # Check that the 2nd lock file was deleted.
+            assert len(_list_existing_del_locks(checkpoint_dir_1)) == 1
+            assert not lock_file_2.exists()
+
+    if not is_local_checkpoint:
+        # Check that the 1st lock file was deleted.
+        assert len(_list_existing_del_locks(checkpoint_dir_1)) == 0
+        assert not lock_file_1.exists()
+
+        # Check that the temp checkpoint directory was deleted.
+        assert not Path(checkpoint_dir_1).exists()
+
+
+def test_as_directory_lock_cleanup(checkpoint: Checkpoint):
+    """Errors when accessing a checkpoint with `as_directory`
+    shouldn't leave behind lock files.
+    """
+    with pytest.raises(RuntimeError):
+        with checkpoint.as_directory() as checkpoint_dir:
+            raise RuntimeError
+
+    assert not _list_existing_del_locks(checkpoint_dir)
+
+    is_local_checkpoint = isinstance(checkpoint.filesystem, pyarrow.fs.LocalFileSystem)
+    if not is_local_checkpoint:
+        assert not Path(checkpoint_dir).exists()
+
+
+def test_as_directory_download_error(checkpoint: Checkpoint, monkeypatch):
+    """Errors during a checkpoint download should be raised directly when accessing
+    it with the `as_directory` context manager."""
+    if isinstance(checkpoint.filesystem, pyarrow.fs.LocalFileSystem):
+        pytest.skip(
+            "Local filesystem checkpoints don't download to a temp dir, so "
+            "there's no error handling to test."
+        )
+
+    error_text = "original error"
+
+    def to_directory_error(*args, **kwargs):
+        raise RuntimeError(error_text)
+
+    monkeypatch.setattr(checkpoint, "to_directory", to_directory_error)
+
+    with pytest.raises(RuntimeError, match=error_text):
+        with checkpoint.as_directory() as _:
+            pass
+
+
+def test_metadata(checkpoint: Checkpoint):
+    assert checkpoint.get_metadata() == {}
+
+    # No metadata file by default.
+    assert not _exists_at_fs_path(
+        fs=checkpoint.filesystem,
+        fs_path=str(Path(checkpoint.path) / _METADATA_FILE_NAME),
+    )
+
+    checkpoint.update_metadata({"foo": "bar"})
+    assert checkpoint.get_metadata() == {"foo": "bar"}
+
+    checkpoint.update_metadata({"foo": "baz"})
+    assert checkpoint.get_metadata() == {"foo": "baz"}
+
+    checkpoint.update_metadata({"x": 1})
+    assert checkpoint.get_metadata() == {"foo": "baz", "x": 1}
+
+    # Set metadata completely resets the metadata.
+    checkpoint.set_metadata({"y": [1, 2, 3]})
+    assert checkpoint.get_metadata() == {"y": [1, 2, 3]}
+
+    # There should be a metadata file in the checkpoint directory.
+    assert _exists_at_fs_path(
+        fs=checkpoint.filesystem,
+        fs_path=str(Path(checkpoint.path) / _METADATA_FILE_NAME),
+    )
+
+    # Non JSON serializable metadata should raise an error.
+    class Test:
+        pass
+
+    with pytest.raises(TypeError):
+        checkpoint.set_metadata({"non_json_serializable": Test()})
+
+
+if __name__ == "__main__":
+    import sys
+
+    sys.exit(pytest.main(["-v", __file__]))

python/ray/train/tests/test_checkpoint_manager.py

@@ -0,0 +1,239 @@
+import random
+from pathlib import Path
+from typing import List
+
+import pytest
+
+from ray.train import Checkpoint, CheckpointConfig
+from ray.train._internal.checkpoint_manager import _CheckpointManager, _TrainingResult
+from ray.train.constants import TUNE_ONLY_STORE_CHECKPOINT_SCORE_ATTRIBUTE
+
+
+@pytest.fixture
+def checkpoint_paths(tmp_path):
+    checkpoint_paths = []
+    for i in range(10):
+        checkpoint_path = tmp_path / f"ckpt_{i}"
+        checkpoint_path.mkdir()
+        (checkpoint_path / "dummy.txt").write_text(f"{i}")
+        checkpoint_paths.append(checkpoint_path)
+
+    yield [str(path) for path in checkpoint_paths]
+
+
+def test_unlimited_checkpoints(checkpoint_paths: List[str]):
+    manager = _CheckpointManager(checkpoint_config=CheckpointConfig(num_to_keep=None))
+
+    for i in range(10):
+        manager.register_checkpoint(
+            _TrainingResult(
+                checkpoint=Checkpoint.from_directory(checkpoint_paths[i]),
+                metrics={"iter": i},
+            )
+        )
+
+    assert len(manager.best_checkpoint_results) == 10
+
+
+def test_limited_checkpoints(checkpoint_paths: List[str]):
+    manager = _CheckpointManager(checkpoint_config=CheckpointConfig(num_to_keep=2))
+
+    for i in range(10):
+        manager.register_checkpoint(
+            _TrainingResult(
+                checkpoint=Checkpoint.from_directory(checkpoint_paths[i]),
+                metrics={"iter": i},
+            )
+        )
+
+    assert len(manager.best_checkpoint_results) == 2
+
+    # Keep the latest checkpoints if no metric is given.
+    assert {
+        tracked_checkpoint.metrics["iter"]
+        for tracked_checkpoint in manager.best_checkpoint_results
+    } == {8, 9}
+
+    # The first 8 checkpoints should be deleted.
+    for i in range(8):
+        assert not Path(checkpoint_paths[i]).exists()
+
+    assert Path(checkpoint_paths[8]).exists()
+    assert Path(checkpoint_paths[9]).exists()
+
+
+@pytest.mark.parametrize("order", ["min", "max"])
+def test_keep_checkpoints_by_score(order, checkpoint_paths):
+    num_to_keep = 2
+    score_attribute = "score"
+
+    manager = _CheckpointManager(
+        checkpoint_config=CheckpointConfig(
+            num_to_keep=num_to_keep,
+            checkpoint_score_attribute=score_attribute,
+            checkpoint_score_order=order,
+        )
+    )
+
+    scores = []
+    for i in range(10):
+        score = random.random()
+        manager.register_checkpoint(
+            _TrainingResult(
+                checkpoint=Checkpoint.from_directory(checkpoint_paths[i]),
+                metrics={"iter": i, score_attribute: score},
+            )
+        )
+        scores.append(score)
+
+    sorted_scores = sorted(scores, reverse=order == "max")
+    assert set(sorted_scores[:num_to_keep]) == {
+        tracked_checkpoint.metrics[score_attribute]
+        for tracked_checkpoint in manager.best_checkpoint_results
+    }
+
+    # Make sure the bottom checkpoints are deleted.
+    best_checkpoint_iters = {
+        tracked_checkpoint.metrics["iter"]
+        for tracked_checkpoint in manager.best_checkpoint_results
+    }
+    for i, checkpoint_path in enumerate(checkpoint_paths):
+        if i in best_checkpoint_iters or i == 9:
+            # The checkpoint should only exist if it's one of the top K or the latest.
+            assert Path(checkpoint_path).exists()
+        else:
+            assert not Path(checkpoint_path).exists()
+
+
+def test_keep_latest_checkpoint(checkpoint_paths):
+    manager = _CheckpointManager(
+        checkpoint_config=CheckpointConfig(
+            num_to_keep=2,
+            checkpoint_score_attribute="score",
+            checkpoint_score_order="max",
+        )
+    )
+
+    manager.register_checkpoint(
+        _TrainingResult(
+            checkpoint=Checkpoint.from_directory(checkpoint_paths[0]),
+            metrics={"score": 3.0},
+        )
+    )
+    manager.register_checkpoint(
+        _TrainingResult(
+            checkpoint=Checkpoint.from_directory(checkpoint_paths[1]),
+            metrics={"score": 2.0},
+        )
+    )
+    manager.register_checkpoint(
+        _TrainingResult(
+            checkpoint=Checkpoint.from_directory(checkpoint_paths[2]),
+            metrics={"score": 1.0},
+        )
+    )
+
+    assert len(manager.best_checkpoint_results) == 2
+
+    # The latest checkpoint with the lowest score should not be deleted yet.
+    assert manager.latest_checkpoint_result.metrics["score"] == 1.0
+
+    # The latest checkpoint with the lowest score should not be deleted yet.
+    assert Path(checkpoint_paths[2]).exists()
+
+    manager.register_checkpoint(
+        _TrainingResult(
+            checkpoint=Checkpoint.from_directory(checkpoint_paths[3]),
+            metrics={"score": 0.0},
+        )
+    )
+    # A newer checkpoint came in. Even though the new one has a lower score, there are
+    # already num_to_keep better checkpoints, so the previous one should be deleted.
+    assert not Path(checkpoint_paths[2]).exists()
+
+    # Quick sanity check to make sure that the new checkpoint is kept.
+    assert manager.latest_checkpoint_result.metrics["score"] == 0.0
+    assert Path(checkpoint_paths[3]).exists()
+
+    # The original 2 checkpoints should still exist
+    assert Path(checkpoint_paths[0]).exists()
+    assert Path(checkpoint_paths[1]).exists()
+
+
+@pytest.mark.parametrize(
+    "metrics",
+    [
+        {"nested": {"sub": {"attr": 5}}},
+        {"nested": {"sub/attr": 5}},
+        {"nested/sub": {"attr": 5}},
+        {"nested/sub/attr": 5},
+    ],
+)
+def test_nested_get_checkpoint_score(metrics):
+    manager = _CheckpointManager(
+        checkpoint_config=CheckpointConfig(
+            num_to_keep=2,
+            checkpoint_score_attribute="nested/sub/attr",
+            checkpoint_score_order="max",
+        )
+    )
+
+    tracked_checkpoint = _TrainingResult(checkpoint=None, metrics=metrics)
+    assert manager._get_checkpoint_score(tracked_checkpoint) == (True, 5.0)
+
+
+@pytest.mark.parametrize("has_score_attr", [True, False])
+def test_only_store_score_attr(has_score_attr, checkpoint_paths, monkeypatch):
+    monkeypatch.setenv(TUNE_ONLY_STORE_CHECKPOINT_SCORE_ATTRIBUTE, "1")
+
+    # Set up CheckpointManager.
+    if has_score_attr:
+        checkpoint_config = CheckpointConfig(
+            num_to_keep=None,
+            checkpoint_score_attribute="score",
+            checkpoint_score_order="max",
+        )
+    else:
+        checkpoint_config = CheckpointConfig(num_to_keep=None)
+    manager = _CheckpointManager(checkpoint_config=checkpoint_config)
+
+    # Ensure we insert TrainingResults with score in the right order.
+    manager.register_checkpoint(
+        _TrainingResult(
+            checkpoint=Checkpoint.from_directory(checkpoint_paths[0]),
+            metrics={"score": 3.0},
+        )
+    )
+    manager.register_checkpoint(
+        _TrainingResult(
+            checkpoint=Checkpoint.from_directory(checkpoint_paths[1]),
+            metrics={"score": 1.0, "another_unsaved_metric": 6.0},
+        )
+    )
+    manager.register_checkpoint(
+        _TrainingResult(
+            checkpoint=Checkpoint.from_directory(checkpoint_paths[2]),
+            metrics={"another_unsaved_metric": 1.0},
+        )
+    )
+    assert len(manager.best_checkpoint_results) == 3
+    if has_score_attr:
+        assert manager.best_checkpoint_results[0].metrics == {"score": 1.0}
+        assert manager.best_checkpoint_results[0].checkpoint.path == checkpoint_paths[1]
+        assert manager.best_checkpoint_results[1].metrics == {"score": 3.0}
+        assert manager.best_checkpoint_results[1].checkpoint.path == checkpoint_paths[0]
+        assert manager.best_checkpoint_results[2].metrics == {}
+        assert manager.best_checkpoint_results[2].checkpoint.path == checkpoint_paths[2]
+    else:
+        assert manager.best_checkpoint_results[0].metrics == {}
+        assert manager.best_checkpoint_results[0].checkpoint.path == checkpoint_paths[0]
+        assert manager.best_checkpoint_results[1].metrics == {}
+        assert manager.best_checkpoint_results[1].checkpoint.path == checkpoint_paths[1]
+        assert manager.best_checkpoint_results[2].metrics == {}
+        assert manager.best_checkpoint_results[2].checkpoint.path == checkpoint_paths[2]
+
+
+if __name__ == "__main__":
+    import sys
+
+    sys.exit(pytest.main(["-v", __file__]))

python/ray/train/tests/test_data_parallel_trainer.py

@@ -0,0 +1,372 @@
+import os
+import time
+from unittest.mock import patch
+
+import pytest
+
+import ray
+from ray import train, tune
+from ray._common.utils import RESOURCE_CONSTRAINT_PREFIX
+from ray.train import RunConfig, ScalingConfig
+from ray.train._internal.backend_executor import BackendExecutor
+from ray.train._internal.worker_group import WorkerGroup
+from ray.train.backend import Backend, BackendConfig
+from ray.train.data_parallel_trainer import DataParallelTrainer
+from ray.train.tests.util import create_dict_checkpoint, load_dict_checkpoint
+from ray.train.utils import _in_ray_train_worker
+from ray.tune.callback import Callback
+from ray.tune.tune_config import TuneConfig
+from ray.tune.tuner import Tuner
+from ray.util.accelerators import NVIDIA_A100
+
+
+@pytest.fixture
+def ray_start_4_cpus():
+    address_info = ray.init(num_cpus=4)
+    yield address_info
+    # The code after the yield will run as teardown code.
+    ray.shutdown()
+
+
+@pytest.fixture
+def ray_start_4_cpus_4_gpus_4_extra():
+    address_info = ray.init(num_cpus=4, num_gpus=4, resources={"extra": 4})
+    yield address_info
+    # The code after the yield will run as teardown code.
+    ray.shutdown()
+
+
+@pytest.fixture
+def ray_start_4_cpus_4_gpus_4_a100():
+    address_info = ray.init(
+        num_cpus=4,
+        num_gpus=4,
+        resources={f"{RESOURCE_CONSTRAINT_PREFIX}{NVIDIA_A100}": 4},
+    )
+    yield address_info
+    ray.shutdown()
+
+
+def gen_execute_single_async_special(special_f):
+    def execute_single_async_special(self, i, f, *args, **kwargs):
+        assert len(self.workers) == 2
+        if i == 0 and hasattr(self, "should_fail") and self.should_fail:
+            kwargs["train_func"] = special_f
+        return (
+            self.workers[i]
+            .actor._RayTrainWorker__execute.options(name=f.__name__)
+            .remote(f, *args, **kwargs)
+        )
+
+    return execute_single_async_special
+
+
+def gen_new_backend_executor(special_f):
+    """Returns a BackendExecutor that runs special_f on worker 0 once."""
+
+    class TestBackendExecutor(BackendExecutor):
+        def __init__(self, *args, **kwargs):
+            super().__init__(*args, **kwargs)
+            self._has_failed = False
+
+        def start_training(self, *args, **kwargs):
+            special_execute = gen_execute_single_async_special(special_f)
+            if not self._has_failed:
+                self.worker_group.should_fail = True
+                self._has_failed = True
+            else:
+                self.worker_group.should_fail = False
+            with patch.object(WorkerGroup, "execute_single_async", special_execute):
+                super().start_training(*args, **kwargs)
+
+    return TestBackendExecutor
+
+
+class CaptureReportCallback(Callback):
+    def __init__(self):
+        self.result_list = []
+
+    def on_trial_result(self, iteration, trials, trial, result, **info):
+        self.result_list.append(result)
+
+
+scale_config = ScalingConfig(num_workers=2)
+
+
+def test_fit_train(ray_start_4_cpus):
+    def train_func():
+        train.report({"loss": 1})
+
+    trainer = DataParallelTrainer(
+        train_loop_per_worker=train_func, scaling_config=scale_config
+    )
+    assert trainer.fit().metrics["loss"] == 1
+
+
+def test_scaling_config(ray_start_4_cpus):
+    def train_func():
+        assert ray.available_resources()["CPU"] == 1
+        train.report({"loss": 1})
+
+    assert ray.available_resources()["CPU"] == 4
+    trainer = DataParallelTrainer(
+        train_loop_per_worker=train_func, scaling_config=ScalingConfig(num_workers=2)
+    )
+    trainer.fit()
+
+
+def test_fit_train_config(ray_start_4_cpus):
+    def train_func(config):
+        train.report({"loss": config["x"]})
+
+    trainer = DataParallelTrainer(
+        train_loop_per_worker=train_func,
+        scaling_config=scale_config,
+        train_loop_config={"x": 100},
+    )
+    assert trainer.fit().metrics["loss"] == 100
+
+
+def test_datasets(ray_start_4_cpus):
+    num_train_data = 10
+    num_val_data = 6
+
+    train_dataset = ray.data.range(num_train_data)
+    val_dataset = ray.data.range(num_val_data)
+
+    def get_dataset():
+        # Train dataset should be sharded.
+        train_dataset = train.get_dataset_shard("train")
+        train_ds_count = len(list(train_dataset.iter_rows()))
+        assert train_ds_count == num_train_data / scale_config.num_workers
+        # All other datasets should not be sharded.
+        val_dataset = train.get_dataset_shard("val")
+        val_ds_count = len(list(val_dataset.iter_rows()))
+        assert val_ds_count == num_val_data / scale_config.num_workers
+
+    trainer = DataParallelTrainer(
+        train_loop_per_worker=get_dataset,
+        scaling_config=scale_config,
+        datasets={"train": train_dataset, "val": val_dataset},
+    )
+    trainer.fit()
+
+
+def test_invalid_train_loop():
+    def train_loop(config, extra_arg):
+        pass
+
+    with pytest.raises(ValueError):
+        DataParallelTrainer(train_loop_per_worker=train_loop)
+
+
+def test_bad_return_in_train_loop(ray_start_4_cpus):
+    """Test to check if returns from train loop are discarded."""
+
+    # Simulates what happens with eg. torch models
+    class FailOnUnpickle:
+        def __reduce__(self):
+            raise RuntimeError("Failing")
+
+    def train_loop(config):
+        train.report({"loss": 1})
+        return FailOnUnpickle()
+
+    trainer = DataParallelTrainer(
+        train_loop_per_worker=train_loop, scaling_config=scale_config
+    )
+
+    # No exception should happen here
+    trainer.fit()
+
+
+def test_tune(ray_start_4_cpus):
+    def train_func(config):
+        train.report({"loss": config["x"]})
+
+    trainer = DataParallelTrainer(
+        train_loop_per_worker=train_func,
+        train_loop_config={"x": 100},
+        scaling_config=scale_config,
+    )
+
+    tuner = Tuner(
+        trainer,
+        param_space={"train_loop_config": {"x": tune.choice([200, 300])}},
+        tune_config=TuneConfig(num_samples=2),
+    )
+    result_grid = tuner.fit()
+    assert result_grid[0].metrics["loss"] in [200, 300]
+
+    # Make sure original Trainer is not affected.
+    assert trainer._train_loop_config["x"] == 100
+
+
+def test_fast_slow(ray_start_4_cpus):
+    def train_func():
+        for i in range(2):
+            with create_dict_checkpoint({"epoch": i}) as checkpoint:
+                train.report(dict(index=i), checkpoint=checkpoint)
+
+    def train_slow():
+        for i in range(2):
+            with create_dict_checkpoint({"epoch": i}) as checkpoint:
+                train.report(dict(index=i), checkpoint=checkpoint)
+            time.sleep(5)
+
+    new_backend_executor_cls = gen_new_backend_executor(train_slow)
+    callback = CaptureReportCallback()
+
+    class DataParallelTrainerPatched(DataParallelTrainer):
+        _backend_executor_cls = new_backend_executor_cls
+
+    trainer = DataParallelTrainerPatched(
+        train_func,
+        scaling_config=scale_config,
+        run_config=RunConfig(callbacks=[callback]),
+    )
+    results = trainer.fit()
+
+    assert load_dict_checkpoint(results.checkpoint)["epoch"] == 1
+
+    result_list = callback.result_list
+    assert len(result_list) == 2
+
+
+def test_mismatch_report(ray_start_4_cpus):
+    def train_func():
+        for _ in range(2):
+            train.report(dict(loss=1))
+
+    def train_mismatch():
+        train.report(dict(loss=1))
+
+    new_backend_executor_cls = gen_new_backend_executor(train_mismatch)
+
+    class DataParallelTrainerPatched(DataParallelTrainer):
+        _backend_executor_cls = new_backend_executor_cls
+
+    trainer = DataParallelTrainerPatched(
+        train_func,
+        scaling_config=scale_config,
+    )
+    with pytest.raises(RuntimeError):
+        trainer.fit()
+
+
+def test_world_rank(ray_start_4_cpus, tmp_path):
+    def train_func():
+        world_rank = train.get_context().get_world_rank()
+        (tmp_path / f"{world_rank}").touch()
+        train.report(dict(world_rank=world_rank))
+
+    trainer = DataParallelTrainer(train_func, scaling_config=scale_config)
+    trainer.fit()
+
+    created_files = list(tmp_path.glob("*"))
+    assert len(created_files) == 2
+    assert {int(file.name) for file in created_files} == {0, 1}
+
+
+def test_gpu_requests(ray_start_4_cpus_4_gpus_4_extra, tmp_path):
+    def get_visible_devices_for_workers():
+        return [file.read_text() for file in tmp_path.glob("*")]
+
+    class CudaTestBackend(Backend):
+        share_cuda_visible_devices = True
+
+    class CudaTestConfig(BackendConfig):
+        @property
+        def backend_cls(self):
+            return CudaTestBackend
+
+    def get_resources():
+        cuda_visible_devices = os.environ.get("CUDA_VISIBLE_DEVICES", "")
+        world_rank = train.get_context().get_world_rank()
+        (tmp_path / f"{world_rank}").write_text(cuda_visible_devices)
+        train.report(dict(devices=cuda_visible_devices))
+
+    # 0 GPUs will be requested and should not raise an error.
+    trainer = DataParallelTrainer(
+        get_resources,
+        backend_config=CudaTestConfig(),
+        scaling_config=ScalingConfig(num_workers=2, use_gpu=False),
+    )
+    trainer.fit()
+    assert get_visible_devices_for_workers() == ["", ""]
+
+    # 1 GPU will be requested and should not raise an error.
+    trainer = DataParallelTrainer(
+        get_resources,
+        backend_config=CudaTestConfig(),
+        scaling_config=ScalingConfig(num_workers=2, use_gpu=True),
+    )
+    trainer.fit()
+    visible_devices = get_visible_devices_for_workers()
+    # Sort the cuda visible devices to have exact match with expected result.
+    visible_devices = [",".join(sorted(r.split(","))) for r in visible_devices]
+    assert visible_devices == ["0,1", "0,1"]
+
+    # Partial GPUs should not raise an error.
+    trainer = DataParallelTrainer(
+        get_resources,
+        backend_config=CudaTestConfig(),
+        scaling_config=ScalingConfig(
+            num_workers=2, use_gpu=True, resources_per_worker={"GPU": 0.1}
+        ),
+    )
+    trainer.fit()
+    visible_devices = get_visible_devices_for_workers()
+    assert visible_devices == ["0", "0"]
+
+    # Multiple GPUs should not raise an error.
+    trainer = DataParallelTrainer(
+        get_resources,
+        backend_config=CudaTestConfig(),
+        scaling_config=ScalingConfig(
+            num_workers=2, use_gpu=True, resources_per_worker={"GPU": 2}
+        ),
+    )
+    trainer.fit()
+    visible_devices = get_visible_devices_for_workers()
+    # Sort the cuda visible devices to have exact match with expected result.
+    visible_devices = [",".join(sorted(r.split(","))) for r in visible_devices]
+    assert visible_devices == ["0,1,2,3", "0,1,2,3"]
+
+
+@pytest.mark.parametrize("accelerator_type", [NVIDIA_A100, None])
+def test_config_accelerator_type(ray_start_4_cpus_4_gpus_4_a100, accelerator_type):
+    def train_func():
+        # Ensure all workers are scheduled on nodes with specified accelerators
+        assigned_resources = ray.get_runtime_context().get_assigned_resources()
+        if accelerator_type:
+            accelerator_key = f"{RESOURCE_CONSTRAINT_PREFIX}{accelerator_type}"
+            assert accelerator_key in assigned_resources
+
+    trainer = DataParallelTrainer(
+        train_func,
+        scaling_config=ScalingConfig(
+            num_workers=4,
+            use_gpu=True,
+            accelerator_type=accelerator_type,
+        ),
+    )
+    trainer.fit()
+
+
+def test_in_ray_train_worker(ray_start_4_cpus):
+    assert not _in_ray_train_worker()
+
+    def train_fn():
+        assert _in_ray_train_worker()
+
+    trainer = DataParallelTrainer(train_fn)
+    trainer.fit()
+
+
+if __name__ == "__main__":
+    import sys
+
+    import pytest
+
+    sys.exit(pytest.main(["-v", "-x", __file__]))

python/ray/train/tests/test_data_parallel_trainer_checkpointing.py

@@ -0,0 +1,122 @@
+import pytest
+
+import ray
+from ray import train
+from ray.train import CheckpointConfig, RunConfig, ScalingConfig
+from ray.train.data_parallel_trainer import DataParallelTrainer
+from ray.train.tests.util import create_dict_checkpoint, load_dict_checkpoint
+
+
+@pytest.fixture
+def ray_start_4_cpus():
+    address_info = ray.init(num_cpus=4)
+    yield address_info
+    # The code after the yield will run as teardown code.
+    ray.shutdown()
+
+
+scale_config = ScalingConfig(num_workers=2)
+NUM_EPOCHS = 3
+
+
+def checkpoint_train_func():
+    for i in range(NUM_EPOCHS):
+        with create_dict_checkpoint({"epoch": i}) as checkpoint:
+            train.report({"epoch": i}, checkpoint=checkpoint)
+
+
+def test_checkpoint(ray_start_4_cpus):
+    """Test that a checkpoint is created and accessible."""
+    trainer = DataParallelTrainer(
+        train_loop_per_worker=checkpoint_train_func,
+        scaling_config=scale_config,
+    )
+    result = trainer.fit()
+    assert load_dict_checkpoint(result.checkpoint)["epoch"] == NUM_EPOCHS - 1
+
+
+def test_resume_from_checkpoint(ray_start_4_cpus, tmpdir):
+    """Test that training can be resumed from a reported checkpoint."""
+
+    def train_func():
+        checkpoint = train.get_checkpoint()
+        if checkpoint:
+            epoch = load_dict_checkpoint(checkpoint)["epoch"]
+        else:
+            epoch = 0
+        for i in range(epoch, epoch + 2):
+            with create_dict_checkpoint({"epoch": i}) as checkpoint:
+                train.report({"epoch": i}, checkpoint=checkpoint)
+
+    trainer = DataParallelTrainer(
+        train_loop_per_worker=train_func, scaling_config=scale_config
+    )
+    result = trainer.fit()
+    assert load_dict_checkpoint(result.checkpoint)["epoch"] == 1
+
+    trainer = DataParallelTrainer(
+        train_loop_per_worker=train_func,
+        scaling_config=scale_config,
+        resume_from_checkpoint=result.checkpoint,
+    )
+    result = trainer.fit()
+    assert load_dict_checkpoint(result.checkpoint)["epoch"] == 2
+
+
+@pytest.mark.parametrize("mode", ["min", "max"])
+def test_checkpoints_to_keep(ray_start_4_cpus, mode):
+    """
+    Test that ``CheckpointConfig`` is respected.
+
+    - Report 4 times with different metrics.
+    - Assert that the kept checkpoints match the expectation.
+    """
+
+    def train_func():
+        with create_dict_checkpoint({"idx": 0}) as checkpoint:
+            train.report(dict(loss=float("nan")), checkpoint=checkpoint)  # nan, deleted
+        with create_dict_checkpoint({"idx": 1}) as checkpoint:
+            train.report(
+                dict(loss=3), checkpoint=checkpoint
+            )  # best for min, worst for max (del)
+        with create_dict_checkpoint({"idx": 2}) as checkpoint:
+            train.report(
+                dict(loss=7), checkpoint=checkpoint
+            )  # worst for min (del), best for max
+        with create_dict_checkpoint({"idx": 3}) as checkpoint:
+            train.report(dict(loss=5), checkpoint=checkpoint)
+
+    checkpoint_config = CheckpointConfig(
+        num_to_keep=2, checkpoint_score_attribute="loss", checkpoint_score_order=mode
+    )
+
+    trainer = DataParallelTrainer(
+        train_func,
+        scaling_config=scale_config,
+        run_config=RunConfig(checkpoint_config=checkpoint_config),
+    )
+    result = trainer.fit()
+    assert len(result.best_checkpoints) == 2
+
+    # Last checkpoint
+    assert load_dict_checkpoint(result.checkpoint)["idx"] == 3
+
+    if mode == "min":
+        indices = [3, 1]
+        losses = [5, 3]
+    else:
+        indices = [3, 2]
+        losses = [5, 7]
+
+    assert load_dict_checkpoint(result.best_checkpoints[0][0])["idx"] == indices[0]
+    assert load_dict_checkpoint(result.best_checkpoints[1][0])["idx"] == indices[1]
+    assert result.best_checkpoints[0][1]["loss"] == losses[0]
+    assert result.best_checkpoints[1][1]["loss"] == losses[1]
+
+
+if __name__ == "__main__":
+    import sys
+
+    import pytest
+
+    sys.exit(pytest.main(["-v", "-x", __file__]))

python/ray/train/tests/test_e2e_wandb_integration.py

@@ -0,0 +1,73 @@
+"""
+If a user uses Trainer API directly with wandb integration, they expect to see
+* train_loop_config to show up in wandb.config.
+
+This test uses mocked call into wandb API.
+"""
+
+import pytest
+
+import ray
+from ray.air.integrations.wandb import WANDB_ENV_VAR
+from ray.air.tests.mocked_wandb_integration import WandbTestExperimentLogger
+from ray.train import RunConfig, ScalingConfig
+from ray.train.examples.pytorch.torch_linear_example import (
+    train_func as linear_train_func,
+)
+from ray.train.torch import TorchTrainer
+
+
+@pytest.fixture
+def ray_start_4_cpus():
+    address_info = ray.init(num_cpus=4)
+    yield address_info
+    # The code after the yield will run as teardown code.
+    ray.shutdown()
+
+
+CONFIG = {"lr": 1e-2, "hidden_size": 1, "batch_size": 4, "epochs": 3}
+
+
+@pytest.mark.parametrize("with_train_loop_config", (True, False))
+def test_trainer_wandb_integration(
+    ray_start_4_cpus, with_train_loop_config, monkeypatch
+):
+    monkeypatch.setenv(WANDB_ENV_VAR, "9012")
+
+    def train_func(config=None):
+        config = config or CONFIG
+        result = linear_train_func(config)
+        assert len(result) == config["epochs"]
+        assert result[-1]["loss"] < result[0]["loss"]
+
+    scaling_config = ScalingConfig(num_workers=2)
+
+    logger = WandbTestExperimentLogger(project="test_project")
+    if with_train_loop_config:
+        trainer = TorchTrainer(
+            train_loop_per_worker=train_func,
+            train_loop_config=CONFIG,
+            scaling_config=scaling_config,
+            run_config=RunConfig(callbacks=[logger]),
+        )
+    else:
+        trainer = TorchTrainer(
+            train_loop_per_worker=train_func,
+            scaling_config=scaling_config,
+            run_config=RunConfig(callbacks=[logger]),
+        )
+    trainer.fit()
+    config = list(logger.trial_logging_actor_states.values())[0].config
+
+    if with_train_loop_config:
+        assert "train_loop_config" in config
+    else:
+        assert "train_loop_config" not in config
+
+
+if __name__ == "__main__":
+    import sys
+
+    import pytest
+
+    sys.exit(pytest.main(["-v", "-x", __file__]))

python/ray/train/tests/test_examples.py

@@ -0,0 +1,98 @@
+import sys
+
+import pytest
+
+from ray.train import ScalingConfig
+from ray.train.examples.pytorch.torch_fashion_mnist_example import (
+    train_func_per_worker as fashion_mnist_train_func,
+)
+from ray.train.examples.pytorch.torch_linear_example import (
+    train_func as linear_train_func,
+)
+from ray.train.examples.pytorch.torch_quick_start import (
+    train_func as torch_quick_start_train_func,
+)
+from ray.train.examples.tf.tensorflow_quick_start import (
+    train_func as tf_quick_start_train_func,
+)
+from ray.train.torch import TorchTrainer
+
+
+@pytest.mark.parametrize("num_workers", [1, 2])
+@pytest.mark.skipif(
+    sys.version_info >= (3, 12), reason="tensorflow is not supported in python 3.12+"
+)
+def test_tensorflow_mnist(ray_start_4_cpus, num_workers):
+    from ray.train.examples.tf.tensorflow_mnist_example import (
+        train_func as tensorflow_mnist_train_func,
+    )
+    from ray.train.tensorflow import TensorflowTrainer
+
+    num_workers = num_workers
+    epochs = 3
+
+    config = {"lr": 1e-3, "batch_size": 64, "epochs": epochs}
+    trainer = TensorflowTrainer(
+        tensorflow_mnist_train_func,
+        train_loop_config=config,
+        scaling_config=ScalingConfig(num_workers=num_workers),
+    )
+    trainer.fit()
+
+
+@pytest.mark.skipif(
+    sys.version_info >= (3, 12), reason="tensorflow is not supported in python 3.12+"
+)
+def test_tf_non_distributed(ray_start_4_cpus):
+    """Make sure Ray Train works without TF MultiWorkerMirroredStrategy."""
+
+    from ray.train.tensorflow import TensorflowTrainer
+
+    trainer = TensorflowTrainer(
+        tf_quick_start_train_func, scaling_config=ScalingConfig(num_workers=1)
+    )
+    trainer.fit()
+
+
+@pytest.mark.parametrize("num_workers", [1, 2])
+def test_torch_linear(ray_start_4_cpus, num_workers):
+    num_workers = num_workers
+    epochs = 3
+
+    config = {"lr": 1e-2, "hidden_size": 1, "batch_size": 4, "epochs": epochs}
+    trainer = TorchTrainer(
+        linear_train_func,
+        train_loop_config=config,
+        scaling_config=ScalingConfig(num_workers=num_workers),
+    )
+    trainer.fit()
+
+
+def test_torch_fashion_mnist(ray_start_4_cpus):
+    num_workers = 2
+    epochs = 3
+
+    config = {"lr": 1e-3, "batch_size_per_worker": 32, "epochs": epochs}
+    trainer = TorchTrainer(
+        fashion_mnist_train_func,
+        train_loop_config=config,
+        scaling_config=ScalingConfig(num_workers=num_workers),
+    )
+    trainer.fit()
+
+
+def test_torch_non_distributed(ray_start_4_cpus):
+    """Make sure Ray Train works without torch DDP."""
+
+    trainer = TorchTrainer(
+        torch_quick_start_train_func, scaling_config=ScalingConfig(num_workers=1)
+    )
+    trainer.fit()
+
+
+if __name__ == "__main__":
+    import sys
+
+    import pytest
+
+    sys.exit(pytest.main(["-v", "-x", __file__]))

python/ray/train/tests/test_gpu.py

@@ -0,0 +1,345 @@
+import json
+import os
+import time
+from pathlib import Path
+from typing import Dict, List, Union
+from unittest.mock import patch
+
+import pytest
+import torch
+import torchvision
+from torch.nn.parallel import DistributedDataParallel
+from torch.utils.data import DataLoader, DistributedSampler
+
+import ray
+import ray.data
+from ray import train
+from ray.exceptions import RayTaskError
+from ray.train import ScalingConfig
+from ray.train.examples.pytorch.torch_linear_example import LinearDataset
+from ray.train.torch.config import TorchConfig
+from ray.train.torch.torch_trainer import TorchTrainer
+from ray.train.trainer import TrainingFailedError
+
+
+class LinearDatasetDict(LinearDataset):
+    """Modifies the LinearDataset to return a Dict instead of a Tuple."""
+
+    def __getitem__(self, index):
+        return {"x": self.x[index, None], "y": self.y[index, None]}
+
+
+class NonTensorDataset(LinearDataset):
+    """Modifies the LinearDataset to also return non-tensor objects."""
+
+    def __getitem__(self, index):
+        return {"x": self.x[index, None], "y": 2}
+
+
+def write_rank_data(tmp_path: Path, data: Union[int, List, Dict]):
+    rank = train.get_context().get_world_rank()
+    with open(tmp_path / f"{rank}.json", "w") as f:
+        json.dump(data, f)
+
+
+def get_data_from_all_ranks(tmp_path: Path) -> Dict[int, Union[int, List, Dict]]:
+    rank_data = {}
+    for rank_file in tmp_path.glob("*.json"):
+        rank = int(rank_file.stem)
+        with open(rank_file, "r") as f:
+            data = json.load(f)
+        rank_data[rank] = data
+    return rank_data
+
+
+@pytest.mark.parametrize("cuda_visible_devices", ["", "1,2"])
+@pytest.mark.parametrize("num_gpus_per_worker", [0.5, 1, 2])
+def test_torch_get_device(
+    shutdown_only, num_gpus_per_worker, cuda_visible_devices, monkeypatch, tmp_path
+):
+    if cuda_visible_devices:
+        # Test if `get_device` is correct even with user specified env var.
+        monkeypatch.setenv("CUDA_VISIBLE_DEVICES", cuda_visible_devices)
+
+    ray.init(num_cpus=4, num_gpus=2)
+
+    def train_fn():
+        # Confirm that the TorchConfig Prologue is effective
+        assert torch.cuda.current_device() == train.torch.get_device().index
+        # Make sure environment variable is being set correctly.
+        if cuda_visible_devices:
+            visible_devices = os.environ["CUDA_VISIBLE_DEVICES"]
+            assert visible_devices == "1,2"
+
+        devices = sorted([device.index for device in train.torch.get_devices()])
+        write_rank_data(tmp_path, devices)
+
+    trainer = TorchTrainer(
+        train_fn,
+        scaling_config=ScalingConfig(
+            num_workers=int(2 / num_gpus_per_worker),
+            use_gpu=True,
+            resources_per_worker={"GPU": num_gpus_per_worker},
+        ),
+    )
+    trainer.fit()
+
+    rank_data = get_data_from_all_ranks(tmp_path)
+    devices = list(rank_data.values())
+
+    if num_gpus_per_worker == 0.5:
+        assert sorted(devices) == [[0], [0], [1], [1]]
+    elif num_gpus_per_worker == 1:
+        assert sorted(devices) == [[0], [1]]
+    elif num_gpus_per_worker == 2:
+        assert sorted(devices[0]) == [0, 1]
+    else:
+        raise RuntimeError(
+            "New parameter for this test has been added without checking that the "
+            "correct devices have been returned."
+        )
+
+
+@pytest.mark.parametrize("num_gpus_per_worker", [0.5, 1, 2])
+def test_torch_get_device_dist(ray_2_node_2_gpu, num_gpus_per_worker, tmp_path):
+    @patch("torch.cuda.is_available", lambda: True)
+    def train_fn():
+        # Confirm that the TorchConfig Prologue is effective
+        assert torch.cuda.current_device() == train.torch.get_device().index
+        devices = sorted([device.index for device in train.torch.get_devices()])
+        write_rank_data(tmp_path, devices)
+
+    trainer = TorchTrainer(
+        train_fn,
+        # use gloo instead of nccl, since nccl is not supported
+        # on this virtual gpu ray environment
+        torch_config=TorchConfig(backend="gloo"),
+        scaling_config=ScalingConfig(
+            num_workers=int(4 / num_gpus_per_worker),
+            use_gpu=True,
+            resources_per_worker={"GPU": num_gpus_per_worker},
+        ),
+    )
+    trainer.fit()
+
+    rank_data = get_data_from_all_ranks(tmp_path)
+    devices = list(rank_data.values())
+
+    # cluster setups: 2 nodes, 2 gpus per node
+    # `CUDA_VISIBLE_DEVICES` is set to "0,1" on node 1 and node 2
+    if num_gpus_per_worker == 0.5:
+        # worker gpu topology:
+        # 4 workers on node 1, 4 workers on node 2
+        # `ray.get_gpu_ids()` returns [0], [0], [1], [1] on node 1
+        # and [0], [0], [1], [1] on node 2
+        assert sorted(devices) == [[0], [0], [0], [0], [1], [1], [1], [1]]
+    elif num_gpus_per_worker == 1:
+        # worker gpu topology:
+        # 2 workers on node 1, 2 workers on node 2
+        # `ray.get_gpu_ids()` returns [0], [1] on node 1 and [0], [1] on node 2
+        assert sorted(devices) == [[0], [0], [1], [1]]
+    elif num_gpus_per_worker == 2:
+        # worker gpu topology:
+        # 1 workers on node 1, 1 workers on node 2
+        # `ray.get_gpu_ids()` returns {0, 1} on node 1 and {0, 1} on node 2
+        # and `device_id` returns the one index from each set.
+        # So total count of devices should be 2.
+        assert devices == [[0, 1], [0, 1]]
+    else:
+        raise RuntimeError(
+            "New parameter for this test has been added without checking that the "
+            "correct devices have been returned."
+        )
+
+
+def test_torch_prepare_model(ray_start_4_cpus_2_gpus):
+    """Tests if ``prepare_model`` correctly wraps in DDP."""
+
+    def train_fn():
+        model = torch.nn.Linear(1, 1)
+
+        # Wrap in DDP.
+        model = train.torch.prepare_model(model)
+
+        # Make sure model is wrapped in DDP.
+        assert isinstance(model, DistributedDataParallel)
+
+        # Make sure model is on cuda.
+        assert next(model.parameters()).is_cuda
+
+    trainer = TorchTrainer(
+        train_fn, scaling_config=ScalingConfig(num_workers=2, use_gpu=True)
+    )
+    trainer.fit()
+
+    def train_fn_manual_override():
+        model = torch.nn.Linear(1, 1)
+
+        # Wrap in DDP and manually specify CPU.
+        model = train.torch.prepare_model(model, device=torch.device("cpu"))
+
+        # Make sure model is wrapped in DDP.
+        assert isinstance(model, DistributedDataParallel)
+
+        # Make sure model is NOT on cuda since we manually specified CPU.
+        assert not next(model.parameters()).is_cuda
+
+    trainer = TorchTrainer(
+        train_fn, scaling_config=ScalingConfig(num_workers=2, use_gpu=True)
+    )
+    trainer.fit()
+
+
+def test_torch_prepare_model_uses_device(ray_start_4_cpus_2_gpus):
+    """Tests if `prepare_model` uses the train.torch.get_device even if it does not
+    match with the local rank."""
+    # The below test should pass without errors.
+
+    @patch.object(
+        ray.train.torch.train_loop_utils,
+        "get_device",
+        lambda: torch.device(f"cuda:{1 - train.get_context().get_local_rank()}"),
+    )
+    def train_func():
+        # These assert statements must hold for prepare_model to wrap with DDP.
+        assert torch.cuda.is_available()
+        assert train.get_context().get_world_size() > 1
+        model = torch.nn.Linear(1, 1)
+        data = torch.ones(1)
+        data = data.to(train.torch.get_device())
+        model = train.torch.prepare_model(model)
+        model(data)
+
+    trainer = TorchTrainer(
+        train_func, scaling_config=ScalingConfig(num_workers=2, use_gpu=True)
+    )
+    trainer.fit()
+
+
+@pytest.mark.parametrize(
+    "dataset", (LinearDataset, LinearDatasetDict, NonTensorDataset)
+)
+def test_torch_prepare_dataloader(ray_start_4_cpus_2_gpus, dataset):
+    data_loader = DataLoader(dataset(a=1, b=2, size=10))
+
+    def train_fn():
+        wrapped_data_loader = train.torch.prepare_data_loader(data_loader)
+
+        # Check that DistributedSampler has been added to the data loader.
+        assert isinstance(wrapped_data_loader.sampler, DistributedSampler)
+
+        # Make sure you can properly iterate through the DataLoader.
+        # Case where the dataset returns a tuple or list from __getitem__.
+        if isinstance(dataset, LinearDataset):
+            for batch in wrapped_data_loader:
+                x = batch[0]
+                y = batch[1]
+
+                # Make sure the data is on the correct device.
+                assert x.is_cuda and y.is_cuda
+        # Case where the dataset returns a dict from __getitem__.
+        elif isinstance(dataset, LinearDatasetDict):
+            for batch in wrapped_data_loader:
+                for x, y in zip(batch["x"], batch["y"]):
+                    # Make sure the data is on the correct device.
+                    assert x.is_cuda and y.is_cuda
+
+        elif isinstance(dataset, NonTensorDataset):
+            for batch in wrapped_data_loader:
+                for x, y in zip(batch["x"], batch["y"]):
+                    # Make sure the data is on the correct device.
+                    assert x.is_cuda and y == 2
+
+    trainer = TorchTrainer(
+        train_fn, scaling_config=ScalingConfig(num_workers=2, use_gpu=True)
+    )
+    trainer.fit()
+
+
+@pytest.mark.parametrize("data_loader_num_workers", (0, 2))
+def test_enable_reproducibility(ray_start_4_cpus_2_gpus, data_loader_num_workers):
+    # NOTE: Reproducible results aren't guaranteed between seeded executions, even with
+    # identical hardware and software dependencies. This test should be okay given that
+    # it only runs for two epochs on a small dataset.
+    # NOTE: I've chosen to use a ResNet model over a more simple model, because
+    # `enable_reproducibility` disables CUDA convolution benchmarking, and a simpler
+    # model (e.g., linear) might not test this feature.
+    def train_func():
+        train.torch.enable_reproducibility()
+
+        model = torchvision.models.resnet18()
+        model = train.torch.prepare_model(model)
+
+        dataset_length = 128
+        dataset = torch.utils.data.TensorDataset(
+            torch.randn(dataset_length, 3, 32, 32),
+            torch.randint(low=0, high=1000, size=(dataset_length,)),
+        )
+
+        # num_workers > 0 tests for https://github.com/ray-project/ray/issues/30247
+        dataloader = torch.utils.data.DataLoader(
+            dataset, batch_size=64, num_workers=data_loader_num_workers
+        )
+        dataloader = train.torch.prepare_data_loader(dataloader)
+
+        optimizer = torch.optim.SGD(model.parameters(), lr=0.001)
+
+        model.train()
+        for epoch in range(2):
+            for images, targets in dataloader:
+                optimizer.zero_grad()
+
+                outputs = model(images)
+                loss = torch.nn.functional.cross_entropy(outputs, targets)
+
+                loss.backward()
+                optimizer.step()
+
+        train.report(dict(loss=loss.item()))
+
+    trainer = TorchTrainer(
+        train_func, scaling_config=ScalingConfig(num_workers=2, use_gpu=True)
+    )
+    result1 = trainer.fit()
+
+    trainer = TorchTrainer(
+        train_func, scaling_config=ScalingConfig(num_workers=2, use_gpu=True)
+    )
+    result2 = trainer.fit()
+
+    assert result1.metrics["loss"] == result2.metrics["loss"]
+
+
+def test_torch_fail_on_nccl_timeout(ray_start_4_cpus_2_gpus):
+    """Tests that TorchTrainer raises exception on NCCL timeouts."""
+
+    def train_fn():
+        model = torch.nn.Linear(1, 1)
+        model = train.torch.prepare_model(model)
+
+        # Rank 0 worker will never reach the collective operation.
+        # NCCL should timeout.
+        if train.get_context().get_world_rank() == 0:
+            while True:
+                time.sleep(100)
+
+        torch.distributed.barrier()
+
+    trainer = TorchTrainer(
+        train_fn,
+        scaling_config=ScalingConfig(num_workers=2, use_gpu=True),
+        torch_config=TorchConfig(timeout_s=5),
+    )
+
+    # Training should fail and not hang.
+    with pytest.raises(TrainingFailedError) as exc_info:
+        trainer.fit()
+    assert isinstance(exc_info.value.__cause__, RayTaskError)
+
+
+if __name__ == "__main__":
+    import sys
+
+    import pytest
+
+    sys.exit(pytest.main(["-v", "-x", "-s", __file__]))

python/ray/train/tests/test_gpu_2.py

@@ -0,0 +1,71 @@
+import numpy as np
+import pytest
+import torch
+
+import ray
+import ray.data
+import ray.train as train
+from ray import tune
+from ray.air.config import ScalingConfig
+from ray.train.examples.pytorch.torch_linear_example import LinearDataset
+from ray.train.torch.torch_trainer import TorchTrainer
+
+
+class LinearDatasetDict(LinearDataset):
+    """Modifies the LinearDataset to return a Dict instead of a Tuple."""
+
+    def __getitem__(self, index):
+        return {"x": self.x[index, None], "y": self.y[index, None]}
+
+
+class NonTensorDataset(LinearDataset):
+    """Modifies the LinearDataset to also return non-tensor objects."""
+
+    def __getitem__(self, index):
+        return {"x": self.x[index, None], "y": 2}
+
+
+# Currently in DataParallelTrainers we only report metrics from rank 0.
+# For testing purposes here, we need to be able to report from all
+# workers.
+class TorchTrainerPatchedMultipleReturns(TorchTrainer):
+    def _report(self, training_iterator) -> None:
+        for results in training_iterator:
+            tune.report(results=results)
+
+
+@pytest.mark.parametrize("use_gpu", (True, False))
+def test_torch_iter_torch_batches_auto_device(ray_start_4_cpus_2_gpus, use_gpu):
+    """
+    Tests that iter_torch_batches in TorchTrainer worker function uses the
+    default device.
+    """
+
+    def train_fn():
+        dataset = train.get_dataset_shard("train")
+        for batch in dataset.iter_torch_batches(dtypes=torch.float, device="cpu"):
+            assert str(batch["data"].device) == "cpu"
+
+        # Autodetect
+        for batch in dataset.iter_torch_batches(dtypes=torch.float):
+            assert str(batch["data"].device) == str(train.torch.get_device())
+
+    dataset = ray.data.from_numpy(np.array([[1, 2, 3, 4, 5], [1, 2, 3, 4, 5]]).T)
+    # Test that this works outside a Train function
+    for batch in dataset.iter_torch_batches(dtypes=torch.float, device="cpu"):
+        assert str(batch["data"].device) == "cpu"
+
+    trainer = TorchTrainer(
+        train_fn,
+        scaling_config=ScalingConfig(num_workers=2, use_gpu=use_gpu),
+        datasets={"train": dataset},
+    )
+    trainer.fit()
+
+
+if __name__ == "__main__":
+    import sys
+
+    import pytest
+
+    sys.exit(pytest.main(["-v", "-x", "-s", __file__]))

python/ray/train/tests/test_gpu_auto_transfer.py

@@ -0,0 +1,62 @@
+import numpy as np
+import pytest
+import torch
+
+from ray.train.torch.train_loop_utils import _WrappedDataLoader
+
+
+@pytest.mark.parametrize(
+    ("device_choice", "auto_transfer"),
+    [
+        ("cpu", True),
+        ("cpu", False),
+        ("cuda", True),
+        ("cuda", False),
+    ],
+)
+def test_auto_transfer_data_from_host_to_device(
+    ray_start_1_cpu_1_gpu, device_choice, auto_transfer
+):
+    def compute_average_runtime(func):
+        device = torch.device(device_choice)
+        start = torch.cuda.Event(enable_timing=True)
+        end = torch.cuda.Event(enable_timing=True)
+        runtime = []
+        for _ in range(10):
+            torch.cuda.synchronize()
+            start.record()
+            func(device)
+            end.record()
+            torch.cuda.synchronize()
+            runtime.append(start.elapsed_time(end))
+        return np.mean(runtime)
+
+    small_dataloader = [
+        (torch.randn((1024 * 4, 1024 * 4), device="cpu"),) for _ in range(10)
+    ]
+
+    def host_to_device(device):
+        for (x,) in small_dataloader:
+            x = x.to(device)
+            torch.matmul(x, x)
+
+    def host_to_device_auto_pipeline(device):
+        wrapped_dataloader = _WrappedDataLoader(small_dataloader, device, auto_transfer)
+        for (x,) in wrapped_dataloader:
+            torch.matmul(x, x)
+
+    # test if all four configurations are okay
+    with_auto_transfer = compute_average_runtime(host_to_device_auto_pipeline)
+
+    if device_choice == "cuda" and auto_transfer:
+        # check if auto transfer is faster than manual transfer
+        without_auto_transfer = compute_average_runtime(host_to_device)
+        assert with_auto_transfer <= without_auto_transfer
+
+
+if __name__ == "__main__":
+    import sys
+
+    import pytest
+
+    sys.exit(pytest.main(["-v", "-x", "-s", __file__]))

python/ray/train/tests/test_gpu_examples.py

@@ -0,0 +1,56 @@
+import sys
+
+import pytest
+
+from ray.train import ScalingConfig
+from ray.train.examples.pytorch.torch_fashion_mnist_example import (
+    train_func_per_worker as fashion_mnist_train_func,
+)
+from ray.train.torch import TorchTrainer
+
+
+@pytest.mark.skipif(
+    sys.version_info >= (3, 12),
+    reason="Tensorflow is not installed in CI for Python 3.12",
+)
+def test_tensorflow_mnist_gpu(ray_start_4_cpus_2_gpus):
+    from ray.train.examples.tf.tensorflow_mnist_example import (
+        train_func as tensorflow_mnist_train_func,
+    )
+    from ray.train.tensorflow import TensorflowTrainer
+
+    num_workers = 2
+    epochs = 3
+
+    config = {"lr": 1e-3, "batch_size": 64, "epochs": epochs}
+    trainer = TensorflowTrainer(
+        tensorflow_mnist_train_func,
+        train_loop_config=config,
+        scaling_config=ScalingConfig(num_workers=num_workers, use_gpu=True),
+    )
+    trainer.fit()
+
+
+def test_torch_fashion_mnist_gpu(ray_start_4_cpus_2_gpus):
+    num_workers = 2
+    epochs = 3
+
+    config = {"lr": 1e-3, "batch_size_per_worker": 32, "epochs": epochs}
+    trainer = TorchTrainer(
+        fashion_mnist_train_func,
+        train_loop_config=config,
+        scaling_config=ScalingConfig(num_workers=num_workers, use_gpu=True),
+    )
+    trainer.fit()
+
+
+def test_train_linear_dataset_gpu(ray_start_4_cpus_2_gpus):
+    from ray.train.examples.pytorch.torch_regression_example import train_regression
+
+    train_regression(num_workers=2, use_gpu=True)
+
+
+if __name__ == "__main__":
+    import sys
+
+    sys.exit(pytest.main(["-v", "-x", "-s", __file__]))

python/ray/train/tests/test_horovod_trainer.py

@@ -0,0 +1,78 @@
+import os
+
+import pytest
+import torch
+import torch.nn
+from torch.utils.data import DataLoader
+from torchvision import datasets
+from torchvision.transforms import transforms
+
+import ray
+from ray.train import ScalingConfig
+from ray.train.examples.horovod.horovod_pytorch_example import (
+    Net,
+    train_func as hvd_train_func,
+)
+from ray.train.horovod import HorovodTrainer
+
+
+@pytest.fixture
+def ray_start_4_cpus():
+    address_info = ray.init(num_cpus=4)
+    yield address_info
+    # The code after the yield will run as teardown code.
+    ray.shutdown()
+
+
+def run_image_prediction(model: torch.nn.Module, images: torch.Tensor) -> torch.Tensor:
+    model.eval()
+    with torch.no_grad():
+        return torch.exp(model(images)).argmax(dim=1)
+
+
+def test_horovod(ray_start_4_cpus):
+    def train_func(config):
+        result = hvd_train_func(config)
+        assert len(result) == epochs
+        assert result[-1] < result[0]
+
+    num_workers = 1
+    epochs = 10
+    scaling_config = ScalingConfig(num_workers=num_workers)
+    config = {"num_epochs": epochs, "save_model_as_dict": False}
+    trainer = HorovodTrainer(
+        train_loop_per_worker=train_func,
+        train_loop_config=config,
+        scaling_config=scaling_config,
+    )
+    result = trainer.fit()
+
+    model = Net()
+    with result.checkpoint.as_directory() as checkpoint_dir:
+        model.load_state_dict(torch.load(os.path.join(checkpoint_dir, "model.pt")))
+
+    # Find some test data to run on.
+    test_set = datasets.MNIST(
+        "./data",
+        train=False,
+        download=True,
+        transform=transforms.Compose(
+            [transforms.ToTensor(), transforms.Normalize((0.1307,), (0.3081,))]
+        ),
+    )
+
+    test_dataloader = DataLoader(test_set, batch_size=10)
+    test_dataloader_iter = iter(test_dataloader)
+    images, labels = next(
+        test_dataloader_iter
+    )  # only running a batch inference of 10 images
+    predicted_labels = run_image_prediction(model, images)
+    assert torch.equal(predicted_labels, labels)
+
+
+if __name__ == "__main__":
+    import sys
+
+    import pytest
+
+    sys.exit(pytest.main(["-v", "-x", __file__]))

python/ray/train/tests/test_iter_torch_batches_gpu.py

@@ -0,0 +1,325 @@
+from typing import Dict, List, Optional, Tuple, Union
+
+import numpy as np
+import pandas as pd
+import pyarrow as pa
+import pytest
+import torch
+
+import ray
+import ray.train.torch
+from ray.data.iterator import (
+    ArrowBatchCollateFn,
+    NumpyBatchCollateFn,
+    PandasBatchCollateFn,
+)
+from ray.data.util.torch_utils import (
+    arrow_batch_to_tensors,
+    convert_ndarray_batch_to_torch_tensor_batch,
+)
+
+
+@pytest.fixture(scope="module")
+def ray_start_4_cpus_1_gpu():
+    address_info = ray.init(num_cpus=4, num_gpus=1)
+    yield address_info
+    ray.shutdown()
+
+
+def _chunk_table_in_half(table: pa.Table) -> pa.Table:
+    num_rows = table.num_rows
+    mid = num_rows // 2
+
+    new_columns = []
+
+    for col in table.itercolumns():
+        # Slice the column in two halves
+        first_half = col.slice(0, mid)
+        second_half = col.slice(mid)
+
+        # Create a chunked array with two chunks
+        chunked_array = pa.chunked_array([first_half, second_half], type=col.type)
+        new_columns.append(chunked_array)
+
+    # Create a new table with the same schema but chunked columns
+    return pa.Table.from_arrays(new_columns, schema=table.schema)
+
+
+class SingleTensorArrowBatchCollateFn(ArrowBatchCollateFn):
+    """Collate function that returns only the id column as a tensor."""
+
+    def __call__(self, batch: pa.Table) -> torch.Tensor:
+        """Return only the id column as a tensor."""
+        assert isinstance(batch, pa.Table)
+        tensor_dict = arrow_batch_to_tensors(batch, combine_chunks=True)
+        return tensor_dict["id"]
+
+
+class TupleArrowBatchCollateFn(ArrowBatchCollateFn):
+    """Collate function that returns id and value as a tuple of tensors."""
+
+    def __call__(self, batch: pa.Table) -> Tuple[torch.Tensor, torch.Tensor]:
+        """Return id and value as a tuple of tensors."""
+        assert isinstance(batch, pa.Table)
+        tensor_dict = arrow_batch_to_tensors(batch, combine_chunks=True)
+        return tensor_dict["id"], tensor_dict["value"]
+
+
+class ListArrowBatchCollateFn(TupleArrowBatchCollateFn):
+    """Collate function that returns id and value as a list of tensors."""
+
+    def __call__(self, batch: pa.Table) -> List[torch.Tensor]:
+        return list(super().__call__(batch))
+
+
+class DictArrowBatchCollateFn(ArrowBatchCollateFn):
+    """Collate function that returns id and value as a dictionary of tensors."""
+
+    def __call__(self, batch: pa.Table) -> Dict[str, torch.Tensor]:
+        """Return id and value as a dictionary of tensors."""
+        assert isinstance(batch, pa.Table)
+        return arrow_batch_to_tensors(batch, combine_chunks=True)
+
+
+class ChunkedDictArrowBatchCollateFn(ArrowBatchCollateFn):
+    """Collate function that returns id and value as a dictionary of chunked tensors."""
+
+    def __call__(self, batch: pa.Table) -> Dict[str, List[torch.Tensor]]:
+        assert isinstance(batch, pa.Table)
+        modified_batch = _chunk_table_in_half(batch)
+        return arrow_batch_to_tensors(modified_batch, combine_chunks=False)
+
+
+class SingleTensorNumpyBatchCollateFn(NumpyBatchCollateFn):
+    """Collate function that returns only the id array as a tensor."""
+
+    def __call__(self, batch: Dict[str, np.ndarray]) -> torch.Tensor:
+        """Return only the id array as a tensor."""
+        assert isinstance(batch, dict)
+        tensor_dict = convert_ndarray_batch_to_torch_tensor_batch(batch)
+        return tensor_dict["id"]
+
+
+class TupleNumpyBatchCollateFn(NumpyBatchCollateFn):
+    """Collate function that returns id and value as a tuple of tensors."""
+
+    def __call__(
+        self, batch: Dict[str, np.ndarray]
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        assert isinstance(batch, dict)
+        tensor_dict = convert_ndarray_batch_to_torch_tensor_batch(batch)
+        return tensor_dict["id"], tensor_dict["value"]
+
+
+class ListNumpyBatchCollateFn(TupleNumpyBatchCollateFn):
+    """Collate function that returns id and value as a list of tensors."""
+
+    def __call__(self, batch: Dict[str, np.ndarray]) -> List[torch.Tensor]:
+        return list(super().__call__(batch))
+
+
+class DictNumpyBatchCollateFn(NumpyBatchCollateFn):
+    """Collate function that returns id and value as a dictionary of tensors."""
+
+    def __call__(self, batch: Dict[str, np.ndarray]) -> Dict[str, torch.Tensor]:
+        assert isinstance(batch, dict)
+        return convert_ndarray_batch_to_torch_tensor_batch(batch)
+
+
+class BasePandasBatchCollateFn(PandasBatchCollateFn):
+    """Base class for Pandas batch collate functions that process and convert to tensors.
+
+    This class provides common functionality for processing Pandas DataFrames and converting
+    them to PyTorch tensors. It handles device placement and dtype conversion.
+
+    Args:
+        device: Optional device to place tensors on. Can be a string (e.g. "cpu", "cuda:0")
+            or a torch.device object.
+    """
+
+    device: Optional[str]
+
+    def __init__(
+        self,
+        device: Optional[Union[str, torch.device]] = None,
+    ) -> None:
+        super().__init__()
+        if isinstance(device, str):
+            self.device = torch.device(device)
+        else:
+            self.device = device
+
+    def _process_batch(self, batch: pd.DataFrame) -> pd.DataFrame:
+        """Process the batch by adding 5 to the id column.
+
+        Args:
+            batch: Input Pandas DataFrame.
+
+        Returns:
+            A new DataFrame with modified "id" column and original "value" column.
+        """
+        return pd.DataFrame({"id": batch["id"] + 5, "value": batch["id"]})
+
+    def _get_tensors(self, batch: pd.DataFrame) -> Dict[str, torch.Tensor]:
+        """Convert batch to tensors.
+
+        Args:
+            batch: Input Pandas DataFrame to convert to tensors.
+
+        Returns:
+            Dictionary mapping column names to PyTorch tensors.
+        """
+        return convert_ndarray_batch_to_torch_tensor_batch(
+            batch.to_dict("series"), dtypes=None, device=None
+        )
+
+
+class SingleTensorPandasBatchCollateFn(PandasBatchCollateFn):
+    """Collate function that returns only the id column as a tensor."""
+
+    def __call__(self, batch: pd.DataFrame) -> torch.Tensor:
+        tensor_dict = convert_ndarray_batch_to_torch_tensor_batch(
+            batch.to_dict("series")
+        )
+        return tensor_dict["id"]
+
+
+class TuplePandasBatchCollateFn(PandasBatchCollateFn):
+    """Collate function that returns id and value as a tuple of tensors."""
+
+    def __call__(self, batch: pd.DataFrame) -> Tuple[torch.Tensor, torch.Tensor]:
+        tensor_dict = convert_ndarray_batch_to_torch_tensor_batch(
+            batch.to_dict("series")
+        )
+        return tensor_dict["id"], tensor_dict["value"]
+
+
+class ListPandasBatchCollateFn(TuplePandasBatchCollateFn):
+    """Collate function that returns id and value as a list of tensors."""
+
+    def __call__(self, batch: pd.DataFrame) -> List[torch.Tensor]:
+        return list(super().__call__(batch))
+
+
+class DictPandasBatchCollateFn(PandasBatchCollateFn):
+    """Collate function that returns id and value as a dictionary of tensors."""
+
+    def __call__(self, batch: pd.DataFrame) -> Dict[str, torch.Tensor]:
+        return convert_ndarray_batch_to_torch_tensor_batch(batch.to_dict("series"))
+
+
+@pytest.fixture
+def collate_fn_map():
+    """Fixture that provides Arrow, Numpy, Pandas custom collate functions."""
+
+    return {
+        "arrow": {
+            "default": None,
+            "single": SingleTensorArrowBatchCollateFn(),
+            "tuple": TupleArrowBatchCollateFn(),
+            "list": ListArrowBatchCollateFn(),
+            "dict": DictArrowBatchCollateFn(),
+            "chunked_dict": ChunkedDictArrowBatchCollateFn(),
+        },
+        "numpy": {
+            "single": SingleTensorNumpyBatchCollateFn(),
+            "tuple": TupleNumpyBatchCollateFn(),
+            "dict": DictNumpyBatchCollateFn(),
+            "list": ListNumpyBatchCollateFn(),
+        },
+        "pandas": {
+            "single": SingleTensorPandasBatchCollateFn(),
+            "tuple": TuplePandasBatchCollateFn(),
+            "dict": DictPandasBatchCollateFn(),
+            "list": ListPandasBatchCollateFn(),
+        },
+    }
+
+
+@pytest.mark.parametrize("collate_batch_type", ["arrow", "numpy", "pandas"])
+@pytest.mark.parametrize(
+    "return_type", ["single", "tuple", "dict", "list", "chunked_dict", "default"]
+)
+@pytest.mark.parametrize("device", ["cpu", "cuda:0"])
+@pytest.mark.parametrize("pin_memory", [True, False])
+def test_custom_batch_collate_fn(
+    ray_start_4_cpus_1_gpu,
+    monkeypatch,
+    collate_batch_type,
+    return_type,
+    device,
+    collate_fn_map,
+    pin_memory,
+):
+    """Tests that custom batch collate functions can be used to modify
+    the batch before it is converted to a PyTorch tensor.
+
+    Note that the collate_fn doesn't move the tensors to the device --
+    that happens in the iterator (finalize_fn).
+    """
+    # Skip GPU tests if CUDA is not available
+    if device == "cuda:0" and not torch.cuda.is_available():
+        pytest.skip("CUDA is not available")
+
+    # Skip pin_memory tests if CUDA is not available
+    if pin_memory and not torch.cuda.is_available():
+        pytest.skip("pin_memory is set to True, but CUDA is not available.")
+
+    # Skip tests if pin_memory is set to True and the collate function is not the
+    # DefaultCollateFn.
+    if pin_memory and not (collate_batch_type == "arrow" and return_type == "default"):
+        pytest.skip(
+            "pin_memory is set to True, but the collate function is not the DefaultCollateFn."
+        )
+
+    collate_fn = collate_fn_map[collate_batch_type].get(return_type)
+    if collate_fn is None:
+        pytest.skip(
+            f"Collate function not found for ({collate_batch_type}, {return_type})"
+        )
+
+    # Set the device that's returned by device="auto" -> get_device()
+    # This is used in `finalize_fn` to move the tensors to the correct device.
+    device = torch.device(device)
+    monkeypatch.setattr(ray.train.utils, "_in_ray_train_worker", lambda: True)
+    monkeypatch.setattr(ray.train.torch, "get_device", lambda: device)
+
+    ds = ray.data.from_items(
+        [{"id": i + 5, "value": i} for i in range(5)],
+    )
+    it = ds.iterator()
+
+    for batch in it.iter_torch_batches(collate_fn=collate_fn, pin_memory=pin_memory):
+        if return_type == "single":
+            assert isinstance(batch, torch.Tensor)
+            assert sorted(batch.tolist()) == list(range(5, 10))
+            assert batch.device == device
+            if pin_memory and device.type == "cpu":
+                assert batch.is_pinned()
+        elif return_type == "dict" or return_type == "chunked_dict":
+            # Chunked dicts get concatenated to single Tensors on the device,
+            # so the assertions are shared with the dict case.
+            assert isinstance(batch, dict)
+            assert sorted(batch["id"].tolist()) == list(range(5, 10))
+            assert sorted(batch["value"].tolist()) == list(range(5))
+            assert batch["id"].device == device
+            assert batch["value"].device == device
+            if pin_memory and device.type == "cpu":
+                assert batch["id"].is_pinned()
+                assert batch["value"].is_pinned()
+        else:  # tuple or list
+            assert isinstance(batch, (tuple, list))
+            assert len(batch) == 2
+            assert sorted(batch[0].tolist()) == list(range(5, 10))
+            assert sorted(batch[1].tolist()) == list(range(5))
+            assert batch[0].device == device
+            assert batch[1].device == device
+            if pin_memory and device.type == "cpu":
+                assert batch[0].is_pinned()
+                assert batch[1].is_pinned()
+
+
+if __name__ == "__main__":
+    import sys
+
+    sys.exit(pytest.main(["-v", __file__]))

python/ray/train/tests/test_lightgbm_trainer.py

@@ -0,0 +1,218 @@
+import math
+from unittest import mock
+
+import lightgbm as lgbm
+import pandas as pd
+import pytest
+from sklearn.datasets import load_breast_cancer
+from sklearn.model_selection import train_test_split
+
+import ray
+from ray import tune
+from ray.train import ScalingConfig
+from ray.train.constants import TRAIN_DATASET_KEY
+from ray.train.lightgbm import LightGBMTrainer, RayTrainReportCallback
+
+
+@pytest.fixture
+def ray_start_6_cpus():
+    address_info = ray.init(num_cpus=6)
+    yield address_info
+    # The code after the yield will run as teardown code.
+    ray.shutdown()
+
+
+@pytest.fixture
+def ray_start_8_cpus():
+    address_info = ray.init(num_cpus=8)
+    yield address_info
+    # The code after the yield will run as teardown code.
+    ray.shutdown()
+
+
+scale_config = ScalingConfig(num_workers=2)
+
+data_raw = load_breast_cancer()
+dataset_df = pd.DataFrame(data_raw["data"], columns=data_raw["feature_names"])
+dataset_df["target"] = data_raw["target"]
+train_df, test_df = train_test_split(dataset_df, test_size=0.3)
+
+params = {
+    "objective": "binary",
+    "metric": ["binary_logloss", "binary_error"],
+}
+
+
+def get_num_trees(booster: lgbm.Booster) -> int:
+    return booster.current_iteration()
+
+
+def test_fit_with_categoricals(ray_start_6_cpus):
+    train_df_with_cat = train_df.copy()
+    test_df_with_cat = test_df.copy()
+    train_df_with_cat["categorical_column"] = pd.Series(
+        (["A", "B"] * math.ceil(len(train_df_with_cat) / 2))[: len(train_df_with_cat)]
+    ).astype("category")
+    test_df_with_cat["categorical_column"] = pd.Series(
+        (["A", "B"] * math.ceil(len(test_df_with_cat) / 2))[: len(test_df_with_cat)]
+    ).astype("category")
+
+    train_dataset = ray.data.from_pandas(train_df_with_cat)
+    valid_dataset = ray.data.from_pandas(test_df_with_cat)
+    trainer = LightGBMTrainer(
+        scaling_config=scale_config,
+        label_column="target",
+        params=params,
+        datasets={TRAIN_DATASET_KEY: train_dataset, "valid": valid_dataset},
+    )
+    result = trainer.fit()
+    checkpoint = result.checkpoint
+    model = LightGBMTrainer.get_model(checkpoint)
+    assert model.pandas_categorical == [["A", "B"]]
+
+
+def test_resume_from_checkpoint(ray_start_6_cpus, tmpdir):
+    train_dataset = ray.data.from_pandas(train_df)
+    valid_dataset = ray.data.from_pandas(test_df)
+    trainer = LightGBMTrainer(
+        scaling_config=scale_config,
+        label_column="target",
+        params=params,
+        num_boost_round=5,
+        datasets={TRAIN_DATASET_KEY: train_dataset, "valid": valid_dataset},
+    )
+    result = trainer.fit()
+    model = LightGBMTrainer.get_model(result.checkpoint)
+    assert get_num_trees(model) == 5
+
+    trainer = LightGBMTrainer(
+        scaling_config=scale_config,
+        label_column="target",
+        params=params,
+        num_boost_round=10,
+        datasets={TRAIN_DATASET_KEY: train_dataset, "valid": valid_dataset},
+        resume_from_checkpoint=result.checkpoint,
+    )
+    result = trainer.fit()
+    checkpoint = result.checkpoint
+    model = LightGBMTrainer.get_model(checkpoint)
+    assert get_num_trees(model) == 10
+
+
+def test_fit_with_arrow_backed_pandas_dtypes(ray_start_6_cpus):
+    # `from_items` produces Arrow-backed blocks, so `to_pandas()` inside the
+    # trainer returns Arrow-backed dtypes — the regression path this test guards.
+    train_dataset = ray.data.from_items(train_df.to_dict("records"))
+    valid_dataset = ray.data.from_items(test_df.to_dict("records"))
+
+    trainer = LightGBMTrainer(
+        scaling_config=scale_config,
+        label_column="target",
+        params=params,
+        datasets={TRAIN_DATASET_KEY: train_dataset, "valid": valid_dataset},
+    )
+    result = trainer.fit()
+    model = LightGBMTrainer.get_model(result.checkpoint)
+    assert get_num_trees(model) == 10
+
+
+@pytest.mark.parametrize(
+    "freq_end_expected",
+    [
+        # With num_boost_round=25 with 0 indexing, the checkpoints will be at:
+        (4, True, 7),  # 3, 7, 11, 15, 19, 23, 24 (end)
+        (4, False, 6),  # 3, 7, 11, 15, 19, 23
+        (5, True, 5),  # 4, 9, 14, 19, 24
+        (0, True, 1),  # 24 (end)
+        (0, False, 0),
+    ],
+)
+def test_checkpoint_freq(ray_start_6_cpus, freq_end_expected):
+    freq, end, expected = freq_end_expected
+
+    train_dataset = ray.data.from_pandas(train_df)
+    valid_dataset = ray.data.from_pandas(test_df)
+    trainer = LightGBMTrainer(
+        run_config=ray.train.RunConfig(
+            checkpoint_config=ray.train.CheckpointConfig(
+                checkpoint_frequency=freq, checkpoint_at_end=end
+            )
+        ),
+        scaling_config=scale_config,
+        label_column="target",
+        params=params,
+        num_boost_round=25,
+        datasets={TRAIN_DATASET_KEY: train_dataset, "valid": valid_dataset},
+    )
+    result = trainer.fit()
+
+    # Assert number of checkpoints
+    assert len(result.best_checkpoints) == expected, str(
+        [(metrics["training_iteration"], cp) for cp, metrics in result.best_checkpoints]
+    )
+
+    # Assert checkpoint numbers are increasing
+    cp_paths = [cp.path for cp, _ in result.best_checkpoints]
+    assert cp_paths == sorted(cp_paths), str(cp_paths)
+
+
+def test_tune(ray_start_8_cpus):
+    train_dataset = ray.data.from_pandas(train_df)
+    valid_dataset = ray.data.from_pandas(test_df)
+    trainer = LightGBMTrainer(
+        scaling_config=ScalingConfig(num_workers=2, resources_per_worker={"CPU": 1}),
+        label_column="target",
+        params={**params, "max_depth": 1},
+        datasets={TRAIN_DATASET_KEY: train_dataset, "valid": valid_dataset},
+    )
+
+    tuner = tune.Tuner(
+        trainer,
+        param_space={"params": {"max_depth": tune.grid_search([2, 4])}},
+    )
+    results = tuner.fit()
+    assert sorted([r.config["params"]["max_depth"] for r in results]) == [2, 4]
+
+
+def test_validation(ray_start_6_cpus):
+    valid_dataset = ray.data.from_pandas(test_df)
+    with pytest.raises(ValueError, match=TRAIN_DATASET_KEY):
+        LightGBMTrainer(
+            scaling_config=ScalingConfig(num_workers=2),
+            label_column="target",
+            params=params,
+            datasets={"valid": valid_dataset},
+        )
+    with pytest.raises(ValueError, match="label_column"):
+        LightGBMTrainer(
+            scaling_config=ScalingConfig(num_workers=2),
+            datasets={"train": valid_dataset},
+        )
+
+
+@pytest.mark.parametrize("rank", [None, 0, 1])
+def test_checkpoint_only_on_rank0(rank):
+    """Tests that the callback only reports checkpoints on rank 0,
+    or if the rank is not available (Tune usage)."""
+    callback = RayTrainReportCallback(frequency=2, checkpoint_at_end=True)
+
+    booster = mock.MagicMock()
+
+    with mock.patch("ray.train.get_context") as mock_get_context:
+        mock_context = mock.MagicMock()
+        mock_context.get_world_rank.return_value = rank
+        mock_get_context.return_value = mock_context
+
+        with callback._get_checkpoint(booster) as checkpoint:
+            if rank in (0, None):
+                assert checkpoint
+            else:
+                assert not checkpoint
+
+
+if __name__ == "__main__":
+    import sys
+
+    import pytest
+
+    sys.exit(pytest.main(["-v", "-x", __file__]))

python/ray/train/tests/test_minimal.py

@@ -0,0 +1,90 @@
+import pytest
+
+import ray
+from ray import train
+from ray.train import ScalingConfig
+from ray.train._internal.worker_group import WorkerGroup
+from ray.train.backend import Backend, BackendConfig
+from ray.train.data_parallel_trainer import DataParallelTrainer
+from ray.train.tests.util import create_dict_checkpoint, load_dict_checkpoint
+
+
+@pytest.fixture
+def ray_start_4_cpus():
+    address_info = ray.init(num_cpus=4)
+    yield address_info
+    # The code after the yield will run as teardown code.
+    ray.shutdown()
+
+
+class TestConfig(BackendConfig):
+    @property
+    def backend_cls(self):
+        return TestBackend
+
+
+class TestBackend(Backend):
+    def on_start(self, worker_group: WorkerGroup, backend_config: TestConfig):
+        pass
+
+    def on_shutdown(self, worker_group: WorkerGroup, backend_config: TestConfig):
+        pass
+
+
+def test_run(ray_start_4_cpus):
+    """Tests that Train can be run without any specific backends."""
+    num_workers = 2
+    key = "value"
+    value = 1
+    config = TestConfig()
+
+    def train_func():
+        checkpoint = train.get_checkpoint()
+        checkpoint_dict = load_dict_checkpoint(checkpoint)
+        if train.get_context().get_world_rank() == 0:
+            train.report(metrics=checkpoint_dict, checkpoint=checkpoint)
+        else:
+            train.report(metrics=checkpoint_dict)
+        return checkpoint_dict[key]
+
+    with create_dict_checkpoint({key: value}) as checkpoint:
+        trainer = DataParallelTrainer(
+            train_func,
+            backend_config=config,
+            resume_from_checkpoint=checkpoint,
+            scaling_config=ScalingConfig(num_workers=num_workers),
+        )
+        results = trainer.fit()
+
+        assert load_dict_checkpoint(results.checkpoint) == load_dict_checkpoint(
+            checkpoint
+        )
+
+
+def test_failure():
+    """Tests that backend frameworks and non-critical libraries are not imported."""
+    with pytest.raises(ModuleNotFoundError):
+        import torch  # noqa: F401
+
+    with pytest.raises(ModuleNotFoundError):
+        import tensorflow  # noqa: F401
+
+    with pytest.raises(ModuleNotFoundError):
+        import horovod  # noqa: F401
+
+    with pytest.raises(ModuleNotFoundError):
+        import accelerate  # noqa: F401
+
+    with pytest.raises(ModuleNotFoundError):
+        import transformers  # noqa: F401
+
+    with pytest.raises(ModuleNotFoundError):
+        import xgboost  # noqa: F401
+
+
+if __name__ == "__main__":
+    import sys
+
+    import pytest
+
+    sys.exit(pytest.main(["-v", "-x", __file__]))

python/ray/train/tests/test_new_persistence.py

@@ -0,0 +1,636 @@
+import logging
+import os
+import pickle
+import re
+import tempfile
+import time
+import uuid
+from contextlib import contextmanager
+from pathlib import Path
+from typing import List, Optional, Tuple
+
+import pyarrow.fs
+import pytest
+
+import ray
+from ray import train, tune
+from ray._common.test_utils import simulate_s3_bucket
+from ray.air._internal.uri_utils import URI
+from ray.air.constants import EXPR_RESULT_FILE
+from ray.train._checkpoint import Checkpoint
+from ray.train._internal.storage import (
+    StorageContext,
+    _delete_fs_path,
+    _download_from_fs_path,
+)
+from ray.train.base_trainer import TrainingFailedError
+from ray.train.data_parallel_trainer import DataParallelTrainer
+from ray.tune.trainable.trainable import _DICT_CHECKPOINT_FILE_NAME
+
+
+class TestConstants:
+    NUM_ITERATIONS = 6  # == num_checkpoints == num_artifacts
+    NUM_TRIALS = 2
+    NUM_WORKERS = 3
+
+    SCORE_KEY = "score"
+
+
+@contextmanager
+def mock_s3_bucket_uri():
+    port = 5002
+    region = "us-west-2"
+    with simulate_s3_bucket(port=port, region=region) as s3_uri:
+        import boto3
+
+        s3 = boto3.client(
+            "s3", region_name=region, endpoint_url=f"http://localhost:{port}"
+        )
+        # Bucket name will be autogenerated/unique per test
+        bucket_name = URI(s3_uri).name
+        s3.create_bucket(
+            Bucket=bucket_name,
+            CreateBucketConfiguration={"LocationConstraint": region},
+        )
+        # Disable server HTTP request logging
+        logging.getLogger("werkzeug").setLevel(logging.WARNING)
+        yield URI(s3_uri)
+        logging.getLogger("werkzeug").setLevel(logging.INFO)
+
+
+@contextmanager
+def dummy_context_manager(*args, **kwargs):
+    yield "dummy value"
+
+
+@pytest.fixture(autouse=True, scope="module")
+def ray_start_4_cpus():
+    ray.init(num_cpus=4)
+    yield
+    ray.shutdown()
+
+
+def _create_mock_custom_fs(custom_fs_root_dir: Path) -> pyarrow.fs.FileSystem:
+    from fsspec.implementations.dirfs import DirFileSystem
+    from fsspec.implementations.local import LocalFileSystem
+
+    custom_fs_root_dir.mkdir(parents=True, exist_ok=True)
+    storage_filesystem = pyarrow.fs.PyFileSystem(
+        pyarrow.fs.FSSpecHandler(
+            DirFileSystem(path=str(custom_fs_root_dir), fs=LocalFileSystem())
+        )
+    )
+    return storage_filesystem
+
+
+@contextmanager
+def _resolve_storage_type(
+    storage_path_type: str, tmp_path: Path
+) -> Tuple[str, Optional[pyarrow.fs.FileSystem]]:
+    storage_path, storage_filesystem = None, None
+
+    context_manager = (
+        mock_s3_bucket_uri if storage_path_type == "cloud" else dummy_context_manager
+    )
+
+    with context_manager() as cloud_storage_path:
+        if storage_path_type == "nfs":
+            storage_path = str(tmp_path / "fake_nfs")
+        elif storage_path_type == "cloud":
+            storage_path = str(cloud_storage_path)
+        elif storage_path_type == "custom_fs":
+            storage_path = "mock_bucket"
+            storage_filesystem = _create_mock_custom_fs(tmp_path / "custom_fs")
+
+        yield storage_path, storage_filesystem
+
+
+def _get_local_inspect_dir(
+    root_local_path: Path,
+    storage_path: str,
+    storage_filesystem: Optional[pyarrow.fs.FileSystem],
+    storage_local_path: Path = None,
+) -> Tuple[Path, str]:
+    """Downloads the storage path -> local dir for inspecting contents.
+
+    Args:
+        root_local_path: Local directory to use as the inspect root.
+        storage_path: The storage path or URI to download from.
+        storage_filesystem: Optional custom filesystem to use.
+        storage_local_path: Local path that ``storage_path`` mirrors on disk
+            when no remote storage is configured.
+
+    Returns:
+        Tuple: (local_inspect_dir, storage_fs_path), where storage_fs_path
+            is the path to the storage path on the filesystem (e.g., prefix stripped).
+            This is used to check the correctness of paths returned from `Result`'s,
+            since URIs are hard to do comparisons with.
+    """
+    local_inspect_dir = root_local_path / "inspect"
+    if storage_path:
+        if storage_filesystem:
+            fs, storage_fs_path = storage_filesystem, storage_path
+        else:
+            fs, storage_fs_path = pyarrow.fs.FileSystem.from_uri(storage_path)
+        _download_from_fs_path(
+            fs=fs, fs_path=storage_fs_path, local_path=str(local_inspect_dir)
+        )
+    else:
+        fs, storage_fs_path = pyarrow.fs.LocalFileSystem(), str(storage_local_path)
+        local_inspect_dir = storage_local_path
+
+    return local_inspect_dir, storage_fs_path
+
+
+def _get_checkpoint_index(checkpoint_dir_name: str) -> int:
+    """Gets the checkpoint index from the checkpoint directory name."""
+    return int(checkpoint_dir_name.split("_")[-1])
+
+
+def _create_checkpoint_shard_filename(rank_str: str) -> str:
+    return f"checkpoint_shard-rank={rank_str}.pkl"
+
+
+def _get_checkpoint_shard_rank(checkpoint_shard_filename: str) -> int:
+    """Get the checkpoint shard rank from the filename."""
+    pattern = _create_checkpoint_shard_filename(r"(\d+)")
+    match = re.search(pattern, checkpoint_shard_filename)
+    assert match
+    return int(match.group(1))
+
+
+def train_fn(config):
+    in_trainer = config.get("in_trainer", False)
+    if in_trainer:
+        from ray.train._internal.session import _TrainSession, get_session
+
+        train_session = get_session()
+
+        assert isinstance(train_session, _TrainSession)
+        assert train_session.storage
+        assert train_session.storage.checkpoint_fs_path
+
+        # Check that the working dir for each worker is the shared trial dir.
+        assert (
+            Path.cwd() == Path(train_session.storage.trial_working_directory).resolve()
+        )
+
+    start = 0
+
+    checkpoint = train.get_checkpoint()
+    if checkpoint:
+        custom_restore_fn = config.get("custom_restore_fn")
+        if custom_restore_fn:
+            state = custom_restore_fn(checkpoint)
+        else:
+            with checkpoint.as_directory() as checkpoint_dir:
+                with open(os.path.join(checkpoint_dir, "checkpoint.pkl"), "rb") as f:
+                    state = pickle.load(f)
+        print("Loaded back state from checkpoint:", state)
+        start = state["iter"] + 1
+
+    for i in range(start, config.get("num_iterations", 5)):
+        time.sleep(config.get("time_per_iter", 0.25))
+
+        metrics = {"iter": i, TestConstants.SCORE_KEY: i}
+
+        # Save an artifact in the local trial dir.
+        rank = train.get_context().get_world_rank()
+        artifact_file_name = (
+            f"artifact-rank={rank}-iter={i}.txt"
+            if in_trainer
+            else f"artifact-iter={i}.txt"
+        )
+        with open(artifact_file_name, "w") as f:
+            f.write(f"{i}")
+
+        if in_trainer and train.get_context().get_world_rank() in config.get(
+            "no_checkpoint_ranks", []
+        ):
+            train.report(metrics)
+        else:
+            with tempfile.TemporaryDirectory() as temp_dir:
+                with open(os.path.join(temp_dir, "checkpoint.pkl"), "wb") as f:
+                    pickle.dump({"iter": i}, f)
+
+                if in_trainer:
+                    checkpoint_file_name = _create_checkpoint_shard_filename(str(rank))
+                    with open(os.path.join(temp_dir, checkpoint_file_name), "wb") as f:
+                        pickle.dump({"iter": i}, f)
+
+                with config.get("custom_save_fn", dummy_context_manager)(temp_dir):
+                    train.report(
+                        metrics, checkpoint=Checkpoint.from_directory(temp_dir)
+                    )
+                # `train.report` should not have deleted this!
+                assert os.path.exists(temp_dir)
+
+        if i in config.get("fail_iters", []):
+            raise RuntimeError(f"Failing on iter={i}!!")
+
+
+class ClassTrainable(tune.Trainable):
+    """Implement (almost) the same thing as `train_fn` but as a class."""
+
+    def setup(self, config):
+        # Save some markers in the trial dir.
+        tmp_path = config.get("tmp_path")
+        self.fail_markers = {
+            i: tmp_path / f"fail_marker_{self.trial_id}_iter={i}"
+            for i in config.get("fail_iters", [])
+        }
+        setup_marker = tmp_path / f"setup_marker_{self.trial_id}"
+        if not setup_marker.exists():
+            for marker in self.fail_markers.values():
+                marker.touch()
+            setup_marker.touch()
+
+        self.save_as_dict = config.get("save_checkpoint_as_dict", False)
+
+    def step(self) -> dict:
+        if self.iteration in self.fail_markers:
+            marker = self.fail_markers[self.iteration]
+            if marker.exists():
+                marker.unlink()
+                raise RuntimeError(f"Failing on iter={self.iteration}")
+
+        # Save an artifact in the local trial dir.
+        artifact_file_name = f"artifact-iter={self.iteration}.txt"
+        with open(artifact_file_name, "w") as f:
+            f.write(f"{self.iteration}")
+
+        return {
+            "score": 1,
+            "done": self.iteration >= self.config.get("num_iterations") - 1,
+            "should_checkpoint": True,
+        }
+
+    def save_checkpoint(self, temp_checkpoint_dir) -> str:
+        if self.save_as_dict:
+            return {"dummy": "data"}
+        (Path(temp_checkpoint_dir) / "checkpoint.pkl").write_text("dummy")
+        return temp_checkpoint_dir
+
+    def load_checkpoint(self, checkpoint_dict_or_path):
+        print("Loading state from:", checkpoint_dict_or_path)
+        print("At iteration =", self.iteration)
+        if self.save_as_dict:
+            assert checkpoint_dict_or_path == {"dummy": "data"}
+        else:
+            assert (
+                Path(checkpoint_dict_or_path) / "checkpoint.pkl"
+            ).read_text() == "dummy"
+
+
+def _resume_from_checkpoint(
+    checkpoint: Checkpoint,
+    expected_state: dict,
+    storage_path: Optional[str] = None,
+    storage_filesystem: Optional[pyarrow.fs.FileSystem] = None,
+):
+    print(f"\nStarting run with `resume_from_checkpoint`: {checkpoint}\n")
+
+    def assert_fn(config):
+        checkpoint_to_check = train.get_checkpoint()
+        with checkpoint_to_check.as_directory() as checkpoint_dir:
+            with open(os.path.join(checkpoint_dir, "checkpoint.pkl"), "rb") as f:
+                state = pickle.load(f)
+
+        print("Loaded state from `resume_from_checkpoint`:", state)
+        print("Expected state:", expected_state)
+        assert state == expected_state, (state, expected_state)
+
+        dummy_ckpt = tempfile.mkdtemp()
+        with open(os.path.join(dummy_ckpt, "dummy.txt"), "w") as f:
+            f.write("data")
+        train.report({"dummy": 1}, checkpoint=Checkpoint.from_directory(dummy_ckpt))
+
+    trainer = DataParallelTrainer(
+        assert_fn,
+        scaling_config=train.ScalingConfig(num_workers=2),
+        run_config=train.RunConfig(
+            name="test_resume_from_checkpoint",
+            storage_path=storage_path,
+            storage_filesystem=storage_filesystem,
+        ),
+        resume_from_checkpoint=checkpoint,
+    )
+    result = trainer.fit()
+
+    # Make sure that the checkpoint indexing starts from scratch.
+    assert Path(
+        result.checkpoint.path
+    ).name == StorageContext._make_checkpoint_dir_name(0)
+
+    # Clean up this run's experiment directory immediately after.
+    _delete_fs_path(result.filesystem, Path(result.path).parent.as_posix())
+
+
+def _assert_storage_contents(
+    local_inspect_dir: Path,
+    exp_name: str,
+    checkpoint_config: train.CheckpointConfig,
+    trainable_name: str,
+    test_trainer: bool,
+    no_checkpoint_ranks: List[int] = None,
+    constants: type = TestConstants,
+):
+    no_checkpoint_ranks = no_checkpoint_ranks or []
+
+    # Second, inspect the contents of the storage path
+    storage_path_ls = list(local_inspect_dir.glob("*"))
+    assert len(storage_path_ls) == 1  # Only expect 1 experiment dir
+    exp_dir = storage_path_ls[0]
+    assert exp_dir.name == exp_name
+
+    # Files synced by the driver
+    assert len(list(exp_dir.glob("tuner.pkl"))) == 1
+    if test_trainer:
+        assert len(list(exp_dir.glob("trainer.pkl"))) == 1
+    # 2 copies of these files:
+    # 1 for the initial run, and 1 for the manually restored run.
+    assert len(list(exp_dir.glob("basic-variant-state-*"))) == 2
+    assert len(list(exp_dir.glob("experiment_state-*"))) == 2
+
+    # Files synced by the worker
+    assert (
+        len(list(exp_dir.glob(f"{trainable_name}*"))) == 1
+        if test_trainer
+        else constants.NUM_TRIALS
+    )
+    for trial_dir in exp_dir.glob(f"{trainable_name}*"):
+        # If set, expect num_to_keep. Otherwise, expect to see all of them.
+        expected_num_checkpoints = (
+            checkpoint_config.num_to_keep or constants.NUM_ITERATIONS
+        )
+
+        assert len(list(trial_dir.glob("checkpoint_*"))) == expected_num_checkpoints
+        checkpoint_idxs = sorted(
+            [
+                _get_checkpoint_index(checkpoint_dir.name)
+                for checkpoint_dir in trial_dir.glob("checkpoint_*")
+            ]
+        )
+        # Ex: If num_to_keep=2 out of 6 total checkpoints,
+        # expect checkpoint_004 and checkpoint_005.
+        assert checkpoint_idxs == list(
+            range(
+                constants.NUM_ITERATIONS - expected_num_checkpoints,
+                constants.NUM_ITERATIONS,
+            )
+        )
+
+        for checkpoint_dir in trial_dir.glob("checkpoint_*"):
+            # 1 shared checkpoint.pkl file, written by the trainable / all workers.
+            assert (
+                len(list(checkpoint_dir.glob("checkpoint.pkl"))) == 1
+                # NOTE: Dict checkpoint is only for the ClassTrainable.
+                or len(list(checkpoint_dir.glob(_DICT_CHECKPOINT_FILE_NAME))) == 1
+            )
+            if test_trainer:
+                # 1 checkpoint shard per worker.
+                # Unless the worker did not report a checkpoint (no_checkpoint_ranks).
+                assert {
+                    _get_checkpoint_shard_rank(checkpoint_shard.name)
+                    for checkpoint_shard in checkpoint_dir.glob(
+                        "checkpoint_shard-*.pkl"
+                    )
+                } == {
+                    i
+                    for i in range(constants.NUM_WORKERS)
+                    if i not in no_checkpoint_ranks
+                }
+
+        if test_trainer:
+            expected_num_artifacts = constants.NUM_ITERATIONS * constants.NUM_WORKERS
+        else:
+            expected_num_artifacts = constants.NUM_ITERATIONS
+        assert len(list(trial_dir.glob("artifact-*"))) == expected_num_artifacts
+
+        # NOTE: This result file is synced by the driver.
+        assert len(list(trial_dir.glob(EXPR_RESULT_FILE))) == 1
+
+
+@pytest.mark.parametrize("trainable", [train_fn, ClassTrainable])
+@pytest.mark.parametrize("storage_path_type", ["nfs", "cloud", "custom_fs"])
+@pytest.mark.parametrize(
+    "checkpoint_config",
+    [train.CheckpointConfig(), train.CheckpointConfig(num_to_keep=2)],
+)
+def test_tuner(
+    tmp_path,
+    trainable,
+    storage_path_type,
+    checkpoint_config: train.CheckpointConfig,
+):
+    """End-to-end test that the new persistence mode works with the Tuner API.
+    This test covers many `storage_path_type` options:
+    - storage_path=None --> save locally to the default local path (e.g., ~/ray_results)
+    - storage_path="nfs" --> save locally to a fake NFS path
+    - storage_path="cloud" --> save to a mock S3 bucket
+    - storage_path="custom_fs" --> save to a custom pyarrow filesystem
+        - The custom fs is a local filesystem that appends a path prefix to every path.
+
+    This is the expected output at the storage path:
+
+    {storage_path}/{exp_name}
+    ├── tuner.pkl                   <- Driver artifacts (global experiment state)
+    ├── basic-variant-state.json
+    ├── experiment_state.json
+    ├── train_fn_a2b9e_00000_0_...
+    │   ├── artifact-iter=0.txt     <- Trial artifacts
+    │   ├── ...
+    │   ├── checkpoint_000000       <- Trial checkpoints
+    │   │   └── checkpoint.pkl
+    │   ├── ...
+    │   ├── events.out.tfevents...  <- Driver artifacts (trial results)
+    │   ├── params.json
+    │   ├── params.pkl
+    │   ├── progress.csv
+    │   └── result.json
+    └── train_fn_a2b9e_00001_1_...
+        └── ...                     <- Same as above
+    """
+    exp_name = f"tuner_persistence_test-{uuid.uuid4().hex}"
+
+    with _resolve_storage_type(storage_path_type, tmp_path) as (
+        storage_path,
+        storage_filesystem,
+    ):
+        run_config = train.RunConfig(
+            storage_path=storage_path,
+            storage_filesystem=storage_filesystem,
+            name=exp_name,
+            verbose=0,
+            failure_config=train.FailureConfig(max_failures=1),
+            checkpoint_config=checkpoint_config,
+            sync_config=train.SyncConfig(sync_artifacts=True),
+        )
+        tuner = tune.Tuner(
+            trainable,
+            param_space={
+                "num_iterations": TestConstants.NUM_ITERATIONS,
+                "fail_iters": [2, 4],
+                # NOTE: This param is only used in the ClassTrainable.
+                "save_checkpoint_as_dict": tune.grid_search([True, False]),
+                "tmp_path": tmp_path,
+            },
+            run_config=run_config,
+            # 2 samples (from the grid search). Run 1 at at time to test actor reuse
+            tune_config=tune.TuneConfig(num_samples=1, max_concurrent_trials=1),
+        )
+        result_grid = tuner.fit()
+        assert result_grid.errors
+
+        restored_tuner = tune.Tuner.restore(
+            path=str(URI(run_config.storage_path) / exp_name),
+            trainable=trainable,
+            storage_filesystem=storage_filesystem,
+            resume_errored=True,
+        )
+        result_grid = restored_tuner.fit()
+        assert not result_grid.errors
+
+        local_inspect_dir, storage_fs_path = _get_local_inspect_dir(
+            root_local_path=tmp_path,
+            storage_path=run_config.storage_path,
+            storage_filesystem=storage_filesystem,
+        )
+
+    # First, check that the ResultGrid returns the correct paths.
+    print(result_grid)
+    experiment_fs_path = result_grid.experiment_path
+    assert isinstance(result_grid.filesystem, pyarrow.fs.FileSystem), result_grid
+    assert experiment_fs_path == os.path.join(storage_fs_path, exp_name)
+    assert len(result_grid) == TestConstants.NUM_TRIALS
+    for result in result_grid:
+        trial_fs_path = result.path
+        assert isinstance(result.filesystem, pyarrow.fs.FileSystem), result
+        assert trial_fs_path.startswith(experiment_fs_path)
+        for checkpoint, _ in result.best_checkpoints:
+            assert checkpoint.path.startswith(trial_fs_path)
+
+    # Next, inspect the storage path contents.
+    _assert_storage_contents(
+        local_inspect_dir,
+        exp_name,
+        checkpoint_config,
+        trainable_name=trainable.__name__,
+        test_trainer=False,
+    )
+
+
+@pytest.mark.parametrize("storage_path_type", ["nfs", "cloud", "custom_fs"])
+@pytest.mark.parametrize(
+    "checkpoint_config",
+    [
+        train.CheckpointConfig(),
+        train.CheckpointConfig(
+            num_to_keep=1,
+            checkpoint_score_attribute=TestConstants.SCORE_KEY,
+            checkpoint_score_order="max",
+        ),
+    ],
+)
+def test_trainer(
+    tmp_path, storage_path_type, checkpoint_config: train.CheckpointConfig
+):
+    """Same end-to-end test as `test_tuner`, but also includes a
+    `DataParallelTrainer(resume_from_checkpoint)` test at the end.
+
+    {storage_path}/{exp_name}
+    ├── experiment_state-2023-07-28_10-00-38.json       <- Initial exp state
+    ├── basic-variant-state-2023-07-28_10-00-38.json
+    ├── experiment_state-2023-07-28_10-01-38.json       <- Restored exp state
+    ├── basic-variant-state-2023-07-28_10-01-38.json
+    ├── trainer.pkl
+    ├── tuner.pkl
+    └── DataParallelTrainer_46367_00000_0_...
+        ├── events.out.tfevents...
+        ├── params.json
+        ├── params.pkl
+        ├── progress.csv
+        ├── result.json
+        ├── checkpoint_000000
+        │   ├── checkpoint.pkl                    <- Shared checkpoint file
+        │   ├── checkpoint_shard-rank=0.pkl       <- Worker checkpoint shards
+        │   └── checkpoint_shard-rank=1.pkl
+        ├── ...
+        ├── artifact-rank=0-iter=0.txt            <- Worker artifacts
+        ├── artifact-rank=1-iter=0.txt
+        ├── ...
+        ├── artifact-rank=0-iter=1.txt
+        ├── artifact-rank=1-iter=1.txt
+        └── ...
+    """
+    exp_name = f"trainer_persistence_test-{uuid.uuid4().hex}"
+    no_checkpoint_ranks = [0]
+
+    with _resolve_storage_type(storage_path_type, tmp_path) as (
+        storage_path,
+        storage_filesystem,
+    ):
+        run_config = train.RunConfig(
+            storage_path=storage_path,
+            storage_filesystem=storage_filesystem,
+            name=exp_name,
+            verbose=0,
+            checkpoint_config=checkpoint_config,
+            failure_config=train.FailureConfig(max_failures=1),
+            sync_config=train.SyncConfig(sync_artifacts=True),
+        )
+        trainer = DataParallelTrainer(
+            train_fn,
+            train_loop_config={
+                "in_trainer": True,
+                "num_iterations": TestConstants.NUM_ITERATIONS,
+                "fail_iters": [2, 4],
+                # Test that global rank 0 is not required to checkpoint.
+                "no_checkpoint_ranks": no_checkpoint_ranks,
+            },
+            scaling_config=train.ScalingConfig(num_workers=TestConstants.NUM_WORKERS),
+            run_config=run_config,
+        )
+        print("\nStarting initial run.\n")
+        with pytest.raises(TrainingFailedError):
+            result = trainer.fit()
+
+        print("\nStarting manually restored run.\n")
+        restored_trainer = DataParallelTrainer.restore(
+            path=str(URI(run_config.storage_path) / exp_name),
+            storage_filesystem=storage_filesystem,
+        )
+        result = restored_trainer.fit()
+
+        _resume_from_checkpoint(
+            result.checkpoint,
+            expected_state={"iter": TestConstants.NUM_ITERATIONS - 1},
+        )
+
+        local_inspect_dir, storage_fs_path = _get_local_inspect_dir(
+            root_local_path=tmp_path,
+            storage_path=run_config.storage_path,
+            storage_filesystem=storage_filesystem,
+        )
+
+    # First, inspect that the result object returns the correct paths.
+    print(result)
+    trial_fs_path = result.path
+    assert trial_fs_path.startswith(storage_fs_path)
+    for checkpoint, _ in result.best_checkpoints:
+        assert checkpoint.path.startswith(trial_fs_path)
+
+    _assert_storage_contents(
+        local_inspect_dir,
+        exp_name,
+        checkpoint_config,
+        trainable_name="DataParallelTrainer",
+        test_trainer=True,
+        no_checkpoint_ranks=no_checkpoint_ranks,
+    )
+
+
+if __name__ == "__main__":
+    import sys
+
+    sys.exit(pytest.main(["-v", __file__]))

python/ray/train/tests/test_result.py

@@ -0,0 +1,154 @@
+import pyarrow
+import pytest
+
+import ray
+from ray import train
+from ray.air._internal.uri_utils import URI
+from ray.air.constants import EXPR_RESULT_FILE
+from ray.train import CheckpointConfig, Result, RunConfig, ScalingConfig
+from ray.train.base_trainer import TrainingFailedError
+from ray.train.tests.util import create_dict_checkpoint, load_dict_checkpoint
+from ray.train.torch import TorchTrainer
+from ray.tune import TuneConfig, Tuner
+
+_PARAM_SPACE = {"a": 1, "b": 2}
+
+
+@pytest.fixture
+def ray_start_4_cpus():
+    address_info = ray.init(num_cpus=4)
+    yield address_info
+    # The code after the yield will run as teardown code.
+    ray.shutdown()
+
+
+def build_dummy_trainer(configs):
+    def worker_loop(_config):
+        for i in range(configs["NUM_ITERATIONS"]):
+            # Do some random reports in between checkpoints.
+            train.report({"metric_a": -100, "metric_b": -100})
+
+            if ray.train.get_context().get_world_rank() == 0:
+                with create_dict_checkpoint({"iter": i}) as checkpoint:
+                    train.report(
+                        metrics={"metric_a": i, "metric_b": -i},
+                        checkpoint=checkpoint,
+                    )
+            else:
+                train.report(metrics={"metric_a": i, "metric_b": -i})
+        raise RuntimeError()
+
+    trainer = TorchTrainer(
+        train_loop_per_worker=worker_loop,
+        train_loop_config=_PARAM_SPACE,
+        scaling_config=ScalingConfig(num_workers=2, use_gpu=False),
+        run_config=RunConfig(
+            name=configs["EXP_NAME"],
+            storage_path=configs["STORAGE_PATH"],
+            checkpoint_config=CheckpointConfig(
+                num_to_keep=configs["NUM_CHECKPOINTS"],
+                checkpoint_score_attribute="metric_a",
+                checkpoint_score_order="max",
+            ),
+        ),
+    )
+    return trainer
+
+
+def build_dummy_tuner(configs):
+    return Tuner(
+        build_dummy_trainer(configs),
+        param_space={"train_loop_config": _PARAM_SPACE},
+        tune_config=TuneConfig(num_samples=1),
+    )
+
+
+@pytest.mark.parametrize("storage", ["local", "remote"])
+@pytest.mark.parametrize("mode", ["trainer", "tuner"])
+def test_result_restore(ray_start_4_cpus, tmpdir, mock_s3_bucket_uri, storage, mode):
+    NUM_ITERATIONS = 5
+    NUM_CHECKPOINTS = 3
+    if storage == "local":
+        storage_path = str(tmpdir)
+    elif storage == "remote":
+        storage_path = str(URI(mock_s3_bucket_uri))
+
+    exp_name = "test_result_restore"
+
+    configs = {
+        "EXP_NAME": exp_name,
+        "STORAGE_PATH": storage_path,
+        "NUM_ITERATIONS": NUM_ITERATIONS,
+        "NUM_CHECKPOINTS": NUM_CHECKPOINTS,
+    }
+
+    if mode == "trainer":
+        trainer = build_dummy_trainer(configs)
+        with pytest.raises(TrainingFailedError):
+            trainer.fit()
+    elif mode == "tuner":
+        tuner = build_dummy_tuner(configs)
+        tuner.fit()
+
+    # Find the trial directory to restore
+    exp_dir = str(URI(storage_path) / exp_name)
+    fs, fs_exp_dir = pyarrow.fs.FileSystem.from_uri(exp_dir)
+    for item in fs.get_file_info(pyarrow.fs.FileSelector(fs_exp_dir)):
+        if item.type == pyarrow.fs.FileType.Directory and item.base_name.startswith(
+            "TorchTrainer"
+        ):
+            trial_dir = str(URI(exp_dir) / item.base_name)
+            break
+
+    # [1] Restore from path
+    result = Result.from_path(trial_dir)
+
+    # Check if we restored all checkpoints
+    assert result.checkpoint
+    assert len(result.best_checkpoints) == NUM_CHECKPOINTS
+
+    """
+    Top-3 checkpoints with metrics:
+
+                        | iter   | metric_a    metric_b
+    checkpoint_000004        4            4          -4
+    checkpoint_000003        3            3          -3
+    checkpoint_000002        2            2          -2
+    """
+    # Check if the checkpoints bounded with correct metrics
+    best_ckpt_a = result.get_best_checkpoint(metric="metric_a", mode="max")
+    assert load_dict_checkpoint(best_ckpt_a)["iter"] == NUM_ITERATIONS - 1
+
+    best_ckpt_b = result.get_best_checkpoint(metric="metric_b", mode="max")
+    assert load_dict_checkpoint(best_ckpt_b)["iter"] == NUM_ITERATIONS - NUM_CHECKPOINTS
+
+    with pytest.raises(RuntimeError, match="Invalid metric name.*"):
+        result.get_best_checkpoint(metric="invalid_metric", mode="max")
+
+    # Check if we properly restored errors
+    assert isinstance(result.error, RuntimeError)
+
+    # Check that the config is properly formatted in the result metrics
+    assert result.metrics.get("config") == {"train_loop_config": _PARAM_SPACE}
+
+    # [2] Restore from path without result.json
+    fs.delete_file((URI(trial_dir) / EXPR_RESULT_FILE).path)
+    result = Result.from_path(trial_dir)
+
+    # Do the same checks as above
+    assert result.checkpoint
+    assert len(result.best_checkpoints) == NUM_CHECKPOINTS
+
+    best_ckpt_a = result.get_best_checkpoint(metric="metric_a", mode="max")
+    assert load_dict_checkpoint(best_ckpt_a)["iter"] == NUM_ITERATIONS - 1
+
+    best_ckpt_b = result.get_best_checkpoint(metric="metric_b", mode="max")
+    assert load_dict_checkpoint(best_ckpt_b)["iter"] == NUM_ITERATIONS - NUM_CHECKPOINTS
+
+    assert isinstance(result.error, RuntimeError)
+
+
+if __name__ == "__main__":
+    import sys
+
+    sys.exit(pytest.main(["-v", "-x", __file__]))

python/ray/train/tests/test_session.py

@@ -0,0 +1,437 @@
+import tempfile
+import time
+import warnings
+
+import pytest
+
+import ray
+from ray.air._internal.util import StartTraceback
+from ray.air.constants import SESSION_MISUSE_LOG_ONCE_KEY
+from ray.air.session import (
+    get_checkpoint,
+    get_dataset_shard,
+    get_local_rank,
+    get_world_rank,
+    get_world_size,
+    report,
+)
+from ray.train._internal.accelerator import Accelerator
+from ray.train._internal.session import (
+    _TrainingResult,
+    get_accelerator,
+    get_session,
+    init_session,
+    set_accelerator,
+    shutdown_session,
+)
+from ray.train._internal.storage import StorageContext
+from ray.train.error import SessionMisuseError
+from ray.train.tests.util import create_dict_checkpoint, load_dict_checkpoint
+
+storage = StorageContext(
+    storage_path=tempfile.mkdtemp(),
+    experiment_dir_name="exp_name",
+    trial_dir_name="trial_name",
+)
+
+
+@pytest.fixture(autouse=True, scope="module")
+def ray_start_4_cpus():
+    ray.init(num_cpus=4)
+    yield
+    ray.shutdown()
+
+
+@pytest.fixture(scope="function")
+def session():
+    def f():
+        return 1
+
+    init_session(
+        training_func=f,
+        world_rank=0,
+        local_rank=0,
+        node_rank=0,
+        local_world_size=1,
+        world_size=1,
+        storage=storage,
+    )
+    yield get_session()
+    shutdown_session()
+
+
+@pytest.fixture(autouse=True)
+def shutdown():
+    if get_session():
+        shutdown_session()
+
+
+def test_init_fail(session):
+    with pytest.raises(ValueError):
+        init_session(lambda: 1, 0)
+
+
+def test_shutdown(session):
+    shutdown_session()
+    assert not get_session()
+
+
+def test_world_rank(session):
+    assert get_world_rank() == 0
+    shutdown_session()
+    # Make sure default to 0.
+    assert get_world_rank() == 0
+
+
+def test_local_rank(session):
+    assert get_local_rank() == 0
+    shutdown_session()
+    # Make sure default to 0.
+    assert get_local_rank() == 0
+
+
+def test_world_size(session):
+    assert get_world_size() == 1
+    shutdown_session()
+    # Make sure default to 1.
+    assert get_world_size() == 1
+
+
+def test_train(session):
+    session.start()
+    session.finish()
+
+
+def test_get_dataset_shard():
+    dataset = ray.data.from_items([1, 2, 3])
+    init_session(
+        training_func=lambda: 1,
+        world_rank=0,
+        local_rank=0,
+        node_rank=0,
+        local_world_size=1,
+        world_size=1,
+        dataset_shard=dataset,
+        storage=storage,
+    )
+    assert get_dataset_shard() == dataset
+    shutdown_session()
+
+
+def test_report():
+    def train_func():
+        for i in range(2):
+            report(dict(loss=i))
+
+    init_session(
+        training_func=train_func,
+        world_rank=0,
+        local_rank=0,
+        node_rank=0,
+        local_world_size=1,
+        world_size=1,
+        storage=storage,
+    )
+    session = get_session()
+    session.start()
+    assert session.get_next().metrics["loss"] == 0
+    assert session.get_next().metrics["loss"] == 1
+    shutdown_session()
+
+
+def test_report_fail():
+    def train_func():
+        for i in range(2):
+            report(i)
+        return 1
+
+    init_session(
+        training_func=train_func,
+        world_rank=0,
+        local_rank=0,
+        node_rank=0,
+        local_world_size=1,
+        world_size=1,
+        storage=storage,
+    )
+    session = get_session()
+    session.start()
+    with pytest.raises(StartTraceback):
+        session.get_next()
+    shutdown_session()
+
+
+def test_report_after_finish(session):
+    session.start()
+    session.pause_reporting()
+    session.finish()
+    for _ in range(2):
+        report(dict(loss=1))
+    assert session.get_next() is None
+    shutdown_session()
+
+
+@pytest.mark.parametrize(
+    "block,put_result_queue,put_actor_queue",
+    [
+        (False, False, False),
+        (False, False, True),
+        (False, True, False),
+        (True, False, False),
+        (True, False, True),
+        (True, True, False),
+    ],
+)
+def test_get_result_from_queues(session, block, put_result_queue, put_actor_queue):
+    """Verify that we get the expected _TrainingResult from each result queue.
+
+    `block` describes whether we wait for a result or return after a timeout.
+        This argument should have no impact on this unit test.
+    `put_result_queue` and `put_actor_queue` are mutually exclusive and describe
+        which queue has results to process. The returned _TrainingResult should be
+        from the expected queue.
+    """
+    result_queue_training_result = _TrainingResult(
+        checkpoint=None,
+        metrics={"result_queue_metric_key": "result_queue_metric_value"},
+    )
+    if put_result_queue:
+        session.result_queue.put(result_queue_training_result, block=True)
+    inter_actor_result = {"inter_actor_metric_key": "inter_actor_metric_value"}
+    if put_actor_queue:
+        session._get_or_create_inter_actor_queue().put(inter_actor_result, block=True)
+    result = session._get_result_from_queues(block=block)
+    if put_result_queue:
+        assert result == result_queue_training_result
+    elif put_actor_queue:
+        assert (
+            result.metrics["inter_actor_metric_key"]
+            == inter_actor_result["inter_actor_metric_key"]
+        )
+    else:
+        assert result is None
+
+
+def test_no_start(session):
+    with pytest.raises(RuntimeError):
+        session.get_next()
+    shutdown_session()
+
+
+def test_checkpoint():
+    def train_func():
+        for i in range(2):
+            with create_dict_checkpoint(dict(epoch=i)) as checkpoint:
+                report({}, checkpoint=checkpoint)
+
+    def validate_zero(expected):
+        next = session.get_next()
+        assert next is not None and next.checkpoint is not None
+        assert load_dict_checkpoint(next.checkpoint)["epoch"] == expected
+
+    init_session(
+        training_func=train_func,
+        world_rank=0,
+        local_rank=0,
+        node_rank=0,
+        local_world_size=1,
+        world_size=1,
+        storage=storage,
+    )
+    session = get_session()
+    session.start()
+    validate_zero(0)
+    validate_zero(1)
+    session.finish()
+    shutdown_session()
+
+
+def test_load_checkpoint_after_save():
+    def train_func():
+        for i in range(2):
+            with create_dict_checkpoint(dict(epoch=i)) as checkpoint:
+                report(dict(epoch=i), checkpoint=checkpoint)
+            checkpoint = get_checkpoint()
+            assert load_dict_checkpoint(checkpoint)["epoch"] == i
+
+    init_session(
+        training_func=train_func,
+        world_rank=0,
+        local_rank=0,
+        node_rank=0,
+        local_world_size=1,
+        world_size=1,
+        storage=storage,
+    )
+    session = get_session()
+    session.start()
+    for i in range(2):
+        session.get_next()
+    session.finish()
+    shutdown_session()
+
+
+def test_locking():
+    """Tests that report pauses training until fetch_next or finish."""
+
+    def train_1():
+        import _thread
+
+        _thread.interrupt_main()
+
+    init_session(
+        training_func=train_1,
+        world_rank=0,
+        local_rank=0,
+        node_rank=0,
+        local_world_size=1,
+        world_size=1,
+        storage=storage,
+    )
+    session = get_session()
+    with pytest.raises(KeyboardInterrupt):
+        session.start()
+    shutdown_session()
+
+    def train_2():
+        for i in range(2):
+            report(dict(loss=i))
+        train_1()
+
+    init_session(
+        training_func=train_2,
+        world_rank=0,
+        local_rank=0,
+        node_rank=0,
+        local_world_size=1,
+        world_size=1,
+        storage=storage,
+    )
+    session = get_session()
+    session.start()
+    time.sleep(3)
+
+    session.pause_reporting()
+    # Releases session.continue_lock to resume the training thread.
+    session.get_next()
+
+    with pytest.raises(KeyboardInterrupt):
+        session.get_next()
+    session.finish()
+    shutdown_session()
+
+
+def reset_log_once_with_str(str_to_append=None):
+    key = SESSION_MISUSE_LOG_ONCE_KEY
+    if str_to_append:
+        key += f"-{str_to_append}"
+    ray.util.debug.reset_log_once(key)
+
+
+@pytest.mark.parametrize("fn", [get_checkpoint, get_dataset_shard])
+def test_warn(fn):
+    """Checks if calling session functions outside of session raises warning."""
+
+    with warnings.catch_warnings(record=True) as record:
+        warnings.simplefilter("always")
+        # Ignore Deprecation warnings.
+        warnings.filterwarnings("ignore", category=DeprecationWarning)
+        assert not fn()
+
+    assert fn.__name__ in record[0].message.args[0]
+
+    reset_log_once_with_str(fn.__name__)
+
+
+def test_warn_report():
+    """Checks if calling session.report function outside of session raises warning."""
+
+    fn = report
+
+    with warnings.catch_warnings(record=True) as record:
+        warnings.simplefilter("always")
+        # Ignore Deprecation warnings.
+        warnings.filterwarnings("ignore", category=DeprecationWarning)
+        assert not fn(dict())
+
+    assert fn.__name__ in record[0].message.args[0]
+
+    reset_log_once_with_str(fn.__name__)
+
+
+def test_warn_once():
+    """Checks if session misuse warning is only shown once per function."""
+
+    with warnings.catch_warnings(record=True) as record:
+        # Ignore Deprecation warnings.
+        warnings.simplefilter("always")
+        warnings.filterwarnings("ignore", category=DeprecationWarning)
+        assert not get_checkpoint()
+        assert not get_checkpoint()
+        assert not report(dict(x=2))
+        assert not report(dict(x=2))
+        assert not get_dataset_shard()
+        assert not get_dataset_shard()
+
+    # Should only warn once.
+    assert len(record) == 3
+
+
+class FakeAccelerator(Accelerator):
+    pass
+
+
+def test_set_and_get_accelerator(session):
+    accelerator = FakeAccelerator()
+    set_accelerator(accelerator)
+    assert get_accelerator(FakeAccelerator) is accelerator
+
+
+def test_get_accelerator_constructs_default_accelerator(session):
+    assert isinstance(get_accelerator(FakeAccelerator), FakeAccelerator)
+
+
+def test_get_accelerator_raises_error_outside_session():
+    with pytest.raises(SessionMisuseError):
+        get_accelerator(FakeAccelerator)
+
+
+def test_set_accelerator_raises_error_if_accelerator_already_set(session):
+    accelerator1, accelerator2 = FakeAccelerator(), FakeAccelerator()
+    set_accelerator(accelerator1)
+    with pytest.raises(RuntimeError):
+        set_accelerator(accelerator2)
+
+
+def test_set_accelerator_raises_error_outside_session():
+    accelerator = FakeAccelerator()
+    with pytest.raises(SessionMisuseError):
+        set_accelerator(accelerator)
+
+
+def test_application_error_raised():
+    def f():
+        raise ValueError
+
+    init_session(
+        training_func=f,
+        world_rank=0,
+        local_rank=0,
+        node_rank=0,
+        local_world_size=1,
+        world_size=1,
+        storage=storage,
+    )
+    session = get_session()
+    session.start()
+    with pytest.raises(StartTraceback):
+        session.get_next()
+    shutdown_session()
+
+
+if __name__ == "__main__":
+    import sys
+
+    import pytest
+
+    sys.exit(pytest.main(["-v", "-x", __file__]))

python/ray/train/tests/test_state.py

@@ -0,0 +1,351 @@
+import json
+import os
+import time
+
+import pytest
+
+import ray
+from ray.cluster_utils import Cluster
+from ray.train import RunConfig, ScalingConfig
+from ray.train._internal.state.schema import (
+    ActorStatusEnum,
+    RunStatusEnum,
+    TrainDatasetInfo,
+    TrainRunInfo,
+    TrainWorkerInfo,
+)
+from ray.train._internal.state.state_actor import (
+    TRAIN_STATE_ACTOR_NAME,
+    TRAIN_STATE_ACTOR_NAMESPACE,
+    get_or_create_state_actor,
+)
+from ray.train._internal.state.state_manager import TrainRunStateManager
+from ray.train._internal.worker_group import WorkerGroup
+from ray.train.data_parallel_trainer import DataParallelTrainer
+
+
+@pytest.fixture
+def ray_start_gpu_cluster():
+    cluster = Cluster()
+    cluster.add_node(num_gpus=8, num_cpus=9)
+
+    ray.shutdown()
+    ray.init(
+        address=cluster.address,
+        runtime_env={"env_vars": {"RAY_TRAIN_ENABLE_STATE_TRACKING": "1"}},
+        ignore_reinit_error=True,
+    )
+
+    yield
+
+    ray.shutdown()
+    cluster.shutdown()
+
+
+RUN_INFO_JSON_SAMPLE = """{
+    "name": "default_run",
+    "id": "ad5256bc64c04c83833a8b006f531799",
+    "job_id": "0000000001",
+    "controller_actor_id": "3abd1972a19148d78acc78dd9414736e",
+    "start_time_ms": 1717448423000,
+    "run_status": "RUNNING",
+    "status_detail": "",
+    "end_time_ms": null,
+    "resources": [{"CPU": 1}, {"CPU": 1}],
+    "workers": [
+        {
+        "actor_id": "3d86c25634a71832dac32c8802000000",
+        "world_rank": 0,
+        "local_rank": 0,
+        "node_rank": 0,
+        "node_id": "b1e6cbed8533ae2def4e7e7ced9d19858ceb1ed8ab9ba81ab9c07825",
+        "node_ip": "10.0.208.100",
+        "pid": 76071,
+        "gpu_ids": [0],
+        "status": "ALIVE",
+        "resources": {"CPU": 1}
+        },
+        {
+        "actor_id": "8f162dd8365346d1b5c98ebd7338c4f9",
+        "world_rank": 1,
+        "local_rank": 1,
+        "node_rank": 0,
+        "node_id": "b1e6cbed8533ae2def4e7e7ced9d19858ceb1ed8ab9ba81ab9c07825",
+        "node_ip": "10.0.208.100",
+        "pid": 76072,
+        "gpu_ids": [1],
+        "status": "ALIVE",
+        "resources": {"CPU": 1}
+        }
+    ],
+    "datasets": [
+        {
+        "name": "train",
+        "dataset_name": "train_dataset",
+        "dataset_uuid": "1"
+        }
+    ]
+}"""
+
+
+def _get_run_info_sample(run_id=None, run_name=None) -> TrainRunInfo:
+    dataset_info = TrainDatasetInfo(
+        name="train", dataset_name="train_dataset", dataset_uuid="1"
+    )
+
+    worker_info_0 = TrainWorkerInfo(
+        actor_id="3d86c25634a71832dac32c8802000000",
+        world_rank=0,
+        local_rank=0,
+        node_rank=0,
+        node_id="b1e6cbed8533ae2def4e7e7ced9d19858ceb1ed8ab9ba81ab9c07825",
+        node_ip="10.0.208.100",
+        pid=76071,
+        gpu_ids=[0],
+        status=ActorStatusEnum.ALIVE,
+        resources={"CPU": 1},
+    )
+
+    worker_info_1 = TrainWorkerInfo(
+        actor_id="8f162dd8365346d1b5c98ebd7338c4f9",
+        world_rank=1,
+        local_rank=1,
+        node_rank=0,
+        node_id="b1e6cbed8533ae2def4e7e7ced9d19858ceb1ed8ab9ba81ab9c07825",
+        node_ip="10.0.208.100",
+        pid=76072,
+        gpu_ids=[1],
+        status=ActorStatusEnum.ALIVE,
+        resources={"CPU": 1},
+    )
+
+    run_info = TrainRunInfo(
+        name=run_name if run_name else "default_run",
+        id=run_id if run_id else "ad5256bc64c04c83833a8b006f531799",
+        job_id="0000000001",
+        controller_actor_id="3abd1972a19148d78acc78dd9414736e",
+        workers=[worker_info_0, worker_info_1],
+        datasets=[dataset_info],
+        start_time_ms=1717448423000,
+        run_status=RunStatusEnum.RUNNING,
+        status_detail="",
+        resources=[{"CPU": 1}, {"CPU": 1}],
+    )
+    return run_info
+
+
+def test_schema_equivalance():
+    json_sample = RUN_INFO_JSON_SAMPLE
+    dict_sample = json.loads(RUN_INFO_JSON_SAMPLE)
+
+    run_info_from_json = TrainRunInfo.parse_raw(json_sample)
+    run_info_from_obj = TrainRunInfo.parse_obj(dict_sample)
+
+    # Test serialization equivalence
+    assert run_info_from_json == run_info_from_obj
+
+    # Test dict deserialization equivalence
+    assert run_info_from_json.dict() == dict_sample
+
+    # Test json deserialization equivalence
+    assert json.loads(run_info_from_json.json()) == json.loads(json_sample)
+
+    # Test constructors equivalence
+    assert _get_run_info_sample() == run_info_from_json
+
+
+def test_state_actor_api(ray_start_4_cpus):
+    state_actor = get_or_create_state_actor()
+    named_actors = ray.util.list_named_actors(all_namespaces=True)
+    assert {
+        "name": TRAIN_STATE_ACTOR_NAME,
+        "namespace": TRAIN_STATE_ACTOR_NAMESPACE,
+    } in named_actors
+
+    # Concurrently register 100 runs
+    num_runs = 100
+    info_list = [_get_run_info_sample(run_id=str(i)) for i in range(num_runs)]
+    ray.get([state_actor.register_train_run.remote(run) for run in info_list])
+
+    # Test get all runs
+    train_runs = ray.get(state_actor.get_all_train_runs.remote())
+    assert len(train_runs) == num_runs
+
+    # Test get a single run by run_id
+    for i in range(num_runs):
+        run_info = ray.get(state_actor.get_train_run.remote(run_id=str(i)))
+        assert run_info == info_list[i]
+
+
+def test_state_manager(ray_start_gpu_cluster):
+    worker_group = WorkerGroup(num_workers=4, resources_per_worker={"GPU": 1})
+
+    # No errors raised if TrainStateActor is not started
+    state_manager = TrainRunStateManager(state_actor=None)
+    state_manager.register_train_run(
+        run_id="run_id",
+        run_name="run_name",
+        job_id="0000000001",
+        controller_actor_id="3abd1972a19148d78acc78dd9414736e",
+        datasets={},
+        worker_group=worker_group,
+        start_time_ms=int(time.time() * 1000),
+        run_status=RunStatusEnum.RUNNING,
+        resources=[{"CPU": 1}, {"CPU": 1}],
+    )
+
+    # Register 100 runs with 10 TrainRunStateManagers
+    state_actor = get_or_create_state_actor()
+    for i in range(10):
+        state_manager = TrainRunStateManager(state_actor=state_actor)
+        for j in range(10):
+            run_id = i * 10 + j
+            state_manager.register_train_run(
+                run_id=str(run_id),
+                run_name="run_name",
+                job_id="0000000001",
+                controller_actor_id="3abd1972a19148d78acc78dd9414736e",
+                datasets={
+                    "train": ray.data.from_items(list(range(4))),
+                    "eval": ray.data.from_items(list(range(4))),
+                },
+                worker_group=worker_group,
+                start_time_ms=int(time.time() * 1000),
+                run_status=RunStatusEnum.RUNNING,
+                resources=[{"CPU": 1}, {"CPU": 1}],
+            )
+
+    runs = ray.get(state_actor.get_all_train_runs.remote())
+    assert len(runs) == 100
+
+    for i in range(100):
+        run_id = str(i)
+        run_info = ray.get(state_actor.get_train_run.remote(run_id=run_id))
+        assert run_info and run_info.id == run_id
+
+
+@pytest.mark.parametrize("gpus_per_worker", [0, 1, 2])
+def test_track_e2e_training(ray_start_gpu_cluster, gpus_per_worker):
+    os.environ["RAY_TRAIN_ENABLE_STATE_TRACKING"] = "1"
+    num_workers = 4
+    run_name = "test"
+    datasets = {
+        "train": ray.data.from_items(list(range(4))),
+        "eval": ray.data.from_items(list(range(4))),
+    }
+
+    if gpus_per_worker == 0:
+        use_gpu = False
+        resources_per_worker = {"CPU": 1}
+    else:
+        use_gpu = True
+        resources_per_worker = {"GPU": gpus_per_worker}
+
+    trainer = DataParallelTrainer(
+        train_loop_per_worker=lambda: None,
+        run_config=RunConfig(name=run_name),
+        scaling_config=ScalingConfig(
+            num_workers=num_workers,
+            use_gpu=use_gpu,
+            resources_per_worker=resources_per_worker,
+        ),
+        datasets=datasets,
+    )
+
+    trainer.fit()
+
+    state_actor = ray.get_actor(
+        name=TRAIN_STATE_ACTOR_NAME, namespace=TRAIN_STATE_ACTOR_NAMESPACE
+    )
+
+    runs = ray.get(state_actor.get_all_train_runs.remote())
+    run_id = next(iter(runs.keys()))
+    run = next(iter(runs.values()))
+
+    # Check Run Info
+    assert run.id == run_id
+    assert run.name == run_name
+    assert len(run.workers) == num_workers
+    assert run.controller_actor_id and run.job_id
+
+    world_ranks = [worker.world_rank for worker in run.workers]
+    local_ranks = [worker.local_rank for worker in run.workers]
+    node_ranks = [worker.node_rank for worker in run.workers]
+
+    # Ensure that the workers are sorted by global rank
+    assert world_ranks == [0, 1, 2, 3]
+    assert local_ranks == [0, 1, 2, 3]
+    assert node_ranks == [0, 0, 0, 0]
+
+    # Check GPU ids
+    gpu_ids = [worker.gpu_ids for worker in run.workers]
+    if gpus_per_worker == 0:
+        assert gpu_ids == [[], [], [], []]
+    elif gpus_per_worker == 1:
+        assert gpu_ids == [[0], [1], [2], [3]]
+    elif gpus_per_worker == 2:
+        flat_gpu_ids = set()
+        for ids in gpu_ids:
+            flat_gpu_ids.update(ids)
+        assert flat_gpu_ids == set(range(8))
+
+    # Check Datasets
+    for dataset_info in run.datasets:
+        dataset = datasets[dataset_info.name]
+        # DataConfig will automatically set the dataset_name to the key of the dataset dict.
+        assert dataset_info.dataset_name == dataset_info.name
+        assert dataset_info.dataset_uuid == dataset._uuid
+
+
+@pytest.mark.parametrize("raise_error", [True, False])
+def test_train_run_status(ray_start_gpu_cluster, raise_error):
+    os.environ["RAY_TRAIN_ENABLE_STATE_TRACKING"] = "1"
+
+    def get_train_run():
+        state_actor = ray.get_actor(
+            name=TRAIN_STATE_ACTOR_NAME, namespace=TRAIN_STATE_ACTOR_NAMESPACE
+        )
+        runs = ray.get(state_actor.get_all_train_runs.remote())
+        return next(iter(runs.values()))
+
+    def check_run_status(expected_status):
+        run = get_train_run()
+        assert run.run_status == expected_status
+
+    def check_run_error(failed_rank, error_message):
+        run = get_train_run()
+        assert run.status_detail
+        assert f"Rank {failed_rank} worker raised an error" in run.status_detail
+        assert error_message in run.status_detail
+
+    failed_rank = 0
+    error_message = "User Application Error"
+
+    def train_func():
+        check_run_status(expected_status=RunStatusEnum.RUNNING)
+        if raise_error and ray.train.get_context().get_world_rank() == failed_rank:
+            raise RuntimeError(error_message)
+
+    trainer = DataParallelTrainer(
+        train_loop_per_worker=train_func,
+        scaling_config=ScalingConfig(num_workers=4, use_gpu=False),
+    )
+
+    try:
+        trainer.fit()
+    except Exception:
+        pass
+
+    if raise_error:
+        check_run_status(expected_status=RunStatusEnum.ERRORED)
+        check_run_error(failed_rank=failed_rank, error_message=error_message)
+    else:
+        check_run_status(expected_status=RunStatusEnum.FINISHED)
+
+    ray.shutdown()
+
+
+if __name__ == "__main__":
+    import sys
+
+    sys.exit(pytest.main(["-v", "-x", __file__]))

python/ray/train/tests/test_state_export.py

@@ -0,0 +1,126 @@
+import uuid
+
+import pytest
+
+from ray.core.generated.export_train_state_pb2 import (
+    ExportTrainRunAttemptEventData as ProtoTrainRunAttempt,
+    ExportTrainRunEventData as ProtoTrainRun,
+)
+from ray.train._internal.state.export import (
+    train_run_info_to_proto_attempt,
+    train_run_info_to_proto_run,
+)
+from ray.train._internal.state.schema import (
+    ActorStatusEnum,
+    RunStatusEnum,
+    TrainRunInfo,
+    TrainWorkerInfo,
+)
+
+
+def create_mock_train_run_info(run_status=RunStatusEnum.RUNNING):
+    """Create a minimal mock TrainRunInfo."""
+    worker = TrainWorkerInfo(
+        actor_id=uuid.uuid4().hex,
+        world_rank=0,
+        local_rank=0,
+        node_rank=0,
+        node_id=uuid.uuid4().hex,
+        node_ip="127.0.0.1",
+        pid=1234,
+        gpu_ids=[0],
+        status=ActorStatusEnum.ALIVE,
+        resources={"CPU": 1},
+    )
+
+    return TrainRunInfo(
+        name="test_run",
+        id=uuid.uuid4().hex,
+        job_id=uuid.uuid4().hex,
+        controller_actor_id=uuid.uuid4().hex,
+        workers=[worker],
+        datasets=[],
+        run_status=run_status,
+        status_detail="Error details" if run_status == RunStatusEnum.ERRORED else "",
+        start_time_ms=1000,
+        end_time_ms=2000
+        if run_status in [RunStatusEnum.FINISHED, RunStatusEnum.ERRORED]
+        else None,
+        resources=[{"CPU": 1}],
+    )
+
+
+def test_train_run_info_to_proto_run():
+    """Test that run info is correctly exported to proto run."""
+    run_info = create_mock_train_run_info()
+    proto_run = train_run_info_to_proto_run(run_info)
+
+    assert proto_run.ray_train_version == 1
+    assert proto_run.id == run_info.id
+    assert proto_run.name == run_info.name
+    assert proto_run.job_id == bytes.fromhex(run_info.job_id)
+    assert proto_run.status == ProtoTrainRun.RunStatus.RUNNING
+
+
+def test_train_run_info_to_proto_attempt():
+    """Test that run info is correctly exported to proto attempt."""
+    run_info = create_mock_train_run_info()
+    proto_attempt = train_run_info_to_proto_attempt(run_info)
+
+    assert proto_attempt.ray_train_version == 1
+    assert proto_attempt.run_id == run_info.id
+    assert proto_attempt.status == ProtoTrainRunAttempt.RunAttemptStatus.RUNNING
+
+    # Verify worker fields
+    assert len(proto_attempt.workers) == 1
+    proto_worker = proto_attempt.workers[0]
+    worker_info = run_info.workers[0]
+
+    assert proto_worker.world_rank == worker_info.world_rank
+    assert proto_worker.actor_id == bytes.fromhex(worker_info.actor_id)
+
+
+def test_status_mapping():
+    """Test that status is correctly mapped between schemas."""
+    status_pairs = [
+        (
+            RunStatusEnum.STARTED,
+            ProtoTrainRun.RunStatus.INITIALIZING,
+            ProtoTrainRunAttempt.RunAttemptStatus.PENDING,
+        ),
+        (
+            RunStatusEnum.RUNNING,
+            ProtoTrainRun.RunStatus.RUNNING,
+            ProtoTrainRunAttempt.RunAttemptStatus.RUNNING,
+        ),
+        (
+            RunStatusEnum.FINISHED,
+            ProtoTrainRun.RunStatus.FINISHED,
+            ProtoTrainRunAttempt.RunAttemptStatus.FINISHED,
+        ),
+        (
+            RunStatusEnum.ERRORED,
+            ProtoTrainRun.RunStatus.ERRORED,
+            ProtoTrainRunAttempt.RunAttemptStatus.ERRORED,
+        ),
+        (
+            RunStatusEnum.ABORTED,
+            ProtoTrainRun.RunStatus.ABORTED,
+            ProtoTrainRunAttempt.RunAttemptStatus.ABORTED,
+        ),
+    ]
+
+    for run_status, expected_run_status, expected_attempt_status in status_pairs:
+        run_info = create_mock_train_run_info(run_status=run_status)
+
+        proto_run = train_run_info_to_proto_run(run_info)
+        assert proto_run.status == expected_run_status
+
+        proto_attempt = train_run_info_to_proto_attempt(run_info)
+        assert proto_attempt.status == expected_attempt_status
+
+
+if __name__ == "__main__":
+    import sys
+
+    sys.exit(pytest.main(["-v", "-x", __file__]))

python/ray/train/tests/test_storage.py

@@ -0,0 +1,207 @@
+import os
+import uuid
+from pathlib import Path
+
+import pyarrow.fs
+import pytest
+
+import ray
+import ray.cloudpickle as ray_pickle
+from ray.train import Checkpoint, SyncConfig
+from ray.train._internal.storage import (
+    _VALIDATE_STORAGE_MARKER_FILENAME,
+    StorageContext,
+    _list_at_fs_path,
+)
+from ray.train.tests.test_new_persistence import _resolve_storage_type
+
+
+@pytest.fixture(params=["nfs", "cloud", "custom_fs"])
+def storage(request, tmp_path) -> StorageContext:
+    storage_type = request.param
+    with _resolve_storage_type(storage_type, tmp_path) as (
+        storage_path,
+        storage_filesystem,
+    ):
+        yield StorageContext(
+            storage_path=storage_path,
+            experiment_dir_name=f"storage_type={storage_type}-{uuid.uuid4().hex}",
+            storage_filesystem=storage_filesystem,
+            trial_dir_name="trial_name",
+            sync_config=SyncConfig(
+                sync_artifacts=True, sync_artifacts_on_checkpoint=True, sync_period=1000
+            ),
+        )
+
+
+@pytest.fixture(autouse=True, scope="module")
+def ray_init():
+    # NOTE: This is needed to set the `/tmp/ray/session_*` directory.
+    ray.init()
+    yield
+    ray.shutdown()
+
+
+def test_custom_fs_validation(tmp_path):
+    """Tests that invalid storage_path inputs give reasonable errors when a
+    custom filesystem is used."""
+    exp_name = "test"
+    StorageContext(
+        storage_path=str(tmp_path),
+        experiment_dir_name=exp_name,
+        storage_filesystem=pyarrow.fs.LocalFileSystem(),
+    )
+
+    mock_fs, _ = pyarrow.fs.FileSystem.from_uri("mock://a")
+    with pytest.raises(pyarrow.lib.ArrowInvalid) as excinfo:
+        StorageContext(
+            storage_path="mock:///a",
+            experiment_dir_name=exp_name,
+            storage_filesystem=mock_fs,
+        )
+    print("Custom fs with URI storage path error: ", excinfo.value)
+
+    StorageContext(
+        storage_path="a",
+        experiment_dir_name=exp_name,
+        storage_filesystem=mock_fs,
+    )
+
+
+def test_storage_path_inputs():
+    """Tests storage path input edge cases."""
+    exp_name = "test_storage_path"
+
+    # Relative paths don't work
+    with pytest.raises(pyarrow.lib.ArrowInvalid) as excinfo:
+        StorageContext(storage_path="./results", experiment_dir_name=exp_name)
+    assert "URI has empty scheme" in str(excinfo.value)
+
+    with pytest.raises(pyarrow.lib.ArrowInvalid) as excinfo:
+        StorageContext(storage_path="results", experiment_dir_name=exp_name)
+    assert "URI has empty scheme" in str(excinfo.value)
+
+    # Tilde paths sometimes raise... They do not work on the CI machines.
+    try:
+        StorageContext(storage_path="~/ray_results", experiment_dir_name=exp_name)
+    except pyarrow.lib.ArrowInvalid as e:
+        print(e)
+        pass
+
+    # Paths with lots of extra . .. and /
+    path = os.path.expanduser("~/ray_results")
+    path = os.path.join(path, ".", "..", "ray_results", ".")
+    path = path.replace(os.path.sep, os.path.sep * 2)
+    storage = StorageContext(storage_path=path, experiment_dir_name=exp_name)
+
+    storage.storage_filesystem.create_dir(
+        os.path.join(storage.storage_fs_path, "test_dir")
+    )
+    assert Path("~/ray_results/test_dir").expanduser().exists()
+
+    # Path objects work
+    StorageContext(storage_path=Path(path), experiment_dir_name=exp_name)
+
+
+def test_storage_validation_marker(storage: StorageContext):
+    # A marker should have been created at initialization
+    storage._check_validation_file()
+
+    # Remove the marker to simulate being on a new node w/o access to the shared storage
+    storage.storage_filesystem.delete_file(
+        os.path.join(storage.experiment_fs_path, _VALIDATE_STORAGE_MARKER_FILENAME)
+    )
+
+    # Simulate passing the storage context around through the object store
+    # The constructor is NOT called again -- so the marker should not be checked here
+    # and we shouldn't raise an error
+    storage = ray_pickle.loads(ray_pickle.dumps(storage))
+
+    # We should raise an error when we try to checkpoint now.
+    with pytest.raises(RuntimeError) as excinfo:
+        storage.persist_current_checkpoint(Checkpoint.from_directory("/tmp/dummy"))
+    assert "Unable to set up cluster storage" in str(excinfo.value)
+
+
+def test_persist_current_checkpoint(storage: StorageContext, tmp_path):
+    # Uploading a non-existent checkpoint directory should fail
+    with pytest.raises(FileNotFoundError):
+        storage.persist_current_checkpoint(
+            Checkpoint.from_directory("/tmp/nonexistent/checkpoint")
+        )
+
+    # Uploading an empty checkpoint directory
+    (tmp_path / "empty").mkdir()
+    storage.persist_current_checkpoint(Checkpoint.from_directory(tmp_path / "empty"))
+    assert (
+        _list_at_fs_path(storage.storage_filesystem, storage.checkpoint_fs_path) == []
+    )
+
+    # Normal use case: Uploading an checkpoint directory with files
+    (tmp_path / "regular").mkdir()
+    (tmp_path / "regular" / "1.txt").touch()
+    storage.persist_current_checkpoint(Checkpoint.from_directory(tmp_path / "regular"))
+    assert _list_at_fs_path(storage.storage_filesystem, storage.checkpoint_fs_path) == [
+        "1.txt"
+    ]
+
+    storage.current_checkpoint_index += 1
+
+    # Persisting a checkpoint that is already at the correct path (for local fs case)
+    if isinstance(storage.storage_filesystem, pyarrow.fs.LocalFileSystem):
+        final_checkpoint_dir = Path(storage.checkpoint_fs_path)
+        final_checkpoint_dir.mkdir(parents=True)
+        (final_checkpoint_dir / "2.txt").touch()
+        storage.persist_current_checkpoint(
+            Checkpoint.from_directory(final_checkpoint_dir)
+        )
+
+
+def test_persist_artifacts(storage: StorageContext):
+    """Tests typical `StorageContext.persist_artifacts(force=True/False)` usage."""
+    trial_working_dir = Path(storage.trial_working_directory)
+    trial_working_dir.mkdir(parents=True)
+    trial_working_dir.joinpath("1.txt").touch()
+
+    storage.persist_artifacts()
+    storage.syncer.wait()
+
+    assert sorted(
+        _list_at_fs_path(storage.storage_filesystem, storage.trial_fs_path)
+    ) == ["1.txt"]
+
+    trial_working_dir.joinpath("2.txt").touch()
+
+    # A new sync should not be triggered because sync_period is 1000 seconds
+    storage.persist_artifacts()
+    storage.syncer.wait()
+
+    # -> No change in the persisted files
+    assert sorted(
+        _list_at_fs_path(storage.storage_filesystem, storage.trial_fs_path)
+    ) == ["1.txt"]
+
+    # This is what happens on `train.report` when a checkpoint is reported
+    # and `sync_artifacts_on_checkpoint=True`
+    storage.persist_artifacts(force=storage.sync_config.sync_artifacts_on_checkpoint)
+    assert sorted(
+        _list_at_fs_path(storage.storage_filesystem, storage.trial_fs_path)
+    ) == ["1.txt", "2.txt"]
+
+
+def test_persist_artifacts_failures(storage: StorageContext):
+    """Tests `StorageContext.persist_artifacts` edge cases (empty directory)."""
+    # Uploading before the trial directory has been created should fail
+    with pytest.raises(FileNotFoundError):
+        storage.persist_artifacts()
+        if storage.syncer:
+            storage.syncer.wait()
+
+    with pytest.raises(FileNotFoundError):
+        storage.persist_artifacts(force=True)
+
+
+if __name__ == "__main__":
+    import sys
+
+    sys.exit(pytest.main(["-v", "-x", __file__]))

python/ray/train/tests/test_telemetry.py

@@ -0,0 +1,50 @@
+import sys
+
+import pytest
+
+import ray
+import ray._common.usage.usage_lib as ray_usage_lib
+from ray._common.test_utils import TelemetryCallsite, check_library_usage_telemetry
+from ray.train.data_parallel_trainer import DataParallelTrainer
+
+
+@pytest.fixture
+def reset_usage_lib():
+    yield
+    ray.shutdown()
+    ray_usage_lib.reset_global_state()
+
+
+@pytest.mark.parametrize("callsite", list(TelemetryCallsite))
+def test_not_used_on_import(reset_usage_lib, callsite: TelemetryCallsite):
+    def _import_ray_train():
+        from ray import train  # noqa: F401
+
+    check_library_usage_telemetry(
+        _import_ray_train, callsite=callsite, expected_library_usages=[set(), {"core"}]
+    )
+
+
+@pytest.mark.parametrize("callsite", list(TelemetryCallsite))
+def test_used_on_train_fit(reset_usage_lib, callsite: TelemetryCallsite):
+    def _call_train_fit():
+        def train_fn():
+            pass
+
+        trainer = DataParallelTrainer(train_fn)
+        trainer.fit()
+
+    check_library_usage_telemetry(
+        _call_train_fit,
+        callsite=callsite,
+        expected_library_usages=[{"train", "tune"}, {"core", "train", "tune"}],
+        expected_extra_usage_tags={
+            "air_entrypoint": "Trainer.fit",
+            "air_storage_configuration": "local",
+            "air_trainer": "Custom",
+        },
+    )
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main(["-v", "-s", __file__]))

python/ray/train/tests/test_tensorflow_checkpoint.py

@@ -0,0 +1,151 @@
+import os
+
+os.environ["TF_USE_LEGACY_KERAS"] = "1"
+
+import os.path
+import sys
+import tempfile
+import unittest
+from typing import List
+
+import pytest
+from numpy import ndarray
+
+from ray import train
+from ray.data import Preprocessor
+from ray.train import ScalingConfig
+
+if sys.version_info >= (3, 12):
+    # Tensorflow is not installed for Python 3.12 because of keras compatibility.
+    sys.exit(0)
+else:
+    import tensorflow as tf
+
+    from ray.train.tensorflow import TensorflowCheckpoint, TensorflowTrainer
+
+
+class DummyPreprocessor(Preprocessor):
+    def __init__(self, multiplier):
+        self.multiplier = multiplier
+
+    def transform_batch(self, df):
+        return df * self.multiplier
+
+
+def get_model():
+    return tf.keras.Sequential(
+        [
+            tf.keras.layers.InputLayer(input_shape=()),
+            tf.keras.layers.Flatten(),
+            tf.keras.layers.Dense(10),
+            tf.keras.layers.Dense(1),
+        ]
+    )
+
+
+def compare_weights(w1: List[ndarray], w2: List[ndarray]) -> bool:
+    if not len(w1) == len(w2):
+        return False
+    size = len(w1)
+    for i in range(size):
+        comparison = w1[i] == w2[i]
+        if not comparison.all():
+            return False
+
+    return True
+
+
+class TestFromModel(unittest.TestCase):
+    def setUp(self):
+        self.model = get_model()
+        self.preprocessor = DummyPreprocessor(1)
+
+    def test_from_model(self):
+        checkpoint = TensorflowCheckpoint.from_model(
+            self.model, preprocessor=DummyPreprocessor(1)
+        )
+        loaded_model = checkpoint.get_model()
+        preprocessor = checkpoint.get_preprocessor()
+
+        assert compare_weights(loaded_model.get_weights(), self.model.get_weights())
+        assert preprocessor.multiplier == 1
+
+    def test_from_saved_model(self):
+        with tempfile.TemporaryDirectory() as tmp_dir:
+            model_dir_path = os.path.join(tmp_dir, "my_model")
+            self.model.save(model_dir_path, save_format="tf")
+            checkpoint = TensorflowCheckpoint.from_saved_model(
+                model_dir_path, preprocessor=DummyPreprocessor(1)
+            )
+            loaded_model = checkpoint.get_model()
+            preprocessor = checkpoint.get_preprocessor()
+            assert compare_weights(self.model.get_weights(), loaded_model.get_weights())
+            assert preprocessor.multiplier == 1
+
+    def test_from_h5_model(self):
+        with tempfile.TemporaryDirectory() as tmp_dir:
+            model_file_path = os.path.join(tmp_dir, "my_model.h5")
+            self.model.save(model_file_path)
+            checkpoint = TensorflowCheckpoint.from_h5(
+                model_file_path, preprocessor=DummyPreprocessor(1)
+            )
+            loaded_model = checkpoint.get_model()
+            preprocessor = checkpoint.get_preprocessor()
+            assert compare_weights(self.model.get_weights(), loaded_model.get_weights())
+            assert preprocessor.multiplier == 1
+
+
+def test_tensorflow_checkpoint_saved_model():
+    # The test passes if the following can run successfully.
+
+    def train_fn():
+        model = tf.keras.Sequential(
+            [
+                tf.keras.layers.InputLayer(input_shape=()),
+                tf.keras.layers.Flatten(),
+                tf.keras.layers.Dense(10),
+                tf.keras.layers.Dense(1),
+            ]
+        )
+        with tempfile.TemporaryDirectory() as tempdir:
+            model.save(tempdir)
+            checkpoint = TensorflowCheckpoint.from_saved_model(tempdir)
+            train.report({"my_metric": 1}, checkpoint=checkpoint)
+
+    trainer = TensorflowTrainer(
+        train_loop_per_worker=train_fn, scaling_config=ScalingConfig(num_workers=2)
+    )
+
+    assert trainer.fit().checkpoint
+
+
+def test_tensorflow_checkpoint_h5():
+    # The test passes if the following can run successfully.
+
+    def train_func():
+        model = tf.keras.Sequential(
+            [
+                tf.keras.layers.InputLayer(input_shape=()),
+                tf.keras.layers.Flatten(),
+                tf.keras.layers.Dense(10),
+                tf.keras.layers.Dense(1),
+            ]
+        )
+        with tempfile.TemporaryDirectory() as tempdir:
+            model.save(os.path.join(tempdir, "my_model.h5"))
+            checkpoint = TensorflowCheckpoint.from_h5(
+                os.path.join(tempdir, "my_model.h5")
+            )
+            train.report({"my_metric": 1}, checkpoint=checkpoint)
+
+    trainer = TensorflowTrainer(
+        train_loop_per_worker=train_func, scaling_config=ScalingConfig(num_workers=2)
+    )
+
+    assert trainer.fit().checkpoint
+
+
+if __name__ == "__main__":
+    import sys
+
+    sys.exit(pytest.main(["-v", "-x", __file__]))

python/ray/train/tests/test_tensorflow_trainer.py

@@ -0,0 +1,125 @@
+import os
+
+os.environ["TF_USE_LEGACY_KERAS"] = "1"
+
+import sys
+import tempfile
+
+import pytest
+
+import ray
+from ray import train
+from ray.data.preprocessors import Concatenator
+from ray.train import ScalingConfig
+from ray.train.constants import TRAIN_DATASET_KEY
+
+if sys.version_info >= (3, 12):
+    # Tensorflow is not installed for Python 3.12 because of keras compatibility.
+    sys.exit(0)
+else:
+    from ray.train.examples.tf.tensorflow_regression_example import (
+        train_func as tensorflow_linear_train_func,
+    )
+    from ray.train.tensorflow import TensorflowCheckpoint, TensorflowTrainer
+
+
+@pytest.fixture
+def ray_start_4_cpus():
+    address_info = ray.init(num_cpus=4)
+    yield address_info
+    # The code after the yield will run as teardown code.
+    ray.shutdown()
+
+
+def build_model():
+    import tensorflow as tf
+
+    model = tf.keras.Sequential(
+        [
+            tf.keras.layers.InputLayer(input_shape=()),
+            tf.keras.layers.Flatten(),
+            tf.keras.layers.Dense(1),
+        ]
+    )
+
+    return model
+
+
+@pytest.mark.parametrize("num_workers", [1, 2])
+def test_tensorflow_linear(ray_start_4_cpus, num_workers):
+    """Also tests air Keras callback."""
+    epochs = 3
+
+    def train_func(config):
+        result = tensorflow_linear_train_func(config)
+        assert len(result) == epochs
+        assert result[-1]["loss"] < result[0]["loss"]
+
+    train_loop_config = {
+        "lr": 1e-3,
+        "batch_size": 32,
+        "epochs": epochs,
+    }
+    scaling_config = ScalingConfig(num_workers=num_workers)
+    dataset = ray.data.read_csv("s3://anonymous@air-example-data/regression.csv")
+    columns_to_concatenate = [f"x{i:03}" for i in range(100)]
+    preprocessor = Concatenator(columns=columns_to_concatenate, output_column_name="x")
+    dataset = preprocessor.transform(dataset)
+
+    trainer = TensorflowTrainer(
+        train_loop_per_worker=train_func,
+        train_loop_config=train_loop_config,
+        scaling_config=scaling_config,
+        datasets={TRAIN_DATASET_KEY: dataset},
+    )
+    result = trainer.fit()
+    assert result.checkpoint
+
+
+def test_report_and_load_using_ml_session(ray_start_4_cpus):
+    def train_func():
+        checkpoint = train.get_checkpoint()
+        if checkpoint:
+            with checkpoint.as_directory() as checkpoint_dir:
+                import tensorflow as tf
+
+                model = tf.keras.models.load_model(checkpoint_dir)
+        else:
+            model = build_model()
+
+        if train.get_context().get_world_rank() == 0:
+            with tempfile.TemporaryDirectory() as tmp_dir:
+                model.save(tmp_dir)
+                train.report(
+                    metrics={"iter": 1},
+                    checkpoint=TensorflowCheckpoint.from_saved_model(tmp_dir),
+                )
+        else:
+            train.report(metrics={"iter": 1})
+
+    scaling_config = ScalingConfig(num_workers=2)
+    trainer = TensorflowTrainer(
+        train_loop_per_worker=train_func,
+        scaling_config=scaling_config,
+    )
+    result = trainer.fit()
+    checkpoint = result.checkpoint
+
+    trainer2 = TensorflowTrainer(
+        train_loop_per_worker=train_func,
+        scaling_config=scaling_config,
+        resume_from_checkpoint=checkpoint,
+    )
+    result = trainer2.fit()
+    checkpoint = result.checkpoint
+    with checkpoint.as_directory() as ckpt_dir:
+        assert os.path.exists(os.path.join(ckpt_dir, "saved_model.pb"))
+    assert result.metrics["iter"] == 1
+
+
+if __name__ == "__main__":
+    import sys
+
+    import pytest
+
+    sys.exit(pytest.main(["-v", "-x", __file__]))

python/ray/train/tests/test_torch_accelerate.py

@@ -0,0 +1,225 @@
+import os
+from tempfile import TemporaryDirectory
+
+import pytest
+import torch
+import torch.nn as nn
+from accelerate import Accelerator
+
+import ray
+import ray.train as train
+from ray.train import Checkpoint, ScalingConfig
+from ray.train.examples.pytorch.torch_linear_example import LinearDataset
+from ray.train.torch import TorchTrainer
+
+DEEPSPEED_CONFIG = {
+    "fp16": {
+        "enabled": "auto",
+        "loss_scale": 0,
+        "loss_scale_window": 1000,
+        "initial_scale_power": 16,
+        "hysteresis": 2,
+        "min_loss_scale": 1,
+    },
+    "bf16": {"enabled": "auto"},
+    "optimizer": {
+        "type": "AdamW",
+        "params": {
+            "lr": "auto",
+            "weight_decay": "auto",
+            "torch_adam": True,
+            "adam_w_mode": True,
+        },
+    },
+    "zero_optimization": {
+        "stage": 2,
+        "offload_optimizer": {"device": "cpu", "pin_memory": True},
+        "allgather_partitions": True,
+        "allgather_bucket_size": 2e8,
+        "overlap_comm": True,
+        "reduce_scatter": True,
+        "contiguous_gradients": True,
+    },
+    "gradient_accumulation_steps": 1,
+    "gradient_clipping": "auto",
+    "steps_per_print": 2000,
+    "train_batch_size": "auto",
+    "train_micro_batch_size_per_gpu": "auto",
+    "wall_clock_breakdown": False,
+}
+
+
+@pytest.fixture
+def ray_start_4_cpus():
+    address_info = ray.init(num_cpus=4)
+    yield address_info
+    # The code after the yield will run as teardown code.
+    ray.shutdown()
+
+
+def linear_train_func(accelerator: Accelerator, config):
+    from accelerate.utils import DummyOptim
+    from deepspeed.ops.adam import DeepSpeedCPUAdam
+
+    data_size = config.get("data_size", 1000)
+    val_size = config.get("val_size", 400)
+    batch_size = config.get("batch_size", 32)
+    hidden_size = config.get("hidden_size", 1)
+    lr = config.get("lr", 1e-2)
+    epochs = config.get("epochs", 3)
+
+    train_dataset = LinearDataset(2, 5, size=data_size)
+    val_dataset = LinearDataset(2, 5, size=val_size)
+    train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=batch_size)
+    validation_loader = torch.utils.data.DataLoader(val_dataset, batch_size=batch_size)
+
+    model = nn.Linear(1, hidden_size)
+
+    loss_fn = nn.MSELoss()
+    if (
+        accelerator.state.deepspeed_plugin
+        and "optimizer" in accelerator.state.deepspeed_plugin.deepspeed_config
+    ):
+        optimizer_cls = DummyOptim
+    elif accelerator.state.deepspeed_plugin:
+        optimizer_cls = DeepSpeedCPUAdam
+    else:
+        optimizer_cls = torch.optim.SGD
+
+    # Accelerate boilerplate
+    no_decay = ["bias", "LayerNorm.weight"]
+    optimizer_grouped_parameters = [
+        {
+            "params": [
+                p
+                for n, p in model.named_parameters()
+                if not any(nd in n for nd in no_decay)
+            ],
+            "weight_decay": 0.0,
+        },
+        {
+            "params": [
+                p
+                for n, p in model.named_parameters()
+                if any(nd in n for nd in no_decay)
+            ],
+            "weight_decay": 0.0,
+        },
+    ]
+    optimizer = optimizer_cls(optimizer_grouped_parameters, lr=lr)
+    train_loader, validation_loader, model, optimizer = accelerator.prepare(
+        train_loader, validation_loader, model, optimizer
+    )
+
+    results = []
+    for _ in range(epochs):
+        for X, y in train_loader:
+            # Compute prediction error
+            pred = model(X)
+            loss = loss_fn(pred, y)
+
+            # Backpropagation
+            accelerator.backward(loss)
+            optimizer.step()
+            optimizer.zero_grad()
+
+        num_batches = len(validation_loader)
+        model.eval()
+        loss = 0
+        with torch.no_grad():
+            for X, y in validation_loader:
+                pred = model(X)
+                loss += loss_fn(pred, y).item()
+        loss /= num_batches
+        import copy
+
+        model_copy = copy.deepcopy(accelerator.unwrap_model(model))
+        state_dict, loss = model_copy.cpu().state_dict(), loss
+
+        result = dict(loss=loss)
+        results.append(result)
+
+        with TemporaryDirectory() as tmpdir:
+            torch.save(state_dict, os.path.join(tmpdir, "checkpoint.pt"))
+            train.report(result, checkpoint=Checkpoint.from_directory(tmpdir))
+
+    return results
+
+
+@pytest.mark.parametrize("use_gpu", [True, False])
+def test_accelerate_base(ray_2_node_2_gpu, use_gpu):
+    def train_func(config):
+        accelerator = Accelerator(cpu=not use_gpu)
+        assert accelerator.device == train.torch.get_device()
+        assert accelerator.process_index == train.get_context().get_world_rank()
+        if accelerator.device.type != "cpu":
+            assert (
+                accelerator.local_process_index == train.get_context().get_local_rank()
+            )
+        result = linear_train_func(accelerator, config)
+        assert len(result) == epochs
+        assert result[-1]["loss"] < result[0]["loss"]
+
+    epochs = 3
+    scaling_config = ScalingConfig(num_workers=2, use_gpu=use_gpu)
+    config = {"lr": 1e-2, "hidden_size": 1, "batch_size": 4, "epochs": epochs}
+
+    trainer = TorchTrainer(
+        train_loop_per_worker=train_func,
+        train_loop_config=config,
+        scaling_config=scaling_config,
+    )
+    trainer.fit()
+
+
+def test_accelerate_deepspeed(ray_2_node_2_gpu):
+    from accelerate import DeepSpeedPlugin
+
+    def train_func(config):
+        deepspeed_plugin = DeepSpeedPlugin(hf_ds_config=DEEPSPEED_CONFIG)
+        accelerator = Accelerator(deepspeed_plugin=deepspeed_plugin)
+
+        assert accelerator.device == train.torch.get_device()
+        assert accelerator.process_index == train.get_context().get_world_rank()
+        assert accelerator.local_process_index == train.get_context().get_local_rank()
+        result = linear_train_func(accelerator, config)
+        assert len(result) == epochs
+        assert result[-1]["loss"] < result[0]["loss"]
+
+    epochs = 3
+    scaling_config = ScalingConfig(num_workers=2, use_gpu=True)
+    config = {"lr": 1e-2, "hidden_size": 1, "batch_size": 4, "epochs": epochs}
+
+    trainer = TorchTrainer(
+        train_loop_per_worker=train_func,
+        train_loop_config=config,
+        scaling_config=scaling_config,
+    )
+    trainer.fit()
+
+
+# Using CPU on purpose
+@pytest.mark.parametrize("num_workers", [1, 2])
+def test_accelerate_e2e(ray_start_4_cpus, num_workers):
+    def train_func():
+        accelerator = Accelerator(cpu=True)
+        assert accelerator.device == train.torch.get_device()
+        assert accelerator.process_index == train.get_context().get_world_rank()
+        model = torch.nn.Linear(3, 1)
+        model = accelerator.prepare(model)
+        with TemporaryDirectory() as tmpdir:
+            torch.save(model, os.path.join(tmpdir, "checkpoint.pt"))
+            train.report({}, checkpoint=Checkpoint.from_directory(tmpdir))
+
+    scaling_config = ScalingConfig(num_workers=num_workers)
+    trainer = TorchTrainer(
+        train_loop_per_worker=train_func,
+        scaling_config=scaling_config,
+    )
+    trainer.fit()
+
+
+if __name__ == "__main__":
+    import sys
+
+    sys.exit(pytest.main(["-v", "-x", __file__]))

python/ray/train/tests/test_torch_checkpoint.py

@@ -0,0 +1,41 @@
+import torch
+
+from ray.train.torch import TorchCheckpoint
+
+
+def assert_equal_torch_models(model1, model2):
+    # Check equality by comparing their `state_dict`
+    model1_state = model1.state_dict()
+    model2_state = model2.state_dict()
+    assert len(model1_state.keys()) == len(model2_state.keys())
+    for key in model1_state:
+        assert key in model2_state
+        assert torch.equal(model1_state[key], model2_state[key])
+
+
+def test_from_model():
+    model = torch.nn.Linear(1, 1)
+    checkpoint = TorchCheckpoint.from_model(model)
+    assert_equal_torch_models(checkpoint.get_model(), model)
+
+    with checkpoint.as_directory() as path:
+        checkpoint = TorchCheckpoint.from_directory(path)
+        checkpoint_model = checkpoint.get_model()
+
+    assert_equal_torch_models(checkpoint_model, model)
+
+
+def test_from_state_dict():
+    model = torch.nn.Linear(1, 1)
+    expected_state_dict = model.state_dict()
+    checkpoint = TorchCheckpoint.from_state_dict(expected_state_dict)
+    actual_state_dict = checkpoint.get_model(torch.nn.Linear(1, 1)).state_dict()
+    assert actual_state_dict == expected_state_dict
+
+
+if __name__ == "__main__":
+    import sys
+
+    import pytest
+
+    sys.exit(pytest.main(["-v", "-x", __file__]))

python/ray/train/tests/test_torch_device_manager.py

@@ -0,0 +1,95 @@
+import pytest
+import torch
+
+import ray
+from ray.air._internal.device_manager import (
+    CUDATorchDeviceManager,
+    NPUTorchDeviceManager,
+    get_torch_device_manager_by_context,
+)
+from ray.air._internal.device_manager.npu import NPU_TORCH_PACKAGE_AVAILABLE
+from ray.cluster_utils import Cluster
+from ray.train import ScalingConfig
+from ray.train.torch import TorchTrainer
+
+if NPU_TORCH_PACKAGE_AVAILABLE:
+    import torch_npu  # noqa: F401
+
+
+@pytest.fixture
+def ray_2_node_2_npus():
+    cluster = Cluster()
+    for _ in range(2):
+        cluster.add_node(num_cpus=4, resources={"NPU": 2})
+
+    ray.init(address=cluster.address)
+
+    yield
+
+    ray.shutdown()
+    cluster.shutdown()
+
+
+@pytest.fixture
+def ray_1_node_1_gpu_1_npu():
+    cluster = Cluster()
+    cluster.add_node(num_cpus=4, num_gpus=1, resources={"NPU": 1})
+    ray.init(address=cluster.address)
+
+    yield
+
+    ray.shutdown()
+    cluster.shutdown()
+
+
+def test_cuda_device_manager(ray_2_node_2_gpu):
+    def train_fn():
+        assert isinstance(get_torch_device_manager_by_context(), CUDATorchDeviceManager)
+
+    trainer = TorchTrainer(
+        train_loop_per_worker=train_fn,
+        scaling_config=ScalingConfig(
+            num_workers=1, use_gpu=True, resources_per_worker={"GPU": 1}
+        ),
+    )
+
+    trainer.fit()
+
+
+def test_npu_device_manager(ray_2_node_2_npus):
+    def train_fn():
+        assert isinstance(get_torch_device_manager_by_context(), NPUTorchDeviceManager)
+
+    trainer = TorchTrainer(
+        train_loop_per_worker=train_fn,
+        scaling_config=ScalingConfig(num_workers=1, resources_per_worker={"NPU": 1}),
+    )
+
+    if NPU_TORCH_PACKAGE_AVAILABLE and torch.npu.is_available():
+        # Except test run successfully when torch npu is available.
+        trainer.fit()
+    else:
+        # A RuntimeError will be triggered when NPU resources are declared
+        # but the torch npu is actually not available
+        with pytest.raises(RuntimeError):
+            trainer.fit()
+
+
+def test_device_manager_conflict(ray_1_node_1_gpu_1_npu):
+    trainer = TorchTrainer(
+        train_loop_per_worker=lambda: None,
+        scaling_config=ScalingConfig(
+            num_workers=1, use_gpu=True, resources_per_worker={"GPU": 1, "NPU": 1}
+        ),
+    )
+    # TODO: Do validation at the `ScalingConfig.__post_init__` level instead.
+    with pytest.raises(RuntimeError):
+        trainer.fit()
+
+
+if __name__ == "__main__":
+    import sys
+
+    import pytest
+
+    sys.exit(pytest.main(["-v", "-x", __file__]))

python/ray/train/tests/test_torch_fsdp.py

@@ -0,0 +1,45 @@
+import pytest
+import torch
+from torch.distributed.fsdp import FullyShardedDataParallel
+
+import ray
+from ray import train
+from ray.train import ScalingConfig
+from ray.train.torch import TorchTrainer
+
+
+@pytest.fixture
+def ray_start_4_cpus_2_gpus():
+    address_info = ray.init(num_cpus=4, num_gpus=2)
+    yield address_info
+    # The code after the yield will run as teardown code.
+    ray.shutdown()
+
+
+def test_torch_fsdp(ray_start_4_cpus_2_gpus):
+    """Tests if ``prepare_model`` correctly wraps in FSDP."""
+
+    def train_fn():
+        model = torch.nn.Linear(1, 1)
+
+        # Wrap in FSDP.
+        model = train.torch.prepare_model(model, parallel_strategy="fsdp")
+
+        # Make sure model is wrapped in FSDP.
+        assert isinstance(model, FullyShardedDataParallel)
+
+        # Make sure the model is on cuda.
+        assert next(model.parameters()).is_cuda
+
+    trainer = TorchTrainer(
+        train_fn, scaling_config=ScalingConfig(num_workers=2, use_gpu=True)
+    )
+    trainer.fit()
+
+
+if __name__ == "__main__":
+    import sys
+
+    import pytest
+
+    sys.exit(pytest.main(["-v", "-x", "-s", __file__]))

python/ray/train/tests/test_torch_lightning_train.py

@@ -0,0 +1,213 @@
+import os
+
+import numpy as np
+import pytest
+
+import ray
+from ray.train import ScalingConfig
+from ray.train.lightning import (
+    RayDDPStrategy,
+    RayDeepSpeedStrategy,
+    RayFSDPStrategy,
+    RayLightningEnvironment,
+    RayTrainReportCallback,
+)
+from ray.train.lightning._lightning_utils import import_lightning
+from ray.train.tests.lightning_test_utils import DummyDataModule, LinearModule
+from ray.train.torch import TorchTrainer
+
+pl = import_lightning()
+
+
+@pytest.fixture
+def ray_start_6_cpus_2_gpus():
+    address_info = ray.init(num_cpus=6, num_gpus=2)
+    yield address_info
+    # The code after the yield will run as teardown code.
+    ray.shutdown()
+
+
+@pytest.fixture
+def ray_start_6_cpus_4_gpus():
+    address_info = ray.init(num_cpus=6, num_gpus=4)
+    yield address_info
+    # The code after the yield will run as teardown code.
+    ray.shutdown()
+
+
+@pytest.mark.parametrize("strategy_name", ["ddp", "fsdp"])
+@pytest.mark.parametrize("accelerator", ["cpu", "gpu"])
+@pytest.mark.parametrize("datasource", ["dataloader", "datamodule"])
+def test_trainer_with_native_dataloader(
+    ray_start_6_cpus_2_gpus, strategy_name, accelerator, datasource
+):
+    """Test basic ddp and fsdp training with dataloader and datamodule."""
+
+    if accelerator == "cpu" and strategy_name == "fsdp":
+        return
+
+    num_workers = 2
+    num_epochs = 4
+    batch_size = 8
+    dataset_size = 256
+
+    strategy_map = {"ddp": RayDDPStrategy(), "fsdp": RayFSDPStrategy()}
+
+    def train_loop():
+        model = LinearModule(input_dim=32, output_dim=4, strategy=strategy_name)
+
+        strategy = strategy_map[strategy_name]
+
+        trainer = pl.Trainer(
+            max_epochs=num_epochs,
+            devices="auto",
+            accelerator=accelerator,
+            strategy=strategy,
+            plugins=[RayLightningEnvironment()],
+            callbacks=[RayTrainReportCallback()],
+        )
+
+        datamodule = DummyDataModule(batch_size, dataset_size)
+
+        if datasource == "dataloader":
+            trainer.fit(
+                model,
+                train_dataloaders=datamodule.train_dataloader(),
+                val_dataloaders=datamodule.val_dataloader(),
+            )
+        if datasource == "datamodule":
+            trainer.fit(model, datamodule=datamodule)
+
+    trainer = TorchTrainer(
+        train_loop_per_worker=train_loop,
+        scaling_config=ScalingConfig(num_workers=2, use_gpu=(accelerator == "gpu")),
+    )
+
+    results = trainer.fit()
+    assert results.metrics["epoch"] == num_epochs - 1
+    assert (
+        results.metrics["step"] == num_epochs * dataset_size / num_workers / batch_size
+    )
+    assert "loss" in results.metrics
+    assert "val_loss" in results.metrics
+
+
+@pytest.mark.parametrize("strategy_name", ["ddp", "fsdp"])
+@pytest.mark.parametrize("accelerator", ["cpu", "gpu"])
+def test_trainer_with_ray_data(ray_start_6_cpus_2_gpus, strategy_name, accelerator):
+    """Test Data integration with ddp and fsdp."""
+
+    if accelerator == "cpu" and strategy_name == "fsdp":
+        return
+
+    num_epochs = 4
+    batch_size = 8
+    num_workers = 2
+    dataset_size = 256
+
+    strategy_map = {"ddp": RayDDPStrategy(), "fsdp": RayFSDPStrategy()}
+
+    dataset = np.random.rand(dataset_size, 32).astype(np.float32)
+    train_dataset = ray.data.from_numpy(dataset)
+    val_dataset = ray.data.from_numpy(dataset)
+
+    def train_loop():
+        model = LinearModule(input_dim=32, output_dim=4, strategy=strategy_name)
+
+        strategy = strategy_map[strategy_name]
+
+        trainer = pl.Trainer(
+            max_epochs=num_epochs,
+            devices="auto",
+            accelerator=accelerator,
+            strategy=strategy,
+            plugins=[RayLightningEnvironment()],
+            callbacks=[RayTrainReportCallback()],
+        )
+
+        train_data_iterable = ray.train.get_dataset_shard("train").iter_torch_batches(
+            batch_size=batch_size
+        )
+        val_data_iterable = ray.train.get_dataset_shard("val").iter_torch_batches(
+            batch_size=batch_size
+        )
+
+        trainer.fit(
+            model,
+            train_dataloaders=train_data_iterable,
+            val_dataloaders=val_data_iterable,
+        )
+
+    trainer = TorchTrainer(
+        train_loop_per_worker=train_loop,
+        scaling_config=ScalingConfig(num_workers=2, use_gpu=(accelerator == "gpu")),
+        datasets={"train": train_dataset, "val": val_dataset},
+    )
+
+    results = trainer.fit()
+    assert results.metrics["epoch"] == num_epochs - 1
+    assert (
+        results.metrics["step"] == num_epochs * dataset_size / num_workers / batch_size
+    )
+    assert "loss" in results.metrics
+    assert "val_loss" in results.metrics
+
+
+@pytest.mark.parametrize("stage", [1, 2, 3])
+def test_deepspeed_zero_stages(ray_start_6_cpus_4_gpus, tmpdir, stage):
+    num_epochs = 5
+    batch_size = 8
+    num_workers = 4
+    dataset_size = 256
+
+    def train_loop():
+        model = LinearModule(input_dim=32, output_dim=4, strategy="deepspeed")
+
+        strategy = RayDeepSpeedStrategy(stage=stage)
+
+        trainer = pl.Trainer(
+            max_epochs=num_epochs,
+            devices="auto",
+            accelerator="gpu",
+            strategy=strategy,
+            plugins=[RayLightningEnvironment()],
+            callbacks=[RayTrainReportCallback()],
+        )
+
+        datamodule = DummyDataModule(batch_size, dataset_size)
+        trainer.fit(model, datamodule=datamodule)
+
+    trainer = TorchTrainer(
+        train_loop_per_worker=train_loop,
+        scaling_config=ScalingConfig(num_workers=num_workers, use_gpu=True),
+    )
+
+    result = trainer.fit()
+
+    # Check all deepspeed model/optimizer shards are saved
+    all_files = os.listdir(f"{result.checkpoint.path}/checkpoint.ckpt/checkpoint")
+    for rank in range(num_workers):
+        full_model = "mp_rank_00_model_states.pt"
+        model_shard = f"zero_pp_rank_{rank}_mp_rank_00_model_states.pt"
+        optim_shard = f"zero_pp_rank_{rank}_mp_rank_00_optim_states.pt"
+
+        assert (
+            optim_shard in all_files
+        ), f"[stage-{stage}] Optimizer states `{optim_shard}` doesn't exist!"
+
+        if stage == 3:
+            assert (
+                model_shard in all_files
+            ), f"[stage-{stage}] Model states {model_shard} doesn't exist!"
+        else:
+            assert (
+                full_model in all_files
+            ), f"[stage-{stage}] Model states {full_model} doesn't exist!"
+
+
+if __name__ == "__main__":
+    import sys
+
+    import pytest
+
+    sys.exit(pytest.main(["-v", "-x", __file__]))

python/ray/train/tests/test_torch_trainer.py

@@ -0,0 +1,306 @@
+import contextlib
+import os
+import time
+import uuid
+from unittest.mock import patch
+
+import pytest
+import torch
+
+import ray
+import ray.train as train
+from ray.cluster_utils import Cluster
+from ray.train import RunConfig, ScalingConfig
+from ray.train.examples.pytorch.torch_linear_example import (
+    train_func as linear_train_func,
+)
+from ray.train.torch import TorchCheckpoint, TorchConfig, TorchTrainer
+from ray.train.trainer import TrainingFailedError
+
+
+@pytest.fixture
+def ray_start_4_cpus():
+    address_info = ray.init(num_cpus=4)
+    yield address_info
+    # The code after the yield will run as teardown code.
+    ray.shutdown()
+
+
+@contextlib.contextmanager
+def ray_start_2_node_cluster(num_cpus_per_node: int, num_gpus_per_node: int):
+    cluster = Cluster()
+    for _ in range(2):
+        cluster.add_node(num_cpus=num_cpus_per_node, num_gpus=num_gpus_per_node)
+
+    ray.init(address=cluster.address)
+
+    yield
+
+    ray.shutdown()
+    cluster.shutdown()
+
+
+@pytest.mark.parametrize("num_workers", [1, 2])
+def test_torch_linear(ray_start_4_cpus, num_workers):
+    def train_func(config):
+        result = linear_train_func(config)
+        assert len(result) == epochs
+        assert result[-1]["loss"] < result[0]["loss"]
+
+    num_workers = num_workers
+    epochs = 3
+    scaling_config = ScalingConfig(num_workers=num_workers)
+    config = {"lr": 1e-2, "hidden_size": 1, "batch_size": 4, "epochs": epochs}
+    trainer = TorchTrainer(
+        train_loop_per_worker=train_func,
+        train_loop_config=config,
+        scaling_config=scaling_config,
+    )
+    trainer.fit()
+
+
+@pytest.mark.parametrize("prepare_model", (True, False))
+def test_torch_e2e(ray_start_4_cpus, prepare_model):
+    def train_func():
+        model = torch.nn.Linear(3, 1)
+        if prepare_model:
+            model = train.torch.prepare_model(model)
+        train.report({}, checkpoint=TorchCheckpoint.from_model(model))
+
+    scaling_config = ScalingConfig(num_workers=2)
+    trainer = TorchTrainer(
+        train_loop_per_worker=train_func,
+        scaling_config=scaling_config,
+    )
+    trainer.fit()
+
+
+@pytest.mark.parametrize("prepare_model", (True, False))
+def test_torch_e2e_state_dict(ray_start_4_cpus, prepare_model):
+    def train_func():
+        model = torch.nn.Linear(3, 1)
+        if prepare_model:
+            model = train.torch.prepare_model(model)
+        train.report({}, checkpoint=TorchCheckpoint.from_state_dict(model.state_dict()))
+
+    scaling_config = ScalingConfig(num_workers=2)
+    trainer = TorchTrainer(
+        train_loop_per_worker=train_func,
+        scaling_config=scaling_config,
+    )
+    result = trainer.fit()
+
+    # If loading from a state dict, a model definition must be passed in.
+    with pytest.raises(ValueError):
+        torch_checkpoint = TorchCheckpoint(
+            path=result.checkpoint.path, filesystem=result.checkpoint.filesystem
+        )
+        torch_checkpoint.get_model()
+
+
+def test_checkpoint_freq(ray_start_4_cpus):
+    # checkpoint_freq is not supported so raise an error
+    trainer = TorchTrainer(
+        train_loop_per_worker=lambda config: None,
+        scaling_config=train.ScalingConfig(num_workers=1),
+        run_config=train.RunConfig(
+            checkpoint_config=train.CheckpointConfig(
+                checkpoint_frequency=2,
+            ),
+        ),
+    )
+    with pytest.raises(ValueError):
+        trainer.fit()
+
+
+def test_torch_session_errors(ray_start_4_cpus):
+    """Test fail-fast behavior when reporting dicts with Torch tensors"""
+
+    def train_func():
+        model = torch.nn.Linear(1, 1).state_dict()
+        with pytest.raises(ValueError):
+            train.report(model)
+
+    scaling_config = ScalingConfig(num_workers=2)
+    trainer = TorchTrainer(
+        train_loop_per_worker=train_func,
+        scaling_config=scaling_config,
+    )
+    trainer.fit()
+
+
+def test_single_worker_failure(ray_start_4_cpus):
+    """Tests if training fails upon any worker failure."""
+
+    def single_worker_fail():
+        if train.get_context().get_world_rank() == 0:
+            raise RuntimeError
+        else:
+            time.sleep(1000000)
+
+    scaling_config = ScalingConfig(num_workers=2)
+    trainer = TorchTrainer(
+        train_loop_per_worker=single_worker_fail,
+        scaling_config=scaling_config,
+    )
+    with pytest.raises(TrainingFailedError) as exc_info:
+        trainer.fit()
+    assert isinstance(exc_info.value.__cause__, RuntimeError)
+
+
+@pytest.mark.parametrize("num_gpus_per_worker", [0.5, 1, 2])
+def test_tune_torch_get_device_gpu(num_gpus_per_worker):
+    """Tests if GPU ids are set correctly when running train concurrently in nested actors
+    (for example when used with Tune).
+    """
+    from ray.train import ScalingConfig
+
+    num_samples = 2
+    num_workers = 2
+
+    # We should have exactly enough resources in the cluster to run both samples
+    # concurrently.
+    total_gpus_required = num_workers * num_gpus_per_worker * num_samples
+    # Divide by two because of a 2 node cluster.
+    gpus_per_node = total_gpus_required // 2
+
+    exception = None
+    # Use the same number of cpus per node as gpus per node.
+    with ray_start_2_node_cluster(
+        num_cpus_per_node=gpus_per_node, num_gpus_per_node=gpus_per_node
+    ):
+
+        @patch("torch.cuda.is_available", lambda: True)
+        def train_fn():
+            # We use STRICT_SPREAD strategy to force multiple samples on the same node.
+            # For single or fractional GPU case, each worker has only 1 visible device (
+            # the other is taken by the other sample) so device index should be 0.
+            # For the multiple GPU case, each worker has 2 visible devices so device
+            # index should be either 0 or 1. It doesn't matter which.
+            device_ids = sorted([device.index for device in train.torch.get_devices()])
+            assert device_ids in [[0], [0, 1]]
+
+        @ray.remote(num_cpus=0)
+        class TrialActor:
+            def __init__(self, warmup_steps):
+                self.trainer = TorchTrainer(
+                    train_fn,
+                    torch_config=TorchConfig(backend="gloo"),
+                    run_config=RunConfig(
+                        # Use a unique name to avoid using the same
+                        # experiment directory
+                        name=f"test_tune_torch_get_device_gpu_{uuid.uuid4()}"
+                    ),
+                    scaling_config=ScalingConfig(
+                        num_workers=num_workers,
+                        use_gpu=True,
+                        resources_per_worker={"CPU": 1, "GPU": num_gpus_per_worker},
+                        # Need to specify 0 trainer resources so STRICT_SPREAD
+                        # will work.
+                        trainer_resources={"CPU": 0},
+                        placement_strategy="STRICT_SPREAD",
+                        # Each gpu worker will be spread onto separate nodes. This
+                        # forces different samples to run concurrently on the same
+                        # node.
+                    ),
+                )
+
+            def run(self):
+                return self.trainer.fit()
+
+        try:
+            actors = [TrialActor.remote(1) for _ in range(num_samples)]
+            ray.get([actor.run.remote() for actor in actors])
+        except Exception as exc:
+            exception = exc
+
+    # Raise exception after Ray cluster has been shutdown to avoid corrupted state
+    if exception:
+        raise exception
+
+
+def test_torch_amp(ray_start_4_cpus):
+    def train_fn():
+        train.torch.accelerate(amp=True)
+        model = torch.nn.Linear(1, 1)
+        model = train.torch.prepare_model(model)
+
+        train.report({}, checkpoint=TorchCheckpoint.from_model(model))
+
+    trainer = TorchTrainer(
+        train_fn,
+        scaling_config=ScalingConfig(num_workers=2),
+    )
+    results = trainer.fit()
+    assert results.checkpoint
+
+
+def test_torch_amp_with_custom_get_state(ray_start_4_cpus):
+    """Tests amp with a model that has a custom __getstate__ method defined.
+
+    See https://discuss.ray.io/t/ray-train-hangs-for-long-time/6333/7
+    """
+
+    def train_fn():
+        train.torch.accelerate(amp=True)
+
+        class CustomLinear(torch.nn.Linear):
+            def __getstate__(self):
+                return self.__dict__.copy()
+
+        model = CustomLinear(1, 1)
+        model = train.torch.prepare_model(model)
+
+        # TorchCheckpoint.from_model fails, so just save the state dict only.
+        train.report(
+            {}, checkpoint=TorchCheckpoint.from_state_dict(model.module.state_dict())
+        )
+
+    trainer = TorchTrainer(
+        train_fn,
+        scaling_config=ScalingConfig(num_workers=2),
+    )
+    results = trainer.fit()
+    assert results.checkpoint
+
+
+def test_torch_env_vars(ray_start_4_cpus):
+    """Check that env vars are set as expected."""
+
+    def train_func(config):
+        context = train.get_context()
+        assert os.environ["LOCAL_RANK"] == str(context.get_local_rank())
+        assert os.environ["RANK"] == str(context.get_world_rank())
+        assert os.environ["LOCAL_WORLD_SIZE"] == str(context.get_local_world_size())
+        assert os.environ["WORLD_SIZE"] == str(context.get_world_size())
+        assert os.environ["NODE_RANK"] == str(context.get_node_rank())
+
+        assert os.environ["ACCELERATE_TORCH_DEVICE"] == str(train.torch.get_device())
+
+    num_workers = 1
+    scaling_config = ScalingConfig(num_workers=num_workers)
+    trainer = TorchTrainer(
+        train_loop_per_worker=train_func,
+        scaling_config=scaling_config,
+    )
+    trainer.fit()
+
+
+def test_nonserializable_train_function(ray_start_4_cpus):
+    import threading
+
+    lock = threading.Lock()
+
+    def train_func():
+        print(lock)
+
+    trainer = TorchTrainer(train_func)
+    # Check that the `inspect_serializability` trace was printed
+    with pytest.raises(TypeError, match=r".*was found to be non-serializable.*"):
+        trainer.fit()
+
+
+if __name__ == "__main__":
+    import sys
+
+    sys.exit(pytest.main(["-v", "-x", __file__]))

python/ray/train/tests/test_torch_transformers_train.py

@@ -0,0 +1,304 @@
+import pandas as pd
+import pytest
+from datasets import Dataset
+from transformers import AutoConfig, AutoModelForCausalLM, Trainer, TrainingArguments
+
+import ray.data
+from ray import tune
+from ray.train import Checkpoint, ScalingConfig
+from ray.train.huggingface.transformers import RayTrainReportCallback, prepare_trainer
+from ray.train.tests._huggingface_data import train_data, validation_data
+from ray.train.torch import TorchTrainer
+from ray.tune import Tuner
+from ray.tune.schedulers.async_hyperband import ASHAScheduler
+from ray.tune.schedulers.resource_changing_scheduler import (
+    DistributeResources,
+    ResourceChangingScheduler,
+)
+
+
+@pytest.fixture
+def ray_start_6_cpus_2_gpus():
+    address_info = ray.init(num_cpus=6, num_gpus=2)
+    yield address_info
+    # The code after the yield will run as teardown code.
+    ray.shutdown()
+
+
+@pytest.fixture
+def ray_start_8_cpus():
+    address_info = ray.init(num_cpus=8)
+    yield address_info
+    # The code after the yield will run as teardown code.
+    ray.shutdown()
+
+
+# We are only testing Causal Language Modeling here
+MODEL_NAME = "hf-internal-testing/tiny-random-BloomForCausalLM"
+
+# Training Loop Configurations
+NUM_WORKERS = 2
+BATCH_SIZE_PER_WORKER = 2
+TRAIN_DATASET_SIZE = 16
+MAX_EPOCHS = 4
+
+STEPS_PER_EPOCH = TRAIN_DATASET_SIZE // (BATCH_SIZE_PER_WORKER * NUM_WORKERS)
+MAX_STEPS = MAX_EPOCHS * STEPS_PER_EPOCH
+
+# Transformers Traienr Configurations
+CONFIGURATIONS = {
+    "epoch_gpu": {
+        "evaluation_strategy": "epoch",
+        "save_strategy": "epoch",
+        "logging_strategy": "epoch",
+        "eval_steps": None,
+        "save_steps": None,
+        "logging_steps": None,
+        "no_cuda": False,
+        "use_dict_eval_datasets": False,
+    },
+    "steps_gpu": {
+        "evaluation_strategy": "steps",
+        "save_strategy": "steps",
+        "logging_strategy": "steps",
+        "eval_steps": STEPS_PER_EPOCH,
+        "save_steps": STEPS_PER_EPOCH * 2,
+        "logging_steps": 1,
+        "no_cuda": False,
+        "use_dict_eval_datasets": False,
+    },
+    "steps_cpu": {
+        "evaluation_strategy": "steps",
+        "save_strategy": "steps",
+        "logging_strategy": "steps",
+        "eval_steps": STEPS_PER_EPOCH,
+        "save_steps": STEPS_PER_EPOCH,
+        "logging_steps": 1,
+        "no_cuda": True,
+        "use_dict_eval_datasets": False,
+    },
+}
+
+
+def train_func(config):
+    # Datasets
+    if config["use_ray_data"]:
+        train_ds_shard = ray.train.get_dataset_shard("train")
+        train_dataset = train_ds_shard.iter_torch_batches(
+            batch_size=BATCH_SIZE_PER_WORKER
+        )
+        if config["use_dict_eval_datasets"]:
+            eval_ds_shard_1 = ray.train.get_dataset_shard("eval_1")
+            eval_ds_shard_2 = ray.train.get_dataset_shard("eval_2")
+
+            eval_dataset = {
+                "eval_1": eval_ds_shard_1.iter_torch_batches(
+                    batch_size=BATCH_SIZE_PER_WORKER
+                ),
+                "eval_2": eval_ds_shard_2.iter_torch_batches(
+                    batch_size=BATCH_SIZE_PER_WORKER
+                ),
+            }
+        else:
+            eval_ds_shard = ray.train.get_dataset_shard("eval")
+
+            eval_dataset = eval_ds_shard.iter_torch_batches(
+                batch_size=BATCH_SIZE_PER_WORKER
+            )
+    else:
+        train_df = pd.read_json(train_data)
+        validation_df = pd.read_json(validation_data)
+
+        train_dataset = Dataset.from_pandas(train_df)
+        eval_dataset = Dataset.from_pandas(validation_df)
+
+    # Model
+    model_config = AutoConfig.from_pretrained(MODEL_NAME)
+    model = AutoModelForCausalLM.from_config(model_config)
+
+    # HF Transformers Trainer
+    training_args = TrainingArguments(
+        f"{MODEL_NAME}-wikitext2",
+        eval_strategy=config["evaluation_strategy"],
+        logging_strategy=config["logging_strategy"],
+        save_strategy=config["save_strategy"],
+        eval_steps=config["eval_steps"],
+        save_steps=config["save_steps"],
+        logging_steps=config["logging_steps"],
+        num_train_epochs=config.get("num_train_epochs", MAX_EPOCHS),
+        max_steps=config.get("max_steps", -1),
+        learning_rate=config.get("learning_rate", 2e-5),
+        per_device_train_batch_size=BATCH_SIZE_PER_WORKER,
+        per_device_eval_batch_size=BATCH_SIZE_PER_WORKER,
+        weight_decay=0.01,
+        disable_tqdm=True,
+        use_cpu=config["no_cuda"],
+        report_to="none",
+    )
+    trainer = Trainer(
+        model=model,
+        args=training_args,
+        train_dataset=train_dataset,
+        eval_dataset=eval_dataset,
+    )
+
+    # Report to Ray Train
+    trainer.add_callback(RayTrainReportCallback())
+    trainer = prepare_trainer(trainer)
+
+    # Start Training
+    trainer.train()
+
+
+@pytest.mark.parametrize("config_id", ["epoch_gpu", "steps_gpu", "steps_cpu"])
+def test_e2e_hf_data(ray_start_6_cpus_2_gpus, config_id):
+    train_loop_config = CONFIGURATIONS[config_id]
+
+    # Specify `num_train_epochs` for Map-style Dataset
+    train_loop_config["use_ray_data"] = False
+    train_loop_config["num_train_epochs"] = MAX_EPOCHS
+
+    # Calculate the num of Ray training iterations
+    if train_loop_config["save_strategy"] == "steps":
+        num_iterations = MAX_STEPS // train_loop_config["save_steps"]
+    else:
+        num_iterations = MAX_EPOCHS
+
+    use_gpu = not train_loop_config["no_cuda"]
+
+    trainer = TorchTrainer(
+        train_func,
+        train_loop_config=train_loop_config,
+        scaling_config=ScalingConfig(num_workers=NUM_WORKERS, use_gpu=use_gpu),
+    )
+    result = trainer.fit()
+
+    assert result.metrics["epoch"] == MAX_EPOCHS
+    assert result.metrics["step"] == MAX_STEPS
+    assert result.metrics["training_iteration"] == num_iterations
+    assert result.checkpoint
+    assert isinstance(result.checkpoint, Checkpoint)
+    assert len(result.best_checkpoints) == num_iterations
+    assert "eval_loss" in result.metrics
+
+
+@pytest.mark.parametrize("config_id", ["steps_gpu", "steps_cpu"])
+def test_e2e_ray_data(ray_start_6_cpus_2_gpus, config_id):
+    train_loop_config = CONFIGURATIONS[config_id]
+
+    # Must specify `max_steps` for Iterable Dataset
+    train_loop_config["use_ray_data"] = True
+    train_loop_config["max_steps"] = MAX_STEPS
+
+    # Calculate the num of Ray training iterations
+    num_iterations = MAX_STEPS // train_loop_config["save_steps"]
+
+    train_df = pd.read_json(train_data)
+    validation_df = pd.read_json(validation_data)
+
+    ray_train_ds = ray.data.from_pandas(train_df)
+    ray_eval_ds = ray.data.from_pandas(validation_df)
+
+    use_gpu = not train_loop_config["no_cuda"]
+
+    trainer = TorchTrainer(
+        train_func,
+        train_loop_config=train_loop_config,
+        scaling_config=ScalingConfig(num_workers=NUM_WORKERS, use_gpu=use_gpu),
+        datasets={"train": ray_train_ds, "eval": ray_eval_ds},
+    )
+    result = trainer.fit()
+
+    assert result.metrics["step"] == MAX_STEPS
+    assert result.metrics["training_iteration"] == num_iterations
+    assert result.checkpoint
+    assert isinstance(result.checkpoint, Checkpoint)
+    assert len(result.best_checkpoints) == num_iterations
+    assert "eval_loss" in result.metrics
+
+
+@pytest.mark.parametrize("config_id", ["steps_gpu", "steps_cpu"])
+def test_e2e_dict_eval_ray_data(ray_start_6_cpus_2_gpus, config_id):
+    train_loop_config = CONFIGURATIONS[config_id]
+
+    # Must specify `max_steps` for Iterable Dataset
+    train_loop_config["use_ray_data"] = True
+    train_loop_config["use_dict_eval_datasets"] = True
+    train_loop_config["max_steps"] = MAX_STEPS
+
+    # Calculate the num of Ray training iterations
+    num_iterations = MAX_STEPS // train_loop_config["save_steps"]
+
+    train_df = pd.read_json(train_data)
+    validation_df = pd.read_json(validation_data)
+
+    ray_train_ds = ray.data.from_pandas(train_df)
+    ray_eval_ds_1 = ray.data.from_pandas(validation_df)
+    ray_eval_ds_2 = ray.data.from_pandas(validation_df)
+
+    use_gpu = not train_loop_config["no_cuda"]
+
+    trainer = TorchTrainer(
+        train_func,
+        train_loop_config=train_loop_config,
+        scaling_config=ScalingConfig(num_workers=NUM_WORKERS, use_gpu=use_gpu),
+        datasets={
+            "train": ray_train_ds,
+            "eval_1": ray_eval_ds_1,
+            "eval_2": ray_eval_ds_2,
+        },
+    )
+    result = trainer.fit()
+
+    assert result.metrics["step"] == MAX_STEPS
+    assert result.metrics["training_iteration"] == num_iterations
+    assert result.checkpoint
+    assert isinstance(result.checkpoint, Checkpoint)
+    assert len(result.best_checkpoints) == num_iterations
+    assert "eval_eval_1_loss" in result.metrics
+    assert "eval_eval_2_loss" in result.metrics
+
+
+# Tests if Ray Tune works correctly.
+def test_tune(ray_start_8_cpus):
+    train_loop_config = CONFIGURATIONS["steps_cpu"]
+    train_loop_config["use_ray_data"] = False
+
+    use_gpu = not train_loop_config["no_cuda"]
+    trainer = TorchTrainer(
+        train_func,
+        train_loop_config=train_loop_config,
+        scaling_config=ScalingConfig(num_workers=NUM_WORKERS, use_gpu=use_gpu),
+    )
+
+    tuner = Tuner(
+        trainer,
+        param_space={
+            "train_loop_config": {
+                "learning_rate": tune.loguniform(2e-6, 2e-5),
+            }
+        },
+        tune_config=tune.TuneConfig(
+            metric="eval_loss",
+            mode="min",
+            num_samples=3,
+            scheduler=ResourceChangingScheduler(
+                ASHAScheduler(
+                    max_t=MAX_EPOCHS,
+                ),
+                resources_allocation_function=DistributeResources(
+                    add_bundles=True, reserve_resources={"CPU": 1}
+                ),
+            ),
+        ),
+    )
+    tune_results = tuner.fit()
+    assert not tune_results.errors
+
+
+if __name__ == "__main__":
+    import sys
+
+    import pytest
+
+    sys.exit(pytest.main(["-v", "-x", __file__]))

python/ray/train/tests/test_torch_utils.py

@@ -0,0 +1,42 @@
+import pytest
+import torch
+
+from ray.air._internal.torch_utils import (
+    contains_tensor,
+    load_torch_model,
+)
+
+torch_module = torch.nn.Linear(1, 1)
+
+
+class TestLoadTorchModel:
+    def test_load_module(self):
+        assert load_torch_model(torch_module) == torch_module
+
+    def test_load_state_dict(self):
+        state_dict = torch_module.state_dict()
+        model_definition = torch.nn.Linear(1, 1)
+        assert model_definition.state_dict() != state_dict
+
+        assert load_torch_model(state_dict, model_definition).state_dict() == state_dict
+
+    def test_load_state_dict_fail(self):
+        with pytest.raises(ValueError):
+            # model_definition is required to load state dict.
+            load_torch_model(torch_module.state_dict())
+
+
+def test_contains_tensor():
+    t = torch.tensor([0])
+    assert contains_tensor(t)
+    assert contains_tensor([1, 2, 3, t, 5, 6])
+    assert contains_tensor([1, 2, 3, {"dict": t}, 5, 6])
+    assert contains_tensor({"outer": [1, 2, 3, {"dict": t}, 5, 6]})
+    assert contains_tensor({t: [1, 2, 3, {"dict": 2}, 5, 6]})
+    assert not contains_tensor([4, 5, 6])
+
+
+if __name__ == "__main__":
+    import sys
+
+    sys.exit(pytest.main(["-sv", __file__]))

python/ray/train/tests/test_train_head.py

@@ -0,0 +1,74 @@
+import sys
+import time
+
+import pytest
+import requests
+
+import ray
+from ray.train import RunConfig, ScalingConfig
+from ray.train.torch import TorchTrainer
+
+
+@pytest.fixture
+def ray_start_8_cpus(monkeypatch):
+    monkeypatch.setenv("RAY_TRAIN_ENABLE_STATE_TRACKING", "1")
+    address_info = ray.init(num_cpus=8)
+    yield address_info
+    # The code after the yield will run as teardown code.
+    ray.shutdown()
+
+
+def test_get_train_runs(ray_start_8_cpus):
+    def train_func():
+        print("Training Starts")
+        time.sleep(0.5)
+
+    datasets = {"train": ray.data.range(100), "val": ray.data.range(100)}
+
+    trainer = TorchTrainer(
+        train_func,
+        run_config=RunConfig(name="my_train_run", storage_path="/tmp/cluster_storage"),
+        scaling_config=ScalingConfig(num_workers=4, use_gpu=False),
+        datasets=datasets,
+    )
+    trainer.fit()
+
+    # Call the train run api
+    url = ray._private.worker.get_dashboard_url()
+    resp = requests.get("http://" + url + "/api/train/v2/runs")
+    assert resp.status_code == 200
+    body = resp.json()
+    assert len(body["train_runs"]) == 1
+    assert body["train_runs"][0]["name"] == "my_train_run"
+    assert len(body["train_runs"][0]["workers"]) == 4
+
+
+def test_add_actor_status(ray_start_8_cpus):
+    from ray.train._internal.state.schema import ActorStatusEnum
+
+    def check_actor_status(expected_actor_status):
+        url = ray._private.worker.get_dashboard_url()
+        resp = requests.get("http://" + url + "/api/train/v2/runs")
+        assert resp.status_code == 200
+        body = resp.json()
+
+        for worker_info in body["train_runs"][0]["workers"]:
+            assert worker_info["status"] == expected_actor_status
+
+    def train_func():
+        print("Training Starts")
+        time.sleep(0.5)
+        check_actor_status(expected_actor_status=ActorStatusEnum.ALIVE)
+
+    trainer = TorchTrainer(
+        train_func,
+        run_config=RunConfig(name="my_train_run", storage_path="/tmp/cluster_storage"),
+        scaling_config=ScalingConfig(num_workers=4, use_gpu=False),
+    )
+    trainer.fit()
+
+    check_actor_status(expected_actor_status=ActorStatusEnum.DEAD)
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main(["-sv", __file__]))

python/ray/train/tests/test_train_usage.py

@@ -0,0 +1,148 @@
+import pytest
+import torch
+
+import ray
+from ray.train import ScalingConfig
+from ray.train.torch import TorchTrainer
+
+
+@pytest.fixture
+def shutdown_only():
+    yield None
+    ray.shutdown()
+
+
+def run_torch():
+    from torch.utils.data import DataLoader, TensorDataset
+
+    from ray.train.torch import (
+        get_device,
+        get_devices,
+        prepare_data_loader,
+        prepare_model,
+    )
+
+    def train_func():
+        # Create dummy model and data loader
+        model = torch.nn.Linear(10, 10)
+        inputs, targets = torch.randn(128, 10), torch.randn(128, 1)
+        dataloader = DataLoader(TensorDataset(inputs, targets), batch_size=32)
+
+        # Test Torch Utilities
+        prepare_data_loader(dataloader)
+        prepare_model(model)
+        get_device()
+        get_devices()
+
+    trainer = TorchTrainer(
+        train_func, scaling_config=ScalingConfig(num_workers=2, use_gpu=False)
+    )
+    trainer.fit()
+
+
+def run_lightning():
+    import lightning.pytorch as pl
+
+    from ray.train.lightning import (
+        RayDDPStrategy,
+        RayDeepSpeedStrategy,
+        RayFSDPStrategy,
+        RayLightningEnvironment,
+        RayTrainReportCallback,
+        prepare_trainer,
+    )
+
+    def train_func():
+        # Test Lighting utilites
+        strategy = RayFSDPStrategy()
+        strategy = RayDeepSpeedStrategy()
+        strategy = RayDDPStrategy()
+        ray_environment = RayLightningEnvironment()
+        report_callback = RayTrainReportCallback()
+
+        trainer = pl.Trainer(
+            devices="auto",
+            accelerator="auto",
+            strategy=strategy,
+            plugins=[ray_environment],
+            callbacks=[report_callback],
+        )
+        trainer = prepare_trainer(trainer)
+
+    trainer = TorchTrainer(
+        train_func, scaling_config=ScalingConfig(num_workers=2, use_gpu=False)
+    )
+
+    trainer.fit()
+
+
+def run_transformers():
+    from datasets import Dataset
+    from transformers import Trainer, TrainingArguments
+
+    from ray.train.huggingface.transformers import (
+        RayTrainReportCallback,
+        prepare_trainer,
+    )
+
+    def train_func():
+        # Create dummy model and datasets
+        dataset = Dataset.from_dict({"text": ["text1", "text2"], "label": [0, 1]})
+        model = torch.nn.Linear(10, 10)
+
+        # Test Transformers utilites
+        training_args = TrainingArguments(output_dir="./results", use_cpu=True)
+        trainer = Trainer(model=model, args=training_args, train_dataset=dataset)
+
+        trainer.add_callback(RayTrainReportCallback())
+        trainer = prepare_trainer(trainer)
+
+    trainer = TorchTrainer(
+        train_func, scaling_config=ScalingConfig(num_workers=2, use_gpu=False)
+    )
+
+    trainer.fit()
+
+
+@pytest.mark.parametrize("framework", ["torch", "lightning", "transformers"])
+def test_torch_utility_usage_tags(shutdown_only, framework):
+    from ray._common.usage.usage_lib import TagKey, get_extra_usage_tags_to_report
+
+    ctx = ray.init()
+    gcs_client = ray._raylet.GcsClient(address=ctx.address_info["gcs_address"])
+
+    if framework == "torch":
+        run_torch()
+        expected_tags = [
+            TagKey.TRAIN_TORCH_GET_DEVICE,
+            TagKey.TRAIN_TORCH_GET_DEVICES,
+            TagKey.TRAIN_TORCH_PREPARE_MODEL,
+            TagKey.TRAIN_TORCH_PREPARE_DATALOADER,
+        ]
+    elif framework == "lightning":
+        run_lightning()
+        expected_tags = [
+            TagKey.TRAIN_LIGHTNING_PREPARE_TRAINER,
+            TagKey.TRAIN_LIGHTNING_RAYTRAINREPORTCALLBACK,
+            TagKey.TRAIN_LIGHTNING_RAYDDPSTRATEGY,
+            TagKey.TRAIN_LIGHTNING_RAYFSDPSTRATEGY,
+            TagKey.TRAIN_LIGHTNING_RAYDEEPSPEEDSTRATEGY,
+            TagKey.TRAIN_LIGHTNING_RAYLIGHTNINGENVIRONMENT,
+        ]
+    elif framework == "transformers":
+        run_transformers()
+        expected_tags = [
+            TagKey.TRAIN_TRANSFORMERS_PREPARE_TRAINER,
+            TagKey.TRAIN_TRANSFORMERS_RAYTRAINREPORTCALLBACK,
+        ]
+
+    result = get_extra_usage_tags_to_report(gcs_client)
+    assert set(result.keys()).issuperset(
+        {TagKey.Name(tag).lower() for tag in expected_tags}
+    )
+
+
+if __name__ == "__main__":
+    import sys
+
+    sys.exit(pytest.main(["-v", "-x", __file__]))

python/ray/train/tests/test_trainer_restore.py

@@ -0,0 +1,359 @@
+from functools import partial
+from pathlib import Path
+from typing import Dict, List
+
+import pandas as pd
+import pyarrow.fs
+import pytest
+
+import ray
+from ray import train
+from ray.air._internal.uri_utils import URI
+from ray.train import CheckpointConfig, RunConfig, ScalingConfig
+from ray.train.base_trainer import BaseTrainer
+from ray.train.data_parallel_trainer import DataParallelTrainer
+from ray.train.lightgbm import LightGBMTrainer
+from ray.train.tests.util import create_dict_checkpoint, load_dict_checkpoint
+from ray.train.trainer import TrainingFailedError
+from ray.train.xgboost import XGBoostTrainer
+from ray.tune import Callback
+from ray.tune.experiment import Trial
+
+
+@pytest.fixture
+def ray_start_4_cpus():
+    address_info = ray.init(num_cpus=4)
+    yield address_info
+    # The code after the yield will run as teardown code.
+    if ray.is_initialized():
+        ray.shutdown()
+
+
+@pytest.fixture
+def ray_start_6_cpus():
+    address_info = ray.init(num_cpus=6)
+    yield address_info
+    # The code after the yield will run as teardown code.
+    if ray.is_initialized():
+        ray.shutdown()
+
+
+class _TestSpecificError(RuntimeError):
+    pass
+
+
+def _failing_train_fn(config):
+    checkpoint = train.get_checkpoint()
+    it = 1
+    if checkpoint:
+        it = load_dict_checkpoint(checkpoint)["it"] + 1
+        print(f"\nLoading from checkpoint, which is at iteration {it}...\n")
+    with create_dict_checkpoint({"it": it}) as checkpoint:
+        train.report({"it": it}, checkpoint=checkpoint)
+    if it == 1:
+        raise _TestSpecificError
+
+
+class FailureInjectionCallback(Callback):
+    """Inject failure at the configured iteration number."""
+
+    def __init__(self, fail_marker_path: Path, num_iters: int = 2):
+        self.num_iters = num_iters
+        self.fail_marker_path = fail_marker_path
+
+    def on_trial_result(
+        self, iteration: int, trials: List[Trial], trial: Trial, result: Dict, **info
+    ):
+        if not self.fail_marker_path.exists():
+            return
+
+        if trial.last_result.get("training_iteration", -1) >= self.num_iters:
+            print(f"Failing after {self.num_iters} iters...")
+            self.fail_marker_path.unlink()
+            raise _TestSpecificError
+
+
+def test_data_parallel_trainer_restore(ray_start_4_cpus, tmpdir):
+    """Restoring a DataParallelTrainer with object refs captured in the train fn
+    or config works by re-specifying them.
+    Success criteria:
+    - Restored to the correct iteration. (1 iteration before crash, 1 after restore).
+    - Results are being logged to the same directory as before.
+    """
+    dataset_size = 10
+    num_workers = 2
+
+    def create_train_fn_and_config():
+        obj_ref = ray.put({"test": 1})
+
+        def train_fn(config):
+            assert ray.get(obj_ref)["test"] == 1
+            assert ray.get(config["obj_ref"])["test"] == 1
+            ds = train.get_dataset_shard("train")
+            assert (
+                sum([len(batch["feature"]) for batch in ds.iter_batches()])
+                == dataset_size // num_workers
+            )
+            _failing_train_fn(config)
+
+        train_loop_config = {"obj_ref": obj_ref}
+        return train_fn, train_loop_config
+
+    datasets = {"train": ray.data.from_items([{"feature": i} for i in range(10)])}
+    train_fn, train_loop_config = create_train_fn_and_config()
+    trainer = DataParallelTrainer(
+        train_loop_per_worker=train_fn,
+        train_loop_config=train_loop_config,
+        datasets=datasets,
+        scaling_config=ScalingConfig(num_workers=num_workers),
+        run_config=RunConfig(
+            name="data_parallel_restore_test",
+            storage_path=str(tmpdir),
+            checkpoint_config=CheckpointConfig(num_to_keep=1),
+        ),
+    )
+    with pytest.raises(TrainingFailedError) as exc_info:
+        result = trainer.fit()
+    assert isinstance(exc_info.value.__cause__, _TestSpecificError)
+
+    # Include an explicit cluster shutdown.
+    # Otherwise, the previously registered object references will still exist,
+    # and the test may trivially pass.
+    ray.shutdown()
+    ray.init(num_cpus=4)
+
+    train_fn, train_loop_config = create_train_fn_and_config()
+    datasets = {"train": ray.data.from_items([{"feature": i} for i in range(10)])}
+    trainer = DataParallelTrainer.restore(
+        str(tmpdir / "data_parallel_restore_test"),
+        train_loop_per_worker=train_fn,
+        train_loop_config=train_loop_config,
+        datasets=datasets,
+    )
+    result = trainer.fit()
+    assert not result.error
+    assert result.metrics["training_iteration"] == 2
+    assert result.metrics["iterations_since_restore"] == 1
+    assert tmpdir / "data_parallel_restore_test" in Path(result.path).parents
+
+
+@pytest.mark.parametrize("trainer_cls", [XGBoostTrainer, LightGBMTrainer])
+def test_gbdt_trainer_restore(ray_start_6_cpus, tmp_path, trainer_cls, monkeypatch):
+    """Tests restoring gradient boosted decision tree trainers.
+    Success criteria:
+    - Picks up at the right iteration. 2 before crash. 3 after. 5 total trees.
+    - Results are being logged to the same directory as before.
+    """
+    monkeypatch.setenv("TUNE_GLOBAL_CHECKPOINT_S", "0")
+    exp_name = f"{trainer_cls.__name__}_restore_test"
+    datasets = {
+        "train": ray.data.from_pandas(
+            pd.DataFrame({"x": range(100), "y": range(1, 101)})
+        )
+    }
+
+    fail_marker_path = tmp_path / "fail_marker"
+    fail_marker_path.touch()
+
+    trainer = trainer_cls(
+        label_column="y",
+        params={
+            "objective": (
+                "reg:squarederror" if trainer_cls == XGBoostTrainer else "regression"
+            )
+        },
+        datasets=datasets,
+        scaling_config=ScalingConfig(
+            num_workers=2, trainer_resources={"CPU": 0}, resources_per_worker={"CPU": 1}
+        ),
+        run_config=RunConfig(
+            storage_path=str(tmp_path),
+            name=exp_name,
+            checkpoint_config=CheckpointConfig(
+                num_to_keep=1, checkpoint_frequency=1, checkpoint_at_end=False
+            ),
+            callbacks=[FailureInjectionCallback(fail_marker_path, num_iters=2)],
+        ),
+        num_boost_round=5,
+    )
+    with pytest.raises(TrainingFailedError):
+        result = trainer.fit()
+
+    trainer = trainer_cls.restore(str(tmp_path / exp_name), datasets=datasets)
+    result = trainer.fit()
+    assert not result.error
+    assert result.metrics["training_iteration"] == 5
+    assert result.metrics["iterations_since_restore"] == 3
+    assert tmp_path / exp_name in Path(result.path).parents
+
+
+@pytest.mark.parametrize("name", [None, "restore_from_uri"])
+def test_restore_from_uri_s3(
+    ray_start_4_cpus, tmp_path, monkeypatch, mock_s3_bucket_uri, name
+):
+    """Restoration from S3 should work."""
+    trainer = DataParallelTrainer(
+        train_loop_per_worker=lambda config: train.report({"score": 1}),
+        scaling_config=ScalingConfig(num_workers=2),
+        run_config=RunConfig(name=name, storage_path=mock_s3_bucket_uri),
+    )
+    result = trainer.fit()
+
+    if name is None:
+        name = Path(result.path).parent.name
+
+    # Restore from S3
+    assert DataParallelTrainer.can_restore(str(URI(mock_s3_bucket_uri) / name))
+    DataParallelTrainer.restore(str(URI(mock_s3_bucket_uri) / name))
+
+
+def test_restore_with_datasets(ray_start_4_cpus, tmpdir):
+    """Datasets are required to re-specify if they were originally provided."""
+    datasets = {
+        "train": ray.data.from_items([{"x": x, "y": x + 1} for x in range(8)]),
+        "valid": ray.data.from_items([{"x": x, "y": x + 1} for x in range(8)]),
+    }
+
+    trainer = DataParallelTrainer(
+        train_loop_per_worker=lambda config: train.report({"score": 1}),
+        datasets=datasets,
+        scaling_config=ScalingConfig(num_workers=2),
+        run_config=RunConfig(name="datasets_respecify_test"),
+    )
+    trainer._save(pyarrow.fs.LocalFileSystem(), str(tmpdir))
+
+    # Restore should complain, if all the datasets don't get passed in again
+    with pytest.raises(ValueError):
+        DataParallelTrainer.restore(str(tmpdir))
+
+    with pytest.raises(ValueError):
+        DataParallelTrainer.restore(str(tmpdir), datasets={"train": datasets["train"]})
+
+    with pytest.raises(ValueError):
+        DataParallelTrainer.restore(
+            str(tmpdir),
+            datasets={"train": datasets["train"], "invalid_key": datasets["valid"]},
+        )
+
+    trainer = DataParallelTrainer.restore(str(tmpdir), datasets=datasets)
+
+
+def test_restore_from_invalid_dir(tmpdir):
+    """Should raise an error if the restore directory doesn't exist or is invalid."""
+    with pytest.raises(ValueError):
+        BaseTrainer.restore(str(tmpdir))
+
+    with pytest.raises(ValueError):
+        BaseTrainer.restore("mock:///not/found")
+
+
+def test_trainer_can_restore_utility(tmp_path):
+    """Make sure that `can_restore` detects an existing experiment at a
+    local/remote path and only returns True if it's at the Train experiment dir root.
+    """
+    name = "exp_name"
+    path = tmp_path / name
+
+    assert not DataParallelTrainer.can_restore(path)
+
+    trainer = DataParallelTrainer(
+        train_loop_per_worker=lambda config: train.report({"score": 1}),
+        scaling_config=ScalingConfig(num_workers=1),
+    )
+    (tmp_path / name).mkdir(exist_ok=True)
+    trainer._save(pyarrow.fs.LocalFileSystem(), str(tmp_path / name))
+
+    assert DataParallelTrainer.can_restore(path)
+
+
+@pytest.mark.parametrize("eventual_success", [True, False])
+def test_retry_with_max_failures(ray_start_4_cpus, eventual_success):
+    """Test auto-resume of a Train run when setting max_failures > 0."""
+
+    num_failures = 2 if eventual_success else 3
+    max_retries = 2
+    final_iter = 10
+
+    def train_func():
+        ckpt = train.get_checkpoint()
+        itr = 1
+        restore_count = 0
+        if ckpt:
+            ckpt = load_dict_checkpoint(ckpt)
+            itr = ckpt["iter"] + 1
+            restore_count = ckpt["restore_count"] + 1
+
+        for i in range(itr, final_iter + 1):
+            with create_dict_checkpoint(
+                dict(iter=i, restore_count=restore_count)
+            ) as checkpoint:
+                train.report(dict(test=i, training_iteration=i), checkpoint=checkpoint)
+            if restore_count < num_failures:
+                raise RuntimeError("try to fail me")
+
+    trainer = DataParallelTrainer(
+        train_func,
+        scaling_config=ScalingConfig(num_workers=2),
+        run_config=RunConfig(
+            failure_config=train.FailureConfig(max_failures=max_retries)
+        ),
+    )
+
+    if not eventual_success:
+        # If we gave up due to hitting our max retry attempts,
+        # then `trainer.fit` should raise the last error we encountered.
+        with pytest.raises(TrainingFailedError):
+            trainer.fit()
+    else:
+        # If we encounter errors but eventually succeed, `trainer.fit` should NOT
+        # raise any of those errors.
+        result = trainer.fit()
+        assert not result.error
+        checkpoint = load_dict_checkpoint(result.checkpoint)
+        assert checkpoint["iter"] == final_iter
+
+
+def test_restoration_after_termination(tmp_path):
+    """Test that the train loop can be run again if restoring the trainer
+    after the run finished running successfully."""
+
+    def train_func_per_worker(config, num_epochs=5):
+        ckpt = train.get_checkpoint()
+        start_iter = 1
+        if ckpt:
+            ckpt = load_dict_checkpoint(ckpt)
+            start_iter = ckpt["iter"] + 1
+
+        for i in range(start_iter, num_epochs + 1):
+            with create_dict_checkpoint(dict(iter=i)) as checkpoint:
+                train.report(dict(iter=i), checkpoint=checkpoint)
+
+    name = "exp_name"
+    path = tmp_path / name
+
+    trainer = DataParallelTrainer(
+        train_loop_per_worker=train_func_per_worker,
+        scaling_config=ScalingConfig(num_workers=1),
+        run_config=RunConfig(
+            name=name,
+            storage_path=tmp_path,
+            checkpoint_config=CheckpointConfig(num_to_keep=2),
+        ),
+    )
+    result = trainer.fit()
+    assert result.metrics["iter"] == 5
+
+    restored_trainer = DataParallelTrainer.restore(
+        str(path), train_loop_per_worker=partial(train_func_per_worker, num_epochs=10)
+    )
+    new_result = restored_trainer.fit()
+    assert new_result.metrics["iter"] == 10
+
+    assert new_result.path == result.path
+    assert len(list(Path(new_result.path).glob("checkpoint*"))) == 2
+
+
+if __name__ == "__main__":
+    import sys
+
+    sys.exit(pytest.main(["-v", "-x", __file__]))

python/ray/train/tests/test_training_iterator.py

@@ -0,0 +1,395 @@
+import functools
+import sys
+import time
+from contextlib import nullcontext
+from unittest.mock import patch
+
+import pytest
+
+import ray
+from ray import train
+from ray.air._internal.util import StartTraceback
+from ray.train import DataConfig
+from ray.train._internal.backend_executor import BackendExecutor
+from ray.train._internal.session import get_session, init_session
+from ray.train._internal.utils import construct_train_func
+from ray.train._internal.worker_group import WorkerGroup
+from ray.train.backend import BackendConfig
+from ray.train.examples.pytorch.torch_linear_example import (
+    train_func as linear_train_func,
+)
+from ray.train.tests.util import mock_storage_context
+from ray.train.trainer import TrainingIterator
+
+MAX_RETRIES = 3
+
+
+@pytest.fixture(autouse=True, scope="module")
+def patch_tune_session():
+    if not get_session():
+        init_session(
+            training_func=None,
+            world_rank=None,
+            local_rank=None,
+            node_rank=None,
+            local_world_size=None,
+            world_size=None,
+            storage=mock_storage_context(),
+        )
+    yield
+
+
+@pytest.fixture
+def ray_start_4_cpus():
+    address_info = ray.init(num_cpus=4)
+    yield address_info
+    # The code after the yield will run as teardown code.
+    ray.shutdown()
+
+
+def gen_execute_single_async_special(special_f):
+    def execute_single_async_special(self, i, f, *args, **kwargs):
+        assert len(self.workers) == 2
+        if i == 0 and hasattr(self, "should_fail") and self.should_fail:
+            kwargs["train_func"] = special_f
+        return (
+            self.workers[i]
+            .actor._RayTrainWorker__execute.options(name=f.__name__)
+            .remote(f, *args, **kwargs)
+        )
+
+    return execute_single_async_special
+
+
+def gen_new_backend_executor(special_f):
+    """Returns a BackendExecutor that runs special_f on worker 0 once."""
+
+    class TestBackendExecutor(BackendExecutor):
+        _has_failed = False
+
+        def start_training(self, *args, **kwargs):
+            special_execute = gen_execute_single_async_special(special_f)
+            if not self._has_failed:
+                self.worker_group.should_fail = True
+                self._has_failed = True
+            else:
+                self.worker_group.should_fail = False
+            with patch.object(WorkerGroup, "execute_single_async", special_execute):
+                super().start_training(*args, **kwargs)
+
+    return TestBackendExecutor
+
+
+def create_iterator(
+    train_func,
+    backend_config,
+    *,
+    num_workers=2,
+    backend_executor_cls=BackendExecutor,
+    init_hook=None,
+):
+    # Similar logic to the old Trainer.run_iterator().
+
+    train_func = construct_train_func(train_func, None, train_func_context=nullcontext)
+
+    backend_executor = backend_executor_cls(
+        backend_config=backend_config, num_workers=num_workers, max_retries=MAX_RETRIES
+    )
+    backend_executor.start(init_hook)
+
+    return TrainingIterator(
+        backend_executor=backend_executor,
+        backend_config=backend_config,
+        train_func=train_func,
+        datasets={},
+        metadata={},
+        data_config=DataConfig(),
+        checkpoint=None,
+    )
+
+
+def test_run_iterator(ray_start_4_cpus):
+    config = BackendConfig()
+
+    def train_func():
+        for i in range(3):
+            train.report(dict(index=i))
+        return 1
+
+    iterator = create_iterator(train_func, config)
+
+    count = 0
+    for results in iterator:
+        assert all(value.metrics["index"] == count for value in results)
+        count += 1
+
+    assert count == 3
+    assert iterator.is_finished()
+
+    with pytest.raises(StopIteration):
+        next(iterator)
+
+
+def test_run_iterator_error(ray_start_4_cpus):
+    config = BackendConfig()
+
+    def fail_train():
+        raise NotImplementedError
+
+    iterator = create_iterator(fail_train, config)
+
+    with pytest.raises(StartTraceback) as exc:
+        next(iterator)
+
+    assert isinstance(exc.value.__cause__, NotImplementedError), (
+        exc.value,
+        exc.value.__cause__,
+    )
+
+    assert iterator.is_finished()
+
+
+def test_worker_failure_1(ray_start_4_cpus):
+    def train_func():
+        train.report({"test": 1})
+
+    def train_actor_failure():
+        import sys
+
+        sys.exit(1)
+
+    new_backend_executor_cls = gen_new_backend_executor(train_actor_failure)
+
+    config = BackendConfig()
+
+    iterator = create_iterator(
+        train_func, config, backend_executor_cls=new_backend_executor_cls
+    )
+    for worker_results in iterator:
+        assert all(result.metrics["test"] == 1 for result in worker_results)
+
+
+def test_worker_failure_2(ray_start_4_cpus):
+    def train_func():
+        for _ in range(2):
+            train.report(dict(loss=1))
+
+    def train_actor_failure():
+        for _ in range(2):
+            train.report(dict(loss=1))
+        import sys
+
+        sys.exit(1)
+
+    new_backend_executor_cls = gen_new_backend_executor(train_actor_failure)
+
+    config = BackendConfig()
+
+    iterator = create_iterator(
+        train_func, config, backend_executor_cls=new_backend_executor_cls
+    )
+    for worker_results in iterator:
+        assert all(result.metrics["loss"] == 1 for result in worker_results)
+
+
+def test_worker_failure_local_rank(ray_start_4_cpus):
+    def train_func():
+        train.report({"rank": train.get_context().get_local_rank()})
+
+    def train_actor_failure():
+        import sys
+
+        sys.exit(1)
+
+    new_backend_executor_cls = gen_new_backend_executor(train_actor_failure)
+
+    config = BackendConfig()
+
+    iterator = create_iterator(
+        train_func, config, backend_executor_cls=new_backend_executor_cls
+    )
+    for worker_results in iterator:
+        assert {result.metrics["rank"] for result in worker_results} == {0, 1}
+
+
+def test_worker_start_failure(ray_start_4_cpus):
+    def init_hook():
+        pass
+
+    def init_hook_fail():
+        ray.actor.exit_actor()
+
+    class TestBackendExecutor(BackendExecutor):
+        def __init__(self, *args, **kwargs):
+            super().__init__(*args, **kwargs)
+
+        def _restart(self):
+            self._initialization_hook = init_hook
+            super()._restart()
+
+    config = BackendConfig()
+
+    iterator = create_iterator(
+        lambda x: 1,
+        config,
+        backend_executor_cls=TestBackendExecutor,
+        init_hook=init_hook_fail,
+    )
+    assert len(iterator._backend_executor.get_worker_group()) == 2
+
+
+def test_max_failures(ray_start_4_cpus):
+    def train_func():
+        import sys
+
+        sys.exit(1)
+
+    config = BackendConfig()
+
+    iterator = create_iterator(train_func, config)
+    with pytest.raises(RuntimeError):
+        for _ in iterator:
+            pass
+    assert iterator._backend_executor._get_num_failures() == MAX_RETRIES
+
+
+def test_start_max_failures(ray_start_4_cpus):
+    def init_hook_fail():
+        import sys
+
+        sys.exit(1)
+
+    config = BackendConfig()
+
+    with pytest.raises(RuntimeError):
+        create_iterator(lambda x: 1, config, init_hook=init_hook_fail)
+
+
+class KillCallback:
+    def __init__(self, fail_on, backend_executor):
+        self.counter = 0
+        self.fail_on = fail_on
+        self.worker_group = backend_executor.get_worker_group()
+        self.results = []
+
+    def handle_result(self, intermiedate_results=None):
+        if intermiedate_results:
+            self.results.append(intermiedate_results)
+        if self.counter == self.fail_on:
+            print("killing")
+            self.results = []
+            ray.kill(self.worker_group.workers[0].actor)
+            time.sleep(3)
+        self.counter += 1
+
+
+@pytest.mark.parametrize(
+    "backend",
+    ["test", "torch", "tf"] if sys.version_info < (3, 12) else ["test", "torch"],
+)
+def test_worker_kill(ray_start_4_cpus, backend):
+    if backend == "test":
+        test_config = BackendConfig()
+    elif backend == "torch":
+        from ray.train.torch import TorchConfig
+
+        test_config = TorchConfig()
+    elif backend == "tf":
+        from ray.train.tensorflow import TensorflowConfig
+
+        test_config = TensorflowConfig()
+
+    def train_func():
+        for i in range(2):
+            train.report(dict(loss=1, iter=i))
+
+    iterator = create_iterator(train_func, test_config)
+    kill_callback = KillCallback(fail_on=0, backend_executor=iterator._backend_executor)
+
+    for intermediate_result in iterator:
+        # Run 1: iter=0, counter=1, Successful
+        # Run 2: iter=1, counter=1, Unsuccessful, starts training from beginning
+        # Run 3: iter=0, counter=2, Successful
+        # Run 4: iter=1, counter=3, Successful
+        kill_callback.handle_result()
+    assert kill_callback.counter == 3
+
+    iterator = create_iterator(train_func, test_config)
+    kill_callback = KillCallback(fail_on=1, backend_executor=iterator._backend_executor)
+    for intermediate_result in iterator:
+        # Run 1: iter=0, counter=1, Successful
+        # Run 2: iter=1, counter=2, Successful
+        # Run 3: None, counter=2, Unsuccessful, starts training from beginning.
+        # Run 4: iter=0, counter=3, Successful
+        # Run 5: iter=1, counter=4, Successful
+        kill_callback.handle_result()
+    assert kill_callback.counter == 4
+
+
+@pytest.mark.skipif(
+    sys.version_info >= (3, 12), reason="tensorflow is not installed in python 3.12+"
+)
+def test_tensorflow_mnist_fail(ray_start_4_cpus):
+    """Tests if tensorflow example works even with worker failure."""
+    epochs = 3
+    num_workers = 2
+
+    from ray.train.examples.tf.tensorflow_mnist_example import (
+        train_func as tensorflow_mnist_train_func,
+    )
+    from ray.train.tensorflow import TensorflowConfig
+
+    test_config = TensorflowConfig()
+
+    train_func = functools.partial(
+        tensorflow_mnist_train_func, {"lr": 1e-3, "batch_size": 64, "epochs": epochs}
+    )
+    iterator = create_iterator(train_func, test_config, num_workers=num_workers)
+    kill_callback = KillCallback(fail_on=0, backend_executor=iterator._backend_executor)
+
+    for intermediate_result in iterator:
+        assert len(intermediate_result) == num_workers
+        kill_callback.handle_result(intermediate_result)
+
+    results = kill_callback.results
+    assert len(results) == epochs
+    last_iter_result = results[-1][0].metrics
+    first_iter_result = results[0][0].metrics
+
+    assert last_iter_result["loss"] < first_iter_result["loss"]
+    assert last_iter_result["accuracy"] > first_iter_result["accuracy"]
+
+
+def test_torch_linear_failure(ray_start_4_cpus):
+    num_workers = 2
+    epochs = 3
+
+    from ray.train.torch import TorchConfig
+
+    test_config = TorchConfig()
+
+    train_func = functools.partial(
+        linear_train_func, {"lr": 1e-3, "batch_size": 64, "epochs": epochs}
+    )
+
+    iterator = create_iterator(train_func, test_config, num_workers=num_workers)
+    kill_callback = KillCallback(fail_on=1, backend_executor=iterator._backend_executor)
+
+    for intermediate_result in iterator:
+        assert len(intermediate_result) == num_workers
+        kill_callback.handle_result(intermediate_result)
+
+    results = kill_callback.results
+    assert len(results) == epochs
+    for i in range(num_workers):
+        last_result = results[-1][i].metrics
+        first_result = results[0][i].metrics
+        assert last_result["loss"] < first_result["loss"]
+
+
+if __name__ == "__main__":
+    import sys
+
+    import pytest
+
+    sys.exit(pytest.main(sys.argv[1:] + ["-v", "-x", __file__]))

python/ray/train/tests/test_tune.py

@@ -0,0 +1,331 @@
+import logging
+import sys
+
+import pytest
+
+import ray
+from ray import train, tune
+from ray.air.constants import TRAINING_ITERATION
+from ray.train._internal.worker_group import WorkerGroup
+from ray.train.backend import Backend, BackendConfig
+from ray.train.data_parallel_trainer import DataParallelTrainer
+from ray.train.examples.pytorch.torch_fashion_mnist_example import (
+    train_func_per_worker as fashion_mnist_train_func,
+)
+from ray.train.tests.util import create_dict_checkpoint, load_dict_checkpoint
+from ray.train.torch import TorchTrainer
+from ray.tune.tune_config import TuneConfig
+from ray.tune.tuner import Tuner
+
+
+@pytest.fixture(scope="module")
+def ray_start_4_cpus():
+    address_info = ray.init(num_cpus=4)
+    yield address_info
+    # The code after the yield will run as teardown code.
+    ray.shutdown()
+
+
+@pytest.fixture
+def ray_start_8_cpus():
+    address_info = ray.init(num_cpus=8)
+    yield address_info
+    # The code after the yield will run as teardown code.
+    ray.shutdown()
+
+
+class TestConfig(BackendConfig):
+    @property
+    def backend_cls(self):
+        return TestBackend
+
+
+class TestBackend(Backend):
+    def on_start(self, worker_group: WorkerGroup, backend_config: TestConfig):
+        pass
+
+    def on_shutdown(self, worker_group: WorkerGroup, backend_config: TestConfig):
+        pass
+
+
+def torch_fashion_mnist(num_workers, use_gpu, num_samples):
+    trainer = TorchTrainer(
+        fashion_mnist_train_func,
+        scaling_config=train.ScalingConfig(num_workers=num_workers, use_gpu=use_gpu),
+    )
+    tuner = Tuner(
+        trainer,
+        param_space={
+            "train_loop_config": {
+                "lr": tune.loguniform(1e-4, 1e-1),
+                "batch_size_per_worker": tune.choice([32, 64, 128]),
+                "epochs": 2,
+            }
+        },
+        tune_config=TuneConfig(
+            num_samples=num_samples,
+        ),
+    )
+    analysis = tuner.fit()._experiment_analysis
+
+    # Check that loss decreases in each trial.
+    for df in analysis.trial_dataframes.values():
+        assert df.loc[1, "loss"] < df.loc[0, "loss"]
+
+
+def test_tune_torch_fashion_mnist(ray_start_8_cpus):
+    torch_fashion_mnist(num_workers=2, use_gpu=False, num_samples=2)
+
+
+@pytest.mark.skipif(
+    sys.version_info >= (3, 12), reason="tensorflow is not installed in python 3.12+"
+)
+def tune_tensorflow_mnist(num_workers, use_gpu, num_samples):
+    from ray.train.examples.tf.tensorflow_mnist_example import (
+        train_func as tensorflow_mnist_train_func,
+    )
+    from ray.train.tensorflow import TensorflowTrainer
+
+    trainer = TensorflowTrainer(
+        tensorflow_mnist_train_func,
+        scaling_config=train.ScalingConfig(num_workers=num_workers, use_gpu=use_gpu),
+    )
+    tuner = Tuner(
+        trainer,
+        param_space={
+            "train_loop_config": {
+                "lr": tune.loguniform(1e-4, 1e-1),
+                "batch_size": tune.choice([32, 64, 128]),
+                "epochs": 2,
+            }
+        },
+        tune_config=TuneConfig(
+            num_samples=num_samples,
+        ),
+    )
+    analysis = tuner.fit()._experiment_analysis
+
+    # Check that loss decreases in each trial.
+    for df in analysis.trial_dataframes.values():
+        assert df.loc[1, "loss"] < df.loc[0, "loss"]
+
+
+@pytest.mark.skipif(
+    sys.version_info >= (3, 12), reason="tensorflow is not installed in python 3.12+"
+)
+def test_tune_tensorflow_mnist(ray_start_8_cpus):
+    tune_tensorflow_mnist(num_workers=2, use_gpu=False, num_samples=2)
+
+
+def test_tune_error(ray_start_4_cpus):
+    def train_func(config):
+        raise RuntimeError("Error in training function!")
+
+    trainer = DataParallelTrainer(
+        train_func,
+        backend_config=TestConfig(),
+        scaling_config=train.ScalingConfig(num_workers=1),
+    )
+    tuner = Tuner(
+        trainer,
+    )
+
+    result_grid = tuner.fit()
+    with pytest.raises(RuntimeError):
+        raise result_grid[0].error
+
+
+def test_tune_checkpoint(ray_start_4_cpus):
+    def train_func():
+        for i in range(9):
+            train.report(dict(test=i))
+        with create_dict_checkpoint(dict(hello="world")) as checkpoint:
+            train.report(dict(test=i + 1), checkpoint=checkpoint)
+
+    trainer = DataParallelTrainer(
+        train_func,
+        backend_config=TestConfig(),
+        scaling_config=train.ScalingConfig(num_workers=1),
+    )
+    tuner = Tuner(
+        trainer,
+        param_space={"train_loop_config": {"max_iter": 5}},
+    )
+
+    result_grid = tuner.fit()
+    assert len(result_grid) == 1
+    result = result_grid[0]
+    assert result.checkpoint
+    assert load_dict_checkpoint(result.checkpoint)["hello"] == "world"
+
+
+def test_reuse_checkpoint(ray_start_4_cpus):
+    def train_func(config):
+        itr = 0
+        ckpt = train.get_checkpoint()
+        if ckpt is not None:
+            ckpt = load_dict_checkpoint(ckpt)
+            itr = ckpt["iter"] + 1
+
+        for i in range(itr, config["max_iter"]):
+            with create_dict_checkpoint(dict(iter=i)) as checkpoint:
+                train.report(dict(test=i, training_iteration=i), checkpoint=checkpoint)
+
+    trainer = DataParallelTrainer(
+        train_func,
+        backend_config=TestConfig(),
+        scaling_config=train.ScalingConfig(num_workers=1),
+    )
+    tuner = Tuner(
+        trainer,
+        param_space={"train_loop_config": {"max_iter": 5}},
+    )
+    result_grid = tuner.fit()
+    assert len(result_grid) == 1
+    result = result_grid[0]
+    assert result.checkpoint
+    assert load_dict_checkpoint(result.checkpoint)["iter"] == 4
+
+    tuner = Tuner.restore(result_grid.experiment_path, trainable=trainer)
+    result_grid = tuner.fit()
+    assert len(result_grid) == 1
+    assert len(result_grid[0].metrics_dataframe) == 5
+
+
+def test_retry_with_max_failures(ray_start_4_cpus):
+    """Tests trainer retry with max_failures > 0 when integrating with Tune."""
+
+    def train_func():
+        ckpt = train.get_checkpoint()
+        restored = bool(ckpt)  # Does a previous checkpoint exist?
+        itr = 0
+        if ckpt:
+            ckpt = load_dict_checkpoint(ckpt)
+            itr = ckpt["iter"] + 1
+
+        for i in range(itr, 4):
+            if i == 2 and not restored:
+                raise Exception("try to fail me")
+            with create_dict_checkpoint(dict(iter=i)) as checkpoint:
+                train.report(dict(test=i, training_iteration=i), checkpoint=checkpoint)
+
+    trainer = DataParallelTrainer(
+        train_func,
+        backend_config=TestConfig(),
+        scaling_config=train.ScalingConfig(num_workers=1),
+    )
+    tuner = Tuner(
+        trainer,
+        run_config=tune.RunConfig(failure_config=tune.FailureConfig(max_failures=3)),
+    )
+
+    result_grid = tuner.fit()
+    checkpoint = load_dict_checkpoint(result_grid[0].checkpoint)
+    assert checkpoint["iter"] == 3
+    df = result_grid[0].metrics_dataframe
+    assert len(df[TRAINING_ITERATION]) == 4
+
+
+def test_restore_with_new_trainer(ray_start_4_cpus, tmpdir, propagate_logs, caplog):
+    def train_func(config):
+        raise RuntimeError("failing!")
+
+    trainer = DataParallelTrainer(
+        train_func,
+        backend_config=TestConfig(),
+        scaling_config=train.ScalingConfig(num_workers=1),
+        run_config=train.RunConfig(
+            name="restore_new_trainer", storage_path=str(tmpdir)
+        ),
+        datasets={"train": ray.data.from_items([{"a": i} for i in range(10)])},
+    )
+    results = Tuner(trainer).fit()
+    assert results.errors
+
+    def train_func(config):
+        dataset = train.get_dataset_shard("train")
+        assert train.get_context().get_world_size() == 2
+        rows = 0
+        for _ in dataset.iter_rows():
+            rows += 1
+        assert rows == 10
+
+    trainer = DataParallelTrainer(
+        # Training function can be modified
+        train_func,
+        backend_config=TestConfig(),
+        # ScalingConfig can be modified
+        scaling_config=train.ScalingConfig(num_workers=2),
+        # New RunConfig will be ignored
+        run_config=train.RunConfig(name="ignored"),
+        # Datasets and preprocessors can be re-specified
+        datasets={"train": ray.data.from_items([{"a": i} for i in range(20)])},
+    )
+    caplog.clear()
+    with caplog.at_level(logging.WARNING, logger="ray.tune.impl.tuner_internal"):
+        tuner = Tuner.restore(
+            str(tmpdir / "restore_new_trainer"),
+            trainable=trainer,
+            resume_errored=True,
+        )
+        assert "they will be ignored in the resumed run" in caplog.text
+
+    results = tuner.fit()
+    assert not results.errors
+
+
+@pytest.mark.parametrize("in_trainer", [True, False])
+@pytest.mark.parametrize("in_tuner", [True, False])
+def test_run_config_in_trainer_and_tuner(
+    propagate_logs, tmp_path, caplog, in_trainer, in_tuner
+):
+    trainer_run_config = (
+        train.RunConfig(name="trainer", storage_path=str(tmp_path))
+        if in_trainer
+        else None
+    )
+    tuner_run_config = (
+        tune.RunConfig(name="tuner", storage_path=str(tmp_path)) if in_tuner else None
+    )
+    trainer = DataParallelTrainer(
+        lambda config: None,
+        backend_config=TestConfig(),
+        scaling_config=train.ScalingConfig(num_workers=1),
+        run_config=trainer_run_config,
+    )
+    with caplog.at_level(logging.INFO, logger="ray.tune.impl.tuner_internal"):
+        tuner = Tuner(trainer, run_config=tuner_run_config)
+
+    both_msg = (
+        "`RunConfig` was passed to both the `Tuner` and the `DataParallelTrainer`"
+    )
+    run_config = tuner._local_tuner.get_run_config()
+    if in_trainer and in_tuner:
+        assert run_config.name == "tuner"
+        assert both_msg in caplog.text
+    elif in_trainer and not in_tuner:
+        assert run_config.name == "trainer"
+        assert both_msg not in caplog.text
+    elif not in_trainer and in_tuner:
+        assert run_config.name == "tuner"
+        assert both_msg not in caplog.text
+    else:
+        assert both_msg not in caplog.text
+
+
+def test_run_config_in_param_space():
+    trainer = DataParallelTrainer(
+        lambda config: None,
+        backend_config=TestConfig(),
+        scaling_config=train.ScalingConfig(num_workers=1),
+    )
+    with pytest.raises(ValueError):
+        Tuner(trainer, param_space={"run_config": train.RunConfig(name="ignored")})
+
+
+if __name__ == "__main__":
+    import sys
+
+    import pytest
+
+    sys.exit(pytest.main(["-v", "-x", __file__]))

python/ray/train/tests/test_windows.py

@@ -0,0 +1,51 @@
+"""This is a very minimal set of windows tests for Train/Tune."""
+
+import os
+
+import pytest
+
+import ray
+from ray import train
+from ray.train.tests.util import create_dict_checkpoint
+from ray.train.v2.api.data_parallel_trainer import DataParallelTrainer
+
+
+@pytest.fixture
+def ray_start_4_cpus():
+    address_info = ray.init(num_cpus=4)
+    yield address_info
+    # The code after the yield will run as teardown code.
+    ray.shutdown()
+
+
+@pytest.fixture
+def chdir_tmpdir(tmp_path):
+    original_path = os.getcwd()
+    os.chdir(tmp_path)
+    yield
+    os.chdir(original_path)
+
+
+def test_storage_path(ray_start_4_cpus, chdir_tmpdir):
+    """Tests that Train with a local storage path works on Windows."""
+
+    def train_fn(config):
+        for i in range(5):
+            if train.get_context().get_world_rank() == 0:
+                with create_dict_checkpoint({"dummy": "data"}) as checkpoint:
+                    train.report({"loss": i}, checkpoint=checkpoint)
+            else:
+                train.report({"loss": i})
+
+    trainer = DataParallelTrainer(
+        train_fn,
+        scaling_config=train.ScalingConfig(num_workers=2),
+        run_config=train.RunConfig(storage_path=os.getcwd()),
+    )
+    trainer.fit()
+
+
+if __name__ == "__main__":
+    import sys
+
+    sys.exit(pytest.main(["-v", "-x", __file__]))

python/ray/train/tests/test_worker_group.py

@@ -0,0 +1,352 @@
+import time
+from collections import defaultdict
+from typing import Dict, List, Optional
+
+import pytest
+
+import ray
+import ray._private.ray_constants as ray_constants
+from ray.cluster_utils import Cluster
+from ray.train._internal.worker_group import Worker, WorkerGroup, WorkerMetadata
+from ray.util.state import list_actors
+
+
+@pytest.fixture
+def ray_start_2_cpus():
+    address_info = ray.init(num_cpus=2)
+    yield address_info
+    # The code after the yield will run as teardown code.
+    ray.shutdown()
+
+
+@pytest.fixture
+def ray_start_2_cpus_and_gpus():
+    address_info = ray.init(num_cpus=2, num_gpus=2)
+    yield address_info
+    # The code after the yield will run as teardown code.
+    ray.shutdown()
+
+
+@pytest.fixture
+def ray_start_2_cpus_and_neuron_core_accelerator():
+    address_info = ray.init(num_cpus=2, resources={ray_constants.NEURON_CORES: 2})
+    yield address_info
+    # The code after the yield will run as teardown code.
+    ray.shutdown()
+
+
+@pytest.fixture
+def ray_start_2_cpus_and_10kb_memory():
+    address_info = ray.init(num_cpus=2, _memory=10_000)
+    yield address_info
+    # The code after the yield will run as teardown code.
+    ray.shutdown()
+
+
+@pytest.fixture
+def ray_start_5_nodes_with_memory():
+    cluster = Cluster()
+    for _ in range(4):
+        cluster.add_node(num_cpus=4, memory=500)
+    cluster.add_node(num_cpus=4, memory=2_000)
+
+    ray.init(address=cluster.address)
+
+    yield
+
+    ray.shutdown()
+    cluster.shutdown()
+
+
+def test_worker_creation(ray_start_2_cpus):
+    assert ray.available_resources()["CPU"] == 2
+    wg = WorkerGroup(num_workers=2)
+    assert len(wg.workers) == 2
+    time.sleep(1)
+    # Make sure both CPUs are being used by the actors.
+    assert "CPU" not in ray.available_resources()
+    wg.shutdown()
+
+
+def test_worker_creation_num_cpus(ray_start_2_cpus):
+    assert ray.available_resources()["CPU"] == 2
+    wg = WorkerGroup(resources_per_worker={"CPU": 2})
+    time.sleep(1)
+    assert len(wg.workers) == 1
+    # Make sure both CPUs are being used by the actor.
+    assert "CPU" not in ray.available_resources()
+    wg.shutdown()
+
+
+def test_worker_creation_with_memory(ray_start_5_nodes_with_memory):
+    resources_per_worker = {"memory": 1_000}
+    wg = WorkerGroup(num_workers=2, resources_per_worker=resources_per_worker)
+    assert len(wg.workers) == 2
+
+    nodes = ray.nodes()
+    large_node = [node for node in nodes if node["Resources"]["memory"] == 2_000][0]
+    large_node_id = large_node["NodeID"]
+
+    def validate_scheduling():
+        resources = ray.get_runtime_context().get_assigned_resources()
+        assert resources == resources_per_worker, "Resources should include memory."
+
+        node_id = ray.get_runtime_context().get_node_id()
+        assert (
+            node_id == large_node_id
+        ), "Workers should be scheduled on the large node."
+
+    wg.execute(validate_scheduling)
+
+
+def test_worker_shutdown(ray_start_2_cpus):
+    assert ray.available_resources()["CPU"] == 2
+    wg = WorkerGroup(num_workers=2)
+    time.sleep(1)
+    assert "CPU" not in ray.available_resources()
+    assert len(list_actors()) == 2
+    wg.shutdown()
+    time.sleep(1)
+    assert ray.available_resources()["CPU"] == 2
+
+    with pytest.raises(RuntimeError):
+        wg.execute(lambda: 1)
+
+
+def test_worker_restart(ray_start_2_cpus):
+    wg = WorkerGroup(num_workers=2)
+    with pytest.raises(RuntimeError):
+        wg.start()
+    # Avoid race condition.
+    time.sleep(1)
+    wg.shutdown(0)
+    wg.start()
+    wg.execute(lambda: 1)
+
+
+def test_worker_with_gpu_ids(ray_start_2_cpus_and_gpus):
+    num_gpus = 2
+    wg = WorkerGroup(num_workers=2, resources_per_worker={"GPU": 1})
+    assert len(wg.workers) == 2
+    time.sleep(1)
+    assert ray_constants.GPU not in ray.available_resources()
+    wg.execute(lambda: 1)
+    assert len(wg.workers) == 2
+    for w in wg.workers:
+        resource_ids = w.metadata.resource_ids
+        gpu_ids = resource_ids[ray_constants.GPU]
+        for gpu_id in gpu_ids:
+            assert gpu_id in [str(i) for i in range(num_gpus)]
+        assert len(resource_ids[ray_constants.NEURON_CORES]) == 0
+
+
+def test_worker_with_neuron_core_accelerator_ids(
+    ray_start_2_cpus_and_neuron_core_accelerator,
+):
+    num_nc = 2
+    wg = WorkerGroup(
+        num_workers=2, resources_per_worker={ray_constants.NEURON_CORES: 1}
+    )
+    assert len(wg.workers) == 2
+    time.sleep(1)
+    assert ray_constants.NEURON_CORES not in ray.available_resources()
+    wg.execute(lambda: 1)
+    assert len(wg.workers) == 2
+    for w in wg.workers:
+        resource_ids = w.metadata.resource_ids
+        assert len(resource_ids[ray_constants.GPU]) == 0
+        neuron_core_ids = resource_ids[ray_constants.NEURON_CORES]
+        for neuron_core_id in neuron_core_ids:
+            assert neuron_core_id in [str(i) for i in range(num_nc)]
+
+
+def test_execute_async(ray_start_2_cpus):
+    wg = WorkerGroup(num_workers=2)
+    futures = wg.execute_async(lambda: 1)
+    assert len(futures) == 2
+    outputs = ray.get(futures)
+    assert all(o == 1 for o in outputs)
+
+
+def test_execute(ray_start_2_cpus):
+    wg = WorkerGroup(num_workers=2)
+    outputs = wg.execute(lambda: 1)
+    assert len(outputs) == 2
+    assert all(o == 1 for o in outputs)
+
+
+def test_execute_args(ray_start_2_cpus):
+    wg = WorkerGroup(num_workers=2)
+    outputs = wg.execute(lambda x: x, 1)
+    assert len(outputs) == 2
+    assert all(o == 1 for o in outputs)
+
+
+def test_group_workers_by_node_id(ray_start_2_cpus):
+    def create_worker_group(node_ids):
+        wg = WorkerGroup(num_workers=2)
+        wg.workers = [
+            Worker(
+                actor=None,
+                metadata=WorkerMetadata(
+                    node_id=node_id,
+                    node_ip="dummy",
+                    hostname="dummy",
+                    resource_ids={},
+                    pid=0,
+                ),
+            )
+            for node_id in node_ids
+        ]
+        return wg
+
+    wg = create_worker_group(["2", "3", "1", "4", "2", "1", "3", "3", "4", "2"])
+    wg.sort_workers_by_node_id_and_gpu_id()
+    expected = ["2", "2", "2", "3", "3", "3", "1", "1", "4", "4"]
+    node_ids = [w.metadata.node_id for w in wg.workers]
+    assert node_ids == expected, (
+        "Workers should be grouped by Node ID "
+        "and follow the same original order of IDs encountered (2, 3, 1, 4)."
+    )
+
+    wg = create_worker_group(["2", "3", "1", "4", "2", "1", "3", "3", "4", "2"])
+    wg.sort_workers_by_node_id_and_gpu_id(_first_node_id="1")
+    expected = ["1", "1", "2", "2", "2", "3", "3", "3", "4", "4"]
+    node_ids = [w.metadata.node_id for w in wg.workers]
+    assert (
+        node_ids == expected
+    ), "Workers should be grouped by Node ID, with the first ID being 1."
+
+
+def test_sort_local_workers_by_gpu_id(ray_start_2_cpus):
+    def create_worker_group(pids, node_ids, gpu_ids):
+        wg = WorkerGroup(num_workers=2)
+        wg.workers = [
+            Worker(
+                actor=None,
+                metadata=WorkerMetadata(
+                    node_id=node_id,
+                    node_ip="dummy",
+                    hostname="dummy",
+                    resource_ids={"GPU": gpu_id.split() if gpu_id else []},
+                    pid=pid,
+                ),
+            )
+            for pid, node_id, gpu_id in zip(pids, node_ids, gpu_ids)
+        ]
+        return wg
+
+    def setup_and_check_worker_group(
+        pids: List[int],
+        node_ids: List[str],
+        gpu_ids: List[Optional[str]],
+        expected_local_ranks: Dict[int, int],
+    ):
+        """
+        Create a worker group, group workers by Node ID,
+        and check local ranks assignment.
+
+        Args:
+            pids: List of unique process IDs.
+            node_ids: List of Node IDs corresponding to each PID.
+            gpu_ids: List of GPU IDs or None for each PID.
+            expected_local_ranks: Dictionary mapping PID to the
+                expected local rank.
+        """
+        wg = create_worker_group(pids=pids, node_ids=node_ids, gpu_ids=gpu_ids)
+        wg.sort_workers_by_node_id_and_gpu_id()
+
+        # Build local ranks according to the logics in
+        # `BackendExecutor._create_rank_world_size_mappings()`
+        node_id_dict = defaultdict(int)
+        local_ranks_map = defaultdict(int)
+        for w in wg.workers:
+            local_ranks_map[w.metadata.pid] = node_id_dict[w.metadata.node_id]
+            node_id_dict[w.metadata.node_id] += 1
+
+        local_ranks = [local_ranks_map[pid] for pid in pids]
+
+        assert (
+            local_ranks == expected_local_ranks
+        ), "Incorrect local ranks allocation!\n"
+        f"Expect: {expected_local_ranks}\nGot: {local_ranks}"
+
+    # Define the worker configurations for different scenarios
+    # For workers without GPU resources, their original order will be preserved
+    cpu_workers_config = {
+        "pids": [0, 1, 2, 3, 4, 5, 6, 7],
+        "node_ids": ["2", "2", "1", "1", "2", "1", "1", "2"],
+        "gpu_ids": [None] * 8,
+        "expected_local_ranks": [0, 1, 0, 1, 2, 2, 3, 3],
+    }
+
+    gpu_workers_single_gpu_config = {
+        "pids": [0, 1, 2, 3, 4, 5, 6, 7],
+        "node_ids": ["2", "2", "1", "1", "2", "1", "1", "2"],
+        "gpu_ids": ["1", "0", "3", "2", "2", "0", "1", "3"],
+        "expected_local_ranks": [1, 0, 3, 2, 2, 0, 1, 3],
+    }
+
+    # For workers with multiple gpus, sort by their lowest gpu id
+    gpu_workers_multiple_gpus_config = {
+        "pids": [0, 1, 2, 3],
+        "node_ids": ["2", "1", "1", "2"],
+        "gpu_ids": ["1,3", "2,1", "0,3", "0,2"],
+        "expected_local_ranks": [1, 1, 0, 0],
+    }
+
+    # Setup and check worker groups for each configuration
+    setup_and_check_worker_group(**cpu_workers_config)
+    setup_and_check_worker_group(**gpu_workers_single_gpu_config)
+    setup_and_check_worker_group(**gpu_workers_multiple_gpus_config)
+
+
+def test_execute_single(ray_start_2_cpus):
+    wg = WorkerGroup(num_workers=2)
+
+    def f():
+        import os
+
+        os.environ["TEST"] = "1"
+
+    wg.execute_single(1, f)
+
+    def check():
+        import os
+
+        return os.environ.get("TEST", "0")
+
+    assert wg.execute(check) == ["0", "1"]
+
+
+def test_bad_resources(ray_start_2_cpus):
+    with pytest.raises(ValueError):
+        WorkerGroup(num_workers=-1)
+
+    with pytest.raises(ValueError):
+        WorkerGroup(resources_per_worker={"CPU": -1})
+
+    with pytest.raises(ValueError):
+        WorkerGroup(resources_per_worker={"GPU": -1})
+
+    with pytest.raises(ValueError):
+        WorkerGroup(resources_per_worker={"memory": -1})
+
+
+def test_placement_group(ray_start_2_cpus):
+    """Tests that workers can be removed and added to a placement group."""
+    num_workers = 2
+    bundle = {"CPU": 1}
+    bundles = [bundle.copy() for _ in range(num_workers)]
+    placement_group = ray.util.placement_group(bundles)
+    wg = WorkerGroup(num_workers=num_workers, placement_group=placement_group)
+    wg.remove_workers([0])
+    wg.add_workers(1)
+
+
+if __name__ == "__main__":
+    import sys
+
+    import pytest
+
+    sys.exit(pytest.main(["-v", "-x", __file__]))

python/ray/train/tests/test_xgboost_trainer.py

@@ -0,0 +1,223 @@
+from unittest import mock
+
+import pandas as pd
+import pytest
+import xgboost as xgb
+from sklearn.datasets import load_breast_cancer
+from sklearn.model_selection import train_test_split
+
+import ray
+from ray import train, tune
+from ray.train import ScalingConfig
+from ray.train.constants import TRAIN_DATASET_KEY
+from ray.train.xgboost import RayTrainReportCallback, XGBoostTrainer
+
+
+@pytest.fixture
+def ray_start_4_cpus():
+    address_info = ray.init(num_cpus=4)
+    yield address_info
+    # The code after the yield will run as teardown code.
+    ray.shutdown()
+
+
+@pytest.fixture
+def ray_start_8_cpus():
+    address_info = ray.init(num_cpus=8)
+    yield address_info
+    # The code after the yield will run as teardown code.
+    ray.shutdown()
+
+
+scale_config = ScalingConfig(num_workers=2)
+
+data_raw = load_breast_cancer()
+dataset_df = pd.DataFrame(data_raw["data"], columns=data_raw["feature_names"])
+dataset_df["target"] = data_raw["target"]
+train_df, test_df = train_test_split(dataset_df, test_size=0.3)
+
+params = {
+    "tree_method": "approx",
+    "objective": "binary:logistic",
+    "eval_metric": ["logloss", "error"],
+}
+
+
+def test_fit(ray_start_8_cpus):
+    train_dataset = ray.data.from_pandas(train_df)
+    valid_dataset = ray.data.from_pandas(test_df)
+    trainer = XGBoostTrainer(
+        scaling_config=scale_config,
+        label_column="target",
+        params=params,
+        datasets={TRAIN_DATASET_KEY: train_dataset, "valid": valid_dataset},
+    )
+    trainer.fit()
+
+
+class ScalingConfigAssertingXGBoostTrainer(XGBoostTrainer):
+    def training_loop(self) -> None:
+        pgf = train.get_context().get_trial_resources()
+        assert pgf.strategy == "SPREAD"
+        return super().training_loop()
+
+
+def test_fit_with_advanced_scaling_config(ray_start_8_cpus):
+    """Ensure that extra ScalingConfig arguments are respected."""
+    train_dataset = ray.data.from_pandas(train_df)
+    valid_dataset = ray.data.from_pandas(test_df)
+    trainer = ScalingConfigAssertingXGBoostTrainer(
+        scaling_config=ScalingConfig(
+            num_workers=2,
+            placement_strategy="SPREAD",
+        ),
+        label_column="target",
+        params=params,
+        datasets={TRAIN_DATASET_KEY: train_dataset, "valid": valid_dataset},
+    )
+    trainer.fit()
+
+
+def test_resume_from_checkpoint(ray_start_8_cpus, tmpdir):
+    train_dataset = ray.data.from_pandas(train_df)
+    valid_dataset = ray.data.from_pandas(test_df)
+    trainer = XGBoostTrainer(
+        scaling_config=scale_config,
+        label_column="target",
+        params=params,
+        num_boost_round=5,
+        datasets={TRAIN_DATASET_KEY: train_dataset, "valid": valid_dataset},
+    )
+    result = trainer.fit()
+    checkpoint = result.checkpoint
+    xgb_model = XGBoostTrainer.get_model(checkpoint)
+    assert xgb_model.num_boosted_rounds() == 5
+
+    trainer = XGBoostTrainer(
+        scaling_config=scale_config,
+        label_column="target",
+        params=params,
+        num_boost_round=10,
+        datasets={TRAIN_DATASET_KEY: train_dataset, "valid": valid_dataset},
+        resume_from_checkpoint=result.checkpoint,
+    )
+    result = trainer.fit()
+    model = XGBoostTrainer.get_model(result.checkpoint)
+    assert model.num_boosted_rounds() == 10
+
+
+@pytest.mark.parametrize(
+    "freq_end_expected",
+    [
+        # With num_boost_round=25 with 0 indexing, the checkpoints will be at:
+        (4, True, 7),  # 3, 7, 11, 15, 19, 23, 24 (end)
+        (4, False, 6),  # 3, 7, 11, 15, 19, 23
+        (5, True, 5),  # 4, 9, 14, 19, 24
+        (0, True, 1),  # 24 (end)
+        (0, False, 0),
+    ],
+)
+def test_checkpoint_freq(ray_start_8_cpus, freq_end_expected):
+    freq, end, expected = freq_end_expected
+
+    train_dataset = ray.data.from_pandas(train_df)
+    valid_dataset = ray.data.from_pandas(test_df)
+    trainer = XGBoostTrainer(
+        run_config=ray.train.RunConfig(
+            checkpoint_config=ray.train.CheckpointConfig(
+                checkpoint_frequency=freq, checkpoint_at_end=end
+            )
+        ),
+        scaling_config=scale_config,
+        label_column="target",
+        params=params,
+        num_boost_round=25,
+        datasets={TRAIN_DATASET_KEY: train_dataset, "valid": valid_dataset},
+    )
+    result = trainer.fit()
+
+    # Assert number of checkpoints
+    assert len(result.best_checkpoints) == expected, str(
+        [(metrics["training_iteration"], cp) for cp, metrics in result.best_checkpoints]
+    )
+
+    # Assert checkpoint numbers are increasing
+    cp_paths = [cp.path for cp, _ in result.best_checkpoints]
+    assert cp_paths == sorted(cp_paths), str(cp_paths)
+
+
+@pytest.mark.parametrize("rank", [None, 0, 1])
+def test_checkpoint_only_on_rank0(rank):
+    """Tests that the callback only reports checkpoints on rank 0,
+    or if the rank is not available (Tune usage)."""
+    callback = RayTrainReportCallback(frequency=2, checkpoint_at_end=True)
+
+    booster = mock.MagicMock()
+
+    with mock.patch("ray.train.get_context") as mock_get_context:
+        mock_context = mock.MagicMock()
+        mock_context.get_world_rank.return_value = rank
+        mock_get_context.return_value = mock_context
+
+        with callback._get_checkpoint(booster) as checkpoint:
+            if rank in (0, None):
+                assert checkpoint
+            else:
+                assert not checkpoint
+
+
+def test_tune(ray_start_8_cpus):
+    train_dataset = ray.data.from_pandas(train_df)
+    valid_dataset = ray.data.from_pandas(test_df)
+    trainer = XGBoostTrainer(
+        scaling_config=scale_config,
+        label_column="target",
+        params={**params, "max_depth": 1},
+        datasets={TRAIN_DATASET_KEY: train_dataset, "valid": valid_dataset},
+    )
+
+    tuner = tune.Tuner(
+        trainer,
+        param_space={"params": {"max_depth": tune.grid_search([2, 4])}},
+    )
+    results = tuner.fit()
+    assert sorted([r.config["params"]["max_depth"] for r in results]) == [2, 4]
+
+
+def test_validation(ray_start_4_cpus):
+    valid_dataset = ray.data.from_pandas(test_df)
+    with pytest.raises(ValueError, match=TRAIN_DATASET_KEY):
+        XGBoostTrainer(
+            scaling_config=ScalingConfig(num_workers=2),
+            label_column="target",
+            params=params,
+            datasets={"valid": valid_dataset},
+        )
+
+    with pytest.raises(ValueError, match="label_column"):
+        XGBoostTrainer(
+            scaling_config=ScalingConfig(num_workers=2),
+            datasets={"train": valid_dataset},
+        )
+
+
+def test_callback_get_model(tmp_path):
+    custom_filename = "custom.json"
+
+    bst = xgb.train(
+        params,
+        dtrain=xgb.DMatrix(train_df, label=train_df["target"]),
+        num_boost_round=1,
+    )
+    bst.save_model(tmp_path.joinpath(custom_filename).as_posix())
+    checkpoint = train.Checkpoint.from_directory(tmp_path.as_posix())
+
+    RayTrainReportCallback.get_model(checkpoint, filename=custom_filename)
+
+
+if __name__ == "__main__":
+    import sys
+
+    import pytest
+
+    sys.exit(pytest.main(["-v", "-x", __file__]))

python/ray/train/tests/util.py

@@ -0,0 +1,52 @@
+import contextlib
+import os
+import tempfile
+from typing import Any, Dict, Optional, Type
+
+import ray.cloudpickle as ray_pickle
+from ray.train import Checkpoint, SyncConfig
+from ray.train._internal.storage import StorageContext
+
+
+@contextlib.contextmanager
+def create_dict_checkpoint(
+    data: Dict[str, Any], checkpoint_cls: Type[Checkpoint] = None
+) -> Checkpoint:
+    with tempfile.TemporaryDirectory() as tmpdir:
+        with open(os.path.join(tmpdir, "data.pkl"), "wb") as f:
+            ray_pickle.dump(data, f)
+        checkpoint_cls = checkpoint_cls or Checkpoint
+        yield checkpoint_cls.from_directory(tmpdir)
+
+
+def load_dict_checkpoint(checkpoint: Checkpoint) -> Dict[str, Any]:
+    with checkpoint.as_directory() as checkpoint_dir:
+        with open(os.path.join(checkpoint_dir, "data.pkl"), "rb") as f:
+            return ray_pickle.load(f)
+
+
+def mock_storage_context(
+    exp_name: str = "exp_name",
+    storage_path: Optional[str] = None,
+    storage_context_cls: Type = StorageContext,
+    sync_config: Optional[SyncConfig] = None,
+) -> StorageContext:
+    storage_path = storage_path or tempfile.mkdtemp()
+    exp_name = exp_name
+    trial_name = "trial_name"
+
+    storage = storage_context_cls(
+        storage_path=storage_path,
+        experiment_dir_name=exp_name,
+        trial_dir_name=trial_name,
+        sync_config=sync_config,
+    )
+    # Patch the default /tmp/ray/session_* so we don't require ray
+    # to be initialized in unit tests.
+    session_path = tempfile.mkdtemp()
+    storage._get_session_path = lambda: session_path
+
+    os.makedirs(storage.trial_fs_path, exist_ok=True)
+    os.makedirs(storage.trial_driver_staging_path, exist_ok=True)
+
+    return storage