wassname/ray
1
0
Fork 0
mirror of https://github.com/wassname/ray.git synced 2026-07-15 11:25:40 +08:00
Code Issues Packages Projects Releases Wiki Activity

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

Browse Source
This commit is contained in:
Eric Liang
2020-02-18 13:43:19 -08:00
committed by GitHub
parent fae99ecb8e
commit 5df801605e
113 changed files with 305 additions and 637 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/ray/experimental/__init__.py
-20
View File
@@ -1,28 +1,8 @@
from .gcs_flush_policy import (set_flushing_policy, GcsFlushPolicy,
SimpleGcsFlushPolicy)
from .named_actors import get_actor, register_actor
from .api import get, wait
from .actor_pool import ActorPool
from .dynamic_resources import set_resource
from . import iter
def TensorFlowVariables(*args, **kwargs):
raise DeprecationWarning(
"'ray.experimental.TensorFlowVariables' is deprecated. Instead, please"
" do 'from ray.experimental.tf_utils import TensorFlowVariables'.")
__all__ = [
"TensorFlowVariables",
"get_actor",
"register_actor",
"get",
"wait",
"set_flushing_policy",
"GcsFlushPolicy",
"SimpleGcsFlushPolicy",
"set_resource",
"ActorPool",
"iter",
]
python/ray/experimental/actor_pool.py
-212
View File
@@ -1,212 +0,0 @@
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/experimental/gcs_flush_policy.py
-89
View File
@@ -1,89 +0,0 @@
import os
import time
import ray
import ray.cloudpickle as pickle
class GcsFlushPolicy:
"""Experimental: a policy to control GCS flushing.
Used by Monitor to enable automatic control of memory usage.
"""
def should_flush(self, redis_client):
"""Returns a bool, whether a flush request should be issued."""
pass
def num_entries_to_flush(self):
"""Returns an upper bound for number of entries to flush next."""
pass
def record_flush(self):
"""Must be called after a flush has been performed."""
pass
class SimpleGcsFlushPolicy(GcsFlushPolicy):
"""A simple policy with constant flush rate, after a warmup period.
Example policy values:
flush_when_at_least_bytes 2GB
flush_period_secs 10s
flush_num_entries_each_time 10k
This means: (1) If the GCS shard uses less than 2GB of memory,
no flushing would take place. This should cover most Ray runs. (2) The
GCS shard will only honor a flush request, if it's issued after 10
seconds since the last processed flush. In particular this means it's
okay for the Monitor to issue requests more frequently than this param.
(3) When processing a flush, the shard will flush at most 10k entries.
This is to control the latency of each request.
Note, flush rate == (flush period) * (num entries each time). So
applications that have a heavier GCS load can tune these params.
"""
def __init__(self,
flush_when_at_least_bytes=(1 << 31),
flush_period_secs=10,
flush_num_entries_each_time=10000):
self.flush_when_at_least_bytes = flush_when_at_least_bytes
self.flush_period_secs = flush_period_secs
self.flush_num_entries_each_time = flush_num_entries_each_time
self.last_flush_timestamp = time.time()
def should_flush(self, redis_client):
if time.time() - self.last_flush_timestamp < self.flush_period_secs:
return False
used_memory = redis_client.info("memory")["used_memory"]
assert used_memory > 0
return used_memory >= self.flush_when_at_least_bytes
def num_entries_to_flush(self):
return self.flush_num_entries_each_time
def record_flush(self):
self.last_flush_timestamp = time.time()
def serialize(self):
return pickle.dumps(self)
def set_flushing_policy(flushing_policy):
"""Serialize this policy for Monitor to pick up."""
if "RAY_USE_NEW_GCS" not in os.environ:
raise Exception(
"set_flushing_policy() is only available when environment "
"variable RAY_USE_NEW_GCS is present at both compile and run time."
)
ray.worker.global_worker.check_connected()
redis_client = ray.worker.global_worker.redis_client
serialized = pickle.dumps(flushing_policy)
redis_client.set("gcs_flushing_policy", serialized)
python/ray/experimental/iter.py
-749
View File
@@ -1,749 +0,0 @@
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.experimental.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.experimental.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.experimental.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.experimental.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/experimental/joblib/__init__.py
-17
View File
@@ -1,17 +0,0 @@
from joblib.parallel import register_parallel_backend
def register_ray():
""" Register Ray Backend to be called with parallel_backend("ray"). """
try:
from ray.experimental.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/experimental/joblib/ray_backend.py
-58
View File
@@ -1,58 +0,0 @@
from joblib._parallel_backends import MultiprocessingBackend
from joblib.pool import PicklingPool
import logging
from ray.experimental.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.experimental.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/experimental/named_actors.py
-62
View File
@@ -1,62 +0,0 @@
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/experimental/serve/BUILD
-28
View File
@@ -1,28 +0,0 @@
# This is a dummy test dependency that causes the above tests to be
# re-run if any of these files changes.
py_library(
name = "serve_lib",
srcs = glob(["**/*.py"], exclude=["tests/*.py"]),
)
# This test aggregates all serve tests and run them in a single session
# similar to `pytest .`
# Serve tests need to run in a single session because starting and stopping
# serve cluster take a large chunk of time. All serve tests use a shared
# cluster.
py_test(
name = "test_serve",
size = "medium",
srcs = glob(["tests/*.py"]),
tags = ["exclusive"],
deps = [":serve_lib"],
)
# Make sure the example showing in doc is tested
py_test(
name = "echo_full",
size = "small",
srcs = glob(["examples/*.py"]),
tags = ["exclusive"],
deps = [":serve_lib"]
)
python/ray/experimental/serve/__init__.py
-11
View File
@@ -1,11 +0,0 @@
from ray.experimental.serve.backend_config import BackendConfig
from ray.experimental.serve.policy import RoutePolicy
from ray.experimental.serve.api import (
init, create_backend, create_endpoint, link, split, get_handle, stat,
set_backend_config, get_backend_config, accept_batch, route) # noqa: E402
__all__ = [
"init", "create_backend", "create_endpoint", "link", "split", "get_handle",
"stat", "set_backend_config", "get_backend_config", "BackendConfig",
"RoutePolicy", "accept_batch", "route"
]
python/ray/experimental/serve/api.py
-475
View File
@@ -1,475 +0,0 @@
import inspect
from functools import wraps
from tempfile import mkstemp
from multiprocessing import cpu_count
import numpy as np
import ray
from ray.experimental.serve.constants import (
DEFAULT_HTTP_HOST, DEFAULT_HTTP_PORT, SERVE_NURSERY_NAME)
from ray.experimental.serve.global_state import (GlobalState,
start_initial_state)
from ray.experimental.serve.kv_store_service import SQLiteKVStore
from ray.experimental.serve.task_runner import RayServeMixin, TaskRunnerActor
from ray.experimental.serve.utils import (block_until_http_ready,
get_random_letters, expand)
from ray.experimental.serve.exceptions import RayServeException
from ray.experimental.serve.backend_config import BackendConfig
from ray.experimental.serve.policy import RoutePolicy
from ray.experimental.serve.queues import Query
global_state = None
def _get_global_state():
"""Used for internal purpose. Because just import serve.global_state
will always reference the original None object
"""
return global_state
def _ensure_connected(f):
@wraps(f)
def check(*args, **kwargs):
if _get_global_state() is None:
raise RayServeException("Please run serve.init to initialize or "
"connect to existing ray serve cluster.")
return f(*args, **kwargs)
return check
def accept_batch(f):
"""Annotation to mark a serving function that batch is accepted.
This annotation need to be used to mark a function expect all arguments
to be passed into a list.
Example:
>>> @serve.accept_batch
def serving_func(flask_request):
assert isinstance(flask_request, list)
...
>>> class ServingActor:
@serve.accept_batch
def __call__(self, *, python_arg=None):
assert isinstance(python_arg, list)
"""
f.serve_accept_batch = True
return f
def init(kv_store_connector=None,
kv_store_path=None,
blocking=False,
start_server=True,
http_host=DEFAULT_HTTP_HOST,
http_port=DEFAULT_HTTP_PORT,
ray_init_kwargs={
"object_store_memory": int(1e8),
"num_cpus": max(cpu_count(), 8)
},
gc_window_seconds=3600,
queueing_policy=RoutePolicy.Random,
policy_kwargs={}):
"""Initialize a serve cluster.
If serve cluster has already initialized, this function will just return.
Calling `ray.init` before `serve.init` is optional. When there is not a ray
cluster initialized, serve will call `ray.init` with `object_store_memory`
requirement.
Args:
kv_store_connector (callable): Function of (namespace) => TableObject.
We will use a SQLite connector that stores to /tmp by default.
kv_store_path (str, path): Path to the SQLite table.
blocking (bool): If true, the function will wait for the HTTP server to
be healthy, and other components to be ready before returns.
start_server (bool): If true, `serve.init` starts http server.
(Default: True)
http_host (str): Host for HTTP server. Default to "0.0.0.0".
http_port (int): Port for HTTP server. Default to 8000.
ray_init_kwargs (dict): Argument passed to ray.init, if there is no ray
connection. Default to {"object_store_memory": int(1e8)} for
performance stability reason
gc_window_seconds(int): How long will we keep the metric data in
memory. Data older than the gc_window will be deleted. The default
is 3600 seconds, which is 1 hour.
queueing_policy(RoutePolicy): Define the queueing policy for selecting
the backend for a service. (Default: RoutePolicy.Random)
policy_kwargs: Arguments required to instantiate a queueing policy
"""
global global_state
# Noop if global_state is no longer None
if global_state is not None:
return
# Initialize ray if needed.
if not ray.is_initialized():
ray.init(**ray_init_kwargs)
# Try to get serve nursery if there exists
try:
ray.experimental.get_actor(SERVE_NURSERY_NAME)
global_state = GlobalState()
return
except ValueError:
pass
# Register serialization context once
ray.register_custom_serializer(Query, Query.ray_serialize,
Query.ray_deserialize)
if kv_store_path is None:
_, kv_store_path = mkstemp()
# Serve has not been initialized, perform init sequence
# Todo, move the db to session_dir
# ray.worker._global_node.address_info["session_dir"]
def kv_store_connector(namespace):
return SQLiteKVStore(namespace, db_path=kv_store_path)
nursery = start_initial_state(kv_store_connector)
global_state = GlobalState(nursery)
if start_server:
global_state.init_or_get_http_server(host=http_host, port=http_port)
global_state.init_or_get_router(
queueing_policy=queueing_policy, policy_kwargs=policy_kwargs)
global_state.init_or_get_metric_monitor(
gc_window_seconds=gc_window_seconds)
if start_server and blocking:
block_until_http_ready("http://{}:{}".format(http_host, http_port))
@_ensure_connected
def create_endpoint(endpoint_name, route=None, blocking=True):
"""Create a service endpoint given route_expression.
Args:
endpoint_name (str): A name to associate to the endpoint. It will be
used as key to set traffic policy.
route (str): A string begin with "/". HTTP server will use
the string to match the path.
blocking (bool): If true, the function will wait for service to be
registered before returning
"""
global_state.route_table.register_service(route, endpoint_name)
@_ensure_connected
def set_backend_config(backend_tag, backend_config):
"""Set a backend configuration for a backend tag
Args:
backend_tag(str): A registered backend.
backend_config(BackendConfig) : Desired backend configuration.
"""
assert backend_tag in global_state.backend_table.list_backends(), (
"Backend {} is not registered.".format(backend_tag))
assert isinstance(backend_config,
BackendConfig), ("backend_config must be"
" of instance BackendConfig")
backend_config_dict = dict(backend_config)
old_backend_config_dict = global_state.backend_table.get_info(backend_tag)
global_state.backend_table.register_info(backend_tag, backend_config_dict)
# inform the router about change in configuration
# particularly for setting max_batch_size
ray.get(global_state.init_or_get_router().set_backend_config.remote(
backend_tag, backend_config_dict))
# checking if replicas need to be restarted
# Replicas are restarted if there is any change in the backend config
# related to restart_configs
# TODO(alind) : have replica restarting policies selected by the user
need_to_restart_replicas = any(
old_backend_config_dict[k] != backend_config_dict[k]
for k in BackendConfig.restart_on_change_fields)
if need_to_restart_replicas:
# kill all the replicas for restarting with new configurations
scale(backend_tag, 0)
# scale the replicas with new configuration
scale(backend_tag, backend_config_dict["num_replicas"])
@_ensure_connected
def get_backend_config(backend_tag):
"""get the backend configuration for a backend tag
Args:
backend_tag(str): A registered backend.
"""
assert backend_tag in global_state.backend_table.list_backends(), (
"Backend {} is not registered.".format(backend_tag))
backend_config_dict = global_state.backend_table.get_info(backend_tag)
return BackendConfig(**backend_config_dict)
@_ensure_connected
def create_backend(func_or_class,
backend_tag,
*actor_init_args,
backend_config=BackendConfig()):
"""Create a backend using func_or_class and assign backend_tag.
Args:
func_or_class (callable, class): a function or a class implements
__call__ protocol.
backend_tag (str): a unique tag assign to this backend. It will be used
to associate services in traffic policy.
backend_config (BackendConfig): An object defining backend properties
for starting a backend.
*actor_init_args (optional): the argument to pass to the class
initialization method.
"""
assert isinstance(backend_config,
BackendConfig), ("backend_config must be"
" of instance BackendConfig")
backend_config_dict = dict(backend_config)
should_accept_batch = (True if backend_config.max_batch_size is not None
else False)
batch_annotation_not_found = RayServeException(
"max_batch_size is set in config but the function or method does not "
"accept batching. Please use @serve.accept_batch to explicitly mark "
"the function or method as batchable and takes in list as arguments.")
arg_list = []
if inspect.isfunction(func_or_class):
if should_accept_batch and not hasattr(func_or_class,
"serve_accept_batch"):
raise batch_annotation_not_found
# arg list for a fn is function itself
arg_list = [func_or_class]
# ignore lint on lambda expression
creator = lambda kwrgs: TaskRunnerActor._remote(**kwrgs) # noqa: E731
elif inspect.isclass(func_or_class):
if should_accept_batch and not hasattr(func_or_class.__call__,
"serve_accept_batch"):
raise batch_annotation_not_found
# Python inheritance order is right-to-left. We put RayServeMixin
# on the left to make sure its methods are not overriden.
@ray.remote
class CustomActor(RayServeMixin, func_or_class):
pass
arg_list = actor_init_args
# ignore lint on lambda expression
creator = lambda kwargs: CustomActor._remote(**kwargs) # noqa: E731
else:
raise TypeError(
"Backend must be a function or class, it is {}.".format(
type(func_or_class)))
# save creator which starts replicas
global_state.backend_table.register_backend(backend_tag, creator)
# save information about configurations needed to start the replicas
global_state.backend_table.register_info(backend_tag, backend_config_dict)
# save the initial arguments needed by replicas
global_state.backend_table.save_init_args(backend_tag, arg_list)
# set the backend config inside the router
# particularly for max-batch-size
ray.get(global_state.init_or_get_router().set_backend_config.remote(
backend_tag, backend_config_dict))
scale(backend_tag, backend_config_dict["num_replicas"])
def _start_replica(backend_tag):
assert backend_tag in global_state.backend_table.list_backends(), (
"Backend {} is not registered.".format(backend_tag))
replica_tag = "{}#{}".format(backend_tag, get_random_letters(length=6))
# get the info which starts the replicas
creator = global_state.backend_table.get_backend_creator(backend_tag)
backend_config_dict = global_state.backend_table.get_info(backend_tag)
backend_config = BackendConfig(**backend_config_dict)
init_args = global_state.backend_table.get_init_args(backend_tag)
# get actor creation kwargs
actor_kwargs = backend_config.get_actor_creation_args(init_args)
# Create the runner in the nursery
[runner_handle] = ray.get(
global_state.actor_nursery_handle.start_actor_with_creator.remote(
creator, actor_kwargs, replica_tag))
# Setup the worker
ray.get(
runner_handle._ray_serve_setup.remote(
backend_tag, global_state.init_or_get_router(), runner_handle))
runner_handle._ray_serve_fetch.remote()
# Register the worker in config tables as well as metric monitor
global_state.backend_table.add_replica(backend_tag, replica_tag)
global_state.init_or_get_metric_monitor().add_target.remote(runner_handle)
def _remove_replica(backend_tag):
assert backend_tag in global_state.backend_table.list_backends(), (
"Backend {} is not registered.".format(backend_tag))
assert len(global_state.backend_table.list_replicas(backend_tag)) > 0, (
"Backend {} does not have enough replicas to be removed.".format(
backend_tag))
replica_tag = global_state.backend_table.remove_replica(backend_tag)
[replica_handle] = ray.get(
global_state.actor_nursery_handle.get_handle.remote(replica_tag))
# Remove the replica from metric monitor.
ray.get(global_state.init_or_get_metric_monitor().remove_target.remote(
replica_handle))
# Remove the replica from actor nursery.
ray.get(
global_state.actor_nursery_handle.remove_handle.remote(replica_tag))
# Remove the replica from router.
# This will also destory the actor handle.
ray.get(global_state.init_or_get_router()
.remove_and_destory_replica.remote(backend_tag, replica_handle))
@_ensure_connected
def scale(backend_tag, num_replicas):
"""Set the number of replicas for backend_tag.
Args:
backend_tag (str): A registered backend.
num_replicas (int): Desired number of replicas
"""
assert backend_tag in global_state.backend_table.list_backends(), (
"Backend {} is not registered.".format(backend_tag))
assert num_replicas >= 0, ("Number of replicas must be"
" greater than or equal to 0.")
replicas = global_state.backend_table.list_replicas(backend_tag)
current_num_replicas = len(replicas)
delta_num_replicas = num_replicas - current_num_replicas
if delta_num_replicas > 0:
for _ in range(delta_num_replicas):
_start_replica(backend_tag)
elif delta_num_replicas < 0:
for _ in range(-delta_num_replicas):
_remove_replica(backend_tag)
@_ensure_connected
def link(endpoint_name, backend_tag):
"""Associate a service endpoint with backend tag.
Example:
>>> serve.link("service-name", "backend:v1")
Note:
This is equivalent to
>>> serve.split("service-name", {"backend:v1": 1.0})
"""
split(endpoint_name, {backend_tag: 1.0})
@_ensure_connected
def split(endpoint_name, traffic_policy_dictionary):
"""Associate a service endpoint with traffic policy.
Example:
>>> serve.split("service-name", {
"backend:v1": 0.5,
"backend:v2": 0.5
})
Args:
endpoint_name (str): A registered service endpoint.
traffic_policy_dictionary (dict): a dictionary maps backend names
to their traffic weights. The weights must sum to 1.
"""
assert endpoint_name in expand(
global_state.route_table.list_service(include_headless=True).values())
assert isinstance(traffic_policy_dictionary,
dict), "Traffic policy must be dictionary"
prob = 0
for backend, weight in traffic_policy_dictionary.items():
prob += weight
assert (backend in global_state.backend_table.list_backends()
), "backend {} is not registered".format(backend)
assert np.isclose(
prob, 1,
atol=0.02), "weights must sum to 1, currently it sums to {}".format(
prob)
global_state.policy_table.register_traffic_policy(
endpoint_name, traffic_policy_dictionary)
ray.get(global_state.init_or_get_router().set_traffic.remote(
endpoint_name, traffic_policy_dictionary))
@_ensure_connected
def get_handle(endpoint_name, relative_slo_ms=None, absolute_slo_ms=None):
"""Retrieve RayServeHandle for service endpoint to invoke it from Python.
Args:
endpoint_name (str): A registered service endpoint.
relative_slo_ms(float): Specify relative deadline in milliseconds for
queries fired using this handle. (Default: None)
absolute_slo_ms(float): Specify absolute deadline in milliseconds for
queries fired using this handle. (Default: None)
Returns:
RayServeHandle
"""
assert endpoint_name in expand(
global_state.route_table.list_service(include_headless=True).values())
# Delay import due to it's dependency on global_state
from ray.experimental.serve.handle import RayServeHandle
return RayServeHandle(global_state.init_or_get_router(), endpoint_name,
relative_slo_ms, absolute_slo_ms)
@_ensure_connected
def stat(percentiles=[50, 90, 95],
agg_windows_seconds=[10, 60, 300, 600, 3600]):
"""Retrieve metric statistics about ray serve system.
Args:
percentiles(List[int]): The percentiles for aggregation operations.
Default is 50th, 90th, 95th percentile.
agg_windows_seconds(List[int]): The aggregation windows in seconds.
The longest aggregation window must be shorter or equal to the
gc_window_seconds.
"""
return ray.get(global_state.init_or_get_metric_monitor().collect.remote(
percentiles, agg_windows_seconds))
class route:
def __init__(self, url_route):
self.route = url_route
def __call__(self, func_or_class):
name = func_or_class.__name__
backend_tag = "{}:v0".format(name)
create_backend(func_or_class, backend_tag)
create_endpoint(name, self.route)
link(name, backend_tag)
python/ray/experimental/serve/backend_config.py
-58
View File
@@ -1,58 +0,0 @@
from copy import deepcopy
class BackendConfig:
# configs not needed for actor creation when
# instantiating a replica
_serve_configs = ["_num_replicas", "max_batch_size"]
# configs which when changed leads to restarting
# the existing replicas.
restart_on_change_fields = ["resources", "num_cpus", "num_gpus"]
def __init__(self,
num_replicas=1,
resources=None,
max_batch_size=None,
num_cpus=None,
num_gpus=None,
memory=None,
object_store_memory=None):
"""
Class for defining backend configuration.
"""
# serve configs
self.num_replicas = num_replicas
self.max_batch_size = max_batch_size
# ray actor configs
self.resources = resources
self.num_cpus = num_cpus
self.num_gpus = num_gpus
self.memory = memory
self.object_store_memory = object_store_memory
@property
def num_replicas(self):
return self._num_replicas
@num_replicas.setter
def num_replicas(self, val):
if not (val > 0):
raise Exception("num_replicas must be greater than zero")
self._num_replicas = val
def __iter__(self):
for k in self.__dict__.keys():
key, val = k, self.__dict__[k]
if key == "_num_replicas":
key = "num_replicas"
yield key, val
def get_actor_creation_args(self, init_args):
ret_d = deepcopy(self.__dict__)
for k in self._serve_configs:
ret_d.pop(k)
ret_d["args"] = init_args
return ret_d
python/ray/experimental/serve/constants.py
-26
View File
@@ -1,26 +0,0 @@
#: The interval which http server refreshes its routing table
HTTP_ROUTER_CHECKER_INTERVAL_S = 2
#: Actor name used to register actor nursery
SERVE_NURSERY_NAME = "SERVE_ACTOR_NURSERY"
#: KVStore connector key in bootstrap config
BOOTSTRAP_KV_STORE_CONN_KEY = "kv_store_connector"
#: HTTP Address
DEFAULT_HTTP_ADDRESS = "http://127.0.0.1:8000"
#: HTTP Host
DEFAULT_HTTP_HOST = "127.0.0.1"
#: HTTP Port
DEFAULT_HTTP_PORT = 8000
#: Max concurrency
ASYNC_CONCURRENCY = int(1e6)
#: Default latency SLO
DEFAULT_LATENCY_SLO_MS = 1e9
#: Key for storing no http route services
NO_ROUTE_KEY = "NO_ROUTE"
python/ray/experimental/serve/context.py
-34
View File
@@ -1,34 +0,0 @@
from enum import IntEnum
from ray.experimental.serve.exceptions import RayServeException
class TaskContext(IntEnum):
"""TaskContext constants for queue.enqueue method"""
Web = 1
Python = 2
# Global variable will be modified in worker
# web == True: currrently processing a request from web server
# web == False: currently processing a request from python
web = False
# batching information in serve context
# batch_size == None : the backend doesn't support batching
# batch_size(int) : the number of elements of input list
batch_size = None
_not_in_web_context_error = """
Accessing the request object outside of the web context. Please use
"serve.context.web" to determine when the function is called within
a web context.
"""
class FakeFlaskRequest:
def __getattribute__(self, name):
raise RayServeException(_not_in_web_context_error)
def __setattr__(self, name, value):
raise RayServeException(_not_in_web_context_error)
python/ray/experimental/serve/examples/benchmark.py
-33
View File
@@ -1,33 +0,0 @@
from ray.experimental import serve
from ray.experimental.serve.constants import DEFAULT_HTTP_ADDRESS
import requests
import time
import pandas as pd
from tqdm import tqdm
serve.init(blocking=True)
@serve.route("/noop")
def noop(_):
return ""
url = "{}/noop".format(DEFAULT_HTTP_ADDRESS)
while requests.get(url).status_code == 404:
time.sleep(1)
print("Waiting for noop route to showup.")
latency = []
for _ in tqdm(range(5200)):
start = time.perf_counter()
resp = requests.get(url)
end = time.perf_counter()
latency.append(end - start)
# Remove initial samples
latency = latency[200:]
series = pd.Series(latency) * 1000
print("Latency for single noop backend (ms)")
print(series.describe(percentiles=[0.5, 0.9, 0.95, 0.99]))
python/ray/experimental/serve/examples/echo.py
-28
View File
@@ -1,28 +0,0 @@
"""
Example service that prints out http context.
"""
import time
import requests
from ray.experimental import serve
from ray.experimental.serve.utils import pformat_color_json
def echo(flask_request):
return "hello " + flask_request.args.get("name", "serve!")
serve.init(blocking=True)
serve.create_endpoint("my_endpoint", "/echo", blocking=True)
serve.create_backend(echo, "echo:v1")
serve.link("my_endpoint", "echo:v1")
while True:
resp = requests.get("http://127.0.0.1:8000/echo").json()
print(pformat_color_json(resp))
print("...Sleeping for 2 seconds...")
time.sleep(2)
python/ray/experimental/serve/examples/echo_actor.py
-46
View File
@@ -1,46 +0,0 @@
"""
Example actor that adds an increment to a number. This number can
come from either web (parsing Flask request) or python call.
This actor can be called from HTTP as well as from Python.
"""
import time
import requests
import ray
from ray.experimental import serve
from ray.experimental.serve.utils import pformat_color_json
class MagicCounter:
def __init__(self, increment):
self.increment = increment
def __call__(self, flask_request, base_number=None):
if serve.context.web:
base_number = int(flask_request.args.get("base_number", "0"))
return base_number + self.increment
serve.init(blocking=True)
serve.create_endpoint("magic_counter", "/counter", blocking=True)
serve.create_backend(MagicCounter, "counter:v1", 42) # increment=42
serve.link("magic_counter", "counter:v1")
print("Sending ten queries via HTTP")
for i in range(10):
url = "http://127.0.0.1:8000/counter?base_number={}".format(i)
print("> Pinging {}".format(url))
resp = requests.get(url).json()
print(pformat_color_json(resp))
time.sleep(0.2)
print("Sending ten queries via Python")
handle = serve.get_handle("magic_counter")
for i in range(10):
print("> Pinging handle.remote(base_number={})".format(i))
result = ray.get(handle.remote(base_number=i))
print("< Result {}".format(result))
python/ray/experimental/serve/examples/echo_actor_batch.py
-60
View File
@@ -1,60 +0,0 @@
"""
Example actor that adds an increment to a number. This number can
come from either web (parsing Flask request) or python call.
The queries incoming to this actor are batched.
This actor can be called from HTTP as well as from Python.
"""
import time
import requests
import ray
from ray.experimental import serve
from ray.experimental.serve.utils import pformat_color_json
from ray.experimental.serve import BackendConfig
class MagicCounter:
def __init__(self, increment):
self.increment = increment
@serve.accept_batch
def __call__(self, flask_request_list, base_number=None):
# batch_size = serve.context.batch_size
if serve.context.web:
result = []
for flask_request in flask_request_list:
base_number = int(flask_request.args.get("base_number", "0"))
result.append(base_number)
return list(map(lambda x: x + self.increment, result))
else:
result = []
for b in base_number:
ans = b + self.increment
result.append(ans)
return result
serve.init(blocking=True)
serve.create_endpoint("magic_counter", "/counter", blocking=True)
b_config = BackendConfig(max_batch_size=5)
serve.create_backend(
MagicCounter, "counter:v1", 42, backend_config=b_config) # increment=42
serve.link("magic_counter", "counter:v1")
print("Sending ten queries via HTTP")
for i in range(10):
url = "http://127.0.0.1:8000/counter?base_number={}".format(i)
print("> Pinging {}".format(url))
resp = requests.get(url).json()
print(pformat_color_json(resp))
time.sleep(0.2)
print("Sending ten queries via Python")
handle = serve.get_handle("magic_counter")
for i in range(10):
print("> Pinging handle.remote(base_number={})".format(i))
result = ray.get(handle.remote(base_number=i))
print("< Result {}".format(result))
python/ray/experimental/serve/examples/echo_batching.py
-56
View File
@@ -1,56 +0,0 @@
"""
This example has backend which has batching functionality enabled.
"""
import ray
from ray.experimental import serve
from ray.experimental.serve import BackendConfig
class MagicCounter:
def __init__(self, increment):
self.increment = increment
@serve.accept_batch
def __call__(self, flask_request, base_number=None):
# __call__ fn should preserve the batch size
# base_number is a python list
if serve.context.batch_size is not None:
batch_size = serve.context.batch_size
result = []
for base_num in base_number:
ret_str = "Number: {} Batch size: {}".format(
base_num, batch_size)
result.append(ret_str)
return result
return ""
serve.init(blocking=True)
serve.create_endpoint("magic_counter", "/counter", blocking=True)
# specify max_batch_size in BackendConfig
b_config = BackendConfig(max_batch_size=5)
serve.create_backend(
MagicCounter, "counter:v1", 42, backend_config=b_config) # increment=42
print("Backend Config for backend: 'counter:v1'")
print(b_config)
serve.link("magic_counter", "counter:v1")
handle = serve.get_handle("magic_counter")
future_list = []
# fire 30 requests
for r in range(30):
print("> [REMOTE] Pinging handle.remote(base_number={})".format(r))
f = handle.remote(base_number=r)
future_list.append(f)
# get results of queries as they complete
left_futures = future_list
while left_futures:
completed_futures, remaining_futures = ray.wait(left_futures, timeout=0.05)
if len(completed_futures) > 0:
result = ray.get(completed_futures[0])
print("< " + result)
left_futures = remaining_futures
python/ray/experimental/serve/examples/echo_error.py
-44
View File
@@ -1,44 +0,0 @@
"""
Example of error handling mechanism in ray serve.
We are going to define a buggy function that raise some exception:
>>> def echo(_):
raise Exception("oh no")
The expected behavior is:
- HTTP server should respond with "internal error" in the response JSON
- ray.get(handle.remote()) should raise RayTaskError with traceback.
This shows that error is hidden from HTTP side but always visible when calling
from Python.
"""
import time
import requests
import ray
from ray.experimental import serve
from ray.experimental.serve.utils import pformat_color_json
def echo(_):
raise Exception("Something went wrong...")
serve.init(blocking=True)
serve.create_endpoint("my_endpoint", "/echo", blocking=True)
serve.create_backend(echo, "echo:v1")
serve.link("my_endpoint", "echo:v1")
for _ in range(2):
resp = requests.get("http://127.0.0.1:8000/echo").json()
print(pformat_color_json(resp))
print("...Sleeping for 2 seconds...")
time.sleep(2)
handle = serve.get_handle("my_endpoint")
print("Invoke from python will raise exception with traceback:")
ray.get(handle.remote())
python/ray/experimental/serve/examples/echo_fixed_packing.py
-48
View File
@@ -1,48 +0,0 @@
"""
Example showing fixed packing policy. The outputs from
v1 and v2 will be coming according to packing_num specified!
This is a packed round robin example. First batch of packing_num
(five in this example) queries would go to 'echo:v1' backend and
then next batch of packing_num queries would go to 'echo:v2'
backend.
"""
import time
import requests
from ray.experimental import serve
from ray.experimental.serve.utils import pformat_color_json
def echo_v1(_):
return "v1"
def echo_v2(_):
return "v2"
# specify the router policy as FixedPacking with packing num as 5
serve.init(
blocking=True,
queueing_policy=serve.RoutePolicy.FixedPacking,
policy_kwargs={"packing_num": 5})
# create a service
serve.create_endpoint("my_endpoint", "/echo", blocking=True)
# create first backend
serve.create_backend(echo_v1, "echo:v1")
# create second backend
serve.create_backend(echo_v2, "echo:v2")
# link and split the service to two backends
serve.split("my_endpoint", {"echo:v1": 0.5, "echo:v2": 0.5})
while True:
resp = requests.get("http://127.0.0.1:8000/echo").json()
print(pformat_color_json(resp))
print("...Sleeping for 2 seconds...")
time.sleep(2)
python/ray/experimental/serve/examples/echo_full.py
-69
View File
@@ -1,69 +0,0 @@
"""
Full example of ray.serve module
"""
import time
import requests
import ray
import ray.experimental.serve as serve
from ray.experimental.serve.utils import pformat_color_json
# initialize ray serve system.
# blocking=True will wait for HTTP server to be ready to serve request.
serve.init(blocking=True)
# an endpoint is associated with an http URL.
serve.create_endpoint("my_endpoint", "/echo")
# a backend can be a function or class.
# it can be made to be invoked from web as well as python.
def echo_v1(flask_request, response="hello from python!"):
if serve.context.web:
response = flask_request.url
return response
serve.create_backend(echo_v1, "echo:v1")
backend_config_v1 = serve.get_backend_config("echo:v1")
# We can link an endpoint to a backend, the means all the traffic
# goes to my_endpoint will now goes to echo:v1 backend.
serve.link("my_endpoint", "echo:v1")
print(requests.get("http://127.0.0.1:8000/echo", timeout=0.5).json())
# The service will be reachable from http
print(ray.get(serve.get_handle("my_endpoint").remote(response="hello")))
# as well as within the ray system.
# We can also add a new backend and split the traffic.
def echo_v2(flask_request):
# magic, only from web.
return "something new"
serve.create_backend(echo_v2, "echo:v2")
backend_config_v2 = serve.get_backend_config("echo:v2")
# The two backend will now split the traffic 50%-50%.
serve.split("my_endpoint", {"echo:v1": 0.5, "echo:v2": 0.5})
# Observe requests are now split between two backends.
for _ in range(10):
print(requests.get("http://127.0.0.1:8000/echo").json())
time.sleep(0.5)
# You can also change number of replicas
# for each backend independently.
backend_config_v1.num_replicas = 2
serve.set_backend_config("echo:v1", backend_config_v1)
backend_config_v2.num_replicas = 2
serve.set_backend_config("echo:v2", backend_config_v2)
# As well as retrieving relevant system metrics
print(pformat_color_json(serve.stat()))
python/ray/experimental/serve/examples/echo_pipeline.py
-70
View File
@@ -1,70 +0,0 @@
"""
Ray serve pipeline example
"""
import ray
import ray.experimental.serve as serve
import time
# initialize ray serve system.
# blocking=True will wait for HTTP server to be ready to serve request.
serve.init(blocking=True)
# a backend can be a function or class.
# it can be made to be invoked from web as well as python.
@serve.route("/echo_v1")
def echo_v1(_, response="hello from python!"):
return f"echo_v1({response})"
@serve.route("/echo_v2")
def echo_v2(_, relay=""):
return f"echo_v2({relay})"
@serve.route("/echo_v3")
def echo_v3(_, relay=""):
return f"echo_v3({relay})"
@serve.route("/echo_v4")
def echo_v4(_, relay1="", relay2=""):
return f"echo_v4({relay1} , {relay2})"
"""
The pipeline created is as follows -
"my_endpoint1"
/\
/ \
/ \
/ \
/ \
/ \
"my_endpoint2" "my_endpoint3"
\ /
\ /
\ /
\ /
\ /
\ /
\/
"my_endpoint4"
"""
# get the handle of the endpoints
handle1 = serve.get_handle("echo_v1")
handle2 = serve.get_handle("echo_v2")
handle3 = serve.get_handle("echo_v3")
handle4 = serve.get_handle("echo_v4")
start = time.time()
print("Start firing to the pipeline: {} s".format(time.time()))
handle1_oid = handle1.remote(response="hello")
handle4_oid = handle4.remote(
relay1=handle2.remote(relay=handle1_oid),
relay2=handle3.remote(relay=handle1_oid))
print("Firing ended now waiting for the result,"
"time taken: {} s".format(time.time() - start))
result = ray.get(handle4_oid)
print("Result: {}, time taken: {} s".format(result, time.time() - start))
python/ray/experimental/serve/examples/echo_round_robin.py
-41
View File
@@ -1,41 +0,0 @@
"""
Example showing round robin policy. The outputs from
v1 and v2 will be (almost) interleaved as queries get processed.
"""
import time
import requests
from ray.experimental import serve
from ray.experimental.serve.utils import pformat_color_json
def echo_v1(_):
return "v1"
def echo_v2(_):
return "v2"
# specify the router policy as RoundRobin
serve.init(blocking=True, queueing_policy=serve.RoutePolicy.RoundRobin)
# create a service
serve.create_endpoint("my_endpoint", "/echo", blocking=True)
# create first backend
serve.create_backend(echo_v1, "echo:v1")
# create second backend
serve.create_backend(echo_v2, "echo:v2")
# link and split the service to two backends
serve.split("my_endpoint", {"echo:v1": 0.5, "echo:v2": 0.5})
while True:
resp = requests.get("http://127.0.0.1:8000/echo").json()
print(pformat_color_json(resp))
print("...Sleeping for 2 seconds...")
time.sleep(2)
python/ray/experimental/serve/examples/echo_slo_reverse.py
-76
View File
@@ -1,76 +0,0 @@
"""
SLO [reverse] example of ray.serve module
"""
import time
import requests
import ray
import ray.experimental.serve as serve
# initialize ray serve system.
# blocking=True will wait for HTTP server to be ready to serve request.
serve.init(blocking=True)
# an endpoint is associated with an http URL.
serve.create_endpoint("my_endpoint", "/echo")
# a backend can be a function or class.
# it can be made to be invoked from web as well as python.
def echo_v1(flask_request, response="hello from python!"):
if serve.context.web:
response = flask_request.url
return response
serve.create_backend(echo_v1, "echo:v1")
serve.link("my_endpoint", "echo:v1")
# wait for routing table to get populated
time.sleep(2)
# relative slo (10 ms deadline) can be specified via http
slo_ms = 10.0
# absolute slo (10 ms deadline) can be specified via http
abs_slo_ms = 11.9
print("> [HTTP] Pinging http://127.0.0.1:8000/"
"echo?relative_slo_ms={}".format(slo_ms))
print(
requests.get("http://127.0.0.1:8000/"
"echo?relative_slo_ms={}".format(slo_ms)).json())
print("> [HTTP] Pinging http://127.0.0.1:8000/"
"echo?absolute_slo_ms={}".format(abs_slo_ms))
print(
requests.get("http://127.0.0.1:8000/"
"echo?absolute_slo_ms={}".format(abs_slo_ms)).json())
# get the handle of the endpoint
handle = serve.get_handle("my_endpoint")
future_list = []
# fire 10 requests with slo's in the (almost) reverse order of the order in
# which remote procedure call is done
for r in range(10):
slo_ms = 1000 - 100 * r
response = "hello from request: {} slo: {}".format(r, slo_ms)
print("> [REMOTE] Pinging handle.remote(response='{}',slo_ms={})".format(
response, slo_ms))
# overriding slo for each query.
# Generally slo is specified for a service handle but it can
# be overrided using options for query specific demands
f = handle.options(relative_slo_ms=slo_ms).remote(response=response)
future_list.append(f)
# get results of queries as they complete
# should be completed (almost) according to the order of their slo time
left_futures = future_list
while left_futures:
completed_futures, remaining_futures = ray.wait(left_futures, timeout=0.05)
if len(completed_futures) > 0:
result = ray.get(completed_futures[0])
print(result)
left_futures = remaining_futures
python/ray/experimental/serve/examples/echo_split.py
-41
View File
@@ -1,41 +0,0 @@
"""
Example of traffic splitting. We will first use echo:v1. Then v1 and v2
will split the incoming traffic evenly.
"""
import time
import requests
from ray.experimental import serve
from ray.experimental.serve.utils import pformat_color_json
def echo_v1(_):
return "v1"
def echo_v2(_):
return "v2"
serve.init(blocking=True)
serve.create_endpoint("my_endpoint", "/echo", blocking=True)
serve.create_backend(echo_v1, "echo:v1")
serve.link("my_endpoint", "echo:v1")
for _ in range(3):
resp = requests.get("http://127.0.0.1:8000/echo").json()
print(pformat_color_json(resp))
print("...Sleeping for 2 seconds...")
time.sleep(2)
serve.create_backend(echo_v2, "echo:v2")
serve.split("my_endpoint", {"echo:v1": 0.5, "echo:v2": 0.5})
while True:
resp = requests.get("http://127.0.0.1:8000/echo").json()
print(pformat_color_json(resp))
print("...Sleeping for 2 seconds...")
time.sleep(2)
python/ray/experimental/serve/exceptions.py
-2
View File
@@ -1,2 +0,0 @@
class RayServeException(Exception):
pass
python/ray/experimental/serve/global_state.py
-172
View File
@@ -1,172 +0,0 @@
import ray
from ray.experimental.serve.constants import (
BOOTSTRAP_KV_STORE_CONN_KEY, DEFAULT_HTTP_HOST, DEFAULT_HTTP_PORT,
SERVE_NURSERY_NAME, ASYNC_CONCURRENCY)
from ray.experimental.serve.kv_store_service import (
BackendTable, RoutingTable, TrafficPolicyTable)
from ray.experimental.serve.metric import (MetricMonitor,
start_metric_monitor_loop)
from ray.experimental.serve.policy import RoutePolicy
from ray.experimental.serve.server import HTTPActor
def start_initial_state(kv_store_connector):
nursery_handle = ActorNursery.remote()
ray.experimental.register_actor(SERVE_NURSERY_NAME, nursery_handle)
ray.get(
nursery_handle.store_bootstrap_state.remote(
BOOTSTRAP_KV_STORE_CONN_KEY, kv_store_connector))
return nursery_handle
@ray.remote
class ActorNursery:
"""Initialize and store all actor handles.
Note:
This actor is necessary because ray will destory actors when the
original actor handle goes out of scope (when driver exit). Therefore
we need to initialize and store actor handles in a seperate actor.
"""
def __init__(self):
self.tag_to_actor_handles = dict()
self.bootstrap_state = dict()
def start_actor(self,
actor_cls,
tag,
init_args=(),
init_kwargs={},
is_asyncio=False):
"""Start an actor and add it to the nursery"""
# Avoid double initialization
if tag in self.tag_to_actor_handles.keys():
return [self.tag_to_actor_handles[tag]]
max_concurrency = ASYNC_CONCURRENCY if is_asyncio else None
handle = (actor_cls.options(max_concurrency=max_concurrency).remote(
*init_args, **init_kwargs))
self.tag_to_actor_handles[tag] = handle
return [handle]
def start_actor_with_creator(self, creator, kwargs, tag):
"""
Args:
creator (Callable[Dict]): a closure that should return
a newly created actor handle when called with kwargs.
The kwargs input is passed to `ActorCls_remote` method.
"""
handle = creator(kwargs)
self.tag_to_actor_handles[tag] = handle
return [handle]
def get_all_handles(self):
return self.tag_to_actor_handles
def get_handle(self, actor_tag):
return [self.tag_to_actor_handles[actor_tag]]
def remove_handle(self, actor_tag):
if actor_tag in self.tag_to_actor_handles.keys():
self.tag_to_actor_handles.pop(actor_tag)
def store_bootstrap_state(self, key, value):
self.bootstrap_state[key] = value
def get_bootstrap_state_dict(self):
return self.bootstrap_state
class GlobalState:
"""Encapsulate all global state in the serving system.
The information is fetch lazily from
1. A collection of namespaced key value stores
2. A actor supervisor service
"""
def __init__(self, actor_nursery_handle=None):
# Get actor nursery handle
if actor_nursery_handle is None:
actor_nursery_handle = ray.experimental.get_actor(
SERVE_NURSERY_NAME)
self.actor_nursery_handle = actor_nursery_handle
# Connect to all the table
bootstrap_config = ray.get(
self.actor_nursery_handle.get_bootstrap_state_dict.remote())
kv_store_connector = bootstrap_config[BOOTSTRAP_KV_STORE_CONN_KEY]
self.route_table = RoutingTable(kv_store_connector)
self.backend_table = BackendTable(kv_store_connector)
self.policy_table = TrafficPolicyTable(kv_store_connector)
self.refresh_actor_handle_cache()
def refresh_actor_handle_cache(self):
self.actor_handle_cache = ray.get(
self.actor_nursery_handle.get_all_handles.remote())
def init_or_get_http_server(self,
host=DEFAULT_HTTP_HOST,
port=DEFAULT_HTTP_PORT):
if "http_server" not in self.actor_handle_cache:
[handle] = ray.get(
self.actor_nursery_handle.start_actor.remote(
HTTPActor, tag="http_server"))
handle.run.remote(host=host, port=port)
self.refresh_actor_handle_cache()
return self.actor_handle_cache["http_server"]
def _get_queueing_policy(self, default_policy):
return_policy = default_policy
# check if there is already a queue_actor running
# with policy as p.name for the case where
# serve nursery exists: ray.experimental.get_actor(SERVE_NURSERY_NAME)
for p in RoutePolicy:
queue_actor_tag = "queue_actor::" + p.name
if queue_actor_tag in self.actor_handle_cache:
return_policy = p
break
return return_policy
def init_or_get_router(self,
queueing_policy=RoutePolicy.Random,
policy_kwargs={}):
# get queueing policy
self.queueing_policy = self._get_queueing_policy(
default_policy=queueing_policy)
queue_actor_tag = "queue_actor::" + self.queueing_policy.name
if queue_actor_tag not in self.actor_handle_cache:
[handle] = ray.get(
self.actor_nursery_handle.start_actor.remote(
self.queueing_policy.value,
init_kwargs=policy_kwargs,
tag=queue_actor_tag,
is_asyncio=True))
# handle.register_self_handle.remote(handle)
self.refresh_actor_handle_cache()
return self.actor_handle_cache[queue_actor_tag]
def init_or_get_metric_monitor(self, gc_window_seconds=3600):
if "metric_monitor" not in self.actor_handle_cache:
[handle] = ray.get(
self.actor_nursery_handle.start_actor.remote(
MetricMonitor,
init_args=(gc_window_seconds, ),
tag="metric_monitor"))
start_metric_monitor_loop.remote(handle)
if "queue_actor" in self.actor_handle_cache:
handle.add_target.remote(
self.actor_handle_cache["queue_actor"])
self.refresh_actor_handle_cache()
return self.actor_handle_cache["metric_monitor"]
python/ray/experimental/serve/handle.py
-106
View File
@@ -1,106 +0,0 @@
from ray.experimental import serve
from ray.experimental.serve.context import TaskContext
from ray.experimental.serve.exceptions import RayServeException
from ray.experimental.serve.constants import DEFAULT_HTTP_ADDRESS
from ray.experimental.serve.request_params import RequestMetadata
class RayServeHandle:
"""A handle to a service endpoint.
Invoking this endpoint with .remote is equivalent to pinging
an HTTP endpoint.
Example:
>>> handle = serve.get_handle("my_endpoint")
>>> handle
RayServeHandle(
Endpoint="my_endpoint",
URL="...",
Traffic=...
)
>>> handle.remote(my_request_content)
ObjectID(...)
>>> ray.get(handle.remote(...))
# result
>>> ray.get(handle.remote(let_it_crash_request))
# raises RayTaskError Exception
"""
def __init__(self,
router_handle,
endpoint_name,
relative_slo_ms=None,
absolute_slo_ms=None):
self.router_handle = router_handle
self.endpoint_name = endpoint_name
assert (relative_slo_ms is None
or absolute_slo_ms is None), ("Can't specify both "
"relative and absolute "
"slo's together!")
self.relative_slo_ms = self._check_slo_ms(relative_slo_ms)
self.absolute_slo_ms = self._check_slo_ms(absolute_slo_ms)
def _check_slo_ms(self, slo_value):
if slo_value is not None:
try:
slo_value = float(slo_value)
if slo_value < 0:
raise ValueError(
"Request SLO must be positive, it is {}".format(
slo_value))
return slo_value
except ValueError as e:
raise RayServeException(str(e))
return None
def remote(self, *args, **kwargs):
if len(args) != 0:
raise RayServeException(
"handle.remote must be invoked with keyword arguments.")
# create RequestMetadata instance
request_in_object = RequestMetadata(
self.endpoint_name, TaskContext.Python, self.relative_slo_ms,
self.absolute_slo_ms)
return self.router_handle.enqueue_request.remote(
request_in_object, **kwargs)
def options(self, relative_slo_ms=None, absolute_slo_ms=None):
# If both the slo's are None then then we use a high default
# value so other queries can be prioritize and put in front of these
# queries.
assert (relative_slo_ms is None
or absolute_slo_ms is None), ("Can't specify both "
"relative and absolute "
"slo's together!")
return RayServeHandle(self.router_handle, self.endpoint_name,
relative_slo_ms, absolute_slo_ms)
def get_traffic_policy(self):
# TODO(simon): This method is implemented via checking global state
# because we are sure handle and global_state are in the same process.
# However, once global_state is deprecated, this method need to be
# updated accordingly.
history = serve.global_state.policy_action_history[self.endpoint_name]
if len(history):
return history[-1]
else:
return None
def get_http_endpoint(self):
return DEFAULT_HTTP_ADDRESS
def __repr__(self):
return """
RayServeHandle(
Endpoint="{endpoint_name}",
URL="{http_endpoint}/{endpoint_name}",
Traffic={traffic_policy}
)
""".format(endpoint_name=self.endpoint_name,
http_endpoint=self.get_http_endpoint(),
traffic_policy=self.get_traffic_policy())
# TODO(simon): a convenience function that dumps equivalent requests
# code for a given call.
python/ray/experimental/serve/http_util.py
-69
View File
@@ -1,69 +0,0 @@
import io
import flask
def build_flask_request(asgi_scope_dict, request_body):
"""Build and return a flask request from ASGI payload
This function is indented to be used immediately before task invocation
happen.
"""
wsgi_environ = build_wsgi_environ(asgi_scope_dict, request_body)
return flask.Request(wsgi_environ)
def build_wsgi_environ(scope, body):
"""
Builds a scope and request body into a WSGI environ object.
This code snippet is taken from https://github.com/django/asgiref/blob
/36c3e8dc70bf38fe2db87ac20b514f21aaf5ea9d/asgiref/wsgi.py#L52
WSGI specification can be found at
https://www.python.org/dev/peps/pep-0333/
This function helps translate ASGI scope and body into a flask request.
"""
environ = {
"REQUEST_METHOD": scope["method"],
"SCRIPT_NAME": scope.get("root_path", ""),
"PATH_INFO": scope["path"],
"QUERY_STRING": scope["query_string"].decode("ascii"),
"SERVER_PROTOCOL": "HTTP/{}".format(scope["http_version"]),
"wsgi.version": (1, 0),
"wsgi.url_scheme": scope.get("scheme", "http"),
"wsgi.input": body,
"wsgi.errors": io.BytesIO(),
"wsgi.multithread": True,
"wsgi.multiprocess": True,
"wsgi.run_once": False,
}
# Get server name and port - required in WSGI, not in ASGI
environ["SERVER_NAME"] = scope["server"][0]
environ["SERVER_PORT"] = str(scope["server"][1])
environ["REMOTE_ADDR"] = scope["client"][0]
# Transforms headers into environ entries.
for name, value in scope.get("headers", []):
# name, values are both bytes, we need to decode them to string
name = name.decode("latin1")
value = value.decode("latin1")
# Handle name correction to conform to WSGI spec
# https://www.python.org/dev/peps/pep-0333/#environ-variables
if name == "content-length":
corrected_name = "CONTENT_LENGTH"
elif name == "content-type":
corrected_name = "CONTENT_TYPE"
else:
corrected_name = "HTTP_%s" % name.upper().replace("-", "_")
# If the header value repeated,
# we will just concatenate it to the field.
if corrected_name in environ:
value = environ[corrected_name] + "," + value
environ[corrected_name] = value
return environ
python/ray/experimental/serve/kv_store_service.py
-284
View File
@@ -1,284 +0,0 @@
import json
import sqlite3
from abc import ABC
from ray import cloudpickle as pickle
import ray.experimental.internal_kv as ray_kv
from ray.experimental.serve.utils import logger
from typing import Union
from ray.experimental.serve.constants import NO_ROUTE_KEY
class NamespacedKVStore(ABC):
"""Abstract base class for a namespaced key-value store.
The idea is that multiple key-value stores can be created while sharing
the same storage system. The keys of each instance are namespaced to avoid
object_id key collision.
Example:
>>> store_ns1 = NamespacedKVStore(namespace="ns1")
>>> store_ns2 = NamespacedKVStore(namespace="ns2")
# Two stores can share the same connection like Redis or SQL Table
>>> store_ns1.put("same-key", 1)
>>> store_ns1.get("same-key")
1
>>> store_ns2.put("same-key", 2)
>>> store_ns2.get("same-key", 2)
2
"""
def __init__(self, namespace):
raise NotImplementedError()
def get(self, key):
"""Retrieve the value for the given key.
Args:
key (str)
"""
raise NotImplementedError()
def put(self, key, value):
"""Serialize the value and store it under the given key.
Args:
key (str)
value (object): any serializable object. The serialization method
is determined by the subclass implementation.
"""
raise NotImplementedError()
def as_dict(self):
"""Return the entire namespace as a dictionary.
Returns:
data (dict): key value pairs in current namespace
"""
raise NotImplementedError()
class InMemoryKVStore(NamespacedKVStore):
"""A reference implementation used for testing."""
def __init__(self, namespace):
self.data = dict()
# Namespace is ignored, because each namespace is backed by
# an in-memory Python dictionary.
self.namespace = namespace
def get(self, key):
return self.data[key]
def put(self, key, value):
self.data[key] = value
def as_dict(self):
return self.data.copy()
class RayInternalKVStore(NamespacedKVStore):
"""A NamespacedKVStore implementation using ray's `internal_kv`."""
def __init__(self, namespace):
assert ray_kv._internal_kv_initialized()
self.index_key = "RAY_SERVE_INDEX"
self.namespace = namespace
self._put(self.index_key, [])
def _format_key(self, key):
return "{ns}-{key}".format(ns=self.namespace, key=key)
def _remove_format_key(self, formatted_key):
return formatted_key.replace(self.namespace + "-", "", 1)
def _serialize(self, obj):
return json.dumps(obj)
def _deserialize(self, buffer):
return json.loads(buffer)
def _put(self, key, value):
ray_kv._internal_kv_put(
self._format_key(self._serialize(key)),
self._serialize(value),
overwrite=True,
)
def _get(self, key):
return self._deserialize(
ray_kv._internal_kv_get(self._format_key(self._serialize(key))))
def get(self, key):
return self._get(key)
def put(self, key, value):
assert isinstance(key, str), "Key must be a string."
self._put(key, value)
all_keys = set(self._get(self.index_key))
all_keys.add(key)
self._put(self.index_key, list(all_keys))
def as_dict(self):
data = {}
all_keys = self._get(self.index_key)
for key in all_keys:
data[self._remove_format_key(key)] = self._get(key)
return data
class SQLiteKVStore(NamespacedKVStore):
def __init__(self, namespace, db_path):
self.namespace = namespace
self.conn = sqlite3.connect(db_path)
cursor = self.conn.cursor()
cursor.execute(
"CREATE TABLE IF NOT EXISTS {} (key TEXT UNIQUE, value TEXT)".
format(self.namespace))
self.conn.commit()
def put(self, key, value):
cursor = self.conn.cursor()
cursor.execute(
"INSERT OR REPLACE INTO {} (key, value) VALUES (?,?)".format(
self.namespace), (key, value))
self.conn.commit()
def get(self, key, default=None):
cursor = self.conn.cursor()
result = list(
cursor.execute(
"SELECT value FROM {} WHERE key = (?)".format(self.namespace),
(key, )))
if len(result) == 0:
return default
else:
# Due to UNIQUE constraint, there can only be one value.
value, *_ = result[0]
return value
def as_dict(self):
cursor = self.conn.cursor()
result = list(
cursor.execute("SELECT key, value FROM {}".format(self.namespace)))
return dict(result)
# Tables
class RoutingTable:
def __init__(self, kv_connector):
self.routing_table = kv_connector("routing_table")
self.request_count = 0
def register_service(self, route: Union[str, None], service: str):
"""Create an entry in the routing table
Args:
route: http path name. Must begin with '/'.
service: service name. This is the name http actor will push
the request to.
"""
logger.debug("[KV] Registering route {} to service {}.".format(
route, service))
# put no route services in default key
if route is None:
no_http_services = json.loads(
self.routing_table.get(NO_ROUTE_KEY, "[]"))
no_http_services.append(service)
self.routing_table.put(NO_ROUTE_KEY, json.dumps(no_http_services))
else:
self.routing_table.put(route, service)
def list_service(self, include_headless=False):
"""Returns the routing table.
Args:
include_headless: If True, returns a no route services (headless)
services with normal services. (Default: False)
"""
table = self.routing_table.as_dict()
if include_headless:
table[NO_ROUTE_KEY] = json.loads(table.get(NO_ROUTE_KEY, "[]"))
else:
table.pop(NO_ROUTE_KEY, None)
return table
def get_request_count(self):
"""Return the number of requests that fetched the routing table.
This method is used for two purpose:
1. Make sure HTTP server has started and healthy. Incremented request
count means HTTP server is actively fetching routing table.
2. Make sure HTTP server does not have stale routing table. This number
should be incremented every HTTP_ROUTER_CHECKER_INTERVAL_S seconds.
Supervisor should check this number as indirect indicator of http
server's health.
"""
return self.request_count
class BackendTable:
def __init__(self, kv_connector):
self.backend_table = kv_connector("backend_creator")
self.replica_table = kv_connector("replica_table")
self.backend_info = kv_connector("backend_info")
self.backend_init_args = kv_connector("backend_init_args")
def register_backend(self, backend_tag: str, backend_creator):
backend_creator_serialized = pickle.dumps(backend_creator)
self.backend_table.put(backend_tag, backend_creator_serialized)
def save_init_args(self, backend_tag: str, arg_list):
serialized_arg_list = pickle.dumps(arg_list)
self.backend_init_args.put(backend_tag, serialized_arg_list)
def get_init_args(self, backend_tag):
return pickle.loads(self.backend_init_args.get(backend_tag))
def register_info(self, backend_tag: str, backend_info_d):
self.backend_info.put(backend_tag, json.dumps(backend_info_d))
def get_info(self, backend_tag):
return json.loads(self.backend_info.get(backend_tag, "{}"))
def get_backend_creator(self, backend_tag):
return pickle.loads(self.backend_table.get(backend_tag))
def list_backends(self):
return list(self.backend_table.as_dict().keys())
def list_replicas(self, backend_tag: str):
return json.loads(self.replica_table.get(backend_tag, "[]"))
def add_replica(self, backend_tag: str, new_replica_tag: str):
replica_tags = self.list_replicas(backend_tag)
replica_tags.append(new_replica_tag)
self.replica_table.put(backend_tag, json.dumps(replica_tags))
def remove_replica(self, backend_tag):
replica_tags = self.list_replicas(backend_tag)
removed_replica = replica_tags.pop()
self.replica_table.put(backend_tag, json.dumps(replica_tags))
return removed_replica
class TrafficPolicyTable:
def __init__(self, kv_connector):
self.traffic_policy_table = kv_connector("traffic_policy")
def register_traffic_policy(self, service_name, policy_dict):
self.traffic_policy_table.put(service_name, json.dumps(policy_dict))
def list_traffic_policy(self):
return {
service: json.loads(policy)
for service, policy in self.traffic_policy_table.as_dict()
}
python/ray/experimental/serve/metric.py
-157
View File
@@ -1,157 +0,0 @@
import time
import numpy as np
import pandas as pd
import ray
@ray.remote(num_cpus=0)
class MetricMonitor:
def __init__(self, gc_window_seconds=3600):
"""Metric monitor scrapes metrics from ray serve actors
and allow windowed query operations.
Args:
gc_window_seconds(int): How long will we keep the metric data in
memory. Data older than the gc_window will be deleted.
"""
#: Mapping actor ID (hex) -> actor handle
self.actor_handles = dict()
self.data_entries = []
self.gc_window_seconds = gc_window_seconds
self.latest_gc_time = time.time()
def is_ready(self):
return True
def add_target(self, target_handle):
hex_id = target_handle._actor_id.hex()
self.actor_handles[hex_id] = target_handle
def remove_target(self, target_handle):
hex_id = target_handle._actor_id.hex()
self.actor_handles.pop(hex_id)
def scrape(self):
# If expected gc time has passed, we will perform metric value GC.
expected_gc_time = self.latest_gc_time + self.gc_window_seconds
if expected_gc_time < time.time():
self._perform_gc()
self.latest_gc_time = time.time()
curr_time = time.time()
result = [
handle._serve_metric.remote()
for handle in self.actor_handles.values()
]
# TODO(simon): handle the possibility that an actor_handle is removed
for handle_result in ray.get(result):
for metric_name, metric_info in handle_result.items():
data_entry = {
"retrieved_at": curr_time,
"name": metric_name,
"type": metric_info["type"],
}
if metric_info["type"] == "counter":
data_entry["value"] = metric_info["value"]
self.data_entries.append(data_entry)
elif metric_info["type"] == "list":
for metric_value in metric_info["value"]:
new_entry = data_entry.copy()
new_entry["value"] = metric_value
self.data_entries.append(new_entry)
def _perform_gc(self):
curr_time = time.time()
earliest_time_allowed = curr_time - self.gc_window_seconds
# If we don"t have any data at hand, no need to gc.
if len(self.data_entries) == 0:
return
df = pd.DataFrame(self.data_entries)
df = df[df["retrieved_at"] >= earliest_time_allowed]
self.data_entries = df.to_dict(orient="record")
def _get_dataframe(self):
return pd.DataFrame(self.data_entries)
def collect(self,
percentiles=[50, 90, 95],
agg_windows_seconds=[10, 60, 300, 600, 3600]):
"""Collect and perform aggregation on all metrics.
Args:
percentiles(List[int]): The percentiles for aggregation operations.
Default is 50th, 90th, 95th percentile.
agg_windows_seconds(List[int]): The aggregation windows in seconds.
The longest aggregation window must be shorter or equal to the
gc_window_seconds.
"""
result = {}
df = pd.DataFrame(self.data_entries)
if len(df) == 0: # no metric to report
return {}
# Retrieve the {metric_name -> metric_type} mapping
metric_types = df[["name",
"type"]].set_index("name").squeeze().to_dict()
for metric_name, metric_type in metric_types.items():
if metric_type == "counter":
result[metric_name] = df.loc[df["name"] == metric_name,
"value"].tolist()[-1]
if metric_type == "list":
result.update(
self._aggregate(metric_name, percentiles,
agg_windows_seconds))
return result
def _aggregate(self, metric_name, percentiles, agg_windows_seconds):
"""Perform aggregation over a metric.
Note:
This metric must have type `list`.
"""
assert max(agg_windows_seconds) <= self.gc_window_seconds, (
"Aggregation window exceeds gc window. You should set a longer gc "
"window or shorter aggregation window.")
curr_time = time.time()
df = pd.DataFrame(self.data_entries)
filtered_df = df[df["name"] == metric_name]
if len(filtered_df) == 0:
return dict()
data_types = filtered_df["type"].unique().tolist()
assert data_types == [
"list"
], ("Can't aggreagte over non-list type. {} has type {}".format(
metric_name, data_types))
aggregated_metric = {}
for window in agg_windows_seconds:
earliest_time = curr_time - window
windowed_df = filtered_df[
filtered_df["retrieved_at"] > earliest_time]
percentile_values = np.percentile(windowed_df["value"],
percentiles)
for percentile, value in zip(percentiles, percentile_values):
result_key = "{name}_{perc}th_perc_{window}_window".format(
name=metric_name, perc=percentile, window=window)
aggregated_metric[result_key] = value
return aggregated_metric
@ray.remote(num_cpus=0)
def start_metric_monitor_loop(monitor_handle, duration_s=5):
while True:
ray.get(monitor_handle.scrape.remote())
time.sleep(duration_s)
python/ray/experimental/serve/policy.py
-188
View File
@@ -1,188 +0,0 @@
from enum import Enum
import itertools
import numpy as np
import ray
from ray.experimental.serve.queues import (CentralizedQueues)
from ray.experimental.serve.utils import logger
class RandomPolicyQueue(CentralizedQueues):
"""
A wrapper class for Random policy.This backend selection policy is
`Stateless` meaning the current decisions of selecting backend are
not dependent on previous decisions. Random policy (randomly) samples
backends based on backend weights for every query. This policy uses the
weights assigned to backends.
"""
async def _flush_service_queues(self):
# perform traffic splitting for requests
for service, queue in self.service_queues.items():
# while there are incoming requests and there are backends
while queue.qsize() and len(self.traffic[service]):
backend_names = list(self.traffic[service].keys())
backend_weights = list(self.traffic[service].values())
# randomly choose a backend for every query
chosen_backend = np.random.choice(
backend_names, replace=False, p=backend_weights).squeeze()
logger.debug("Matching service {} to backend {}".format(
service, chosen_backend))
request = await queue.get()
self.buffer_queues[chosen_backend].add(request)
@ray.remote
class RandomPolicyQueueActor(RandomPolicyQueue):
pass
class RoundRobinPolicyQueue(CentralizedQueues):
"""
A wrapper class for RoundRobin policy. This backend selection policy
is `Stateful` meaning the current decisions of selecting backend are
dependent on previous decisions. RoundRobinPolicy assigns queries in
an interleaved manner to every backend serving for a service. Consider
backend A,B linked to a service. Now queries will be assigned to backends
in the following order - [ A, B, A, B ... ] . This policy doesn't use the
weights assigned to backends.
"""
# Saves the information about last assigned
# backend for every service
round_robin_iterator_map = {}
async def set_traffic(self, service, traffic_dict):
logger.debug("Setting traffic for service %s to %s", service,
traffic_dict)
self.traffic[service] = traffic_dict
backend_names = list(self.traffic[service].keys())
self.round_robin_iterator_map[service] = itertools.cycle(backend_names)
await self.flush()
async def _flush_service_queues(self):
# perform traffic splitting for requests
for service, queue in self.service_queues.items():
# if there are incoming requests and there are backends
if queue.qsize() and len(self.traffic[service]):
while queue.qsize():
# choose the next backend available from persistent
# information
chosen_backend = next(
self.round_robin_iterator_map[service])
request = await queue.get()
self.buffer_queues[chosen_backend].add(request)
@ray.remote
class RoundRobinPolicyQueueActor(RoundRobinPolicyQueue):
pass
class PowerOfTwoPolicyQueue(CentralizedQueues):
"""
A wrapper class for powerOfTwo policy. This backend selection policy is
`Stateless` meaning the current decisions of selecting backend are
dependent on previous decisions. PowerOfTwo policy (randomly) samples two
backends (say Backend A,B among A,B,C) based on the backend weights
specified and chooses the backend which is less loaded. This policy uses
the weights assigned to backends.
"""
async def _flush_service_queues(self):
# perform traffic splitting for requests
for service, queue in self.service_queues.items():
# while there are incoming requests and there are backends
while queue.qsize() and len(self.traffic[service]):
backend_names = list(self.traffic[service].keys())
backend_weights = list(self.traffic[service].values())
if len(self.traffic[service]) >= 2:
# randomly pick 2 backends
backend1, backend2 = np.random.choice(
backend_names, 2, replace=False, p=backend_weights)
# see the length of buffer queues of the two backends
# and pick the one which has less no. of queries
# in the buffer
if (len(self.buffer_queues[backend1]) <= len(
self.buffer_queues[backend2])):
chosen_backend = backend1
else:
chosen_backend = backend2
logger.debug("[Power of two chocies] found two backends "
"{} and {}: choosing {}.".format(
backend1, backend2, chosen_backend))
else:
chosen_backend = np.random.choice(
backend_names, replace=False,
p=backend_weights).squeeze()
request = await queue.get()
self.buffer_queues[chosen_backend].add(request)
@ray.remote
class PowerOfTwoPolicyQueueActor(PowerOfTwoPolicyQueue):
pass
class FixedPackingPolicyQueue(CentralizedQueues):
"""
A wrapper class for FixedPacking policy. This backend selection policy is
`Stateful` meaning the current decisions of selecting backend are dependent
on previous decisions. FixedPackingPolicy is k RoundRobin policy where
first packing_num queries are handled by 'backend-1' and next k queries are
handled by 'backend-2' and so on ... where 'backend-1' and 'backend-2' are
served by the same service. This policy doesn't use the weights assigned to
backends.
"""
def __init__(self, packing_num=3):
# Saves the information about last assigned
# backend for every service
self.fixed_packing_iterator_map = {}
self.packing_num = packing_num
super().__init__()
async def set_traffic(self, service, traffic_dict):
logger.debug("Setting traffic for service %s to %s", service,
traffic_dict)
self.traffic[service] = traffic_dict
backend_names = list(self.traffic[service].keys())
self.fixed_packing_iterator_map[service] = itertools.cycle(
itertools.chain.from_iterable(
itertools.repeat(x, self.packing_num) for x in backend_names))
await self.flush()
async def _flush_service_queues(self):
# perform traffic splitting for requests
for service, queue in self.service_queues.items():
# if there are incoming requests and there are backends
if queue.qsize() and len(self.traffic[service]):
while queue.qsize():
# choose the next backend available from persistent
# information
chosen_backend = next(
self.fixed_packing_iterator_map[service])
request = await queue.get()
self.buffer_queues[chosen_backend].add(request)
@ray.remote
class FixedPackingPolicyQueueActor(FixedPackingPolicyQueue):
pass
class RoutePolicy(Enum):
"""
A class for registering the backend selection policy.
Add a name and the corresponding class.
Serve will support the added policy and policy can be accessed
in `serve.init` method through name provided here.
"""
Random = RandomPolicyQueueActor
RoundRobin = RoundRobinPolicyQueueActor
PowerOfTwo = PowerOfTwoPolicyQueueActor
FixedPacking = FixedPackingPolicyQueueActor
python/ray/experimental/serve/queues.py
-305
View File
@@ -1,305 +0,0 @@
import asyncio
import copy
from collections import defaultdict
from typing import DefaultDict, List
import pickle
# Note on choosing blist instead of stdlib heapq
# 1. pop operation should be O(1) (amortized)
# (helpful even for batched pop)
# 2. There should not be significant overhead in
# maintaining the sorted list.
# 3. The blist implementation is fast and uses C extensions.
import blist
import ray
from ray.experimental.serve.utils import logger
class Query:
def __init__(self, request_args, request_kwargs, request_context,
request_slo_ms):
self.request_args = request_args
self.request_kwargs = request_kwargs
self.request_context = request_context
self.async_future = asyncio.get_event_loop().create_future()
# Service level objective in milliseconds. This is expected to be the
# absolute time since unix epoch.
self.request_slo_ms = request_slo_ms
def ray_serialize(self):
# NOTE: this method is needed because Query need to be serialized and
# sent to the replica worker. However, after we send the query to
# replica worker the async_future is still needed to retrieve the final
# result. Therefore we need a way to pass the information to replica
# worker without removing async_future.
clone = copy.copy(self)
clone.async_future = None
# We can't use cloudpickle due to a recursion issue
return pickle.dumps(clone)
@staticmethod
def ray_deserialize(value):
return pickle.loads(value)
# adding comparator fn for maintaining an
# ascending order sorted list w.r.t request_slo_ms
def __lt__(self, other):
return self.request_slo_ms < other.request_slo_ms
def __repr__(self):
return "<Query args={} kwargs={}>".format(self.request_args,
self.request_kwargs)
def _make_future_unwrapper(client_futures: List[asyncio.Future],
host_future: asyncio.Future):
"""Distribute the result of host_future to each of client_future"""
for client_future in client_futures:
# Keep a reference to host future so the host future won't get
# garbage collected.
client_future.host_ref = host_future
def unwrap_future(_):
result = host_future.result()
if isinstance(result, list):
for client_future, result_item in zip(client_futures, result):
client_future.set_result(result_item)
else: # Result is an exception.
for client_future in client_futures:
client_future.set_result(result)
return unwrap_future
class CentralizedQueues:
"""A router that routes request to available workers.
Router aceepts each request from the `enqueue_request` method and enqueues
it. It also accepts worker request to work (called work_intention in code)
from workers via the `dequeue_request` method. The traffic policy is used
to match requests with their corresponding workers.
Behavior:
>>> # psuedo-code
>>> queue = CentralizedQueues()
>>> queue.enqueue_request(
"service-name", request_args, request_kwargs, request_context)
# nothing happens, request is queued.
# returns result ObjectID, which will contains the final result
>>> queue.dequeue_request('backend-1', replica_handle)
# nothing happens, work intention is queued.
# return work ObjectID, which will contains the future request payload
>>> queue.link('service-name', 'backend-1')
# here the enqueue_requester is matched with replica, request
# data is put into work ObjectID, and the replica processes the request
# and store the result into result ObjectID
Traffic policy splits the traffic among different replicas
probabilistically:
1. When all backends are ready to receive traffic, we will randomly
choose a backend based on the weights assigned by the traffic policy
dictionary.
2. When more than 1 but not all backends are ready, we will normalize the
weights of the ready backends to 1 and choose a backend via sampling.
3. When there is only 1 backend ready, we will only use that backend.
"""
def __init__(self):
# Note: Several queues are used in the router
# - When a request come in, it's placed inside its corresponding
# service_queue.
# - The service_queue is dequed during flush operation, which moves
# the queries to backend buffer_queue. Here we match a request
# for a service to a backend given some policy.
# - The worker_queue is used to collect idle actor handle. These
# handles are dequed during the second stage of flush operation,
# which assign queries in buffer_queue to actor handle.
# -- Queues -- #
# service_name -> request queue
self.service_queues: DefaultDict[asyncio.Queue[Query]] = defaultdict(
asyncio.Queue)
# backend_name -> worker request queue
self.worker_queues: DefaultDict[asyncio.Queue[
ray.actor.ActorHandle]] = defaultdict(asyncio.Queue)
# backend_name -> worker payload queue
self.buffer_queues = defaultdict(blist.sortedlist)
# -- Metadata -- #
# service_name -> traffic_policy
self.traffic = defaultdict(dict)
# backend_name -> backend_config
self.backend_info = dict()
# -- Synchronization -- #
# This lock guarantee that only one flush operation can happen at a
# time. Without the lock, multiple flush operation can pop from the
# same buffer_queue and worker_queue and create deadlock. For example,
# an operation holding the only query and the other flush operation
# holding the only idle replica. Additionally, allowing only one flush
# operation at a time simplifies design overhead for custom queuing and
# batching polcies.
self.flush_lock = asyncio.Lock()
def is_ready(self):
return True
def _serve_metric(self):
return {
"backend_{}_queue_size".format(backend_name): {
"value": len(queue),
"type": "counter",
}
for backend_name, queue in self.buffer_queues.items()
}
async def enqueue_request(self, request_in_object, *request_args,
**request_kwargs):
service = request_in_object.service
logger.debug("Received a request for service {}".format(service))
# check if the slo specified is directly the
# wall clock time
if request_in_object.absolute_slo_ms is not None:
request_slo_ms = request_in_object.absolute_slo_ms
else:
request_slo_ms = request_in_object.adjust_relative_slo_ms()
request_context = request_in_object.request_context
query = Query(request_args, request_kwargs, request_context,
request_slo_ms)
await self.service_queues[service].put(query)
await self.flush()
# Note: a future change can be to directly return the ObjectID from
# replica task submission
result = await query.async_future
return result
async def dequeue_request(self, backend, replica_handle):
logger.debug(
"Received a dequeue request for backend {}".format(backend))
await self.worker_queues[backend].put(replica_handle)
await self.flush()
async def remove_and_destory_replica(self, backend, replica_handle):
# We need this lock because we modify worker_queue here.
async with self.flush_lock:
old_queue = self.worker_queues[backend]
new_queue = asyncio.Queue()
target_id = replica_handle._actor_id
while not old_queue.empty():
replica_handle = await old_queue.get()
if replica_handle._actor_id != target_id:
await new_queue.put(replica_handle)
self.worker_queues[backend] = new_queue
# TODO: consider await this with timeout, or use ray_kill
replica_handle.__ray_terminate__.remote()
async def link(self, service, backend):
logger.debug("Link %s with %s", service, backend)
await self.set_traffic(service, {backend: 1.0})
async def set_traffic(self, service, traffic_dict):
logger.debug("Setting traffic for service %s to %s", service,
traffic_dict)
self.traffic[service] = traffic_dict
await self.flush()
async def set_backend_config(self, backend, config_dict):
logger.debug("Setting backend config for "
"backend {} to {}".format(backend, config_dict))
self.backend_info[backend] = config_dict
async def flush(self):
"""In the default case, flush calls ._flush.
When this class is a Ray actor, .flush can be scheduled as a remote
method invocation.
"""
async with self.flush_lock:
await self._flush_service_queues()
await self._flush_buffer_queues()
def _get_available_backends(self, service):
backends_in_policy = set(self.traffic[service].keys())
available_workers = {
backend
for backend, queues in self.worker_queues.items()
if queues.qsize() > 0
}
return list(backends_in_policy.intersection(available_workers))
async def _flush_service_queues(self):
"""Selects the backend and puts the service queue query to the buffer
Expected Implementation:
The implementer is expected to access and manipulate
self.service_queues : dict[str,Deque]
self.buffer_queues : dict[str,sortedlist]
For registering the implemented policies register at policy.py
Expected Behavior:
the Deque of all services in self.service_queues linked with
atleast one backend must be empty irrespective of whatever
backend policy is implemented.
"""
raise NotImplementedError(
"This method should be implemented by child class.")
# flushes the buffer queue and assigns work to workers
async def _flush_buffer_queues(self):
for service in self.traffic.keys():
ready_backends = self._get_available_backends(service)
for backend in ready_backends:
# no work available
if len(self.buffer_queues[backend]) == 0:
continue
buffer_queue = self.buffer_queues[backend]
worker_queue = self.worker_queues[backend]
logger.debug("Assigning queries for backend {} with buffer "
"queue size {} and worker queue size {}".format(
backend, len(buffer_queue),
worker_queue.qsize()))
max_batch_size = None
if backend in self.backend_info:
max_batch_size = self.backend_info[backend][
"max_batch_size"]
await self._assign_query_to_worker(buffer_queue, worker_queue,
max_batch_size)
async def _assign_query_to_worker(self,
buffer_queue,
worker_queue,
max_batch_size=None):
while len(buffer_queue) and worker_queue.qsize():
worker = await worker_queue.get()
if max_batch_size is None: # No batching
request = buffer_queue.pop(0)
future = worker._ray_serve_call.remote(request).as_future()
# chaining satisfies request.async_future with future result.
asyncio.futures._chain_future(future, request.async_future)
else:
real_batch_size = min(len(buffer_queue), max_batch_size)
requests = [
buffer_queue.pop(0) for _ in range(real_batch_size)
]
future = worker._ray_serve_call.remote(requests).as_future()
future.add_done_callback(
_make_future_unwrapper(
client_futures=[req.async_future for req in requests],
host_future=future))
python/ray/experimental/serve/request_params.py
-37
View File
@@ -1,37 +0,0 @@
import time
from ray.experimental.serve.constants import DEFAULT_LATENCY_SLO_MS
class RequestMetadata:
"""
Request Arguments required for enqueuing a request to the service
queue.
Args:
service(str): A registered service endpoint.
request_context(TaskContext): Context of a request.
request_slo_ms(float): Expected time for the query to get
completed.
is_wall_clock_time(bool): if True, router won't add wall clock
time to `request_slo_ms`.
"""
def __init__(self,
service,
request_context,
relative_slo_ms=None,
absolute_slo_ms=None):
self.service = service
self.request_context = request_context
self.relative_slo_ms = relative_slo_ms
self.absolute_slo_ms = absolute_slo_ms
def adjust_relative_slo_ms(self) -> float:
"""Normalize the input latency objective to absoluate timestamp.
"""
slo_ms = self.relative_slo_ms
if slo_ms is None:
slo_ms = DEFAULT_LATENCY_SLO_MS
current_time_ms = time.time() * 1000
return current_time_ms + slo_ms
python/ray/experimental/serve/scripts.py
-50
View File
@@ -1,50 +0,0 @@
import json
import click
import ray
import ray.experimental.serve as serve
@click.group("serve", help="Commands working with ray serve")
def serve_cli():
pass
@serve_cli.command(help="Initialize ray serve components")
def init():
ray.init(address="auto")
serve.init(blocking=True)
@serve_cli.command(help="Split traffic for a endpoint")
@click.argument("endpoint", required=True, type=str)
# TODO(simon): Make traffic dictionary more ergonomic. e.g.
# --traffic backend1=0.5 --traffic backend2=0.5
@click.option(
"--traffic",
required=True,
type=str,
help="Traffic dictionary in JSON format")
def split(endpoint, traffic):
ray.init(address="auto")
serve.init()
serve.split(endpoint, json.loads(traffic))
@serve_cli.command(help="Scale the number of replicas for a backend")
@click.argument("backend", required=True, type=str)
@click.option(
"--num-replicas",
required=True,
type=int,
help="New number of replicas to set")
def scale(backend_tag, num_replicas):
if num_replicas <= 0:
click.Abort(
"Cannot set number of replicas to be smaller or equal to 0.")
ray.init(address="auto")
serve.init()
serve.scale(backend_tag, num_replicas)
python/ray/experimental/serve/server.py
-196
View File
@@ -1,196 +0,0 @@
import asyncio
import json
import uvicorn
import ray
from ray.experimental.async_api import _async_init
from ray.experimental.serve.constants import HTTP_ROUTER_CHECKER_INTERVAL_S
from ray.experimental.serve.context import TaskContext
from ray.experimental.serve.utils import BytesEncoder
from ray.experimental.serve.request_params import RequestMetadata
from urllib.parse import parse_qs
class JSONResponse:
"""ASGI compliant response class.
It is expected to be called in async context and pass along
`scope, receive, send` as in ASGI spec.
>>> await JSONResponse({"k": "v"})(scope, receive, send)
"""
def __init__(self, content=None, status_code=200):
"""Construct a JSON HTTP Response.
Args:
content (optional): Any JSON serializable object.
status_code (int, optional): Default status code is 200.
"""
self.body = self.render(content)
self.status_code = status_code
self.raw_headers = [[b"content-type", b"application/json"]]
def render(self, content):
if content is None:
return b""
if isinstance(content, bytes):
return content
return json.dumps(content, cls=BytesEncoder, indent=2).encode()
async def __call__(self, scope, receive, send):
await send({
"type": "http.response.start",
"status": self.status_code,
"headers": self.raw_headers,
})
await send({"type": "http.response.body", "body": self.body})
class HTTPProxy:
"""
This class should be instantiated and ran by ASGI server.
>>> import uvicorn
>>> uvicorn.run(HTTPProxy(kv_store_actor_handle, router_handle))
# blocks forever
"""
def __init__(self):
assert ray.is_initialized()
# Delay import due to GlobalState depends on HTTP actor
from ray.experimental.serve.global_state import GlobalState
self.serve_global_state = GlobalState()
self.route_table_cache = dict()
self.route_checker_task = None
self.route_checker_should_shutdown = False
async def route_checker(self, interval):
while True:
if self.route_checker_should_shutdown:
return
self.route_table_cache = (
self.serve_global_state.route_table.list_service())
await asyncio.sleep(interval)
async def handle_lifespan_message(self, scope, receive, send):
assert scope["type"] == "lifespan"
message = await receive()
if message["type"] == "lifespan.startup":
await _async_init()
self.route_checker_task = asyncio.get_event_loop().create_task(
self.route_checker(interval=HTTP_ROUTER_CHECKER_INTERVAL_S))
await send({"type": "lifespan.startup.complete"})
elif message["type"] == "lifespan.shutdown":
self.route_checker_task.cancel()
self.route_checker_should_shutdown = True
await send({"type": "lifespan.shutdown.complete"})
async def receive_http_body(self, scope, receive, send):
body_buffer = []
more_body = True
while more_body:
message = await receive()
assert message["type"] == "http.request"
more_body = message["more_body"]
body_buffer.append(message["body"])
return b"".join(body_buffer)
def _check_slo_ms(self, request_slo_ms):
if request_slo_ms is not None:
if len(request_slo_ms) != 1:
raise ValueError(
"Multiple SLO specified, please specific only one.")
request_slo_ms = request_slo_ms[0]
request_slo_ms = float(request_slo_ms)
if request_slo_ms < 0:
raise ValueError(
"Request SLO must be positive, it is {}".format(
request_slo_ms))
return request_slo_ms
return None
async def __call__(self, scope, receive, send):
# NOTE: This implements ASGI protocol specified in
# https://asgi.readthedocs.io/en/latest/specs/index.html
if scope["type"] == "lifespan":
await self.handle_lifespan_message(scope, receive, send)
return
assert scope["type"] == "http"
current_path = scope["path"]
if current_path == "/":
await JSONResponse(self.route_table_cache)(scope, receive, send)
return
# TODO(simon): Use werkzeug route mapper to support variable path
if current_path not in self.route_table_cache:
error_message = ("Path {} not found. "
"Please ping http://.../ for routing table"
).format(current_path)
await JSONResponse(
{
"error": error_message
}, status_code=404)(scope, receive, send)
return
endpoint_name = self.route_table_cache[current_path]
http_body_bytes = await self.receive_http_body(scope, receive, send)
# get slo_ms before enqueuing the query
query_string = scope["query_string"].decode("ascii")
query_kwargs = parse_qs(query_string)
relative_slo_ms = query_kwargs.pop("relative_slo_ms", None)
absolute_slo_ms = query_kwargs.pop("absolute_slo_ms", None)
try:
relative_slo_ms = self._check_slo_ms(relative_slo_ms)
absolute_slo_ms = self._check_slo_ms(absolute_slo_ms)
if relative_slo_ms is not None and absolute_slo_ms is not None:
raise ValueError("Both relative and absolute slo's"
"cannot be specified.")
except ValueError as e:
await JSONResponse({"error": str(e)})(scope, receive, send)
return
# create objects necessary for enqueue
# enclosing http_body_bytes to list due to
# https://github.com/ray-project/ray/issues/6944
# TODO(alind): remove list enclosing after issue is fixed
args = (scope, [http_body_bytes])
request_in_object = RequestMetadata(
endpoint_name,
TaskContext.Web,
relative_slo_ms=relative_slo_ms,
absolute_slo_ms=absolute_slo_ms)
actual_result = await (self.serve_global_state.init_or_get_router()
.enqueue_request.remote(request_in_object,
*args))
result = actual_result
if isinstance(result, ray.exceptions.RayTaskError):
await JSONResponse({
"error": "internal error, please use python API to debug"
})(scope, receive, send)
else:
await JSONResponse({"result": result})(scope, receive, send)
@ray.remote
class HTTPActor:
def __init__(self):
self.app = HTTPProxy()
def run(self, host="0.0.0.0", port=8000):
uvicorn.run(
self.app, host=host, port=port, lifespan="on", access_log=False)
python/ray/experimental/serve/task_runner.py
-215
View File
@@ -1,215 +0,0 @@
import time
import traceback
import ray
from ray.experimental.serve import context as serve_context
from ray.experimental.serve.context import FakeFlaskRequest
from collections import defaultdict
from ray.experimental.serve.utils import parse_request_item
from ray.experimental.serve.exceptions import RayServeException
class TaskRunner:
"""A simple class that runs a function.
The purpose of this class is to model what the most basic actor could be.
That is, a ray serve actor should implement the TaskRunner interface.
"""
def __init__(self, func_to_run):
self.func = func_to_run
def __call__(self, *args, **kwargs):
return self.func(*args, **kwargs)
def wrap_to_ray_error(exception):
"""Utility method that catch and seal exceptions in execution"""
try:
# Raise and catch so we can access traceback.format_exc()
raise exception
except Exception as e:
traceback_str = ray.utils.format_error_message(traceback.format_exc())
return ray.exceptions.RayTaskError(str(e), traceback_str, e.__class__)
class RayServeMixin:
"""This mixin class adds the functionality to fetch from router queues.
Warning:
It assumes the main execution method is `__call__` of the user defined
class. This means that serve will call `your_instance.__call__` when
each request comes in. This behavior will be fixed in the future to
allow assigning artibrary methods.
Example:
>>> # Use ray.remote decorator and RayServeMixin
>>> # to make MyClass servable.
>>> @ray.remote
class RayServeActor(RayServeMixin, MyClass):
pass
"""
_ray_serve_self_handle = None
_ray_serve_router_handle = None
_ray_serve_setup_completed = False
_ray_serve_dequeue_requester_name = None
# Work token can be unfullfilled from last iteration.
# This cache will be used to determine whether or not we should
# work on the same task as previous iteration or we are ready to
# move on.
_ray_serve_cached_work_token = None
_serve_metric_error_counter = 0
_serve_metric_latency_list = []
def _serve_metric(self):
# Make a copy of the latency list and clear current list
latency_lst = self._serve_metric_latency_list[:]
self._serve_metric_latency_list = []
my_name = self._ray_serve_dequeue_requester_name
return {
"{}_error_counter".format(my_name): {
"value": self._serve_metric_error_counter,
"type": "counter",
},
"{}_latency_s".format(my_name): {
"value": latency_lst,
"type": "list",
},
}
def _ray_serve_setup(self, my_name, router_handle, my_handle):
self._ray_serve_dequeue_requester_name = my_name
self._ray_serve_router_handle = router_handle
self._ray_serve_self_handle = my_handle
self._ray_serve_setup_completed = True
def _ray_serve_fetch(self):
assert self._ray_serve_setup_completed
self._ray_serve_router_handle.dequeue_request.remote(
self._ray_serve_dequeue_requester_name,
self._ray_serve_self_handle)
def invoke_single(self, request_item):
args, kwargs, is_web_context = parse_request_item(request_item)
serve_context.web = is_web_context
start_timestamp = time.time()
try:
result = self.__call__(*args, **kwargs)
except Exception as e:
result = wrap_to_ray_error(e)
self._serve_metric_error_counter += 1
self._serve_metric_latency_list.append(time.time() - start_timestamp)
return result
def invoke_batch(self, request_item_list):
# TODO(alind) : create no-http services. The enqueues
# from such services will always be TaskContext.Python.
# Assumption : all the requests in a bacth
# have same serve context.
# For batching kwargs are modified as follows -
# kwargs [Python Context] : key,val
# kwargs_list : key, [val1,val2, ... , valn]
# or
# args[Web Context] : val
# args_list : [val1,val2, ...... , valn]
# where n (current batch size) <= max_batch_size of a backend
arg_list = []
kwargs_list = defaultdict(list)
context_flags = set()
batch_size = len(request_item_list)
for item in request_item_list:
args, kwargs, is_web_context = parse_request_item(item)
context_flags.add(is_web_context)
if is_web_context:
# Python context only have kwargs
flask_request = args[0]
arg_list.append(flask_request)
else:
# Web context only have one positional argument
for k, v in kwargs.items():
kwargs_list[k].append(v)
# Set the flask request as a list to conform
# with batching semantics: when in batching
# mode, each argument it turned into list.
arg_list.append(FakeFlaskRequest())
try:
# check mixing of query context
# unified context needed
if len(context_flags) != 1:
raise RayServeException(
"Batched queries contain mixed context. Please only send "
"the same type of requests in batching mode.")
serve_context.web = context_flags.pop()
serve_context.batch_size = batch_size
# Flask requests are passed to __call__ as a list
arg_list = [arg_list]
start_timestamp = time.time()
result_list = self.__call__(*arg_list, **kwargs_list)
self._serve_metric_latency_list.append(time.time() -
start_timestamp)
if (not isinstance(result_list,
list)) or (len(result_list) != batch_size):
raise RayServeException("__call__ function "
"doesn't preserve batch-size. "
"Please return a list of result "
"with length equals to the batch "
"size.")
return result_list
except Exception as e:
wrapped_exception = wrap_to_ray_error(e)
self._serve_metric_error_counter += batch_size
return [wrapped_exception for _ in range(batch_size)]
def _ray_serve_call(self, request):
# check if work_item is a list or not
# if it is list: then batching supported
if not isinstance(request, list):
result = self.invoke_single(request)
else:
result = self.invoke_batch(request)
# re-assign to default values
serve_context.web = False
serve_context.batch_size = None
# Tell router that current actor is idle
self._ray_serve_fetch()
return result
class TaskRunnerBackend(TaskRunner, RayServeMixin):
"""A simple function serving backend
Note that this is not yet an actor. To make it an actor:
>>> @ray.remote
class TaskRunnerActor(TaskRunnerBackend):
pass
Note:
This class is not used in the actual ray serve system. It exists
for documentation purpose.
"""
@ray.remote
class TaskRunnerActor(TaskRunnerBackend):
pass
python/ray/experimental/serve/tests/conftest.py
-28
View File
@@ -1,28 +0,0 @@
import os
import tempfile
import pytest
import ray
from ray.experimental import serve
@pytest.fixture(scope="session")
def serve_instance():
_, new_db_path = tempfile.mkstemp(suffix=".test.db")
serve.init(
kv_store_path=new_db_path,
blocking=True,
ray_init_kwargs={"num_cpus": 36})
yield
os.remove(new_db_path)
@pytest.fixture(scope="session")
def ray_instance():
ray_already_initialized = ray.is_initialized()
if not ray_already_initialized:
ray.init(object_store_memory=int(1e8))
yield
if not ray_already_initialized:
ray.shutdown()
python/ray/experimental/serve/tests/test_api.py
-221
View File
@@ -1,221 +0,0 @@
import time
import pytest
import requests
from ray.experimental import serve
from ray.experimental.serve import BackendConfig
import ray
from ray.experimental.serve.constants import NO_ROUTE_KEY
def test_e2e(serve_instance):
serve.init() # so we have access to global state
serve.create_endpoint("endpoint", "/api", blocking=True)
result = serve.api._get_global_state().route_table.list_service()
assert result["/api"] == "endpoint"
retry_count = 5
timeout_sleep = 0.5
while True:
try:
resp = requests.get("http://127.0.0.1:8000/", timeout=0.5).json()
assert resp == result
break
except Exception:
time.sleep(timeout_sleep)
timeout_sleep *= 2
retry_count -= 1
if retry_count == 0:
assert False, "Route table hasn't been updated after 3 tries."
def function(flask_request):
return "OK"
serve.create_backend(function, "echo:v1")
serve.link("endpoint", "echo:v1")
resp = requests.get("http://127.0.0.1:8000/api").json()["result"]
assert resp == "OK"
def test_no_route(serve_instance):
serve.create_endpoint("noroute-endpoint", blocking=True)
global_state = serve.api._get_global_state()
result = global_state.route_table.list_service(include_headless=True)
assert result[NO_ROUTE_KEY] == ["noroute-endpoint"]
without_headless_result = global_state.route_table.list_service()
assert NO_ROUTE_KEY not in without_headless_result
def func(_, i=1):
return 1
serve.create_backend(func, "backend:1")
serve.link("noroute-endpoint", "backend:1")
service_handle = serve.get_handle("noroute-endpoint")
result = ray.get(service_handle.remote(i=1))
assert result == 1
def test_scaling_replicas(serve_instance):
class Counter:
def __init__(self):
self.count = 0
def __call__(self, _):
self.count += 1
return self.count
serve.create_endpoint("counter", "/increment")
# Keep checking the routing table until /increment is populated
while "/increment" not in requests.get("http://127.0.0.1:8000/").json():
time.sleep(0.2)
b_config = BackendConfig(num_replicas=2)
serve.create_backend(Counter, "counter:v1", backend_config=b_config)
serve.link("counter", "counter:v1")
counter_result = []
for _ in range(10):
resp = requests.get("http://127.0.0.1:8000/increment").json()["result"]
counter_result.append(resp)
# If the load is shared among two replicas. The max result cannot be 10.
assert max(counter_result) < 10
b_config = serve.get_backend_config("counter:v1")
b_config.num_replicas = 1
serve.set_backend_config("counter:v1", b_config)
counter_result = []
for _ in range(10):
resp = requests.get("http://127.0.0.1:8000/increment").json()["result"]
counter_result.append(resp)
# Give some time for a replica to spin down. But majority of the request
# should be served by the only remaining replica.
assert max(counter_result) - min(counter_result) > 6
def test_batching(serve_instance):
class BatchingExample:
def __init__(self):
self.count = 0
@serve.accept_batch
def __call__(self, flask_request, temp=None):
self.count += 1
batch_size = serve.context.batch_size
return [self.count] * batch_size
serve.create_endpoint("counter1", "/increment")
# Keep checking the routing table until /increment is populated
while "/increment" not in requests.get("http://127.0.0.1:8000/").json():
time.sleep(0.2)
# set the max batch size
b_config = BackendConfig(max_batch_size=5)
serve.create_backend(
BatchingExample, "counter:v11", backend_config=b_config)
serve.link("counter1", "counter:v11")
future_list = []
handle = serve.get_handle("counter1")
for _ in range(20):
f = handle.remote(temp=1)
future_list.append(f)
counter_result = ray.get(future_list)
# since count is only updated per batch of queries
# If there atleast one __call__ fn call with batch size greater than 1
# counter result will always be less than 20
assert max(counter_result) < 20
def test_batching_exception(serve_instance):
class NoListReturned:
def __init__(self):
self.count = 0
@serve.accept_batch
def __call__(self, flask_request, temp=None):
batch_size = serve.context.batch_size
return batch_size
serve.create_endpoint("exception-test", "/noListReturned")
# set the max batch size
b_config = BackendConfig(max_batch_size=5)
serve.create_backend(
NoListReturned, "exception:v1", backend_config=b_config)
serve.link("exception-test", "exception:v1")
handle = serve.get_handle("exception-test")
with pytest.raises(ray.exceptions.RayTaskError):
assert ray.get(handle.remote(temp=1))
def test_killing_replicas(serve_instance):
class Simple:
def __init__(self):
self.count = 0
def __call__(self, flask_request, temp=None):
return temp
serve.create_endpoint("simple", "/simple")
b_config = BackendConfig(num_replicas=3, num_cpus=2)
serve.create_backend(Simple, "simple:v1", backend_config=b_config)
global_state = serve.api._get_global_state()
old_replica_tag_list = global_state.backend_table.list_replicas(
"simple:v1")
bnew_config = serve.get_backend_config("simple:v1")
# change the config
bnew_config.num_cpus = 1
# set the config
serve.set_backend_config("simple:v1", bnew_config)
new_replica_tag_list = global_state.backend_table.list_replicas(
"simple:v1")
global_state.refresh_actor_handle_cache()
new_all_tag_list = list(global_state.actor_handle_cache.keys())
# the new_replica_tag_list must be subset of all_tag_list
assert set(new_replica_tag_list) <= set(new_all_tag_list)
# the old_replica_tag_list must not be subset of all_tag_list
assert not set(old_replica_tag_list) <= set(new_all_tag_list)
def test_not_killing_replicas(serve_instance):
class BatchSimple:
def __init__(self):
self.count = 0
@serve.accept_batch
def __call__(self, flask_request, temp=None):
batch_size = serve.context.batch_size
return [1] * batch_size
serve.create_endpoint("bsimple", "/bsimple")
b_config = BackendConfig(num_replicas=3, max_batch_size=2)
serve.create_backend(BatchSimple, "bsimple:v1", backend_config=b_config)
global_state = serve.api._get_global_state()
old_replica_tag_list = global_state.backend_table.list_replicas(
"bsimple:v1")
bnew_config = serve.get_backend_config("bsimple:v1")
# change the config
bnew_config.max_batch_size = 5
# set the config
serve.set_backend_config("bsimple:v1", bnew_config)
new_replica_tag_list = global_state.backend_table.list_replicas(
"bsimple:v1")
global_state.refresh_actor_handle_cache()
new_all_tag_list = list(global_state.actor_handle_cache.keys())
# the old and new replica tag list should be identical
# and should be subset of all_tag_list
assert set(old_replica_tag_list) <= set(new_all_tag_list)
assert set(old_replica_tag_list) == set(new_replica_tag_list)
python/ray/experimental/serve/tests/test_metric.py
-74
View File
@@ -1,74 +0,0 @@
import numpy as np
import pytest
import ray
from ray.experimental.serve.metric import MetricMonitor
@pytest.fixture(scope="session")
def start_target_actor(ray_instance):
@ray.remote
class Target:
def __init__(self):
self.counter_value = 0
def _serve_metric(self):
self.counter_value += 1
return {
"latency_list": {
"type": "list",
# Generate 0 to 100 inclusive.
# This means total of 101 items.
"value": np.arange(101).tolist()
},
"counter": {
"type": "counter",
"value": self.counter_value
}
}
def get_counter_value(self):
return self.counter_value
yield Target.remote()
def test_metric_gc(ray_instance, start_target_actor):
target_actor = start_target_actor
# this means when new scrapes are invoked, the
metric_monitor = MetricMonitor.remote(gc_window_seconds=0)
ray.get(metric_monitor.add_target.remote(target_actor))
ray.get(metric_monitor.scrape.remote())
df = ray.get(metric_monitor._get_dataframe.remote())
assert len(df) == 102
# Old metric sould be cleared. So only 1 counter + 101 list values left.
ray.get(metric_monitor.scrape.remote())
df = ray.get(metric_monitor._get_dataframe.remote())
assert len(df) == 102
def test_metric_system(ray_instance, start_target_actor):
target_actor = start_target_actor
metric_monitor = MetricMonitor.remote()
ray.get(metric_monitor.add_target.remote(target_actor))
# Scrape once
ray.get(metric_monitor.scrape.remote())
percentiles = [50, 90, 95]
agg_windows_seconds = [60]
result = ray.get(
metric_monitor.collect.remote(percentiles, agg_windows_seconds))
real_counter_value = ray.get(target_actor.get_counter_value.remote())
expected_result = {
"counter": real_counter_value,
"latency_list_50th_perc_60_window": 50.0,
"latency_list_90th_perc_60_window": 90.0,
"latency_list_95th_perc_60_window": 95.0,
}
assert result == expected_result
python/ray/experimental/serve/tests/test_persistence.py
-33
View File
@@ -1,33 +0,0 @@
import os
import subprocess
import tempfile
import ray
from ray.experimental import serve
def test_new_driver(serve_instance):
script = """
import ray
ray.init(address="auto")
from ray.experimental import serve
serve.init()
@serve.route("/driver")
def driver(flask_request):
return "OK!"
"""
with tempfile.NamedTemporaryFile(mode="w", delete=False) as f:
path = f.name
f.write(script)
proc = subprocess.Popen(["python", path])
return_code = proc.wait(timeout=10)
assert return_code == 0
handle = serve.get_handle("driver")
assert ray.get(handle.remote()) == "OK!"
os.remove(path)
python/ray/experimental/serve/tests/test_queue.py
-192
View File
@@ -1,192 +0,0 @@
import asyncio
import pytest
import ray
from ray.experimental.serve.policy import (
RandomPolicyQueue, RandomPolicyQueueActor, RoundRobinPolicyQueueActor,
PowerOfTwoPolicyQueueActor, FixedPackingPolicyQueueActor)
from ray.experimental.serve.request_params import RequestMetadata
pytestmark = pytest.mark.asyncio
def make_task_runner_mock():
@ray.remote(num_cpus=0)
class TaskRunnerMock:
def __init__(self):
self.query = None
self.queries = []
async def _ray_serve_call(self, request_item):
self.query = request_item
self.queries.append(request_item)
return "DONE"
def get_recent_call(self):
return self.query
def get_all_calls(self):
return self.queries
return TaskRunnerMock.remote()
@pytest.fixture(scope="session")
def task_runner_mock_actor():
yield make_task_runner_mock()
async def test_single_prod_cons_queue(serve_instance, task_runner_mock_actor):
q = RandomPolicyQueueActor.remote()
q.link.remote("svc", "backend")
q.dequeue_request.remote("backend", task_runner_mock_actor)
# Make sure we get the request result back
result = await q.enqueue_request.remote(RequestMetadata("svc", None), 1)
assert result == "DONE"
# Make sure it's the right request
got_work = await task_runner_mock_actor.get_recent_call.remote()
assert got_work.request_args[0] == 1
assert got_work.request_kwargs == {}
async def test_slo(serve_instance, task_runner_mock_actor):
q = RandomPolicyQueueActor.remote()
await q.link.remote("svc", "backend")
all_request_sent = []
for i in range(10):
slo_ms = 1000 - 100 * i
all_request_sent.append(
q.enqueue_request.remote(
RequestMetadata("svc", None, relative_slo_ms=slo_ms), i))
for i in range(10):
await q.dequeue_request.remote("backend", task_runner_mock_actor)
await asyncio.gather(*all_request_sent)
i_should_be = 9
all_calls = await task_runner_mock_actor.get_all_calls.remote()
all_calls = all_calls[-10:]
for call in all_calls:
assert call.request_args[0] == i_should_be
i_should_be -= 1
async def test_alter_backend(serve_instance, task_runner_mock_actor):
q = RandomPolicyQueueActor.remote()
await q.set_traffic.remote("svc", {"backend-1": 1})
await q.dequeue_request.remote("backend-1", task_runner_mock_actor)
await q.enqueue_request.remote(RequestMetadata("svc", None), 1)
got_work = await task_runner_mock_actor.get_recent_call.remote()
assert got_work.request_args[0] == 1
await q.set_traffic.remote("svc", {"backend-2": 1})
await q.dequeue_request.remote("backend-2", task_runner_mock_actor)
await q.enqueue_request.remote(RequestMetadata("svc", None), 2)
got_work = await task_runner_mock_actor.get_recent_call.remote()
assert got_work.request_args[0] == 2
async def test_split_traffic_random(serve_instance, task_runner_mock_actor):
q = RandomPolicyQueueActor.remote()
await q.set_traffic.remote("svc", {"backend-1": 0.5, "backend-2": 0.5})
runner_1, runner_2 = [make_task_runner_mock() for _ in range(2)]
for _ in range(20):
await q.dequeue_request.remote("backend-1", runner_1)
await q.dequeue_request.remote("backend-2", runner_2)
# assume 50% split, the probability of all 20 requests goes to a
# single queue is 0.5^20 ~ 1-6
for _ in range(20):
await q.enqueue_request.remote(RequestMetadata("svc", None), 1)
got_work = [
await runner.get_recent_call.remote()
for runner in (runner_1, runner_2)
]
assert [g.request_args[0] for g in got_work] == [1, 1]
async def test_round_robin(serve_instance, task_runner_mock_actor):
q = RoundRobinPolicyQueueActor.remote()
await q.set_traffic.remote("svc", {"backend-1": 0.5, "backend-2": 0.5})
runner_1, runner_2 = [make_task_runner_mock() for _ in range(2)]
# NOTE: this is the only difference between the
# test_split_traffic_random and test_round_robin
for _ in range(10):
await q.dequeue_request.remote("backend-1", runner_1)
await q.dequeue_request.remote("backend-2", runner_2)
for _ in range(20):
await q.enqueue_request.remote(RequestMetadata("svc", None), 1)
got_work = [
await runner.get_recent_call.remote()
for runner in (runner_1, runner_2)
]
assert [g.request_args[0] for g in got_work] == [1, 1]
async def test_fixed_packing(serve_instance):
packing_num = 4
q = FixedPackingPolicyQueueActor.remote(packing_num=packing_num)
await q.set_traffic.remote("svc", {"backend-1": 0.5, "backend-2": 0.5})
runner_1, runner_2 = (make_task_runner_mock() for _ in range(2))
# both the backends will get equal number of queries
# as it is packed round robin
for _ in range(packing_num):
await q.dequeue_request.remote("backend-1", runner_1)
await q.dequeue_request.remote("backend-2", runner_2)
for backend, runner in zip(["1", "2"], [runner_1, runner_2]):
for _ in range(packing_num):
input_value = "should-go-to-backend-{}".format(backend)
await q.enqueue_request.remote(
RequestMetadata("svc", None), input_value)
all_calls = await runner.get_all_calls.remote()
for call in all_calls:
assert call.request_args[0] == input_value
async def test_power_of_two_choices(serve_instance):
q = PowerOfTwoPolicyQueueActor.remote()
enqueue_futures = []
# First, fill the queue for backend-1 with 3 requests
await q.set_traffic.remote("svc", {"backend-1": 1.0})
for _ in range(3):
future = q.enqueue_request.remote(RequestMetadata("svc", None), "1")
enqueue_futures.append(future)
# Then, add a new backend, this backend should be filled next
await q.set_traffic.remote("svc", {"backend-1": 0.5, "backend-2": 0.5})
for _ in range(2):
future = q.enqueue_request.remote(RequestMetadata("svc", None), "2")
enqueue_futures.append(future)
runner_1, runner_2 = (make_task_runner_mock() for _ in range(2))
for _ in range(3):
await q.dequeue_request.remote("backend-1", runner_1)
await q.dequeue_request.remote("backend-2", runner_2)
await asyncio.gather(*enqueue_futures)
assert len(await runner_1.get_all_calls.remote()) == 3
assert len(await runner_2.get_all_calls.remote()) == 2
async def test_queue_remove_replicas(serve_instance):
temp_actor = make_task_runner_mock()
q = RandomPolicyQueue()
await q.dequeue_request("backend", temp_actor)
await q.remove_and_destory_replica("backend", temp_actor)
assert q.worker_queues["backend"].qsize() == 0
python/ray/experimental/serve/tests/test_routing.py
-54
View File
@@ -1,54 +0,0 @@
import os
import tempfile
from ray.experimental.serve.kv_store_service import (
InMemoryKVStore, RayInternalKVStore, SQLiteKVStore)
def test_default_in_memory_kv():
kv = InMemoryKVStore("")
kv.put("1", 2)
assert kv.get("1") == 2
kv.put("1", 3)
assert kv.get("1") == 3
assert kv.as_dict() == {"1": 3}
def test_ray_interal_kv(ray_instance):
kv = RayInternalKVStore("")
kv.put("1", 2)
assert kv.get("1") == 2
kv.put("1", 3)
assert kv.get("1") == 3
assert kv.as_dict() == {"1": 3}
kv = RayInternalKVStore("othernamespace")
kv.put("1", 2)
assert kv.get("1") == 2
kv.put("1", 3)
assert kv.get("1") == 3
assert kv.as_dict() == {"1": 3}
def test_sqlite_kv():
_, path = tempfile.mkstemp()
# Test get
kv = SQLiteKVStore("routing_table", db_path=path)
kv.put("/api", "api-endpoint")
assert kv.get("/api") == "api-endpoint"
assert kv.get("not-exist") is None
# Test namespace
kv2 = SQLiteKVStore("other_table", db_path=path)
kv2.put("/api", "api-endpoint-two")
assert kv2.get("/api") == "api-endpoint-two"
# Test as dict
assert kv.as_dict() == {"/api": "api-endpoint"}
# Test override
kv.put("/api", "api-new")
assert kv.get("/api") == "api-new"
os.remove(path)
python/ray/experimental/serve/tests/test_serve.py
-15
View File
@@ -1,15 +0,0 @@
import pytest
from pathlib import Path
import sys
if __name__ == "__main__":
curr_dir = Path(__file__).parent
test_paths = curr_dir.rglob("test_*.py")
sorted_path = sorted(map(lambda path: str(path.absolute()), test_paths))
serve_tests_files = list(sorted_path)
print("Testing the following files")
for test_file in serve_tests_files:
print("->", test_file.split("/")[-1])
sys.exit(pytest.main(["-v", "-s"] + serve_tests_files))
python/ray/experimental/serve/tests/test_task_runner.py
-99
View File
@@ -1,99 +0,0 @@
import pytest
import ray
import ray.experimental.serve.context as context
from ray.experimental.serve.policy import RoundRobinPolicyQueueActor
from ray.experimental.serve.task_runner import (
RayServeMixin, TaskRunner, TaskRunnerActor, wrap_to_ray_error)
from ray.experimental.serve.request_params import RequestMetadata
pytestmark = pytest.mark.asyncio
async def test_runner_basic():
def echo(i):
return i
r = TaskRunner(echo)
assert r(1) == 1
async def test_runner_wraps_error():
wrapped = wrap_to_ray_error(Exception())
assert isinstance(wrapped, ray.exceptions.RayTaskError)
async def test_runner_actor(serve_instance):
q = RoundRobinPolicyQueueActor.remote()
def echo(flask_request, i=None):
return i
CONSUMER_NAME = "runner"
PRODUCER_NAME = "prod"
runner = TaskRunnerActor.remote(echo)
runner._ray_serve_setup.remote(CONSUMER_NAME, q, runner)
runner._ray_serve_fetch.remote()
q.link.remote(PRODUCER_NAME, CONSUMER_NAME)
for query in [333, 444, 555]:
query_param = RequestMetadata(PRODUCER_NAME,
context.TaskContext.Python)
result = await q.enqueue_request.remote(query_param, i=query)
assert result == query
async def test_ray_serve_mixin(serve_instance):
q = RoundRobinPolicyQueueActor.remote()
CONSUMER_NAME = "runner-cls"
PRODUCER_NAME = "prod-cls"
class MyAdder:
def __init__(self, inc):
self.increment = inc
def __call__(self, flask_request, i=None):
return i + self.increment
@ray.remote
class CustomActor(MyAdder, RayServeMixin):
pass
runner = CustomActor.remote(3)
runner._ray_serve_setup.remote(CONSUMER_NAME, q, runner)
runner._ray_serve_fetch.remote()
q.link.remote(PRODUCER_NAME, CONSUMER_NAME)
for query in [333, 444, 555]:
query_param = RequestMetadata(PRODUCER_NAME,
context.TaskContext.Python)
result = await q.enqueue_request.remote(query_param, i=query)
assert result == query + 3
async def test_task_runner_check_context(serve_instance):
q = RoundRobinPolicyQueueActor.remote()
def echo(flask_request, i=None):
# Accessing the flask_request without web context should throw.
return flask_request.args["i"]
CONSUMER_NAME = "runner"
PRODUCER_NAME = "producer"
runner = TaskRunnerActor.remote(echo)
runner._ray_serve_setup.remote(CONSUMER_NAME, q, runner)
runner._ray_serve_fetch.remote()
q.link.remote(PRODUCER_NAME, CONSUMER_NAME)
query_param = RequestMetadata(PRODUCER_NAME, context.TaskContext.Python)
result_oid = q.enqueue_request.remote(query_param, i=42)
with pytest.raises(ray.exceptions.RayTaskError):
await result_oid
python/ray/experimental/serve/tests/test_util.py
-9
View File
@@ -1,9 +0,0 @@
import json
from ray.experimental.serve.utils import BytesEncoder
def test_bytes_encoder():
data_before = {"inp": {"nest": b"bytes"}}
data_after = {"inp": {"nest": "bytes"}}
assert json.loads(json.dumps(data_before, cls=BytesEncoder)) == data_after
python/ray/experimental/serve/utils.py
-112
View File
@@ -1,112 +0,0 @@
import json
import logging
import random
import string
import time
import io
import os
import requests
from pygments import formatters, highlight, lexers
from ray.experimental.serve.context import FakeFlaskRequest, TaskContext
from ray.experimental.serve.http_util import build_flask_request
import itertools
def expand(l):
"""
Implements a nested flattening of a list.
Example:
>>> serve.utils.expand([1,2,[3,4,5],6])
[1,2,3,4,5,6]
>>> serve.utils.expand(["a", ["b", "c"], "d", ["e", "f"]])
["a", "b", "c", "d", "e", "f"]
"""
return list(
itertools.chain.from_iterable(
[x if isinstance(x, list) else [x] for x in l]))
def parse_request_item(request_item):
if request_item.request_context == TaskContext.Web:
is_web_context = True
asgi_scope, body_bytes = request_item.request_args
# http_body_bytes enclosed in list due to
# https://github.com/ray-project/ray/issues/6944
# TODO(alind): remove list enclosing after issue is fixed
flask_request = build_flask_request(asgi_scope,
io.BytesIO(body_bytes[0]))
args = (flask_request, )
kwargs = {}
else:
is_web_context = False
args = (FakeFlaskRequest(), )
kwargs = request_item.request_kwargs
return args, kwargs, is_web_context
def _get_logger():
logger = logging.getLogger("ray.serve")
# TODO(simon): Make logging level configurable.
if os.environ.get("SERVE_LOG_DEBUG"):
logger.setLevel(logging.DEBUG)
else:
logger.setLevel(logging.INFO)
return logger
logger = _get_logger()
class BytesEncoder(json.JSONEncoder):
"""Allow bytes to be part of the JSON document.
BytesEncoder will walk the JSON tree and decode bytes with utf-8 codec.
Example:
>>> json.dumps({b'a': b'c'}, cls=BytesEncoder)
'{"a":"c"}'
"""
def default(self, o): # pylint: disable=E0202
if isinstance(o, bytes):
return o.decode("utf-8")
return super().default(o)
def pformat_color_json(d):
"""Use pygments to pretty format and colroize dictionary"""
formatted_json = json.dumps(d, sort_keys=True, indent=4)
colorful_json = highlight(formatted_json, lexers.JsonLexer(),
formatters.TerminalFormatter())
return colorful_json
def block_until_http_ready(http_endpoint, num_retries=5, backoff_time_s=1):
http_is_ready = False
retries = num_retries
while not http_is_ready:
try:
resp = requests.get(http_endpoint)
assert resp.status_code == 200
http_is_ready = True
except Exception:
pass
# Exponential backoff
time.sleep(backoff_time_s)
backoff_time_s *= 2
retries -= 1
if retries == 0:
raise Exception(
"HTTP server not ready after {} retries.".format(num_retries))
def get_random_letters(length=6):
return "".join(random.choices(string.ascii_letters, k=length))
python/ray/experimental/sgd/__init__.py
-4
View File
@@ -1,4 +0,0 @@
from ray.experimental.sgd.pytorch import PyTorchTrainer
from ray.experimental.sgd.tf import TFTrainer
__all__ = ["PyTorchTrainer", "TFTrainer"]
python/ray/experimental/sgd/pytorch/__init__.py
-15
View File
@@ -1,15 +0,0 @@
import logging
logger = logging.getLogger(__name__)
PyTorchTrainer = None
PyTorchTrainable = None
try:
import torch # noqa: F401
from ray.experimental.sgd.pytorch.pytorch_trainer import (PyTorchTrainer,
PyTorchTrainable)
__all__ = ["PyTorchTrainer", "PyTorchTrainable"]
except ImportError:
logger.warning("PyTorch not found. PyTorchTrainer will not be available")
python/ray/experimental/sgd/pytorch/distributed_pytorch_runner.py
-134
View File
@@ -1,134 +0,0 @@
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.experimental.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/experimental/sgd/pytorch/examples/cifar_pytorch_example.py
-161
View File
@@ -1,161 +0,0 @@
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.experimental.sgd.pytorch import (PyTorchTrainer, PyTorchTrainable)
from ray.experimental.sgd.pytorch.resnet import ResNet18
from ray.experimental.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/experimental/sgd/pytorch/examples/dcgan.py
-267
View File
@@ -1,267 +0,0 @@
#!/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.experimental.sgd import PyTorchTrainer
from ray.experimental.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, "experimental/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/experimental/sgd/pytorch/examples/example-sgd.yaml
-71
View File
@@ -1,71 +0,0 @@
# 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/experimental/sgd/pytorch/examples/mnist_cnn.pt
BIN
View File
Binary file not shown.
python/ray/experimental/sgd/pytorch/examples/train_example.py
-112
View File
@@ -1,112 +0,0 @@
"""
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.experimental.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/experimental/sgd/pytorch/examples/tune_example.py
-96
View File
@@ -1,96 +0,0 @@
# 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.experimental.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/experimental/sgd/pytorch/pytorch_runner.py
-301
View File
@@ -1,301 +0,0 @@
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.experimental.sgd.pytorch import utils as pytorch_utils
from ray.experimental.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/experimental/sgd/pytorch/pytorch_trainer.py
-464
View File
@@ -1,464 +0,0 @@
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.experimental.sgd.pytorch.distributed_pytorch_runner import (
DistributedPyTorchRunner)
from ray.experimental.sgd import utils
from ray.experimental.sgd.pytorch.pytorch_runner import PyTorchRunner
from ray.experimental.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/experimental/sgd/pytorch/resnet.py
-134
View File
@@ -1,134 +0,0 @@
"""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/experimental/sgd/pytorch/utils.py
-229
View File
@@ -1,229 +0,0 @@
import collections
import time
import torch
from ray.experimental.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/experimental/sgd/tests/test_pytorch.py
-385
View File
@@ -1,385 +0,0 @@
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.experimental.sgd.pytorch import PyTorchTrainer, PyTorchTrainable
from ray.experimental.sgd.pytorch.utils import (train, BATCH_COUNT, TEST_MODE,
SCHEDULER_STEP)
from ray.experimental.sgd.utils import check_for_failure
from ray.experimental.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/experimental/sgd/tests/test_pytorch_runner.py
-151
View File
@@ -1,151 +0,0 @@
import numpy as np
import torch
import torch.nn as nn
import unittest
from unittest.mock import MagicMock
from ray.experimental.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/experimental/sgd/tests/test_tensorflow.py
-131
View File
@@ -1,131 +0,0 @@
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.experimental.sgd.tf import TFTrainer, TFTrainable
from ray.experimental.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/experimental/sgd/tf/__init__.py
-3
View File
@@ -1,3 +0,0 @@
from ray.experimental.sgd.tf.tf_trainer import (TFTrainer, TFTrainable)
__all__ = ["TFTrainer", "TFTrainable"]
python/ray/experimental/sgd/tf/examples/cifar_tf_example.py
-227
View File
@@ -1,227 +0,0 @@
"""
#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.experimental.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/experimental/sgd/tf/examples/tensorflow_train_example.py
-143
View File
@@ -1,143 +0,0 @@
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.experimental.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/experimental/sgd/tf/examples/tf-example-sgd.yaml
-70
View File
@@ -1,70 +0,0 @@
# 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/experimental/sgd/tf/tf_runner.py
-163
View File
@@ -1,163 +0,0 @@
import logging
import json
import os
import numpy as np
import ray
import ray.services
from ray.experimental.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/experimental/sgd/tf/tf_trainer.py
-203
View File
@@ -1,203 +0,0 @@
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.experimental.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/experimental/sgd/utils.py
-151
View File
@@ -1,151 +0,0 @@
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