Add ray.util package and move libraries from experimental (#7100)

python/ray/util/__init__.py

@@ -0,0 +1,10 @@
+from .named_actors import get_actor, register_actor
+from .actor_pool import ActorPool
+from . import iter
+
+__all__ = [
+    "ActorPool",
+    "iter",
+    "get_actor",
+    "register_actor",
+]

python/ray/util/actor_pool.py

@@ -0,0 +1,212 @@
+import ray
+
+
+class ActorPool:
+    """Utility class to operate on a fixed pool of actors.
+
+    Arguments:
+        actors (list): List of Ray actor handles to use in this pool.
+
+    Examples:
+        >>> a1, a2 = Actor.remote(), Actor.remote()
+        >>> pool = ActorPool([a1, a2])
+        >>> print(pool.map(lambda a, v: a.double.remote(v), [1, 2, 3, 4]))
+        [2, 4, 6, 8]
+    """
+
+    def __init__(self, actors):
+        # actors to be used
+        self._idle_actors = list(actors)
+
+        # get actor from future
+        self._future_to_actor = {}
+
+        # get future from index
+        self._index_to_future = {}
+
+        # next task to do
+        self._next_task_index = 0
+
+        # next task to return
+        self._next_return_index = 0
+
+        # next work depending when actors free
+        self._pending_submits = []
+
+    def map(self, fn, values):
+        """Apply the given function in parallel over the actors and values.
+
+        This returns an ordered iterator that will return results of the map
+        as they finish. Note that you must iterate over the iterator to force
+        the computation to finish.
+
+        Arguments:
+            fn (func): Function that takes (actor, value) as argument and
+                returns an ObjectID computing the result over the value. The
+                actor will be considered busy until the ObjectID completes.
+            values (list): List of values that fn(actor, value) should be
+                applied to.
+
+        Returns:
+            Iterator over results from applying fn to the actors and values.
+
+        Examples:
+            >>> pool = ActorPool(...)
+            >>> print(pool.map(lambda a, v: a.double.remote(v), [1, 2, 3, 4]))
+            [2, 4, 6, 8]
+        """
+        for v in values:
+            self.submit(fn, v)
+        while self.has_next():
+            yield self.get_next()
+
+    def map_unordered(self, fn, values):
+        """Similar to map(), but returning an unordered iterator.
+
+        This returns an unordered iterator that will return results of the map
+        as they finish. This can be more efficient that map() if some results
+        take longer to compute than others.
+
+        Arguments:
+            fn (func): Function that takes (actor, value) as argument and
+                returns an ObjectID computing the result over the value. The
+                actor will be considered busy until the ObjectID completes.
+            values (list): List of values that fn(actor, value) should be
+                applied to.
+
+        Returns:
+            Iterator over results from applying fn to the actors and values.
+
+        Examples:
+            >>> pool = ActorPool(...)
+            >>> print(pool.map(lambda a, v: a.double.remote(v), [1, 2, 3, 4]))
+            [6, 2, 4, 8]
+        """
+        for v in values:
+            self.submit(fn, v)
+        while self.has_next():
+            yield self.get_next_unordered()
+
+    def submit(self, fn, value):
+        """Schedule a single task to run in the pool.
+
+        This has the same argument semantics as map(), but takes on a single
+        value instead of a list of values. The result can be retrieved using
+        get_next() / get_next_unordered().
+
+        Arguments:
+            fn (func): Function that takes (actor, value) as argument and
+                returns an ObjectID computing the result over the value. The
+                actor will be considered busy until the ObjectID completes.
+            value (object): Value to compute a result for.
+
+        Examples:
+            >>> pool = ActorPool(...)
+            >>> pool.submit(lambda a, v: a.double.remote(v), 1)
+            >>> pool.submit(lambda a, v: a.double.remote(v), 2)
+            >>> print(pool.get_next(), pool.get_next())
+            2, 4
+        """
+        if self._idle_actors:
+            actor = self._idle_actors.pop()
+            future = fn(actor, value)
+            self._future_to_actor[future] = (self._next_task_index, actor)
+            self._index_to_future[self._next_task_index] = future
+            self._next_task_index += 1
+        else:
+            self._pending_submits.append((fn, value))
+
+    def has_next(self):
+        """Returns whether there are any pending results to return.
+
+        Returns:
+            True if there are any pending results not yet returned.
+
+        Examples:
+            >>> pool = ActorPool(...)
+            >>> pool.submit(lambda a, v: a.double.remote(v), 1)
+            >>> print(pool.has_next())
+            True
+            >>> print(pool.get_next())
+            2
+            >>> print(pool.has_next())
+            False
+        """
+        return bool(self._future_to_actor)
+
+    def get_next(self, timeout=None):
+        """Returns the next pending result in order.
+
+        This returns the next result produced by submit(), blocking for up to
+        the specified timeout until it is available.
+
+        Returns:
+            The next result.
+
+        Raises:
+            TimeoutError if the timeout is reached.
+
+        Examples:
+            >>> pool = ActorPool(...)
+            >>> pool.submit(lambda a, v: a.double.remote(v), 1)
+            >>> print(pool.get_next())
+            2
+        """
+        if not self.has_next():
+            raise StopIteration("No more results to get")
+        if self._next_return_index >= self._next_task_index:
+            raise ValueError("It is not allowed to call get_next() after "
+                             "get_next_unordered().")
+        future = self._index_to_future[self._next_return_index]
+        if timeout is not None:
+            res, _ = ray.wait([future], timeout=timeout)
+            if not res:
+                raise TimeoutError("Timed out waiting for result")
+        del self._index_to_future[self._next_return_index]
+        self._next_return_index += 1
+        i, a = self._future_to_actor.pop(future)
+        self._return_actor(a)
+        return ray.get(future)
+
+    def get_next_unordered(self, timeout=None):
+        """Returns any of the next pending results.
+
+        This returns some result produced by submit(), blocking for up to
+        the specified timeout until it is available. Unlike get_next(), the
+        results are not always returned in same order as submitted, which can
+        improve performance.
+
+        Returns:
+            The next result.
+
+        Raises:
+            TimeoutError if the timeout is reached.
+
+        Examples:
+            >>> pool = ActorPool(...)
+            >>> pool.submit(lambda a, v: a.double.remote(v), 1)
+            >>> pool.submit(lambda a, v: a.double.remote(v), 2)
+            >>> print(pool.get_next_unordered())
+            4
+            >>> print(pool.get_next_unordered())
+            2
+        """
+        if not self.has_next():
+            raise StopIteration("No more results to get")
+        # TODO(ekl) bulk wait for performance
+        res, _ = ray.wait(
+            list(self._future_to_actor), num_returns=1, timeout=timeout)
+        if res:
+            [future] = res
+        else:
+            raise TimeoutError("Timed out waiting for result")
+        i, a = self._future_to_actor.pop(future)
+        self._return_actor(a)
+        del self._index_to_future[i]
+        self._next_return_index = max(self._next_return_index, i + 1)
+        return ray.get(future)
+
+    def _return_actor(self, actor):
+        self._idle_actors.append(actor)
+        if self._pending_submits:
+            self.submit(*self._pending_submits.pop(0))

python/ray/util/iter.py

@@ -0,0 +1,749 @@
+from typing import TypeVar, Generic, Iterable, List, Callable, Any
+import random
+
+import ray
+
+# The type of an iterator element.
+T = TypeVar("T")
+U = TypeVar("U")
+
+
+def from_items(items: List[T], num_shards: int = 2,
+               repeat: bool = False) -> "ParallelIterator[T]":
+    """Create a parallel iterator from an existing set of objects.
+
+    The objects will be divided round-robin among the number of shards.
+
+    Args:
+        items (list): The list of items to iterate over.
+        num_shards (int): The number of worker actors to create.
+        repeat (bool): Whether to cycle over the items forever.
+    """
+    shards = [[] for _ in range(num_shards)]
+    for i, item in enumerate(items):
+        shards[i % num_shards].append(item)
+    name = "from_items[{}, {}, shards={}{}]".format(
+        items and type(items[0]).__name__ or "None", len(items), num_shards,
+        ", repeat=True" if repeat else "")
+    return from_iterators(shards, repeat=repeat, name=name)
+
+
+def from_range(n: int, num_shards: int = 2,
+               repeat: bool = False) -> "ParallelIterator[int]":
+    """Create a parallel iterator over the range 0..n.
+
+    The range will be partitioned sequentially among the number of shards.
+
+    Args:
+        n (int): The max end of the range of numbers.
+        num_shards (int): The number of worker actors to create.
+        repeat (bool): Whether to cycle over the range forever.
+    """
+    generators = []
+    shard_size = n // num_shards
+    for i in range(num_shards):
+        start = i * shard_size
+        if i == num_shards - 1:
+            end = n
+        else:
+            end = (i + 1) * shard_size
+        generators.append(range(start, end))
+    name = "from_range[{}, shards={}{}]".format(
+        n, num_shards, ", repeat=True" if repeat else "")
+    return from_iterators(generators, repeat=repeat, name=name)
+
+
+def from_iterators(generators: List[Iterable[T]],
+                   repeat: bool = False,
+                   name=None) -> "ParallelIterator[T]":
+    """Create a parallel iterator from a set of iterators.
+
+    An actor will be created for each iterator.
+
+    Examples:
+        >>> # Create using a list of generators.
+        >>> from_iterators([range(100), range(100)])
+
+        >>> # Equivalent to the above.
+        >>> from_iterators([lambda: range(100), lambda: range(100)])
+
+    Args:
+        generators (list): A list of Python generator objects or lambda
+            functions that produced a generator when called. We allow lambda
+            functions since the generator itself might not be serializable,
+            but a lambda that returns it can be.
+        repeat (bool): Whether to cycle over the iterators forever.
+        name (str): Optional name to give the iterator.
+    """
+    worker_cls = ray.remote(ParallelIteratorWorker)
+    actors = [worker_cls.remote(g, repeat) for g in generators]
+    if not name:
+        name = "from_iterators[shards={}{}]".format(
+            len(generators), ", repeat=True" if repeat else "")
+    return from_actors(actors, name=name)
+
+
+def from_actors(actors: List["ray.actor.ActorHandle"],
+                name=None) -> "ParallelIterator[T]":
+    """Create a parallel iterator from an existing set of actors.
+
+    Each actor must subclass the ParallelIteratorWorker interface.
+
+    Args:
+        actors (list): List of actors that each implement
+            ParallelIteratorWorker.
+        name (str): Optional name to give the iterator.
+    """
+    if not name:
+        name = "from_actors[shards={}]".format(len(actors))
+    return ParallelIterator([_ActorSet(actors, [])], name)
+
+
+class ParallelIterator(Generic[T]):
+    """A parallel iterator over a set of remote actors.
+
+    This can be used to iterate over a fixed set of task results
+    (like an actor pool), or a stream of data (e.g., a fixed range of numbers,
+    an infinite stream of RLlib rollout results).
+
+    This class is **serializable** and can be passed to other remote
+    tasks and actors. However, each shard should be read from at most one
+    process at a time.
+
+    Examples:
+        >>> # Applying a function over items in parallel.
+        >>> it = ray.util.iter.from_items([1, 2, 3], num_shards=2)
+        ... <__main__.ParallelIterator object>
+        >>> it = it.for_each(lambda x: x * 2).gather_sync()
+        ... <__main__.LocalIterator object>
+        >>> print(list(it))
+        ... [2, 4, 6]
+
+        >>> # Creating from generators.
+        >>> it = ray.util.iter.from_iterators([range(3), range(3)])
+        ... <__main__.ParallelIterator object>
+        >>> print(list(it.gather_sync()))
+        ... [0, 0, 1, 1, 2, 2]
+
+        >>> # Accessing the individual shards of an iterator.
+        >>> it = ray.util.iter.from_range(10, num_shards=2)
+        ... <__main__.ParallelIterator object>
+        >>> it0 = it.get_shard(0)
+        ... <__main__.LocalIterator object>
+        >>> print(list(it0))
+        ... [0, 1, 2, 3, 4]
+        >>> it1 = it.get_shard(1)
+        ... <__main__.LocalIterator object>
+        >>> print(list(it1))
+        ... [5, 6, 7, 8, 9]
+
+        >>> # Gathering results from actors synchronously in parallel.
+        >>> it = ray.util.iter.from_actors(workers)
+        ... <__main__.ParallelIterator object>
+        >>> it = it.batch_across_shards()
+        ... <__main__.LocalIterator object>
+        >>> print(next(it))
+        ... [worker_1_result_1, worker_2_result_1]
+        >>> print(next(it))
+        ... [worker_1_result_2, worker_2_result_2]
+    """
+
+    def __init__(self, actor_sets: List["_ActorSet"], name: str):
+        # We track multiple sets of actors to support parallel .union().
+        self.actor_sets = actor_sets
+        self.name = name
+
+    def __iter__(self):
+        raise TypeError(
+            "You must use it.gather_sync() or it.gather_async() to "
+            "iterate over the results of a ParallelIterator.")
+
+    def __str__(self):
+        return repr(self)
+
+    def __repr__(self):
+        return "ParallelIterator[{}]".format(self.name)
+
+    def for_each(self, fn: Callable[[T], U]) -> "ParallelIterator[U]":
+        """Remotely apply fn to each item in this iterator.
+
+        Args:
+            fn (func): function to apply to each item.
+
+        Examples:
+            >>> next(from_range(4).for_each(lambda x: x * 2).gather_sync())
+            ... [0, 2, 4, 8]
+        """
+        return ParallelIterator(
+            [
+                a.with_transform(lambda local_it: local_it.for_each(fn))
+                for a in self.actor_sets
+            ],
+            name=self.name + ".for_each()")
+
+    def filter(self, fn: Callable[[T], bool]) -> "ParallelIterator[T]":
+        """Remotely filter items from this iterator.
+
+        Args:
+            fn (func): returns False for items to drop from the iterator.
+
+        Examples:
+            >>> it = from_items([0, 1, 2]).filter(lambda x: x > 0)
+            >>> next(it.gather_sync())
+            ... [1, 2]
+        """
+        return ParallelIterator(
+            [
+                a.with_transform(lambda local_it: local_it.filter(fn))
+                for a in self.actor_sets
+            ],
+            name=self.name + ".filter()")
+
+    def batch(self, n: int) -> "ParallelIterator[List[T]]":
+        """Remotely batch together items in this iterator.
+
+        Args:
+            n (int): Number of items to batch together.
+
+        Examples:
+            >>> next(from_range(10, 1).batch(4).gather_sync())
+            ... [0, 1, 2, 3]
+        """
+        return ParallelIterator(
+            [
+                a.with_transform(lambda local_it: local_it.batch(n))
+                for a in self.actor_sets
+            ],
+            name=self.name + ".batch({})".format(n))
+
+    def flatten(self) -> "ParallelIterator[T[0]]":
+        """Flatten batches of items into individual items.
+
+        Examples:
+            >>> next(from_range(10, 1).batch(4).flatten())
+            ... 0
+        """
+        return ParallelIterator(
+            [
+                a.with_transform(lambda local_it: local_it.flatten())
+                for a in self.actor_sets
+            ],
+            name=self.name + ".flatten()")
+
+    def combine(self, fn: Callable[[T], List[U]]) -> "ParallelIterator[U]":
+        """Transform and then combine items horizontally.
+
+        This is the equivalent of for_each(fn).flatten() (flat map).
+        """
+        it = self.for_each(fn).flatten()
+        it.name = self.name + ".combine()"
+        return it
+
+    def local_shuffle(self, shuffle_buffer_size: int,
+                      seed: int = None) -> "ParallelIterator[T]":
+        """Remotely shuffle items of each shard independently
+
+        Args:
+            shuffle_buffer_size (int): The algorithm fills a buffer with
+                shuffle_buffer_size elements and randomly samples elements from
+                this buffer, replacing the selected elements with new elements.
+                For perfect shuffling, this argument should be greater than or
+                equal to the largest iterator size.
+            seed (int): Seed to use for
+                randomness. Default value is None.
+
+        Returns:
+            Returns a ParallelIterator with a local shuffle applied on the
+                base iterator
+
+        Examples:
+            >>> it = from_range(10, 1).local_shuffle(shuffle_buffer_size=2)
+            >>> it = it.gather_sync()
+            >>> next(it)
+            0
+            >>> next(it)
+            2
+            >>> next(it)
+            3
+            >>> next(it)
+            1
+        """
+        return ParallelIterator(
+            [
+                a.with_transform(
+                    lambda localit: localit.shuffle(shuffle_buffer_size, seed))
+                for a in self.actor_sets
+            ],
+            name=self.name +
+            ".local_shuffle(shuffle_buffer_size={}, seed={})".format(
+                shuffle_buffer_size,
+                str(seed) if seed is not None else "None"))
+
+    def gather_sync(self) -> "LocalIterator[T]":
+        """Returns a local iterable for synchronous iteration.
+
+        New items will be fetched from the shards on-demand as the iterator
+        is stepped through.
+
+        This is the equivalent of batch_across_shards().flatten().
+
+        Examples:
+            >>> it = from_range(100, 1).gather_sync()
+            >>> next(it)
+            ... 0
+            >>> next(it)
+            ... 1
+            >>> next(it)
+            ... 2
+        """
+        it = self.batch_across_shards().flatten()
+        it.name = "{}.gather_sync()".format(self)
+        return it
+
+    def batch_across_shards(self) -> "LocalIterator[List[T]]":
+        """Iterate over the results of multiple shards in parallel.
+
+        Examples:
+            >>> it = from_iterators([range(3), range(3)])
+            >>> next(it.batch_across_shards())
+            ... [0, 0]
+        """
+
+        def base_iterator(timeout=None):
+            active = []
+            for actor_set in self.actor_sets:
+                actor_set.init_actors()
+                active.extend(actor_set.actors)
+            futures = [a.par_iter_next.remote() for a in active]
+            while active:
+                try:
+                    yield ray.get(futures, timeout=timeout)
+                    futures = [a.par_iter_next.remote() for a in active]
+                    # Always yield after each round of gets with timeout.
+                    if timeout is not None:
+                        yield _NextValueNotReady()
+                except TimeoutError:
+                    yield _NextValueNotReady()
+                except StopIteration:
+                    # Find and remove the actor that produced StopIteration.
+                    results = []
+                    for a, f in zip(list(active), futures):
+                        try:
+                            results.append(ray.get(f))
+                        except StopIteration:
+                            active.remove(a)
+                    if results:
+                        yield results
+                    futures = [a.par_iter_next.remote() for a in active]
+
+        name = "{}.batch_across_shards()".format(self)
+        return LocalIterator(base_iterator, name=name)
+
+    def gather_async(self) -> "LocalIterator[T]":
+        """Returns a local iterable for asynchronous iteration.
+
+        New items will be fetched from the shards asynchronously as soon as
+        the previous one is computed. Items arrive in non-deterministic order.
+
+        Examples:
+            >>> it = from_range(100, 1).gather_async()
+            >>> next(it)
+            ... 3
+            >>> next(it)
+            ... 0
+            >>> next(it)
+            ... 1
+        """
+
+        def base_iterator(timeout=None):
+            all_actors = []
+            for actor_set in self.actor_sets:
+                actor_set.init_actors()
+                all_actors.extend(actor_set.actors)
+            futures = {}
+            for a in all_actors:
+                futures[a.par_iter_next.remote()] = a
+            while futures:
+                pending = list(futures)
+                if timeout is None:
+                    # First try to do a batch wait for efficiency.
+                    ready, _ = ray.wait(
+                        pending, num_returns=len(pending), timeout=0)
+                    # Fall back to a blocking wait.
+                    if not ready:
+                        ready, _ = ray.wait(pending, num_returns=1)
+                else:
+                    ready, _ = ray.wait(
+                        pending, num_returns=len(pending), timeout=timeout)
+                for obj_id in ready:
+                    actor = futures.pop(obj_id)
+                    try:
+                        yield ray.get(obj_id)
+                        futures[actor.par_iter_next.remote()] = actor
+                    except StopIteration:
+                        pass
+                # Always yield after each round of wait with timeout.
+                if timeout is not None:
+                    yield _NextValueNotReady()
+
+        name = "{}.gather_async()".format(self)
+        return LocalIterator(base_iterator, name=name)
+
+    def take(self, n: int) -> List[T]:
+        """Return up to the first n items from this iterator."""
+        return self.gather_sync().take(n)
+
+    def show(self, n: int = 20):
+        """Print up to the first n items from this iterator."""
+        return self.gather_sync().show(n)
+
+    def union(self, other: "ParallelIterator[T]") -> "ParallelIterator[T]":
+        """Return an iterator that is the union of this and the other."""
+        if not isinstance(other, ParallelIterator):
+            raise ValueError(
+                "other must be of type ParallelIterator, got {}".format(
+                    type(other)))
+        actor_sets = []
+        actor_sets.extend(self.actor_sets)
+        actor_sets.extend(other.actor_sets)
+        return ParallelIterator(actor_sets, "ParallelUnion[{}, {}]".format(
+            self, other))
+
+    def num_shards(self) -> int:
+        """Return the number of worker actors backing this iterator."""
+        return sum(len(a.actors) for a in self.actor_sets)
+
+    def shards(self) -> List["LocalIterator[T]"]:
+        """Return the list of all shards."""
+        return [self.get_shard(i) for i in range(self.num_shards())]
+
+    def get_shard(self, shard_index: int) -> "LocalIterator[T]":
+        """Return a local iterator for the given shard.
+
+        The iterator is guaranteed to be serializable and can be passed to
+        remote tasks or actors.
+        """
+        a, t = None, None
+        i = shard_index
+        for actor_set in self.actor_sets:
+            if i < len(actor_set.actors):
+                a = actor_set.actors[i]
+                t = actor_set.transforms
+                break
+            else:
+                i -= len(actor_set.actors)
+        if a is None:
+            raise ValueError("Shard index out of range", shard_index,
+                             self.num_shards())
+
+        def base_iterator(timeout=None):
+            ray.get(a.par_iter_init.remote(t))
+            while True:
+                try:
+                    yield ray.get(a.par_iter_next.remote(), timeout=timeout)
+                    # Always yield after each round of gets with timeout.
+                    if timeout is not None:
+                        yield _NextValueNotReady()
+                except TimeoutError:
+                    yield _NextValueNotReady()
+                except StopIteration:
+                    break
+
+        name = self.name + ".shard[{}]".format(shard_index)
+        return LocalIterator(base_iterator, name=name)
+
+
+class LocalIterator(Generic[T]):
+    """An iterator over a single shard of data.
+
+    It implements similar transformations as ParallelIterator[T], but the
+    transforms will be applied locally and not remotely in parallel.
+
+    This class is **serializable** and can be passed to other remote
+    tasks and actors. However, it should be read from at most one process at
+    a time."""
+
+    def __init__(self,
+                 base_iterator: Callable[[], Iterable[T]],
+                 local_transforms: List[Callable[[Iterable], Any]] = None,
+                 timeout: int = None,
+                 name=None):
+        """Create a local iterator (this is an internal function).
+
+        Args:
+            base_iterator (func): A function that produces the base iterator.
+                This is a function so that we can ensure LocalIterator is
+                serializable.
+            local_transforms (list): A list of transformation functions to be
+                applied on top of the base iterator. When iteration begins, we
+                create the base iterator and apply these functions. This lazy
+                creation ensures LocalIterator is serializable until you start
+                iterating over it.
+            timeout (int): Optional timeout in seconds for this iterator, after
+                which _NextValueNotReady will be returned. This avoids
+                blocking.
+            name (str): Optional name for this iterator.
+        """
+        self.base_iterator = base_iterator
+        self.built_iterator = None
+        self.local_transforms = local_transforms or []
+        self.timeout = timeout
+        self.name = name or "unknown"
+
+    def _build_once(self):
+        if self.built_iterator is None:
+            it = iter(self.base_iterator(self.timeout))
+            for fn in self.local_transforms:
+                it = fn(it)
+            self.built_iterator = it
+
+    def __iter__(self):
+        self._build_once()
+        return self.built_iterator
+
+    def __next__(self):
+        self._build_once()
+        return next(self.built_iterator)
+
+    def __str__(self):
+        return repr(self)
+
+    def __repr__(self):
+        return "LocalIterator[{}]".format(self.name)
+
+    def for_each(self, fn: Callable[[T], U]) -> "LocalIterator[U]":
+        def apply_foreach(it):
+            for item in it:
+                if isinstance(item, _NextValueNotReady):
+                    yield item
+                else:
+                    yield fn(item)
+
+        return LocalIterator(
+            self.base_iterator,
+            self.local_transforms + [apply_foreach],
+            name=self.name + ".for_each()")
+
+    def filter(self, fn: Callable[[T], bool]) -> "LocalIterator[T]":
+        def apply_filter(it):
+            for item in it:
+                if isinstance(item, _NextValueNotReady) or fn(item):
+                    yield item
+
+        return LocalIterator(
+            self.base_iterator,
+            self.local_transforms + [apply_filter],
+            name=self.name + ".filter()")
+
+    def batch(self, n: int) -> "LocalIterator[List[T]]":
+        def apply_batch(it):
+            batch = []
+            for item in it:
+                if isinstance(item, _NextValueNotReady):
+                    yield item
+                else:
+                    batch.append(item)
+                    if len(batch) >= n:
+                        yield batch
+                        batch = []
+            if batch:
+                yield batch
+
+        return LocalIterator(
+            self.base_iterator,
+            self.local_transforms + [apply_batch],
+            name=self.name + ".batch({})".format(n))
+
+    def flatten(self) -> "LocalIterator[T[0]]":
+        def apply_flatten(it):
+            for item in it:
+                if isinstance(item, _NextValueNotReady):
+                    yield item
+                else:
+                    for subitem in item:
+                        yield subitem
+
+        return LocalIterator(
+            self.base_iterator,
+            self.local_transforms + [apply_flatten],
+            name=self.name + ".flatten()")
+
+    def shuffle(self, shuffle_buffer_size: int,
+                seed: int = None) -> "LocalIterator[T]":
+        """Shuffle items of this iterator
+
+        Args:
+            shuffle_buffer_size (int): The algorithm fills a buffer with
+                shuffle_buffer_size elements and randomly samples elements from
+                this buffer, replacing the selected elements with new elements.
+                For perfect shuffling, this argument should be greater than or
+                equal to the largest iterator size.
+            seed (int): Seed to use for
+                randomness. Default value is None.
+
+        Returns:
+            A new LocalIterator with shuffling applied
+        """
+        shuffle_random = random.Random(seed)
+
+        def apply_shuffle(it):
+            buffer = []
+            for item in it:
+                if isinstance(item, _NextValueNotReady):
+                    yield item
+                else:
+                    buffer.append(item)
+                    if len(buffer) >= shuffle_buffer_size:
+                        yield buffer.pop(
+                            shuffle_random.randint(0,
+                                                   len(buffer) - 1))
+            while len(buffer) > 0:
+                yield buffer.pop(shuffle_random.randint(0, len(buffer) - 1))
+
+        return LocalIterator(
+            self.base_iterator,
+            self.local_transforms + [apply_shuffle],
+            name=self.name +
+            ".shuffle(shuffle_buffer_size={}, seed={})".format(
+                shuffle_buffer_size,
+                str(seed) if seed is not None else "None"))
+
+    def combine(self, fn: Callable[[T], List[U]]) -> "LocalIterator[U]":
+        it = self.for_each(fn).flatten()
+        it.name = self.name + ".combine()"
+        return it
+
+    def take(self, n: int) -> List[T]:
+        """Return up to the first n items from this iterator."""
+        out = []
+        for item in self:
+            out.append(item)
+            if len(out) >= n:
+                break
+        return out
+
+    def show(self, n: int = 20):
+        """Print up to the first n items from this iterator."""
+        i = 0
+        for item in self:
+            print(item)
+            i += 1
+            if i >= n:
+                break
+
+    def union(self, other: "LocalIterator[T]") -> "LocalIterator[T]":
+        """Return an iterator that is the union of this and the other.
+
+        There are no ordering guarantees between the two iterators. We make a
+        best-effort attempt to return items from both as they become ready,
+        preventing starvation of any particular iterator.
+        """
+
+        if not isinstance(other, LocalIterator):
+            raise ValueError(
+                "other must be of type LocalIterator, got {}".format(
+                    type(other)))
+
+        it1 = LocalIterator(
+            self.base_iterator, self.local_transforms, timeout=0)
+        it2 = LocalIterator(
+            other.base_iterator, other.local_transforms, timeout=0)
+        active = [it1, it2]
+
+        def build_union(timeout=None):
+            while True:
+                for it in list(active):
+                    # Yield items from the iterator until _NextValueNotReady is
+                    # found, then switch to the next iterator.
+                    try:
+                        while True:
+                            item = next(it)
+                            if isinstance(item, _NextValueNotReady):
+                                break
+                            else:
+                                yield item
+                    except StopIteration:
+                        active.remove(it)
+                if not active:
+                    break
+
+        return LocalIterator(
+            build_union, [], name="LocalUnion[{}, {}]".format(self, other))
+
+
+class ParallelIteratorWorker(object):
+    """Worker actor for a ParallelIterator.
+
+    Actors that are passed to iter.from_actors() must subclass this interface.
+    """
+
+    def __init__(self, item_generator: Any, repeat: bool):
+        """Create an iterator worker.
+
+        Subclasses must call this init function.
+
+        Args:
+            item_generator (obj): A Python generator objects or lambda function
+                that produces a generator when called. We allow lambda
+                functions since the generator itself might not be serializable,
+                but a lambda that returns it can be.
+            repeat (bool): Whether to loop over the iterator forever.
+        """
+
+        def make_iterator():
+            if callable(item_generator):
+                return item_generator()
+            else:
+                return item_generator
+
+        if repeat:
+
+            def cycle():
+                while True:
+                    it = make_iterator()
+                    for item in it:
+                        yield item
+
+            self.item_generator = cycle()
+        else:
+            self.item_generator = make_iterator()
+
+        self.transforms = []
+        self.local_it = None
+
+    def par_iter_init(self, transforms):
+        """Implements ParallelIterator worker init."""
+        it = LocalIterator(lambda timeout: self.item_generator)
+        for fn in transforms:
+            it = fn(it)
+            assert it is not None, fn
+        self.local_it = iter(it)
+
+    def par_iter_next(self):
+        """Implements ParallelIterator worker item fetch."""
+        assert self.local_it is not None, "must call par_iter_init()"
+        return next(self.local_it)
+
+
+class _NextValueNotReady(Exception):
+    """Indicates that a local iterator has no value currently available.
+
+    This is used internally to implement the union() of multiple blocking
+    local generators."""
+    pass
+
+
+class _ActorSet(object):
+    """Helper class that represents a set of actors and transforms."""
+
+    def __init__(
+            self, actors: List["ray.actor.ActorHandle"],
+            transforms: List[Callable[["LocalIterator"], "LocalIterator"]]):
+        self.actors = actors
+        self.transforms = transforms
+
+    def init_actors(self):
+        ray.get([a.par_iter_init.remote(self.transforms) for a in self.actors])
+
+    def with_transform(self, fn):
+        return _ActorSet(self.actors, self.transforms + [fn])

python/ray/util/joblib/__init__.py

@@ -0,0 +1,17 @@
+from joblib.parallel import register_parallel_backend
+
+
+def register_ray():
+    """ Register Ray Backend to be called with parallel_backend("ray"). """
+    try:
+        from ray.util.joblib.ray_backend import RayBackend
+        register_parallel_backend("ray", RayBackend)
+    except ImportError:
+        msg = ("To use the ray backend you must install ray."
+               "Try running 'pip install ray'."
+               "See https://ray.readthedocs.io/en/latest/installation.html"
+               "for more information.")
+        raise ImportError(msg)
+
+
+__all__ = ["register_ray"]

python/ray/util/joblib/ray_backend.py

@@ -0,0 +1,58 @@
+from joblib._parallel_backends import MultiprocessingBackend
+from joblib.pool import PicklingPool
+import logging
+
+from ray.util.multiprocessing.pool import Pool
+import ray
+
+RAY_ADDRESS_ENV = "RAY_ADDRESS"
+
+logger = logging.getLogger(__name__)
+
+
+class RayBackend(MultiprocessingBackend):
+    """Ray backend uses ray, a system for scalable distributed computing.
+    More info about Ray is available here: https://ray.readthedocs.io.
+    """
+
+    def configure(self,
+                  n_jobs=1,
+                  parallel=None,
+                  prefer=None,
+                  require=None,
+                  **memmappingpool_args):
+        """Make Ray Pool the father class of PicklingPool. PicklingPool is a
+        father class that inherits Pool from multiprocessing.pool. The next
+        line is a patch, which changes the inheritance of Pool to be from
+        ray.util.multiprocessing.pool.
+        """
+        PicklingPool.__bases__ = (Pool, )
+        """Use all available resources when n_jobs == -1. Must set RAY_ADDRESS
+        variable in the environment or run ray.init(address=..) to run on
+        multiple nodes.
+        """
+        if n_jobs == -1:
+            if not ray.is_initialized():
+                import os
+                if RAY_ADDRESS_ENV in os.environ:
+                    ray_address = os.environ[RAY_ADDRESS_ENV]
+                    logger.info(
+                        "Connecting to ray cluster at address='{}'".format(
+                            ray_address))
+                    ray.init(address=ray_address)
+                else:
+                    logger.info("Starting local ray cluster")
+                    ray.init()
+            ray_cpus = int(ray.state.cluster_resources()["CPU"])
+            n_jobs = ray_cpus
+
+        eff_n_jobs = super(RayBackend, self).configure(
+            n_jobs, parallel, prefer, require, **memmappingpool_args)
+        return eff_n_jobs
+
+    def effective_n_jobs(self, n_jobs):
+        eff_n_jobs = super(RayBackend, self).effective_n_jobs(n_jobs)
+        if n_jobs == -1:
+            ray_cpus = int(ray.state.cluster_resources()["CPU"])
+            eff_n_jobs = ray_cpus
+        return eff_n_jobs

python/ray/util/named_actors.py

@@ -0,0 +1,62 @@
+import ray
+import ray.cloudpickle as pickle
+from ray.experimental.internal_kv import _internal_kv_get, _internal_kv_put
+
+
+def _calculate_key(name):
+    """Generate a Redis key with the given name.
+
+    Args:
+        name: The name of the named actor.
+
+    Returns:
+        The key to use for storing a named actor in Redis.
+    """
+    return b"Actor:" + name.encode("ascii")
+
+
+def get_actor(name):
+    """Get a named actor which was previously created.
+
+    If the actor doesn't exist, an exception will be raised.
+
+    Args:
+        name: The name of the named actor.
+
+    Returns:
+        The ActorHandle object corresponding to the name.
+    """
+    actor_name = _calculate_key(name)
+    pickled_state = _internal_kv_get(actor_name)
+    if pickled_state is None:
+        raise ValueError("The actor with name={} doesn't exist".format(name))
+    handle = pickle.loads(pickled_state)
+    return handle
+
+
+def register_actor(name, actor_handle):
+    """Register a named actor under a string key.
+
+   Args:
+       name: The name of the named actor.
+       actor_handle: The actor object to be associated with this name
+   """
+    if not isinstance(name, str):
+        raise TypeError("The name argument must be a string.")
+    if not isinstance(actor_handle, ray.actor.ActorHandle):
+        raise TypeError("The actor_handle argument must be an ActorHandle "
+                        "object.")
+    actor_name = _calculate_key(name)
+
+    # First check if the actor already exists.
+    try:
+        get_actor(name)
+        exists = True
+    except ValueError:
+        exists = False
+
+    if exists:
+        raise ValueError("An actor with name={} already exists".format(name))
+
+    # Add the actor to Redis if it does not already exist.
+    _internal_kv_put(actor_name, pickle.dumps(actor_handle))

python/ray/util/sgd/__init__.py

@@ -0,0 +1,4 @@
+from ray.util.sgd.pytorch import PyTorchTrainer
+from ray.util.sgd.tf import TFTrainer
+
+__all__ = ["PyTorchTrainer", "TFTrainer"]

python/ray/util/sgd/pytorch/__init__.py

@@ -0,0 +1,15 @@
+import logging
+logger = logging.getLogger(__name__)
+
+PyTorchTrainer = None
+PyTorchTrainable = None
+
+try:
+    import torch  # noqa: F401
+
+    from ray.util.sgd.pytorch.pytorch_trainer import (PyTorchTrainer,
+                                                      PyTorchTrainable)
+
+    __all__ = ["PyTorchTrainer", "PyTorchTrainable"]
+except ImportError:
+    logger.warning("PyTorch not found. PyTorchTrainer will not be available")

python/ray/util/sgd/pytorch/distributed_pytorch_runner.py

@@ -0,0 +1,134 @@
+import collections
+from filelock import FileLock
+import logging
+import os
+import torch.nn as nn
+import torch.distributed as dist
+import torch.utils.data
+from torch.nn.parallel import DistributedDataParallel
+
+from ray.util.sgd.pytorch.pytorch_runner import PyTorchRunner
+
+logger = logging.getLogger(__name__)
+
+
+class DistributedPyTorchRunner(PyTorchRunner):
+    """Manages a distributed PyTorch model replica.
+
+
+    Args:
+        args: Arguments for PyTorchRunner.
+        backend (string): backend used by distributed PyTorch.
+        kwargs: Keyword arguments for PyTorchRunner.
+
+    """
+
+    def __init__(self, *args, backend="gloo", **kwargs):
+        super(DistributedPyTorchRunner, self).__init__(*args, **kwargs)
+        self.backend = backend
+
+    def setup(self, url, world_rank, world_size):
+        """Connects to the distributed PyTorch backend and initializes the model.
+
+        Args:
+            url (str): the URL used to connect to distributed PyTorch.
+            world_rank (int): the index of the runner.
+            world_size (int): the total number of runners.
+        """
+        self._setup_distributed_pytorch(url, world_rank, world_size)
+        self._setup_training()
+
+    def _setup_distributed_pytorch(self, url, world_rank, world_size):
+        with self._timers["setup_proc"]:
+            self.world_rank = world_rank
+            logger.debug(
+                "Connecting to {} world_rank: {} world_size: {}".format(
+                    url, world_rank, world_size))
+            logger.debug("using {}".format(self.backend))
+            dist.init_process_group(
+                backend=self.backend,
+                init_method=url,
+                rank=world_rank,
+                world_size=world_size)
+
+    def _setup_training(self):
+        logger.debug("Creating model")
+        self.models = self.model_creator(self.config)
+        if not isinstance(self.models, collections.Iterable):
+            self.models = [self.models]
+        assert all(isinstance(model, nn.Module) for model in self.models), (
+            "All models must be PyTorch models: {}.".format(self.models))
+        if torch.cuda.is_available():
+            self.models = [model.cuda() for model in self.models]
+
+        logger.debug("Creating optimizer.")
+        self.optimizers = self.optimizer_creator(self.given_models,
+                                                 self.config)
+        if not isinstance(self.optimizers, collections.Iterable):
+            self.optimizers = [self.optimizers]
+
+        self._create_schedulers_if_available()
+
+        self._try_setup_apex()
+
+        # This needs to happen after apex
+        self.models = [DistributedDataParallel(model) for model in self.models]
+
+        logger.debug("Creating loss.")
+        self._create_loss()
+
+        logger.debug("Creating dataset.")
+        with FileLock(os.path.expanduser("~/.ray_data.lock")):
+            datasets = self.data_creator(self.config)
+            train_set, val_set = self._validate_datasets(datasets)
+
+        train_loader_config = self.dataloader_config.copy()
+        train_loader_config.update(
+            sampler=torch.utils.data.distributed.DistributedSampler(train_set),
+            shuffle=False)
+
+        self.train_loader = torch.utils.data.DataLoader(
+            train_set, batch_size=self.batch_size, **train_loader_config)
+
+        self.validation_loader = None
+        if val_set:
+            self.validation_loader = torch.utils.data.DataLoader(
+                val_set, batch_size=self.batch_size, **self.dataloader_config)
+
+    def step(self):
+        """Runs a training epoch and updates the model parameters.
+
+        Automatically sets epoch of sampler if possible.
+        """
+        logger.debug("Starting step")
+        if hasattr(self.train_loader.sampler, "set_epoch"):
+            self.train_loader.sampler.set_epoch(self.epoch)
+        return super(DistributedPyTorchRunner, self).step()
+
+    def _get_model_state_dicts(self):
+        """Fetch state from ``model.module`` instead of ``model``.
+
+        This is needed for PyTorch DistributedDataParallel models.
+        """
+        cpu_state_dicts = []
+        for model in self.models:
+            state_dict = model.module.state_dict()
+            # This is so that we create a duplicate of weights into CPU rather
+            # than move the model weights out of the GPU so that we can
+            # resume training while saving intermediate checkpoints.
+            cpu_state_dicts += [{k: v.cpu() for k, v in state_dict.items()}]
+        return cpu_state_dicts
+
+    def _set_model_state_dicts(self, model_state_dicts):
+        for model, model_state_dict in zip(self.models, model_state_dicts):
+            model.module.load_state_dict(model_state_dict)
+
+    # def shutdown(self):
+        """Attempts to shut down the worker."""
+        # super(DistributedPyTorchRunner, self).shutdown()
+        # TODO: Temporarily removing since it causes hangs on MacOSX.
+        # However, it seems to be harmless to remove permanently
+        # since the processes are shutdown anyways. This comment can be
+        # removed in a future release if it is still not documented
+        # the stable Pytorch docs.
+        # dist.destroy_process_group()

python/ray/util/sgd/pytorch/examples/cifar_pytorch_example.py

@@ -0,0 +1,161 @@
+import os
+import torch
+import torch.nn as nn
+import argparse
+from ray import tune
+import torch.utils.data
+import torchvision
+import torchvision.transforms as transforms
+
+import ray
+from ray.util.sgd.pytorch import (PyTorchTrainer, PyTorchTrainable)
+from ray.util.sgd.pytorch.resnet import ResNet18
+from ray.util.sgd.pytorch.utils import TEST_MODE
+
+
+def initialization_hook(runner):
+    print("NCCL DEBUG SET")
+    # Need this for avoiding a connection restart issue
+    os.environ["NCCL_SOCKET_IFNAME"] = "^docker0,lo"
+    os.environ["NCCL_LL_THRESHOLD"] = "0"
+    os.environ["NCCL_DEBUG"] = "INFO"
+
+
+def cifar_creator(config):
+    transform_train = transforms.Compose([
+        transforms.RandomCrop(32, padding=4),
+        transforms.RandomHorizontalFlip(),
+        transforms.ToTensor(),
+        transforms.Normalize((0.4914, 0.4822, 0.4465),
+                             (0.2023, 0.1994, 0.2010)),
+    ])  # meanstd transformation
+
+    transform_test = transforms.Compose([
+        transforms.ToTensor(),
+        transforms.Normalize((0.4914, 0.4822, 0.4465),
+                             (0.2023, 0.1994, 0.2010)),
+    ])
+    train_dataset = torchvision.datasets.CIFAR10(
+        root="~/data", train=True, download=True, transform=transform_train)
+    validation_dataset = torchvision.datasets.CIFAR10(
+        root="~/data", train=False, download=False, transform=transform_test)
+
+    return train_dataset, validation_dataset
+
+
+def optimizer_creator(model, config):
+    """Returns optimizer"""
+    return torch.optim.SGD(model.parameters(), lr=config.get("lr", 0.1))
+
+
+def scheduler_creator(optimizer, config):
+    return torch.optim.lr_scheduler.MultiStepLR(
+        optimizer, milestones=[150, 250, 350], gamma=0.1)
+
+
+def train_example(num_replicas=1,
+                  num_epochs=5,
+                  use_gpu=False,
+                  use_fp16=False,
+                  test_mode=False):
+    config = {TEST_MODE: test_mode}
+    trainer1 = PyTorchTrainer(
+        ResNet18,
+        cifar_creator,
+        optimizer_creator,
+        nn.CrossEntropyLoss,
+        scheduler_creator=scheduler_creator,
+        initialization_hook=initialization_hook,
+        num_replicas=num_replicas,
+        config=config,
+        use_gpu=use_gpu,
+        batch_size=16 if test_mode else 512,
+        backend="nccl" if use_gpu else "gloo",
+        scheduler_step_freq="epoch",
+        use_fp16=use_fp16)
+    for i in range(num_epochs):
+        # Increase `max_retries` to turn on fault tolerance.
+        stats = trainer1.train(max_retries=0)
+        print(stats)
+
+    print(trainer1.validate())
+    trainer1.shutdown()
+    print("success!")
+
+
+def tune_example(num_replicas=1, use_gpu=False, test_mode=False):
+    config = {
+        "model_creator": ResNet18,
+        "data_creator": cifar_creator,
+        "optimizer_creator": optimizer_creator,
+        "loss_creator": lambda config: nn.CrossEntropyLoss(),
+        "num_replicas": num_replicas,
+        "initialization_hook": initialization_hook,
+        "use_gpu": use_gpu,
+        "batch_size": 16 if test_mode else 512,
+        "config": {
+            "lr": tune.choice([1e-4, 1e-3, 5e-3, 1e-2]),
+            TEST_MODE: test_mode
+        },
+        "backend": "nccl" if use_gpu else "gloo"
+    }
+
+    analysis = tune.run(
+        PyTorchTrainable,
+        num_samples=2,
+        config=config,
+        stop={"training_iteration": 2},
+        verbose=2)
+
+    return analysis.get_best_config(metric="mean_accuracy", mode="max")
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--address",
+        required=False,
+        type=str,
+        help="the address to use for Redis")
+    parser.add_argument(
+        "--num-replicas",
+        "-n",
+        type=int,
+        default=1,
+        help="Sets number of replicas for training.")
+    parser.add_argument(
+        "--num-epochs", type=int, default=5, help="Number of epochs to train.")
+    parser.add_argument(
+        "--use-gpu",
+        action="store_true",
+        default=False,
+        help="Enables GPU training")
+    parser.add_argument(
+        "--fp16",
+        action="store_true",
+        default=False,
+        help="Enables FP16 training with apex. Requires `use-gpu`.")
+    parser.add_argument(
+        "--smoke-test",
+        action="store_true",
+        default=False,
+        help="Finish quickly for testing.")
+    parser.add_argument(
+        "--tune", action="store_true", default=False, help="Tune training")
+
+    args, _ = parser.parse_known_args()
+
+    ray.init(address=args.address, log_to_driver=True)
+
+    if args.tune:
+        tune_example(
+            num_replicas=args.num_replicas,
+            use_gpu=args.use_gpu,
+            test_mode=args.smoke_test)
+    else:
+        train_example(
+            num_replicas=args.num_replicas,
+            num_epochs=args.num_epochs,
+            use_gpu=args.use_gpu,
+            use_fp16=args.fp16,
+            test_mode=args.smoke_test)

