[tune] horovod release test (#12495)

python/ray/tune/integration/horovod.py

@@ -1,6 +1,10 @@
+from typing import Callable, Dict, Type
+
+from contextlib import contextmanager
 import os
 import logging
-from typing import Callable, Dict, Type
+import shutil
+import tempfile
 
 from filelock import FileLock
 
@@ -30,6 +34,45 @@ def logger_creator(log_config: Dict, logdir: str) -> NoopLogger:
     return NoopLogger(log_config, worker_dir)
 
 
+@contextmanager
+def distributed_checkpoint_dir(step: int, disable: bool = False):
+    """ContextManager for creating a distributed checkpoint.
+
+    Only checkpoints a file on the "main" training actor, avoiding
+    redundant work.
+
+    Args:
+        step (int): Used to label the checkpoint
+        disable (bool): Disable for prototyping.
+
+    Yields:
+        str: A path to a directory. This path will be used
+        again when invoking the training_function.
+
+    Example:
+
+    .. code-block:: python
+
+        def train_func(config, checkpoint_dir):
+            if checkpoint_dir:
+                path = os.path.join(checkpoint_dir, "checkpoint")
+                model_state_dict = torch.load(path)
+
+            if epoch % 3 == 0:
+                with distributed_checkpoint_dir(step=epoch) as checkpoint_dir:
+                    path = os.path.join(checkpoint_dir, "checkpoint")
+                    torch.save(model.state_dict(), path)
+    """
+
+    if int(get_rank()) == 0 and not disable:
+        with tune.checkpoint_dir(step=step) as checkpoint_dir:
+            yield checkpoint_dir
+    else:
+        path = tempfile.mkdtemp()
+        yield path
+        shutil.rmtree(path)
+
+
 class _HorovodTrainable(tune.Trainable):
     """Abstract Trainable class for Horovod."""
     # Callable function for training.
@@ -103,7 +146,8 @@ class _HorovodTrainable(tune.Trainable):
     def load_checkpoint(self, checkpoint_dir: str):
         checkpoint_obj = TrainableUtil.checkpoint_to_object(checkpoint_dir)
         x_id = ray.put(checkpoint_obj)
-        return self.executor.execute(lambda w: w.restore_from_object(x_id))
+        return self.executor.execute(
+            lambda w: w.restore_from_object(ray.get(x_id)))
 
     def stop(self):
         self.executor.execute(lambda w: w.stop())
@@ -227,4 +271,10 @@ def _train_simple(config: Dict):
     for i in range(config.get("epochs", 2)):
         import time
         time.sleep(1)
+        if config.get("enable_checkpoint", True):
+            with distributed_checkpoint_dir(step=i) as checkpoint_dir:
+                path = os.path.join(checkpoint_dir, "checkpoint")
+                import pickle
+                with open(path, "wb") as f:
+                    pickle.dump("hi", f)
         tune.report(test=1, rank=hvd.rank())

python/ray/tune/tests/test_horovod.py

@@ -73,11 +73,16 @@ def test_set_global(ray_start_2_cpus):
     assert result["rank"] == 0
 
 
-def test_simple_tune(ray_start_4_cpus):
+@pytest.mark.parametrize("enabled_checkpoint", [True, False])
+def test_simple_tune(ray_start_4_cpus, enabled_checkpoint):
     trainable_cls = DistributedTrainableCreator(_train_simple, num_slots=2)
     analysis = tune.run(
-        trainable_cls, num_samples=2, stop={"training_iteration": 2})
+        trainable_cls,
+        config={"enable_checkpoint": enabled_checkpoint},
+        num_samples=2,
+        stop={"training_iteration": 2})
     assert analysis.trials[0].last_result["training_iteration"] == 2
+    assert analysis.trials[0].has_checkpoint() == enabled_checkpoint
 
 
 @pytest.mark.parametrize("use_gpu", [True, False])

python/ray/tune/utils/mock.py

@@ -1,6 +1,7 @@
 import os
 import numpy as np
 import json
+import random
 
 import ray.utils
 
@@ -8,6 +9,7 @@ from ray.rllib.agents.mock import _MockTrainer
 from ray.tune import DurableTrainable, Trainable
 from ray.tune.sync_client import get_sync_client
 from ray.tune.syncer import NodeSyncer
+from ray.tune.callback import Callback
 
 MOCK_REMOTE_DIR = os.path.join(ray.utils.get_user_temp_dir(),
                                "mock-tune-remote") + os.sep
