[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)