python/ray/util/sgd/pytorch/examples/dcgan.py

@@ -0,0 +1,266 @@
+#!/usr/bin/env python
+
+import argparse
+import os
+import torch
+import torch.nn as nn
+import torch.optim as optim
+import torch.utils.data
+import torchvision.datasets as dset
+import torchvision.transforms as transforms
+import numpy as np
+
+from torch.autograd import Variable
+from torch.nn import functional as F
+from scipy.stats import entropy
+
+import ray
+from ray.util.sgd import PyTorchTrainer
+from ray.util.sgd.pytorch.utils import TEST_MODE
+
+# Training parameters
+TRAIN_BATCHES = 5
+# Number of channels in the training images. For color images this is 3
+num_channels = 1
+
+# Size of z latent vector (i.e. size of generator input)
+latent_vector_size = 100
+
+# Size of feature maps in generator
+features_g = 32
+
+# Size of feature maps in discriminator
+features_d = 32
+
+
+def data_creator(config):
+    return dset.MNIST(
+        root="~/mnist/",
+        download=True,
+        transform=transforms.Compose([
+            transforms.Resize(32),
+            transforms.ToTensor(),
+            transforms.Normalize((0.5, ), (0.5, )),
+        ]))
+
+
+def weights_init(m):
+    classname = m.__class__.__name__
+    if classname.find("Conv") != -1:
+        nn.init.normal_(m.weight.data, 0.0, 0.02)
+    elif classname.find("BatchNorm") != -1:
+        nn.init.normal_(m.weight.data, 1.0, 0.02)
+        nn.init.constant_(m.bias.data, 0)
+
+
+class Generator(nn.Module):
+    def __init__(self):
+        super(Generator, self).__init__()
+        self.main = nn.Sequential(
+            # input is Z, going into a convolution
+            nn.ConvTranspose2d(
+                latent_vector_size, features_g * 4, 4, 1, 0, bias=False),
+            nn.BatchNorm2d(features_g * 4),
+            nn.ReLU(True),
+            nn.ConvTranspose2d(
+                features_g * 4, features_g * 2, 4, 2, 1, bias=False),
+            nn.BatchNorm2d(features_g * 2),
+            nn.ReLU(True),
+            nn.ConvTranspose2d(
+                features_g * 2, features_g, 4, 2, 1, bias=False),
+            nn.BatchNorm2d(features_g),
+            nn.ReLU(True),
+            nn.ConvTranspose2d(features_g, num_channels, 4, 2, 1, bias=False),
+            nn.Tanh())
+
+    def forward(self, input):
+        return self.main(input)
+
+
+class Discriminator(nn.Module):
+    def __init__(self):
+        super(Discriminator, self).__init__()
+        self.main = nn.Sequential(
+            nn.Conv2d(num_channels, features_d, 4, 2, 1, bias=False),
+            nn.LeakyReLU(0.2, inplace=True),
+            nn.Conv2d(features_d, features_d * 2, 4, 2, 1, bias=False),
+            nn.BatchNorm2d(features_d * 2), nn.LeakyReLU(0.2, inplace=True),
+            nn.Conv2d(features_d * 2, features_d * 4, 4, 2, 1, bias=False),
+            nn.BatchNorm2d(features_d * 4), nn.LeakyReLU(0.2, inplace=True),
+            nn.Conv2d(features_d * 4, 1, 4, 1, 0, bias=False), nn.Sigmoid())
+
+    def forward(self, input):
+        return self.main(input)
+
+
+class Net(nn.Module):
+    """LeNet for MNist classification, used for inception_score."""
+
+    def __init__(self):
+        super(Net, self).__init__()
+        self.conv1 = nn.Conv2d(1, 10, kernel_size=5)
+        self.conv2 = nn.Conv2d(10, 20, kernel_size=5)
+        self.conv2_drop = nn.Dropout2d()
+        self.fc1 = nn.Linear(320, 50)
+        self.fc2 = nn.Linear(50, 10)
+
+    def forward(self, x):
+        x = F.relu(F.max_pool2d(self.conv1(x), 2))
+        x = F.relu(F.max_pool2d(self.conv2_drop(self.conv2(x)), 2))
+        x = x.view(-1, 320)
+        x = F.relu(self.fc1(x))
+        x = F.dropout(x, training=self.training)
+        x = self.fc2(x)
+        return F.log_softmax(x, dim=1)
+
+
+def inception_score(imgs, batch_size=32, splits=1):
+    N = len(imgs)
+    dtype = torch.FloatTensor
+    dataloader = torch.utils.data.DataLoader(imgs, batch_size=batch_size)
+    cm = Net()
+    cm.load_state_dict(torch.load(model_path))
+    cm.eval()
+    up = nn.Upsample(size=(28, 28), mode="bilinear").type(dtype)
+
+    def get_pred(x):
+        x = up(x)
+        x = cm(x)
+        return F.softmax(x).data.cpu().numpy()
+
+    preds = np.zeros((N, 10))
+    for i, batch in enumerate(dataloader, 0):
+        batch = batch.type(dtype)
+        batchv = Variable(batch)
+        batch_size_i = batch.size()[0]
+        preds[i * batch_size:i * batch_size + batch_size_i] = get_pred(batchv)
+
+    # Now compute the mean kl-div
+    split_scores = []
+    for k in range(splits):
+        part = preds[k * (N // splits):(k + 1) * (N // splits), :]
+        py = np.mean(part, axis=0)
+        scores = []
+        for i in range(part.shape[0]):
+            pyx = part[i, :]
+            scores.append(entropy(pyx, py))
+        split_scores.append(np.exp(np.mean(scores)))
+
+    return np.mean(split_scores), np.std(split_scores)
+
+
+def model_creator(config):
+    netD = Discriminator()
+    netD.apply(weights_init)
+
+    netG = Generator()
+    netG.apply(weights_init)
+    return netD, netG
+
+
+def train(config, models, dataloader, criterion, optimizers, **kwargs):
+    netD, netG = models
+    optimD, optimG = optimizers
+    real_label = 1
+    fake_label = 0
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+    for i, data in enumerate(dataloader, 0):
+        if i >= TRAIN_BATCHES and config.get(TEST_MODE):
+            break
+
+        netD.zero_grad()
+        real_cpu = data[0].to(device)
+        b_size = real_cpu.size(0)
+        label = torch.full((b_size, ), real_label, device=device)
+        output = netD(real_cpu).view(-1)
+        errD_real = criterion(output, label)
+        errD_real.backward()
+
+        noise = torch.randn(b_size, latent_vector_size, 1, 1, device=device)
+        fake = netG(noise)
+        label.fill_(fake_label)
+        output = netD(fake.detach()).view(-1)
+        errD_fake = criterion(output, label)
+        errD_fake.backward()
+        errD = errD_real + errD_fake
+        optimD.step()
+
+        netG.zero_grad()
+        label.fill_(real_label)
+        output = netD(fake).view(-1)
+        errG = criterion(output, label)
+        errG.backward()
+        optimG.step()
+
+        is_score, is_std = inception_score(fake)
+
+    return {
+        "loss_g": errG.item(),
+        "loss_d": errD.item(),
+        "inception": is_score
+    }
+
+
+def optimizer_creator(models, config):
+    net_d, net_g = models
+    discriminator_opt = optim.Adam(
+        net_d.parameters(), lr=config.get("lr", 0.01), betas=(0.5, 0.999))
+    generator_opt = optim.Adam(
+        net_g.parameters(), lr=config.get("lr", 0.01), betas=(0.5, 0.999))
+    return discriminator_opt, generator_opt
+
+
+def train_example(num_replicas=1, use_gpu=False, test_mode=False):
+    config = {TEST_MODE: test_mode}
+    trainer = PyTorchTrainer(
+        model_creator,
+        data_creator,
+        optimizer_creator,
+        nn.BCELoss,
+        train_function=train,
+        validation_function=False,
+        num_replicas=num_replicas,
+        config=config,
+        use_gpu=use_gpu,
+        batch_size=16 if test_mode else 512,
+        backend="nccl" if use_gpu else "gloo")
+    for i in range(10):
+        stats = trainer.train(max_retries=3)
+        print(stats)
+
+    return trainer
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--smoke-test", action="store_true", help="Finish quickly for testing")
+    parser.add_argument(
+        "--address",
+        required=False,
+        type=str,
+        help="the address to use for Redis")
+    parser.add_argument(
+        "--num-replicas",
+        "-n",
+        type=int,
+        default=1,
+        help="Sets number of replicas for training.")
+    parser.add_argument(
+        "--use-gpu",
+        action="store_true",
+        default=False,
+        help="Enables GPU training")
+    args, _ = parser.parse_known_args()
+    ray.init(address=args.address)
+
+    path = os.path.dirname(ray.__file__)
+    model_path = os.path.join(path, "util/sgd/pytorch/examples/mnist_cnn.pt")
+    # load the pretrained mnist classification model for inception_score
+
+    trainer = train_example(
+        num_replicas=args.num_replicas,
+        use_gpu=args.use_gpu,
+        test_mode=args.smoke_test)
+    models = trainer.get_model()

python/ray/util/sgd/pytorch/examples/example-sgd.yaml

@@ -0,0 +1,71 @@
+# An unique identifier for the head node and workers of this cluster.
+cluster_name: sgd-pytorch
+
+# The maximum number of workers nodes to launch in addition to the head
+# node. This takes precedence over min_workers. min_workers default to 0.
+min_workers: 3
+initial_workers: 3
+max_workers: 3
+
+target_utilization_fraction: 0.9
+
+# If a node is idle for this many minutes, it will be removed.
+idle_timeout_minutes: 20
+# docker:
+#     image: tensorflow/tensorflow:1.5.0-py3
+#     container_name: ray_docker
+
+# Cloud-provider specific configuration.
+provider:
+    type: aws
+    region: us-east-1
+
+# How Ray will authenticate with newly launched nodes.
+auth:
+    ssh_user: ubuntu
+
+head_node:
+    InstanceType: p3.8xlarge
+    ImageId: ami-0757fc5a639fe7666
+    InstanceMarketOptions:
+        MarketType: spot
+           # SpotOptions:
+           #     MaxPrice: "9.0"
+
+
+worker_nodes:
+    InstanceType: p3.8xlarge
+    ImageId: ami-0757fc5a639fe7666
+    # Run workers on spot by default. Comment this out to use on-demand.
+    InstanceMarketOptions:
+        MarketType: spot
+        # SpotOptions:
+        #     MaxPrice: "9.0"
+
+setup_commands:
+    - ray || pip install -U https://s3-us-west-2.amazonaws.com/ray-wheels/latest/ray-0.9.0.dev0-cp36-cp36m-manylinux1_x86_64.whl
+    - pip install -U ipdb ray[rllib] torch torchvision
+    # Install apex.
+    # - rm -rf apex || true
+    # - git clone https://github.com/NVIDIA/apex && cd apex && pip install -v --no-cache-dir  ./ || true
+
+
+file_mounts: {
+}
+
+# Custom commands that will be run on the head node after common setup.
+head_setup_commands: []
+
+# Custom commands that will be run on worker nodes after common setup.
+worker_setup_commands: []
+
+# # Command to start ray on the head node. You don't need to change this.
+head_start_ray_commands:
+    - ray stop
+    - ray start --head --redis-port=6379 --object-manager-port=8076 --autoscaling-config=~/ray_bootstrap_config.yaml --object-store-memory=1000000000
+
+# Command to start ray on worker nodes. You don't need to change this.
+worker_start_ray_commands:
+    - ray stop
+    - ray start --address=$RAY_HEAD_IP:6379 --object-manager-port=8076 --object-store-memory=1000000000
+

python/ray/util/sgd/pytorch/examples/train_example.py

@@ -0,0 +1,112 @@
+"""
+This file holds code for a Training guide for PytorchSGD in the documentation.
+
+It ignores yapf because yapf doesn't allow comments right after code blocks,
+but we put comments right after code blocks to prevent large white spaces
+in the documentation.
+"""
+
+# yapf: disable
+# __torch_train_example__
+import argparse
+import numpy as np
+import torch
+import torch.nn as nn
+
+from ray.util.sgd import PyTorchTrainer
+
+
+class LinearDataset(torch.utils.data.Dataset):
+    """y = a * x + b"""
+
+    def __init__(self, a, b, size=1000):
+        x = np.arange(0, 10, 10 / size, dtype=np.float32)
+        self.x = torch.from_numpy(x)
+        self.y = torch.from_numpy(a * x + b)
+
+    def __getitem__(self, index):
+        return self.x[index, None], self.y[index, None]
+
+    def __len__(self):
+        return len(self.x)
+
+
+def model_creator(config):
+    """Returns a torch.nn.Module object."""
+    return nn.Linear(1, config.get("hidden_size", 1))
+
+
+def optimizer_creator(model, config):
+    """Returns optimizer defined upon the model parameters."""
+    return torch.optim.SGD(model.parameters(), lr=config.get("lr", 1e-2))
+
+
+def scheduler_creator(optimizer, config):
+    """Returns a learning rate scheduler wrapping the optimizer.
+
+    You will need to set ``PyTorchTrainer(scheduler_step_freq="epoch")``
+    for the scheduler to be incremented correctly.
+
+    If using a scheduler for validation loss, be sure to call
+    ``trainer.update_scheduler(validation_loss)``.
+    """
+    return torch.optim.lr_scheduler.StepLR(optimizer, step_size=5, gamma=0.9)
+
+
+def data_creator(config):
+    """Returns training dataloader, validation dataloader."""
+    return LinearDataset(2, 5),  LinearDataset(2, 5, size=400)
+
+
+def train_example(num_replicas=1, use_gpu=False):
+    trainer1 = PyTorchTrainer(
+        model_creator,
+        data_creator,
+        optimizer_creator,
+        loss_creator=nn.MSELoss,
+        scheduler_creator=scheduler_creator,
+        num_replicas=num_replicas,
+        use_gpu=use_gpu,
+        batch_size=num_replicas * 4,
+        config={"lr": 1e-2, "hidden_size": 1},
+        backend="gloo",
+        scheduler_step_freq="epoch")
+    for i in range(5):
+        stats = trainer1.train()
+        print(stats)
+
+    print(trainer1.validate())
+    m = trainer1.get_model()
+    print("trained weight: % .2f, bias: % .2f" % (
+        m.weight.item(), m.bias.item()))
+    trainer1.shutdown()
+    print("success!")
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--address",
+        required=False,
+        type=str,
+        help="the address to use for Ray")
+    parser.add_argument(
+        "--num-replicas",
+        "-n",
+        type=int,
+        default=1,
+        help="Sets number of replicas for training.")
+    parser.add_argument(
+        "--use-gpu",
+        action="store_true",
+        default=False,
+        help="Enables GPU training")
+    parser.add_argument(
+        "--tune", action="store_true", default=False, help="Tune training")
+
+    args, _ = parser.parse_known_args()
+
+    import ray
+
+    ray.init(address=args.address)
+    train_example(num_replicas=args.num_replicas, use_gpu=args.use_gpu)

python/ray/util/sgd/pytorch/examples/tune_example.py

@@ -0,0 +1,96 @@
+# yapf: disable
+"""
+This file holds code for a Distributed Pytorch + Tune page in the docs.
+
+It ignores yapf because yapf doesn't allow comments right after code blocks,
+but we put comments right after code blocks to prevent large white spaces
+in the documentation.
+"""
+
+# __torch_tune_example__
+import numpy as np
+import torch
+import torch.nn as nn
+
+import ray
+from ray import tune
+from ray.util.sgd.pytorch.pytorch_trainer import PyTorchTrainable
+
+
+class LinearDataset(torch.utils.data.Dataset):
+    """y = a * x + b"""
+
+    def __init__(self, a, b, size=1000):
+        x = np.random.random(size).astype(np.float32) * 10
+        x = np.arange(0, 10, 10 / size, dtype=np.float32)
+        self.x = torch.from_numpy(x)
+        self.y = torch.from_numpy(a * x + b)
+
+    def __getitem__(self, index):
+        return self.x[index, None], self.y[index, None]
+
+    def __len__(self):
+        return len(self.x)
+
+
+def model_creator(config):
+    return nn.Linear(1, 1)
+
+
+def optimizer_creator(model, config):
+    """Returns optimizer."""
+    return torch.optim.SGD(model.parameters(), lr=config.get("lr", 1e-4))
+
+
+def data_creator(config):
+    """Returns training dataloader, validation dataloader."""
+    return LinearDataset(2, 5), LinearDataset(2, 5, size=400)
+
+
+def tune_example(num_replicas=1, use_gpu=False):
+    config = {
+        "model_creator": tune.function(model_creator),
+        "data_creator": tune.function(data_creator),
+        "optimizer_creator": tune.function(optimizer_creator),
+        "loss_creator": tune.function(nn.MSELoss),
+        "num_replicas": num_replicas,
+        "use_gpu": use_gpu,
+        "batch_size": 512,
+        "backend": "gloo"
+    }
+
+    analysis = tune.run(
+        PyTorchTrainable,
+        num_samples=12,
+        config=config,
+        stop={"training_iteration": 2},
+        verbose=1)
+
+    return analysis.get_best_config(metric="validation_loss", mode="min")
+
+
+if __name__ == "__main__":
+    import argparse
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--address",
+        type=str,
+        help="the address to use for Ray")
+    parser.add_argument(
+        "--num-replicas",
+        "-n",
+        type=int,
+        default=1,
+        help="Sets number of replicas for training.")
+    parser.add_argument(
+        "--use-gpu",
+        action="store_true",
+        default=False,
+        help="Enables GPU training")
+    parser.add_argument(
+        "--tune", action="store_true", default=False, help="Tune training")
+
+    args, _ = parser.parse_known_args()
+
+    ray.init(address=args.address)
+    tune_example(num_replicas=args.num_replicas, use_gpu=args.use_gpu)

