From da42bf29d07141c61948bd78cad936548fc47911 Mon Sep 17 00:00:00 2001 From: Richard Liaw Date: Wed, 2 Dec 2020 12:04:54 -0800 Subject: [PATCH] [tune] horovod release test (#12495) --- python/ray/tune/integration/horovod.py | 54 ++++++- python/ray/tune/tests/test_horovod.py | 9 +- python/ray/tune/utils/mock.py | 26 ++++ release/horovod_tests/cluster.yaml | 69 +++++++++ release/horovod_tests/run.sh | 58 +++++++ .../horovod_tests/workloads/horovod_test.py | 142 ++++++++++++++++++ .../workloads/pytorch_pbt_failure.py | 27 +--- 7 files changed, 357 insertions(+), 28 deletions(-) create mode 100644 release/horovod_tests/cluster.yaml create mode 100755 release/horovod_tests/run.sh create mode 100644 release/horovod_tests/workloads/horovod_test.py diff --git a/python/ray/tune/integration/horovod.py b/python/ray/tune/integration/horovod.py index b89412e0a..8d85141f8 100644 --- a/python/ray/tune/integration/horovod.py +++ b/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()) diff --git a/python/ray/tune/tests/test_horovod.py b/python/ray/tune/tests/test_horovod.py index fc71fc399..e4f4bcd52 100644 --- a/python/ray/tune/tests/test_horovod.py +++ b/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]) diff --git a/python/ray/tune/utils/mock.py b/python/ray/tune/utils/mock.py index cdc614b93..337482fe9 100644 --- a/python/ray/tune/utils/mock.py +++ b/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) diff --git a/release/horovod_tests/cluster.yaml b/release/horovod_tests/cluster.yaml new file mode 100644 index 000000000..880ebdba2 --- /dev/null +++ b/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 diff --git a/release/horovod_tests/run.sh b/release/horovod_tests/run.sh new file mode 100755 index 000000000..9ca69a44d --- /dev/null +++ b/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" diff --git a/release/horovod_tests/workloads/horovod_test.py b/release/horovod_tests/workloads/horovod_test.py new file mode 100644 index 000000000..dda5ec1a2 --- /dev/null +++ b/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) diff --git a/release/long_running_distributed_tests/workloads/pytorch_pbt_failure.py b/release/long_running_distributed_tests/workloads/pytorch_pbt_failure.py index 696855e28..0fa94cb44 100644 --- a/release/long_running_distributed_tests/workloads/pytorch_pbt_failure.py +++ b/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"))