chore: import upstream snapshot with attribution

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__]))