python/ray/util/sgd/pytorch/pytorch_runner.py

@@ -0,0 +1,301 @@
+import collections
+from filelock import FileLock
+import logging
+import inspect
+import os
+import torch
+import torch.utils.data
+from torch.utils.data import Dataset
+
+import ray
+from ray.util.sgd.pytorch import utils as pytorch_utils
+from ray.util.sgd import utils
+
+logger = logging.getLogger(__name__)
+amp = None
+
+try:
+    from apex import amp
+except ImportError:
+    logger.debug("apex is not installed.")
+    pass
+
+
+class PyTorchRunner:
+    """Manages a PyTorch model for training.
+
+    Args:
+        model_creator (dict -> *): see pytorch_trainer.py
+        data_creator (dict -> Dataset, Dataset): see pytorch_trainer.py.
+        optimizer_creator (models, dict -> optimizers): see pytorch_trainer.py.
+        loss_creator (dict -> loss | Loss class): see pytorch_trainer.py.
+        scheduler_creator (optimizers, dict -> schedulers): see
+            pytorch_trainer.py.
+        train_function: see pytorch_trainer.py
+        validation_function: see pytorch_trainer.py
+        config (dict): see pytorch_trainer.py.
+        dataloader_config (dict): See pytorch_trainer.py.
+        batch_size (int): see pytorch_trainer.py.
+        use_fp16 (bool): see pytorch_trainer.py.
+        apex_args (dict|None): see pytorch_trainer.py.
+        scheduler_step_freq (str): see pytorch_trainer.py.
+    """
+
+    def __init__(self,
+                 model_creator,
+                 data_creator,
+                 optimizer_creator,
+                 loss_creator,
+                 scheduler_creator=None,
+                 train_function=None,
+                 validation_function=None,
+                 config=None,
+                 dataloader_config=None,
+                 batch_size=16,
+                 use_fp16=False,
+                 apex_args=None,
+                 scheduler_step_freq="batch"):
+        self.model_creator = model_creator
+        self.data_creator = data_creator
+        self.optimizer_creator = optimizer_creator
+        self.loss_creator = loss_creator
+        self.scheduler_creator = scheduler_creator
+        self.config = {} if config is None else config
+        self.dataloader_config = {
+            "num_workers": 2
+        } if dataloader_config is None else dataloader_config
+        self.train_function = train_function or pytorch_utils.train
+        self.validation_function = (validation_function
+                                    or pytorch_utils.validate)
+        self.batch_size = batch_size
+        self.verbose = True
+
+        self.epoch = 0
+        self._timers = {
+            k: utils.TimerStat(window_size=1)
+            for k in [
+                "setup_proc", "setup_model", "get_state", "set_state",
+                "validation", "training"
+            ]
+        }
+        self.models = None
+        self.optimizers = None
+        self.criterion = None
+        self.schedulers = None
+        self.train_loader = None
+        self.validation_loader = None
+        self.use_fp16 = use_fp16
+        self.apex_args = apex_args or {}
+        if use_fp16 and not amp:
+            raise ImportError(
+                "Please install apex from "
+                "https://www.github.com/nvidia/apex to use fp16 training.")
+        self.scheduler_step_freq = scheduler_step_freq
+
+    def _validate_datasets(self, dataset):
+        assert dataset, "Datasets need to be returned in data_creator."
+        if issubclass(type(dataset), Dataset):
+            return dataset, None
+        elif len(dataset) == 2 and issubclass(type(dataset[0]), Dataset):
+            return dataset
+        else:
+            raise ValueError("Datasets must be <= 2. Got {}".format(dataset))
+
+    def _create_loss(self):
+        if inspect.isclass(self.loss_creator) and issubclass(
+                self.loss_creator, torch.nn.modules.loss._Loss):
+            self.criterion = self.loss_creator()
+        else:
+            self.criterion = self.loss_creator(self.config)
+
+        if torch.cuda.is_available():
+            self.criterion = self.criterion.cuda()
+
+    def _create_schedulers_if_available(self):
+        # Learning rate schedules are optional.
+        if not self.scheduler_creator:
+            return
+        self.schedulers = self.scheduler_creator(self.given_optimizers,
+                                                 self.config)
+
+        if not isinstance(self.schedulers, collections.Iterable):
+            self.schedulers = [self.schedulers]
+
+    def _try_setup_apex(self):
+        """Sets up the model for fp16 training via apex if available."""
+        if self.use_fp16 and amp:
+            self.models, self.optimizers = amp.initialize(
+                self.models, self.optimizers, **self.apex_args)
+
+    def setup(self):
+        """Initializes the model."""
+        logger.debug("Creating model")
+        self.models = self.model_creator(self.config)
+        if not isinstance(self.models, collections.Iterable):
+            self.models = [self.models]
+        if torch.cuda.is_available():
+            self.models = [model.cuda() for model in self.models]
+
+        logger.debug("Creating optimizer")
+        self.optimizers = self.optimizer_creator(self.given_models,
+                                                 self.config)
+        if not isinstance(self.optimizers, collections.Iterable):
+            self.optimizers = [self.optimizers]
+        self._create_schedulers_if_available()
+        self._try_setup_apex()
+        self._create_loss()
+
+        logger.debug("Creating dataset")
+        # When creating datasets, a filelock will be used to ensure no
+        # race conditions in data downloading among different workers.
+        with FileLock(os.path.expanduser("~/.ray_data.lock")):
+            datasets = self.data_creator(self.config)
+            train_set, val_set = self._validate_datasets(datasets)
+
+        self.train_loader = torch.utils.data.DataLoader(
+            train_set, batch_size=self.batch_size, **self.dataloader_config)
+
+        self.validation_loader = None
+        if val_set:
+            self.validation_loader = torch.utils.data.DataLoader(
+                val_set, batch_size=self.batch_size, **self.dataloader_config)
+
+    def get_node_ip(self):
+        """Returns the IP address of the current node."""
+        return ray.services.get_node_ip_address()
+
+    def find_free_port(self):
+        """Finds a free port on the current node."""
+        return utils.find_free_port()
+
+    def step(self):
+        """Runs a training epoch and updates the model parameters."""
+        logger.debug("Begin Training Epoch {}".format(self.epoch + 1))
+        training_config = self.config.copy()
+        training_config.update({
+            pytorch_utils.USE_FP16: self.use_fp16,
+            pytorch_utils.SCHEDULER_STEP: self.scheduler_step_freq
+        })
+        with self._timers["training"]:
+            train_stats = self.train_function(
+                training_config,
+                self.given_models,
+                self.train_loader,
+                self.criterion,
+                self.given_optimizers,
+                scheduler=self.given_schedulers)
+            train_stats["epoch"] = self.epoch
+
+        self.epoch += 1
+
+        train_stats.update(self.stats())
+        return train_stats
+
+    def validate(self):
+        """Evaluates the model on the validation data set."""
+        if self.validation_loader is None:
+            raise ValueError("No validation dataloader provided.")
+        with self._timers["validation"]:
+            validation_stats = self.validation_function(
+                self.config,
+                self.given_models,
+                self.validation_loader,
+                self.criterion,
+                scheduler=self.given_schedulers)
+
+        validation_stats.update(self.stats())
+        return validation_stats
+
+    def stats(self):
+        """Returns a dictionary of statistics collected."""
+        stats = {"epoch": self.epoch}
+        for k, t in self._timers.items():
+            stats[k + "_time_mean"] = t.mean
+            stats[k + "_time_total"] = t.sum
+            t.reset()
+        return stats
+
+    def _get_model_state_dicts(self):
+        # This is so that we create a duplicate of weights into CPU rather than
+        # move the model weights entirely out of the GPU, so that we can
+        # resume training while saving intermediate checkpoints.
+        cpu_state_dicts = []
+        for model in self.models:
+            state_dict = model.state_dict()
+            cpu_state_dicts += [{k: v.cpu() for k, v in state_dict.items()}]
+        return cpu_state_dicts
+
+    def _set_model_state_dicts(self, models_state_dicts):
+        for model, state_dict in zip(self.models, models_state_dicts):
+            model.load_state_dict(state_dict)
+
+    def get_state(self):
+        """Returns the state of the runner."""
+
+        state = {
+            "epoch": self.epoch,
+            "models": self._get_model_state_dicts(),
+            "optimizers": [opt.state_dict() for opt in self.optimizers],
+            "stats": self.stats()
+        }
+        if self.schedulers:
+            state.update({
+                "schedulers": [
+                    scheduler.state_dict() for scheduler in self.schedulers
+                ]
+            })
+        # Check if fp16 is True and if NVIDIA Apex is imported.
+        if self.use_fp16 and amp:
+            state.update({"amp": amp.state_dict()})
+        return state
+
+    def set_state(self, state):
+        """Sets the state of the model."""
+        # TODO: restore timer stats
+        self._set_model_state_dicts(state["models"])
+        for optimizer, state_dict in zip(self.optimizers, state["optimizers"]):
+            optimizer.load_state_dict(state_dict)
+        if self.schedulers:
+            for scheduler, state_dict in zip(self.schedulers,
+                                             state["schedulers"]):
+                scheduler.load_state_dict(state_dict)
+
+        if self.use_fp16 and "amp" in state and amp:
+            amp.load_state_dict(state["amp"])
+        self.epoch = state["stats"]["epoch"]
+
+    def apply_fn(self, fn):
+        return fn(self)
+
+    def shutdown(self):
+        """Attempts to shut down the worker."""
+        del self.validation_loader
+        del self.train_loader
+        del self.criterion
+        del self.optimizers
+        del self.models
+        if torch.cuda.is_available():
+            torch.cuda.empty_cache()
+
+    @property
+    def given_models(self):
+        if len(self.models) > 1:
+            return self.models
+        else:
+            return self.models[0]
+
+    @property
+    def given_optimizers(self):
+        if len(self.optimizers) > 1:
+            return self.optimizers
+        else:
+            return self.optimizers[0]
+
+    @property
+    def given_schedulers(self):
+        if not self.schedulers:
+            return self.schedulers
+        if len(self.schedulers) > 1:
+            return self.schedulers
+        else:
+            return self.schedulers[0]