@@ -88,3 +90,27 @@ class MyTrainableClass(Trainable):
     def load_checkpoint(self, checkpoint_path):
         with open(checkpoint_path) as f:
             self.timestep = json.loads(f.read())["timestep"]
+
+
+class FailureInjectorCallback(Callback):
+    """Adds random failure injection to the TrialExecutor."""
+
+    def __init__(self,
+                 config_path="/home/ubuntu/ray_bootstrap_config.yaml",
+                 probability=0.1,
+                 disable=False):
+        self.probability = probability
+        self.config_path = config_path
+        self.disable = disable
+
+    def on_step_begin(self, **info):
+        from ray.autoscaler._private.commands import kill_node
+        # With 10% probability inject failure to a worker.
+        if random.random() < self.probability and not self.disable:
+            # With 10% probability fully terminate the node.
+            should_terminate = random.random() < self.probability
+            kill_node(
+                self.config_path,
+                yes=True,
+                hard=should_terminate,
+                override_cluster_name=None)

release/horovod_tests/cluster.yaml

@@ -0,0 +1,69 @@
+# This file is generated by `ray project create`.
+
+# A unique identifier for the head node and workers of this cluster.
+cluster_name: horovod-release-tests
+
+# The minimum number of workers nodes to launch in addition to the head
+# node. This number should be >= 0.
+min_workers: 3
+# The maximum number of workers nodes to launch in addition to the head
+# node. This takes precedence over min_workers. min_workers defaults to 0.
+max_workers: 3
+
+target_utilization_fraction: 0.8
+
+# If a node is idle for this many minutes, it will be removed.
+idle_timeout_minutes: 5
+
+# Cloud-provider specific configuration.
+provider:
+    type: aws
+    region: us-west-2
+    availability_zone: us-west-2a
+    cache_stopped_nodes: False
+
+# How Ray will authenticate with newly launched nodes.
+auth:
+    ssh_user: ubuntu
+
+# By default Ray creates a new private keypair, but you can also use your own.
+# If you do so, make sure to also set "KeyName" in the head and worker node
+# configurations below.
+#    ssh_private_key: /path/to/your/key.pem
+
+# Provider-specific config for the head node, e.g. instance type. By default
+# Ray will auto-configure unspecified fields such as SubnetId and KeyName.
+# For more documentation on available fields, see:
+# http://boto3.readthedocs.io/en/latest/reference/services/ec2.html#EC2.ServiceResource.create_instances
+head_node:
+    InstanceType: g3.8xlarge
+    ImageId: latest_dlami
+    BlockDeviceMappings:
+        - DeviceName: /dev/sda1
+          Ebs:
+              VolumeSize: 150
+
+worker_nodes:
+  InstanceType: g3.8xlarge
+  ImageId: latest_dlami
+  BlockDeviceMappings:
+        - DeviceName: /dev/sda1
+          Ebs:
+              VolumeSize: 150
+  InstanceMarketOptions:
+    MarketType: spot
+
+
+file_mounts: {}
+
+# Command to start ray on the head node. You don't need to change this.
+head_start_ray_commands:
+    - ray stop
+    - export RAY_BACKEND_LOG_LEVEL=debug
+    - ray start --head --port=6379 --object-manager-port=8076 --autoscaling-config=~/ray_bootstrap_config.yaml
+
+# Command to start ray on worker nodes. You don't need to change this.
+worker_start_ray_commands:
+    - ray stop
+    - export RAY_BACKEND_LOG_LEVEL=debug
+    - ray start --address=$RAY_HEAD_IP:6379 --object-manager-port=8076

release/horovod_tests/run.sh

