[tune] distributed torch wrapper (#9550)

* changes

* add-working

* checkpoint

* ccleanu

* fix

* ok

* formatting

* ok

* tests

* some-good-stuff

* fix-torch

* ddp-torch

* torch-test

* sessions

* add-small-test

* fix

* remove

* gpu-working

* update-tests

* ok

* try-test

* formgat

* ok

* ok

python/ray/tune/BUILD

@@ -287,6 +287,15 @@ py_test(
     args = ["--smoke-test"]
 )
 
+py_test(
+    name = "ddp_mnist_torch",
+    size = "small",
+    srcs = ["examples/ddp_mnist_torch.py"],
+    deps = [":tune_lib"],
+    tags = ["exclusive", "example", "pytorch"],
+    args = ["--num-workers=2"]
+)
+
 py_test(
     name = "dragonfly_example",
     size = "medium",

python/ray/tune/examples/ddp_mnist_torch.py

@@ -0,0 +1,73 @@
+# Original Code here:
+# https://github.com/pytorch/examples/blob/master/mnist/main.py
+import argparse
+import logging
+import torch
+import torch.optim as optim
+from torch.nn.parallel import DistributedDataParallel
+
+import ray
+from ray import tune
+from ray.tune.examples.mnist_pytorch import (train, test, get_data_loaders,
+                                             ConvNet)
+from ray.util.sgd.torch.func_trainable import (DistributedTrainableCreator,
+                                               distributed_checkpoint)
+
+logger = logging.getLogger(__name__)
+
+
+def train_mnist(config, checkpoint=False):
+    use_cuda = torch.cuda.is_available()
+    device = torch.device("cuda" if use_cuda else "cpu")
+    train_loader, test_loader = get_data_loaders()
+    model = ConvNet().to(device)
+    optimizer = optim.SGD(model.parameters(), lr=0.1)
+
+    if checkpoint:
+        with open(checkpoint) as f:
+            model_state, optimizer_state = torch.load(f)
+
+        model.load_state_dict(model_state)
+        optimizer.load_state_dict(optimizer_state)
+
+    model = DistributedDataParallel(model)
+
+    for epoch in range(40):
+        train(model, optimizer, train_loader, device)
+        acc = test(model, test_loader, device)
+
+        if epoch % 3 == 0:
+            with distributed_checkpoint(label=epoch) as path:
+                torch.save((model.state_dict(), optimizer.state_dict()), path)
+        tune.report(mean_accuracy=acc)
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--num-workers",
+        "-n",
+        type=int,
+        default=2,
+        help="Sets number of workers for training.")
+    parser.add_argument(
+        "--use-gpu",
+        action="store_true",
+        default=False,
+        help="enables CUDA training")
+    parser.add_argument(
+        "--cluster",
+        action="store_true",
+        default=False,
+        help="enables multi-node tuning")
+
+    args = parser.parse_args()
+
+    if args.cluster:
+        options = dict(address="auto")
+    else:
+        options = dict(num_cpus=2)
+    ray.init(**options)
+    trainable_cls = DistributedTrainableCreator(
+        train_mnist, num_workers=args.num_workers, use_gpu=args.use_gpu)
+    tune.run(trainable_cls, num_samples=4, stop={"training_iteration": 10})

python/ray/tune/function_runner.py

@@ -1,6 +1,5 @@
 import logging
 import os
-import io
 import time
 import inspect
 import shutil
@@ -87,6 +86,7 @@ class StatusReporter:
     def make_checkpoint_dir(self, step=None):
         checkpoint_dir = TrainableUtil.make_checkpoint_dir(
             self.logdir, index=step)
+        logger.debug("Making checkpoint dir at %s", checkpoint_dir)
         return checkpoint_dir
 
     def save_checkpoint(self, checkpoint):
@@ -279,6 +279,9 @@ class FunctionRunner(Trainable):
             result[SHOULD_CHECKPOINT] = True
         return result
 
+    def execute(self, fn):
+        return fn(self)
+
     def create_default_checkpoint_dir(self):
         self.default_checkpoint_dir = TrainableUtil.make_checkpoint_dir(
             self.logdir, index="default")
@@ -306,12 +309,8 @@ class FunctionRunner(Trainable):
 
     def save_to_object(self):
         checkpoint_path = self.save()
-        data_dict = TrainableUtil.pickle_checkpoint(checkpoint_path)
-        out = io.BytesIO()
-        if len(data_dict) > 10e6:  # getting pretty large
-            logger.info("Checkpoint size is {} bytes".format(len(data_dict)))
-        out.write(data_dict)
-        return out.getvalue()
+        obj = TrainableUtil.checkpoint_to_object(checkpoint_path)
+        return obj
 
     def load_checkpoint(self, checkpoint):
         # This should be removed once Trainables are refactored.

python/ray/tune/ray_trial_executor.py

@@ -670,9 +670,9 @@ class RayTrialExecutor(TrialExecutor):
                 # This provides FT backwards compatibility in the
                 # case where a DurableTrainable is not provided.
                 logger.debug("Trial %s: Reading checkpoint into memory", trial)
-                data_dict = TrainableUtil.pickle_checkpoint(value)
+                obj = TrainableUtil.checkpoint_to_object(value)
                 with self._change_working_directory(trial):
-                    remote = trial.runner.restore_from_object.remote(data_dict)
+                    remote = trial.runner.restore_from_object.remote(obj)
             else:
                 raise AbortTrialExecution(
                     "Pass in `sync_on_checkpoint=True` for driver-based trial"

python/ray/tune/session.py

@@ -1,3 +1,4 @@
+import os
 import logging
 
 logger = logging.getLogger(__name__)
@@ -7,8 +8,8 @@ _session = None
 
 def get_session():
     global _session
-    if _session is None:
-        raise ValueError(
+    if not _session:
+        logger.warning(
             "Session not detected. You should not be calling this function "
             "outside `tune.run` or while using the class API. ")
     return _session
@@ -67,7 +68,8 @@ def report(**kwargs):
             metrics can be used for early stopping or optimization.
     """
     _session = get_session()
-    return _session(**kwargs)
+    if _session:
+        return _session(**kwargs)
 
 
 def make_checkpoint_dir(step=None):
@@ -106,7 +108,10 @@ def make_checkpoint_dir(step=None):
 
     """
     _session = get_session()
-    return _session.make_checkpoint_dir(step=step)
+    if _session:
+        return _session.make_checkpoint_dir(step=step)
+    else:
+        return os.path.abspath("./")
 
 
 def save_checkpoint(checkpoint):
@@ -149,7 +154,8 @@ def save_checkpoint(checkpoint):
     .. versionadded:: 0.8.6
     """
     _session = get_session()
-    return _session.save_checkpoint(checkpoint)
+    if _session:
+        return _session.save_checkpoint(checkpoint)
 
 
 def get_trial_dir():
@@ -158,7 +164,8 @@ def get_trial_dir():
     For function API use only.
     """
     _session = get_session()
-    return _session.logdir
+    if _session:
+        return _session.logdir
 
 
 def get_trial_name():
@@ -167,7 +174,8 @@ def get_trial_name():
     For function API use only.
     """
     _session = get_session()
-    return _session.trial_name
+    if _session:
+        return _session.trial_name
 
 
 def get_trial_id():
@@ -176,7 +184,8 @@ def get_trial_id():
     For function API use only.
     """
     _session = get_session()
-    return _session.trial_id
+    if _session:
+        return _session.trial_id
 
 
 __all__ = ["report", "get_trial_dir", "get_trial_name", "get_trial_id"]

python/ray/tune/trainable.py

@@ -77,6 +77,15 @@ class TrainableUtil:
         })
         return data_dict
 
+    @staticmethod
+    def checkpoint_to_object(checkpoint_path):
+        data_dict = TrainableUtil.pickle_checkpoint(checkpoint_path)
+        out = io.BytesIO()
+        if len(data_dict) > 10e6:  # getting pretty large
+            logger.info("Checkpoint size is {} bytes".format(len(data_dict)))
+        out.write(data_dict)
+        return out.getvalue()
+
     @staticmethod
     def find_checkpoint_dir(checkpoint_path):
         """Returns the directory containing the checkpoint path.
@@ -424,14 +433,10 @@ class Trainable:
         """
         tmpdir = tempfile.mkdtemp("save_to_object", dir=self.logdir)
         checkpoint_path = self.save(tmpdir)
-        # Save all files in subtree.
-        data_dict = TrainableUtil.pickle_checkpoint(checkpoint_path)
-        out = io.BytesIO()
-        if len(data_dict) > 10e6:  # getting pretty large
-            logger.info("Checkpoint size is {} bytes".format(len(data_dict)))
-        out.write(data_dict)
+        # Save all files in subtree and delete the tmpdir.
+        obj = TrainableUtil.checkpoint_to_object(checkpoint_path)
         shutil.rmtree(tmpdir)
-        return out.getvalue()
+        return obj
 
     def restore(self, checkpoint_path):
         """Restores training state from a given model checkpoint.