python/ray/util/sgd/pytorch/pytorch_trainer.py

@@ -0,0 +1,464 @@
+import numpy as np
+import os
+import logging
+import numbers
+import tempfile
+import time
+import torch
+import torch.distributed as dist
+
+import ray
+
+from ray.tune import Trainable
+from ray.tune.trial import Resources
+from ray.util.sgd.pytorch.distributed_pytorch_runner import (
+    DistributedPyTorchRunner)
+from ray.util.sgd import utils
+from ray.util.sgd.pytorch.pytorch_runner import PyTorchRunner
+from ray.util.sgd.pytorch import utils as pytorch_utils
+
+logger = logging.getLogger(__name__)
+RESIZE_COOLDOWN_S = 10
+
+
+class PyTorchTrainer:
+    """Train a PyTorch model using distributed PyTorch.
+
+    Launches a set of actors which connect via distributed PyTorch and
+    coordinate gradient updates to train the provided model.
+
+    .. code-block:: python
+
+        def model_creator(config):
+            return nn.Linear(1, 1)
+
+
+        def optimizer_creator(model, config):
+            return torch.optim.SGD(
+                model.parameters(), lr=config.get("lr", 1e-4))
+
+
+        def data_creator(config):
+            return LinearDataset(2, 5), LinearDataset(2, 5, size=400)
+
+        trainer = PyTorchTrainer(
+            model_creator,
+            data_creator,
+            optimizer_creator,
+            loss_creator=nn.MSELoss,
+            use_gpu=True
+        )
+        trainer.train()
+
+
+    Args:
+        model_creator (dict -> Model(s)): Constructor function that takes in
+            config and returns the model(s) to be optimized. These must be
+            ``torch.nn.Module`` objects. If multiple models are returned,
+            a ``train_function`` must be specified. You do not need to
+            handle GPU/devices in this function; RaySGD will do that under
+            the hood.
+        data_creator (dict -> Dataset(s)): Constructor function
+            that takes in the passed config and returns one or
+            two ``torch.utils.data.Dataset`` objects.
+            Note that even though two Dataset objects can be returned,
+            only one dataset will be used for training. RaySGD
+            will automatically wrap the objects with a ``DataLoader``.
+        optimizer_creator ((models, dict) -> optimizers): Constructor
+            function that takes in the return values from
+            ``model_creator`` and the passed config and returns One or
+            more Torch optimizer objects. You do not need to handle
+            GPU/devices in this function; ``RaySGD`` will do that for you.
+        loss_creator (torch.nn.*Loss class | dict -> loss): A constructor
+            function for the training loss. This can be either a function that
+            takes in the provided config for customization or a subclass
+            of ``torch.nn.modules.loss._Loss``, which is most Pytorch
+            loss classes. For example, ``loss_creator=torch.nn.BCELoss``.
+        scheduler_creator (optimizers, dict -> loss):
+            A constructor function for the scheduler loss. This is
+            a function that takes in the generated optimizers (from
+            ``optimizer_creator``) provided config for customization.
+            Be sure to set ``scheduler_step_freq`` to increment the
+            scheduler correctly.
+        train_function: Custom function for training. This function
+            will be executed in parallel across all workers at once. The
+            function needs to take in (models, train_dataloader, criterion,
+            optimizers, config), and return a dict of training stats.
+        validation_function: Custom function for validation. This function
+            will be executed in parallel across all workers at once.
+            This takes in (model, val_dataloader, criterion, config)
+            and returns a dict of validation stats.
+        config (dict): Custom configuration value to be passed to
+            "model_creator", "data_creator", "optimizer_creator", and
+            "loss_creator".
+        dataloader_config (dict): Configuration values to be passed into
+            the ``torch.utils.data.DataLoader`` object that wraps
+            the dataset on each parallel worker for both training
+            and validation. Note that if ``num_replicas``
+            is greater than 1, ``shuffle`` and ``sampler`` will be
+            automatically set. See the available arguments
+            here https://pytorch.org/docs/stable/data.html.
+        num_replicas (int): the number of workers used in distributed
+            training.
+        use_gpu (bool): Sets resource allocation for workers to 1 GPU
+            if true, and automatically moves both the model and optimizer
+            to the available CUDA device.
+        batch_size (int): Total batch size for each minibatch. This
+            value is divided among all workers and rounded.
+        backend (string): backend used by distributed PyTorch. Currently
+            support "nccl", "gloo", and "auto". If "auto", RaySGD will
+            automatically use "nccl" if `use_gpu` is True, and "gloo"
+            otherwise.
+        use_fp16 (bool): Enables mixed precision training via apex if apex
+            is installed. This is automatically done after the model and
+            optimizers are constructed and will work for multi-model training.
+            Please see https://github.com/NVIDIA/apex for more details.
+        apex_args (dict|None): Dict containing keyword args for amp.initialize.
+            See https://nvidia.github.io/apex/amp.html#module-apex.amp. By
+            default, the models and optimizers are passed in. Consider using
+            "num_losses" if operating over multiple models and optimizers.
+        scheduler_step_freq: "batch", "epoch", or None. This will
+            determine when ``scheduler.step`` is called. If "batch",
+            ``step`` will be called after every optimizer step. If "epoch",
+            ``step`` will be called after one pass of the DataLoader.
+
+    """
+
+    def __init__(self,
+                 model_creator,
+                 data_creator,
+                 optimizer_creator,
+                 loss_creator,
+                 scheduler_creator=None,
+                 train_function=None,
+                 validation_function=None,
+                 initialization_hook=None,
+                 config=None,
+                 dataloader_config=None,
+                 num_replicas=1,
+                 use_gpu=False,
+                 batch_size=16,
+                 backend="auto",
+                 use_fp16=False,
+                 apex_args=None,
+                 scheduler_step_freq="batch"):
+        if num_replicas > 1 and not dist.is_available():
+            raise ValueError(
+                ("Distributed PyTorch is not supported on macOS. "
+                 "To run without distributed PyTorch, set 'num_replicas=1'. "
+                 "For more information, see "
+                 "https://github.com/pytorch/examples/issues/467."))
+
+        self.model_creator = model_creator
+        self.data_creator = data_creator
+        self.train_function = train_function
+        self.optimizer_creator = optimizer_creator
+        self.loss_creator = loss_creator
+        self.scheduler_creator = scheduler_creator
+        self.validation_function = validation_function
+        self.initialization_hook = initialization_hook
+        self.config = {} if config is None else config
+        self.dataloader_config = dataloader_config
+        self.optimizer_timer = utils.TimerStat(window_size=1)
+
+        if backend == "auto":
+            backend = "nccl" if use_gpu else "gloo"
+
+        logger.info("Using {} as backend.".format(backend))
+        self.backend = backend
+        self.use_gpu = use_gpu
+        self.batch_size = batch_size
+        self.max_replicas = num_replicas
+
+        self.use_fp16 = use_fp16
+
+        if apex_args and not isinstance(apex_args, dict):
+            raise ValueError("apex_args needs to be a dict object.")
+
+        self.apex_args = apex_args
+        self.temp_dir = tempfile.mkdtemp(prefix="raysgd")
+        self._num_failures = 0
+        self._last_resize = float("-inf")
+
+        if scheduler_step_freq and (
+                scheduler_step_freq not in pytorch_utils.VALID_SCHEDULER_STEP):
+            raise ValueError(
+                "Scheduler step freq must be in {}. Got {}".format(
+                    pytorch_utils.VALID_SCHEDULER_STEP, scheduler_step_freq))
+
+        self.scheduler_step_freq = scheduler_step_freq
+
+        self._start_workers(self.max_replicas)
+
+    def _start_workers(self, num_replicas):
+        logger.info(f"start_workers: Setting %d replicas." % num_replicas)
+        if num_replicas == 1:
+            # Generate actor class
+            Runner = ray.remote(
+                num_cpus=1, num_gpus=int(self.use_gpu))(PyTorchRunner)
+            # Start workers
+            self.workers = [
+                Runner.remote(
+                    self.model_creator,
+                    self.data_creator,
+                    self.optimizer_creator,
+                    self.loss_creator,
+                    self.scheduler_creator,
+                    train_function=self.train_function,
+                    validation_function=self.validation_function,
+                    config=self.config,
+                    dataloader_config=self.dataloader_config,
+                    batch_size=self.batch_size,
+                    use_fp16=self.use_fp16,
+                    apex_args=self.apex_args,
+                    scheduler_step_freq=self.scheduler_step_freq,
+                )
+            ]
+            if self.initialization_hook:
+                self.apply_all_workers(self.initialization_hook)
+            # Get setup tasks in order to throw errors on failure
+            ray.get(self.workers[0].setup.remote())
+        else:
+            # Generate actor class
+            Runner = ray.remote(
+                num_cpus=1,
+                num_gpus=int(self.use_gpu))(DistributedPyTorchRunner)
+            # Compute batch size per replica
+            batch_size_per_replica = self.batch_size // num_replicas
+            if self.batch_size % num_replicas > 0:
+                new_batch_size = batch_size_per_replica * num_replicas
+                logger.warning(
+                    ("Changing batch size from {old_batch_size} to "
+                     "{new_batch_size} to evenly distribute batches across "
+                     "{num_replicas} replicas.").format(
+                         old_batch_size=self.batch_size,
+                         new_batch_size=new_batch_size,
+                         num_replicas=num_replicas))
+            # Start workers
+            self.workers = [
+                Runner.remote(
+                    self.model_creator,
+                    self.data_creator,
+                    self.optimizer_creator,
+                    self.loss_creator,
+                    self.scheduler_creator,
+                    backend=self.backend,
+                    train_function=self.train_function,
+                    validation_function=self.validation_function,
+                    config=self.config,
+                    dataloader_config=self.dataloader_config,
+                    batch_size=batch_size_per_replica,
+                    use_fp16=self.use_fp16,
+                    apex_args=self.apex_args,
+                    scheduler_step_freq=self.scheduler_step_freq)
+                for i in range(num_replicas)
+            ]
+            if self.initialization_hook:
+                self.apply_all_workers(self.initialization_hook)
+
+            # Compute URL for initializing distributed PyTorch
+            ip = ray.get(self.workers[0].get_node_ip.remote())
+            port = ray.get(self.workers[0].find_free_port.remote())
+            address = "tcp://{ip}:{port}".format(ip=ip, port=port)
+            # Get setup tasks in order to throw errors on failure
+            ray.get([
+                worker.setup.remote(address, i, len(self.workers))
+                for i, worker in enumerate(self.workers)
+            ])
+
+    def train(self, max_retries=0, checkpoint="auto"):
+        """Runs a training epoch.
+
+        Runs an average over all values returned from workers. Set
+        `max_retries` to enable fault handling in case of instance preemption.
+
+        Args:
+            max_retries (int): Must be non-negative. If set to N, will
+                kill all current workers, query the Ray global state for
+                total available resources, and re-launch up to the
+                available resources. Behavior is not well-defined
+                in case of shared cluster usage.
+            checkpoint (str): Path to checkpoint to restore from if retrying.
+                If max_retries is set and checkpoint == "auto", PyTorchTrainer
+                will save a checkpoint before starting to train.
+        """
+        assert max_retries >= 0, "`max_retries` must be non-negative."
+        if max_retries:
+            if checkpoint == "auto":
+                logger.debug("Retrying detected. Automatically checkpointing.")
+                checkpoint = self.save(
+                    os.path.join(self.temp_dir, "tmp_checkpoint"))
+            elif not checkpoint:
+                raise ValueError("Cannot retry from empty checkpoint.")
+
+        if checkpoint and self._should_resize():
+            logger.info("Resize opportunity detected. Attempting to scale up.")
+            self._resize_workers(checkpoint=checkpoint)
+
+        with self.optimizer_timer:
+            success, worker_stats = self._train_step()
+            # Fault handling
+            for i in range(max_retries):
+                if success:
+                    break
+                else:
+                    self._num_failures += 1
+                self._resize_workers(checkpoint=checkpoint)
+                logger.info("Retrying training step with %d workers." % len(
+                    self.workers))
+                success, worker_stats = self._train_step()
+        if not success:
+            raise RuntimeError("Training run failed.")
+
+        worker_stats = ray.get(worker_stats)
+
+        train_stats = {}
+        for stat_key in worker_stats[0]:
+            if isinstance(worker_stats[0], numbers.Number):
+                train_stats[stat_key] = np.nanmean(
+                    [s.get(stat_key, np.nan) for s in worker_stats])
+            else:
+                train_stats[stat_key] = worker_stats[0][stat_key]
+        return train_stats
+
+    def _train_step(self):
+        worker_stats = [w.step.remote() for w in self.workers]
+        success = utils.check_for_failure(worker_stats)
+        return success, worker_stats
+
+    def apply_all_workers(self, fn):
+        return ray.get([w.apply_fn.remote(fn) for w in self.workers])
+
+    def validate(self):
+        """Evaluates the model on the validation data set."""
+        if self.validation_function is False:
+            return {}
+        worker_stats = ray.get([w.validate.remote() for w in self.workers])
+
+        validation_stats = {}
+        for stat_key in worker_stats[0]:
+            validation_stats[stat_key] = np.nanmean(
+                [s.get(stat_key, np.nan) for s in worker_stats])
+        return validation_stats
+
+    def update_scheduler(self, metric):
+        """Calls ``scheduler.step(metric)`` on all schedulers.
+
+        This is useful for lr_schedulers such as ``ReduceLROnPlateau``.
+        """
+        self.apply_all_workers(
+            lambda runner: [sched.step(metric) for sched in runner.schedulers])
+
+    def get_model(self):
+        """Returns the learned model(s)."""
+        models = self.model_creator(self.config)
+        state = ray.get(self.workers[0].get_state.remote())
+        if len(state["models"]) == 1:
+            models.load_state_dict(state["models"][0])
+        else:
+            for model, state_dict in zip(models, state["models"]):
+                model.load_state_dict(state_dict)
+        return models
+
+    def save(self, checkpoint):
+        """Saves the model(s) to the provided checkpoint.
+
+        Args:
+            checkpoint (str): Path to target checkpoint file.
+
+        """
+        state = ray.get(self.workers[0].get_state.remote())
+        torch.save(state, checkpoint)
+        return checkpoint
+
+    def restore(self, checkpoint):
+        """Restores the model from the provided checkpoint.
+
+        Args:
+            checkpoint (str): Path to target checkpoint file.
+
+        """
+        state = torch.load(checkpoint)
+        state_id = ray.put(state)
+        ray.get([worker.set_state.remote(state_id) for worker in self.workers])
+
+    def shutdown(self, force=False):
+        """Shuts down workers and releases resources."""
+        if not force:
+            cleanup = [worker.shutdown.remote() for worker in self.workers]
+            ray.get(cleanup)
+            [worker.__ray_terminate__.remote() for worker in self.workers]
+        else:
+            for worker in self.workers:
+                logger.warning("Killing worker {}.".format(worker))
+                worker.__ray_kill__()
+
+        self.workers = []
+
+    def _resize_workers(self, checkpoint, max_retries=10):
+        # check available resources
+        self.shutdown(force=True)
+        assert checkpoint, "Cannot restore without checkpoint."
+
+        time.sleep(1)
+        for i in range(max_retries):
+            resources = ray.available_resources()
+            new_workers = min(resources.get("CPU", 0), self.max_replicas)
+            if self.use_gpu:
+                new_workers = min(resources.get("GPU", 0), new_workers)
+            if new_workers:
+                self._last_resize = time.time()
+                self._start_workers(int(new_workers))
+                self.restore(checkpoint)
+                return
+            else:
+                delay = 2**i
+                logger.info("Resources: {}".format(resources))
+                logger.warning(
+                    "No new workers found. Retrying in %d sec." % delay)
+                time.sleep(delay)
+        raise RuntimeError("Exceeded max_retries for relaunching workers.")
+
+    def _should_resize(self):
+        """Returns True if past cooldown and exists resources to scale up."""
+        worker_gap = self.max_replicas - len(self.workers)
+        past_cooldown = (time.time() - self._last_resize) > RESIZE_COOLDOWN_S
+        if past_cooldown and worker_gap:
+            resources = ray.available_resources()
+            potential_workers = min(resources.get("CPU", 0), self.max_replicas)
+            if self.use_gpu:
+                potential_workers = min(
+                    resources.get("GPU", 0), potential_workers)
+            return potential_workers > 0
+        return False
+
+
+class PyTorchTrainable(Trainable):
+    @classmethod
+    def default_resource_request(cls, config):
+        return Resources(
+            cpu=0,
+            gpu=0,
+            extra_cpu=config["num_replicas"],
+            extra_gpu=int(config["use_gpu"]) * config["num_replicas"])
+
+    def _setup(self, config):
+        self._trainer = PyTorchTrainer(**config)
+
+    def _train(self):
+        train_stats = self._trainer.train()
+        validation_stats = self._trainer.validate()
+
+        train_stats.update(validation_stats)
+
+        # output {"mean_loss": test_loss, "mean_accuracy": accuracy}
+        return train_stats
+
+    def _save(self, checkpoint_dir):
+        return self._trainer.save(os.path.join(checkpoint_dir, "model.pth"))
+
+    def _restore(self, checkpoint_path):
+        return self._trainer.restore(checkpoint_path)
+
+    def _stop(self):
+        self._trainer.shutdown()