@@ -0,0 +1,58 @@
+#!/usr/bin/env bash
+
+ray_version=""
+commit=""
+ray_branch=""
+workload=""
+
+usage() {
+    echo "Run a horovod test."
+}
+
+for i in "$@"
+do
+echo "$i"
+case "$i" in
+    --ray-version=*)
+    ray_version="${i#*=}"
+
+    ;;
+    --commit=*)
+    commit="${i#*=}"
+    ;;
+    --ray-branch=*)
+    ray_branch="${i#*=}"
+    ;;
+    --workload=*)
+    workload="${i#*=}"
+    ;;
+    --help)
+    usage
+    exit
+    ;;
+    *)
+    echo "unknown arg, $i"
+    exit 1
+    ;;
+esac
+done
+
+echo "version: $ray_version"
+echo "commit: $commit"
+echo "branch: $ray_branch"
+echo "workload: $workload"
+
+wheel="https://s3-us-west-2.amazonaws.com/ray-wheels/$ray_branch/$commit/ray-$ray_version-cp37-cp37m-manylinux2014_x86_64.whl"
+
+conda uninstall -y terminado || true
+pip install -U pip
+pip install terminado
+pip install -U "$wheel"
+pip install "ray[rllib]"
+pip install -q -U ipdb torch torchvision
+HOROVOD_WITH_GLOO=1 HOROVOD_WITHOUT_MPI=1 HOROVOD_WITHOUT_MXNET=1 HOROVOD_WITH_PYTORCH=1  pip install -U git+https://github.com/horovod/horovod.git
+
+echo set-window-option -g mouse on > ~/.tmux.conf
+echo 'termcapinfo xterm* ti@:te@' > ~/.screenrc
+
+python "workloads/$workload.py"

release/horovod_tests/workloads/horovod_test.py

@@ -0,0 +1,142 @@
+import argparse
+import torch
+import torch.nn as nn
+import os
+import numpy as np
+import torchvision
+from torch.utils.data import DataLoader
+
+import torchvision.transforms as transforms
+
+import ray
+from ray import tune
+from ray.tune.schedulers import create_scheduler
+from ray.tune.integration.horovod import (DistributedTrainableCreator,
+                                          distributed_checkpoint_dir)
+from ray.util.sgd.torch.resnet import ResNet18
+
+parser = argparse.ArgumentParser()
+parser.add_argument("--gpu", action="store_true")
+parser.add_argument(
+    "--smoke-test", action="store_true", help=("Finish quickly for testing."))
+args = parser.parse_args()
+
+CIFAR10_STATS = {
+    "mean": (0.4914, 0.4822, 0.4465),
+    "std": (0.2023, 0.1994, 0.2010),
+}
+
+
+def train(config, checkpoint_dir=None):
+    import horovod.torch as hvd
+    hvd.init()
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    net = ResNet18(None).to(device)
+    optimizer = torch.optim.SGD(
+        net.parameters(),
+        lr=config["lr"],
+    )
+    epoch = 0
+
+    if checkpoint_dir:
+        with open(os.path.join(checkpoint_dir, "checkpoint")) as f:
+            model_state, optimizer_state, epoch = torch.load(f)
+
+        net.load_state_dict(model_state)
+        optimizer.load_state_dict(optimizer_state)
+
+    criterion = nn.CrossEntropyLoss()
+    optimizer = hvd.DistributedOptimizer(optimizer)
+    np.random.seed(1 + hvd.rank())
+    torch.manual_seed(1234)
+    # To ensure consistent initialization across slots,
+    hvd.broadcast_parameters(net.state_dict(), root_rank=0)
+    hvd.broadcast_optimizer_state(optimizer, root_rank=0)
+
+    trainset = ray.get(config["data"])
+    trainloader = DataLoader(
+        trainset,
+        batch_size=int(config["batch_size"]),
+        shuffle=True,
+        num_workers=4)
+
+    for epoch in range(epoch, 40):  # loop over the dataset multiple times
+        running_loss = 0.0
+        epoch_steps = 0
+        for i, data in enumerate(trainloader):
+            # get the inputs; data is a list of [inputs, labels]
+            inputs, labels = data
+            inputs, labels = inputs.to(device), labels.to(device)
+
+            # zero the parameter gradients
+            optimizer.zero_grad()
+
+            # forward + backward + optimize
+            outputs = net(inputs)
+            loss = criterion(outputs, labels)
+            loss.backward()
+            optimizer.step()
+
+            # print statistics
+            running_loss += loss.item()
+            epoch_steps += 1
+            tune.report(loss=running_loss / epoch_steps)
+            if i % 2000 == 1999:  # print every 2000 mini-batches
+                print("[%d, %5d] loss: %.3f" % (epoch + 1, i + 1,
+                                                running_loss / epoch_steps))
+
+        with distributed_checkpoint_dir(step=epoch) as checkpoint_dir:
+            print("this checkpoint dir: ", checkpoint_dir)
+            path = os.path.join(checkpoint_dir, "checkpoint")
+            torch.save((net.state_dict(), optimizer.state_dict(), epoch), path)
+
+
+if __name__ == "__main__":
+    if args.smoke_test:
+        ray.init()
+    else:
+        ray.init(address="auto")  # assumes ray is started with ray up
+
+    horovod_trainable = DistributedTrainableCreator(
+        train,
+        use_gpu=True,
+        num_hosts=1 if args.smoke_test else 2,
+        num_slots=2 if args.smoke_test else 2,
+        replicate_pem=False)
+
+    transform_train = transforms.Compose([
+        transforms.RandomCrop(32, padding=4),
+        transforms.RandomHorizontalFlip(),
+        transforms.ToTensor(),
+        transforms.Normalize(CIFAR10_STATS["mean"], CIFAR10_STATS["std"]),
+    ])  # meanstd transformation
+
+    dataset = torchvision.datasets.CIFAR10(
+        root="/tmp/data_cifar",
+        train=True,
+        download=True,
+        transform=transform_train)
+
+    # ensure that checkpointing works.
+    pbt = create_scheduler(
+        "pbt",
+        perturbation_interval=2,
+        hyperparam_mutations={
+            "lr": tune.uniform(0.001, 0.1),
+        })
+
+    analysis = tune.run(
+        horovod_trainable,
+        metric="loss",
+        mode="min",
+        keep_checkpoints_num=1,
+        scheduler=pbt,
+        config={
+            "lr": tune.grid_search([0.1 * i for i in range(1, 10)]),
+            "batch_size": 64,
+            "data": ray.put(dataset)
+        },
+        num_samples=1,
+        fail_fast=True)
+    # callbacks=[FailureInjectorCallback()])
+    print("Best hyperparameters found were: ", analysis.best_config)