python/ray/util/sgd/pytorch/resnet.py

@@ -0,0 +1,134 @@
+"""ResNet in PyTorch.
+
+Copied from https://github.com/kuangliu/pytorch-cifar/
+            blob/ab908327d44bf9b1d22cd333a4466e85083d3f21/models/resnet.py
+"""
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+class BasicBlock(nn.Module):
+    expansion = 1
+
+    def __init__(self, in_planes, planes, stride=1):
+        super(BasicBlock, self).__init__()
+        self.conv1 = nn.Conv2d(
+            in_planes,
+            planes,
+            kernel_size=3,
+            stride=stride,
+            padding=1,
+            bias=False)
+        self.bn1 = nn.BatchNorm2d(planes)
+        self.conv2 = nn.Conv2d(
+            planes, planes, kernel_size=3, stride=1, padding=1, bias=False)
+        self.bn2 = nn.BatchNorm2d(planes)
+
+        self.shortcut = nn.Sequential()
+        if stride != 1 or in_planes != self.expansion * planes:
+            self.shortcut = nn.Sequential(
+                nn.Conv2d(
+                    in_planes,
+                    self.expansion * planes,
+                    kernel_size=1,
+                    stride=stride,
+                    bias=False), nn.BatchNorm2d(self.expansion * planes))
+
+    def forward(self, x):
+        out = F.relu(self.bn1(self.conv1(x)))
+        out = self.bn2(self.conv2(out))
+        out += self.shortcut(x)
+        out = F.relu(out)
+        return out
+
+
+class Bottleneck(nn.Module):
+    expansion = 4
+
+    def __init__(self, in_planes, planes, stride=1):
+        super(Bottleneck, self).__init__()
+        self.conv1 = nn.Conv2d(in_planes, planes, kernel_size=1, bias=False)
+        self.bn1 = nn.BatchNorm2d(planes)
+        self.conv2 = nn.Conv2d(
+            planes,
+            planes,
+            kernel_size=3,
+            stride=stride,
+            padding=1,
+            bias=False)
+        self.bn2 = nn.BatchNorm2d(planes)
+        self.conv3 = nn.Conv2d(
+            planes, self.expansion * planes, kernel_size=1, bias=False)
+        self.bn3 = nn.BatchNorm2d(self.expansion * planes)
+
+        self.shortcut = nn.Sequential()
+        if stride != 1 or in_planes != self.expansion * planes:
+            self.shortcut = nn.Sequential(
+                nn.Conv2d(
+                    in_planes,
+                    self.expansion * planes,
+                    kernel_size=1,
+                    stride=stride,
+                    bias=False), nn.BatchNorm2d(self.expansion * planes))
+
+    def forward(self, x):
+        out = F.relu(self.bn1(self.conv1(x)))
+        out = F.relu(self.bn2(self.conv2(out)))
+        out = self.bn3(self.conv3(out))
+        out += self.shortcut(x)
+        out = F.relu(out)
+        return out
+
+
+class ResNet(nn.Module):
+    def __init__(self, block, num_blocks, num_classes=10):
+        super(ResNet, self).__init__()
+        self.in_planes = 64
+
+        self.conv1 = nn.Conv2d(
+            3, 64, kernel_size=3, stride=1, padding=1, bias=False)
+        self.bn1 = nn.BatchNorm2d(64)
+        self.layer1 = self._make_layer(block, 64, num_blocks[0], stride=1)
+        self.layer2 = self._make_layer(block, 128, num_blocks[1], stride=2)
+        self.layer3 = self._make_layer(block, 256, num_blocks[2], stride=2)
+        self.layer4 = self._make_layer(block, 512, num_blocks[3], stride=2)
+        self.linear = nn.Linear(512 * block.expansion, num_classes)
+
+    def _make_layer(self, block, planes, num_blocks, stride):
+        strides = [stride] + [1] * (num_blocks - 1)
+        layers = []
+        for stride in strides:
+            layers.append(block(self.in_planes, planes, stride))
+            self.in_planes = planes * block.expansion
+        return nn.Sequential(*layers)
+
+    def forward(self, x):
+        out = F.relu(self.bn1(self.conv1(x)))
+        out = self.layer1(out)
+        out = self.layer2(out)
+        out = self.layer3(out)
+        out = self.layer4(out)
+        out = F.avg_pool2d(out, 4)
+        out = out.view(out.size(0), -1)
+        out = self.linear(out)
+        return out
+
+
+def ResNet18(_):
+    return ResNet(BasicBlock, [2, 2, 2, 2])
+
+
+def ResNet34(_):
+    return ResNet(BasicBlock, [3, 4, 6, 3])
+
+
+def ResNet50(_):
+    return ResNet(Bottleneck, [3, 4, 6, 3])
+
+
+def ResNet101(_):
+    return ResNet(Bottleneck, [3, 4, 23, 3])
+
+
+def ResNet152(_):
+    return ResNet(Bottleneck, [3, 8, 36, 3])

python/ray/util/sgd/pytorch/utils.py

@@ -0,0 +1,229 @@
+import collections
+import time
+import torch
+
+from ray.util.sgd.utils import TimerStat
+
+amp = None
+
+try:
+    from apex import amp
+except ImportError:
+    # Apex library is not installed, so we cannot enable mixed precision.
+    # We don't log here because logging happens in the pytorch_runner,
+    # where amp is initialized.
+    pass
+
+USE_FP16 = "__use_fp16__"
+TEST_MODE = "__test_mode__"
+BATCH_COUNT = "batch_processed"
+SCHEDULER_STEP = "scheduler_step"
+SCHEDULER_STEP_BATCH = "batch"
+SCHEDULER_STEP_EPOCH = "epoch"
+
+VALID_SCHEDULER_STEP = {SCHEDULER_STEP_BATCH, SCHEDULER_STEP_EPOCH}
+
+
+def train(config, model, train_iterator, criterion, optimizer, scheduler=None):
+    """Runs one standard training pass over the train_iterator.
+
+    This function automatically measures timing for various operations such
+    as host to device transfer, gradient calculation, and gradient application.
+
+    It also automatically detects and places the data on the given GPU device
+    if available.
+
+    The scheduler will only be called at a batch or epoch frequency, depending
+    on the user parameter. Be sure to set ``scheduler_step_freq`` in
+    ``PyTorchTrainer`` to either "batch" or "epoch" to increment the scheduler
+    correctly during training. If using a learning rate scheduler
+    that depends on validation loss, you can use ``trainer.update_scheduler``.
+
+    Raises:
+        ValueError if multiple models/optimizers/schedulers are provided. You
+            are expected to have a custom training function if you wish
+            to use multiple models/optimizers/schedulers.
+
+    Args:
+        config: (dict): A user configuration provided into the Trainer
+            constructor.
+        model: The model as created by the model_creator.
+        train_iterator: An iterator created from the DataLoader which
+            wraps the provided Dataset.
+        criterion: The loss object created by the loss_creator.
+        optimizer: The torch.optim.Optimizer object as created by the
+            optimizer_creator.
+        scheduler (optional): The torch.optim.lr_scheduler object
+            as created by the scheduler_creator. Be sure to set
+            ``scheduler_step_freq`` in ``PyTorchTrainer``
+             to increment the scheduler correctly.
+
+    Returns:
+        A dict of metrics from training.
+    """
+    if isinstance(model, collections.Iterable) or isinstance(
+            optimizer, collections.Iterable) or isinstance(
+                scheduler, collections.Iterable):
+        raise ValueError(
+            "Need to provide custom training function if using multi-model "
+            "or multi-scheduler or multi-optimizer training.")
+
+    batch_time = AverageMeter()
+    data_time = AverageMeter()
+    losses = AverageMeter()
+
+    timers = {k: TimerStat() for k in ["h2d", "fwd", "grad", "apply"]}
+
+    # switch to train mode
+    model.train()
+
+    end = time.time()
+
+    for batch_idx, (features, target) in enumerate(train_iterator):
+        # measure data loading time
+        data_time.update(time.time() - end)
+
+        # Create non_blocking tensors for distributed training
+        with timers["h2d"]:
+            if torch.cuda.is_available():
+                features = features.cuda(non_blocking=True)
+                target = target.cuda(non_blocking=True)
+
+        # compute output
+        with timers["fwd"]:
+            output = model(features)
+            loss = criterion(output, target)
+
+            # measure accuracy and record loss
+            losses.update(loss.item(), features.size(0))
+
+        with timers["grad"]:
+            # compute gradients in a backward pass
+            optimizer.zero_grad()
+
+            if config.get(USE_FP16):
+                with amp.scale_loss(loss, optimizer) as scaled_loss:
+                    scaled_loss.backward()
+            else:
+                loss.backward()
+
+        with timers["apply"]:
+            # Call step of optimizer to update model params
+            optimizer.step()
+
+        if scheduler and config.get(SCHEDULER_STEP) == SCHEDULER_STEP_BATCH:
+            scheduler.step()
+
+        # measure elapsed time
+        batch_time.update(time.time() - end)
+        end = time.time()
+
+        if config.get(TEST_MODE) and batch_idx == 0:
+            break
+
+    if scheduler and config.get(SCHEDULER_STEP) == SCHEDULER_STEP_EPOCH:
+        scheduler.step()
+
+    stats = {
+        "batch_time": batch_time.avg,
+        BATCH_COUNT: batch_idx + 1,
+        "train_loss": losses.avg,
+        "data_time": data_time.avg,
+    }
+    stats.update({k: t.mean for k, t in timers.items()})
+    return stats
+
+
+def validate(config, model, val_iterator, criterion, scheduler=None):
+    """Runs one standard validation pass over the val_iterator.
+
+    This function automatically measures timing for various operations such
+    as host to device transfer and processing time for the batch.
+
+    It also automatically detects and places the data on the given GPU device
+    if available.
+
+    Raises:
+        ValueError if multiple models/schedulers are provided. You
+            are expected to have a custom validation function if you wish
+            to use multiple models/schedulers.
+
+    Args:
+        config: (dict): A user configuration provided into the Trainer
+            constructor.
+        model: The model as created by the model_creator.
+        train_iterator: An iterator created from the DataLoader which
+            wraps the provided Dataset.
+        criterion: The loss object created by the loss_creator.
+        scheduler (optional): The torch.optim.lr_scheduler object
+            as created by the scheduler_creator. By default,
+            this is not used in this function.
+
+    Returns:
+        A dict of metrics from the evaluation.
+    """
+
+    if isinstance(model, collections.Iterable) or isinstance(
+            scheduler, collections.Iterable):
+        raise ValueError(
+            "Need to provide custom validation function if using multi-model "
+            "or multi-scheduler training.")
+    batch_time = AverageMeter()
+    losses = AverageMeter()
+
+    # switch to evaluate mode
+    model.eval()
+    correct = 0
+    total = 0
+    batch_idx = 0
+    with torch.no_grad():
+        end = time.time()
+        for batch_idx, (features, target) in enumerate(val_iterator):
+            if torch.cuda.is_available():
+                features = features.cuda(non_blocking=True)
+                target = target.cuda(non_blocking=True)
+
+            # compute output
+            output = model(features)
+            loss = criterion(output, target)
+            _, predicted = torch.max(output.data, 1)
+            total += target.size(0)
+            correct += (predicted == target).sum().item()
+
+            # measure accuracy and record loss
+            losses.update(loss.item(), features.size(0))
+
+            # measure elapsed time
+            batch_time.update(time.time() - end)
+            end = time.time()
+
+            if config.get(TEST_MODE) and batch_idx == 0:
+                break
+
+    stats = {
+        BATCH_COUNT: batch_idx + 1,
+        "batch_time": batch_time.avg,
+        "validation_loss": losses.avg,
+        "mean_accuracy": correct / total,
+        "mean_loss": losses.sum / total,
+    }
+    return stats
+
+
+class AverageMeter:
+    """Computes and stores the average and current value."""
+
+    def __init__(self):
+        self.reset()
+
+    def reset(self):
+        self.val = 0
+        self.avg = 0
+        self.sum = 0
+        self.count = 0
+
+    def update(self, val, n=1):
+        self.val = val
+        self.sum += val * n
+        self.count += n
+        self.avg = self.sum / self.count

python/ray/util/sgd/tests/test_pytorch.py

@@ -0,0 +1,385 @@
+import os
+import tempfile
+from unittest.mock import patch
+
+import pytest
+import time
+import torch
+import torch.nn as nn
+import torch.distributed as dist
+
+import ray
+from ray import tune
+from ray.tests.conftest import ray_start_2_cpus  # noqa: F401
+from ray.util.sgd.pytorch import PyTorchTrainer, PyTorchTrainable
+from ray.util.sgd.pytorch.utils import (train, BATCH_COUNT, TEST_MODE,
+                                        SCHEDULER_STEP)
+from ray.util.sgd.utils import check_for_failure
+
+from ray.util.sgd.pytorch.examples.train_example import (
+    model_creator, optimizer_creator, data_creator, LinearDataset)
+
+
+def test_test_mode(ray_start_2_cpus):  # noqa: F811
+    trainer = PyTorchTrainer(
+        model_creator,
+        data_creator,
+        optimizer_creator,
+        loss_creator=lambda config: nn.MSELoss(),
+        config={TEST_MODE: True},
+        num_replicas=1)
+    metrics = trainer.train()
+    assert metrics[BATCH_COUNT] == 1
+
+    val_metrics = trainer.validate()
+    assert val_metrics[BATCH_COUNT] == 1
+
+
+@pytest.mark.parametrize("num_replicas", [1, 2]
+                         if dist.is_available() else [1])
+def test_train(ray_start_2_cpus, num_replicas):  # noqa: F811
+    trainer = PyTorchTrainer(
+        model_creator,
+        data_creator,
+        optimizer_creator,
+        loss_creator=lambda config: nn.MSELoss(),
+        num_replicas=num_replicas)
+    for i in range(3):
+        train_loss1 = trainer.train()["train_loss"]
+    validation_loss1 = trainer.validate()["validation_loss"]
+
+    for i in range(3):
+        train_loss2 = trainer.train()["train_loss"]
+    validation_loss2 = trainer.validate()["validation_loss"]
+
+    print(train_loss1, train_loss2)
+    print(validation_loss1, validation_loss2)
+
+    assert train_loss2 <= train_loss1
+    assert validation_loss2 <= validation_loss1
+
+
+@pytest.mark.parametrize("num_replicas", [1, 2]
+                         if dist.is_available() else [1])
+def test_multi_model(ray_start_2_cpus, num_replicas):  # noqa: F811
+    def custom_train(config, models, dataloader, criterion, optimizers,
+                     **kwargs):
+        result = {}
+        for i, (model, optimizer) in enumerate(zip(models, optimizers)):
+            result["model_{}".format(i)] = train(config, model, dataloader,
+                                                 criterion, optimizer)
+        return result
+
+    def multi_model_creator(config):
+        return nn.Linear(1, 1), nn.Linear(1, 1)
+
+    def multi_optimizer_creator(models, config):
+        opts = [
+            torch.optim.SGD(model.parameters(), lr=0.0001) for model in models
+        ]
+        return opts[0], opts[1]
+
+    trainer1 = PyTorchTrainer(
+        multi_model_creator,
+        data_creator,
+        multi_optimizer_creator,
+        loss_creator=lambda config: nn.MSELoss(),
+        train_function=custom_train,
+        num_replicas=num_replicas)
+    trainer1.train()
+
+    filename = os.path.join(tempfile.mkdtemp(), "checkpoint")
+    trainer1.save(filename)
+
+    models1 = trainer1.get_model()
+
+    trainer1.shutdown()
+
+    trainer2 = PyTorchTrainer(
+        multi_model_creator,
+        data_creator,
+        multi_optimizer_creator,
+        loss_creator=lambda config: nn.MSELoss(),
+        num_replicas=num_replicas)
+    trainer2.restore(filename)
+
+    os.remove(filename)
+
+    models2 = trainer2.get_model()
+
+    for model_1, model_2 in zip(models1, models2):
+
+        model1_state_dict = model_1.state_dict()
+        model2_state_dict = model_2.state_dict()
+
+        assert set(model1_state_dict.keys()) == set(model2_state_dict.keys())
+
+        for k in model1_state_dict:
+            assert torch.equal(model1_state_dict[k], model2_state_dict[k])
+
+    trainer2.shutdown()
+
+
+@pytest.mark.parametrize("num_replicas", [1, 2]
+                         if dist.is_available() else [1])
+def test_multi_model_matrix(ray_start_2_cpus, num_replicas):  # noqa: F811
+    def custom_train(config, model, dataloader, criterion, optimizer,
+                     scheduler):
+        if config.get("models", 1) > 1:
+            assert len(model) == config["models"], config
+
+        if config.get("optimizers", 1) > 1:
+            assert len(optimizer) == config["optimizers"], config
+
+        if config.get("schedulers", 1) > 1:
+            assert len(scheduler) == config["schedulers"], config
+        return {"done": 1}
+
+    def multi_model_creator(config):
+        models = []
+        for i in range(config.get("models", 1)):
+            models += [nn.Linear(1, 1)]
+        return models[0] if len(models) == 1 else models
+
+    def multi_optimizer_creator(models, config):
+        optimizers = []
+        main_model = models[0] if type(models) is list else models
+        for i in range(config.get("optimizers", 1)):
+            optimizers += [torch.optim.SGD(main_model.parameters(), lr=0.0001)]
+        return optimizers[0] if len(optimizers) == 1 else optimizers
+
+    def multi_scheduler_creator(optimizer, config):
+        schedulers = []
+        main_opt = optimizer[0] if type(optimizer) is list else optimizer
+        for i in range(config.get("schedulers", 1)):
+            schedulers += [
+                torch.optim.lr_scheduler.StepLR(
+                    main_opt, step_size=30, gamma=0.1)
+            ]
+        return schedulers[0] if len(schedulers) == 1 else schedulers
+
+    for model_count in range(1, 3):
+        for optimizer_count in range(1, 3):
+            for scheduler_count in range(1, 3):
+                trainer = PyTorchTrainer(
+                    multi_model_creator,
+                    data_creator,
+                    multi_optimizer_creator,
+                    loss_creator=nn.MSELoss,
+                    scheduler_creator=multi_scheduler_creator,
+                    train_function=custom_train,
+                    num_replicas=num_replicas,
+                    config={
+                        "models": model_count,
+                        "optimizers": optimizer_count,
+                        "schedulers": scheduler_count
+                    })
+                trainer.train()
+                trainer.shutdown()
+
+
+@pytest.mark.parametrize("scheduler_freq", ["epoch", "batch"])
+def test_scheduler_freq(ray_start_2_cpus, scheduler_freq):  # noqa: F811
+    def custom_train(config, model, dataloader, criterion, optimizer,
+                     scheduler):
+        assert config[SCHEDULER_STEP] == scheduler_freq
+        return {"done": 1}
+
+    def scheduler_creator(optimizer, config):
+        return torch.optim.lr_scheduler.StepLR(
+            optimizer, step_size=30, gamma=0.1)
+
+    trainer = PyTorchTrainer(
+        model_creator,
+        data_creator,
+        optimizer_creator,
+        loss_creator=lambda config: nn.MSELoss(),
+        scheduler_creator=scheduler_creator)
+
+    for i in range(3):
+        trainer.train()["train_loss"]
+    trainer.shutdown()
+
+
+def test_scheduler_validate(ray_start_2_cpus):  # noqa: F811
+    def custom_train(config, model, dataloader, criterion, optimizer,
+                     scheduler):
+        return {"done": 1}
+
+    from torch.optim.lr_scheduler import ReduceLROnPlateau
+
+    trainer = PyTorchTrainer(
+        model_creator,
+        data_creator,
+        optimizer_creator,
+        loss_creator=lambda config: nn.MSELoss(),
+        scheduler_creator=lambda optimizer, cfg: ReduceLROnPlateau(optimizer))
+    trainer.update_scheduler(0.5)
+    trainer.update_scheduler(0.5)
+    assert all(
+        trainer.apply_all_workers(lambda r: r.schedulers[0].last_epoch == 2))
+    trainer.shutdown()
+
+
+@pytest.mark.parametrize("num_replicas", [1, 2]
+                         if dist.is_available() else [1])
+def test_tune_train(ray_start_2_cpus, num_replicas):  # noqa: F811
+
+    config = {
+        "model_creator": model_creator,
+        "data_creator": data_creator,
+        "optimizer_creator": optimizer_creator,
+        "loss_creator": lambda config: nn.MSELoss(),
+        "num_replicas": num_replicas,
+        "use_gpu": False,
+        "batch_size": 512,
+        "backend": "gloo",
+        "config": {
+            "lr": 0.001
+        }
+    }
+
+    analysis = tune.run(
+        PyTorchTrainable,
+        num_samples=2,
+        config=config,
+        stop={"training_iteration": 2},
+        verbose=1)
+
+    # checks loss decreasing for every trials
+    for path, df in analysis.trial_dataframes.items():
+        train_loss1 = df.loc[0, "train_loss"]
+        train_loss2 = df.loc[1, "train_loss"]
+        validation_loss1 = df.loc[0, "validation_loss"]
+        validation_loss2 = df.loc[1, "validation_loss"]
+
+        assert train_loss2 <= train_loss1
+        assert validation_loss2 <= validation_loss1
+
+
+@pytest.mark.parametrize("num_replicas", [1, 2]
+                         if dist.is_available() else [1])
+def test_save_and_restore(ray_start_2_cpus, num_replicas):  # noqa: F811
+    trainer1 = PyTorchTrainer(
+        model_creator,
+        data_creator,
+        optimizer_creator,
+        loss_creator=lambda config: nn.MSELoss(),
+        num_replicas=num_replicas)
+    trainer1.train()
+
+    filename = os.path.join(tempfile.mkdtemp(), "checkpoint")
+    trainer1.save(filename)
+
+    model1 = trainer1.get_model()
+
+    trainer1.shutdown()
+
+    trainer2 = PyTorchTrainer(
+        model_creator,
+        data_creator,
+        optimizer_creator,
+        loss_creator=lambda config: nn.MSELoss(),
+        num_replicas=num_replicas)
+    trainer2.restore(filename)
+
+    os.remove(filename)
+
+    model2 = trainer2.get_model()
+
+    model1_state_dict = model1.state_dict()
+    model2_state_dict = model2.state_dict()
+
+    assert set(model1_state_dict.keys()) == set(model2_state_dict.keys())
+
+    for k in model1_state_dict:
+        assert torch.equal(model1_state_dict[k], model2_state_dict[k])
+
+
+def test_fail_with_recover(ray_start_2_cpus):  # noqa: F811
+    if not dist.is_available():
+        return
+
+    def single_loader(config):
+        return LinearDataset(2, 5, size=1000000)
+
+    def step_with_fail(self):
+        worker_stats = [w.step.remote() for w in self.workers]
+        if self._num_failures < 3:
+            time.sleep(1)  # Make the batch will fail correctly.
+            self.workers[0].__ray_kill__()
+        success = check_for_failure(worker_stats)
+        return success, worker_stats
+
+    with patch.object(PyTorchTrainer, "_train_step", step_with_fail):
+        trainer1 = PyTorchTrainer(
+            model_creator,
+            single_loader,
+            optimizer_creator,
+            batch_size=100000,
+            loss_creator=lambda config: nn.MSELoss(),
+            num_replicas=2)
+
+        with pytest.raises(RuntimeError):
+            trainer1.train(max_retries=1)
+
+
+def test_resize(ray_start_2_cpus):  # noqa: F811
+    if not dist.is_available():
+        return
+
+    def single_loader(config):
+        return LinearDataset(2, 5, size=1000000)
+
+    def step_with_fail(self):
+        worker_stats = [w.step.remote() for w in self.workers]
+        if self._num_failures < 1:
+            time.sleep(1)  # Make the batch will fail correctly.
+            self.workers[0].__ray_kill__()
+        success = check_for_failure(worker_stats)
+        return success, worker_stats
+
+    with patch.object(PyTorchTrainer, "_train_step", step_with_fail):
+        trainer1 = PyTorchTrainer(
+            model_creator,
+            single_loader,
+            optimizer_creator,
+            batch_size=100000,
+            loss_creator=lambda config: nn.MSELoss(),
+            num_replicas=2)
+
+        @ray.remote
+        def try_test():
+            import time
+            time.sleep(100)
+
+        try_test.remote()
+        trainer1.train(max_retries=1)
+        assert len(trainer1.workers) == 1
+
+
+def test_fail_twice(ray_start_2_cpus):  # noqa: F811
+    if not dist.is_available():
+        return
+
+    def single_loader(config):
+        return LinearDataset(2, 5, size=1000000)
+
+    def step_with_fail(self):
+        worker_stats = [w.step.remote() for w in self.workers]
+        if self._num_failures < 2:
+            time.sleep(1)
+            self.workers[0].__ray_kill__()
+        success = check_for_failure(worker_stats)
+        return success, worker_stats
+
+    with patch.object(PyTorchTrainer, "_train_step", step_with_fail):
+        trainer1 = PyTorchTrainer(
+            model_creator,
+            single_loader,
+            optimizer_creator,
+            batch_size=100000,
+            loss_creator=lambda config: nn.MSELoss(),
+            num_replicas=2)
+
+        trainer1.train(max_retries=2)

python/ray/util/sgd/tests/test_pytorch_runner.py