release/long_running_distributed_tests/workloads/pytorch_pbt_failure.py

@@ -1,7 +1,6 @@
 import argparse
 import numpy as np
 import os
-import random
 import torch
 import torch.nn as nn
 from torch.utils.data import DataLoader, Subset
@@ -10,11 +9,10 @@ import torchvision.transforms as transforms
 
 import ray
 from ray import tune
-from ray.autoscaler._private.commands import kill_node
 from ray.tune import CLIReporter
-from ray.tune.ray_trial_executor import RayTrialExecutor
 from ray.tune.schedulers import PopulationBasedTraining
 from ray.tune.utils.util import merge_dicts
+from ray.tune.utils.mock import FailureInjectorCallback
 from ray.util.sgd.torch import TorchTrainer
 from ray.util.sgd.torch.resnet import ResNet18
 from ray.util.sgd.utils import BATCH_SIZE
@@ -28,23 +26,6 @@ parser.add_argument(
 args = parser.parse_args()
 
 
-class FailureInjectorExecutor(RayTrialExecutor):
-    """Adds random failure injection to the TrialExecutor."""
-
-    def on_step_begin(self, trial_runner):
-        """Before step(), update available resources and inject failure."""
-        self._update_avail_resources()
-        # With 10% probability inject failure to a worker.
-        if random.random() < 0.1 and not args.smoke_test:
-            # With 10% probability fully terminate the node.
-            should_terminate = random.random() < 0.1
-            kill_node(
-                "/home/ubuntu/ray_bootstrap_config.yaml",
-                yes=True,
-                hard=should_terminate,
-                override_cluster_name=None)
-
-
 def initialization_hook():
     # Need this for avoiding a connection restart issue on AWS.
     os.environ["NCCL_SOCKET_IFNAME"] = "^docker0,lo"
@@ -95,9 +76,6 @@ def optimizer_creator(model, config):
 
 ray.init(address="auto" if not args.smoke_test else None, _log_to_driver=True)
 num_training_workers = 1 if args.smoke_test else 3
-
-executor = FailureInjectorExecutor(queue_trials=True)
-
 TorchTrainable = TorchTrainer.as_trainable(
     model_creator=ResNet18,
     data_creator=cifar_creator,
@@ -149,7 +127,8 @@ analysis = tune.run(
     checkpoint_freq=2,  # used for fault tolerance
     progress_reporter=reporter,
     scheduler=pbt_scheduler,
-    trial_executor=executor,
+    callbacks=[FailureInjectorCallback()],
+    queue_trials=True,
     stop={"training_iteration": 1} if args.smoke_test else None)
 
 print(analysis.get_best_config(metric="val_loss", mode="min"))