@@ -0,0 +1,151 @@
+import numpy as np
+import torch
+import torch.nn as nn
+import unittest
+from unittest.mock import MagicMock
+
+from ray.util.sgd.pytorch.pytorch_runner import PyTorchRunner
+
+
+class LinearDataset(torch.utils.data.Dataset):
+    """y = a * x + b"""
+
+    def __init__(self, a, b, size=1000):
+        x = np.random.random(size).astype(np.float32) * 10
+        x = np.arange(0, 10, 10 / size, dtype=np.float32)
+        self.x = torch.from_numpy(x)
+        self.y = torch.from_numpy(a * x + b)
+
+    def __getitem__(self, index):
+        return self.x[index, None], self.y[index, None]
+
+    def __len__(self):
+        return len(self.x)
+
+
+def model_creator(config):
+    return nn.Linear(1, 1)
+
+
+def optimizer_creator(models, config):
+    """Returns optimizer."""
+    return torch.optim.SGD(models.parameters(), lr=0.1)
+
+
+def loss_creator(config):
+    return nn.MSELoss()
+
+
+def single_loader(config):
+    return LinearDataset(2, 5)
+
+
+def create_dataloaders(config):
+    return LinearDataset(2, 5), LinearDataset(2, 5, size=400)
+
+
+class TestPyTorchRunner(unittest.TestCase):
+    def testValidate(self):
+        mock_function = MagicMock(returns=dict(mean_accuracy=10))
+        runner = PyTorchRunner(
+            model_creator,
+            create_dataloaders,
+            optimizer_creator,
+            loss_creator,
+            validation_function=mock_function)
+        runner.setup()
+        runner.step()
+        runner.step()
+        runner.step()
+        self.assertEqual(mock_function.call_count, 0)
+        runner.validate()
+        self.assertTrue(mock_function.called)
+        self.assertEqual(runner.stats()["epoch"], 3)
+
+    def testStep(self):
+        mock_function = MagicMock(return_value=dict(mean_accuracy=10))
+        runner = PyTorchRunner(
+            model_creator,
+            create_dataloaders,
+            optimizer_creator,
+            loss_creator,
+            train_function=mock_function)
+        runner.setup()
+        runner.step()
+        runner.step()
+        result = runner.step()
+        self.assertEqual(mock_function.call_count, 3)
+        self.assertEqual(result["epoch"], 3)
+        self.assertEqual(runner.stats()["epoch"], 3)
+
+    def testGivens(self):
+        def three_model_creator(config):
+            return nn.Linear(1, 1), nn.Linear(1, 1), nn.Linear(1, 1)
+
+        def three_optimizer_creator(models, config):
+            opts = [
+                torch.optim.SGD(model.parameters(), lr=0.1) for model in models
+            ]
+            return opts[0], opts[1], opts[2]
+
+        runner = PyTorchRunner(three_model_creator, single_loader,
+                               three_optimizer_creator, loss_creator)
+        runner.setup()
+
+        self.assertEqual(len(runner.given_models), 3)
+        self.assertEqual(len(runner.given_optimizers), 3)
+
+        runner2 = PyTorchRunner(model_creator, single_loader,
+                                optimizer_creator, loss_creator)
+        runner2.setup()
+
+        self.assertNotEqual(runner2.given_models, runner2.models)
+        self.assertNotEqual(runner2.given_optimizers, runner2.optimizers)
+
+    def testMultiLoaders(self):
+        def three_data_loader(config):
+            return (LinearDataset(2, 5), LinearDataset(2, 5, size=400),
+                    LinearDataset(2, 5, size=400))
+
+        runner = PyTorchRunner(model_creator, three_data_loader,
+                               optimizer_creator, loss_creator)
+        with self.assertRaises(ValueError):
+            runner.setup()
+
+        runner2 = PyTorchRunner(model_creator, three_data_loader,
+                                optimizer_creator, loss_creator)
+        with self.assertRaises(ValueError):
+            runner2.setup()
+
+    def testSingleLoader(self):
+        runner = PyTorchRunner(model_creator, single_loader, optimizer_creator,
+                               loss_creator)
+        runner.setup()
+        runner.step()
+        with self.assertRaises(ValueError):
+            runner.validate()
+
+    def testNativeLoss(self):
+        runner = PyTorchRunner(
+            model_creator,
+            single_loader,
+            optimizer_creator,
+            loss_creator=nn.MSELoss)
+        runner.setup()
+        runner.step()
+
+    def testMultiModel(self):
+        def multi_model_creator(config):
+            return nn.Linear(1, 1), nn.Linear(1, 1), nn.Linear(1, 1)
+
+        def multi_optimizer_creator(models, config):
+            opts = [
+                torch.optim.SGD(model.parameters(), lr=0.1) for model in models
+            ]
+            return opts[0], opts[1], opts[2]
+
+        runner = PyTorchRunner(multi_model_creator, single_loader,
+                               multi_optimizer_creator, loss_creator)
+        runner.setup()
+        with self.assertRaises(ValueError):
+            runner.step()

python/ray/util/sgd/tests/test_tensorflow.py

@@ -0,0 +1,131 @@
+import os
+import pytest
+import tempfile
+import numpy as np
+import shutil
+
+from ray import tune
+from ray.tests.conftest import ray_start_2_cpus  # noqa: F401
+from ray.util.sgd.tf import TFTrainer, TFTrainable
+
+from ray.util.sgd.tf.examples.tensorflow_train_example import (simple_model,
+                                                               simple_dataset)
+
+SIMPLE_CONFIG = {
+    "batch_size": 128,
+    "fit_config": {
+        "steps_per_epoch": 3,
+    },
+    "evaluate_config": {
+        "steps": 3,
+    }
+}
+
+
+@pytest.mark.parametrize(  # noqa: F811
+    "num_replicas", [1, 2])
+def test_train(ray_start_2_cpus, num_replicas):  # noqa: F811
+    trainer = TFTrainer(
+        model_creator=simple_model,
+        data_creator=simple_dataset,
+        num_replicas=num_replicas,
+        config=SIMPLE_CONFIG)
+
+    train_stats1 = trainer.train()
+    train_stats1.update(trainer.validate())
+
+    train_stats2 = trainer.train()
+    train_stats2.update(trainer.validate())
+
+
+@pytest.mark.parametrize(  # noqa: F811
+    "num_replicas", [1, 2])
+def test_tune_train(ray_start_2_cpus, num_replicas):  # noqa: F811
+
+    config = {
+        "model_creator": tune.function(simple_model),
+        "data_creator": tune.function(simple_dataset),
+        "num_replicas": num_replicas,
+        "use_gpu": False,
+        "trainer_config": SIMPLE_CONFIG
+    }
+
+    tune.run(
+        TFTrainable,
+        num_samples=2,
+        config=config,
+        stop={"training_iteration": 2},
+        verbose=1)
+
+
+@pytest.mark.parametrize(  # noqa: F811
+    "num_replicas", [1, 2])
+def test_save_and_restore(ray_start_2_cpus, num_replicas):  # noqa: F811
+    trainer1 = TFTrainer(
+        model_creator=simple_model,
+        data_creator=simple_dataset,
+        num_replicas=num_replicas,
+        config=SIMPLE_CONFIG)
+    trainer1.train()
+
+    tmpdir = tempfile.mkdtemp()
+    filename = os.path.join(tmpdir, "checkpoint")
+    trainer1.save(filename)
+
+    model1 = trainer1.get_model()
+    trainer1.shutdown()
+
+    trainer2 = TFTrainer(
+        model_creator=simple_model,
+        data_creator=simple_dataset,
+        num_replicas=num_replicas,
+        config=SIMPLE_CONFIG)
+    trainer2.restore(filename)
+
+    model2 = trainer2.get_model()
+    trainer2.shutdown()
+
+    shutil.rmtree(tmpdir)
+
+    model1_config = model1.get_config()
+    model2_config = model2.get_config()
+    assert _compare(model1_config, model2_config, skip_keys=["name"])
+
+    model1_weights = model1.get_weights()
+    model2_weights = model2.get_weights()
+    assert _compare(model1_weights, model2_weights)
+
+    model1_opt_weights = model1.optimizer.get_weights()
+    model2_opt_weights = model2.optimizer.get_weights()
+    assert _compare(model1_opt_weights, model2_opt_weights)
+
+
+def _compare(d1, d2, skip_keys=None):
+    """Compare two lists or dictionaries or array"""
+    if type(d1) != type(d2):
+        return False
+
+    if isinstance(d1, dict):
+        if set(d1) != set(d2):
+            return False
+
+        for key in d1:
+            if skip_keys is not None and key in skip_keys:
+                continue
+
+            if not _compare(d1[key], d2[key], skip_keys=skip_keys):
+                return False
+
+    elif isinstance(d1, list):
+        for i, _ in enumerate(d1):
+            if not _compare(d1[i], d2[i], skip_keys=skip_keys):
+                return False
+
+    elif isinstance(d1, np.ndarray):
+        if not np.array_equal(d1, d2):
+            return False
+    else:
+        if d1 != d2:
+            return False
+
+    return True

python/ray/util/sgd/tf/__init__.py

@@ -0,0 +1,3 @@
+from ray.util.sgd.tf.tf_trainer import (TFTrainer, TFTrainable)
+
+__all__ = ["TFTrainer", "TFTrainable"]

python/ray/util/sgd/tf/examples/cifar_tf_example.py

@@ -0,0 +1,227 @@
+"""
+#Train a simple deep CNN on the CIFAR10 small images dataset.
+
+It gets to 75% validation accuracy in 25 epochs, and 79% after 50 epochs.
+(it"s still underfitting at that point, though).
+"""
+import argparse
+import time
+
+from tensorflow.keras.datasets import cifar10
+from tensorflow.keras.preprocessing.image import ImageDataGenerator
+from tensorflow.keras.models import Sequential
+from tensorflow.keras.layers import Dense, Dropout, Activation, Flatten
+from tensorflow.keras.layers import Conv2D, MaxPooling2D
+import os
+from filelock import FileLock
+
+import ray
+from ray.util.sgd.tf.tf_trainer import TFTrainer
+
+num_classes = 10
+
+
+def fetch_keras_data():
+    import tensorflow as tf
+    # The data, split between train and test sets:
+    with FileLock(os.path.expanduser("~/.cifar.lock")):
+        (x_train, y_train), (x_test, y_test) = cifar10.load_data()
+
+    # Convert class vectors to binary class matrices.
+    y_train = tf.keras.utils.to_categorical(y_train, num_classes)
+    y_test = tf.keras.utils.to_categorical(y_test, num_classes)
+
+    x_train = x_train.astype("float32")
+    x_test = x_test.astype("float32")
+    x_train /= 255
+    x_test /= 255
+    return (x_train, y_train), (x_test, y_test)
+
+
+(x_train, y_train), (x_test, y_test) = fetch_keras_data()
+input_shape = x_train.shape[1:]
+
+
+def create_model(config):
+    import tensorflow as tf
+    model = Sequential()
+    model.add(Conv2D(32, (3, 3), padding="same", input_shape=input_shape))
+    model.add(Activation("relu"))
+    model.add(Conv2D(32, (3, 3)))
+    model.add(Activation("relu"))
+    model.add(MaxPooling2D(pool_size=(2, 2)))
+    model.add(Dropout(0.25))
+
+    model.add(Conv2D(64, (3, 3), padding="same"))
+    model.add(Activation("relu"))
+    model.add(Conv2D(64, (3, 3)))
+    model.add(Activation("relu"))
+    model.add(MaxPooling2D(pool_size=(2, 2)))
+    model.add(Dropout(0.25))
+
+    model.add(Flatten())
+    model.add(Dense(64))
+    model.add(Activation("relu"))
+    model.add(Dropout(0.5))
+    model.add(Dense(num_classes))
+    model.add(Activation("softmax"))
+
+    # initiate RMSprop optimizer
+    opt = tf.keras.optimizers.RMSprop(lr=0.001, decay=1e-6)
+
+    # Let"s train the model using RMSprop
+    model.compile(
+        loss="categorical_crossentropy", optimizer=opt, metrics=["accuracy"])
+    return model
+
+
+def data_creator(config):
+    import tensorflow as tf
+    batch_size = config["batch_size"]
+    (x_train, y_train), (x_test, y_test) = fetch_keras_data()
+    train_dataset = tf.data.Dataset.from_tensor_slices((x_train, y_train))
+    test_dataset = tf.data.Dataset.from_tensor_slices((x_test, y_test))
+
+    # Repeat is needed to avoid
+    train_dataset = train_dataset.repeat().shuffle(
+        len(x_train)).batch(batch_size)
+    test_dataset = test_dataset.repeat().batch(batch_size)
+    return train_dataset, test_dataset
+
+
+def _make_generator(x_train, y_train, batch_size):
+    # This will do preprocessing and realtime data augmentation:
+    datagen = ImageDataGenerator(
+        featurewise_center=False,  # set input mean to 0 over the dataset
+        samplewise_center=False,  # set each sample mean to 0
+        # divide inputs by std of the dataset
+        featurewise_std_normalization=False,
+        samplewise_std_normalization=False,  # divide each input by its std
+        zca_whitening=False,  # apply ZCA whitening
+        zca_epsilon=1e-06,  # epsilon for ZCA whitening
+        # randomly rotate images in the range (degrees, 0 to 180)
+        rotation_range=0,
+        # randomly shift images horizontally (fraction of total width)
+        width_shift_range=0.1,
+        # randomly shift images vertically (fraction of total height)
+        height_shift_range=0.1,
+        shear_range=0.,  # set range for random shear
+        zoom_range=0.,  # set range for random zoom
+        channel_shift_range=0.,  # set range for random channel shifts
+        # set mode for filling points outside the input boundaries
+        fill_mode="nearest",
+        cval=0.,  # value used for fill_mode = "constant"
+        horizontal_flip=True,  # randomly flip images
+        vertical_flip=False,  # randomly flip images
+        # set rescaling factor (applied before any other transformation)
+        rescale=None,
+        # set function that will be applied on each input
+        preprocessing_function=None,
+        # image data format, either "channels_first" or "channels_last"
+        data_format=None,
+        # fraction of images reserved for validation (strictly between 0 and 1)
+        validation_split=0.0)
+
+    # Compute quantities required for feature-wise normalization
+    # (std, mean, and principal components if ZCA whitening is applied).
+    datagen.fit(x_train)
+    return datagen.flow(x_train, y_train, batch_size=batch_size)
+
+
+def data_augmentation_creator(config):
+    import tensorflow as tf
+    batch_size = config["batch_size"]
+    (x_train, y_train), (x_test, y_test) = fetch_keras_data()
+    trainset = tf.data.Dataset.from_generator(
+        lambda: _make_generator(x_train, y_train, batch_size),
+        output_types=(tf.float32, tf.float32),
+        # https://github.com/tensorflow/tensorflow/issues/24520
+        output_shapes=(tf.TensorShape((None, None, None, None)),
+                       tf.TensorShape((None, 10))))
+    trainset = trainset.repeat()
+
+    test_dataset = tf.data.Dataset.from_tensor_slices((x_test, y_test))
+    test_dataset = test_dataset.repeat().batch(batch_size)
+    return trainset, test_dataset
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--address",
+        required=False,
+        type=str,
+        help="the address to use for Ray")
+    parser.add_argument(
+        "--num-replicas",
+        "-n",
+        type=int,
+        default=1,
+        help="Sets number of replicas for training.")
+    parser.add_argument(
+        "--batch-size", type=int, default=32, help="Sets batch size.")
+    parser.add_argument(
+        "--use-gpu",
+        action="store_true",
+        default=False,
+        help="Enables GPU training")
+    parser.add_argument(
+        "--augment-data",
+        action="store_true",
+        default=False,
+        help="Sets data augmentation.")
+    parser.add_argument(
+        "--smoke-test",
+        action="store_true",
+        default=False,
+        help="Finish quickly for testing. Assume False for users.")
+
+    args, _ = parser.parse_known_args()
+    ray.init(address=args.address)
+    data_size = 60000
+    test_size = 10000
+    batch_size = args.batch_size
+
+    num_train_steps = 10 if args.smoke_test else data_size // batch_size
+    num_eval_steps = 10 if args.smoke_test else test_size // batch_size
+
+    trainer = TFTrainer(
+        model_creator=create_model,
+        data_creator=(data_augmentation_creator
+                      if args.augment_data else data_creator),
+        num_replicas=args.num_replicas,
+        use_gpu=args.use_gpu,
+        verbose=True,
+        config={
+            "batch_size": batch_size,
+            "fit_config": {
+                "steps_per_epoch": num_train_steps,
+            },
+            "evaluate_config": {
+                "steps": num_eval_steps,
+            }
+        })
+
+    training_start = time.time()
+    for i in range(3):
+        # Trains num epochs
+        train_stats1 = trainer.train()
+        train_stats1.update(trainer.validate())
+        print("iter {}:".format(i), train_stats1)
+
+    dt = (time.time() - training_start) / 3
+    print(f"Training on workers takes: {dt:.3f} seconds/epoch")
+
+    model = trainer.get_model()
+    trainer.shutdown()
+    dataset, test_dataset = data_augmentation_creator(
+        dict(batch_size=batch_size))
+
+    training_start = time.time()
+    model.fit(dataset, steps_per_epoch=num_train_steps, epochs=1)
+    dt = (time.time() - training_start)
+    print(f"Training on workers takes: {dt:.3f} seconds/epoch")
+
+    scores = model.evaluate(test_dataset, steps=num_eval_steps)
+    print("Test loss:", scores[0])
+    print("Test accuracy:", scores[1])

python/ray/util/sgd/tf/examples/tensorflow_train_example.py

@@ -0,0 +1,143 @@
+import argparse
+import tensorflow as tf
+from tensorflow.keras.models import Sequential
+from tensorflow.keras.layers import Dense
+import numpy as np
+
+import ray
+from ray import tune
+from ray.util.sgd.tf.tf_trainer import TFTrainer, TFTrainable
+
+NUM_TRAIN_SAMPLES = 1000
+NUM_TEST_SAMPLES = 400
+
+
+def create_config(batch_size):
+    return {
+        # todo: batch size needs to scale with # of workers
+        "batch_size": batch_size,
+        "fit_config": {
+            "steps_per_epoch": NUM_TRAIN_SAMPLES // batch_size
+        },
+        "evaluate_config": {
+            "steps": NUM_TEST_SAMPLES // batch_size,
+        }
+    }
+
+
+def linear_dataset(a=2, size=1000):
+    x = np.random.rand(size)
+    y = x / 2
+
+    x = x.reshape((-1, 1))
+    y = y.reshape((-1, 1))
+
+    return x, y
+
+
+def simple_dataset(config):
+    batch_size = config["batch_size"]
+    x_train, y_train = linear_dataset(size=NUM_TRAIN_SAMPLES)
+    x_test, y_test = linear_dataset(size=NUM_TEST_SAMPLES)
+
+    train_dataset = tf.data.Dataset.from_tensor_slices((x_train, y_train))
+    test_dataset = tf.data.Dataset.from_tensor_slices((x_test, y_test))
+    train_dataset = train_dataset.shuffle(NUM_TRAIN_SAMPLES).repeat().batch(
+        batch_size)
+    test_dataset = test_dataset.repeat().batch(batch_size)
+
+    return train_dataset, test_dataset
+
+
+def simple_model(config):
+    model = Sequential([Dense(10, input_shape=(1, )), Dense(1)])
+
+    model.compile(
+        optimizer="sgd",
+        loss="mean_squared_error",
+        metrics=["mean_squared_error"])
+
+    return model
+
+
+def train_example(num_replicas=1, batch_size=128, use_gpu=False):
+    trainer = TFTrainer(
+        model_creator=simple_model,
+        data_creator=simple_dataset,
+        num_replicas=num_replicas,
+        use_gpu=use_gpu,
+        verbose=True,
+        config=create_config(batch_size))
+
+    # model baseline performance
+    start_stats = trainer.validate()
+    print(start_stats)
+
+    # train for 2 epochs
+    trainer.train()
+    trainer.train()
+
+    # model performance after training (should improve)
+    end_stats = trainer.validate()
+    print(end_stats)
+
+    # sanity check that training worked
+    dloss = end_stats["validation_loss"] - start_stats["validation_loss"]
+    dmse = (end_stats["validation_mean_squared_error"] -
+            start_stats["validation_mean_squared_error"])
+    print(f"dLoss: {dloss}, dMSE: {dmse}")
+
+    if dloss > 0 or dmse > 0:
+        print("training sanity check failed. loss increased!")
+    else:
+        print("success!")
+
+
+def tune_example(num_replicas=1, use_gpu=False):
+    config = {
+        "model_creator": tune.function(simple_model),
+        "data_creator": tune.function(simple_dataset),
+        "num_replicas": num_replicas,
+        "use_gpu": use_gpu,
+        "trainer_config": create_config(batch_size=128)
+    }
+
+    analysis = tune.run(
+        TFTrainable,
+        num_samples=2,
+        config=config,
+        stop={"training_iteration": 2},
+        verbose=1)
+
+    return analysis.get_best_config(metric="validation_loss", mode="min")
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--address",
+        required=False,
+        type=str,
+        help="the address to use for Ray")
+    parser.add_argument(
+        "--num-replicas",
+        "-n",
+        type=int,
+        default=1,
+        help="Sets number of replicas for training.")
+    parser.add_argument(
+        "--use-gpu",
+        action="store_true",
+        default=False,
+        help="Enables GPU training")
+    parser.add_argument(
+        "--tune", action="store_true", default=False, help="Tune training")
+
+    args, _ = parser.parse_known_args()
+
+    ray.init(address=args.address)
+
+    if args.tune:
+        tune_example(num_replicas=args.num_replicas, use_gpu=args.use_gpu)
+    else:
+        train_example(num_replicas=args.num_replicas, use_gpu=args.use_gpu)

python/ray/util/sgd/tf/examples/tf-example-sgd.yaml

@@ -0,0 +1,70 @@
+# An unique identifier for the head node and workers of this cluster.
+cluster_name: sgd-tf
+
+# The maximum number of workers nodes to launch in addition to the head
+# node. This takes precedence over min_workers. min_workers default to 0.
+min_workers: 3
+initial_workers: 3
+max_workers: 3
+
+target_utilization_fraction: 0.9
+
+# If a node is idle for this many minutes, it will be removed.
+idle_timeout_minutes: 20
+# docker:
+#     image: tensorflow/tensorflow:1.5.0-py3
+#     container_name: ray_docker
+
+# Cloud-provider specific configuration.
+provider:
+    type: aws
+    region: us-west-1
+    availability_zone: us-west-1a
+
+# How Ray will authenticate with newly launched nodes.
+auth:
+    ssh_user: ubuntu
+
+head_node:
+    InstanceType: g4dn.xlarge
+    ImageId: ami-074c29e29c500f623 # latest_dlami on 01/28/20
+    #    InstanceMarketOptions:
+    #        MarketType: spot
+    #        SpotOptions:
+    #            MaxPrice: "9.0"
+
+
+worker_nodes:
+    InstanceType: g4dn.xlarge
+    ImageId: ami-074c29e29c500f623 # latest_dlami on 01/28/20
+    # InstanceMarketOptions:
+    #     MarketType: spot
+        # SpotOptions:
+        #     MaxPrice: "9.0"
+
+    #     # Run workers on spot by default. Comment this out to use on-demand.
+    #     InstanceMarketOptions:
+    #         MarketType: spot
+
+setup_commands:
+    - conda install setuptools=45.1.0=py36_0 wrapt=1.11.2 --yes # workaround to fix wrapt error
+    - ray &> /dev/null || pip install -U https://s3-us-west-2.amazonaws.com/ray-wheels/latest/ray-0.9.0.dev0-cp36-cp36m-manylinux1_x86_64.whl
+    - pip install -U ray[tune]
+    - pip install tensorflow-gpu==2.1.0
+
+# Custom commands that will be run on the head node after common setup.
+head_setup_commands: []
+
+# Custom commands that will be run on worker nodes after common setup.
+worker_setup_commands: []
+
+# # Command to start ray on the head node. You don't need to change this.
+head_start_ray_commands:
+    - ray stop
+    - ray start --head --redis-port=6379 --object-manager-port=8076 --autoscaling-config=~/ray_bootstrap_config.yaml --object-store-memory=1000000000
+
+# Command to start ray on worker nodes. You don't need to change this.
+worker_start_ray_commands:
+    - ray stop
+    - ray start --address=$RAY_HEAD_IP:6379 --object-manager-port=8076 --object-store-memory=1000000000
+

python/ray/util/sgd/tf/tf_runner.py

@@ -0,0 +1,163 @@
+import logging
+import json
+import os
+import numpy as np
+
+import ray
+import ray.services
+from ray.util.sgd import utils
+
+logger = logging.getLogger(__name__)
+
+
+def _try_import_strategy():
+    """Late import for Tesnorflow"""
+    import tensorflow as tf
+    return tf.distribute.experimental.MultiWorkerMirroredStrategy
+
+
+class TFRunner:
+    """Manages a TensorFlow model for training."""
+
+    def __init__(self, model_creator, data_creator, config=None,
+                 verbose=False):
+        """Initializes the runner.
+
+        Args:
+            model_creator (dict -> Model): see tf_trainer.py.
+            data_creator (dict -> tf.Dataset, tf.Dataset): see tf_trainer.py.
+            config (dict): see tf_trainer.py.
+            verbose (bool): Outputs training data if true.
+        """
+
+        self.model_creator = model_creator
+        self.data_creator = data_creator
+        self.config = {} if config is None else config
+        self.epoch = 0
+        self.verbose = verbose
+
+    def setup(self):
+        """Initializes the model."""
+        logger.debug("Creating dataset")
+        self.train_dataset, self.test_dataset = self.data_creator(self.config)
+
+        logger.debug("Creating model")
+        self.model = self.model_creator(self.config)
+
+    def setup_distributed(self, urls, world_rank, world_size):
+        """Sets up TensorFLow distributed environment and initializes the model.
+
+        Args:
+            urls (str): the URLs that each node uses to connect.
+            world_rank (int): the index of the runner.
+            world_size (int): the total number of runners.
+        """
+        assert len(urls) == world_size
+        tf_config = {
+            "cluster": {
+                "worker": urls
+            },
+            "task": {
+                "index": world_rank,
+                "type": "worker"
+            }
+        }
+        os.environ["TF_CONFIG"] = json.dumps(tf_config)
+
+        MultiWorkerMirroredStrategy = _try_import_strategy()
+
+        # MultiWorkerMirroredStrategy handles everything for us, from
+        # sharding the dataset (or even sharding the data itself if the loader
+        # reads files from disk) to merging the metrics and weight updates
+        #
+        # worker 0 is the "chief" worker and will handle the map-reduce
+        # every worker ends up with the exact same metrics and model
+        # after model.fit
+        #
+        # because of this, we only really ever need to query its state
+        self.strategy = MultiWorkerMirroredStrategy()
+
+        self.train_dataset, self.test_dataset = self.data_creator(self.config)
+
+        logger.debug("Creating model with MultiWorkerMirroredStrategy")
+        with self.strategy.scope():
+            self.model = self.model_creator(self.config)
+
+        # For use in model.evaluate()
+        self.local_model = None
+
+    def step(self):
+        """Runs a training epoch and updates the model parameters."""
+        fit_default_config = {"verbose": self.verbose}
+        fit_default_config.update(self.config.get("fit_config", {}))
+
+        history = self.model.fit(self.train_dataset, **fit_default_config)
+        if history is None:
+            stats = {}
+        else:
+            stats = {"train_" + k: v[-1] for k, v in history.history.items()}
+
+        self.epoch += 1
+        return stats
+
+    def validate(self):
+        """Evaluates the model on the validation data set."""
+        stats = {}
+        evaluate_config = {"verbose": self.verbose}
+        evaluate_config.update(self.config.get("evaluate_config", {}))
+
+        results = self.model.evaluate(self.test_dataset, **evaluate_config)
+        if results is None:
+            # Using local Model since model.evaluate() returns None
+            # for MultiWorkerMirroredStrategy
+            logger.warning("Running a local model to get validation score.")
+            self.local_model = self.model_creator(self.config)
+            self.local_model.set_weights(self.model.get_weights())
+            results = self.local_model.evaluate(self.test_dataset,
+                                                **evaluate_config)
+
+        if isinstance(results, list):
+            stats = {
+                "validation_" + k: v
+                for k, v in zip(self.model.metrics_names, results)
+            }
+        else:
+            stats = {"loss": results}
+
+        return stats
+
+    def get_state(self):
+        """Returns the state of the runner."""
+        return {
+            "epoch": self.epoch,
+            "weights": self.model.get_weights(),
+            "optimizer_weights": self.model.optimizer.get_weights()
+        }
+
+    def set_state(self, state):
+        """Sets the state of the model."""
+
+        self.model = self.model_creator(self.config)
+        self.epoch = state["epoch"]
+        self.model.set_weights(state["weights"])
+        # This part is due to ray.get() changing scalar np.int64 object to int
+        state["optimizer_weights"][0] = np.array(
+            state["optimizer_weights"][0], dtype=np.int64)
+
+        if self.model.optimizer.weights == []:
+            self.model._make_train_function()
+        self.model.optimizer.set_weights(state["optimizer_weights"])
+
+    def shutdown(self):
+        """Attempts to shut down the worker."""
+        del self.model
+        del self.train_dataset
+        del self.test_dataset
+
+    def get_node_ip(self):
+        """Returns the IP address of the current node."""
+        return ray.services.get_node_ip_address()
+
+    def find_free_port(self):
+        """Finds a free port on the current node."""
+        return utils.find_free_port()

python/ray/util/sgd/tf/tf_trainer.py

@@ -0,0 +1,203 @@
+import numpy as np
+import os
+import logging
+import pickle
+
+import ray
+
+from ray.tune import Trainable
+from ray.tune.resources import Resources
+from ray.util.sgd.tf.tf_runner import TFRunner
+
+logger = logging.getLogger(__name__)
+
+
+class TFTrainer:
+    def __init__(self,
+                 model_creator,
+                 data_creator,
+                 config=None,
+                 num_replicas=1,
+                 use_gpu=False,
+                 verbose=False):
+        """Sets up the TensorFlow trainer.
+
+        Args:
+            model_creator (dict -> Model): This function takes in the `config`
+                dict and returns a compiled TF model.
+            data_creator (dict -> tf.Dataset, tf.Dataset): Creates
+                the training and validation data sets using the config.
+                `config` dict is passed into the function.
+            config (dict): configuration passed to 'model_creator',
+                'data_creator'. Also contains `fit_config`, which is passed
+                into `model.fit(data, **fit_config)` and
+                `evaluate_config` which is passed into `model.evaluate`.
+            num_replicas (int): Sets number of workers used in distributed
+                training. Workers will be placed arbitrarily across the
+                cluster.
+            use_gpu (bool): Enables all workers to use GPU.
+            verbose (bool): Prints output of one model if true.
+        """
+        self.model_creator = model_creator
+        self.data_creator = data_creator
+        self.config = {} if config is None else config
+        self.use_gpu = use_gpu
+        self.num_replicas = num_replicas
+        self.verbose = verbose
+
+        # Generate actor class
+        # todo: are these resource quotas right?
+        # should they be exposed to the client codee?
+        Runner = ray.remote(num_cpus=1, num_gpus=int(use_gpu))(TFRunner)
+
+        # todo: should we warn about using
+        # distributed training on one device only?
+        # it's likely that whenever this happens it's a mistake
+        if num_replicas == 1:
+            # Start workers
+            self.workers = [
+                Runner.remote(
+                    model_creator,
+                    data_creator,
+                    config=self.config,
+                    verbose=self.verbose)
+            ]
+            # Get setup tasks in order to throw errors on failure
+            ray.get(self.workers[0].setup.remote())
+        else:
+            # Start workers
+            self.workers = [
+                Runner.remote(
+                    model_creator,
+                    data_creator,
+                    config=self.config,
+                    verbose=self.verbose and i == 0)
+                for i in range(num_replicas)
+            ]
+
+            # Compute URL for initializing distributed setup
+            ips = ray.get(
+                [worker.get_node_ip.remote() for worker in self.workers])
+            ports = ray.get(
+                [worker.find_free_port.remote() for worker in self.workers])
+
+            urls = [
+                "{ip}:{port}".format(ip=ips[i], port=ports[i])
+                for i in range(len(self.workers))
+            ]
+
+            # Get setup tasks in order to throw errors on failure
+            ray.get([
+                worker.setup_distributed.remote(urls, i, len(self.workers))
+                for i, worker in enumerate(self.workers)
+            ])
+
+    def train(self):
+        """Runs a training epoch."""
+
+        # see ./tf_runner.py:setup_distributed
+        # for an explanation of only taking the first worker's data
+        worker_stats = ray.get([w.step.remote() for w in self.workers])
+        stats = worker_stats[0].copy()
+        return stats
+
+    def validate(self):
+        """Evaluates the model on the validation data set."""
+        logger.info("Starting validation step.")
+
+        # see ./tf_runner.py:setup_distributed
+        # for an explanation of only taking the first worker's data
+        stats = ray.get([w.validate.remote() for w in self.workers])
+        stats = stats[0].copy()
+        return stats
+
+    def get_model(self):
+        """Returns the learned model."""
+        state = ray.get(self.workers[0].get_state.remote())
+        return self._get_model_from_state(state)
+
+    def save(self, checkpoint):
+        """Saves the model at the provided checkpoint.
+
+        Args:
+            checkpoint (str): Path to target checkpoint file.
+
+        """
+
+        state = ray.get(self.workers[0].get_state.remote())
+
+        with open(checkpoint, "wb") as f:
+            pickle.dump(state, f)
+
+        return checkpoint
+
+    def restore(self, checkpoint):
+        """Restores the model from the provided checkpoint.
+
+        Args:
+            checkpoint (str): Path to target checkpoint file.
+
+        """
+        with open(checkpoint, "rb") as f:
+            state = pickle.load(f)
+
+        state_id = ray.put(state)
+        ray.get([worker.set_state.remote(state_id) for worker in self.workers])
+
+    def shutdown(self):
+        """Shuts down workers and releases resources."""
+        for worker in self.workers:
+            worker.shutdown.remote()
+            worker.__ray_terminate__.remote()
+
+    def _get_model_from_state(self, state):
+        """Creates model and load weights from state"""
+
+        model = self.model_creator(self.config)
+        model.set_weights(state["weights"])
+
+        # This part is due to ray.get() changing scalar np.int64 object to int
+        state["optimizer_weights"][0] = np.array(
+            state["optimizer_weights"][0], dtype=np.int64)
+
+        if model.optimizer.weights == []:
+            model._make_train_function()
+        model.optimizer.set_weights(state["optimizer_weights"])
+
+        return model
+
+
+class TFTrainable(Trainable):
+    @classmethod
+    def default_resource_request(cls, config):
+        return Resources(
+            cpu=0,
+            gpu=0,
+            extra_cpu=config["num_replicas"],
+            extra_gpu=int(config["use_gpu"]) * config["num_replicas"])
+
+    def _setup(self, config):
+        self._trainer = TFTrainer(
+            model_creator=config["model_creator"],
+            data_creator=config["data_creator"],
+            config=config.get("trainer_config", {}),
+            num_replicas=config["num_replicas"],
+            use_gpu=config["use_gpu"])
+
+    def _train(self):
+
+        train_stats = self._trainer.train()
+        validation_stats = self._trainer.validate()
+
+        train_stats.update(validation_stats)
+
+        return train_stats
+
+    def _save(self, checkpoint_dir):
+        return self._trainer.save(os.path.join(checkpoint_dir, "model"))
+
+    def _restore(self, checkpoint_path):
+        return self._trainer.restore(checkpoint_path)
+
+    def _stop(self):
+        self._trainer.shutdown()

python/ray/util/sgd/utils.py

@@ -0,0 +1,151 @@
+from contextlib import closing
+import logging
+import numpy as np
+import socket
+import time
+
+import ray
+from ray.exceptions import RayActorError
+
+logger = logging.getLogger(__name__)
+
+
+class TimerStat:
+    """A running stat for conveniently logging the duration of a code block.
+
+    Note that this class is *not* thread-safe.
+
+    Examples:
+        Time a call to 'time.sleep'.
+
+        >>> import time
+        >>> sleep_timer = TimerStat()
+        >>> with sleep_timer:
+        ...     time.sleep(1)
+        >>> round(sleep_timer.mean)
+        1
+    """
+
+    def __init__(self, window_size=10):
+        self._window_size = window_size
+        self._samples = []
+        self._units_processed = []
+        self._start_time = None
+        self._total_time = 0.0
+        self.count = 0
+
+    def __enter__(self):
+        assert self._start_time is None, "concurrent updates not supported"
+        self._start_time = time.time()
+
+    def __exit__(self, type, value, tb):
+        assert self._start_time is not None
+        time_delta = time.time() - self._start_time
+        self.push(time_delta)
+        self._start_time = None
+
+    def push(self, time_delta):
+        self._samples.append(time_delta)
+        if len(self._samples) > self._window_size:
+            self._samples.pop(0)
+        self.count += 1
+        self._total_time += time_delta
+
+    def push_units_processed(self, n):
+        self._units_processed.append(n)
+        if len(self._units_processed) > self._window_size:
+            self._units_processed.pop(0)
+
+    @property
+    def mean(self):
+        return np.mean(self._samples)
+
+    @property
+    def median(self):
+        return np.median(self._samples)
+
+    @property
+    def sum(self):
+        return np.sum(self._samples)
+
+    @property
+    def max(self):
+        return np.max(self._samples)
+
+    @property
+    def first(self):
+        return self._samples[0] if self._samples else None
+
+    @property
+    def last(self):
+        return self._samples[-1] if self._samples else None
+
+    @property
+    def size(self):
+        return len(self._samples)
+
+    @property
+    def mean_units_processed(self):
+        return float(np.mean(self._units_processed))
+
+    @property
+    def mean_throughput(self):
+        time_total = sum(self._samples)
+        if not time_total:
+            return 0.0
+        return sum(self._units_processed) / time_total
+
+    def reset(self):
+        self._samples = []
+        self._units_processed = []
+        self._start_time = None
+        self._total_time = 0.0
+        self.count = 0
+
+
+def find_free_port():
+    with closing(socket.socket(socket.AF_INET, socket.SOCK_STREAM)) as s:
+        s.bind(("", 0))
+        s.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
+        return s.getsockname()[1]
+
+
+class AverageMeter:
+    """Computes and stores the average and current value."""
+
+    def __init__(self):
+        self.reset()
+
+    def reset(self):
+        self.val = 0
+        self.avg = 0
+        self.sum = 0
+        self.count = 0
+
+    def update(self, val, n=1):
+        self.val = val
+        self.sum += val * n
+        self.count += n
+        self.avg = self.sum / self.count
+
+
+def check_for_failure(remote_values):
+    """Checks remote values for any that returned and failed.
+
+    Args:
+        remote_values (list): List of object IDs representing functions
+            that may fail in the middle of execution. For example, running
+            a SGD training loop in multiple parallel actor calls.
+
+    Returns:
+        Bool for success in executing given remote tasks.
+    """
+    unfinished = remote_values
+    try:
+        while len(unfinished) > 0:
+            finished, unfinished = ray.wait(unfinished)
+            finished = ray.get(finished)
+        return True
+    except RayActorError as exc:
+        logger.exception(str(exc))
+    return False