wehub/ray-project--ray
1
0
Fork 0
Code Issues Pull Requests Actions Deployments Models Agent Notes Packages Projects Releases Wiki Activity

chore: import upstream snapshot with attribution

Browse Source
This commit is contained in:
wehub-resource-sync
2026-07-13 13:17:40 +08:00
commit f1825c8ceb
10096 changed files with 2364182 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/ray/util/BUILD.bazel
+15
View File
@@ -0,0 +1,15 @@
load("//bazel:python.bzl", "doctest")
# NOTE: `iter.py`, `iter_metrics.py`, and `client_connect.py` only contain deprecated
# APIs, so we're excluding them.
doctest(
files = glob(
["*.py"],
exclude = [
"client_connect.py",
"iter.py",
"iter_metrics.py",
],
),
tags = ["team:core"],
)
python/ray/util/__init__.py
+75
View File
@@ -0,0 +1,75 @@
from typing import List
import ray
from ray._private.auto_init_hook import wrap_auto_init
from ray._private.client_mode_hook import client_mode_hook
from ray._private.services import get_node_instance_id, get_node_ip_address
from ray.util import accelerators, debugpy as ray_debugpy, iter, rpdb as pdb
from ray.util.actor_pool import ActorPool
from ray.util.annotations import PublicAPI
from ray.util.check_serialize import inspect_serializability
from ray.util.client_connect import connect, disconnect
from ray.util.debug import disable_log_once_globally, enable_periodic_logging, log_once
from ray.util.helpers import as_completed, map_unordered
from ray.util.placement_group import (
get_current_placement_group,
get_placement_group,
placement_group,
placement_group_table,
remove_placement_group,
)
from ray.util.serialization import deregister_serializer, register_serializer
@PublicAPI(stability="beta")
@wrap_auto_init
@client_mode_hook
def list_named_actors(all_namespaces: bool = False) -> List[str]:
"""List all named actors in the system.
Actors must have been created with Actor.options(name="name").remote().
This works for both detached & non-detached actors.
By default, only actors in the current namespace will be returned
and the returned entries will simply be their name.
If `all_namespaces` is set to True, all actors in the cluster will be
returned regardless of namespace, and the returned entries will be of the
form {"namespace": namespace, "name": name}.
"""
worker = ray._private.worker.global_worker
worker.check_connected()
actors = worker.core_worker.list_named_actors(all_namespaces)
if all_namespaces:
return [{"name": name, "namespace": namespace} for namespace, name in actors]
else:
return [name for _, name in actors]
__all__ = [
"accelerators",
"ActorPool",
"as_completed",
"disable_log_once_globally",
"enable_periodic_logging",
"iter",
"log_once",
"pdb",
"placement_group",
"placement_group_table",
"get_placement_group",
"get_current_placement_group",
"get_node_instance_id",
"get_node_ip_address",
"map_unordered",
"remove_placement_group",
"ray_debugpy",
"inspect_serializability",
"collective",
"connect",
"disconnect",
"register_serializer",
"deregister_serializer",
"list_named_actors",
]
python/ray/util/accelerators/__init__.py
+84
View File
@@ -0,0 +1,84 @@
import warnings
from ray.util import tpu
from ray.util.accelerators.accelerators import (
AMD_INSTINCT_MI100,
AMD_INSTINCT_MI210,
AMD_INSTINCT_MI250,
AMD_RADEON_HD_7900,
AMD_RADEON_R9_200_HD_7900,
AWS_NEURON_CORE,
GOOGLE_TPU_V2,
GOOGLE_TPU_V3,
GOOGLE_TPU_V4,
GOOGLE_TPU_V5LITEPOD,
GOOGLE_TPU_V5P,
GOOGLE_TPU_V6E,
GOOGLE_TPU_V7X,
INTEL_GAUDI,
INTEL_MAX_1100,
INTEL_MAX_1550,
METAX_C500,
METAX_C550,
NVIDIA_A100,
NVIDIA_H100,
NVIDIA_L4,
NVIDIA_TESLA_A10G,
NVIDIA_TESLA_K80,
NVIDIA_TESLA_P4,
NVIDIA_TESLA_P100,
NVIDIA_TESLA_T4,
NVIDIA_TESLA_V100,
AMD_INSTINCT_MI250x,
AMD_INSTINCT_MI300x,
)
__all__ = [
"tpu",
"NVIDIA_TESLA_V100",
"NVIDIA_TESLA_P100",
"NVIDIA_TESLA_T4",
"NVIDIA_TESLA_P4",
"NVIDIA_TESLA_K80",
"NVIDIA_TESLA_A10G",
"NVIDIA_L4",
"NVIDIA_A100",
"NVIDIA_A100_40G",
"NVIDIA_A100_80G",
"NVIDIA_H100",
"INTEL_MAX_1550",
"INTEL_MAX_1100",
"INTEL_GAUDI",
"AMD_INSTINCT_MI100",
"AMD_INSTINCT_MI210",
"AMD_INSTINCT_MI250",
"AMD_INSTINCT_MI250x",
"AMD_INSTINCT_MI300x",
"AMD_RADEON_R9_200_HD_7900",
"AMD_RADEON_HD_7900",
"AWS_NEURON_CORE",
"GOOGLE_TPU_V2",
"GOOGLE_TPU_V3",
"GOOGLE_TPU_V4",
"GOOGLE_TPU_V5P",
"GOOGLE_TPU_V5LITEPOD",
"GOOGLE_TPU_V6E",
"GOOGLE_TPU_V7X",
"METAX_C500",
"METAX_C550",
# Deprecated
"NVIDIA_TESLA_A100",
]
def __getattr__(name: str):
if name == "NVIDIA_TESLA_A100":
from ray.util.annotations import RayDeprecationWarning
warnings.warn(
"NVIDIA_TESLA_A100 is deprecated, use NVIDIA_A100 instead.",
RayDeprecationWarning,
stacklevel=2,
)
return NVIDIA_A100
raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
python/ray/util/accelerators/accelerators.py
+50
View File
@@ -0,0 +1,50 @@
NVIDIA_TESLA_V100 = "V100"
NVIDIA_TESLA_P100 = "P100"
NVIDIA_TESLA_T4 = "T4"
NVIDIA_TESLA_P4 = "P4"
NVIDIA_TESLA_K80 = "K80"
NVIDIA_TESLA_A10G = "A10G"
NVIDIA_L4 = "L4"
NVIDIA_L40S = "L40S"
NVIDIA_A100 = "A100"
NVIDIA_H100 = "H100"
NVIDIA_H200 = "H200"
NVIDIA_H20 = "H20"
NVIDIA_B200 = "B200"
NVIDIA_B300 = "B300"
NVIDIA_RTX_PRO_6000 = "RTX-PRO-6000"
INTEL_MAX_1550 = "Intel-GPU-Max-1550"
INTEL_MAX_1100 = "Intel-GPU-Max-1100"
INTEL_GAUDI = "Intel-GAUDI"
AMD_INSTINCT_MI100 = "AMD-Instinct-MI100"
AMD_INSTINCT_MI250x = "AMD-Instinct-MI250X"
AMD_INSTINCT_MI250 = "AMD-Instinct-MI250X-MI250"
AMD_INSTINCT_MI210 = "AMD-Instinct-MI210"
AMD_INSTINCT_MI300A = "AMD-Instinct-MI300A"
AMD_INSTINCT_MI300x = "AMD-Instinct-MI300X-OAM"
AMD_INSTINCT_MI300x_HF = "AMD-Instinct-MI300X-HF"
AMD_INSTINCT_MI308x = "AMD-Instinct-MI308X"
AMD_INSTINCT_MI325x = "AMD-Instinct-MI325X-OAM"
AMD_INSTINCT_MI350x = "AMD-Instinct-MI350X-OAM"
AMD_INSTINCT_MI355x = "AMD-Instinct-MI355X-OAM"
AMD_RADEON_R9_200_HD_7900 = "AMD-Radeon-R9-200-HD-7900"
AMD_RADEON_HD_7900 = "AMD-Radeon-HD-7900"
AWS_NEURON_CORE = "aws-neuron-core"
GOOGLE_TPU_V2 = "TPU-V2"
GOOGLE_TPU_V3 = "TPU-V3"
GOOGLE_TPU_V4 = "TPU-V4"
GOOGLE_TPU_V5P = "TPU-V5P"
GOOGLE_TPU_V5LITEPOD = "TPU-V5LITEPOD"
GOOGLE_TPU_V6E = "TPU-V6E"
GOOGLE_TPU_V7X = "TPU-V7X"
HUAWEI_NPU_910B = "Ascend910B"
HUAWEI_NPU_910B4 = "Ascend910B4"
METAX_C500 = "MXC500"
METAX_C550 = "MXC550"
FURIOSA_RNGD = "FURIOSA_RNGD"
# Use these instead of NVIDIA_A100 if you need a specific accelerator size. Note that
# these labels are not auto-added to nodes, you'll have to add them manually in
# addition to the default A100 label if needed.
NVIDIA_A100_40G = "A100-40G"
NVIDIA_A100_80G = "A100-80G"
python/ray/util/actor_group.py
+231
View File
@@ -0,0 +1,231 @@
import logging
import weakref
from dataclasses import dataclass
from typing import Dict, List, Optional, Tuple, Type, TypeVar
import ray
from ray._private.utils import get_ray_doc_version
from ray.actor import ActorHandle
from ray.util.annotations import Deprecated
T = TypeVar("T")
ActorMetadata = TypeVar("ActorMetadata")
logger = logging.getLogger(__name__)
@dataclass
class ActorWrapper:
"""Class containing an actor and its metadata."""
actor: ActorHandle
metadata: ActorMetadata
@dataclass
class ActorConfig:
num_cpus: float
num_gpus: float
resources: Optional[Dict[str, float]]
init_args: Tuple
init_kwargs: Dict
class ActorGroupMethod:
def __init__(self, actor_group: "ActorGroup", method_name: str):
self.actor_group = weakref.ref(actor_group)
self._method_name = method_name
def __call__(self, *args, **kwargs):
raise TypeError(
"ActorGroup methods cannot be called directly. "
"Instead "
f"of running 'object.{self._method_name}()', try "
f"'object.{self._method_name}.remote()'."
)
def remote(self, *args, **kwargs):
return [
getattr(a.actor, self._method_name).remote(*args, **kwargs)
for a in self.actor_group().actors
]
@Deprecated(
message="For stateless/task processing, use ray.util.multiprocessing, see details "
f"in https://docs.ray.io/en/{get_ray_doc_version()}/ray-more-libs/multiprocessing.html. " # noqa: E501
"For stateful/actor processing such as batch prediction, use "
"Datasets.map_batches(compute=ActorPoolStrategy, ...), see details in "
f"https://docs.ray.io/en/{get_ray_doc_version()}/data/api/dataset.html#ray.data.Dataset.map_batches.", # noqa: E501
warning=True,
)
class ActorGroup:
"""Group of Ray Actors that can execute arbitrary functions.
``ActorGroup`` launches Ray actors according to the given
specification. It can then execute arbitrary Python functions in each of
these actors.
If not enough resources are available to launch the actors, the Ray
cluster will automatically scale up if autoscaling is enabled.
Args:
actor_cls: The class to use as the remote actors.
num_actors: The number of the provided Ray actors to
launch. Defaults to 1.
num_cpus_per_actor: The number of CPUs to reserve for each
actor. Fractional values are allowed. Defaults to 1.
num_gpus_per_actor: The number of GPUs to reserve for each
actor. Fractional values are allowed. Defaults to 0.
resources_per_actor: Dictionary specifying the resources that will be
requested for each actor in addition to ``num_cpus_per_actor``
and ``num_gpus_per_actor``.
init_args: Positional arguments forwarded to ``actor_cls`` when each
actor is created.
init_kwargs: Keyword arguments forwarded to ``actor_cls`` when each
actor is created.
"""
def __init__(
self,
actor_cls: Type,
num_actors: int = 1,
num_cpus_per_actor: float = 1,
num_gpus_per_actor: float = 0,
resources_per_actor: Optional[Dict[str, float]] = None,
init_args: Optional[Tuple] = None,
init_kwargs: Optional[Dict] = None,
):
from ray._common.usage.usage_lib import record_library_usage
record_library_usage("util.ActorGroup")
if num_actors <= 0:
raise ValueError(
"The provided `num_actors` must be greater "
f"than 0. Received num_actors={num_actors} "
f"instead."
)
if num_cpus_per_actor < 0 or num_gpus_per_actor < 0:
raise ValueError(
"The number of CPUs and GPUs per actor must "
"not be negative. Received "
f"num_cpus_per_actor={num_cpus_per_actor} and "
f"num_gpus_per_actor={num_gpus_per_actor}."
)
self.actors = []
self.num_actors = num_actors
self.actor_config = ActorConfig(
num_cpus=num_cpus_per_actor,
num_gpus=num_gpus_per_actor,
resources=resources_per_actor,
init_args=init_args or (),
init_kwargs=init_kwargs or {},
)
self._remote_cls = ray.remote(
num_cpus=self.actor_config.num_cpus,
num_gpus=self.actor_config.num_gpus,
resources=self.actor_config.resources,
)(actor_cls)
self.start()
def __getattr__(self, item):
if len(self.actors) == 0:
raise RuntimeError(
"This ActorGroup has been shutdown. Please start it again."
)
# Same implementation as actor.py
return ActorGroupMethod(self, item)
def __len__(self):
return len(self.actors)
def __getitem__(self, item):
return self.actors[item]
def start(self):
"""Starts all the actors in this actor group."""
if self.actors and len(self.actors) > 0:
raise RuntimeError(
"The actors have already been started. "
"Please call `shutdown` first if you want to "
"restart them."
)
logger.debug(f"Starting {self.num_actors} actors.")
self.add_actors(self.num_actors)
logger.debug(f"{len(self.actors)} actors have successfully started.")
def shutdown(self, patience_s: float = 5):
"""Shutdown all the actors in this actor group.
Args:
patience_s: Attempt a graceful shutdown
of the actors for this many seconds. Fallback to force kill
if graceful shutdown is not complete after this time. If
this is less than or equal to 0, immediately force kill all
actors.
"""
logger.debug(f"Shutting down {len(self.actors)} actors.")
if patience_s <= 0:
for actor in self.actors:
ray.kill(actor.actor)
else:
done_refs = [w.actor.__ray_terminate__.remote() for w in self.actors]
# Wait for actors to die gracefully.
done, not_done = ray.wait(done_refs, timeout=patience_s)
if not_done:
logger.debug("Graceful termination failed. Falling back to force kill.")
# If all actors are not able to die gracefully, then kill them.
for actor in self.actors:
ray.kill(actor.actor)
logger.debug("Shutdown successful.")
self.actors = []
def remove_actors(self, actor_indexes: List[int]):
"""Removes the actors with the specified indexes.
Args:
actor_indexes: The indexes of the actors to remove.
"""
new_actors = []
for i in range(len(self.actors)):
if i not in actor_indexes:
new_actors.append(self.actors[i])
self.actors = new_actors
def add_actors(self, num_actors: int):
"""Adds ``num_actors`` to this ActorGroup.
Args:
num_actors: The number of actors to add.
"""
new_actors = []
new_actor_metadata = []
for _ in range(num_actors):
actor = self._remote_cls.remote(
*self.actor_config.init_args, **self.actor_config.init_kwargs
)
new_actors.append(actor)
if hasattr(actor, "get_actor_metadata"):
new_actor_metadata.append(actor.get_actor_metadata.remote())
# Get metadata from all actors.
metadata = ray.get(new_actor_metadata)
if len(metadata) == 0:
metadata = [None] * len(new_actors)
for i in range(len(new_actors)):
self.actors.append(ActorWrapper(actor=new_actors[i], metadata=metadata[i]))
@property
def actor_metadata(self):
return [a.metadata for a in self.actors]
python/ray/util/actor_pool.py
+488
View File
@@ -0,0 +1,488 @@
from typing import TYPE_CHECKING, Any, Callable, List, Optional, TypeVar
import ray
from ray.util.annotations import DeveloperAPI
if TYPE_CHECKING:
import ray.actor
V = TypeVar("V")
@DeveloperAPI
class ActorPool:
"""Utility class to operate on a fixed pool of actors.
Arguments:
actors: List of Ray actor handles to use in this pool.
Examples:
.. testcode::
import ray
from ray.util.actor_pool import ActorPool
@ray.remote
class Actor:
def double(self, v):
return 2 * v
a1, a2 = Actor.remote(), Actor.remote()
pool = ActorPool([a1, a2])
print(list(pool.map(lambda a, v: a.double.remote(v),
[1, 2, 3, 4])))
.. testoutput::
[2, 4, 6, 8]
"""
def __init__(self, actors: list):
from ray._common.usage.usage_lib import record_library_usage
record_library_usage("util.ActorPool")
# 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: Callable[["ray.actor.ActorHandle", V], Any], values: List[V]):
"""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: Function that takes (actor, value) as argument and
returns an ObjectRef computing the result over the value. The
actor will be considered busy until the ObjectRef completes.
values: List of values that fn(actor, value) should be
applied to.
Returns:
Iterator over results from applying fn to the actors and values.
Examples:
.. testcode::
import ray
from ray.util.actor_pool import ActorPool
@ray.remote
class Actor:
def double(self, v):
return 2 * v
a1, a2 = Actor.remote(), Actor.remote()
pool = ActorPool([a1, a2])
print(list(pool.map(lambda a, v: a.double.remote(v),
[1, 2, 3, 4])))
.. testoutput::
[2, 4, 6, 8]
"""
# Ignore/Cancel all the previous submissions
# by calling `has_next` and `gen_next` repeteadly.
while self.has_next():
try:
self.get_next(timeout=0, ignore_if_timedout=True)
except TimeoutError:
pass
for v in values:
self.submit(fn, v)
def get_generator():
while self.has_next():
yield self.get_next()
return get_generator()
def map_unordered(
self, fn: Callable[["ray.actor.ActorHandle", V], Any], values: List[V]
):
"""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: Function that takes (actor, value) as argument and
returns an ObjectRef computing the result over the value. The
actor will be considered busy until the ObjectRef completes.
values: List of values that fn(actor, value) should be
applied to.
Returns:
Iterator over results from applying fn to the actors and values.
Examples:
.. testcode::
import ray
from ray.util.actor_pool import ActorPool
@ray.remote
class Actor:
def double(self, v):
return 2 * v
a1, a2 = Actor.remote(), Actor.remote()
pool = ActorPool([a1, a2])
print(list(pool.map_unordered(lambda a, v: a.double.remote(v),
[1, 2, 3, 4])))
.. testoutput::
:options: +MOCK
[6, 8, 4, 2]
"""
# Ignore/Cancel all the previous submissions
# by calling `has_next` and `gen_next_unordered` repeteadly.
while self.has_next():
try:
self.get_next_unordered(timeout=0)
except TimeoutError:
pass
for v in values:
self.submit(fn, v)
def get_generator():
while self.has_next():
yield self.get_next_unordered()
return get_generator()
def submit(self, fn: Callable[["ray.actor.ActorHandle", V], Any], value: V):
"""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: Function that takes (actor, value) as argument and
returns an ObjectRef computing the result over the value. The
actor will be considered busy until the ObjectRef completes.
value: Value to compute a result for.
Examples:
.. testcode::
import ray
from ray.util.actor_pool import ActorPool
@ray.remote
class Actor:
def double(self, v):
return 2 * v
a1, a2 = Actor.remote(), Actor.remote()
pool = ActorPool([a1, a2])
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())
.. testoutput::
2 4
"""
if self._idle_actors:
actor = self._idle_actors.pop()
future = fn(actor, value)
future_key = tuple(future) if isinstance(future, list) else future
self._future_to_actor[future_key] = (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:
.. testcode::
import ray
from ray.util.actor_pool import ActorPool
@ray.remote
class Actor:
def double(self, v):
return 2 * v
a1, a2 = Actor.remote(), Actor.remote()
pool = ActorPool([a1, a2])
pool.submit(lambda a, v: a.double.remote(v), 1)
print(pool.has_next())
print(pool.get_next())
print(pool.has_next())
.. testoutput::
True
2
False
"""
return bool(self._future_to_actor)
def get_next(
self,
timeout: Optional[float] = None,
ignore_if_timedout: bool = False,
):
"""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.
Arguments:
timeout: Max seconds to wait for the next result. ``None`` waits
indefinitely.
ignore_if_timedout: When True, drop the timed-out task and raise
``TimeoutError`` after advancing past it instead of leaving it
in place.
Returns:
The next result.
Raises:
TimeoutError: if the timeout is reached.
Examples:
.. testcode::
import ray
from ray.util.actor_pool import ActorPool
@ray.remote
class Actor:
def double(self, v):
return 2 * v
a1, a2 = Actor.remote(), Actor.remote()
pool = ActorPool([a1, a2])
pool.submit(lambda a, v: a.double.remote(v), 1)
print(pool.get_next())
.. testoutput::
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]
timeout_msg = "Timed out waiting for result"
raise_timeout_after_ignore = False
if timeout is not None:
res, _ = ray.wait([future], timeout=timeout)
if not res:
if not ignore_if_timedout:
raise TimeoutError(timeout_msg)
else:
raise_timeout_after_ignore = True
del self._index_to_future[self._next_return_index]
self._next_return_index += 1
future_key = tuple(future) if isinstance(future, list) else future
i, a = self._future_to_actor.pop(future_key)
self._return_actor(a)
if raise_timeout_after_ignore:
raise TimeoutError(
timeout_msg + ". The task {} has been ignored.".format(future)
)
return ray.get(future)
def get_next_unordered(
self,
timeout: Optional[float] = None,
ignore_if_timedout: bool = False,
):
"""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.
Arguments:
timeout: Max seconds to wait for the next result. ``None`` waits
indefinitely.
ignore_if_timedout: When True, drop the timed-out task and raise
``TimeoutError`` after advancing past it instead of leaving it
in place.
Returns:
The next result.
Raises:
TimeoutError: if the timeout is reached.
Examples:
.. testcode::
import ray
from ray.util.actor_pool import ActorPool
@ray.remote
class Actor:
def double(self, v):
return 2 * v
a1, a2 = Actor.remote(), Actor.remote()
pool = ActorPool([a1, a2])
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())
print(pool.get_next_unordered())
.. testoutput::
:options: +MOCK
4
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)
timeout_msg = "Timed out waiting for result"
raise_timeout_after_ignore = False
if res:
[future] = res
else:
if not ignore_if_timedout:
raise TimeoutError(timeout_msg)
else:
raise_timeout_after_ignore = True
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)
if raise_timeout_after_ignore:
raise TimeoutError(
timeout_msg + ". The task {} has been ignored.".format(future)
)
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))
def has_free(self):
"""Returns whether there are any idle actors available.
Returns:
True if there are any idle actors and no pending submits.
Examples:
.. testcode::
import ray
from ray.util.actor_pool import ActorPool
@ray.remote
class Actor:
def double(self, v):
return 2 * v
a1 = Actor.remote()
pool = ActorPool([a1])
pool.submit(lambda a, v: a.double.remote(v), 1)
print(pool.has_free())
print(pool.get_next())
print(pool.has_free())
.. testoutput::
False
2
True
"""
return len(self._idle_actors) > 0 and len(self._pending_submits) == 0
def pop_idle(self):
"""Removes an idle actor from the pool.
Returns:
An idle actor if one is available.
None if no actor was free to be removed.
Examples:
.. testcode::
import ray
from ray.util.actor_pool import ActorPool
@ray.remote
class Actor:
def double(self, v):
return 2 * v
a1 = Actor.remote()
pool = ActorPool([a1])
pool.submit(lambda a, v: a.double.remote(v), 1)
assert pool.pop_idle() is None
assert pool.get_next() == 2
assert pool.pop_idle() == a1
"""
if self.has_free():
return self._idle_actors.pop()
return None
def push(self, actor: "ray.actor.ActorHandle"):
"""Pushes a new actor into the current list of idle actors.
Arguments:
actor: The Ray actor handle to add to the pool's idle set.
Examples:
.. testcode::
import ray
from ray.util.actor_pool import ActorPool
@ray.remote
class Actor:
def double(self, v):
return 2 * v
a1, a2 = Actor.remote(), Actor.remote()
pool = ActorPool([a1])
pool.push(a2)
"""
busy_actors = []
if self._future_to_actor.values():
_, busy_actors = zip(*self._future_to_actor.values())
if actor in self._idle_actors or actor in busy_actors:
raise ValueError("Actor already belongs to current ActorPool")
else:
self._return_actor(actor)
python/ray/util/annotations.py
+346
View File
@@ -0,0 +1,346 @@
import inspect
import sys
import warnings
from enum import Enum
from functools import wraps
from typing import Any, Callable, Optional, TypeVar, cast, overload
# TypeVar for preserving function/class signatures through decorators.
# Note: These decorators also accept properties, but we use Callable for the
# common case. Properties work at runtime but won't get full type inference.
F = TypeVar("F", bound=Callable[..., Any])
class AnnotationType(Enum):
PUBLIC_API = "PublicAPI"
DEVELOPER_API = "DeveloperAPI"
DEPRECATED = "Deprecated"
UNKNOWN = "Unknown"
@overload
def PublicAPI(obj: F) -> F:
...
@overload
def PublicAPI(
*, stability: str = "stable", api_group: str = "Others"
) -> Callable[[F], F]:
...
def PublicAPI(*args: Any, **kwargs: Any):
"""Annotation for documenting public APIs.
Public APIs are classes and methods exposed to end users of Ray.
If ``stability="alpha"``, the API can be used by advanced users who are
tolerant to and expect breaking changes.
If ``stability="beta"``, the API is still public and can be used by early
users, but are subject to change.
If ``stability="stable"``, the APIs will remain backwards compatible across
minor Ray releases (e.g., Ray 1.4 -> 1.8).
For a full definition of the stability levels, please refer to the
:ref:`Ray API Stability definitions <api-stability>`.
Args:
*args: When used as a bare ``@PublicAPI`` decorator, contains the
wrapped function or class as the single positional argument.
**kwargs: Supported keyword arguments are ``stability`` (one of
``"stable"``, ``"beta"``, ``"alpha"``) and ``api_group`` (used
only for doc rendering; APIs in the same group are grouped
together in the API doc pages).
Returns:
Either the annotated object (when used as ``@PublicAPI``) or a
decorator that annotates an object (when used as
``@PublicAPI(...)``).
Examples:
>>> from ray.util.annotations import PublicAPI
>>> @PublicAPI
... def func(x):
... return x
>>> @PublicAPI(stability="beta")
... def func(y):
... return y
"""
if len(args) == 1 and len(kwargs) == 0 and callable(args[0]):
return PublicAPI(stability="stable", api_group="Others")(args[0])
if "stability" in kwargs:
stability = kwargs["stability"]
assert stability in ["stable", "beta", "alpha"], stability
else:
stability = "stable"
api_group = kwargs.get("api_group", "Others")
def wrap(obj: F) -> F:
if stability in ["alpha", "beta"]:
message = (
f"**PublicAPI ({stability}):** This API is in {stability} "
"and may change before becoming stable."
)
_append_doc(obj, message=message)
_mark_annotated(obj, type=AnnotationType.PUBLIC_API, api_group=api_group)
return obj
return wrap
@overload
def DeveloperAPI(obj: F) -> F:
...
@overload
def DeveloperAPI() -> Callable[[F], F]:
...
def DeveloperAPI(*args: Any, **kwargs: Any):
"""Annotation for documenting developer APIs.
Developer APIs are lower-level methods explicitly exposed to advanced Ray
users and library developers. Their interfaces may change across minor
Ray releases.
Args:
*args: When used as a bare ``@DeveloperAPI`` decorator, contains the
wrapped function or class as the single positional argument.
**kwargs: Reserved for future use; no keyword arguments are currently
supported.
Returns:
Either the annotated object (when used as ``@DeveloperAPI``) or a
decorator that annotates an object (when used as
``@DeveloperAPI()``).
Examples:
>>> from ray.util.annotations import DeveloperAPI
>>> @DeveloperAPI
... def func(x):
... return x
"""
if len(args) == 1 and len(kwargs) == 0 and callable(args[0]):
return DeveloperAPI()(args[0])
def wrap(obj: F) -> F:
_append_doc(
obj,
message="**DeveloperAPI:** This API may change across minor Ray releases.",
)
_mark_annotated(obj, type=AnnotationType.DEVELOPER_API)
return obj
return wrap
class RayDeprecationWarning(DeprecationWarning):
"""Specialized Deprecation Warning for fine grained filtering control"""
pass
# By default, print the first occurrence of matching warnings for
# each module where the warning is issued (regardless of line number)
if not sys.warnoptions:
warnings.filterwarnings("module", category=RayDeprecationWarning)
@overload
def Deprecated(obj: F) -> F:
...
@overload
def Deprecated(*, message: str = ..., warning: bool = False) -> Callable[[F], F]:
...
def Deprecated(*args: Any, **kwargs: Any):
"""Annotation for documenting a deprecated API.
Deprecated APIs may be removed in future releases of Ray.
Args:
*args: When used as a bare ``@Deprecated`` decorator, contains the
wrapped function or class as the single positional argument.
**kwargs: Supported keyword arguments are ``message`` (a string to
help users understand the reason for the deprecation and provide
a migration path) and ``warning`` (whether to also emit a
``RayDeprecationWarning`` at runtime; defaults to ``False``).
Returns:
Either the annotated object (when used as ``@Deprecated``) or a
decorator that annotates an object (when used as
``@Deprecated(...)``).
Examples:
>>> from ray.util.annotations import Deprecated
>>> @Deprecated
... def func(x):
... return x
>>> @Deprecated(message="g() is deprecated because the API is error "
... "prone. Please call h() instead.")
... def g(y):
... return y
"""
if len(args) == 1 and len(kwargs) == 0 and callable(args[0]):
return Deprecated()(args[0])
doc_message = (
"**DEPRECATED**: This API is deprecated and may be removed "
"in future Ray releases."
)
warning_message = (
"This API is deprecated and may be removed in future Ray releases. "
"You could suppress this warning by setting env variable "
'PYTHONWARNINGS="ignore::DeprecationWarning"'
)
warning = kwargs.pop("warning", False)
if "message" in kwargs:
doc_message = doc_message + "\n" + kwargs["message"]
warning_message = warning_message + "\n" + kwargs["message"]
del kwargs["message"]
if kwargs:
raise ValueError("Unknown kwargs: {}".format(kwargs.keys()))
def inner(obj: F) -> F:
_append_doc(obj, message=doc_message, directive="warning")
_mark_annotated(obj, type=AnnotationType.DEPRECATED)
if not warning:
return obj
if inspect.isclass(obj):
obj_init = obj.__init__
def patched_init(*args, **kwargs):
warnings.warn(warning_message, RayDeprecationWarning, stacklevel=2)
return obj_init(*args, **kwargs)
obj.__init__ = patched_init
return obj
else:
# class method or function.
def wrapper(*args, **kwargs):
warnings.warn(warning_message, RayDeprecationWarning, stacklevel=2)
return obj(*args, **kwargs)
# Only apply @wraps for actual callables, not properties/descriptors.
# Setting __wrapped__ on a property causes inspect.unwrap() to return
# the property, which breaks inspect.signature() in the tracing helper.
if callable(obj):
wrapper = wraps(obj)(wrapper)
return cast(F, wrapper)
return inner
def _append_doc(obj, *, message: str, directive: Optional[str] = None) -> None:
if not obj.__doc__:
obj.__doc__ = ""
obj.__doc__ = obj.__doc__.rstrip()
indent = _get_indent(obj.__doc__)
obj.__doc__ += "\n\n"
if directive is not None:
obj.__doc__ += f"{' ' * indent}.. {directive}::\n\n"
message = message.replace("\n", "\n" + " " * (indent + 4))
obj.__doc__ += f"{' ' * (indent + 4)}{message}"
else:
message = message.replace("\n", "\n" + " " * (indent + 4))
obj.__doc__ += f"{' ' * indent}{message}"
obj.__doc__ += f"\n{' ' * indent}"
def _get_indent(docstring: str) -> int:
"""Return the indentation level (in spaces) of the docstring body.
Args:
docstring: The docstring whose body indentation should be measured.
Returns:
The number of leading whitespace characters on the second non-empty
line of the docstring (i.e. the indentation of the body), or 0 if
the docstring is empty or contains only a summary line.
Example:
>>> def f():
... '''Docstring summary.'''
>>> f.__doc__
'Docstring summary.'
>>> _get_indent(f.__doc__)
0
>>> def g(foo):
... '''Docstring summary.
...
... Args:
... foo: Does bar.
... '''
>>> g.__doc__
'Docstring summary.\\n\\n Args:\\n foo: Does bar.\\n '
>>> _get_indent(g.__doc__)
4
>>> class A:
... def h():
... '''Docstring summary.
...
... Returns:
... None.
... '''
>>> A.h.__doc__
'Docstring summary.\\n\\n Returns:\\n None.\\n '
>>> _get_indent(A.h.__doc__)
8
"""
if not docstring:
return 0
non_empty_lines = [line for line in docstring.splitlines() if line]
if len(non_empty_lines) == 1:
# Docstring contains summary only.
return 0
# The docstring summary isn't indented, so check the indentation of the second
# non-empty line.
return len(non_empty_lines[1]) - len(non_empty_lines[1].lstrip())
def _mark_annotated(
obj, type: AnnotationType = AnnotationType.UNKNOWN, api_group="Others"
) -> None:
# Set magic token for check_api_annotations linter.
if hasattr(obj, "__name__"):
obj._annotated = obj.__name__
obj._annotated_type = type
obj._annotated_api_group = api_group
def _is_annotated(obj) -> bool:
# Check the magic token exists and applies to this class (not a subclass).
return hasattr(obj, "_annotated") and obj._annotated == obj.__name__
def _get_annotation_type(obj) -> Optional[str]:
if not _is_annotated(obj):
return None
return obj._annotated_type.value
python/ray/util/check_open_ports.py
+178
View File
@@ -0,0 +1,178 @@
"""A CLI utility for check open ports in the Ray cluster.
See https://www.anyscale.com/blog/update-on-ray-cve-2023-48022-new-verification-tooling-available # noqa: E501
for more details.
"""
import json
import subprocess
import urllib
from typing import List, Tuple
import click
import ray
from ray.autoscaler._private.cli_logger import add_click_logging_options, cli_logger
from ray.autoscaler._private.constants import RAY_PROCESSES
from ray.util.annotations import PublicAPI
import psutil
def _get_ray_ports() -> List[int]:
unique_ports = set()
process_infos = []
for proc in psutil.process_iter(["name", "cmdline"]):
try:
process_infos.append((proc, proc.name(), proc.cmdline()))
except psutil.Error:
pass
for keyword, filter_by_cmd in RAY_PROCESSES:
for candidate in process_infos:
proc, proc_cmd, proc_args = candidate
corpus = proc_cmd if filter_by_cmd else subprocess.list2cmdline(proc_args)
if keyword in corpus:
try:
for connection in proc.connections():
if connection.status == psutil.CONN_LISTEN:
unique_ports.add(connection.laddr.port)
except psutil.AccessDenied:
cli_logger.info(
"Access denied to process connections for process,"
" worker process probably restarted",
proc,
)
return sorted(unique_ports)
def _check_for_open_ports_from_internet(
service_url: str, ports: List[int]
) -> Tuple[List[int], List[int]]:
request = urllib.request.Request(
method="POST",
url=service_url,
headers={
"Content-Type": "application/json",
"X-Ray-Open-Port-Check": "1",
},
data=json.dumps({"ports": ports}).encode("utf-8"),
)
response = urllib.request.urlopen(request)
if response.status != 200:
raise RuntimeError(
f"Failed to check with Ray Open Port Service: {response.status}"
)
response_body = json.load(response)
publicly_open_ports = response_body.get("open_ports", [])
checked_ports = response_body.get("checked_ports", [])
return publicly_open_ports, checked_ports
def _check_if_exposed_to_internet(
service_url: str,
) -> Tuple[List[int], List[int]]:
return _check_for_open_ports_from_internet(service_url, _get_ray_ports())
def _check_ray_cluster(
service_url: str,
) -> List[Tuple[str, Tuple[List[int], List[int]]]]:
ray.init(ignore_reinit_error=True)
@ray.remote(num_cpus=0)
def check(node_id, service_url):
return node_id, _check_if_exposed_to_internet(service_url)
ray_node_ids = [node["NodeID"] for node in ray.nodes() if node["Alive"]]
cli_logger.info(
f"Cluster has {len(ray_node_ids)} node(s)."
" Scheduling tasks on each to check for exposed ports",
)
per_node_tasks = {
node_id: (
check.options(label_selector={ray._raylet.RAY_NODE_ID_KEY: node_id}).remote(
node_id, service_url
)
)
for node_id in ray_node_ids
}
results = []
for node_id, per_node_task in per_node_tasks.items():
try:
results.append(ray.get(per_node_task))
except Exception as e:
cli_logger.info(f"Failed to check on node {node_id}: {e}")
return results
@click.command()
@click.option(
"--yes", "-y", is_flag=True, default=False, help="Don't ask for confirmation."
)
@click.option(
"--service-url",
required=False,
type=str,
default="https://ray-open-port-checker.uc.r.appspot.com/open-port-check",
help="The url of service that checks whether submitted ports are open.",
)
@add_click_logging_options
@PublicAPI
def check_open_ports(yes, service_url):
"""Check open ports in the local Ray cluster."""
if not cli_logger.confirm(
yes=yes,
msg=(
"Do you want to check the local Ray cluster"
" for any nodes with ports accessible to the internet?"
),
_default=True,
):
cli_logger.info("Exiting without checking as instructed")
return
cluster_open_ports = _check_ray_cluster(service_url)
public_nodes = []
for node_id, (open_ports, checked_ports) in cluster_open_ports:
if open_ports:
cli_logger.info(
f"[🛑] open ports detected open_ports={open_ports!r} node={node_id!r}"
)
public_nodes.append((node_id, open_ports, checked_ports))
else:
cli_logger.info(
f"[🟢] No open ports detected "
f"checked_ports={checked_ports!r} node={node_id!r}"
)
cli_logger.info("Check complete, results:")
if public_nodes:
cli_logger.info(
"""
[🛑] An server on the internet was able to open a connection to one of this Ray
cluster's public IP on one of Ray's internal ports. If this is not a false
positive, this is an extremely unsafe configuration for Ray to be running in.
Ray is not meant to be exposed to untrusted clients and will allow them to run
arbitrary code on your machine.
You should take immediate action to validate this result and if confirmed shut
down your Ray cluster immediately and take appropriate action to remediate its
exposure. Anything either running on this Ray cluster or that this cluster has
had access to could be at risk.
For guidance on how to operate Ray safely, please review [Ray's security
documentation](https://docs.ray.io/en/latest/ray-security/index.html).
""".strip()
)
else:
cli_logger.info("[🟢] No open ports detected from any Ray nodes")
python/ray/util/check_serialize.py
+297
View File
@@ -0,0 +1,297 @@
"""A utility for debugging serialization issues."""
import inspect
from contextlib import contextmanager
from typing import Any, Optional, Set, Tuple
import colorama
# Import ray first to use the bundled colorama
import ray # noqa: F401
import ray.cloudpickle as cp
from ray.util.annotations import DeveloperAPI
@contextmanager
def _indent(printer):
printer.level += 1
yield
printer.level -= 1
class _Printer:
def __init__(self, print_file):
self.level = 0
self.print_file = print_file
def indent(self):
return _indent(self)
def print(self, msg):
indent = " " * self.level
print(indent + msg, file=self.print_file)
@DeveloperAPI
class FailureTuple:
"""Represents the serialization 'frame'.
Attributes:
obj: The object that fails serialization.
path: Tuple of variable names representing the traversal path.
parent: The object that references the `obj`.
"""
def __init__(self, obj: Any, path: tuple[str, ...], parent: Any):
"""Initialize FailureTuple.
Args:
obj: The object that fails serialization.
path: Tuple of variable names representing the traversal path.
parent: The object that references the object.
"""
self.obj = obj
self.path = path
self.parent = parent
def __repr__(self):
formatted_path = [str(p) if p is not None else "unknown" for p in self.path]
path_str = " -> ".join(formatted_path)
var_name = formatted_path[-1] if formatted_path else "unknown"
return f"FailTuple({var_name} [path={path_str}, obj={self.obj}, parent={self.parent}])"
def _inspect_func_serialization(base_obj, depth, parent, failure_set, printer, path=()):
"""Adds the first-found non-serializable element to the failure_set."""
assert inspect.isfunction(base_obj)
printer.print(f"Inspecting closure of '{base_obj.__qualname__}':")
closure = inspect.getclosurevars(base_obj)
found = False
if closure.globals:
printer.print(
f"Detected {len(closure.globals)} global variables. "
"Checking serializability..."
)
with printer.indent():
for name, obj in closure.globals.items():
serializable, _ = _inspect_serializability(
obj,
name=name,
depth=depth - 1,
parent=parent,
failure_set=failure_set,
printer=printer,
path=path,
)
found = found or not serializable
if found:
break
if closure.nonlocals:
printer.print(
f"Detected {len(closure.nonlocals)} nonlocal variables. "
"Checking serializability..."
)
with printer.indent():
for name, obj in closure.nonlocals.items():
serializable, _ = _inspect_serializability(
obj,
name=name,
depth=depth - 1,
parent=parent,
failure_set=failure_set,
printer=printer,
path=path,
)
found = found or not serializable
if found:
break
if not found:
obj_name = getattr(base_obj, "__qualname__", None) or repr(base_obj)
printer.print(
f"WARNING: Did not find non-serializable object in {obj_name}. "
"This may be because the object is only non-serializable when "
"combined with other objects, or because cloudpickle and the "
"traversal disagree on which object causes the failure. "
"Try calling `ray.util.inspect_serializability` directly on "
"the suspected object, or simplify the closure to isolate the cause."
)
return found
def _inspect_generic_serialization(
base_obj, depth, parent, failure_set, printer, path=()
):
"""Adds the first-found non-serializable element to the failure_set."""
assert not inspect.isfunction(base_obj)
functions = inspect.getmembers(base_obj, predicate=inspect.isfunction)
found = False
with printer.indent():
for name, obj in functions:
serializable, _ = _inspect_serializability(
obj,
name=name,
depth=depth - 1,
parent=parent,
failure_set=failure_set,
printer=printer,
path=path,
)
found = found or not serializable
if found:
break
with printer.indent():
members = inspect.getmembers(base_obj)
for name, obj in members:
if name.startswith("__") and name.endswith("__") or inspect.isbuiltin(obj):
continue
serializable, _ = _inspect_serializability(
obj,
name=name,
depth=depth - 1,
parent=parent,
failure_set=failure_set,
printer=printer,
path=path,
)
found = found or not serializable
if found:
break
if not found:
obj_name = (
base_obj.__name__ if inspect.isclass(base_obj) else type(base_obj).__name__
)
printer.print(
f"WARNING: Did not find non-serializable object in "
f"{obj_name if inspect.isclass(base_obj) else f'{obj_name} instance'}. "
"This may be because the object is only non-serializable when "
"combined with other objects, or because cloudpickle and the "
"traversal disagree. "
"Try calling `ray.util.inspect_serializability` on individual "
"attributes to isolate the cause."
)
return found
@DeveloperAPI
def inspect_serializability(
base_obj: Any,
name: Optional[str] = None,
depth: int = 3,
print_file: Optional[Any] = None,
) -> Tuple[bool, Set[FailureTuple]]:
"""Identifies what objects are preventing serialization.
Args:
base_obj: Object to be serialized.
name: Optional name of string.
depth: Depth of the scope stack to walk through. Defaults to 3.
print_file: file argument that will be passed to print().
Returns:
bool: True if serializable.
set[FailureTuple]: Set of unserializable objects.
.. versionadded:: 1.1.0
"""
printer = _Printer(print_file)
return _inspect_serializability(base_obj, name, depth, None, None, printer)
def _inspect_serializability(
base_obj, name, depth, parent, failure_set, printer, path=()
) -> Tuple[bool, Set[FailureTuple]]:
colorama.init()
top_level = False
declaration = ""
found = False
if failure_set is None:
top_level = True
failure_set = set()
declaration = f"Checking Serializability of {base_obj}"
printer.print("=" * min(len(declaration), 80))
printer.print(declaration)
printer.print("=" * min(len(declaration), 80))
else:
printer.print(f"Serializing '{name}' {base_obj}...")
try:
cp.dumps(base_obj)
return True, failure_set
except Exception as e:
printer.print(
f"{colorama.Fore.RED}!!! FAIL{colorama.Fore.RESET} " f"serialization: {e}"
)
found = True
try:
if depth == 0:
failure_path = path + ((name,) if name is not None else ())
failure_set.add(FailureTuple(base_obj, failure_path, parent))
except Exception:
pass
if depth <= 0:
return False, failure_set
# TODO: we only differentiate between 'function' and 'object'
# but we should do a better job of diving into something
# more specific like a Type, Object, etc.
if inspect.isfunction(base_obj):
_inspect_func_serialization(
base_obj,
depth=depth,
parent=base_obj,
failure_set=failure_set,
printer=printer,
path=path + ((name,) if name is not None else ()),
)
else:
_inspect_generic_serialization(
base_obj,
depth=depth,
parent=base_obj,
failure_set=failure_set,
printer=printer,
path=path + ((name,) if name is not None else ()),
)
if not failure_set:
failure_path = path + ((name,) if name is not None else ())
failure_set.add(FailureTuple(base_obj, failure_path, parent))
if top_level:
printer.print("=" * min(len(declaration), 80))
if not failure_set:
printer.print(
"Nothing failed the inspect_serialization test, though "
"serialization did not succeed."
)
else:
fail_vars = (
f"\n\n\t{colorama.Style.BRIGHT}"
+ "\n".join(str(k) for k in failure_set)
+ f"{colorama.Style.RESET_ALL}\n\n"
)
printer.print(
f"Variable: {fail_vars}was found to be non-serializable. "
"There may be multiple other undetected variables that were "
"non-serializable. "
)
printer.print(
"Consider either removing the "
"instantiation/imports of these variables or moving the "
"instantiation into the scope of the function/class. "
)
printer.print("=" * min(len(declaration), 80))
printer.print(
"Check https://docs.ray.io/en/master/ray-core/objects/serialization.html#troubleshooting for more information." # noqa
)
printer.print(
"If you have any suggestions on how to improve "
"this error message, please reach out to the "
"Ray developers on github.com/ray-project/ray/issues/"
)
printer.print("=" * min(len(declaration), 80))
return not found, failure_set
python/ray/util/client/ARCHITECTURE.md
+133
View File
@@ -0,0 +1,133 @@
# Ray Client Architecture Guide
A quick development primer on the codebase layout
## General
The Ray client is a gRPC client and server.
The server runs `ray.init()` and acts like a normal Ray driver, controlled by the gRPC connection.
It does all the bookkeeping and keeps things in scope for the clients that connect.
Generally, the client side lives in `ray/util/client` and the server lives in `ray/util/client/server`.
By convention, the `ray/util/client` avoids importing `ray` directly, but the server side, being just another Ray application, is allowed to do so.
This separation exists both for dependency cycle reasons and also to, if desired in the future, pull either portion out into its own repo or sub-installation.
(eg, `pip install ray_client`)
The `ray` global variable of type `RayAPIStub` in [`ray/util/client/__init__.py`](./__init__.py) acts as the equivalent API surface as does the `ray` package.
Functions in the `ray` namespace are methods on the RayAPIStub object.
For many of the objects in the root `ray` namespace, there is an equivalent client object. These are mostly contained in [`ray/util/client/common.py`](./common.py).
These objects are client stand-ins for their server-side objects. For example:
```
ObjectRef <-> ClientObjectRef
ActorID <-> ClientActorRef
RemoteFunc <-> ClientRemoteFunc
```
This means that, if the type of the object you're looking at (say, in a bug report) is a ClientObjectRef, it should have come from the client code (ie, constructed by returning from the server).
How the two interchange is talked about under Protocol.
## Protocol
There's one gRPC spec for the client, and that lives at [`/src/ray/protobuf/ray_client.proto`](/src/ray/protobuf/ray_client.proto).
There's another protocol at play, however, and that is _the way in which functions and data are encoded_.
This is particularly important in the context of the Ray client; since both ends are Python, this is a `pickle`.
The key separation to understand is that `pickle` is how client objects, including the client-side stubs, get serialized and transported, opaquely, to the server.
The gRPC service is the API surface to implement remote, thin, client functionality.
The `pickle` protocol is how the data is encoded for that API.
### gRPC services
The gRPC side is the most straightforward and easiest to follow from the proto file.
#### get, put, and function calls
Client started life as a set of unary RPCs, with just enough functionality to implement the most-used APIs.
As an introduction to the RPC API, they're a good place to start to understand how the protocol works.
The proto file is well-commented with every field describing what it does.
The Unary RPCs are still around, but they are ripe for deprecation.
The problem they have is that they are not tied to a persistent connection.
If you imagine a load balancer in front of a couple client-servers, then any client could hit any state on any server with a unary RPC.
As we need to keep handles to ray ObjectRefs and similar so that they don't go out of scope and dropped, these must stay in sync for connected clients.
With unary RPCs, that means one RPC could go to one server (say, a `x = f.remote()`), and the follow up (`ray.get(x)`) wouldn't have the corresponding ObjectRef on the other server.
Get, Put, and Wait are pretty standard.
The more interesting one is Schedule, which implies a Put before it (the function to execute) and then executes it.
#### Data Channel
Which brings us to the data channel.
The data channel is a bidirectional streaming connection for the client.
It wraps all the same Request/Response patterns as the Unary RPCs.
At the start, the client associates itself with the server with a UUID-generated ClientID.
As long as the channel is open, the client is connected.
Tracking the ClientID then allows us to track all the resources we're holding for a particular client, and we know if the client has disconnected (the channel drops).
It's also through this mechanism we can do reference counting.
The client can keep track of how many references it has to various Ray client objects, and, if they fall out of scope, can send a `ReleaseRequest` to the server to optimistically clean up after itself.
Otherwise, all reference counting is done on the client side and the server only needs to know when to clean up.
The server can also clean up all references held for a client whenever it thinks it's safe to do so.
In the future, having an explicit "ClientDisconnection" message may help here, to delineate between a client that's intentionally done and will never come back and one that's experiencing a connectivity issue.
#### Logs Channel
Similar to the data channel, there's an associated logs channel which will pipe logs back to the client.
It's a separate channel as it's ancillary to the continued connection of the client, and if some logs are dropped due to disconnection, that's generally okay.
It's also then a separate way for a log aggregator to connect without implementing the full API.
This is also a bidirectional stream, where the client sends messages to control the verbosity and type of content, and the server streams back all the logs as they come in.
#### On CloudPickle
As per the introduction to this section, `pickle` and `cloudpickle` are the way to encode to Python-executable data for the generic transport of gRPC.
Ray Client provides its own pickle/unpickle subclasses in `client_pickler.py` and `server/server_pickler.py`.
The reason it has its own subclasses is to solve the problem of mixing `Client*` stub objects.
This is easier to describe by example.
Suppose a `RemoteFunc`, `f()` calls another `RemoteFunc`, `g()`
`f()` has to have a reference to `g()`, and knows it's a `RemoteFunc` (calls `g.remote()`, say) and so the function `f()` gets serialized with pickle, and in that serialization data is an object of class `RemoteFunc` to be deserialized on the worker side.
In Ray client, `f()` and `g()` are `ClientRemoteFunc`s and work the same way.
In early versions of the Ray Client, a `ClientRemoteFunc` had to know whether it was on the server or client side, and either run like a normal `RemoteFunc` (on the server) or a client call on the client.
This led to some interesting bugs where passing, or especially returning, client-side stub objects around.
(Imagine if `f()` calls `g()` which builds and returns a new closure `h()`)
To simplify all of this, the pickler subclasses were written.
Now, whenever a Client-stub-object is serialized, a struct (as of writing, a tuple) is stored in its place, and when deserialized, the server "fills in" the appropriate non-stub object.
And vice versa -- if the server is encoding a return/response that is an `ObjectRef`, a tuple is passed on the wire instead and the deserializer on the client side turns it back into a `ClientObjectRef`
This means that the client side deals as much as possible in its stub objects, and the server side never sees a stub object, and there is a clean separation between the two. Now, a `ClientObjectRef` existing on the server is an error case, and not a case to be handled specially.
It also means that the server side works just like normal Ray and deals in the normal Ray objects and can encode them transparently into client objects as they are sent.
Because never the twain shall meet, it's much easier to model and debug what's going on.
## Integration points with Ray core
In order to provide a seamless client experience with Ray core, we need to wrap some of the core Ray functions (eg, `ray.get()`).
Python's dynamic nature helps us here. As mentioned, the `RayAPIStub` is a class, not a module, which is a subtle difference.
This also allows us to have a `__getattr__` on the API level to redirect whereever we'd like, which we can't do in modules ([at least until Python 3.6 is deprecated](https://www.python.org/dev/peps/pep-0562/))
If the `ray` core were an object instead of functions in a namespace, we wouldn't need to wrap them to integrate, we'd simply swap the implementation.
But we have backwards compatibility to maintain.
All the interesting integration points with Ray core live within `ray/_private/client_mode_hook.py`.
In that file are contextmanagers and decorators meant to wrap Ray core functions.
If `client_mode_should_convert()` returns `True`, based on the environment variables as they've been set, then the decorators spring into action, and forward the calls to the `ray/util/client` object.
## Testing
There are two primary ways to test client code.
The first is to approach it from the context of knowing that we're testing a client/server app.
We can run both ends of the connection, call the client side (that we know to be the client side), and see the effect on the server and vice-versa.
The set of tests of the form `test_client*.py` take this approach.
The other way to approach them is as a fixture where we're testing the API of Ray, with Ray's own tests, _as though it were normal Ray_.
It's a highly powerful pattern, in that a test passing there means that the user can feel confident that things work the way they always have, client or not.
It's also a more difficult pattern to implement, as hooking the setup and shutdown of a client/server pair and a single-node ray instance are different, especially when these tests may make assumptions about how the fixtures were set up in the first place.
It is, however, the only way to test the integration points.
So generally speaking, if it's implementing a feature that makes the client work, it probably should be a test in the `test_client` series where one controls both ends and tests the client code.
If it's fixing a user-side API bug or an integration with Ray core, it's probably adapting or including a pre-existing unit test as part of Ray Client.
python/ray/util/client/__init__.py
+424
View File
@@ -0,0 +1,424 @@
import logging
import os
import threading
from typing import Any, Dict, List, Optional, Tuple
import ray._private.ray_constants as ray_constants
from ray._common.network_utils import build_address, get_localhost_ip
from ray._private.client_mode_hook import (
_explicitly_disable_client_mode,
_explicitly_enable_client_mode,
)
from ray._private.ray_logging import setup_logger
from ray._private.utils import check_version_info
from ray.job_config import JobConfig
from ray.util.annotations import DeveloperAPI
logger = logging.getLogger(__name__)
def _apply_uv_hook_for_client(
runtime_env: Optional[Dict[str, Any]],
) -> Optional[Dict[str, Any]]:
"""Apply UV runtime env hook on client side before connection.
UV (https://docs.astral.sh/uv/) is a modern Python package manager that
manages dependencies via pyproject.toml and uv.lock files. This function
detects when the client is running under 'uv run' and automatically
propagates the UV configuration to cluster workers so they can install
the same dependencies.
How it works:
1. Detects 'uv run' in the parent process tree
2. Extracts UV command-line arguments (e.g., --python, --locked)
3. Sets py_executable to 'uv run [args]' in runtime_env
4. Workers will use this UV command to install dependencies
Precedence rules:
- If user provides py_executable, UV hook is skipped entirely to avoid
unintended side effects (e.g., auto-setting working_dir)
- User-provided working_dir is preserved when UV hook runs
- Other runtime_env settings are merged with UV config
Feature flag:
Controlled by RAY_ENABLE_UV_RUN_RUNTIME_ENV constant (default: enabled)
Args:
runtime_env: The runtime environment dict to potentially modify.
Can be None if no runtime_env was specified.
Returns:
Modified runtime_env dict with UV configuration if detected,
otherwise the original runtime_env unchanged. Returns None if
input was None.
Raises:
RuntimeError: If UV environment is detected but configuration is invalid
(e.g., pyproject.toml not in working_dir, conflicting runtime_env).
Validation errors fail fast to provide clear feedback.
Note:
ImportError and other environmental errors are caught and logged,
allowing connection to proceed without UV propagation.
Example:
Client running under: uv run --python 3.11 my_script.py
>>> runtime_env = {"working_dir": "/tmp/myapp"}
>>> result = _apply_uv_hook_for_client(runtime_env)
>>> result
{'working_dir': '/tmp/myapp', 'py_executable': 'uv run --python 3.11'}
See Also:
- Issue: https://github.com/ray-project/ray/issues/57991
- UV docs: https://docs.astral.sh/uv/
"""
if not ray_constants.RAY_ENABLE_UV_RUN_RUNTIME_ENV:
return runtime_env
# If user provided py_executable, skip UV hook entirely to avoid side effects
# (e.g., auto-setting working_dir which triggers unwanted directory upload)
if runtime_env and "py_executable" in runtime_env:
logger.debug(
"User-provided py_executable found, skipping UV hook to avoid "
"unintended runtime_env modifications"
)
return runtime_env
# Import hook here (not at module level) to:
# 1. Avoid circular import issues with ray._private modules
# 2. Only load UV hook code when feature flag is enabled
from ray._private.runtime_env.uv_runtime_env_hook import hook
try:
result = hook(runtime_env)
except Exception as e:
raise RuntimeError(
f"Failed to apply UV runtime env hook for Ray Client: {e} "
"If you want the driver to use UV without propagating to workers, "
"set RAY_ENABLE_UV_RUN_RUNTIME_ENV=0."
) from e
if "py_executable" in result:
# UV environment was detected and applied by the hook
logger.debug(
f"UV environment detected for Ray Client: "
f"py_executable={result['py_executable']}"
)
return result
return runtime_env
class _ClientContext:
def __init__(self):
from ray.util.client.api import _ClientAPI
self.api = _ClientAPI()
self.client_worker = None
self._server = None
self._connected_with_init = False
self._inside_client_test = False
def connect(
self,
conn_str: str,
job_config: JobConfig = None,
secure: bool = False,
metadata: List[Tuple[str, str]] = None,
connection_retries: int = 3,
namespace: str = None,
*,
ignore_version: bool = False,
_credentials: Optional["grpc.ChannelCredentials"] = None, # noqa: F821
ray_init_kwargs: Optional[Dict[str, Any]] = None,
) -> Dict[str, Any]:
"""Connect the Ray Client to a server.
Args:
conn_str: Connection string, in the form "[host]:port"
job_config: The job config of the server.
secure: Whether to use a TLS secured gRPC channel
metadata: gRPC metadata to send on connect
connection_retries: number of connection attempts to make
namespace: The namespace to connect to.
ignore_version: whether to ignore Python or Ray version mismatches.
This should only be used for debugging purposes.
_credentials: Optional gRPC channel credentials for secure connection.
ray_init_kwargs: Optional additional keyword arguments for ray.init().
Returns:
Dictionary of connection info, e.g., {"num_clients": 1}.
"""
# Delay imports until connect to avoid circular imports.
from ray.util.client.worker import Worker
if self.client_worker is not None:
if self._connected_with_init:
return
raise Exception("ray.init() called, but ray client is already connected")
if not self._inside_client_test:
# If we're calling a client connect specifically and we're not
# currently in client mode, ensure we are.
_explicitly_enable_client_mode()
if namespace is not None:
job_config = job_config or JobConfig()
job_config.set_ray_namespace(namespace)
logging_level = ray_constants.LOGGER_LEVEL
logging_format = ray_constants.LOGGER_FORMAT
if ray_init_kwargs is None:
ray_init_kwargs = {}
# Apply UV hook client-side before connection.
# UV detection must happen on client side where 'uv run' process exists.
# See: https://github.com/ray-project/ray/issues/57991
#
# Runtime env can come from two sources:
# 1. ray_init_kwargs["runtime_env"] - directly passed to connect()
# 2. job_config.runtime_env - passed via JobConfig object
# We need to handle both sources and update them appropriately after UV hook.
runtime_env = ray_init_kwargs.get("runtime_env")
if runtime_env is None and job_config and job_config.runtime_env is not None:
runtime_env = job_config.runtime_env
runtime_env = _apply_uv_hook_for_client(runtime_env)
if runtime_env is not None:
# Update both ray_init_kwargs and job_config with UV modifications.
# This is necessary because _server_init() reads runtime_env from
# job_config.runtime_env, not from ray_init_kwargs["runtime_env"].
ray_init_kwargs["runtime_env"] = runtime_env
if job_config:
job_config.set_runtime_env(runtime_env)
# NOTE(architkulkarni): Custom env_hook is not supported with Ray Client.
# However, UV hook is now applied client-side above.
ray_init_kwargs["_skip_env_hook"] = True
if ray_init_kwargs.get("logging_level") is not None:
logging_level = ray_init_kwargs["logging_level"]
if ray_init_kwargs.get("logging_format") is not None:
logging_format = ray_init_kwargs["logging_format"]
setup_logger(logging_level, logging_format)
try:
self.client_worker = Worker(
conn_str,
secure=secure,
_credentials=_credentials,
metadata=metadata,
connection_retries=connection_retries,
)
self.api.worker = self.client_worker
self.client_worker._server_init(job_config, ray_init_kwargs)
conn_info = self.client_worker.connection_info()
self._check_versions(conn_info, ignore_version)
self._register_serializers()
return conn_info
except Exception:
self.disconnect()
raise
def _register_serializers(self):
"""Register the custom serializer addons at the client side.
The server side should have already registered the serializers via
regular worker's serialization_context mechanism.
"""
import ray.util.serialization_addons
from ray.util.serialization import StandaloneSerializationContext
ctx = StandaloneSerializationContext()
ray.util.serialization_addons.apply(ctx)
def _check_versions(self, conn_info: Dict[str, Any], ignore_version: bool) -> None:
# conn_info has "python_version" and "ray_version" so it can be used to compare.
ignore_version = ignore_version or ("RAY_IGNORE_VERSION_MISMATCH" in os.environ)
check_version_info(
conn_info,
"Ray Client",
raise_on_mismatch=not ignore_version,
python_version_match_level="minor",
)
def disconnect(self):
"""Disconnect the Ray Client."""
from ray.util.client.api import _ClientAPI
if self.client_worker is not None:
self.client_worker.close()
self.api = _ClientAPI()
self.client_worker = None
# remote can be called outside of a connection, which is why it
# exists on the same API layer as connect() itself.
def remote(self, *args: Any, **kwargs: Any):
"""remote is the hook stub passed on to replace `ray.remote`.
This sets up remote functions or actors, as the decorator,
but does not execute them.
Args:
*args: opaque arguments forwarded to ``_ClientAPI.remote``.
**kwargs: opaque keyword arguments forwarded to ``_ClientAPI.remote``.
Returns:
A client-side stub for the remote function or actor, or a
decorator that produces one when applied.
"""
return self.api.remote(*args, **kwargs)
def __getattr__(self, key: str):
if self.is_connected():
return getattr(self.api, key)
elif key in ["is_initialized", "_internal_kv_initialized"]:
# Client is not connected, thus Ray is not considered initialized.
return lambda: False
else:
raise Exception(
"Ray Client is not connected. Please connect by calling `ray.init`."
)
def is_connected(self) -> bool:
if self.client_worker is None:
return False
return self.client_worker.is_connected()
def init(self, *args, **kwargs):
if self._server is not None:
raise Exception("Trying to start two instances of ray via client")
import ray.util.client.server.server as ray_client_server
server_handle, address_info = ray_client_server.init_and_serve(
get_localhost_ip(), 50051, *args, **kwargs
)
self._server = server_handle.grpc_server
self.connect(build_address(get_localhost_ip(), 50051))
self._connected_with_init = True
return address_info
def shutdown(self, _exiting_interpreter=False):
self.disconnect()
import ray.util.client.server.server as ray_client_server
if self._server is None:
return
ray_client_server.shutdown_with_server(self._server, _exiting_interpreter)
self._server = None
# All connected context will be put here
# This struct will be guarded by a lock for thread safety
_all_contexts = set()
_lock = threading.Lock()
# This is the default context which is used when allow_multiple is not True
_default_context = _ClientContext()
@DeveloperAPI
class RayAPIStub:
"""This class stands in as the replacement API for the `import ray` module.
Much like the ray module, this mostly delegates the work to the
_client_worker. As parts of the ray API are covered, they are piped through
here or on the client worker API.
"""
def __init__(self):
self._cxt = threading.local()
self._cxt.handler = _default_context
self._inside_client_test = False
def get_context(self):
try:
return self._cxt.__getattribute__("handler")
except AttributeError:
self._cxt.handler = _default_context
return self._cxt.handler
def set_context(self, cxt):
old_cxt = self.get_context()
if cxt is None:
self._cxt.handler = _ClientContext()
else:
self._cxt.handler = cxt
return old_cxt
def is_default(self):
return self.get_context() == _default_context
def connect(self, *args, **kw_args):
self.get_context()._inside_client_test = self._inside_client_test
conn = self.get_context().connect(*args, **kw_args)
global _lock, _all_contexts
with _lock:
_all_contexts.add(self._cxt.handler)
return conn
def disconnect(self, *args, **kw_args):
global _lock, _all_contexts, _default_context
with _lock:
if _default_context == self.get_context():
for cxt in _all_contexts:
cxt.disconnect(*args, **kw_args)
_all_contexts = set()
else:
self.get_context().disconnect(*args, **kw_args)
if self.get_context() in _all_contexts:
_all_contexts.remove(self.get_context())
if len(_all_contexts) == 0:
_explicitly_disable_client_mode()
def remote(self, *args, **kwargs):
return self.get_context().remote(*args, **kwargs)
def __getattr__(self, name):
return self.get_context().__getattr__(name)
def is_connected(self, *args, **kwargs):
return self.get_context().is_connected(*args, **kwargs)
def init(self, *args, **kwargs):
ret = self.get_context().init(*args, **kwargs)
global _lock, _all_contexts
with _lock:
_all_contexts.add(self._cxt.handler)
return ret
def shutdown(self, *args, **kwargs):
global _lock, _all_contexts
with _lock:
if _default_context == self.get_context():
for cxt in _all_contexts:
cxt.shutdown(*args, **kwargs)
_all_contexts = set()
else:
self.get_context().shutdown(*args, **kwargs)
if self.get_context() in _all_contexts:
_all_contexts.remove(self.get_context())
if len(_all_contexts) == 0:
_explicitly_disable_client_mode()
ray = RayAPIStub()
@DeveloperAPI
def num_connected_contexts():
"""Return the number of client connections active."""
global _lock, _all_contexts
with _lock:
return len(_all_contexts)
# Someday we might add methods in this module so that someone who
# tries to `import ray_client as ray` -- as a module, instead of
# `from ray_client import ray` -- as the API stub
# still gets expected functionality. This is the way the ray package
# worked in the past.
#
# This really calls for PEP 562: https://www.python.org/dev/peps/pep-0562/
# But until Python 3.6 is EOL, here we are.
python/ray/util/client/api.py
+460
View File
@@ -0,0 +1,460 @@
"""This file defines the interface between the ray client worker
and the overall ray module API.
"""
import json
import logging
from concurrent.futures import Future
from typing import TYPE_CHECKING, Any, Callable, List, Optional, Union
from ray._common import ray_option_utils
from ray.util.client.runtime_context import _ClientWorkerPropertyAPI
if TYPE_CHECKING:
from ray.actor import ActorClass
from ray.core.generated.ray_client_pb2 import DataResponse
from ray.remote_function import RemoteFunction
from ray.util.client.common import ClientActorHandle, ClientObjectRef, ClientStub
logger = logging.getLogger(__name__)
def _as_bytes(value):
if isinstance(value, str):
return value.encode("utf-8")
return value
class _ClientAPI:
"""The Client-side methods corresponding to the ray API. Delegates
to the Client Worker that contains the connection to the ClientServer.
"""
def __init__(self, worker=None):
self.worker = worker
def get(
self,
vals: Union["ClientObjectRef", List["ClientObjectRef"]],
*,
timeout: Optional[float] = None,
):
"""get is the hook stub passed on to replace `ray.get`
Args:
vals: [Client]ObjectRef or list of these refs to retrieve.
timeout: Optional timeout in seconds
Returns:
The Python object(s) corresponding to ``vals``.
"""
return self.worker.get(vals, timeout=timeout)
def put(self, *args: Any, **kwargs: Any):
"""put is the hook stub passed on to replace `ray.put`
Args:
*args: opaque arguments forwarded to the worker's ``put``.
**kwargs: opaque keyword arguments forwarded to the worker's ``put``.
Returns:
A ``ClientObjectRef`` for the stored value.
"""
return self.worker.put(*args, **kwargs)
def wait(self, *args: Any, **kwargs: Any):
"""wait is the hook stub passed on to replace `ray.wait`
Args:
*args: opaque arguments forwarded to the worker's ``wait``.
**kwargs: opaque keyword arguments forwarded to the worker's ``wait``.
Returns:
A tuple ``(ready, remaining)`` of object refs, mirroring ``ray.wait``.
"""
return self.worker.wait(*args, **kwargs)
def _wait_generators_bulk(self, *args, **kwargs):
raise RuntimeError(
"ray._private.worker._wait_generators_bulk is not supported on Ray Client. "
"Connect with ray.init(address=...) instead, or use ray.wait."
)
def remote(self, *args: Any, **kwargs: Any):
"""remote is the hook stub passed on to replace `ray.remote`.
This sets up remote functions or actors, as the decorator,
but does not execute them.
Args:
*args: opaque arguments; when used as ``@ray.remote`` with no
parentheses, contains the wrapped function or class.
**kwargs: opaque keyword arguments; the options forwarded to
``ray.remote(...)``.
Returns:
A client-side stub for the remote function or actor, or a
decorator that produces one when applied.
"""
# Delayed import to avoid a cyclic import
from ray.util.client.common import remote_decorator
if len(args) == 1 and len(kwargs) == 0 and callable(args[0]):
# This is the case where the decorator is just @ray.remote.
return remote_decorator(options=None)(args[0])
assert (
len(args) == 0 and len(kwargs) > 0
), ray_option_utils.remote_args_error_string
return remote_decorator(options=kwargs)
# TODO(mwtian): consider adding _internal_ prefix to call_remote /
# call_release / call_retain.
def call_remote(
self, instance: "ClientStub", *args: Any, **kwargs: Any
) -> List[Future]:
"""call_remote is called by stub objects to execute them remotely.
This is used by stub objects in situations where they're called
with .remote, eg, `f.remote()` or `actor_cls.remote()`.
This allows the client stub objects to delegate execution to be
implemented in the most effective way whether it's in the client,
clientserver, or raylet worker.
Args:
instance: The Client-side stub reference to a remote object
*args: opaque arguments forwarded to the remote invocation.
**kwargs: opaque keyword arguments forwarded to the remote invocation.
Returns:
A list of futures, one per return value of the remote call.
"""
return self.worker.call_remote(instance, *args, **kwargs)
def call_release(self, id: bytes) -> None:
"""Attempts to release an object reference.
When client references are destructed, they release their reference,
which can opportunistically send a notification through the datachannel
to release the reference being held for that object on the server.
Args:
id: The id of the reference to release on the server side.
"""
return self.worker.call_release(id)
def call_retain(self, id: bytes) -> None:
"""Attempts to retain a client object reference.
Increments the reference count on the client side, to prevent
the client worker from attempting to release the server reference.
Args:
id: The id of the reference to retain on the client side.
"""
return self.worker.call_retain(id)
def close(self) -> None:
"""close cleans up an API connection by closing any channels or
shutting down any servers gracefully.
"""
return self.worker.close()
def get_actor(
self, name: str, namespace: Optional[str] = None
) -> "ClientActorHandle":
"""Returns a handle to an actor by name.
Args:
name: The name passed to this actor by
Actor.options(name="name").remote()
namespace: The namespace the named actor was created in.
Defaults to the current namespace.
Returns:
A ``ClientActorHandle`` for the named actor.
"""
return self.worker.get_actor(name, namespace)
def list_named_actors(self, all_namespaces: bool = False) -> List[str]:
"""List all named actors in the system.
Actors must have been created with Actor.options(name="name").remote().
This works for both detached & non-detached actors.
By default, only actors in the current namespace will be returned
and the returned entries will simply be their name.
If `all_namespaces` is set to True, all actors in the cluster will be
returned regardless of namespace, and the retunred entries will be of
the form '<namespace>/<name>'.
"""
return self.worker.list_named_actors(all_namespaces)
def kill(self, actor: "ClientActorHandle", *, no_restart: bool = True):
"""kill forcibly stops an actor running in the cluster
Args:
actor: The client-side handle of the actor to kill.
no_restart: Whether this actor should be restarted if it's a
restartable actor.
Returns:
The result of the underlying ``terminate_actor`` call.
"""
return self.worker.terminate_actor(actor, no_restart)
def cancel(
self,
obj: "ClientObjectRef",
*,
force: bool = False,
recursive: bool = True,
):
"""Cancels a task on the cluster.
If the specified task is pending execution, it will not be executed. If
the task is currently executing, the behavior depends on the ``force``
flag, as per `ray.cancel()`
Only non-actor tasks can be canceled. Canceled tasks will not be
retried (max_retries will not be respected).
Args:
obj: ObjectRef returned by the task that should be canceled.
force: Whether to force-kill a running task by killing
the worker that is running the task.
recursive: Whether to try to cancel tasks submitted by
the task specified.
Returns:
The result of the underlying ``terminate_task`` call.
"""
return self.worker.terminate_task(obj, force, recursive)
# Various metadata methods for the client that are defined in the protocol.
def is_initialized(self) -> bool:
"""True if our client is connected, and if the server is initialized.
Returns:
A boolean determining if the client is connected and
server initialized.
"""
return self.worker.is_initialized()
def nodes(self):
"""Get a list of the nodes in the cluster (for debugging only).
Returns:
Information about the Ray clients in the cluster.
"""
# This should be imported here, otherwise, it will error doc build.
import ray.core.generated.ray_client_pb2 as ray_client_pb2
return self.worker.get_cluster_info(ray_client_pb2.ClusterInfoType.NODES)
def method(self, *args: Any, **kwargs: Any):
"""Annotate an actor method
Args:
*args: Positional arguments are not supported; ``@ray.method``
must be invoked with at least one keyword argument.
**kwargs: Supported keyword arguments are ``num_returns`` (the
number of object refs that should be returned by invocations
of this actor method) and ``concurrency_group``.
Returns:
A decorator that annotates an actor method with the supplied
options.
"""
# NOTE: So this follows the same logic as in ray/actor.py::method()
# The reason to duplicate it here is to simplify the client mode
# redirection logic. As the annotated method gets pickled and sent to
# the server from the client it carries this private variable, it
# activates the same logic on the server side; so there's no need to
# pass anything else. It's inside the class definition that becomes an
# actor. Similar annotations would follow the same way.
valid_kwargs = ["num_returns", "concurrency_group"]
error_string = (
"The @ray.method decorator must be applied using at least one of "
f"the arguments in the list {valid_kwargs}, for example "
"'@ray.method(num_returns=2)'."
)
assert len(args) == 0 and len(kwargs) > 0, error_string
for key in kwargs:
key_error_string = (
f'Unexpected keyword argument to @ray.method: "{key}". The '
f"supported keyword arguments are {valid_kwargs}"
)
assert key in valid_kwargs, key_error_string
def annotate_method(method):
if "num_returns" in kwargs:
method.__ray_num_returns__ = kwargs["num_returns"]
if "concurrency_group" in kwargs:
method.__ray_concurrency_group__ = kwargs["concurrency_group"]
return method
return annotate_method
def cluster_resources(self):
"""Get the current total cluster resources.
Note that this information can grow stale as nodes are added to or
removed from the cluster.
Returns:
A dictionary mapping resource name to the total quantity of that
resource in the cluster.
"""
# This should be imported here, otherwise, it will error doc build.
import ray.core.generated.ray_client_pb2 as ray_client_pb2
return self.worker.get_cluster_info(
ray_client_pb2.ClusterInfoType.CLUSTER_RESOURCES
)
def available_resources(self):
"""Get the current available cluster resources.
This is different from `cluster_resources` in that this will return
idle (available) resources rather than total resources.
Note that this information can grow stale as tasks start and finish.
Returns:
A dictionary mapping resource name to the total quantity of that
resource in the cluster.
"""
# This should be imported here, otherwise, it will error doc build.
import ray.core.generated.ray_client_pb2 as ray_client_pb2
return self.worker.get_cluster_info(
ray_client_pb2.ClusterInfoType.AVAILABLE_RESOURCES
)
def get_runtime_context(self):
"""Return a Ray RuntimeContext describing the state on the server
Returns:
A RuntimeContext wrapping a client making get_cluster_info calls.
"""
return _ClientWorkerPropertyAPI(self.worker).build_runtime_context()
# Client process isn't assigned any GPUs.
def get_gpu_ids(self) -> list:
return []
def timeline(self, filename: Optional[str] = None) -> Optional[List[Any]]:
logger.warning(
"Timeline will include events from other clients using this server."
)
# This should be imported here, otherwise, it will error doc build.
import ray.core.generated.ray_client_pb2 as ray_client_pb2
all_events = self.worker.get_cluster_info(
ray_client_pb2.ClusterInfoType.TIMELINE
)
if filename is not None:
with open(filename, "w") as outfile:
json.dump(all_events, outfile)
else:
return all_events
def _internal_kv_initialized(self) -> bool:
"""Hook for internal_kv._internal_kv_initialized."""
# NOTE(edoakes): the kv is always initialized because we initialize it
# manually in the proxier with a GCS client if Ray hasn't been
# initialized yet.
return True
def _internal_kv_exists(
self, key: Union[str, bytes], *, namespace: Optional[Union[str, bytes]] = None
) -> bool:
"""Hook for internal_kv._internal_kv_exists."""
return self.worker.internal_kv_exists(
_as_bytes(key), namespace=_as_bytes(namespace)
)
def _internal_kv_get(
self, key: Union[str, bytes], *, namespace: Optional[Union[str, bytes]] = None
) -> bytes:
"""Hook for internal_kv._internal_kv_get."""
return self.worker.internal_kv_get(
_as_bytes(key), namespace=_as_bytes(namespace)
)
def _internal_kv_put(
self,
key: Union[str, bytes],
value: Union[str, bytes],
overwrite: bool = True,
*,
namespace: Optional[Union[str, bytes]] = None,
) -> bool:
"""Hook for internal_kv._internal_kv_put."""
return self.worker.internal_kv_put(
_as_bytes(key), _as_bytes(value), overwrite, namespace=_as_bytes(namespace)
)
def _internal_kv_del(
self,
key: Union[str, bytes],
*,
del_by_prefix: bool = False,
namespace: Optional[Union[str, bytes]] = None,
) -> int:
"""Hook for internal_kv._internal_kv_del."""
return self.worker.internal_kv_del(
_as_bytes(key), del_by_prefix=del_by_prefix, namespace=_as_bytes(namespace)
)
def _internal_kv_list(
self,
prefix: Union[str, bytes],
*,
namespace: Optional[Union[str, bytes]] = None,
) -> List[bytes]:
"""Hook for internal_kv._internal_kv_list."""
return self.worker.internal_kv_list(
_as_bytes(prefix), namespace=_as_bytes(namespace)
)
def _pin_runtime_env_uri(self, uri: str, expiration_s: int) -> None:
"""Hook for internal_kv._pin_runtime_env_uri."""
return self.worker.pin_runtime_env_uri(uri, expiration_s)
def _convert_actor(self, actor: "ActorClass") -> str:
"""Register a ClientActorClass for the ActorClass and return a UUID"""
return self.worker._convert_actor(actor)
def _convert_function(self, func: "RemoteFunction") -> str:
"""Register a ClientRemoteFunc for the ActorClass and return a UUID"""
return self.worker._convert_function(func)
def _get_converted(self, key: str) -> "ClientStub":
"""Given a UUID, return the converted object"""
return self.worker._get_converted(key)
def _converted_key_exists(self, key: str) -> bool:
"""Check if a key UUID is present in the store of converted objects."""
return self.worker._converted_key_exists(key)
def __getattr__(self, key: str):
if not key.startswith("_"):
raise NotImplementedError(
"Not available in Ray client: `ray.{}`. This method is only "
"available within Ray remote functions and is not yet "
"implemented in the client API.".format(key)
)
return self.__getattribute__(key)
def _register_callback(
self, ref: "ClientObjectRef", callback: Callable[["DataResponse"], None]
) -> None:
self.worker.register_callback(ref, callback)
def _get_dashboard_url(self) -> str:
import ray.core.generated.ray_client_pb2 as ray_client_pb2
return self.worker.get_cluster_info(
ray_client_pb2.ClusterInfoType.DASHBOARD_URL
).get("dashboard_url", "")
python/ray/util/client/client_app.py
+91
View File
@@ -0,0 +1,91 @@
from typing import Tuple
from ray.util.client import ray
ray.connect("localhost:50051")
@ray.remote
class HelloActor:
def __init__(self):
self.count = 0
def say_hello(self, whom: str) -> Tuple[str, int]:
self.count += 1
return ("Hello " + whom, self.count)
actor = HelloActor.remote()
s, count = ray.get(actor.say_hello.remote("you"))
print(s, count)
assert s == "Hello you"
assert count == 1
s, count = ray.get(actor.say_hello.remote("world"))
print(s, count)
assert s == "Hello world"
assert count == 2
@ray.remote
def plus2(x):
return x + 2
@ray.remote
def fact(x):
print(x, type(fact))
if x <= 0:
return 1
# This hits the "nested tasks" issue
# https://github.com/ray-project/ray/issues/3644
# So we're on the right track!
return ray.get(fact.remote(x - 1)) * x
@ray.remote
def get_nodes():
return ray.nodes() # Can access the full Ray API in remote methods.
print("Cluster nodes", ray.get(get_nodes.remote()))
print(ray.nodes())
objectref = ray.put("hello world")
# `ClientObjectRef(...)`
print(objectref)
# `hello world`
print(ray.get(objectref))
ref2 = plus2.remote(234)
# `ClientObjectRef(...)`
print(ref2)
# `236`
print(ray.get(ref2))
ref3 = fact.remote(20)
# `ClientObjectRef(...)`
print(ref3)
# `2432902008176640000`
print(ray.get(ref3))
# Reuse the cached ClientRemoteFunc object
ref4 = fact.remote(5)
# `120`
print(ray.get(ref4))
ref5 = fact.remote(10)
print([ref2, ref3, ref4, ref5])
# should return ref2, ref3, ref4
res = ray.wait([ref5, ref2, ref3, ref4], num_returns=3)
print(res)
assert [ref2, ref3, ref4] == res[0]
assert [ref5] == res[1]
# should return ref2, ref3, ref4, ref5
res = ray.wait([ref2, ref3, ref4, ref5], num_returns=4)
print(res)
assert [ref2, ref3, ref4, ref5] == res[0]
assert [] == res[1]
python/ray/util/client/client_pickler.py
+175
View File
@@ -0,0 +1,175 @@
"""Implements the client side of the client/server pickling protocol.
All ray client client/server data transfer happens through this pickling
protocol. The model is as follows:
* All Client objects (eg ClientObjectRef) always live on the client and
are never represented in the server
* All Ray objects (eg, ray.ObjectRef) always live on the server and are
never returned to the client
* In order to translate between these two references, PickleStub tuples
are generated as persistent ids in the data blobs during the pickling
and unpickling of these objects.
The PickleStubs have just enough information to find or generate their
associated partner object on either side.
This also has the advantage of avoiding predefined pickle behavior for ray
objects, which may include ray internal reference counting.
ClientPickler dumps things from the client into the appropriate stubs
ServerUnpickler loads stubs from the server into their client counterparts.
"""
import io
import pickle # noqa: F401
from typing import Any, Dict, NamedTuple, Optional
import ray.cloudpickle as cloudpickle
import ray.core.generated.ray_client_pb2 as ray_client_pb2
from ray.util.client import RayAPIStub
from ray.util.client.common import (
ClientActorClass,
ClientActorHandle,
ClientActorRef,
ClientObjectRef,
ClientRemoteFunc,
ClientRemoteMethod,
InProgressSentinel,
OptionWrapper,
)
# NOTE(barakmich): These PickleStubs are really close to
# the data for an execution, with no arguments. Combine the two?
class PickleStub(
NamedTuple(
"PickleStub",
[
("type", str),
("client_id", str),
("ref_id", bytes),
("name", Optional[str]),
("baseline_options", Optional[Dict]),
],
)
):
def __reduce__(self):
# PySpark's namedtuple monkey patch breaks compatibility with
# cloudpickle. Thus we revert this patch here if it exists.
return object.__reduce__(self)
class ClientPickler(cloudpickle.CloudPickler):
def __init__(self, client_id, *args, **kwargs):
super().__init__(*args, **kwargs)
self.client_id = client_id
def persistent_id(self, obj):
if isinstance(obj, RayAPIStub):
return PickleStub(
type="Ray",
client_id=self.client_id,
ref_id=b"",
name=None,
baseline_options=None,
)
elif isinstance(obj, ClientObjectRef):
return PickleStub(
type="Object",
client_id=self.client_id,
ref_id=obj.id,
name=None,
baseline_options=None,
)
elif isinstance(obj, ClientActorHandle):
return PickleStub(
type="Actor",
client_id=self.client_id,
ref_id=obj._actor_id.id,
name=None,
baseline_options=None,
)
elif isinstance(obj, ClientRemoteFunc):
if obj._ref is None:
obj._ensure_ref()
if type(obj._ref) is InProgressSentinel:
return PickleStub(
type="RemoteFuncSelfReference",
client_id=self.client_id,
ref_id=obj._client_side_ref.id,
name=None,
baseline_options=None,
)
return PickleStub(
type="RemoteFunc",
client_id=self.client_id,
ref_id=obj._ref.id,
name=None,
baseline_options=obj._options,
)
elif isinstance(obj, ClientActorClass):
if obj._ref is None:
obj._ensure_ref()
if type(obj._ref) is InProgressSentinel:
return PickleStub(
type="RemoteActorSelfReference",
client_id=self.client_id,
ref_id=obj._client_side_ref.id,
name=None,
baseline_options=None,
)
return PickleStub(
type="RemoteActor",
client_id=self.client_id,
ref_id=obj._ref.id,
name=None,
baseline_options=obj._options,
)
elif isinstance(obj, ClientRemoteMethod):
return PickleStub(
type="RemoteMethod",
client_id=self.client_id,
ref_id=obj._actor_handle.actor_ref.id,
name=obj._method_name,
baseline_options=None,
)
elif isinstance(obj, OptionWrapper):
raise NotImplementedError("Sending a partial option is unimplemented")
return None
class ServerUnpickler(pickle.Unpickler):
def persistent_load(self, pid):
assert isinstance(pid, PickleStub)
if pid.type == "Object":
return ClientObjectRef(pid.ref_id)
elif pid.type == "Actor":
return ClientActorHandle(ClientActorRef(pid.ref_id))
else:
raise NotImplementedError("Being passed back an unknown stub")
def dumps_from_client(obj: Any, client_id: str, protocol=None) -> bytes:
with io.BytesIO() as file:
cp = ClientPickler(client_id, file, protocol=protocol)
cp.dump(obj)
return file.getvalue()
def loads_from_server(
data: bytes, *, fix_imports=True, encoding="ASCII", errors="strict"
) -> Any:
if isinstance(data, str):
raise TypeError("Can't load pickle from unicode string")
file = io.BytesIO(data)
return ServerUnpickler(
file, fix_imports=fix_imports, encoding=encoding, errors=errors
).load()
def convert_to_arg(val: Any, client_id: str) -> ray_client_pb2.Arg:
out = ray_client_pb2.Arg()
out.local = ray_client_pb2.Arg.Locality.INTERNED
out.data = dumps_from_client(val, client_id)
return out
python/ray/util/client/common.py
+957
View File
@@ -0,0 +1,957 @@
import inspect
import logging
import os
import pickle
import threading
import uuid
from collections import OrderedDict
from concurrent.futures import Future
from dataclasses import dataclass
from typing import Any, Callable, Dict, List, Optional, Tuple, Union
import grpc
import ray._raylet as raylet
import ray.core.generated.ray_client_pb2 as ray_client_pb2
import ray.core.generated.ray_client_pb2_grpc as ray_client_pb2_grpc
from ray._common.signature import extract_signature, get_signature
from ray._private import ray_constants
from ray._private.inspect_util import (
is_class_method,
is_cython,
is_function_or_method,
is_static_method,
)
from ray._private.utils import check_oversized_function
from ray.util.client import ray
from ray.util.client.options import validate_options
from ray.util.common import INT32_MAX
logger = logging.getLogger(__name__)
# gRPC status codes that the client shouldn't attempt to recover from
# Resource exhausted: Server is low on resources, or has hit the max number
# of client connections
# Invalid argument: Reserved for application errors
# Not found: Set if the client is attempting to reconnect to a session that
# does not exist
# Failed precondition: Reserverd for application errors
# Aborted: Set when an error is serialized into the details of the context,
# signals that error should be deserialized on the client side
GRPC_UNRECOVERABLE_ERRORS = (
grpc.StatusCode.RESOURCE_EXHAUSTED,
grpc.StatusCode.INVALID_ARGUMENT,
grpc.StatusCode.NOT_FOUND,
grpc.StatusCode.FAILED_PRECONDITION,
grpc.StatusCode.ABORTED,
)
# TODO: Instead of just making the max message size large, the right thing to
# do is to split up the bytes representation of serialized data into multiple
# messages and reconstruct them on either end. That said, since clients are
# drivers and really just feed initial things in and final results out, (when
# not going to S3 or similar) then a large limit will suffice for many use
# cases.
#
# Currently, this is 2GiB, the max for a signed int.
GRPC_MAX_MESSAGE_SIZE = (2 * 1024 * 1024 * 1024) - 1
# 30 seconds because ELB timeout is 60 seconds
GRPC_KEEPALIVE_TIME_MS = 1000 * 30
# Long timeout because we do not want gRPC ending a connection.
GRPC_KEEPALIVE_TIMEOUT_MS = 1000 * 600
GRPC_OPTIONS = [
*ray_constants.GLOBAL_GRPC_OPTIONS,
("grpc.max_send_message_length", GRPC_MAX_MESSAGE_SIZE),
("grpc.max_receive_message_length", GRPC_MAX_MESSAGE_SIZE),
("grpc.keepalive_time_ms", GRPC_KEEPALIVE_TIME_MS),
("grpc.keepalive_timeout_ms", GRPC_KEEPALIVE_TIMEOUT_MS),
("grpc.keepalive_permit_without_calls", 1),
# Send an infinite number of pings
("grpc.http2.max_pings_without_data", 0),
("grpc.http2.min_ping_interval_without_data_ms", GRPC_KEEPALIVE_TIME_MS - 50),
# Allow many strikes
("grpc.http2.max_ping_strikes", 0),
]
CLIENT_SERVER_MAX_THREADS = float(os.getenv("RAY_CLIENT_SERVER_MAX_THREADS", 100))
# Large objects are chunked into 5 MiB messages, ref PR #35025
OBJECT_TRANSFER_CHUNK_SIZE = 5 * 2**20
# Warn the user if the object being transferred is larger than 2 GiB
OBJECT_TRANSFER_WARNING_SIZE = 2 * 2**30
class ClientObjectRef(raylet.ObjectRef):
def __init__(self, id: Union[bytes, Future]):
self._mutex = threading.Lock()
self._worker = ray.get_context().client_worker
self._id_future = None
if isinstance(id, bytes):
self._set_id(id)
elif isinstance(id, Future):
self._id_future = id
else:
raise TypeError("Unexpected type for id {}".format(id))
def __del__(self):
if self._worker is not None and self._worker.is_connected():
try:
if not self.is_nil():
self._worker.call_release(self.id)
except Exception:
logger.info(
"Exception in ObjectRef is ignored in destructor. "
"To receive this exception in application code, call "
"a method on the actor reference before its destructor "
"is run."
)
def binary(self):
self._wait_for_id()
return super().binary()
def hex(self):
self._wait_for_id()
return super().hex()
def is_nil(self):
self._wait_for_id()
return super().is_nil()
def __hash__(self):
self._wait_for_id()
return hash(self.id)
def task_id(self):
self._wait_for_id()
return super().task_id()
@property
def id(self):
return self.binary()
def future(self) -> Future:
fut = Future()
def set_future(data: Any) -> None:
"""Schedules a callback to set the exception or result
in the Future."""
if isinstance(data, Exception):
fut.set_exception(data)
else:
fut.set_result(data)
self._on_completed(set_future)
# Prevent this object ref from being released.
fut.object_ref = self
return fut
def _on_completed(self, py_callback: Callable[[Any], None]) -> None:
"""Register a callback that will be called after Object is ready.
If the ObjectRef is already ready, the callback will be called soon.
The callback should take the result as the only argument. The result
can be an exception object in case of task error.
"""
def deserialize_obj(
resp: Union[ray_client_pb2.DataResponse, Exception]
) -> None:
from ray.util.client.client_pickler import loads_from_server
if isinstance(resp, Exception):
data = resp
elif isinstance(resp, bytearray):
data = loads_from_server(resp)
else:
obj = resp.get
data = None
if not obj.valid:
data = loads_from_server(resp.get.error)
else:
data = loads_from_server(resp.get.data)
py_callback(data)
self._worker.register_callback(self, deserialize_obj)
def _set_id(self, id):
super()._set_id(id)
self._worker.call_retain(id)
def _wait_for_id(self, timeout=None):
if self._id_future:
with self._mutex:
if self._id_future:
self._set_id(self._id_future.result(timeout=timeout))
self._id_future = None
class ClientActorRef(raylet.ActorID):
def __init__(
self,
id: Union[bytes, Future],
weak_ref: Optional[bool] = False,
):
self._weak_ref = weak_ref
self._mutex = threading.Lock()
self._worker = ray.get_context().client_worker
if isinstance(id, bytes):
self._set_id(id)
self._id_future = None
elif isinstance(id, Future):
self._id_future = id
else:
raise TypeError("Unexpected type for id {}".format(id))
def __del__(self):
if self._weak_ref:
return
if self._worker is not None and self._worker.is_connected():
try:
if not self.is_nil():
self._worker.call_release(self.id)
except Exception:
logger.debug(
"Exception from actor creation is ignored in destructor. "
"To receive this exception in application code, call "
"a method on the actor reference before its destructor "
"is run."
)
def binary(self):
self._wait_for_id()
return super().binary()
def hex(self):
self._wait_for_id()
return super().hex()
def is_nil(self):
self._wait_for_id()
return super().is_nil()
def __hash__(self):
self._wait_for_id()
return hash(self.id)
@property
def id(self):
return self.binary()
def _set_id(self, id):
super()._set_id(id)
self._worker.call_retain(id)
def _wait_for_id(self, timeout=None):
if self._id_future:
with self._mutex:
if self._id_future:
self._set_id(self._id_future.result(timeout=timeout))
self._id_future = None
class ClientStub:
pass
class ClientRemoteFunc(ClientStub):
"""A stub created on the Ray Client to represent a remote
function that can be exectued on the cluster.
This class is allowed to be passed around between remote functions.
Args:
f: The actual function to execute remotely.
options: Optional ``ray.remote`` options applied to this function.
"""
def __init__(self, f: Callable, options: Optional[Dict[str, Any]] = None):
self._lock = threading.Lock()
self._func = f
self._name = f.__name__
self._signature = get_signature(f)
self._ref = None
self._client_side_ref = ClientSideRefID.generate_id()
self._options = validate_options(options)
def __call__(self, *args, **kwargs):
raise TypeError(
"Remote function cannot be called directly. "
f"Use {self._name}.remote method instead"
)
def remote(self, *args, **kwargs):
# Check if supplied parameters match the function signature. Same case
# at the other callsites.
self._signature.bind(*args, **kwargs)
return return_refs(ray.call_remote(self, *args, **kwargs))
def options(self, **kwargs):
return OptionWrapper(self, kwargs)
def _remote(self, args=None, kwargs=None, **option_args):
if args is None:
args = []
if kwargs is None:
kwargs = {}
return self.options(**option_args).remote(*args, **kwargs)
def __repr__(self):
return "ClientRemoteFunc(%s, %s)" % (self._name, self._ref)
def _ensure_ref(self):
with self._lock:
if self._ref is None:
# While calling ray.put() on our function, if
# our function is recursive, it will attempt to
# encode the ClientRemoteFunc -- itself -- and
# infinitely recurse on _ensure_ref.
#
# So we set the state of the reference to be an
# in-progress self reference value, which
# the encoding can detect and handle correctly.
self._ref = InProgressSentinel()
data = ray.worker._dumps_from_client(self._func)
# Check pickled size before sending it to server, which is more
# efficient and can be done synchronously inside remote() call.
check_oversized_function(data, self._name, "remote function", None)
self._ref = ray.worker._put_pickled(
data, client_ref_id=self._client_side_ref.id
)
def _prepare_client_task(self) -> ray_client_pb2.ClientTask:
self._ensure_ref()
task = ray_client_pb2.ClientTask()
task.type = ray_client_pb2.ClientTask.FUNCTION
task.name = self._name
task.payload_id = self._ref.id
set_task_options(task, self._options, "baseline_options")
return task
def _num_returns(self) -> int:
if not self._options:
return None
return self._options.get("num_returns")
class ClientActorClass(ClientStub):
"""A stub created on the Ray Client to represent an actor class.
It is wrapped by ray.remote and can be executed on the cluster.
Args:
actor_cls: The actual class to execute remotely.
options: Optional ``ray.remote`` options applied to this actor class.
"""
def __init__(self, actor_cls: type, options: Optional[Dict[str, Any]] = None):
self.actor_cls = actor_cls
self._lock = threading.Lock()
self._name = actor_cls.__name__
self._init_signature = inspect.Signature(
parameters=extract_signature(actor_cls.__init__, ignore_first=True)
)
self._ref = None
self._client_side_ref = ClientSideRefID.generate_id()
self._options = validate_options(options)
def __call__(self, *args, **kwargs):
raise TypeError(
"Remote actor cannot be instantiated directly. "
f"Use {self._name}.remote() instead"
)
def _ensure_ref(self):
with self._lock:
if self._ref is None:
# As before, set the state of the reference to be an
# in-progress self reference value, which
# the encoding can detect and handle correctly.
self._ref = InProgressSentinel()
data = ray.worker._dumps_from_client(self.actor_cls)
# Check pickled size before sending it to server, which is more
# efficient and can be done synchronously inside remote() call.
check_oversized_function(data, self._name, "actor", None)
self._ref = ray.worker._put_pickled(
data, client_ref_id=self._client_side_ref.id
)
def remote(self, *args, **kwargs) -> "ClientActorHandle":
self._init_signature.bind(*args, **kwargs)
# Actually instantiate the actor
futures = ray.call_remote(self, *args, **kwargs)
assert len(futures) == 1
return ClientActorHandle(ClientActorRef(futures[0]), actor_class=self)
def options(self, **kwargs):
return ActorOptionWrapper(self, kwargs)
def _remote(self, args=None, kwargs=None, **option_args):
if args is None:
args = []
if kwargs is None:
kwargs = {}
return self.options(**option_args).remote(*args, **kwargs)
def __repr__(self):
return "ClientActorClass(%s, %s)" % (self._name, self._ref)
def __getattr__(self, key):
if key not in self.__dict__:
raise AttributeError("Not a class attribute")
raise NotImplementedError("static methods")
def _prepare_client_task(self) -> ray_client_pb2.ClientTask:
self._ensure_ref()
task = ray_client_pb2.ClientTask()
task.type = ray_client_pb2.ClientTask.ACTOR
task.name = self._name
task.payload_id = self._ref.id
set_task_options(task, self._options, "baseline_options")
return task
@staticmethod
def _num_returns() -> int:
return 1
class ClientActorHandle(ClientStub):
"""Client-side stub for instantiated actor.
A stub created on the Ray Client to represent a remote actor that
has been started on the cluster. This class is allowed to be passed
around between remote functions.
Args:
actor_ref: A reference to the running actor given to the client. This
is a serialized version of the actual handle as an opaque token.
actor_class: Optional ``ClientActorClass`` used to populate method
signatures and ``num_returns`` metadata without a server round-trip.
"""
def __init__(
self,
actor_ref: ClientActorRef,
actor_class: Optional[ClientActorClass] = None,
):
self.actor_ref = actor_ref
self._dir: Optional[List[str]] = None
if actor_class is not None:
self._method_num_returns = {}
self._method_signatures = {}
for method_name, method_obj in inspect.getmembers(
actor_class.actor_cls, is_function_or_method
):
self._method_num_returns[method_name] = getattr(
method_obj, "__ray_num_returns__", None
)
self._method_signatures[method_name] = inspect.Signature(
parameters=extract_signature(
method_obj,
ignore_first=(
not (
is_class_method(method_obj)
or is_static_method(actor_class.actor_cls, method_name)
)
),
)
)
else:
self._method_num_returns = None
self._method_signatures = None
def __dir__(self) -> List[str]:
if self._method_num_returns is not None:
return self._method_num_returns.keys()
if ray.is_connected():
self._init_class_info()
return self._method_num_returns.keys()
return super().__dir__()
# For compatibility with core worker ActorHandle._actor_id which returns
# ActorID
@property
def _actor_id(self) -> ClientActorRef:
return self.actor_ref
def __hash__(self) -> int:
return hash(self._actor_id)
def __eq__(self, __value) -> bool:
return hash(self) == hash(__value)
def __getattr__(self, key):
if key == "_method_num_returns":
# We need to explicitly handle this value since it is used below,
# otherwise we may end up infinitely recursing when deserializing.
# This can happen after unpickling an object but before
# _method_num_returns is correctly populated.
raise AttributeError(f"ClientActorRef has no attribute '{key}'")
if self._method_num_returns is None:
self._init_class_info()
if key not in self._method_signatures:
raise AttributeError(f"ClientActorRef has no attribute '{key}'")
return ClientRemoteMethod(
self,
key,
self._method_num_returns.get(key),
self._method_signatures.get(key),
)
def __repr__(self):
return "ClientActorHandle(%s)" % (self.actor_ref.id.hex())
def _init_class_info(self):
# TODO: fetch Ray method decorators
@ray.remote(num_cpus=0)
def get_class_info(x):
return x._ray_method_num_returns, x._ray_method_signatures
self._method_num_returns, method_parameters = ray.get(
get_class_info.remote(self)
)
self._method_signatures = {}
for method, parameters in method_parameters.items():
self._method_signatures[method] = inspect.Signature(parameters=parameters)
class ClientRemoteMethod(ClientStub):
"""A stub for a method on a remote actor.
Can be annotated with execution options.
Args:
actor_handle: A reference to the ClientActorHandle that generated
this method and will have this method called upon it.
method_name: The name of this method.
num_returns: Number of object refs returned by invocations of this
method.
signature: The method's bound signature, used to validate call args.
"""
def __init__(
self,
actor_handle: ClientActorHandle,
method_name: str,
num_returns: int,
signature: inspect.Signature,
):
self._actor_handle = actor_handle
self._method_name = method_name
self._method_num_returns = num_returns
self._signature = signature
def __call__(self, *args, **kwargs):
raise TypeError(
"Actor methods cannot be called directly. Instead "
f"of running 'object.{self._method_name}()', try "
f"'object.{self._method_name}.remote()'."
)
def remote(self, *args, **kwargs):
self._signature.bind(*args, **kwargs)
return return_refs(ray.call_remote(self, *args, **kwargs))
def __repr__(self):
return "ClientRemoteMethod(%s, %s, %s)" % (
self._method_name,
self._actor_handle,
self._method_num_returns,
)
def options(self, **kwargs):
return OptionWrapper(self, kwargs)
def _remote(self, args=None, kwargs=None, **option_args):
if args is None:
args = []
if kwargs is None:
kwargs = {}
return self.options(**option_args).remote(*args, **kwargs)
def _prepare_client_task(self) -> ray_client_pb2.ClientTask:
task = ray_client_pb2.ClientTask()
task.type = ray_client_pb2.ClientTask.METHOD
task.name = self._method_name
task.payload_id = self._actor_handle.actor_ref.id
return task
def _num_returns(self) -> int:
return self._method_num_returns
class OptionWrapper:
def __init__(self, stub: ClientStub, options: Optional[Dict[str, Any]]):
self._remote_stub = stub
self._options = validate_options(options)
def remote(self, *args, **kwargs):
self._remote_stub._signature.bind(*args, **kwargs)
return return_refs(ray.call_remote(self, *args, **kwargs))
def __getattr__(self, key):
return getattr(self._remote_stub, key)
def _prepare_client_task(self):
task = self._remote_stub._prepare_client_task()
set_task_options(task, self._options)
return task
def _num_returns(self) -> int:
if self._options:
num = self._options.get("num_returns")
if num is not None:
return num
return self._remote_stub._num_returns()
class ActorOptionWrapper(OptionWrapper):
def remote(self, *args, **kwargs):
self._remote_stub._init_signature.bind(*args, **kwargs)
futures = ray.call_remote(self, *args, **kwargs)
assert len(futures) == 1
actor_class = None
if isinstance(self._remote_stub, ClientActorClass):
actor_class = self._remote_stub
return ClientActorHandle(ClientActorRef(futures[0]), actor_class=actor_class)
def set_task_options(
task: ray_client_pb2.ClientTask,
options: Optional[Dict[str, Any]],
field: str = "options",
) -> None:
if options is None:
task.ClearField(field)
return
getattr(task, field).pickled_options = pickle.dumps(options)
def return_refs(
futures: List[Future],
) -> Union[None, ClientObjectRef, List[ClientObjectRef]]:
if not futures:
return None
if len(futures) == 1:
return ClientObjectRef(futures[0])
return [ClientObjectRef(fut) for fut in futures]
class InProgressSentinel:
def __repr__(self) -> str:
return self.__class__.__name__
class ClientSideRefID:
"""An ID generated by the client for objects not yet given an ObjectRef"""
def __init__(self, id: bytes):
assert len(id) != 0
self.id = id
@staticmethod
def generate_id() -> "ClientSideRefID":
tid = uuid.uuid4()
return ClientSideRefID(b"\xcc" + tid.bytes)
def remote_decorator(options: Optional[Dict[str, Any]]):
def decorator(function_or_class) -> ClientStub:
if inspect.isfunction(function_or_class) or is_cython(function_or_class):
return ClientRemoteFunc(function_or_class, options=options)
elif inspect.isclass(function_or_class):
return ClientActorClass(function_or_class, options=options)
else:
raise TypeError(
"The @ray.remote decorator must be applied to "
"either a function or to a class."
)
return decorator
@dataclass
class ClientServerHandle:
"""Holds the handles to the registered gRPC servicers and their server."""
task_servicer: ray_client_pb2_grpc.RayletDriverServicer
data_servicer: ray_client_pb2_grpc.RayletDataStreamerServicer
logs_servicer: ray_client_pb2_grpc.RayletLogStreamerServicer
grpc_server: grpc.Server
def stop(self, grace: int) -> None:
# The data servicer might be sleeping while waiting for clients to
# reconnect. Signal that they no longer have to sleep and can exit
# immediately, since the RPC server is stopped.
self.grpc_server.stop(grace)
self.data_servicer.stopped.set()
# Add a hook for all the cases that previously
# expected simply a gRPC server
def __getattr__(self, attr):
return getattr(self.grpc_server, attr)
def _get_client_id_from_context(context: Any) -> str:
"""
Get `client_id` from gRPC metadata. If the `client_id` is not present,
this function logs an error and sets the status_code.
"""
metadata = dict(context.invocation_metadata())
client_id = metadata.get("client_id") or ""
if client_id == "":
logger.error("Client connecting with no client_id")
context.set_code(grpc.StatusCode.FAILED_PRECONDITION)
return client_id
def _propagate_error_in_context(e: Exception, context: Any) -> bool:
"""
Encode an error into the context of an RPC response. Returns True
if the error can be recovered from, false otherwise
"""
try:
if isinstance(e, grpc.RpcError):
# RPC error, propagate directly by copying details into context
context.set_code(e.code())
context.set_details(e.details())
return e.code() not in GRPC_UNRECOVERABLE_ERRORS
except Exception:
# Extra precaution -- if encoding the RPC directly fails fallback
# to treating it as a regular error
pass
context.set_code(grpc.StatusCode.FAILED_PRECONDITION)
context.set_details(str(e))
return False
def _id_is_newer(id1: int, id2: int) -> bool:
"""
We should only replace cache entries with the responses for newer IDs.
Most of the time newer IDs will be the ones with higher value, except when
the req_id counter rolls over. We check for this case by checking the
distance between the two IDs. If the distance is significant, then it's
likely that the req_id counter rolled over, and the smaller id should
still be used to replace the one in cache.
"""
diff = abs(id2 - id1)
# Int32 max is also the maximum number of simultaneous in-flight requests.
if diff > (INT32_MAX // 2):
# Rollover likely occurred. In this case the smaller ID is newer
return id1 < id2
return id1 > id2
class ResponseCache:
"""
Cache for blocking method calls. Needed to prevent retried requests from
being applied multiple times on the server, for example when the client
disconnects. This is used to cache requests/responses sent through
unary-unary RPCs to the RayletServicer.
Note that no clean up logic is used, the last response for each thread
will always be remembered, so at most the cache will hold N entries,
where N is the number of threads on the client side. This relies on the
assumption that a thread will not make a new blocking request until it has
received a response for a previous one, at which point it's safe to
overwrite the old response.
The high level logic is:
1. Before making a call, check the cache for the current thread.
2. If present in the cache, check the request id of the cached
response.
a. If it matches the current request_id, then the request has been
received before and we shouldn't re-attempt the logic. Wait for
the response to become available in the cache, and then return it
b. If it doesn't match, then this is a new request and we can
proceed with calling the real stub. While the response is still
being generated, temporarily keep (req_id, None) in the cache.
Once the call is finished, update the cache entry with the
new (req_id, response) pair. Notify other threads that may
have been waiting for the response to be prepared.
"""
def __init__(self):
self.cv = threading.Condition()
self.cache: Dict[int, Tuple[int, Any]] = {}
def check_cache(self, thread_id: int, request_id: int) -> Optional[Any]:
"""
Check the cache for a given thread, and see if the entry in the cache
matches the current request_id. Returns None if the request_id has
not been seen yet, otherwise returns the cached result.
Throws an error if the placeholder in the cache doesn't match the
request_id -- this means that a new request evicted the old value in
the cache, and that the RPC for `request_id` is redundant and the
result can be discarded, i.e.:
1. Request A is sent (A1)
2. Channel disconnects
3. Request A is resent (A2)
4. A1 is received
5. A2 is received, waits for A1 to finish
6. A1 finishes and is sent back to client
7. Request B is sent
8. Request B overwrites cache entry
9. A2 wakes up extremely late, but cache is now invalid
In practice this is VERY unlikely to happen, but the error can at
least serve as a sanity check or catch invalid request id's.
"""
with self.cv:
if thread_id in self.cache:
cached_request_id, cached_resp = self.cache[thread_id]
if cached_request_id == request_id:
while cached_resp is None:
# The call was started, but the response hasn't yet
# been added to the cache. Let go of the lock and
# wait until the response is ready.
self.cv.wait()
cached_request_id, cached_resp = self.cache[thread_id]
if cached_request_id != request_id:
raise RuntimeError(
"Cached response doesn't match the id of the "
"original request. This might happen if this "
"request was received out of order. The "
"result of the caller is no longer needed. "
f"({request_id} != {cached_request_id})"
)
return cached_resp
if not _id_is_newer(request_id, cached_request_id):
raise RuntimeError(
"Attempting to replace newer cache entry with older "
"one. This might happen if this request was received "
"out of order. The result of the caller is no "
f"longer needed. ({request_id} != {cached_request_id}"
)
self.cache[thread_id] = (request_id, None)
return None
def update_cache(self, thread_id: int, request_id: int, response: Any) -> None:
"""
Inserts `response` into the cache for `request_id`.
"""
with self.cv:
cached_request_id, cached_resp = self.cache[thread_id]
if cached_request_id != request_id or cached_resp is not None:
# The cache was overwritten by a newer requester between
# our call to check_cache and our call to update it.
# This can't happen if the assumption that the cached requests
# are all blocking on the client side, so if you encounter
# this, check if any async requests are being cached.
raise RuntimeError(
"Attempting to update the cache, but placeholder's "
"do not match the current request_id. This might happen "
"if this request was received out of order. The result "
f"of the caller is no longer needed. ({request_id} != "
f"{cached_request_id})"
)
self.cache[thread_id] = (request_id, response)
self.cv.notify_all()
class OrderedResponseCache:
"""
Cache for streaming RPCs, i.e. the DataServicer. Relies on explicit
ack's from the client to determine when it can clean up cache entries.
"""
def __init__(self):
self.last_received = 0
self.cv = threading.Condition()
self.cache: Dict[int, Any] = OrderedDict()
def check_cache(self, req_id: int) -> Optional[Any]:
"""
Check the cache for a given thread, and see if the entry in the cache
matches the current request_id. Returns None if the request_id has
not been seen yet, otherwise returns the cached result.
"""
with self.cv:
if _id_is_newer(self.last_received, req_id) or self.last_received == req_id:
# Request is for an id that has already been cleared from
# cache/acknowledged.
raise RuntimeError(
"Attempting to accesss a cache entry that has already "
"cleaned up. The client has already acknowledged "
f"receiving this response. ({req_id}, "
f"{self.last_received})"
)
if req_id in self.cache:
cached_resp = self.cache[req_id]
while cached_resp is None:
# The call was started, but the response hasn't yet been
# added to the cache. Let go of the lock and wait until
# the response is ready
self.cv.wait()
if req_id not in self.cache:
raise RuntimeError(
"Cache entry was removed. This likely means that "
"the result of this call is no longer needed."
)
cached_resp = self.cache[req_id]
return cached_resp
self.cache[req_id] = None
return None
def update_cache(self, req_id: int, resp: Any) -> None:
"""
Inserts `response` into the cache for `request_id`.
"""
with self.cv:
self.cv.notify_all()
if req_id not in self.cache:
raise RuntimeError(
"Attempting to update the cache, but placeholder is "
"missing. This might happen on a redundant call to "
f"update_cache. ({req_id})"
)
self.cache[req_id] = resp
def invalidate(self, e: Exception) -> bool:
"""
Invalidate any partially populated cache entries, replacing their
placeholders with the passed in exception. Useful to prevent a thread
from waiting indefinitely on a failed call.
Returns True if the cache contains an error, False otherwise
"""
with self.cv:
invalid = False
for req_id in self.cache:
if self.cache[req_id] is None:
self.cache[req_id] = e
if isinstance(self.cache[req_id], Exception):
invalid = True
self.cv.notify_all()
return invalid
def cleanup(self, last_received: int) -> None:
"""
Cleanup all of the cached requests up to last_received. Assumes that
the cache entries were inserted in ascending order.
"""
with self.cv:
if _id_is_newer(last_received, self.last_received):
self.last_received = last_received
to_remove = []
for req_id in self.cache:
if _id_is_newer(last_received, req_id) or last_received == req_id:
to_remove.append(req_id)
else:
break
for req_id in to_remove:
del self.cache[req_id]
self.cv.notify_all()
python/ray/util/client/dataclient.py
+598
View File
@@ -0,0 +1,598 @@
"""This file implements a threaded stream controller to abstract a data stream
back to the ray clientserver.
"""
import logging
import math
import queue
import threading
import warnings
from collections import OrderedDict
from typing import TYPE_CHECKING, Any, Callable, Dict, Optional, Union
import grpc
import ray.core.generated.ray_client_pb2 as ray_client_pb2
import ray.core.generated.ray_client_pb2_grpc as ray_client_pb2_grpc
from ray.util.client.common import (
INT32_MAX,
OBJECT_TRANSFER_CHUNK_SIZE,
OBJECT_TRANSFER_WARNING_SIZE,
)
from ray.util.debug import log_once
if TYPE_CHECKING:
from ray.util.client.worker import Worker
logger = logging.getLogger(__name__)
ResponseCallable = Callable[[Union[ray_client_pb2.DataResponse, Exception]], None]
# Send an acknowledge on every 32nd response received
ACKNOWLEDGE_BATCH_SIZE = 32
def chunk_put(req: ray_client_pb2.DataRequest):
"""
Chunks a put request. Doing this lazily is important for large objects,
since taking slices of bytes objects does a copy. This means if we
immediately materialized every chunk of a large object and inserted them
into the result_queue, we would effectively double the memory needed
on the client to handle the put.
"""
# When accessing a protobuf field, deserialization is performed, which will
# generate a copy. So we need to avoid accessing the `data` field multiple
# times in the loop
request_data = req.put.data
total_size = len(request_data)
assert total_size > 0, "Cannot chunk object with missing data"
if total_size >= OBJECT_TRANSFER_WARNING_SIZE and log_once(
"client_object_put_size_warning"
):
size_gb = total_size / 2**30
warnings.warn(
"Ray Client is attempting to send a "
f"{size_gb:.2f} GiB object over the network, which may "
"be slow. Consider serializing the object and using a remote "
"URI to transfer via S3 or Google Cloud Storage instead. "
"Documentation for doing this can be found here: "
"https://docs.ray.io/en/latest/handling-dependencies.html#remote-uris",
UserWarning,
)
total_chunks = math.ceil(total_size / OBJECT_TRANSFER_CHUNK_SIZE)
for chunk_id in range(0, total_chunks):
start = chunk_id * OBJECT_TRANSFER_CHUNK_SIZE
end = min(total_size, (chunk_id + 1) * OBJECT_TRANSFER_CHUNK_SIZE)
chunk = ray_client_pb2.PutRequest(
client_ref_id=req.put.client_ref_id,
data=request_data[start:end],
chunk_id=chunk_id,
total_chunks=total_chunks,
total_size=total_size,
)
yield ray_client_pb2.DataRequest(req_id=req.req_id, put=chunk)
def chunk_task(req: ray_client_pb2.DataRequest):
"""
Chunks a client task. Doing this lazily is important with large arguments,
since taking slices of bytes objects does a copy. This means if we
immediately materialized every chunk of a large argument and inserted them
into the result_queue, we would effectively double the memory needed
on the client to handle the task.
"""
# When accessing a protobuf field, deserialization is performed, which will
# generate a copy. So we need to avoid accessing the `data` field multiple
# times in the loop
request_data = req.task.data
total_size = len(request_data)
assert total_size > 0, "Cannot chunk object with missing data"
total_chunks = math.ceil(total_size / OBJECT_TRANSFER_CHUNK_SIZE)
for chunk_id in range(0, total_chunks):
start = chunk_id * OBJECT_TRANSFER_CHUNK_SIZE
end = min(total_size, (chunk_id + 1) * OBJECT_TRANSFER_CHUNK_SIZE)
chunk = ray_client_pb2.ClientTask(
type=req.task.type,
name=req.task.name,
payload_id=req.task.payload_id,
client_id=req.task.client_id,
options=req.task.options,
baseline_options=req.task.baseline_options,
namespace=req.task.namespace,
data=request_data[start:end],
chunk_id=chunk_id,
total_chunks=total_chunks,
)
yield ray_client_pb2.DataRequest(req_id=req.req_id, task=chunk)
class ChunkCollector:
"""
This object collects chunks from async get requests via __call__, and
calls the underlying callback when the object is fully received, or if an
exception while retrieving the object occurs.
This is not used in synchronous gets (synchronous gets interact with the
raylet servicer directly, not through the datapath).
__call__ returns true once the underlying call back has been called.
"""
def __init__(self, callback: ResponseCallable, request: ray_client_pb2.DataRequest):
# Bytearray containing data received so far
self.data = bytearray()
# The callback that will be called once all data is received
self.callback = callback
# The id of the last chunk we've received, or -1 if haven't seen any yet
self.last_seen_chunk = -1
# The GetRequest that initiated the transfer. start_chunk_id will be
# updated as chunks are received to avoid re-requesting chunks that
# we've already received.
self.request = request
def __call__(self, response: Union[ray_client_pb2.DataResponse, Exception]) -> bool:
if isinstance(response, Exception):
self.callback(response)
return True
get_resp = response.get
if not get_resp.valid:
self.callback(response)
return True
if get_resp.total_size > OBJECT_TRANSFER_WARNING_SIZE and log_once(
"client_object_transfer_size_warning"
):
size_gb = get_resp.total_size / 2**30
warnings.warn(
"Ray Client is attempting to retrieve a "
f"{size_gb:.2f} GiB object over the network, which may "
"be slow. Consider serializing the object to a file and "
"using rsync or S3 instead.",
UserWarning,
)
chunk_data = get_resp.data
chunk_id = get_resp.chunk_id
if chunk_id == self.last_seen_chunk + 1:
self.data.extend(chunk_data)
self.last_seen_chunk = chunk_id
# If we disconnect partway through, restart the get request
# at the first chunk we haven't seen
self.request.get.start_chunk_id = self.last_seen_chunk + 1
elif chunk_id > self.last_seen_chunk + 1:
# A chunk was skipped. This shouldn't happen in practice since
# grpc guarantees that chunks will arrive in order.
msg = (
f"Received chunk {chunk_id} when we expected "
f"{self.last_seen_chunk + 1} for request {response.req_id}"
)
logger.warning(msg)
self.callback(RuntimeError(msg))
return True
else:
# We received a chunk that've already seen before. Ignore, since
# it should already be appended to self.data.
logger.debug(
f"Received a repeated chunk {chunk_id} "
f"from request {response.req_id}."
)
if get_resp.chunk_id == get_resp.total_chunks - 1:
self.callback(self.data)
return True
else:
# Not done yet
return False
class DataClient:
def __init__(self, client_worker: "Worker", client_id: str, metadata: list):
"""Initializes a thread-safe datapath over a Ray Client gRPC channel.
Args:
client_worker: The Ray Client worker that manages this client
client_id: the generated ID representing this client
metadata: metadata to pass to gRPC requests
"""
self.client_worker = client_worker
self._client_id = client_id
self._metadata = metadata
self.data_thread = self._start_datathread()
# Track outstanding requests to resend in case of disconnection
self.outstanding_requests: Dict[int, Any] = OrderedDict()
# Serialize access to all mutable internal states: self.request_queue,
# self.ready_data, self.asyncio_waiting_data,
# self._in_shutdown, self._req_id, self.outstanding_requests and
# calling self._next_id()
self.lock = threading.Lock()
# Waiting for response or shutdown.
self.cv = threading.Condition(lock=self.lock)
self.request_queue = self._create_queue()
self.ready_data: Dict[int, Any] = {}
# NOTE: Dictionary insertion is guaranteed to complete before lookup
# and/or removal because of synchronization via the request_queue.
self.asyncio_waiting_data: Dict[int, ResponseCallable] = {}
self._in_shutdown = False
self._req_id = 0
self._last_exception = None
self._acknowledge_counter = 0
self.data_thread.start()
# Must hold self.lock when calling this function.
def _next_id(self) -> int:
assert self.lock.locked()
self._req_id += 1
if self._req_id > INT32_MAX:
self._req_id = 1
# Responses that aren't tracked (like opportunistic releases)
# have req_id=0, so make sure we never mint such an id.
assert self._req_id != 0
return self._req_id
def _start_datathread(self) -> threading.Thread:
return threading.Thread(
target=self._data_main,
name="ray_client_streaming_rpc",
args=(),
daemon=True,
)
# A helper that takes requests from queue. If the request wraps a PutRequest,
# lazily chunks and yields the request. Otherwise, yields the request directly.
def _requests(self):
while True:
req = self.request_queue.get()
if req is None:
# Stop when client signals shutdown.
return
req_type = req.WhichOneof("type")
if req_type == "put":
yield from chunk_put(req)
elif req_type == "task":
yield from chunk_task(req)
else:
yield req
def _data_main(self) -> None:
reconnecting = False
try:
while not self.client_worker._in_shutdown:
stub = ray_client_pb2_grpc.RayletDataStreamerStub(
self.client_worker.channel
)
metadata = self._metadata + [("reconnecting", str(reconnecting))]
resp_stream = stub.Datapath(
self._requests(),
metadata=metadata,
wait_for_ready=True,
)
try:
for response in resp_stream:
self._process_response(response)
return
except grpc.RpcError as e:
reconnecting = self._can_reconnect(e)
if not reconnecting:
self._last_exception = e
return
self._reconnect_channel()
except Exception as e:
self._last_exception = e
finally:
logger.debug("Shutting down data channel.")
self._shutdown()
def _process_response(self, response: Any) -> None:
"""
Process responses from the data servicer.
"""
if response.req_id == 0:
# This is not being waited for.
logger.debug(f"Got unawaited response {response}")
return
if response.req_id in self.asyncio_waiting_data:
can_remove = True
try:
callback = self.asyncio_waiting_data[response.req_id]
if isinstance(callback, ChunkCollector):
can_remove = callback(response)
elif callback:
callback(response)
if can_remove:
# NOTE: calling del self.asyncio_waiting_data results
# in the destructor of ClientObjectRef running, which
# calls ReleaseObject(). So self.asyncio_waiting_data
# is accessed without holding self.lock. Holding the
# lock shouldn't be necessary either.
del self.asyncio_waiting_data[response.req_id]
except Exception:
logger.exception("Callback error:")
with self.lock:
# Update outstanding requests
if response.req_id in self.outstanding_requests and can_remove:
del self.outstanding_requests[response.req_id]
# Acknowledge response
self._acknowledge(response.req_id)
else:
with self.lock:
self.ready_data[response.req_id] = response
self.cv.notify_all()
def _can_reconnect(self, e: grpc.RpcError) -> bool:
"""
Processes RPC errors that occur while reading from data stream.
Returns True if the error can be recovered from, False otherwise.
"""
if not self.client_worker._can_reconnect(e):
logger.error("Unrecoverable error in data channel.")
logger.debug(e)
return False
logger.debug("Recoverable error in data channel.")
logger.debug(e)
return True
def _shutdown(self) -> None:
"""
Shutdown the data channel
"""
with self.lock:
self._in_shutdown = True
self.cv.notify_all()
callbacks = self.asyncio_waiting_data.values()
self.asyncio_waiting_data = {}
if self._last_exception:
# Abort async requests with the error.
err = ConnectionError(
"Failed during this or a previous request. Exception that "
f"broke the connection: {self._last_exception}"
)
else:
err = ConnectionError(
"Request cannot be fulfilled because the data client has "
"disconnected."
)
for callback in callbacks:
if callback:
callback(err)
# Since self._in_shutdown is set to True, no new item
# will be added to self.asyncio_waiting_data
def _acknowledge(self, req_id: int) -> None:
"""
Puts an acknowledge request on the request queue periodically.
Lock should be held before calling this. Used when an async or
blocking response is received.
"""
if not self.client_worker._reconnect_enabled:
# Skip ACKs if reconnect isn't enabled
return
assert self.lock.locked()
self._acknowledge_counter += 1
if self._acknowledge_counter % ACKNOWLEDGE_BATCH_SIZE == 0:
self.request_queue.put(
ray_client_pb2.DataRequest(
acknowledge=ray_client_pb2.AcknowledgeRequest(req_id=req_id)
)
)
def _reconnect_channel(self) -> None:
"""
Attempts to reconnect the gRPC channel and resend outstanding
requests. First, the server is pinged to see if the current channel
still works. If the ping fails, then the current channel is closed
and replaced with a new one.
Once a working channel is available, a new request queue is made
and filled with any outstanding requests to be resent to the server.
"""
try:
# Ping the server to see if the current channel is reuseable, for
# example if gRPC reconnected the channel on its own or if the
# RPC error was transient and the channel is still open
ping_succeeded = self.client_worker.ping_server(timeout=5)
except grpc.RpcError:
ping_succeeded = False
if not ping_succeeded:
# Ping failed, try refreshing the data channel
logger.warning(
"Encountered connection issues in the data channel. "
"Attempting to reconnect."
)
try:
self.client_worker._connect_channel(reconnecting=True)
except ConnectionError:
logger.warning("Failed to reconnect the data channel")
raise
logger.debug("Reconnection succeeded!")
# Recreate the request queue, and resend outstanding requests
with self.lock:
self.request_queue = self._create_queue()
for request in self.outstanding_requests.values():
# Resend outstanding requests
self.request_queue.put(request)
# Use SimpleQueue to avoid deadlocks when appending to queue from __del__()
@staticmethod
def _create_queue():
return queue.SimpleQueue()
def close(self) -> None:
thread = None
with self.lock:
self._in_shutdown = True
# Notify blocking operations to fail.
self.cv.notify_all()
# Add sentinel to terminate streaming RPC.
if self.request_queue is not None:
# Intentional shutdown, tell server it can clean up the
# connection immediately and ignore the reconnect grace period.
cleanup_request = ray_client_pb2.DataRequest(
connection_cleanup=ray_client_pb2.ConnectionCleanupRequest()
)
self.request_queue.put(cleanup_request)
self.request_queue.put(None)
if self.data_thread is not None:
thread = self.data_thread
# Wait until streaming RPCs are done.
if thread is not None:
thread.join()
def _blocking_send(
self, req: ray_client_pb2.DataRequest
) -> ray_client_pb2.DataResponse:
with self.lock:
self._check_shutdown()
req_id = self._next_id()
req.req_id = req_id
self.request_queue.put(req)
self.outstanding_requests[req_id] = req
self.cv.wait_for(lambda: req_id in self.ready_data or self._in_shutdown)
self._check_shutdown()
data = self.ready_data[req_id]
del self.ready_data[req_id]
del self.outstanding_requests[req_id]
self._acknowledge(req_id)
return data
def _async_send(
self,
req: ray_client_pb2.DataRequest,
callback: Optional[ResponseCallable] = None,
) -> None:
with self.lock:
self._check_shutdown()
req_id = self._next_id()
req.req_id = req_id
self.asyncio_waiting_data[req_id] = callback
self.outstanding_requests[req_id] = req
self.request_queue.put(req)
# Must hold self.lock when calling this function.
def _check_shutdown(self):
assert self.lock.locked()
if not self._in_shutdown:
return
self.lock.release()
# Do not try disconnect() or throw exceptions in self.data_thread.
# Otherwise deadlock can occur.
if threading.current_thread().ident == self.data_thread.ident:
return
from ray.util import disconnect
disconnect()
self.lock.acquire()
if self._last_exception is not None:
msg = (
"Request can't be sent because the Ray client has already "
"been disconnected due to an error. Last exception: "
f"{self._last_exception}"
)
else:
msg = (
"Request can't be sent because the Ray client has already "
"been disconnected."
)
raise ConnectionError(msg)
def Init(
self, request: ray_client_pb2.InitRequest, context=None
) -> ray_client_pb2.InitResponse:
datareq = ray_client_pb2.DataRequest(
init=request,
)
resp = self._blocking_send(datareq)
return resp.init
def PrepRuntimeEnv(
self, request: ray_client_pb2.PrepRuntimeEnvRequest, context=None
) -> ray_client_pb2.PrepRuntimeEnvResponse:
datareq = ray_client_pb2.DataRequest(
prep_runtime_env=request,
)
resp = self._blocking_send(datareq)
return resp.prep_runtime_env
def ConnectionInfo(self, context=None) -> ray_client_pb2.ConnectionInfoResponse:
datareq = ray_client_pb2.DataRequest(
connection_info=ray_client_pb2.ConnectionInfoRequest()
)
resp = self._blocking_send(datareq)
return resp.connection_info
def GetObject(
self, request: ray_client_pb2.GetRequest, context=None
) -> ray_client_pb2.GetResponse:
datareq = ray_client_pb2.DataRequest(
get=request,
)
resp = self._blocking_send(datareq)
return resp.get
def RegisterGetCallback(
self, request: ray_client_pb2.GetRequest, callback: ResponseCallable
) -> None:
if len(request.ids) != 1:
raise ValueError(
"RegisterGetCallback() must have exactly 1 Object ID. "
f"Actual: {request}"
)
datareq = ray_client_pb2.DataRequest(
get=request,
)
collector = ChunkCollector(callback=callback, request=datareq)
self._async_send(datareq, collector)
# TODO: convert PutObject to async
def PutObject(
self, request: ray_client_pb2.PutRequest, context=None
) -> ray_client_pb2.PutResponse:
datareq = ray_client_pb2.DataRequest(
put=request,
)
resp = self._blocking_send(datareq)
return resp.put
def ReleaseObject(
self, request: ray_client_pb2.ReleaseRequest, context=None
) -> None:
datareq = ray_client_pb2.DataRequest(
release=request,
)
self._async_send(datareq)
def Schedule(self, request: ray_client_pb2.ClientTask, callback: ResponseCallable):
datareq = ray_client_pb2.DataRequest(task=request)
self._async_send(datareq, callback)
def Terminate(
self, request: ray_client_pb2.TerminateRequest
) -> ray_client_pb2.TerminateResponse:
req = ray_client_pb2.DataRequest(
terminate=request,
)
resp = self._blocking_send(req)
return resp.terminate
def ListNamedActors(
self, request: ray_client_pb2.ClientListNamedActorsRequest
) -> ray_client_pb2.ClientListNamedActorsResponse:
req = ray_client_pb2.DataRequest(
list_named_actors=request,
)
resp = self._blocking_send(req)
return resp.list_named_actors
python/ray/util/client/examples/run_tune.py
+6
View File
@@ -0,0 +1,6 @@
from ray.tune import tune
from ray.util.client import ray
ray.connect("localhost:50051")
tune.run("PG", config={"env": "CartPole-v0"})
python/ray/util/client/logsclient.py
+135
View File
@@ -0,0 +1,135 @@
"""This file implements a threaded stream controller to return logs back from
the ray clientserver.
"""
import logging
import queue
import sys
import threading
import time
from typing import TYPE_CHECKING
import grpc
import ray.core.generated.ray_client_pb2 as ray_client_pb2
import ray.core.generated.ray_client_pb2_grpc as ray_client_pb2_grpc
from ray.util.debug import log_once
if TYPE_CHECKING:
from ray.util.client.worker import Worker
logger = logging.getLogger(__name__)
# TODO(barakmich): Running a logger in a logger causes loopback.
# The client logger need its own root -- possibly this one.
# For the moment, let's just not propagate beyond this point.
logger.propagate = False
class LogstreamClient:
def __init__(self, client_worker: "Worker", metadata: list):
"""Initializes a thread-safe log stream over a Ray Client gRPC channel.
Args:
client_worker: The Ray Client worker that manages this client
metadata: metadata to pass to gRPC requests
"""
self.client_worker = client_worker
self._metadata = metadata
self.request_queue = queue.Queue()
self.log_thread = self._start_logthread()
self.log_thread.start()
self.last_req = None
def _start_logthread(self) -> threading.Thread:
return threading.Thread(target=self._log_main, args=(), daemon=True)
def _log_main(self) -> None:
reconnecting = False
while not self.client_worker._in_shutdown:
if reconnecting:
# Refresh queue and retry last request
self.request_queue = queue.Queue()
if self.last_req:
self.request_queue.put(self.last_req)
stub = ray_client_pb2_grpc.RayletLogStreamerStub(self.client_worker.channel)
try:
log_stream = stub.Logstream(
iter(self.request_queue.get, None), metadata=self._metadata
)
except ValueError:
# Trying to use the stub on a cancelled channel will raise
# ValueError. This should only happen when the data client
# is attempting to reset the connection -- sleep and try
# again.
time.sleep(0.5)
continue
try:
for record in log_stream:
if record.level < 0:
self.stdstream(level=record.level, msg=record.msg)
self.log(level=record.level, msg=record.msg)
return
except grpc.RpcError as e:
reconnecting = self._process_rpc_error(e)
if not reconnecting:
return
def _process_rpc_error(self, e: grpc.RpcError) -> bool:
"""
Processes RPC errors that occur while reading from data stream.
Returns True if the error can be recovered from, False otherwise.
"""
if self.client_worker._can_reconnect(e):
if log_once("lost_reconnect_logs"):
logger.warning(
"Log channel is reconnecting. Logs produced while "
"the connection was down can be found on the head "
"node of the cluster in "
"`ray_client_server_[port].out`"
)
logger.debug("Log channel dropped, retrying.")
time.sleep(0.5)
return True
logger.debug("Shutting down log channel.")
if not self.client_worker._in_shutdown:
logger.exception("Unexpected exception:")
return False
def log(self, level: int, msg: str):
"""Log the message from the log stream.
By default, calls logger.log but this can be overridden.
Args:
level: The loglevel of the received log message
msg: The content of the message
"""
logger.log(level=level, msg=msg)
def stdstream(self, level: int, msg: str):
"""Log the stdout/stderr entry from the log stream.
By default, calls print but this can be overridden.
Args:
level: The loglevel of the received log message
msg: The content of the message
"""
print_file = sys.stderr if level == -2 else sys.stdout
print(msg, file=print_file, end="")
def set_logstream_level(self, level: int):
logger.setLevel(level)
req = ray_client_pb2.LogSettingsRequest()
req.enabled = True
req.loglevel = level
self.request_queue.put(req)
self.last_req = req
def close(self) -> None:
self.request_queue.put(None)
if self.log_thread is not None:
self.log_thread.join()
def disable_logs(self) -> None:
req = ray_client_pb2.LogSettingsRequest()
req.enabled = False
self.request_queue.put(req)
self.last_req = req
python/ray/util/client/options.py
+45
View File
@@ -0,0 +1,45 @@
from typing import Any, Dict, Optional
from ray._common import ray_option_utils
from ray.util.placement_group import PlacementGroup, check_placement_group_index
from ray.util.scheduling_strategies import PlacementGroupSchedulingStrategy
def validate_options(kwargs_dict: Optional[Dict[str, Any]]) -> Optional[Dict[str, Any]]:
if kwargs_dict is None:
return None
if len(kwargs_dict) == 0:
return None
out = {}
for k, v in kwargs_dict.items():
if k not in ray_option_utils.valid_options:
raise ValueError(
f"Invalid option keyword: '{k}'. "
f"{ray_option_utils.remote_args_error_string}"
)
ray_option_utils.valid_options[k].validate(k, v)
out[k] = v
# Validate placement setting similar to the logic in ray/actor.py and
# ray/remote_function.py. The difference is that when
# placement_group = default and placement_group_capture_child_tasks
# specified, placement group cannot be resolved at client. So this check
# skips this case and relies on server to enforce any condition.
bundle_index = out.get("placement_group_bundle_index", None)
pg = out.get("placement_group", None)
scheduling_strategy = out.get("scheduling_strategy", None)
if isinstance(scheduling_strategy, PlacementGroupSchedulingStrategy):
pg = scheduling_strategy.placement_group
bundle_index = scheduling_strategy.placement_group_bundle_index
if bundle_index is not None:
if pg is None:
pg = PlacementGroup.empty()
if pg == "default" and (
out.get("placement_group_capture_child_tasks", None) is None
):
pg = PlacementGroup.empty()
if isinstance(pg, PlacementGroup):
check_placement_group_index(pg, bundle_index)
return out
python/ray/util/client/ray_client_helpers.py
+83
View File
@@ -0,0 +1,83 @@
import time
from contextlib import contextmanager
from typing import Any, Dict
import ray as real_ray
import ray.util.client.server.server as ray_client_server
from ray._common.network_utils import build_address, get_localhost_ip
from ray._private.client_mode_hook import disable_client_hook
from ray.job_config import JobConfig
from ray.util.client import ray
@contextmanager
def ray_start_client_server(metadata=None, ray_connect_handler=None, **kwargs):
with ray_start_client_server_pair(
metadata=metadata, ray_connect_handler=ray_connect_handler, **kwargs
) as pair:
client, server = pair
yield client
@contextmanager
def ray_start_client_server_for_address(address):
"""
Starts a Ray client server that initializes drivers at the specified address.
"""
def connect_handler(
job_config: JobConfig = None, **ray_init_kwargs: Dict[str, Any]
):
import ray
with disable_client_hook():
if not ray.is_initialized():
return ray.init(address, job_config=job_config, **ray_init_kwargs)
with ray_start_client_server(ray_connect_handler=connect_handler) as ray:
yield ray
@contextmanager
def ray_start_client_server_pair(metadata=None, ray_connect_handler=None, **kwargs):
ray._inside_client_test = True
with disable_client_hook():
assert not ray.is_initialized()
server = ray_client_server.serve(
get_localhost_ip(), 50051, ray_connect_handler=ray_connect_handler
)
ray.connect(build_address(get_localhost_ip(), 50051), metadata=metadata, **kwargs)
try:
yield ray, server
finally:
ray._inside_client_test = False
ray.disconnect()
server.stop(0)
del server
start = time.monotonic()
with disable_client_hook():
while ray.is_initialized():
time.sleep(1)
if time.monotonic() - start > 30:
raise RuntimeError("Failed to terminate Ray")
# Allow windows to close processes before moving on
time.sleep(3)
@contextmanager
def ray_start_cluster_client_server_pair(address):
ray._inside_client_test = True
def ray_connect_handler(job_config=None, **ray_init_kwargs):
real_ray.init(address=address)
server = ray_client_server.serve(
get_localhost_ip(), 50051, ray_connect_handler=ray_connect_handler
)
ray.connect(build_address(get_localhost_ip(), 50051))
try:
yield ray, server
finally:
ray._inside_client_test = False
ray.disconnect()
server.stop(0)
python/ray/util/client/runtime_context.py
+77
View File
@@ -0,0 +1,77 @@
from types import SimpleNamespace
from typing import TYPE_CHECKING
if TYPE_CHECKING:
from ray import JobID, NodeID, WorkerID
from ray.runtime_context import RuntimeContext
class _ClientWorkerPropertyAPI:
"""Emulates the properties of the ray._private.worker object for the client"""
def __init__(self, worker):
assert worker is not None
self.worker = worker
def build_runtime_context(self) -> "RuntimeContext":
"""Creates a RuntimeContext backed by the properites of this API"""
# Defer the import of RuntimeContext until needed to avoid cycles
from ray.runtime_context import RuntimeContext
return RuntimeContext(self)
def _fetch_runtime_context(self):
import ray.core.generated.ray_client_pb2 as ray_client_pb2
return self.worker.get_cluster_info(
ray_client_pb2.ClusterInfoType.RUNTIME_CONTEXT
)
@property
def mode(self):
from ray._private.worker import SCRIPT_MODE
return SCRIPT_MODE
@property
def current_job_id(self) -> "JobID":
from ray import JobID
return JobID(self._fetch_runtime_context().job_id)
@property
def current_node_id(self) -> "NodeID":
from ray import NodeID
return NodeID(self._fetch_runtime_context().node_id)
@property
def worker_id(self) -> "WorkerID":
"""Binary worker id for the Ray Client server process (cluster driver worker)."""
from ray import WorkerID
return WorkerID(self._fetch_runtime_context().worker_id)
@property
def namespace(self) -> str:
return self._fetch_runtime_context().namespace
@property
def should_capture_child_tasks_in_placement_group(self) -> bool:
return self._fetch_runtime_context().capture_client_tasks
@property
def runtime_env(self) -> str:
return self._fetch_runtime_context().runtime_env
def check_connected(self) -> bool:
return self.worker.ping_server()
@property
def gcs_client(self) -> str:
return SimpleNamespace(address=self._fetch_runtime_context().gcs_address)
@property
def node(self):
"""Emulates the worker.node property for client mode"""
return SimpleNamespace(session_name=self._fetch_runtime_context().session_name)
python/ray/util/client/server/__init__.py
+1
View File
@@ -0,0 +1 @@
from ray.util.client.server.server import serve # noqa
python/ray/util/client/server/__main__.py
+4
View File
@@ -0,0 +1,4 @@
if __name__ == "__main__":
from ray.util.client.server.server import main
main()
python/ray/util/client/server/dataservicer.py
+415
View File
@@ -0,0 +1,415 @@
import logging
import sys
import time
from collections import defaultdict
from queue import Queue
from threading import Event, Lock, Thread
from typing import TYPE_CHECKING, Any, Dict, Iterator, Union
import grpc
import ray
import ray.core.generated.ray_client_pb2 as ray_client_pb2
import ray.core.generated.ray_client_pb2_grpc as ray_client_pb2_grpc
from ray._private.client_mode_hook import disable_client_hook
from ray.util.client.common import (
CLIENT_SERVER_MAX_THREADS,
OrderedResponseCache,
_propagate_error_in_context,
)
from ray.util.client.server.server_pickler import loads_from_client
from ray.util.debug import log_once
if TYPE_CHECKING:
from ray.util.client.server.server import RayletServicer
logger = logging.getLogger(__name__)
QUEUE_JOIN_SECONDS = 10
def _get_reconnecting_from_context(context: Any) -> bool:
"""
Get `reconnecting` from gRPC metadata, or False if missing.
"""
metadata = dict(context.invocation_metadata())
val = metadata.get("reconnecting")
if val is None or val not in ("True", "False"):
logger.error(
f'Client connecting with invalid value for "reconnecting": {val}, '
"This may be because you have a mismatched client and server "
"version."
)
return False
return val == "True"
def _should_cache(req: ray_client_pb2.DataRequest) -> bool:
"""
Returns True if the response should to the given request should be cached,
false otherwise. At the moment the only requests we do not cache are:
- asynchronous gets: These arrive out of order. Skipping caching here
is fine, since repeating an async get is idempotent
- acks: Repeating acks is idempotent
- clean up requests: Also idempotent, and client has likely already
wrapped up the data connection by this point.
- puts: We should only cache when we receive the final chunk, since
any earlier chunks won't generate a response
- tasks: We should only cache when we receive the final chunk,
since any earlier chunks won't generate a response
"""
req_type = req.WhichOneof("type")
if req_type == "get" and req.get.asynchronous:
return False
if req_type == "put":
return req.put.chunk_id == req.put.total_chunks - 1
if req_type == "task":
return req.task.chunk_id == req.task.total_chunks - 1
return req_type not in ("acknowledge", "connection_cleanup")
def fill_queue(
grpc_input_generator: Iterator[ray_client_pb2.DataRequest],
output_queue: "Queue[Union[ray_client_pb2.DataRequest, ray_client_pb2.DataResponse]]", # noqa: E501
) -> None:
"""
Pushes incoming requests to a shared output_queue.
"""
try:
for req in grpc_input_generator:
output_queue.put(req)
except grpc.RpcError as e:
logger.debug(
"closing dataservicer reader thread "
f"grpc error reading request_iterator: {e}"
)
finally:
# Set the sentinel value for the output_queue
output_queue.put(None)
class ChunkCollector:
"""
Helper class for collecting chunks from PutObject or ClientTask messages
"""
def __init__(self):
self.curr_req_id = None
self.last_seen_chunk_id = -1
self.data = bytearray()
def add_chunk(
self,
req: ray_client_pb2.DataRequest,
chunk: Union[ray_client_pb2.PutRequest, ray_client_pb2.ClientTask],
):
if self.curr_req_id is not None and self.curr_req_id != req.req_id:
raise RuntimeError(
"Expected to receive a chunk from request with id "
f"{self.curr_req_id}, but found {req.req_id} instead."
)
self.curr_req_id = req.req_id
next_chunk = self.last_seen_chunk_id + 1
if chunk.chunk_id < next_chunk:
# Repeated chunk, ignore
return
if chunk.chunk_id > next_chunk:
raise RuntimeError(
f"A chunk {chunk.chunk_id} of request {req.req_id} was "
"received out of order."
)
elif chunk.chunk_id == self.last_seen_chunk_id + 1:
self.data.extend(chunk.data)
self.last_seen_chunk_id = chunk.chunk_id
return chunk.chunk_id + 1 == chunk.total_chunks
def reset(self):
self.curr_req_id = None
self.last_seen_chunk_id = -1
self.data = bytearray()
class DataServicer(ray_client_pb2_grpc.RayletDataStreamerServicer):
def __init__(self, basic_service: "RayletServicer"):
self.basic_service = basic_service
self.clients_lock = Lock()
self.num_clients = 0 # guarded by self.clients_lock
# dictionary mapping client_id's to the last time they connected
self.client_last_seen: Dict[str, float] = {}
# dictionary mapping client_id's to their reconnect grace periods
self.reconnect_grace_periods: Dict[str, float] = {}
# dictionary mapping client_id's to their response cache
self.response_caches: Dict[str, OrderedResponseCache] = defaultdict(
OrderedResponseCache
)
# stopped event, useful for signals that the server is shut down
self.stopped = Event()
# Helper for collecting chunks from PutObject calls. Assumes that
# that put requests from different objects aren't interleaved.
self.put_request_chunk_collector = ChunkCollector()
# Helper for collecting chunks from ClientTask calls. Assumes that
# schedule requests from different remote calls aren't interleaved.
self.client_task_chunk_collector = ChunkCollector()
def Datapath(self, request_iterator, context):
start_time = time.time()
# set to True if client shuts down gracefully
cleanup_requested = False
metadata = dict(context.invocation_metadata())
client_id = metadata.get("client_id")
if client_id is None:
logger.error("Client connecting with no client_id")
return
logger.debug(f"New data connection from client {client_id}: ")
accepted_connection = self._init(client_id, context, start_time)
response_cache = self.response_caches[client_id]
# Set to False if client requests a reconnect grace period of 0
reconnect_enabled = True
if not accepted_connection:
return
try:
request_queue = Queue()
queue_filler_thread = Thread(
target=fill_queue, daemon=True, args=(request_iterator, request_queue)
)
queue_filler_thread.start()
"""For non `async get` requests, this loop yields immediately
For `async get` requests, this loop:
1) does not yield, it just continues
2) When the result is ready, it yields
"""
for req in iter(request_queue.get, None):
if isinstance(req, ray_client_pb2.DataResponse):
# Early shortcut if this is the result of an async get.
yield req
continue
assert isinstance(req, ray_client_pb2.DataRequest)
if _should_cache(req) and reconnect_enabled:
cached_resp = response_cache.check_cache(req.req_id)
if isinstance(cached_resp, Exception):
# Cache state is invalid, raise exception
raise cached_resp
if cached_resp is not None:
yield cached_resp
continue
resp = None
req_type = req.WhichOneof("type")
if req_type == "init":
resp_init = self.basic_service.Init(req.init)
resp = ray_client_pb2.DataResponse(
init=resp_init,
)
with self.clients_lock:
self.reconnect_grace_periods[
client_id
] = req.init.reconnect_grace_period
if req.init.reconnect_grace_period == 0:
reconnect_enabled = False
elif req_type == "get":
if req.get.asynchronous:
get_resp = self.basic_service._async_get_object(
req.get, client_id, req.req_id, request_queue
)
if get_resp is None:
# Skip sending a response for this request and
# continue to the next requst. The response for
# this request will be sent when the object is
# ready.
continue
else:
get_resp = self.basic_service._get_object(req.get, client_id)
resp = ray_client_pb2.DataResponse(get=get_resp)
elif req_type == "put":
if not self.put_request_chunk_collector.add_chunk(req, req.put):
# Put request still in progress
continue
put_resp = self.basic_service._put_object(
self.put_request_chunk_collector.data,
req.put.client_ref_id,
client_id,
)
self.put_request_chunk_collector.reset()
resp = ray_client_pb2.DataResponse(put=put_resp)
elif req_type == "release":
released = []
for rel_id in req.release.ids:
rel = self.basic_service.release(client_id, rel_id)
released.append(rel)
resp = ray_client_pb2.DataResponse(
release=ray_client_pb2.ReleaseResponse(ok=released)
)
elif req_type == "connection_info":
resp = ray_client_pb2.DataResponse(
connection_info=self._build_connection_response()
)
elif req_type == "prep_runtime_env":
with self.clients_lock:
resp_prep = self.basic_service.PrepRuntimeEnv(
req.prep_runtime_env
)
resp = ray_client_pb2.DataResponse(prep_runtime_env=resp_prep)
elif req_type == "connection_cleanup":
cleanup_requested = True
cleanup_resp = ray_client_pb2.ConnectionCleanupResponse()
resp = ray_client_pb2.DataResponse(connection_cleanup=cleanup_resp)
elif req_type == "acknowledge":
# Clean up acknowledged cache entries
response_cache.cleanup(req.acknowledge.req_id)
continue
elif req_type == "task":
with self.clients_lock:
task = req.task
if not self.client_task_chunk_collector.add_chunk(req, task):
# Not all serialized arguments have arrived
continue
arglist, kwargs = loads_from_client(
self.client_task_chunk_collector.data, self.basic_service
)
self.client_task_chunk_collector.reset()
resp_ticket = self.basic_service.Schedule(
req.task, arglist, kwargs, context
)
resp = ray_client_pb2.DataResponse(task_ticket=resp_ticket)
del arglist
del kwargs
elif req_type == "terminate":
with self.clients_lock:
response = self.basic_service.Terminate(req.terminate, context)
resp = ray_client_pb2.DataResponse(terminate=response)
elif req_type == "list_named_actors":
with self.clients_lock:
response = self.basic_service.ListNamedActors(
req.list_named_actors
)
resp = ray_client_pb2.DataResponse(list_named_actors=response)
else:
raise Exception(
f"Unreachable code: Request type "
f"{req_type} not handled in Datapath"
)
resp.req_id = req.req_id
if _should_cache(req) and reconnect_enabled:
response_cache.update_cache(req.req_id, resp)
yield resp
except Exception as e:
logger.exception("Error in data channel:")
recoverable = _propagate_error_in_context(e, context)
invalid_cache = response_cache.invalidate(e)
if not recoverable or invalid_cache:
context.set_code(grpc.StatusCode.FAILED_PRECONDITION)
# Connection isn't recoverable, skip cleanup
cleanup_requested = True
finally:
logger.debug(f"Stream is broken with client {client_id}")
queue_filler_thread.join(QUEUE_JOIN_SECONDS)
if queue_filler_thread.is_alive():
logger.error(
"Queue filler thread failed to join before timeout: {}".format(
QUEUE_JOIN_SECONDS
)
)
cleanup_delay = self.reconnect_grace_periods.get(client_id)
if not cleanup_requested and cleanup_delay is not None:
logger.debug(
"Cleanup wasn't requested, delaying cleanup by"
f"{cleanup_delay} seconds."
)
# Delay cleanup, since client may attempt a reconnect
# Wait on the "stopped" event in case the grpc server is
# stopped and we can clean up earlier.
self.stopped.wait(timeout=cleanup_delay)
else:
logger.debug("Cleanup was requested, cleaning up immediately.")
with self.clients_lock:
if client_id not in self.client_last_seen:
logger.debug("Connection already cleaned up.")
# Some other connection has already cleaned up this
# this client's session. This can happen if the client
# reconnects and then gracefully shut's down immediately.
return
last_seen = self.client_last_seen[client_id]
if last_seen > start_time:
# The client successfully reconnected and updated
# last seen some time during the grace period
logger.debug("Client reconnected, skipping cleanup")
return
# Either the client shut down gracefully, or the client
# failed to reconnect within the grace period. Clean up
# the connection.
self.basic_service.release_all(client_id)
del self.client_last_seen[client_id]
if client_id in self.reconnect_grace_periods:
del self.reconnect_grace_periods[client_id]
if client_id in self.response_caches:
del self.response_caches[client_id]
self.num_clients -= 1
logger.debug(
f"Removed client {client_id}, " f"remaining={self.num_clients}"
)
# It's important to keep the Ray shutdown
# within this locked context or else Ray could hang.
# NOTE: it is strange to start ray in server.py but shut it
# down here. Consider consolidating ray lifetime management.
with disable_client_hook():
if self.num_clients == 0:
logger.debug("Shutting down ray.")
ray.shutdown()
def _init(self, client_id: str, context: Any, start_time: float):
"""
Checks if resources allow for another client.
Returns a boolean indicating if initialization was successful.
"""
with self.clients_lock:
reconnecting = _get_reconnecting_from_context(context)
threshold = int(CLIENT_SERVER_MAX_THREADS / 2)
if self.num_clients >= threshold:
logger.warning(
f"[Data Servicer]: Num clients {self.num_clients} "
f"has reached the threshold {threshold}. "
f"Rejecting client: {client_id}. "
)
if log_once("client_threshold"):
logger.warning(
"You can configure the client connection "
"threshold by setting the "
"RAY_CLIENT_SERVER_MAX_THREADS env var "
f"(currently set to {CLIENT_SERVER_MAX_THREADS})."
)
context.set_code(grpc.StatusCode.RESOURCE_EXHAUSTED)
return False
if reconnecting and client_id not in self.client_last_seen:
# Client took too long to reconnect, session has been
# cleaned up.
context.set_code(grpc.StatusCode.NOT_FOUND)
context.set_details(
"Attempted to reconnect to a session that has already "
"been cleaned up."
)
return False
if client_id in self.client_last_seen:
logger.debug(f"Client {client_id} has reconnected.")
else:
self.num_clients += 1
logger.debug(
f"Accepted data connection from {client_id}. "
f"Total clients: {self.num_clients}"
)
self.client_last_seen[client_id] = start_time
return True
def _build_connection_response(self):
with self.clients_lock:
cur_num_clients = self.num_clients
return ray_client_pb2.ConnectionInfoResponse(
num_clients=cur_num_clients,
python_version="{}.{}.{}".format(
sys.version_info[0], sys.version_info[1], sys.version_info[2]
),
ray_version=ray.__version__,
ray_commit=ray.__commit__,
)
python/ray/util/client/server/logservicer.py
+125
View File
@@ -0,0 +1,125 @@
"""This file responds to log stream requests and forwards logs
with its handler.
"""
import io
import logging
import queue
import threading
import uuid
import grpc
import ray.core.generated.ray_client_pb2 as ray_client_pb2
import ray.core.generated.ray_client_pb2_grpc as ray_client_pb2_grpc
from ray._private.ray_logging import global_worker_stdstream_dispatcher
from ray._private.worker import print_worker_logs
from ray.util.client.common import CLIENT_SERVER_MAX_THREADS
logger = logging.getLogger(__name__)
class LogstreamHandler(logging.Handler):
def __init__(self, queue, level):
super().__init__()
self.queue = queue
self.level = level
def emit(self, record: logging.LogRecord):
logdata = ray_client_pb2.LogData()
logdata.msg = record.getMessage()
logdata.level = record.levelno
logdata.name = record.name
self.queue.put(logdata)
class StdStreamHandler:
def __init__(self, queue):
self.queue = queue
self.id = str(uuid.uuid4())
def handle(self, data):
logdata = ray_client_pb2.LogData()
logdata.level = -2 if data["is_err"] else -1
logdata.name = "stderr" if data["is_err"] else "stdout"
with io.StringIO() as file:
print_worker_logs(data, file)
logdata.msg = file.getvalue()
self.queue.put(logdata)
def register_global(self):
global_worker_stdstream_dispatcher.add_handler(self.id, self.handle)
def unregister_global(self):
global_worker_stdstream_dispatcher.remove_handler(self.id)
def log_status_change_thread(log_queue, request_iterator):
std_handler = StdStreamHandler(log_queue)
current_handler = None
root_logger = logging.getLogger("ray")
default_level = root_logger.getEffectiveLevel()
try:
for req in request_iterator:
if current_handler is not None:
root_logger.setLevel(default_level)
root_logger.removeHandler(current_handler)
std_handler.unregister_global()
if not req.enabled:
current_handler = None
continue
current_handler = LogstreamHandler(log_queue, req.loglevel)
std_handler.register_global()
root_logger.addHandler(current_handler)
root_logger.setLevel(req.loglevel)
except grpc.RpcError as e:
logger.debug(f"closing log thread " f"grpc error reading request_iterator: {e}")
finally:
if current_handler is not None:
root_logger.setLevel(default_level)
root_logger.removeHandler(current_handler)
std_handler.unregister_global()
log_queue.put(None)
class LogstreamServicer(ray_client_pb2_grpc.RayletLogStreamerServicer):
def __init__(self):
super().__init__()
self.num_clients = 0
self.client_lock = threading.Lock()
def Logstream(self, request_iterator, context):
initialized = False
with self.client_lock:
threshold = CLIENT_SERVER_MAX_THREADS / 2
if self.num_clients + 1 >= threshold:
context.set_code(grpc.StatusCode.RESOURCE_EXHAUSTED)
logger.warning(
f"Logstream: Num clients {self.num_clients} has reached "
f"the threshold {threshold}. Rejecting new connection."
)
return
self.num_clients += 1
initialized = True
logger.info(
"New logs connection established. " f"Total clients: {self.num_clients}"
)
log_queue = queue.Queue()
thread = threading.Thread(
target=log_status_change_thread,
args=(log_queue, request_iterator),
daemon=True,
)
thread.start()
try:
queue_iter = iter(log_queue.get, None)
for record in queue_iter:
if record is None:
break
yield record
except grpc.RpcError as e:
logger.debug(f"Closing log channel: {e}")
finally:
thread.join()
with self.client_lock:
if initialized:
self.num_clients -= 1
python/ray/util/client/server/proxier.py
+938
View File
@@ -0,0 +1,938 @@
import atexit
import json
import logging
import socket
import sys
import time
import traceback
import urllib
from concurrent import futures
from dataclasses import dataclass
from itertools import chain
from threading import Event, Lock, RLock, Thread
from typing import Callable, Dict, List, Optional, Tuple
from urllib.parse import urlparse, urlunparse
import grpc
import ray
import ray.core.generated.ray_client_pb2 as ray_client_pb2
import ray.core.generated.ray_client_pb2_grpc as ray_client_pb2_grpc
import ray.core.generated.runtime_env_agent_pb2 as runtime_env_agent_pb2
from ray._common.network_utils import (
build_address,
get_localhost_ip,
is_ipv6,
is_localhost,
)
from ray._common.tls_utils import add_port_to_grpc_server
from ray._common.utils import env_integer
from ray._private.authentication.http_token_authentication import (
format_authentication_http_error,
get_auth_headers_if_auth_enabled,
)
from ray._private.client_mode_hook import disable_client_hook
from ray._private.grpc_utils import init_grpc_channel
from ray._private.parameter import RayParams
from ray._private.runtime_env.context import RuntimeEnvContext
from ray._private.services import (
ProcessInfo,
get_node_with_retry,
start_ray_client_server,
)
from ray._private.utils import detect_fate_sharing_support
from ray._raylet import GcsClient
from ray.cloudpickle.compat import pickle
from ray.exceptions import AuthenticationError
from ray.job_config import JobConfig
from ray.util.client.common import (
CLIENT_SERVER_MAX_THREADS,
GRPC_OPTIONS,
ClientServerHandle,
_get_client_id_from_context,
_propagate_error_in_context,
)
from ray.util.client.server.dataservicer import _get_reconnecting_from_context
# Import psutil after ray so the packaged version is used.
import psutil
logger = logging.getLogger(__name__)
CHECK_PROCESS_INTERVAL_S = 30
MIN_SPECIFIC_SERVER_PORT = 23000
MAX_SPECIFIC_SERVER_PORT = 24000
CHECK_CHANNEL_TIMEOUT_S = env_integer("RAY_CLIENT_SERVER_CHECK_CHANNEL_TIMEOUT_S", 30)
LOGSTREAM_RETRIES = 5
LOGSTREAM_RETRY_INTERVAL_SEC = 2
@dataclass
class SpecificServer:
port: int
process_handle_future: futures.Future
channel: "grpc._channel.Channel"
def is_ready(self) -> bool:
"""Check if the server is ready or not (doesn't block)."""
return self.process_handle_future.done()
def wait_ready(self, timeout: Optional[float] = None) -> None:
"""
Wait for the server to actually start up.
"""
res = self.process_handle_future.result(timeout=timeout)
if res is None:
# This is only set to none when server creation specifically fails.
raise RuntimeError("Server startup failed.")
def poll(self) -> Optional[int]:
"""Check if the process has exited."""
try:
proc = self.process_handle_future.result(timeout=0.1)
if proc is not None:
return proc.process.poll()
except futures.TimeoutError:
return
def kill(self) -> None:
"""Try to send a KILL signal to the process."""
try:
proc = self.process_handle_future.result(timeout=0.1)
if proc is not None:
proc.process.kill()
except futures.TimeoutError:
# Server has not been started yet.
pass
def set_result(self, proc: Optional[ProcessInfo]) -> None:
"""Set the result of the internal future if it is currently unset."""
if not self.is_ready():
self.process_handle_future.set_result(proc)
def _match_running_client_server(command: List[str]) -> bool:
"""
Detects if the main process in the given command is the RayClient Server.
This works by ensuring that the command is of the form:
<py_executable> -m ray.util.client.server <args>
"""
flattened = " ".join(command)
return "-m ray.util.client.server" in flattened
class ProxyManager:
def __init__(
self,
address: Optional[str],
runtime_env_agent_address: str,
*,
session_dir: Optional[str] = None,
redis_username: Optional[str] = None,
redis_password: Optional[str] = None,
node_id: Optional[str] = None,
):
self.servers: Dict[str, SpecificServer] = dict()
self.server_lock = RLock()
self._address = address
self._redis_username = redis_username
self._redis_password = redis_password
self._free_ports: List[int] = list(
range(MIN_SPECIFIC_SERVER_PORT, MAX_SPECIFIC_SERVER_PORT)
)
if runtime_env_agent_address:
parsed = urlparse(runtime_env_agent_address)
# runtime env agent self-assigns a free port, fetch it from GCS
if parsed.port is None or parsed.port == 0:
if node_id is None:
raise ValueError(
"node_id is required when runtime_env_agent_address "
"has no port specified"
)
node_info = get_node_with_retry(address, node_id)
runtime_env_agent_address = urlunparse(
parsed._replace(
netloc=f"{parsed.hostname}:{node_info['runtime_env_agent_port']}"
)
)
self._runtime_env_agent_address = runtime_env_agent_address
self._check_thread = Thread(target=self._check_processes, daemon=True)
self._check_thread.start()
self.fate_share = bool(detect_fate_sharing_support())
self._node: Optional[ray._private.node.Node] = None
atexit.register(self._cleanup)
def _get_unused_port(self, family: int = socket.AF_INET) -> int:
"""
Search for a port in _free_ports that is unused.
"""
with self.server_lock:
num_ports = len(self._free_ports)
for _ in range(num_ports):
port = self._free_ports.pop(0)
s = socket.socket(family, socket.SOCK_STREAM)
try:
s.bind(("", port))
except OSError:
self._free_ports.append(port)
continue
finally:
s.close()
return port
raise RuntimeError("Unable to succeed in selecting a random port.")
@property
def address(self) -> str:
"""
Returns the provided Ray bootstrap address, or creates a new cluster.
"""
if self._address:
return self._address
# Start a new, locally scoped cluster.
connection_tuple = ray.init()
self._address = connection_tuple["address"]
self._session_dir = connection_tuple["session_dir"]
return self._address
@property
def node(self) -> ray._private.node.Node:
"""Gets a 'ray.Node' object for this node (the head node).
If it does not already exist, one is created using the bootstrap
address.
"""
if self._node:
return self._node
ray_params = RayParams(gcs_address=self.address)
self._node = ray._private.node.Node(
ray_params,
head=False,
shutdown_at_exit=False,
spawn_reaper=False,
connect_only=True,
)
return self._node
def create_specific_server(self, client_id: str) -> SpecificServer:
"""
Create, but not start a SpecificServer for a given client. This
method must be called once per client.
"""
with self.server_lock:
assert (
self.servers.get(client_id) is None
), f"Server already created for Client: {client_id}"
host = get_localhost_ip()
port = self._get_unused_port(
socket.AF_INET6 if is_ipv6(host) else socket.AF_INET
)
server = SpecificServer(
port=port,
process_handle_future=futures.Future(),
channel=init_grpc_channel(
build_address(host, port), options=GRPC_OPTIONS
),
)
self.servers[client_id] = server
return server
def _create_runtime_env(
self,
serialized_runtime_env: str,
runtime_env_config: str,
specific_server: SpecificServer,
):
"""Increase the runtime_env reference by sending an RPC to the agent.
Includes retry logic to handle the case when the agent is
temporarily unreachable (e.g., hasn't been started up yet).
"""
logger.info(
f"Increasing runtime env reference for "
f"ray_client_server_{specific_server.port}."
f"Serialized runtime env is {serialized_runtime_env}."
)
assert (
len(self._runtime_env_agent_address) > 0
), "runtime_env_agent_address not set"
create_env_request = runtime_env_agent_pb2.GetOrCreateRuntimeEnvRequest(
serialized_runtime_env=serialized_runtime_env,
runtime_env_config=runtime_env_config,
job_id=f"ray_client_server_{specific_server.port}".encode("utf-8"),
source_process="client_server",
)
retries = 0
max_retries = 5
wait_time_s = 0.5
last_exception = None
while retries <= max_retries:
try:
url = urllib.parse.urljoin(
self._runtime_env_agent_address, "/get_or_create_runtime_env"
)
data = create_env_request.SerializeToString()
headers = {"Content-Type": "application/octet-stream"}
headers.update(**get_auth_headers_if_auth_enabled(headers))
req = urllib.request.Request(
url, data=data, method="POST", headers=headers
)
response = urllib.request.urlopen(req, timeout=None)
response_data = response.read()
r = runtime_env_agent_pb2.GetOrCreateRuntimeEnvReply()
r.ParseFromString(response_data)
if r.status == runtime_env_agent_pb2.AgentRpcStatus.AGENT_RPC_STATUS_OK:
return r.serialized_runtime_env_context
elif (
r.status
== runtime_env_agent_pb2.AgentRpcStatus.AGENT_RPC_STATUS_FAILED
):
raise RuntimeError(
"Failed to create runtime_env for Ray client "
f"server, it is caused by:\n{r.error_message}"
)
else:
assert False, f"Unknown status: {r.status}."
except urllib.error.HTTPError as e:
body = ""
try:
body = e.read().decode("utf-8", "ignore")
except Exception:
body = e.reason if hasattr(e, "reason") else str(e)
formatted_error = format_authentication_http_error(e.code, body or "")
if formatted_error:
raise AuthenticationError(formatted_error) from e
# Treat non-auth HTTP errors like URLError (retry with backoff)
last_exception = e
logger.warning(
f"GetOrCreateRuntimeEnv request failed with HTTP {e.code}: {body or e}. "
f"Retrying after {wait_time_s}s. "
f"{max_retries-retries} retries remaining."
)
except urllib.error.URLError as e:
last_exception = e
logger.warning(
f"GetOrCreateRuntimeEnv request failed: {e}. "
f"Retrying after {wait_time_s}s. "
f"{max_retries-retries} retries remaining."
)
# Exponential backoff.
time.sleep(wait_time_s)
retries += 1
wait_time_s *= 2
raise TimeoutError(
f"GetOrCreateRuntimeEnv request failed after {max_retries} attempts."
f" Last exception: {last_exception}"
)
def start_specific_server(self, client_id: str, job_config: JobConfig) -> bool:
"""
Start up a RayClient Server for an incoming client to
communicate with. Returns whether creation was successful.
"""
specific_server = self._get_server_for_client(client_id)
assert specific_server, f"Server has not been created for: {client_id}"
output, error = self.node.get_log_file_handles(
f"ray_client_server_{specific_server.port}", unique=True
)
serialized_runtime_env = job_config._get_serialized_runtime_env()
runtime_env_config = job_config._get_proto_runtime_env_config()
if not serialized_runtime_env or serialized_runtime_env == "{}":
# TODO(edoakes): can we just remove this case and always send it
# to the agent?
serialized_runtime_env_context = RuntimeEnvContext().serialize()
else:
serialized_runtime_env_context = self._create_runtime_env(
serialized_runtime_env=serialized_runtime_env,
runtime_env_config=runtime_env_config,
specific_server=specific_server,
)
proc = start_ray_client_server(
self.address,
get_localhost_ip(),
specific_server.port,
stdout_file=output,
stderr_file=error,
fate_share=self.fate_share,
server_type="specific-server",
serialized_runtime_env_context=serialized_runtime_env_context,
redis_username=self._redis_username,
redis_password=self._redis_password,
)
# Wait for the process being run transitions from the shim process
# to the actual RayClient Server.
pid = proc.process.pid
if sys.platform != "win32":
psutil_proc = psutil.Process(pid)
else:
psutil_proc = None
# Don't use `psutil` on Win32
while psutil_proc is not None:
if proc.process.poll() is not None:
logger.error(f"SpecificServer startup failed for client: {client_id}")
break
cmd = psutil_proc.cmdline()
if _match_running_client_server(cmd):
break
logger.debug("Waiting for Process to reach the actual client server.")
time.sleep(0.5)
specific_server.set_result(proc)
logger.info(
f"SpecificServer started on port: {specific_server.port} "
f"with PID: {pid} for client: {client_id}"
)
return proc.process.poll() is None
def _get_server_for_client(self, client_id: str) -> Optional[SpecificServer]:
with self.server_lock:
client = self.servers.get(client_id)
if client is None:
logger.error(f"Unable to find channel for client: {client_id}")
return client
def has_channel(self, client_id: str) -> bool:
server = self._get_server_for_client(client_id)
if server is None:
return False
return server.is_ready()
def get_channel(
self,
client_id: str,
) -> Optional["grpc._channel.Channel"]:
"""
Find the gRPC Channel for the given client_id. This will block until
the server process has started.
"""
server = self._get_server_for_client(client_id)
if server is None:
return None
# Wait for the SpecificServer to become ready.
server.wait_ready()
try:
grpc.channel_ready_future(server.channel).result(
timeout=CHECK_CHANNEL_TIMEOUT_S
)
return server.channel
except grpc.FutureTimeoutError:
logger.exception(f"Timeout waiting for channel for {client_id}")
return None
def _check_processes(self):
"""
Keeps the internal servers dictionary up-to-date with running servers.
"""
while True:
with self.server_lock:
for client_id, specific_server in list(self.servers.items()):
if specific_server.poll() is not None:
logger.info(
f"Specific server {client_id} is no longer running"
f", freeing its port {specific_server.port}"
)
del self.servers[client_id]
# Port is available to use again.
self._free_ports.append(specific_server.port)
time.sleep(CHECK_PROCESS_INTERVAL_S)
def _cleanup(self) -> None:
"""
Forcibly kill all spawned RayClient Servers. This ensures cleanup
for platforms where fate sharing is not supported.
"""
for server in self.servers.values():
server.kill()
class RayletServicerProxy(ray_client_pb2_grpc.RayletDriverServicer):
def __init__(self, ray_connect_handler: Callable, proxy_manager: ProxyManager):
self.proxy_manager = proxy_manager
self.ray_connect_handler = ray_connect_handler
def _call_inner_function(
self, request, context, method: str
) -> Optional[ray_client_pb2_grpc.RayletDriverStub]:
client_id = _get_client_id_from_context(context)
chan = self.proxy_manager.get_channel(client_id)
if not chan:
logger.error(f"Channel for Client: {client_id} not found!")
context.set_code(grpc.StatusCode.NOT_FOUND)
return None
stub = ray_client_pb2_grpc.RayletDriverStub(chan)
try:
metadata = [("client_id", client_id)]
if context:
metadata = context.invocation_metadata()
return getattr(stub, method)(request, metadata=metadata)
except Exception as e:
# Error while proxying -- propagate the error's context to user
logger.exception(f"Proxying call to {method} failed!")
_propagate_error_in_context(e, context)
def _has_channel_for_request(self, context):
client_id = _get_client_id_from_context(context)
return self.proxy_manager.has_channel(client_id)
def Init(self, request, context=None) -> ray_client_pb2.InitResponse:
return self._call_inner_function(request, context, "Init")
def KVPut(self, request, context=None) -> ray_client_pb2.KVPutResponse:
"""Proxies internal_kv.put.
This is used by the working_dir code to upload to the GCS before
ray.init is called. In that case (if we don't have a server yet)
we directly make the internal KV call from the proxier.
Otherwise, we proxy the call to the downstream server as usual.
"""
if self._has_channel_for_request(context):
return self._call_inner_function(request, context, "KVPut")
with disable_client_hook():
already_exists = ray.experimental.internal_kv._internal_kv_put(
request.key, request.value, overwrite=request.overwrite
)
return ray_client_pb2.KVPutResponse(already_exists=already_exists)
def KVGet(self, request, context=None) -> ray_client_pb2.KVGetResponse:
"""Proxies internal_kv.get.
This is used by the working_dir code to upload to the GCS before
ray.init is called. In that case (if we don't have a server yet)
we directly make the internal KV call from the proxier.
Otherwise, we proxy the call to the downstream server as usual.
"""
if self._has_channel_for_request(context):
return self._call_inner_function(request, context, "KVGet")
with disable_client_hook():
value = ray.experimental.internal_kv._internal_kv_get(request.key)
return ray_client_pb2.KVGetResponse(value=value)
def KVDel(self, request, context=None) -> ray_client_pb2.KVDelResponse:
"""Proxies internal_kv.delete.
This is used by the working_dir code to upload to the GCS before
ray.init is called. In that case (if we don't have a server yet)
we directly make the internal KV call from the proxier.
Otherwise, we proxy the call to the downstream server as usual.
"""
if self._has_channel_for_request(context):
return self._call_inner_function(request, context, "KVDel")
with disable_client_hook():
ray.experimental.internal_kv._internal_kv_del(request.key)
return ray_client_pb2.KVDelResponse()
def KVList(self, request, context=None) -> ray_client_pb2.KVListResponse:
"""Proxies internal_kv.list.
This is used by the working_dir code to upload to the GCS before
ray.init is called. In that case (if we don't have a server yet)
we directly make the internal KV call from the proxier.
Otherwise, we proxy the call to the downstream server as usual.
"""
if self._has_channel_for_request(context):
return self._call_inner_function(request, context, "KVList")
with disable_client_hook():
keys = ray.experimental.internal_kv._internal_kv_list(request.prefix)
return ray_client_pb2.KVListResponse(keys=keys)
def KVExists(self, request, context=None) -> ray_client_pb2.KVExistsResponse:
"""Proxies internal_kv.exists.
This is used by the working_dir code to upload to the GCS before
ray.init is called. In that case (if we don't have a server yet)
we directly make the internal KV call from the proxier.
Otherwise, we proxy the call to the downstream server as usual.
"""
if self._has_channel_for_request(context):
return self._call_inner_function(request, context, "KVExists")
with disable_client_hook():
exists = ray.experimental.internal_kv._internal_kv_exists(request.key)
return ray_client_pb2.KVExistsResponse(exists=exists)
def PinRuntimeEnvURI(
self, request, context=None
) -> ray_client_pb2.ClientPinRuntimeEnvURIResponse:
"""Proxies internal_kv.pin_runtime_env_uri.
This is used by the working_dir code to upload to the GCS before
ray.init is called. In that case (if we don't have a server yet)
we directly make the internal KV call from the proxier.
Otherwise, we proxy the call to the downstream server as usual.
"""
if self._has_channel_for_request(context):
return self._call_inner_function(request, context, "PinRuntimeEnvURI")
with disable_client_hook():
ray.experimental.internal_kv._pin_runtime_env_uri(
request.uri, expiration_s=request.expiration_s
)
return ray_client_pb2.ClientPinRuntimeEnvURIResponse()
def ListNamedActors(
self, request, context=None
) -> ray_client_pb2.ClientListNamedActorsResponse:
return self._call_inner_function(request, context, "ListNamedActors")
def ClusterInfo(self, request, context=None) -> ray_client_pb2.ClusterInfoResponse:
# NOTE: We need to respond to the PING request here to allow the client
# to continue with connecting.
if request.type == ray_client_pb2.ClusterInfoType.PING:
resp = ray_client_pb2.ClusterInfoResponse(json=json.dumps({}))
return resp
return self._call_inner_function(request, context, "ClusterInfo")
def Terminate(self, req, context=None):
return self._call_inner_function(req, context, "Terminate")
def GetObject(self, request, context=None):
try:
yield from self._call_inner_function(request, context, "GetObject")
except Exception as e:
# Error while iterating over response from GetObject stream
logger.exception("Proxying call to GetObject failed!")
_propagate_error_in_context(e, context)
def PutObject(
self, request: ray_client_pb2.PutRequest, context=None
) -> ray_client_pb2.PutResponse:
return self._call_inner_function(request, context, "PutObject")
def WaitObject(self, request, context=None) -> ray_client_pb2.WaitResponse:
return self._call_inner_function(request, context, "WaitObject")
def Schedule(self, task, context=None) -> ray_client_pb2.ClientTaskTicket:
return self._call_inner_function(task, context, "Schedule")
def ray_client_server_env_prep(job_config: JobConfig) -> JobConfig:
return job_config
def prepare_runtime_init_req(
init_request: ray_client_pb2.DataRequest,
) -> Tuple[ray_client_pb2.DataRequest, JobConfig]:
"""
Extract JobConfig and possibly mutate InitRequest before it is passed to
the specific RayClient Server.
"""
init_type = init_request.WhichOneof("type")
assert init_type == "init", (
"Received initial message of type " f"{init_type}, not 'init'."
)
req = init_request.init
job_config = JobConfig()
if req.job_config:
job_config = pickle.loads(req.job_config)
new_job_config = ray_client_server_env_prep(job_config)
modified_init_req = ray_client_pb2.InitRequest(
job_config=pickle.dumps(new_job_config),
ray_init_kwargs=init_request.init.ray_init_kwargs,
reconnect_grace_period=init_request.init.reconnect_grace_period,
)
init_request.init.CopyFrom(modified_init_req)
return (init_request, new_job_config)
class RequestIteratorProxy:
def __init__(self, request_iterator):
self.request_iterator = request_iterator
def __iter__(self):
return self
def __next__(self):
try:
return next(self.request_iterator)
except grpc.RpcError as e:
# To stop proxying already CANCLLED request stream gracefully,
# we only translate the exact grpc.RpcError to StopIteration,
# not its subsclasses. ex: grpc._Rendezvous
# https://github.com/grpc/grpc/blob/v1.43.0/src/python/grpcio/grpc/_server.py#L353-L354
# This fixes the https://github.com/ray-project/ray/issues/23865
if type(e) is not grpc.RpcError:
raise e # re-raise other grpc exceptions
logger.exception(
"Stop iterating cancelled request stream with the following exception:"
)
raise StopIteration
class DataServicerProxy(ray_client_pb2_grpc.RayletDataStreamerServicer):
def __init__(self, proxy_manager: ProxyManager):
self.num_clients = 0
# dictionary mapping client_id's to the last time they connected
self.clients_last_seen: Dict[str, float] = {}
self.reconnect_grace_periods: Dict[str, float] = {}
self.clients_lock = Lock()
self.proxy_manager = proxy_manager
self.stopped = Event()
def modify_connection_info_resp(
self, init_resp: ray_client_pb2.DataResponse
) -> ray_client_pb2.DataResponse:
"""
Modify the `num_clients` returned the ConnectionInfoResponse because
individual SpecificServers only have **one** client.
"""
init_type = init_resp.WhichOneof("type")
if init_type != "connection_info":
return init_resp
modified_resp = ray_client_pb2.DataResponse()
modified_resp.CopyFrom(init_resp)
with self.clients_lock:
modified_resp.connection_info.num_clients = self.num_clients
return modified_resp
def Datapath(self, request_iterator, context):
request_iterator = RequestIteratorProxy(request_iterator)
cleanup_requested = False
start_time = time.time()
client_id = _get_client_id_from_context(context)
if client_id == "":
return
reconnecting = _get_reconnecting_from_context(context)
if reconnecting:
with self.clients_lock:
if client_id not in self.clients_last_seen:
# Client took too long to reconnect, session has already
# been cleaned up
context.set_code(grpc.StatusCode.NOT_FOUND)
context.set_details(
"Attempted to reconnect a session that has already "
"been cleaned up"
)
return
self.clients_last_seen[client_id] = start_time
server = self.proxy_manager._get_server_for_client(client_id)
channel = self.proxy_manager.get_channel(client_id)
# iterator doesn't need modification on reconnect
new_iter = request_iterator
else:
# Create Placeholder *before* reading the first request.
server = self.proxy_manager.create_specific_server(client_id)
with self.clients_lock:
self.clients_last_seen[client_id] = start_time
self.num_clients += 1
try:
if not reconnecting:
logger.info(f"New data connection from client {client_id}: ")
init_req = next(request_iterator)
with self.clients_lock:
self.reconnect_grace_periods[
client_id
] = init_req.init.reconnect_grace_period
try:
modified_init_req, job_config = prepare_runtime_init_req(init_req)
if not self.proxy_manager.start_specific_server(
client_id, job_config
):
logger.error(
f"Server startup failed for client: {client_id}, "
f"using JobConfig: {job_config}!"
)
raise RuntimeError(
"Starting Ray client server failed. See "
f"ray_client_server_{server.port}.err for "
"detailed logs."
)
channel = self.proxy_manager.get_channel(client_id)
if channel is None:
logger.error(f"Channel not found for {client_id}")
raise RuntimeError(
"Proxy failed to Connect to backend! Check "
"`ray_client_server.err` and "
f"`ray_client_server_{server.port}.err` on the "
"head node of the cluster for the relevant logs. "
"By default these are located at "
"/tmp/ray/session_latest/logs."
)
except Exception:
init_resp = ray_client_pb2.DataResponse(
init=ray_client_pb2.InitResponse(
ok=False, msg=traceback.format_exc()
)
)
init_resp.req_id = init_req.req_id
yield init_resp
return None
new_iter = chain([modified_init_req], request_iterator)
stub = ray_client_pb2_grpc.RayletDataStreamerStub(channel)
metadata = [("client_id", client_id), ("reconnecting", str(reconnecting))]
resp_stream = stub.Datapath(new_iter, metadata=metadata)
for resp in resp_stream:
resp_type = resp.WhichOneof("type")
if resp_type == "connection_cleanup":
# Specific server is skipping cleanup, proxier should too
cleanup_requested = True
yield self.modify_connection_info_resp(resp)
except Exception as e:
logger.exception("Proxying Datapath failed!")
# Propogate error through context
recoverable = _propagate_error_in_context(e, context)
if not recoverable:
# Client shouldn't attempt to recover, clean up connection
cleanup_requested = True
finally:
cleanup_delay = self.reconnect_grace_periods.get(client_id)
if not cleanup_requested and cleanup_delay is not None:
# Delay cleanup, since client may attempt a reconnect
# Wait on stopped event in case the server closes and we
# can clean up earlier
self.stopped.wait(timeout=cleanup_delay)
with self.clients_lock:
if client_id not in self.clients_last_seen:
logger.info(f"{client_id} not found. Skipping clean up.")
# Connection has already been cleaned up
return
last_seen = self.clients_last_seen[client_id]
logger.info(
f"{client_id} last started stream at {last_seen}. Current "
f"stream started at {start_time}."
)
if last_seen > start_time:
logger.info("Client reconnected. Skipping cleanup.")
# Client has reconnected, don't clean up
return
logger.debug(f"Client detached: {client_id}")
self.num_clients -= 1
del self.clients_last_seen[client_id]
if client_id in self.reconnect_grace_periods:
del self.reconnect_grace_periods[client_id]
server.set_result(None)
class LogstreamServicerProxy(ray_client_pb2_grpc.RayletLogStreamerServicer):
def __init__(self, proxy_manager: ProxyManager):
super().__init__()
self.proxy_manager = proxy_manager
def Logstream(self, request_iterator, context):
request_iterator = RequestIteratorProxy(request_iterator)
client_id = _get_client_id_from_context(context)
if client_id == "":
return
logger.debug(f"New logstream connection from client {client_id}: ")
channel = None
# We need to retry a few times because the LogClient *may* connect
# Before the DataClient has finished connecting.
for i in range(LOGSTREAM_RETRIES):
channel = self.proxy_manager.get_channel(client_id)
if channel is not None:
break
logger.warning(f"Retrying Logstream connection. {i+1} attempts failed.")
time.sleep(LOGSTREAM_RETRY_INTERVAL_SEC)
if channel is None:
context.set_code(grpc.StatusCode.NOT_FOUND)
context.set_details(
"Logstream proxy failed to connect. Channel for client "
f"{client_id} not found."
)
return None
stub = ray_client_pb2_grpc.RayletLogStreamerStub(channel)
resp_stream = stub.Logstream(
request_iterator, metadata=[("client_id", client_id)]
)
try:
for resp in resp_stream:
yield resp
except Exception:
logger.exception("Proxying Logstream failed!")
def serve_proxier(
host: str,
port: int,
gcs_address: Optional[str],
*,
redis_username: Optional[str] = None,
redis_password: Optional[str] = None,
session_dir: Optional[str] = None,
runtime_env_agent_address: Optional[str] = None,
node_id: Optional[str] = None,
):
# Initialize internal KV to be used to upload and download working_dir
# before calling ray.init within the RayletServicers.
# NOTE(edoakes): redis_address and redis_password should only be None in
# tests.
if gcs_address is not None:
gcs_cli = GcsClient(address=gcs_address)
ray.experimental.internal_kv._initialize_internal_kv(gcs_cli)
from ray._private.grpc_utils import create_grpc_server_with_interceptors
server = create_grpc_server_with_interceptors(
max_workers=CLIENT_SERVER_MAX_THREADS,
thread_name_prefix="ray_client_proxier",
options=GRPC_OPTIONS,
asynchronous=False,
)
proxy_manager = ProxyManager(
gcs_address,
session_dir=session_dir,
redis_username=redis_username,
redis_password=redis_password,
runtime_env_agent_address=runtime_env_agent_address,
node_id=node_id,
)
task_servicer = RayletServicerProxy(None, proxy_manager)
data_servicer = DataServicerProxy(proxy_manager)
logs_servicer = LogstreamServicerProxy(proxy_manager)
ray_client_pb2_grpc.add_RayletDriverServicer_to_server(task_servicer, server)
ray_client_pb2_grpc.add_RayletDataStreamerServicer_to_server(data_servicer, server)
ray_client_pb2_grpc.add_RayletLogStreamerServicer_to_server(logs_servicer, server)
if not is_localhost(host):
add_port_to_grpc_server(server, build_address(get_localhost_ip(), port))
add_port_to_grpc_server(server, build_address(host, port))
server.start()
return ClientServerHandle(
task_servicer=task_servicer,
data_servicer=data_servicer,
logs_servicer=logs_servicer,
grpc_server=server,
)
python/ray/util/client/server/server.py
+968
View File
@@ -0,0 +1,968 @@
import base64
import functools
import gc
import inspect
import json
import logging
import math
import os
import pickle
import queue
import threading
import time
from collections import defaultdict
from typing import Any, Callable, Dict, List, Optional, Set, Union
import grpc
import ray
import ray._private.state
import ray.core.generated.ray_client_pb2 as ray_client_pb2
import ray.core.generated.ray_client_pb2_grpc as ray_client_pb2_grpc
from ray import cloudpickle
from ray._common.network_utils import (
build_address,
get_all_interfaces_ip,
get_localhost_ip,
is_localhost,
)
from ray._common.tls_utils import add_port_to_grpc_server
from ray._private import ray_constants
from ray._private.client_mode_hook import disable_client_hook
from ray._private.ray_constants import env_integer
from ray._private.ray_logging import setup_logger
from ray._private.ray_logging.logging_config import LoggingConfig
from ray._private.services import canonicalize_bootstrap_address_or_die
from ray._raylet import GcsClient
from ray.job_config import JobConfig
from ray.util.client.common import (
CLIENT_SERVER_MAX_THREADS,
GRPC_OPTIONS,
OBJECT_TRANSFER_CHUNK_SIZE,
ClientServerHandle,
ResponseCache,
)
from ray.util.client.server.dataservicer import DataServicer
from ray.util.client.server.logservicer import LogstreamServicer
from ray.util.client.server.proxier import serve_proxier
from ray.util.client.server.server_pickler import dumps_from_server, loads_from_client
from ray.util.client.server.server_stubs import current_server
logger = logging.getLogger(__name__)
TIMEOUT_FOR_SPECIFIC_SERVER_S = env_integer("TIMEOUT_FOR_SPECIFIC_SERVER_S", 30)
def _use_response_cache(func):
"""
Decorator for gRPC stubs. Before calling the real stubs, checks if there's
an existing entry in the caches. If there is, then return the cached
entry. Otherwise, call the real function and use the real cache
"""
@functools.wraps(func)
def wrapper(self, request, context):
metadata = dict(context.invocation_metadata())
expected_ids = ("client_id", "thread_id", "req_id")
if any(i not in metadata for i in expected_ids):
# Missing IDs, skip caching and call underlying stub directly
return func(self, request, context)
# Get relevant IDs to check cache
client_id = metadata["client_id"]
thread_id = metadata["thread_id"]
req_id = int(metadata["req_id"])
# Check if response already cached
response_cache = self.response_caches[client_id]
cached_entry = response_cache.check_cache(thread_id, req_id)
if cached_entry is not None:
if isinstance(cached_entry, Exception):
# Original call errored, propagate error
context.set_code(grpc.StatusCode.FAILED_PRECONDITION)
context.set_details(str(cached_entry))
raise cached_entry
return cached_entry
try:
# Response wasn't cached, call underlying stub and cache result
resp = func(self, request, context)
except Exception as e:
# Unexpected error in underlying stub -- update cache and
# propagate to user through context
response_cache.update_cache(thread_id, req_id, e)
context.set_code(grpc.StatusCode.FAILED_PRECONDITION)
context.set_details(str(e))
raise
response_cache.update_cache(thread_id, req_id, resp)
return resp
return wrapper
class RayletServicer(ray_client_pb2_grpc.RayletDriverServicer):
def __init__(self, ray_connect_handler: Callable):
"""Construct a raylet service
Args:
ray_connect_handler: Function to connect to ray cluster
"""
# Stores client_id -> (ref_id -> ObjectRef)
self.object_refs: Dict[str, Dict[bytes, ray.ObjectRef]] = defaultdict(dict)
# Stores client_id -> (client_ref_id -> ref_id (in self.object_refs))
self.client_side_ref_map: Dict[str, Dict[bytes, bytes]] = defaultdict(dict)
self.function_refs = {}
self.actor_refs: Dict[bytes, ray.ActorHandle] = {}
self.actor_owners: Dict[str, Set[bytes]] = defaultdict(set)
self.registered_actor_classes = {}
self.named_actors = set()
self.state_lock = threading.Lock()
self.ray_connect_handler = ray_connect_handler
self.response_caches: Dict[str, ResponseCache] = defaultdict(ResponseCache)
def Init(
self, request: ray_client_pb2.InitRequest, context=None
) -> ray_client_pb2.InitResponse:
if request.job_config:
job_config = pickle.loads(request.job_config)
job_config._client_job = True
else:
job_config = None
current_job_config = None
with disable_client_hook():
if ray.is_initialized():
worker = ray._private.worker.global_worker
current_job_config = worker.core_worker.get_job_config()
else:
extra_kwargs = json.loads(request.ray_init_kwargs or "{}")
# Reconstruct LoggingConfig from dict after InitRequest ray_init_kwargs is parsed from JSON on the server.
if "logging_config" in extra_kwargs and isinstance(
extra_kwargs["logging_config"], dict
):
extra_kwargs["logging_config"] = LoggingConfig.from_dict(
extra_kwargs["logging_config"]
)
try:
self.ray_connect_handler(job_config, **extra_kwargs)
except Exception as e:
logger.exception("Running Ray Init failed:")
return ray_client_pb2.InitResponse(
ok=False,
msg=f"Call to `ray.init()` on the server failed with: {e}",
)
if job_config is None:
return ray_client_pb2.InitResponse(ok=True)
# NOTE(edoakes): this code should not be necessary anymore because we
# only allow a single client/job per server. There is an existing test
# that tests the behavior of multiple clients with the same job config
# connecting to one server (test_client_init.py::test_num_clients),
# so I'm leaving it here for now.
job_config = job_config._get_proto_job_config()
# If the server has been initialized, we need to compare whether the
# runtime env is compatible.
if current_job_config:
job_uris = set(job_config.runtime_env_info.uris.working_dir_uri)
job_uris.update(job_config.runtime_env_info.uris.py_modules_uris)
current_job_uris = set(
current_job_config.runtime_env_info.uris.working_dir_uri
)
current_job_uris.update(
current_job_config.runtime_env_info.uris.py_modules_uris
)
if job_uris != current_job_uris and len(job_uris) > 0:
return ray_client_pb2.InitResponse(
ok=False,
msg="Runtime environment doesn't match "
f"request one {job_config.runtime_env_info.uris} "
f"current one {current_job_config.runtime_env_info.uris}",
)
return ray_client_pb2.InitResponse(ok=True)
@_use_response_cache
def KVPut(self, request, context=None) -> ray_client_pb2.KVPutResponse:
try:
with disable_client_hook():
already_exists = ray.experimental.internal_kv._internal_kv_put(
request.key,
request.value,
overwrite=request.overwrite,
namespace=request.namespace,
)
except Exception as e:
return_exception_in_context(e, context)
already_exists = False
return ray_client_pb2.KVPutResponse(already_exists=already_exists)
def KVGet(self, request, context=None) -> ray_client_pb2.KVGetResponse:
try:
with disable_client_hook():
value = ray.experimental.internal_kv._internal_kv_get(
request.key, namespace=request.namespace
)
except Exception as e:
return_exception_in_context(e, context)
value = b""
return ray_client_pb2.KVGetResponse(value=value)
@_use_response_cache
def KVDel(self, request, context=None) -> ray_client_pb2.KVDelResponse:
try:
with disable_client_hook():
deleted_num = ray.experimental.internal_kv._internal_kv_del(
request.key,
del_by_prefix=request.del_by_prefix,
namespace=request.namespace,
)
except Exception as e:
return_exception_in_context(e, context)
deleted_num = 0
return ray_client_pb2.KVDelResponse(deleted_num=deleted_num)
def KVList(self, request, context=None) -> ray_client_pb2.KVListResponse:
try:
with disable_client_hook():
keys = ray.experimental.internal_kv._internal_kv_list(
request.prefix, namespace=request.namespace
)
except Exception as e:
return_exception_in_context(e, context)
keys = []
return ray_client_pb2.KVListResponse(keys=keys)
def KVExists(self, request, context=None) -> ray_client_pb2.KVExistsResponse:
try:
with disable_client_hook():
exists = ray.experimental.internal_kv._internal_kv_exists(
request.key, namespace=request.namespace
)
except Exception as e:
return_exception_in_context(e, context)
exists = False
return ray_client_pb2.KVExistsResponse(exists=exists)
def ListNamedActors(
self, request, context=None
) -> ray_client_pb2.ClientListNamedActorsResponse:
with disable_client_hook():
actors = ray.util.list_named_actors(all_namespaces=request.all_namespaces)
return ray_client_pb2.ClientListNamedActorsResponse(
actors_json=json.dumps(actors)
)
def ClusterInfo(self, request, context=None) -> ray_client_pb2.ClusterInfoResponse:
resp = ray_client_pb2.ClusterInfoResponse()
resp.type = request.type
if request.type == ray_client_pb2.ClusterInfoType.CLUSTER_RESOURCES:
with disable_client_hook():
resources = ray.cluster_resources()
# Normalize resources into floats
# (the function may return values that are ints)
float_resources = {k: float(v) for k, v in resources.items()}
resp.resource_table.CopyFrom(
ray_client_pb2.ClusterInfoResponse.ResourceTable(table=float_resources)
)
elif request.type == ray_client_pb2.ClusterInfoType.AVAILABLE_RESOURCES:
with disable_client_hook():
resources = ray.available_resources()
# Normalize resources into floats
# (the function may return values that are ints)
float_resources = {k: float(v) for k, v in resources.items()}
resp.resource_table.CopyFrom(
ray_client_pb2.ClusterInfoResponse.ResourceTable(table=float_resources)
)
elif request.type == ray_client_pb2.ClusterInfoType.RUNTIME_CONTEXT:
ctx = ray_client_pb2.ClusterInfoResponse.RuntimeContext()
with disable_client_hook():
rtc = ray.get_runtime_context()
ctx.job_id = ray._common.utils.hex_to_binary(rtc.get_job_id())
ctx.node_id = ray._common.utils.hex_to_binary(rtc.get_node_id())
ctx.worker_id = ray._common.utils.hex_to_binary(rtc.get_worker_id())
ctx.namespace = rtc.namespace
ctx.capture_client_tasks = (
rtc.should_capture_child_tasks_in_placement_group
)
ctx.gcs_address = rtc.gcs_address
ctx.runtime_env = rtc.get_runtime_env_string()
ctx.session_name = rtc.get_session_name()
resp.runtime_context.CopyFrom(ctx)
else:
with disable_client_hook():
resp.json = self._return_debug_cluster_info(request, context)
return resp
def _return_debug_cluster_info(self, request, context=None) -> str:
"""Handle ClusterInfo requests that only return a json blob."""
data = None
if request.type == ray_client_pb2.ClusterInfoType.NODES:
data = ray.nodes()
elif request.type == ray_client_pb2.ClusterInfoType.IS_INITIALIZED:
data = ray.is_initialized()
elif request.type == ray_client_pb2.ClusterInfoType.TIMELINE:
data = ray.timeline()
elif request.type == ray_client_pb2.ClusterInfoType.PING:
data = {}
elif request.type == ray_client_pb2.ClusterInfoType.DASHBOARD_URL:
data = {"dashboard_url": ray._private.worker.get_dashboard_url()}
else:
raise TypeError("Unsupported cluster info type")
return json.dumps(data)
def release(self, client_id: str, id: bytes) -> bool:
with self.state_lock:
if client_id in self.object_refs:
if id in self.object_refs[client_id]:
logger.debug(f"Releasing object {id.hex()} for {client_id}")
del self.object_refs[client_id][id]
return True
if client_id in self.actor_owners:
if id in self.actor_owners[client_id]:
logger.debug(f"Releasing actor {id.hex()} for {client_id}")
self.actor_owners[client_id].remove(id)
if self._can_remove_actor_ref(id):
logger.debug(f"Deleting reference to actor {id.hex()}")
del self.actor_refs[id]
return True
return False
def release_all(self, client_id):
with self.state_lock:
self._release_objects(client_id)
self._release_actors(client_id)
# NOTE: Try to actually dereference the object and actor refs.
# Otherwise dereferencing will happen later, which may run concurrently
# with ray.shutdown() and will crash the process. The crash is a bug
# that should be fixed eventually.
gc.collect()
def _can_remove_actor_ref(self, actor_id_bytes):
no_owner = not any(
actor_id_bytes in actor_list for actor_list in self.actor_owners.values()
)
return no_owner and actor_id_bytes not in self.named_actors
def _release_objects(self, client_id):
if client_id not in self.object_refs:
logger.debug(f"Releasing client with no references: {client_id}")
return
count = len(self.object_refs[client_id])
del self.object_refs[client_id]
if client_id in self.client_side_ref_map:
del self.client_side_ref_map[client_id]
if client_id in self.response_caches:
del self.response_caches[client_id]
logger.debug(f"Released all {count} objects for client {client_id}")
def _release_actors(self, client_id):
if client_id not in self.actor_owners:
logger.debug(f"Releasing client with no actors: {client_id}")
return
count = 0
actors_to_remove = self.actor_owners.pop(client_id)
for id_bytes in actors_to_remove:
count += 1
if self._can_remove_actor_ref(id_bytes):
logger.debug(f"Deleting reference to actor {id_bytes.hex()}")
del self.actor_refs[id_bytes]
logger.debug(f"Released all {count} actors for client: {client_id}")
@_use_response_cache
def Terminate(self, req, context=None):
if req.WhichOneof("terminate_type") == "task_object":
try:
object_ref = self.object_refs[req.client_id][req.task_object.id]
with disable_client_hook():
ray.cancel(
object_ref,
force=req.task_object.force,
recursive=req.task_object.recursive,
)
except Exception as e:
return_exception_in_context(e, context)
elif req.WhichOneof("terminate_type") == "actor":
try:
actor_ref = self.actor_refs[req.actor.id]
with disable_client_hook():
ray.kill(actor_ref, no_restart=req.actor.no_restart)
except Exception as e:
return_exception_in_context(e, context)
else:
raise RuntimeError(
"Client requested termination without providing a valid terminate_type"
)
return ray_client_pb2.TerminateResponse(ok=True)
def _async_get_object(
self,
request: ray_client_pb2.GetRequest,
client_id: str,
req_id: int,
result_queue: queue.Queue,
context=None,
) -> Optional[ray_client_pb2.GetResponse]:
"""Attempts to schedule a callback to push the GetResponse to the
main loop when the desired object is ready. If there is some failure
in scheduling, a GetResponse will be immediately returned.
"""
if len(request.ids) != 1:
raise ValueError(
f"Async get() must have exactly 1 Object ID. Actual: {request}"
)
rid = request.ids[0]
ref = self.object_refs[client_id].get(rid, None)
if not ref:
return ray_client_pb2.GetResponse(
valid=False,
error=cloudpickle.dumps(
ValueError(
f"ClientObjectRef with id {rid} not found for "
f"client {client_id}"
)
),
)
try:
logger.debug("async get: %s" % ref)
with disable_client_hook():
def send_get_response(result: Any) -> None:
"""Pushes GetResponses to the main DataPath loop to send
to the client. This is called when the object is ready
on the server side."""
try:
serialized = dumps_from_server(result, client_id, self)
total_size = len(serialized)
assert total_size > 0, "Serialized object cannot be zero bytes"
total_chunks = math.ceil(
total_size / OBJECT_TRANSFER_CHUNK_SIZE
)
for chunk_id in range(request.start_chunk_id, total_chunks):
start = chunk_id * OBJECT_TRANSFER_CHUNK_SIZE
end = min(
total_size, (chunk_id + 1) * OBJECT_TRANSFER_CHUNK_SIZE
)
get_resp = ray_client_pb2.GetResponse(
valid=True,
data=serialized[start:end],
chunk_id=chunk_id,
total_chunks=total_chunks,
total_size=total_size,
)
chunk_resp = ray_client_pb2.DataResponse(
get=get_resp, req_id=req_id
)
result_queue.put(chunk_resp)
except Exception as exc:
get_resp = ray_client_pb2.GetResponse(
valid=False, error=cloudpickle.dumps(exc)
)
resp = ray_client_pb2.DataResponse(get=get_resp, req_id=req_id)
result_queue.put(resp)
ref._on_completed(send_get_response)
return None
except Exception as e:
return ray_client_pb2.GetResponse(valid=False, error=cloudpickle.dumps(e))
def GetObject(self, request: ray_client_pb2.GetRequest, context):
metadata = dict(context.invocation_metadata())
client_id = metadata.get("client_id")
if client_id is None:
yield ray_client_pb2.GetResponse(
valid=False,
error=cloudpickle.dumps(
ValueError("client_id is not specified in request metadata")
),
)
else:
yield from self._get_object(request, client_id)
def _get_object(self, request: ray_client_pb2.GetRequest, client_id: str):
objectrefs = []
for rid in request.ids:
ref = self.object_refs[client_id].get(rid, None)
if ref:
objectrefs.append(ref)
else:
yield ray_client_pb2.GetResponse(
valid=False,
error=cloudpickle.dumps(
ValueError(
f"ClientObjectRef {rid} is not found for client {client_id}"
)
),
)
return
try:
logger.debug("get: %s" % objectrefs)
with disable_client_hook():
items = ray.get(objectrefs, timeout=request.timeout)
except Exception as e:
yield ray_client_pb2.GetResponse(valid=False, error=cloudpickle.dumps(e))
return
serialized = dumps_from_server(items, client_id, self)
total_size = len(serialized)
assert total_size > 0, "Serialized object cannot be zero bytes"
total_chunks = math.ceil(total_size / OBJECT_TRANSFER_CHUNK_SIZE)
for chunk_id in range(request.start_chunk_id, total_chunks):
start = chunk_id * OBJECT_TRANSFER_CHUNK_SIZE
end = min(total_size, (chunk_id + 1) * OBJECT_TRANSFER_CHUNK_SIZE)
yield ray_client_pb2.GetResponse(
valid=True,
data=serialized[start:end],
chunk_id=chunk_id,
total_chunks=total_chunks,
total_size=total_size,
)
def PutObject(
self, request: ray_client_pb2.PutRequest, context=None
) -> ray_client_pb2.PutResponse:
"""gRPC entrypoint for unary PutObject"""
return self._put_object(request.data, request.client_ref_id, "", context)
def _put_object(
self,
data: Union[bytes, bytearray],
client_ref_id: bytes,
client_id: str,
context: Optional[grpc.ServicerContext] = None,
) -> ray_client_pb2.PutResponse:
"""Put an object in the cluster with ray.put() via gRPC.
Args:
data: Pickled data. Can either be bytearray if this is called
from the dataservicer, or bytes if called from PutObject.
client_ref_id: The id associated with this object on the client.
client_id: The client who owns this data, for tracking when to
delete this reference.
context: gRPC context.
Returns:
A ``PutResponse`` containing the resulting object ref id, or an
error payload if the put failed.
"""
try:
obj = loads_from_client(data, self)
with disable_client_hook():
objectref = ray.put(obj)
except Exception as e:
logger.exception("Put failed:")
return ray_client_pb2.PutResponse(
id=b"", valid=False, error=cloudpickle.dumps(e)
)
self.object_refs[client_id][objectref.binary()] = objectref
if len(client_ref_id) > 0:
self.client_side_ref_map[client_id][client_ref_id] = objectref.binary()
logger.debug("put: %s" % objectref)
return ray_client_pb2.PutResponse(id=objectref.binary(), valid=True)
def WaitObject(self, request, context=None) -> ray_client_pb2.WaitResponse:
object_refs = []
for rid in request.object_ids:
if rid not in self.object_refs[request.client_id]:
raise Exception(
"Asking for a ref not associated with this client: %s" % str(rid)
)
object_refs.append(self.object_refs[request.client_id][rid])
num_returns = request.num_returns
timeout = request.timeout
try:
with disable_client_hook():
ready_object_refs, remaining_object_refs = ray.wait(
object_refs,
num_returns=num_returns,
timeout=timeout if timeout != -1 else None,
)
except Exception as e:
# TODO(ameer): improve exception messages.
logger.error(f"Exception {e}")
return ray_client_pb2.WaitResponse(valid=False)
logger.debug(
"wait: %s %s" % (str(ready_object_refs), str(remaining_object_refs))
)
ready_object_ids = [
ready_object_ref.binary() for ready_object_ref in ready_object_refs
]
remaining_object_ids = [
remaining_object_ref.binary()
for remaining_object_ref in remaining_object_refs
]
return ray_client_pb2.WaitResponse(
valid=True,
ready_object_ids=ready_object_ids,
remaining_object_ids=remaining_object_ids,
)
def Schedule(
self,
task: ray_client_pb2.ClientTask,
arglist: List[Any],
kwargs: Dict[str, Any],
context=None,
) -> ray_client_pb2.ClientTaskTicket:
logger.debug(
"schedule: %s %s"
% (task.name, ray_client_pb2.ClientTask.RemoteExecType.Name(task.type))
)
try:
with disable_client_hook():
if task.type == ray_client_pb2.ClientTask.FUNCTION:
result = self._schedule_function(task, arglist, kwargs, context)
elif task.type == ray_client_pb2.ClientTask.ACTOR:
result = self._schedule_actor(task, arglist, kwargs, context)
elif task.type == ray_client_pb2.ClientTask.METHOD:
result = self._schedule_method(task, arglist, kwargs, context)
elif task.type == ray_client_pb2.ClientTask.NAMED_ACTOR:
result = self._schedule_named_actor(task, context)
else:
raise NotImplementedError(
"Unimplemented Schedule task type: %s"
% ray_client_pb2.ClientTask.RemoteExecType.Name(task.type)
)
result.valid = True
return result
except Exception as e:
logger.debug("Caught schedule exception", exc_info=True)
return ray_client_pb2.ClientTaskTicket(
valid=False, error=cloudpickle.dumps(e)
)
def _schedule_method(
self,
task: ray_client_pb2.ClientTask,
arglist: List[Any],
kwargs: Dict[str, Any],
context=None,
) -> ray_client_pb2.ClientTaskTicket:
actor_handle = self.actor_refs.get(task.payload_id)
if actor_handle is None:
raise Exception("Can't run an actor the server doesn't have a handle for")
method = getattr(actor_handle, task.name)
opts = decode_options(task.options)
if opts is not None:
method = method.options(**opts)
output = method.remote(*arglist, **kwargs)
ids = self.unify_and_track_outputs(output, task.client_id)
return ray_client_pb2.ClientTaskTicket(return_ids=ids)
def _schedule_actor(
self,
task: ray_client_pb2.ClientTask,
arglist: List[Any],
kwargs: Dict[str, Any],
context=None,
) -> ray_client_pb2.ClientTaskTicket:
remote_class = self.lookup_or_register_actor(
task.payload_id, task.client_id, decode_options(task.baseline_options)
)
opts = decode_options(task.options)
if opts is not None:
remote_class = remote_class.options(**opts)
with current_server(self):
actor = remote_class.remote(*arglist, **kwargs)
self.actor_refs[actor._actor_id.binary()] = actor
self.actor_owners[task.client_id].add(actor._actor_id.binary())
return ray_client_pb2.ClientTaskTicket(return_ids=[actor._actor_id.binary()])
def _schedule_function(
self,
task: ray_client_pb2.ClientTask,
arglist: List[Any],
kwargs: Dict[str, Any],
context=None,
) -> ray_client_pb2.ClientTaskTicket:
remote_func = self.lookup_or_register_func(
task.payload_id, task.client_id, decode_options(task.baseline_options)
)
opts = decode_options(task.options)
if opts is not None:
remote_func = remote_func.options(**opts)
with current_server(self):
output = remote_func.remote(*arglist, **kwargs)
ids = self.unify_and_track_outputs(output, task.client_id)
return ray_client_pb2.ClientTaskTicket(return_ids=ids)
def _schedule_named_actor(
self, task: ray_client_pb2.ClientTask, context=None
) -> ray_client_pb2.ClientTaskTicket:
assert len(task.payload_id) == 0
# Convert empty string back to None.
actor = ray.get_actor(task.name, task.namespace or None)
bin_actor_id = actor._actor_id.binary()
if bin_actor_id not in self.actor_refs:
self.actor_refs[bin_actor_id] = actor
self.actor_owners[task.client_id].add(bin_actor_id)
self.named_actors.add(bin_actor_id)
return ray_client_pb2.ClientTaskTicket(return_ids=[actor._actor_id.binary()])
def lookup_or_register_func(
self, id: bytes, client_id: str, options: Optional[Dict]
) -> ray.remote_function.RemoteFunction:
with disable_client_hook():
if id not in self.function_refs:
funcref = self.object_refs[client_id][id]
func = ray.get(funcref)
if not inspect.isfunction(func):
raise Exception(
"Attempting to register function that isn't a function."
)
if options is None or len(options) == 0:
self.function_refs[id] = ray.remote(func)
else:
self.function_refs[id] = ray.remote(**options)(func)
return self.function_refs[id]
def lookup_or_register_actor(
self, id: bytes, client_id: str, options: Optional[Dict]
):
with disable_client_hook():
if id not in self.registered_actor_classes:
actor_class_ref = self.object_refs[client_id][id]
actor_class = ray.get(actor_class_ref)
if not inspect.isclass(actor_class):
raise Exception("Attempting to schedule actor that isn't a class.")
if options is None or len(options) == 0:
reg_class = ray.remote(actor_class)
else:
reg_class = ray.remote(**options)(actor_class)
self.registered_actor_classes[id] = reg_class
return self.registered_actor_classes[id]
def unify_and_track_outputs(self, output, client_id):
if output is None:
outputs = []
elif isinstance(output, list):
outputs = output
else:
outputs = [output]
for out in outputs:
if out.binary() in self.object_refs[client_id]:
logger.warning(f"Already saw object_ref {out}")
self.object_refs[client_id][out.binary()] = out
return [out.binary() for out in outputs]
def return_exception_in_context(err, context):
if context is not None:
context.set_details(encode_exception(err))
# Note: https://grpc.github.io/grpc/core/md_doc_statuscodes.html
# ABORTED used here since it should never be generated by the
# grpc lib -- this way we know the error was generated by ray logic
context.set_code(grpc.StatusCode.ABORTED)
def encode_exception(exception) -> str:
data = cloudpickle.dumps(exception)
return base64.standard_b64encode(data).decode()
def decode_options(options: ray_client_pb2.TaskOptions) -> Optional[Dict[str, Any]]:
if not options.pickled_options:
return None
opts = pickle.loads(options.pickled_options)
assert isinstance(opts, dict)
return opts
def serve(host: str, port: int, ray_connect_handler=None):
def default_connect_handler(
job_config: JobConfig = None, **ray_init_kwargs: Dict[str, Any]
):
with disable_client_hook():
if not ray.is_initialized():
return ray.init(job_config=job_config, **ray_init_kwargs)
from ray._private.grpc_utils import create_grpc_server_with_interceptors
ray_connect_handler = ray_connect_handler or default_connect_handler
server = create_grpc_server_with_interceptors(
max_workers=CLIENT_SERVER_MAX_THREADS,
thread_name_prefix="ray_client_server",
options=GRPC_OPTIONS,
asynchronous=False,
)
task_servicer = RayletServicer(ray_connect_handler)
data_servicer = DataServicer(task_servicer)
logs_servicer = LogstreamServicer()
ray_client_pb2_grpc.add_RayletDriverServicer_to_server(task_servicer, server)
ray_client_pb2_grpc.add_RayletDataStreamerServicer_to_server(data_servicer, server)
ray_client_pb2_grpc.add_RayletLogStreamerServicer_to_server(logs_servicer, server)
if not is_localhost(host):
add_port_to_grpc_server(server, build_address(get_localhost_ip(), port))
add_port_to_grpc_server(server, build_address(host, port))
current_handle = ClientServerHandle(
task_servicer=task_servicer,
data_servicer=data_servicer,
logs_servicer=logs_servicer,
grpc_server=server,
)
server.start()
return current_handle
def init_and_serve(host: str, port: int, *args, **kwargs):
with disable_client_hook():
# Disable client mode inside the worker's environment
info = ray.init(*args, **kwargs)
def ray_connect_handler(job_config=None, **ray_init_kwargs):
# Ray client will disconnect from ray when
# num_clients == 0.
if ray.is_initialized():
return info
else:
return ray.init(job_config=job_config, *args, **kwargs)
server_handle = serve(host, port, ray_connect_handler=ray_connect_handler)
return (server_handle, info)
def shutdown_with_server(server, _exiting_interpreter=False):
server.stop(1)
with disable_client_hook():
ray.shutdown(_exiting_interpreter=_exiting_interpreter)
def create_ray_handler(address, redis_password, redis_username=None):
def ray_connect_handler(job_config: JobConfig = None, **ray_init_kwargs):
if address:
if redis_password:
ray.init(
address=address,
_redis_username=redis_username,
_redis_password=redis_password,
job_config=job_config,
**ray_init_kwargs,
)
else:
ray.init(address=address, job_config=job_config, **ray_init_kwargs)
else:
ray.init(job_config=job_config, **ray_init_kwargs)
return ray_connect_handler
def try_create_gcs_client(address: Optional[str]) -> Optional[GcsClient]:
"""
Try to create a gcs client based on the command line args or by
autodetecting a running Ray cluster.
"""
address = canonicalize_bootstrap_address_or_die(address)
return GcsClient(address=address)
def main():
import argparse
parser = argparse.ArgumentParser()
parser.add_argument(
"--host",
type=str,
default=get_all_interfaces_ip(),
help="Host IP to bind to. Defaults to all interfaces (0.0.0.0/::).",
)
parser.add_argument("-p", "--port", type=int, default=10001, help="Port to bind to")
parser.add_argument(
"--mode",
type=str,
choices=["proxy", "legacy", "specific-server"],
default="proxy",
)
parser.add_argument(
"--address", required=False, type=str, help="Address to use to connect to Ray"
)
parser.add_argument(
"--redis-username",
required=False,
type=str,
help="username for connecting to Redis",
)
parser.add_argument(
"--runtime-env-agent-address",
required=False,
type=str,
default=None,
help="The port to use for connecting to the runtime_env_agent.",
)
parser.add_argument(
"--node-id",
required=False,
type=str,
default=None,
help="The hex ID of this node.",
)
args, _ = parser.parse_known_args()
redis_password = os.environ.get(ray_constants.RAY_REDIS_PASSWORD_ENV)
setup_logger(ray_constants.LOGGER_LEVEL, ray_constants.LOGGER_FORMAT)
ray_connect_handler = create_ray_handler(
args.address, redis_password, args.redis_username
)
hostport = build_address(args.host, args.port)
args_str = str(args)
logger.info(f"Starting Ray Client server on {hostport}, args {args_str}")
if args.mode == "proxy":
server = serve_proxier(
args.host,
args.port,
args.address,
redis_username=args.redis_username,
redis_password=redis_password,
runtime_env_agent_address=args.runtime_env_agent_address,
node_id=args.node_id,
)
else:
server = serve(args.host, args.port, ray_connect_handler)
try:
idle_checks_remaining = TIMEOUT_FOR_SPECIFIC_SERVER_S
while True:
health_report = {
"time": time.time(),
}
try:
if not ray.experimental.internal_kv._internal_kv_initialized():
gcs_client = try_create_gcs_client(args.address)
ray.experimental.internal_kv._initialize_internal_kv(gcs_client)
ray.experimental.internal_kv._internal_kv_put(
"ray_client_server",
json.dumps(health_report),
namespace=ray_constants.KV_NAMESPACE_HEALTHCHECK,
)
except Exception as e:
logger.error(
f"[{args.mode}] Failed to put health check on {args.address}"
)
logger.exception(e)
time.sleep(1)
if args.mode == "specific-server":
if server.data_servicer.num_clients > 0:
idle_checks_remaining = TIMEOUT_FOR_SPECIFIC_SERVER_S
else:
idle_checks_remaining -= 1
if idle_checks_remaining == 0:
raise KeyboardInterrupt()
if (
idle_checks_remaining % 5 == 0
and idle_checks_remaining != TIMEOUT_FOR_SPECIFIC_SERVER_S
):
logger.info(f"{idle_checks_remaining} idle checks before shutdown.")
except KeyboardInterrupt:
server.stop(0)
if __name__ == "__main__":
main()
python/ray/util/client/server/server_pickler.py
+124
View File
@@ -0,0 +1,124 @@
"""Implements the client side of the client/server pickling protocol.
These picklers are aware of the server internals and can find the
references held for the client within the server.
More discussion about the client/server pickling protocol can be found in:
ray/util/client/client_pickler.py
ServerPickler dumps ray objects from the server into the appropriate stubs.
ClientUnpickler loads stubs from the client and finds their associated handle
in the server instance.
"""
import io
from typing import TYPE_CHECKING, Any
import ray
import ray.cloudpickle as cloudpickle
from ray._private.client_mode_hook import disable_client_hook
from ray.util.client.client_pickler import PickleStub
from ray.util.client.server.server_stubs import (
ClientReferenceActor,
ClientReferenceFunction,
)
if TYPE_CHECKING:
from ray.util.client.server.server import RayletServicer
import pickle # noqa: F401
class ServerPickler(cloudpickle.CloudPickler):
def __init__(self, client_id: str, server: "RayletServicer", *args, **kwargs):
super().__init__(*args, **kwargs)
self.client_id = client_id
self.server = server
def persistent_id(self, obj):
if isinstance(obj, ray.ObjectRef):
obj_id = obj.binary()
if obj_id not in self.server.object_refs[self.client_id]:
# We're passing back a reference, probably inside a reference.
# Let's hold onto it.
self.server.object_refs[self.client_id][obj_id] = obj
return PickleStub(
type="Object",
client_id=self.client_id,
ref_id=obj_id,
name=None,
baseline_options=None,
)
elif isinstance(obj, ray.actor.ActorHandle):
actor_id = obj._actor_id.binary()
if actor_id not in self.server.actor_refs:
# We're passing back a handle, probably inside a reference.
self.server.actor_refs[actor_id] = obj
if actor_id not in self.server.actor_owners[self.client_id]:
self.server.actor_owners[self.client_id].add(actor_id)
return PickleStub(
type="Actor",
client_id=self.client_id,
ref_id=obj._actor_id.binary(),
name=None,
baseline_options=None,
)
return None
class ClientUnpickler(pickle.Unpickler):
def __init__(self, server, *args, **kwargs):
super().__init__(*args, **kwargs)
self.server = server
def persistent_load(self, pid):
assert isinstance(pid, PickleStub)
if pid.type == "Ray":
return ray
elif pid.type == "Object":
return self.server.object_refs[pid.client_id][pid.ref_id]
elif pid.type == "Actor":
return self.server.actor_refs[pid.ref_id]
elif pid.type == "RemoteFuncSelfReference":
return ClientReferenceFunction(pid.client_id, pid.ref_id)
elif pid.type == "RemoteFunc":
return self.server.lookup_or_register_func(
pid.ref_id, pid.client_id, pid.baseline_options
)
elif pid.type == "RemoteActorSelfReference":
return ClientReferenceActor(pid.client_id, pid.ref_id)
elif pid.type == "RemoteActor":
return self.server.lookup_or_register_actor(
pid.ref_id, pid.client_id, pid.baseline_options
)
elif pid.type == "RemoteMethod":
actor = self.server.actor_refs[pid.ref_id]
return getattr(actor, pid.name)
else:
raise NotImplementedError("Uncovered client data type")
def dumps_from_server(
obj: Any, client_id: str, server_instance: "RayletServicer", protocol=None
) -> bytes:
with io.BytesIO() as file:
sp = ServerPickler(client_id, server_instance, file, protocol=protocol)
sp.dump(obj)
return file.getvalue()
def loads_from_client(
data: bytes,
server_instance: "RayletServicer",
*,
fix_imports=True,
encoding="ASCII",
errors="strict"
) -> Any:
with disable_client_hook():
if isinstance(data, str):
raise TypeError("Can't load pickle from unicode string")
file = io.BytesIO(data)
return ClientUnpickler(
server_instance, file, fix_imports=fix_imports, encoding=encoding
).load()
python/ray/util/client/server/server_stubs.py
+66
View File
@@ -0,0 +1,66 @@
from abc import ABC, abstractmethod
from contextlib import contextmanager
_current_server = None
@contextmanager
def current_server(r):
global _current_server
remote = _current_server
_current_server = r
try:
yield
finally:
_current_server = remote
class ClientReferenceSentinel(ABC):
def __init__(self, client_id, id):
self.client_id = client_id
self.id = id
def __reduce__(self):
remote_obj = self.get_remote_obj()
if remote_obj is None:
return (self.__class__, (self.client_id, self.id))
return (identity, (remote_obj,))
@abstractmethod
def get_remote_obj(self):
pass
def get_real_ref_from_server(self):
global _current_server
if _current_server is None:
return None
client_map = _current_server.client_side_ref_map.get(self.client_id, None)
if client_map is None:
return None
return client_map.get(self.id, None)
class ClientReferenceActor(ClientReferenceSentinel):
def get_remote_obj(self):
global _current_server
real_ref_id = self.get_real_ref_from_server()
if real_ref_id is None:
return None
return _current_server.lookup_or_register_actor(
real_ref_id, self.client_id, None
)
class ClientReferenceFunction(ClientReferenceSentinel):
def get_remote_obj(self):
global _current_server
real_ref_id = self.get_real_ref_from_server()
if real_ref_id is None:
return None
return _current_server.lookup_or_register_func(
real_ref_id, self.client_id, None
)
def identity(x):
return x
python/ray/util/client/worker.py
+1005
View File
File diff suppressed because it is too large Load Diff
python/ray/util/client_connect.py
+81
View File
@@ -0,0 +1,81 @@
import logging
from typing import Any, Dict, List, Optional, Tuple
from ray._private.client_mode_hook import (
_explicitly_enable_client_mode,
_set_client_hook_status,
)
from ray._private.utils import get_ray_doc_version
from ray.job_config import JobConfig
from ray.util.annotations import Deprecated
from ray.util.client import ray
logger = logging.getLogger(__name__)
@Deprecated(
message="Use ray.init(ray://<head_node_ip_address>:<ray_client_server_port>) "
"instead. See detailed usage at {}.".format(
f"https://docs.ray.io/en/{get_ray_doc_version()}/ray-core/package-ref.html#ray-init" # noqa: E501
)
)
def connect(
conn_str: str,
secure: bool = False,
metadata: List[Tuple[str, str]] = None,
connection_retries: int = 3,
job_config: JobConfig = None,
namespace: str = None,
*,
ignore_version: bool = False,
_credentials: Optional["grpc.ChannelCredentials"] = None, # noqa: F821
ray_init_kwargs: Optional[Dict[str, Any]] = None,
) -> Dict[str, Any]:
# Check if Ray is already connected (any type of connection)
if ray.is_connected():
ignore_reinit_error = (
ray_init_kwargs.get("ignore_reinit_error", False)
if ray_init_kwargs is not None
else False
)
if ignore_reinit_error:
logger.info(
"Calling ray.init() again after it has already been called. "
"Reusing the existing Ray client connection."
)
return ray.get_context().client_worker.connection_info()
raise RuntimeError(
"Ray Client is already connected. Maybe you called "
'ray.init("ray://<address>") twice by accident?'
)
# Enable the same hooks that RAY_CLIENT_MODE does, as calling
# ray.init("ray://<address>") is specifically for using client mode.
_set_client_hook_status(True)
_explicitly_enable_client_mode()
# TODO(barakmich): https://github.com/ray-project/ray/issues/13274
# for supporting things like cert_path, ca_path, etc and creating
# the correct metadata
conn = ray.connect(
conn_str,
job_config=job_config,
secure=secure,
metadata=metadata,
connection_retries=connection_retries,
namespace=namespace,
ignore_version=ignore_version,
_credentials=_credentials,
ray_init_kwargs=ray_init_kwargs,
)
return conn
@Deprecated(
message="Use ray.shutdown() instead. See detailed usage at {}.".format(
f"https://docs.ray.io/en/{get_ray_doc_version()}/ray-core/package-ref.html#ray-shutdown" # noqa: E501
)
)
def disconnect():
"""Disconnects from server; is idempotent."""
return ray.disconnect()
python/ray/util/collective/__init__.py
+59
View File
@@ -0,0 +1,59 @@
from ray.util.collective.backend_registry import (
register_collective_backend,
)
from ray.util.collective.collective import (
allgather,
allgather_multigpu,
allreduce,
allreduce_multigpu,
barrier,
broadcast,
broadcast_multigpu,
create_collective_group,
destroy_collective_group,
get_collective_group_size,
get_group_handle,
get_rank,
gloo_available,
init_collective_group,
is_backend_available,
is_group_initialized,
nccl_available,
recv,
recv_multigpu,
reduce,
reduce_multigpu,
reducescatter,
reducescatter_multigpu,
send,
send_multigpu,
)
__all__ = [
"nccl_available",
"gloo_available",
"is_backend_available",
"is_group_initialized",
"init_collective_group",
"destroy_collective_group",
"create_collective_group",
"get_rank",
"get_collective_group_size",
"allreduce",
"allreduce_multigpu",
"barrier",
"reduce",
"reduce_multigpu",
"broadcast",
"broadcast_multigpu",
"allgather",
"allgather_multigpu",
"reducescatter",
"reducescatter_multigpu",
"send",
"send_multigpu",
"recv",
"recv_multigpu",
"get_group_handle",
"register_collective_backend",
]
python/ray/util/collective/backend_registry.py
+122
View File
@@ -0,0 +1,122 @@
from typing import Dict, Type
from .collective_group.base_collective_group import BaseGroup
from ray.util.annotations import PublicAPI
class BackendRegistry:
_instance = None
_map: Dict[str, Type[BaseGroup]]
def __new__(cls):
if cls._instance is None:
cls._instance = super(BackendRegistry, cls).__new__(cls)
cls._instance._map = {}
return cls._instance
def put(self, name: str, group_cls: Type[BaseGroup]) -> None:
if not issubclass(group_cls, BaseGroup):
raise TypeError(f"{group_cls} is not a subclass of BaseGroup")
if name.upper() in self._map:
raise ValueError(f"Backend {name.upper()} already registered")
self._map[name.upper()] = group_cls
def get(self, name: str) -> Type[BaseGroup]:
name = name.upper()
if name not in self._map:
raise ValueError(f"Backend {name} not registered")
return self._map[name]
def is_registered(self, name: str) -> bool:
"""Check if a backend is registered (regardless of availability)."""
return name.upper() in self._map
def check(self, name: str) -> bool:
"""Check if a backend is both registered and available."""
try:
cls = self.get(name)
return cls.check_backend_availability()
except (ValueError, AttributeError):
return False
_global_registry = BackendRegistry()
@PublicAPI(stability="alpha")
def register_collective_backend(name: str, group_cls: Type[BaseGroup]):
"""Register a custom collective backend with Ray.
This function registers a custom backend class that can be used for
collective operations. The backend must be a subclass of
:class:`~ray.util.collective.collective_group.base_collective_group.BaseGroup`
and implement all required collective operations.
Important: The backend must be registered on both the driver and all
actors before creating collective groups. This is because each process
(driver and each actor) needs to know about your backend class to
instantiate it.
Args:
name: The name of the backend (e.g., "MY_BACKEND"). This will be
automatically added to the Backend enum as Backend.MY_BACKEND.
group_cls: The backend class, which must be a subclass of
:class:`~ray.util.collective.collective_group.base_collective_group.BaseGroup`.
Example:
>>> import ray
>>> from ray.util.collective import create_collective_group, init_collective_group
>>> from ray.util.collective.backend_registry import register_collective_backend
>>> from ray.util.collective.collective_group.base_collective_group import BaseGroup
>>>
>>> class MyCustomBackend(BaseGroup):
... def __init__(self, world_size, rank, group_name):
... super().__init__(world_size, rank, group_name)
... @classmethod
... def backend(cls):
... return "MY_BACKEND"
... @classmethod
... def check_backend_availability(cls) -> bool:
... return True
... def allreduce(self, tensor, allreduce_options=None):
... pass
... def broadcast(self, tensor, broadcast_options=None):
... pass
... def barrier(self, barrier_options=None):
... pass
>>>
>>> # Register on the driver
>>> register_collective_backend("MY_BACKEND", MyCustomBackend)
>>>
>>> ray.init()
>>>
>>> @ray.remote
... class Worker:
... def __init__(self, rank):
... self.rank = rank
... def setup(self, world_size):
... # IMPORTANT: Register on each worker too
... register_collective_backend("MY_BACKEND", MyCustomBackend)
... init_collective_group(
... world_size=world_size,
... rank=self.rank,
... backend="MY_BACKEND",
... group_name="default",
... )
>>>
>>> actors = [Worker.remote(rank=i) for i in range(2)]
>>> create_collective_group(
... actors=actors,
... world_size=2,
... ranks=[0, 1],
... backend="MY_BACKEND",
... group_name="default",
... )
>>> ray.get([a.setup.remote(2) for a in actors])
"""
_global_registry.put(name, group_cls)
from . import types
upper_name = name.upper()
if not hasattr(types.Backend, upper_name):
setattr(types.Backend, upper_name, upper_name)
python/ray/util/collective/collective.py
+816
View File
@@ -0,0 +1,816 @@
"""APIs exposed under the namespace ray.util.collective."""
import logging
import os
import socket
import threading
from typing import Any, List, Tuple
import numpy as np
import ray
from . import types
from ray._common.network_utils import find_free_port, is_ipv6
from ray.util.collective.backend_registry import (
_global_registry,
register_collective_backend,
)
logger = logging.getLogger(__name__)
try:
from ray.util.collective.collective_group.nccl_collective_group import NCCLGroup
register_collective_backend("NCCL", NCCLGroup)
except ImportError:
pass
try:
from ray.util.collective.collective_group.torch_gloo_collective_group import (
TorchGLOOGroup,
)
register_collective_backend("GLOO", TorchGLOOGroup)
except ImportError:
pass
def nccl_available():
return is_backend_available("NCCL")
def gloo_available():
return is_backend_available("GLOO")
def is_backend_available(backend: str) -> bool:
"""Check if a collective backend is available.
Args:
backend: The name of the backend to check (e.g., "NCCL", "GLOO").
Returns:
True if the backend is available, False otherwise.
"""
return _global_registry.check(backend)
def get_address_and_port() -> Tuple[str, int]:
"""Returns the IP address and a free port on this node."""
addr = ray.util.get_node_ip_address()
port = find_free_port(socket.AF_INET6 if is_ipv6(addr) else socket.AF_INET)
return addr, port
class GroupManager(object):
"""Use this class to manage the collective groups we created so far.
Each process will have an instance of `GroupManager`. Each process
could belong to multiple collective groups. The membership information
and other metadata are stored in the global `_group_mgr` object.
"""
def __init__(self):
self._name_group_map = {}
self._registry = _global_registry
def create_collective_group(
self, backend, world_size, rank, group_name, gloo_timeout=None
):
"""The entry to create new collective groups in the manager.
Put the registration and the group information into the manager
metadata as well.
"""
backend = backend.upper()
backend_cls = self._registry.get(backend)
if not backend_cls.check_backend_availability():
raise RuntimeError(
f"Backend {backend} is not available. Please check the installation."
)
if backend == "GLOO":
logger.debug(
"Creating torch.distributed GLOO group: '{}'...".format(group_name)
)
g = backend_cls(world_size, rank, group_name, gloo_timeout)
else:
g = backend_cls(world_size, rank, group_name)
self._name_group_map[group_name] = g
return self._name_group_map[group_name]
def is_group_exist(self, group_name):
return group_name in self._name_group_map
def get_group_by_name(self, group_name):
"""Get the collective group handle by its name."""
if not self.is_group_exist(group_name):
logger.warning("The group '{}' is not initialized.".format(group_name))
return None
return self._name_group_map[group_name]
def destroy_collective_group(self, group_name):
"""Group destructor."""
if not self.is_group_exist(group_name):
logger.warning("The group '{}' does not exist.".format(group_name))
return
# release the collective group resource
g = self._name_group_map[group_name]
# clean up the dicts
del self._name_group_map[group_name]
# Release the communicator resources
g.destroy_group()
# Release the detached actors spawned by `create_collective_group()`
name = "info_" + group_name
try:
store = ray.get_actor(name)
ray.kill(store)
except ValueError:
pass
_group_mgr = GroupManager()
# This lock is used to make external calls to the _group_mgr thread-safe.
_group_mgr_lock = threading.Lock()
def is_group_initialized(group_name):
"""Check if the group is initialized in this process by the group name."""
global _group_mgr
global _group_mgr_lock
with _group_mgr_lock:
return _group_mgr.is_group_exist(group_name)
def init_collective_group(
world_size: int,
rank: int,
backend: types.Backend = types.Backend.NCCL,
group_name: str = "default",
gloo_timeout: int = 30000,
):
"""Initialize a collective group inside an actor process.
Args:
world_size: the total number of processes in the group.
rank: the rank of the current process.
backend: the CCL backend to use, NCCL or GLOO.
group_name: the name of the collective group.
gloo_timeout: timeout in milliseconds for GLOO operations.
"""
_check_inside_actor()
global _group_mgr
global _group_mgr_lock
# TODO(Hao): implement a group auto-counter.
if not group_name:
raise ValueError("group_name '{}' needs to be a string.".format(group_name))
with _group_mgr_lock:
if _group_mgr.is_group_exist(group_name):
raise RuntimeError("Trying to initialize a group a second time.")
assert world_size > 0
assert rank >= 0
assert rank < world_size
_group_mgr.create_collective_group(
backend, world_size, rank, group_name, gloo_timeout
)
def create_collective_group(
actors: List[Any],
world_size: int,
ranks: List[int],
backend: types.Backend = types.Backend.NCCL,
group_name: str = "default",
gloo_timeout: int = 30000,
):
"""Declare a list of actors as a collective group.
Note: This function should be called in a driver process.
Args:
actors: a list of actors to be set in a collective group.
world_size: the total number of processes in the group.
ranks: the rank of each actor.
backend: the CCL backend to use, NCCL or GLOO.
group_name: the name of the collective group.
gloo_timeout: timeout in milliseconds for GLOO operations.
"""
name = "info_" + group_name
try:
ray.get_actor(name)
raise RuntimeError("Trying to initialize a group twice.")
except ValueError:
pass
if len(ranks) != len(actors):
raise RuntimeError(
"Each actor should correspond to one rank. Got '{}' "
"ranks but '{}' actors".format(len(ranks), len(actors))
)
if set(ranks) != set(range(len(ranks))):
raise RuntimeError(
"Ranks must be a permutation from 0 to '{}'. Got '{}'.".format(
len(ranks), "".join([str(r) for r in ranks])
)
)
if world_size <= 0:
raise RuntimeError(
"World size must be greater than zero. Got '{}'.".format(world_size)
)
if not all(ranks) >= 0:
raise RuntimeError("Ranks must be non-negative.")
if not all(ranks) < world_size:
raise RuntimeError("Ranks cannot be greater than world_size.")
# Check if backend is registered and available
backend_upper = backend.upper()
if not _global_registry.is_registered(backend_upper):
raise RuntimeError(
f"Backend {backend_upper} is not registered. "
f"Please register it using register_collective_backend('{backend_upper}', YourBackendClass)."
)
if not _global_registry.check(backend_upper):
raise RuntimeError(
f"Backend {backend_upper} is registered but not available. "
f"Please check the installation requirements for this backend."
)
# avoid a circular dependency
from ray.util.collective.util import Info
# store the information into a NamedActor that can be accessed later.
name = "info_" + group_name
actors_id = [a._ray_actor_id for a in actors]
# TODO (Dacheng): how do we recycle this name actor?
info = Info.options(name=name, lifetime="detached").remote()
ray.get([info.set_info.remote(actors_id, world_size, ranks, backend, gloo_timeout)])
# TODO (we need a declarative destroy() API here.)
def destroy_collective_group(group_name: str = "default") -> None:
"""Destroy a collective group given its group name."""
_check_inside_actor()
global _group_mgr
global _group_mgr_lock
with _group_mgr_lock:
_group_mgr.destroy_collective_group(group_name)
def get_rank(group_name: str = "default") -> int:
"""Return the rank of this process in the given group.
Args:
group_name: the name of the group to query
Returns:
the rank of this process in the named group,
-1 if the group does not exist or the process does
not belong to the group.
"""
_check_inside_actor()
global _group_mgr
global _group_mgr_lock
with _group_mgr_lock:
if not _group_mgr.is_group_exist(group_name):
return -1
g = _group_mgr.get_group_by_name(group_name)
return g.rank
def get_collective_group_size(group_name: str = "default") -> int:
"""Return the size of the collective group with the given name.
Args:
group_name: the name of the group to query
Returns:
The world size of the collective group, -1 if the group does
not exist or the process does not belong to the group.
"""
_check_inside_actor()
global _group_mgr
global _group_mgr_lock
with _group_mgr_lock:
if not _group_mgr.is_group_exist(group_name):
return -1
g = _group_mgr.get_group_by_name(group_name)
return g.world_size
def allreduce(
tensor: Any,
group_name: str = "default",
op: types.ReduceOp = types.ReduceOp.SUM,
):
"""Collective allreduce the tensor across the group.
Args:
tensor: the tensor to be all-reduced on this process.
group_name: the collective group name to perform allreduce.
op: The reduce operation.
"""
_check_single_tensor_input(tensor)
g = get_group_handle(group_name)
opts = types.AllReduceOptions
opts.reduceOp = op
g.allreduce([tensor], opts)
def allreduce_multigpu(
tensor_list: list,
group_name: str = "default",
op: types.ReduceOp = types.ReduceOp.SUM,
):
"""Collective allreduce a list of tensors across the group.
Args:
tensor_list: list of tensors to be allreduced, each on a GPU.
group_name: the collective group name to perform allreduce.
op: The reduce operation.
"""
if not types.cupy_available():
raise RuntimeError("Multigpu calls requires NCCL and Cupy.")
_check_tensor_list_input(tensor_list)
g = get_group_handle(group_name)
opts = types.AllReduceOptions
opts.reduceOp = op
g.allreduce(tensor_list, opts)
def barrier(group_name: str = "default"):
"""Barrier all processes in the collective group.
Args:
group_name: the name of the group to barrier.
"""
g = get_group_handle(group_name)
g.barrier()
def reduce(
tensor: Any,
dst_rank: int = 0,
group_name: str = "default",
op: types.ReduceOp = types.ReduceOp.SUM,
):
"""Reduce the tensor across the group to the destination rank.
Args:
tensor: the tensor to be reduced on this process.
dst_rank: the rank of the destination process.
group_name: the collective group name to perform reduce.
op: The reduce operation.
"""
_check_single_tensor_input(tensor)
g = get_group_handle(group_name)
# check dst rank
_check_rank_valid(g, dst_rank)
opts = types.ReduceOptions()
opts.reduceOp = op
opts.root_rank = dst_rank
opts.root_tensor = 0
g.reduce([tensor], opts)
def reduce_multigpu(
tensor_list: list,
dst_rank: int = 0,
dst_tensor: int = 0,
group_name: str = "default",
op: types.ReduceOp = types.ReduceOp.SUM,
):
"""Reduce the tensor across the group to the destination rank
and destination tensor.
Args:
tensor_list: the list of tensors to be reduced on this process;
each tensor located on a GPU.
dst_rank: the rank of the destination process.
dst_tensor: the index of GPU at the destination.
group_name: the collective group name to perform reduce.
op: The reduce operation.
"""
if not types.cupy_available():
raise RuntimeError("Multigpu calls requires NCCL and Cupy.")
_check_tensor_list_input(tensor_list)
g = get_group_handle(group_name)
# check dst rank
_check_rank_valid(g, dst_rank)
_check_root_tensor_valid(len(tensor_list), dst_tensor)
opts = types.ReduceOptions()
opts.reduceOp = op
opts.root_rank = dst_rank
opts.root_tensor = dst_tensor
g.reduce(tensor_list, opts)
def broadcast(tensor: Any, src_rank: int = 0, group_name: str = "default"):
"""Broadcast the tensor from a source process to all others.
Args:
tensor: the tensor to be broadcasted (src) or received (destination).
src_rank: the rank of the source process.
group_name: the collective group name to perform broadcast.
"""
_check_single_tensor_input(tensor)
g = get_group_handle(group_name)
# check src rank
_check_rank_valid(g, src_rank)
opts = types.BroadcastOptions()
opts.root_rank = src_rank
opts.root_tensor = 0
g.broadcast([tensor], opts)
def broadcast_multigpu(
tensor_list: list,
src_rank: int = 0,
src_tensor: int = 0,
group_name: str = "default",
):
"""Broadcast the tensor from a source GPU to all other GPUs.
Args:
tensor_list: the tensors to broadcast (src) or receive (dst).
src_rank: the rank of the source process.
src_tensor: the index of the source GPU on the source process.
group_name: the collective group name to perform broadcast.
"""
if not types.cupy_available():
raise RuntimeError("Multigpu calls requires NCCL and Cupy.")
_check_tensor_list_input(tensor_list)
g = get_group_handle(group_name)
# check src rank
_check_rank_valid(g, src_rank)
_check_root_tensor_valid(len(tensor_list), src_tensor)
opts = types.BroadcastOptions()
opts.root_rank = src_rank
opts.root_tensor = src_tensor
g.broadcast(tensor_list, opts)
def allgather(tensor_list: list, tensor: Any, group_name: str = "default"):
"""Allgather tensors from each process of the group into a list.
Args:
tensor_list: the results, stored as a list of tensors.
tensor: the tensor (to be gathered) in the current process
group_name: the name of the collective group.
"""
_check_single_tensor_input(tensor)
_check_tensor_list_input(tensor_list)
g = get_group_handle(group_name)
if len(tensor_list) != g.world_size:
# Typically CLL lib requires len(tensor_list) >= world_size;
# Here we make it more strict: len(tensor_list) == world_size.
raise RuntimeError(
"The length of the tensor list operands to allgather "
"must be equal to world_size."
)
opts = types.AllGatherOptions()
g.allgather([tensor_list], [tensor], opts)
def allgather_multigpu(
output_tensor_lists: list, input_tensor_list: list, group_name: str = "default"
):
"""Allgather tensors from each gpus of the group into lists.
Args:
output_tensor_lists: gathered results, with shape
must be num_gpus * world_size * shape(tensor).
input_tensor_list: a list of tensors, with shape
num_gpus * shape(tensor).
group_name: the name of the collective group.
"""
if not types.cupy_available():
raise RuntimeError("Multigpu calls requires NCCL and Cupy.")
_check_tensor_lists_input(output_tensor_lists)
_check_tensor_list_input(input_tensor_list)
g = get_group_handle(group_name)
opts = types.AllGatherOptions()
g.allgather(output_tensor_lists, input_tensor_list, opts)
def reducescatter(
tensor: Any,
tensor_list: list,
group_name: str = "default",
op: types.ReduceOp = types.ReduceOp.SUM,
):
"""Reducescatter a list of tensors across the group.
Reduce the list of the tensors across each process in the group, then
scatter the reduced list of tensors -- one tensor for each process.
Args:
tensor: the resulted tensor on this process.
tensor_list: The list of tensors to be reduced and scattered.
group_name: the name of the collective group.
op: The reduce operation.
"""
_check_single_tensor_input(tensor)
_check_tensor_list_input(tensor_list)
g = get_group_handle(group_name)
opts = types.ReduceScatterOptions()
opts.reduceOp = op
if len(tensor_list) != g.world_size:
raise RuntimeError(
"The length of the tensor list operands to reducescatter "
"must not be equal to world_size."
)
g.reducescatter([tensor], [tensor_list], opts)
def reducescatter_multigpu(
output_tensor_list: list,
input_tensor_lists: list,
group_name: str = "default",
op: types.ReduceOp = types.ReduceOp.SUM,
):
"""Reducescatter a list of tensors across all GPUs.
Args:
output_tensor_list: the resulted list of tensors, with
shape: num_gpus * shape(tensor).
input_tensor_lists: the original tensors, with shape:
num_gpus * world_size * shape(tensor).
group_name: the name of the collective group.
op: The reduce operation.
"""
if not types.cupy_available():
raise RuntimeError("Multigpu calls requires NCCL and Cupy.")
_check_tensor_lists_input(input_tensor_lists)
_check_tensor_list_input(output_tensor_list)
g = get_group_handle(group_name)
opts = types.ReduceScatterOptions()
opts.reduceOp = op
g.reducescatter(output_tensor_list, input_tensor_lists, opts)
def send(tensor: Any, dst_rank: int, group_name: str = "default"):
"""Send a tensor to a remote process synchronously.
Args:
tensor: the tensor to send.
dst_rank: the rank of the destination process.
group_name: the name of the collective group.
"""
_check_single_tensor_input(tensor)
g = get_group_handle(group_name)
_check_rank_valid(g, dst_rank)
if dst_rank == g.rank:
raise RuntimeError("The destination rank '{}' is self.".format(dst_rank))
opts = types.SendOptions()
opts.dst_rank = dst_rank
g.send([tensor], opts)
def send_multigpu(
tensor: Any,
dst_rank: int,
dst_gpu_index: int,
group_name: str = "default",
n_elements: int = 0,
):
"""Send a tensor to a remote GPU synchronously.
The function assumes each process owns >1 GPUs, and the sender
process and receiver process has equal number of GPUs.
Args:
tensor: the tensor to send, located on a GPU.
dst_rank: the rank of the destination process.
dst_gpu_index: the destination gpu index.
group_name: the name of the collective group.
n_elements: if specified, send the next n elements
from the starting address of tensor.
"""
if not types.cupy_available():
raise RuntimeError("send_multigpu call requires NCCL.")
_check_single_tensor_input(tensor)
g = get_group_handle(group_name)
_check_rank_valid(g, dst_rank)
if dst_rank == g.rank:
raise RuntimeError(
"The dst_rank '{}' is self. Considering "
"doing GPU to GPU memcpy instead?".format(dst_rank)
)
if n_elements < 0:
raise RuntimeError("The n_elements '{}' should >= 0.".format(n_elements))
opts = types.SendOptions()
opts.dst_rank = dst_rank
opts.dst_gpu_index = dst_gpu_index
opts.n_elements = n_elements
g.send([tensor], opts)
def recv(tensor: Any, src_rank: int, group_name: str = "default"):
"""Receive a tensor from a remote process synchronously.
Args:
tensor: the received tensor.
src_rank: the rank of the source process.
group_name: the name of the collective group.
"""
_check_single_tensor_input(tensor)
g = get_group_handle(group_name)
_check_rank_valid(g, src_rank)
if src_rank == g.rank:
raise RuntimeError("The destination rank '{}' is self.".format(src_rank))
opts = types.RecvOptions()
opts.src_rank = src_rank
g.recv([tensor], opts)
def recv_multigpu(
tensor: Any,
src_rank: int,
src_gpu_index: int,
group_name: str = "default",
n_elements: int = 0,
):
"""Receive a tensor from a remote GPU synchronously.
The function asssume each process owns >1 GPUs, and the sender
process and receiver process has equal nubmer of GPUs.
Args:
tensor: The received tensor, located on a GPU.
src_rank: The rank of the source process.
src_gpu_index: The index of the source GPU on the src process.
group_name: The name of the collective group.
n_elements: if specified, receive the next n elements
from the starting address of tensor.
"""
if not types.cupy_available():
raise RuntimeError("recv_multigpu call requires NCCL.")
_check_single_tensor_input(tensor)
g = get_group_handle(group_name)
_check_rank_valid(g, src_rank)
if src_rank == g.rank:
raise RuntimeError(
"The dst_rank '{}' is self. Considering "
"doing GPU to GPU memcpy instead?".format(src_rank)
)
if n_elements < 0:
raise RuntimeError("The n_elements '{}' should be >= 0.".format(n_elements))
opts = types.RecvOptions()
opts.src_rank = src_rank
opts.src_gpu_index = src_gpu_index
opts.n_elements = n_elements
g.recv([tensor], opts)
def synchronize(gpu_id: int):
"""Synchronize the current process to a give device.
Args:
gpu_id: the GPU device id to synchronize.
"""
if not types.cupy_available():
raise RuntimeError("synchronize call requires CUDA and NCCL.")
import cupy as cp
cp.cuda.Device(gpu_id).synchronize()
def get_group_handle(group_name: str = "default"):
"""Check if the group is initialized and return the group handle.
Args:
group_name: the name of the collective group.
Returns:
The collective group handle.
"""
_check_inside_actor()
global _group_mgr
global _group_mgr_lock
with _group_mgr_lock:
if not _group_mgr.is_group_exist(group_name):
# try loading from remote info store
try:
# if the information is stored in an Info object,
# get and create the group.
name = "info_" + group_name
mgr = ray.get_actor(name=name)
ids, world_size, rank, backend, gloo_timeout = ray.get(
mgr.get_info.remote()
)
worker = ray._private.worker.global_worker
id_ = worker.core_worker.get_actor_id()
r = rank[ids.index(id_)]
_group_mgr.create_collective_group(
backend, world_size, r, group_name, gloo_timeout
)
except ValueError as exc:
# check if this group is initialized using options()
if (
"collective_group_name" in os.environ
and os.environ["collective_group_name"] == group_name
):
rank = int(os.environ["collective_rank"])
world_size = int(os.environ["collective_world_size"])
backend = os.environ["collective_backend"]
gloo_timeout = int(os.getenv("collective_gloo_timeout", 30000))
_group_mgr.create_collective_group(
backend, world_size, rank, group_name, gloo_timeout
)
else:
raise RuntimeError(
"The collective group '{}' is not "
"initialized in the process.".format(group_name)
) from exc
g = _group_mgr.get_group_by_name(group_name)
return g
def _check_single_tensor_input(tensor):
"""Check if the tensor is with a supported type."""
if isinstance(tensor, np.ndarray):
return
if types.cupy_available():
if isinstance(tensor, types.cp.ndarray):
return
if types.torch_available():
if isinstance(tensor, types.th.Tensor):
return
raise RuntimeError(
"Unrecognized tensor type '{}'. Supported types are: "
"np.ndarray, torch.Tensor, cupy.ndarray.".format(type(tensor))
)
def _check_inside_actor():
"""Check if currently it is inside a Ray actor/task."""
worker = ray._private.worker.global_worker
if worker.mode == ray.WORKER_MODE:
return
else:
raise RuntimeError(
"The collective APIs shall be only used inside a Ray actor or task."
)
def _check_rank_valid(g, rank: int):
"""Check the rank: 0 <= rank < world_size."""
if rank < 0:
raise ValueError("rank '{}' is negative.".format(rank))
if rank >= g.world_size:
raise ValueError(
"rank '{}' must be less than world size '{}'".format(rank, g.world_size)
)
def _check_tensor_list_input(tensor_list):
"""Check if the input is a list of supported tensor types."""
if not isinstance(tensor_list, list):
raise RuntimeError(
"The input must be a list of tensors. "
"Got '{}'.".format(type(tensor_list))
)
if not tensor_list:
raise RuntimeError("Got an empty list of tensors.")
for t in tensor_list:
_check_single_tensor_input(t)
def _check_tensor_lists_input(tensor_lists):
"""Check if the input is a list of lists of supported tensor types."""
if not isinstance(tensor_lists, list):
raise RuntimeError(
"The input must be a list of lists of tensors. "
"Got '{}'.".format(type(tensor_lists))
)
if not tensor_lists:
raise RuntimeError(f"Did not receive tensors. Got: {tensor_lists}")
for t in tensor_lists:
_check_tensor_list_input(t)
def _check_root_tensor_valid(length, root_tensor):
"""Check the root_tensor device is 0 <= root_tensor < length"""
if root_tensor < 0:
raise ValueError("root_tensor '{}' is negative.".format(root_tensor))
if root_tensor >= length:
raise ValueError(
"root_tensor '{}' is greater than the number of GPUs: "
"'{}'".format(root_tensor, length)
)
python/ray/util/collective/collective_group/__init__.py
View File
python/ray/util/collective/collective_group/base_collective_group.py
+91
View File
@@ -0,0 +1,91 @@
"""Abstract class for collective groups."""
from abc import ABCMeta, abstractmethod
from ray.util.collective.types import (
AllGatherOptions,
AllReduceOptions,
BarrierOptions,
BroadcastOptions,
RecvOptions,
ReduceOptions,
ReduceScatterOptions,
SendOptions,
)
class BaseGroup(metaclass=ABCMeta):
def __init__(self, world_size: int, rank: int, group_name: str):
"""Init the process group with basic information.
Args:
world_size: The total number of processes in the group.
rank: The rank of the current process.
group_name: The group name.
"""
self._world_size = world_size
self._rank = rank
self._group_name = group_name
@property
def rank(self):
"""Return the rank of the current process."""
return self._rank
@property
def world_size(self):
"""Return the number of processes in this group."""
return self._world_size
@property
def group_name(self):
"""Return the group name of this group."""
return self._group_name
def destroy_group(self):
"""GC the communicators."""
pass
@classmethod
def backend(cls):
"""The backend of this collective group."""
raise NotImplementedError()
@classmethod
@abstractmethod
def check_backend_availability(cls) -> bool:
"""Check if the backend is available."""
raise NotImplementedError()
@abstractmethod
def allreduce(self, tensor, allreduce_options=AllReduceOptions()):
raise NotImplementedError()
@abstractmethod
def barrier(self, barrier_options=BarrierOptions()):
raise NotImplementedError()
@abstractmethod
def reduce(self, tensor, reduce_options=ReduceOptions()):
raise NotImplementedError()
@abstractmethod
def allgather(self, tensor_list, tensor, allgather_options=AllGatherOptions()):
raise NotImplementedError()
@abstractmethod
def broadcast(self, tensor, broadcast_options=BroadcastOptions()):
raise NotImplementedError()
@abstractmethod
def reducescatter(
self, tensor, tensor_list, reducescatter_options=ReduceScatterOptions()
):
raise NotImplementedError()
@abstractmethod
def send(self, tensor, send_options: SendOptions):
raise NotImplementedError()
@abstractmethod
def recv(self, tensor, recv_options: RecvOptions):
raise NotImplementedError()
python/ray/util/collective/collective_group/cuda_stream.py
+95
View File
@@ -0,0 +1,95 @@
import logging
import threading
import cupy
from ray.util.collective.collective_group import nccl_util
from ray.util.collective.const import ENV
NCCL_STREAM_POOL_SIZE = 32
MAX_GPU_PER_ACTOR = 16
logger = logging.getLogger(__name__)
class StreamPool:
"""The class that represents a stream pool associated with a GPU.
When multistream is enabled, we will allocate a pool of streams for each
GPU, and get available stream from this pool when a collective kernel is
initialized. This enables overlapping computation/communication kernels
using multiple CUDA streams, given that the streams a appropriately
synchronized. The class is thread-safe.
Args:
device_idx: the absolute index of the device for this pool.
"""
def __init__(self, device_idx: int):
self.device_idx = device_idx
self._initialized = False
self._initialized_lock = threading.Lock()
self._pool = [None] * NCCL_STREAM_POOL_SIZE
self._counter = 0
self._pool_lock = threading.Lock()
def get_stream(self):
"""Get an available stream from the pool.
The function locks the stream pool and releases the lock before
returning.
Returns:
stream (cupy.cuda.Stream): the returned stream from pool.
"""
# check the flag
self._initialized_lock.acquire()
if not self._initialized:
self._init_once()
self._initialized_lock.release()
# Get the stream from the pool.
self._pool_lock.acquire()
stream = self._pool[self._counter]
self._counter = (self._counter + 1) % NCCL_STREAM_POOL_SIZE
self._pool_lock.release()
return stream
def _init_once(self):
"""Initialize the stream pool only for once."""
with nccl_util.Device(self.device_idx):
for i in range(NCCL_STREAM_POOL_SIZE):
# this is the only place where self._pool will be written.
if ENV.NCCL_USE_MULTISTREAM.val:
logger.debug("NCCL multistream enabled.")
self._pool[i] = cupy.cuda.Stream(null=False, non_blocking=False)
else:
logger.debug("NCCL multistream disabled.")
self._pool[i] = cupy.cuda.Stream.null
self._init_flag = True
# This is a map from GPU index to its stream pool.
# It is supposed to be READ-ONLY out of this file
_device_stream_pool_map = dict()
def _init_stream_pool():
global _device_stream_pool_map
for i in range(MAX_GPU_PER_ACTOR):
_device_stream_pool_map[i] = StreamPool(i)
def get_stream_pool(device_idx):
"""Get the CUDA stream pool of a GPU device."""
# In case there will be multiple threads writing to the pool.
lock = threading.Lock()
lock.acquire()
if not _device_stream_pool_map:
_init_stream_pool()
lock.release()
return _device_stream_pool_map[device_idx]
python/ray/util/collective/collective_group/nccl_collective_group.py
+850
View File
@@ -0,0 +1,850 @@
import datetime
import logging
import time
from typing import Callable, List, Optional
import ray
from ray.util.collective.collective_group import nccl_util
from ray.util.collective.collective_group.base_collective_group import BaseGroup
from ray.util.collective.collective_group.cuda_stream import get_stream_pool
from ray.util.collective.const import ENV, get_store_name
from ray.util.collective.types import (
AllGatherOptions,
AllReduceOptions,
Backend,
BarrierOptions,
BroadcastOptions,
RecvOptions,
ReduceOptions,
ReduceScatterOptions,
SendOptions,
torch_available,
)
logger = logging.getLogger(__name__)
try:
import cupy
import torch
_NCCL_AVAILABLE = True
_LOG_NCCL_WARNING = False
except ImportError:
_NCCL_AVAILABLE = False
_LOG_NCCL_WARNING = True
class Rendezvous:
"""A rendezvous class for different actor/task processes to meet.
To initialize an NCCL collective communication group, different
actors/tasks spawned in Ray in a collective group needs to meet
each other to synchronize the NCCLUniqueID. This class guarantees
they meet via the NCCLUniqueIDStore, initialized on the rank=0
process.
Args:
store_key: the unique store key, usually as a concatanation
of group_name and communicator key. See `get_nccl_communicator`
for more details.
"""
def __init__(self, store_key: str):
if not store_key:
raise ValueError(
"Invalid store_key. The store_key is a concatenation of "
"'group_name' and the 'communicator_key'. See the "
"docstring of `get_nccl_communicator` for details."
)
self._store_key = store_key
self._store_name = None
self._store = None
def meet(self, timeout_s: int = 180):
"""Meet at the named actor store.
Args:
timeout_s: timeout in seconds.
"""
if timeout_s <= 0:
raise ValueError(
"The 'timeout' argument must be positive. "
"Got '{}'.".format(timeout_s)
)
self._store_name = get_store_name(self._store_key)
timeout_delta = datetime.timedelta(seconds=timeout_s)
elapsed = datetime.timedelta(seconds=0)
start_time = datetime.datetime.now()
while elapsed < timeout_delta:
try:
logger.debug(
"Trying to meet at the store '{}'".format(self._store_name)
)
self._store = ray.get_actor(self._store_name)
except ValueError:
logger.debug(
"Failed to meet at the store '{}'."
"Trying again...".format(self._store_name)
)
time.sleep(1)
elapsed = datetime.datetime.now() - start_time
continue
logger.debug("Successful rendezvous!")
break
if not self._store:
raise RuntimeError(
"Unable to meet other processes "
"at the rendezvous store. If you are using "
"P2P communication, please check if tensors "
"are put in the correct GPU. "
)
@property
def store(self):
return self._store
def get_nccl_id(self, timeout_s: int = 180):
"""Get the NCCLUniqueID from the store through Ray.
Args:
timeout_s: timeout in seconds.
Returns:
uid: the NCCLUniqueID if successful.
"""
if not self._store:
raise ValueError("Rendezvous store is not setup.")
try:
uid = ray.get(self._store.wait_and_get_id.remote(), timeout=timeout_s)
except ray.exceptions.GetTimeoutError:
raise RuntimeError(
f"Unable to get the NCCLUniqueID from the store within {timeout_s} seconds."
) from None
return uid
class NCCLGroup(BaseGroup):
def __init__(self, world_size: int, rank: int, group_name: str):
"""Init an NCCL collective group."""
super(NCCLGroup, self).__init__(world_size, rank, group_name)
# communicator and stream cache.
# TODO (Hao): we need a lock here...
self._dev_comm_map = {}
self._dev_streams_map = {}
# record the used GPU IDs.
self._used_gpu_indices = set()
# TODO(Fu): might need an event map
self._dev_event_map = {}
if nccl_util.get_nccl_build_version() < 2000:
raise RuntimeError("NCCL in Ray requires NCCL >= 2.0.")
if nccl_util.get_nccl_runtime_version() < 2704:
logger.warning("NCCL send/recv calls requires NCCL>=2.7.4")
def destroy_group(self):
"""Destroy the group and release NCCL communicators."""
if len(self._dev_comm_map.keys()) > 0:
# TODO(Hao): check this barrier call
# self.barrier()
# Destroy the communicators and streams.
for comm_key, comms in self._dev_comm_map.items():
for c in comms:
c.destroy()
self._dev_comm_map[comm_key] = None
if self.rank == 0:
for comm_key in self._dev_comm_map:
assert not self._dev_comm_map[comm_key]
group_key = self._generate_group_key(comm_key)
self._destroy_store(group_key)
self._barrier_tensor = None
self._dev_comm_map = None
self._dev_streams_map = None
super(NCCLGroup, self).destroy_group()
@classmethod
def backend(cls):
return Backend.NCCL
@classmethod
def check_backend_availability(cls) -> bool:
global _LOG_NCCL_WARNING, _NCCL_AVAILABLE
if _LOG_NCCL_WARNING:
logger.warning(
"NCCL is not available. Please install Cupy "
"following the guide at: "
"https://docs.cupy.dev/en/stable/install.html."
)
_LOG_NCCL_WARNING = False
return _NCCL_AVAILABLE
def allreduce(
self,
tensors: list,
allreduce_options: AllReduceOptions = AllReduceOptions(),
):
"""AllReduce tensors across the collective group following options.
Args:
tensors: the list of tensors to be reduced. Each tensor must
reside on one GPU of the current process.
allreduce_options: allreduce options.
"""
def collective_fn(input_tensor, output_tensor, comm, stream):
comm.allReduce(
nccl_util.get_tensor_ptr(input_tensor),
nccl_util.get_tensor_ptr(output_tensor),
nccl_util.get_tensor_n_elements(input_tensor),
nccl_util.get_nccl_tensor_dtype(input_tensor),
nccl_util.get_nccl_reduce_op(allreduce_options.reduceOp),
stream.ptr,
)
self._collective(tensors, tensors, collective_fn)
def barrier(self, barrier_options: BarrierOptions = BarrierOptions()):
"""Blocks until all processes reach this barrier.
Args:
barrier_options: barrier options.
"""
# Get the device list.
if self._used_gpu_indices:
devices = list(self._used_gpu_indices)
else:
devices = list(range(nccl_util.get_num_gpus()))
barrier_tensors = [None] * len(devices)
for i, d in enumerate(devices):
with nccl_util.Device(d):
barrier_tensors[i] = cupy.array([1])
self.allreduce(barrier_tensors)
def reduce(self, tensors: list, reduce_options: ReduceOptions = ReduceOptions()):
"""Reduce tensors to a destination gpu following options.
Args:
tensors: the list of tensors to be reduced, each tensor
must reside on one gpu of the current process.
reduce_options: reduce options.
"""
root_rank = len(tensors) * reduce_options.root_rank + reduce_options.root_tensor
def collective_fn(input_tensor, output_tensor, comm, stream):
comm.reduce(
nccl_util.get_tensor_ptr(input_tensor),
nccl_util.get_tensor_ptr(output_tensor),
nccl_util.get_tensor_n_elements(input_tensor),
nccl_util.get_nccl_tensor_dtype(input_tensor),
nccl_util.get_nccl_reduce_op(reduce_options.reduceOp),
root_rank,
stream.ptr,
)
self._collective(tensors, tensors, collective_fn)
def broadcast(
self,
tensors: list,
broadcast_options: BroadcastOptions = BroadcastOptions(),
):
"""Broadcast tensors to all other gpus following options.
Args:
tensors: tensors to be broadcast or received.
broadcast_options: broadcast options.
"""
root_rank = (
len(tensors) * broadcast_options.root_rank + broadcast_options.root_tensor
)
def collective_fn(input_tensor, output_tensor, comm, stream):
comm.broadcast(
nccl_util.get_tensor_ptr(input_tensor),
nccl_util.get_tensor_ptr(output_tensor),
nccl_util.get_tensor_n_elements(input_tensor),
nccl_util.get_nccl_tensor_dtype(input_tensor),
root_rank,
stream.ptr,
)
self._collective(tensors, tensors, collective_fn)
def allgather(
self,
tensor_lists: list,
tensors: list,
allgather_options: AllGatherOptions = AllGatherOptions(),
):
"""Allgather tensors across gpus into a list of tensors.
Args:
tensor_lists: allgathered tensors.
tensors: the list of tensors to allgather across the group.
Each tensor must lolcate on a GPU of the process.
allgather_options: allgather options.
"""
def collective_fn(input_tensor, output_tensor, comm, stream):
comm.allGather(
nccl_util.get_tensor_ptr(input_tensor),
nccl_util.get_tensor_ptr(output_tensor),
nccl_util.get_tensor_n_elements(input_tensor),
nccl_util.get_nccl_tensor_dtype(input_tensor),
stream.ptr,
)
_check_inputs_compatibility_for_scatter_gather(tensors, tensor_lists)
output_flattened = [
_flatten_for_scatter_gather(tensor_list, copy=False)
for tensor_list in tensor_lists
]
def postprocess_fn(stream):
# TODO(Hao): designate a copy stream.
for i, tensor_list in enumerate(tensor_lists):
for j, tensor in enumerate(tensor_list):
nccl_util.copy_tensor(tensor, output_flattened[i][j])
self._collective(
tensors, output_flattened, collective_fn, postprocess_fn=postprocess_fn
)
def reducescatter(
self,
tensors: list,
tensor_lists: list,
reducescatter_options: ReduceScatterOptions = ReduceScatterOptions(),
):
"""Reduce then scatter a list of tensors across the group.
Args:
tensors: the output tensors (could be unspecified), each
located on a GPU of the current process.
tensor_lists: the list of tensors to be reduced then
scattered.
reducescatter_options: reduce-scatter options.
"""
def collective_fn(input_tensor, output_tensor, comm, stream):
comm.reduceScatter(
nccl_util.get_tensor_ptr(input_tensor),
nccl_util.get_tensor_ptr(output_tensor),
nccl_util.get_tensor_n_elements(output_tensor),
nccl_util.get_nccl_tensor_dtype(output_tensor),
nccl_util.get_nccl_reduce_op(reducescatter_options.reduceOp),
stream.ptr,
)
_check_inputs_compatibility_for_scatter_gather(tensors, tensor_lists)
input_flattened = [
_flatten_for_scatter_gather(tensor_list, copy=False)
for tensor_list in tensor_lists
]
def preprocess_fn(stream):
for i, tensor_list in enumerate(tensor_lists):
for j, tensor in enumerate(tensor_list):
nccl_util.copy_tensor(input_flattened[i][j], tensor)
self._collective(
input_flattened, tensors, collective_fn, preprocess_fn=preprocess_fn
)
def send(self, tensors: list, send_options: SendOptions = SendOptions()):
"""Send a tensor to a destination gpu in the group.
Args:
tensors: the tensor to send.
send_options: send options.
"""
def p2p_fn(tensor, comm, stream, peer):
comm.send(
nccl_util.get_tensor_ptr(tensor),
send_options.n_elements
if send_options.n_elements > 0
else nccl_util.get_tensor_n_elements(tensor),
nccl_util.get_nccl_tensor_dtype(tensor),
peer,
stream.ptr,
)
self._point2point(
tensors, p2p_fn, send_options.dst_rank, send_options.dst_gpu_index
)
def recv(self, tensors: list, recv_options: RecvOptions = RecvOptions()):
"""Receive a tensor from a source gpu in the group.
Args:
tensors: the received tensor.
recv_options: Receive options.
"""
def p2p_fn(tensor, comm, stream, peer):
comm.recv(
nccl_util.get_tensor_ptr(tensor),
recv_options.n_elements
if recv_options.n_elements > 0
else nccl_util.get_tensor_n_elements(tensor),
nccl_util.get_nccl_tensor_dtype(tensor),
peer,
stream.ptr,
)
self._point2point(
tensors, p2p_fn, recv_options.src_rank, recv_options.src_gpu_index
)
def _get_nccl_collective_communicator(self, comm_key: str, device_list: list):
"""Create or retrieve an NCCL communicator from cache.
If the communicator is found in cache, return the communicator. If not,
a communicator and a stream will be created and put in cache.
TODO(Hao): this function is not thread-safe now.
Args:
comm_key: the key to query the communicator cache.
device_list: a list of GPU devices of the current process
that participates into the collective.
Returns:
communicator: the NCCL communicator corresponded to the devices.
"""
if not comm_key:
raise RuntimeError("Got empty communicator key.")
for d in device_list:
self._used_gpu_indices.add(d)
# TODO(Hao): lock the _dev_comm_map here.
if comm_key in self._dev_comm_map:
return self._dev_comm_map[comm_key]
group_key = self._generate_group_key(comm_key)
if self.rank == 0:
nccl_uid = self._generate_nccl_uid(group_key)
else:
rendezvous = Rendezvous(group_key)
rendezvous.meet()
nccl_uid = rendezvous.get_nccl_id()
# Now create the communicators
actual_world_size = len(device_list) * self.world_size
comms = [None] * len(device_list)
streams = [None] * len(device_list)
events = [None] * len(device_list)
nccl_util.groupStart()
for i, device in enumerate(device_list):
actual_rank = self.rank * len(device_list) + i
with nccl_util.Device(device):
comms[i] = nccl_util.create_nccl_communicator(
actual_world_size, nccl_uid, actual_rank
)
# request a stream from the pool
# note the device_idx is absolute index.
streams[i] = get_stream_pool(device).get_stream()
# TODO(Fu): double check the parameters
events[i] = cupy.cuda.Event()
nccl_util.groupEnd()
# TODO(Fu): lock
self._dev_comm_map[comm_key] = comms
self._dev_streams_map[comm_key] = streams
self._dev_event_map[comm_key] = events
return comms
@staticmethod
def _sync_streams(device_list, events, streams):
"""Let NCCL streams wait for current streams for every device."""
# TODO(Fu): recordStream besides calling this function?
if ENV.NCCL_USE_MULTISTREAM.val:
for i, device in enumerate(device_list):
with nccl_util.Device(device):
events[i].record(cupy.cuda.get_current_stream())
streams[i].wait_event(events[i])
def _get_nccl_p2p_communicator(
self,
comm_key: str,
my_gpu_idx: int,
peer_rank: int,
peer_gpu_idx: int,
):
"""Create or retrieve an NCCL communicator for p2p tasks.
Note(Hao): this function is not thread-safe now.
Args:
comm_key: communicator key.
my_gpu_idx: the gpu index on the current process.
peer_rank: the rank of the destination process.
peer_gpu_idx: the gpu index on the peer process.
Returns:
communicator
"""
if not comm_key:
raise RuntimeError("Got empty communicator key.")
# TODO(Hao): lock the _dev_comm_map here.
if comm_key in self._dev_comm_map:
return self._dev_comm_map[comm_key]
# Note (Hao): This is a bit complex so I decide to take a note here.
# Here we need to consider three cases:
# Case 1: src_rank != dst_rank, hence the send and recv happen on
# different process (actors/tasks); each process makes independent
# collective calls and manages corresponding communicators.
# Case 2: src_rank == dst_rank, src_gpu_idx == dst_gpu_idx; for
# this case, we simply throw a RuntimeError;
# Case 3: src_rank == dst_rank, src_gpu_idx != dst_gpu_idx, which
# means the send and recv will be called on the same process. We
# DO NOT support this case for now. We need to properly scope:
# (1) communicators creation, and
# (2) send/recv calls
# using groupStart(（ and groupEnd() calls to avoid deadlocks.
if self.rank < peer_rank:
my_p2p_rank = 0
elif self.rank > peer_rank:
my_p2p_rank = 1
else:
raise RuntimeError(
"Send and recv happens on the same process! "
"ray.util.collective does not support this case as of now. "
"Alternatively, consider doing GPU to GPU memcpy?"
)
group_key = self._generate_group_key(comm_key)
if my_p2p_rank == 0:
nccl_uid = self._generate_nccl_uid(group_key)
else:
rendezvous = Rendezvous(group_key)
rendezvous.meet()
nccl_uid = rendezvous.get_nccl_id()
# create the p2p communicators
with nccl_util.Device(my_gpu_idx):
comm = nccl_util.create_nccl_communicator(2, nccl_uid, my_p2p_rank)
stream = get_stream_pool(my_gpu_idx).get_stream()
event = cupy.cuda.Event()
# TODO(Fu): lock and might need to add event
self._dev_comm_map[comm_key] = [comm]
self._dev_streams_map[comm_key] = [stream]
self._dev_event_map[comm_key] = [event]
return [comm]
def _generate_group_key(self, comm_key):
"""Generate a unique key used to initialize the KV store.
The group key is a concatenation of the communicator key and
the group name, following: [comm_key]@[group_name].
"""
return comm_key + "@" + self.group_name
@staticmethod
def _destroy_store(group_key: str):
"""Destroy the KV store (Ray named actor).
Args:
group_key: the unique key to retrieve the KV store.
"""
store_name = get_store_name(group_key)
store = ray.get_actor(store_name)
# ray.get([store.__ray_terminate__.remote()])
ray.kill(store)
def _generate_nccl_uid(self, key: str):
"""Generate an NCCL unique ID for initializing communicators.
The method will also create a KV store using Ray named actor and store
the NCCLUniqueID in the store. The store needs to be garbage collected
when destroying the collective group.
Args:
key: the key of the .
Returns:
NCCLUniqueID (str): NCCL unique ID.
"""
group_uid = nccl_util.get_nccl_unique_id()
store_name = get_store_name(key)
# Avoid a potential circular dependency in ray/actor.py
from ray.util.collective.util import NCCLUniqueIDStore
store = NCCLUniqueIDStore.options(name=store_name, lifetime="detached").remote(
store_name
)
ray.get([store.set_id.remote(group_uid)])
return group_uid
def _collective(
self,
input_tensors: list,
output_tensors: list,
collective_fn: Callable,
preprocess_fn: Optional[Callable] = None,
postprocess_fn: Optional[Callable] = None,
):
"""A method to encapsulate all collective calls.
Args:
input_tensors: the list of the input tensors.
output_tensors: the list of the output tensors.
collective_fn: the collective function call.
preprocess_fn: preprocess procedures before collective calls.
postprocess_fn: postprocess procedures after collective calls.
"""
_check_gpu_tensors(input_tensors)
_check_gpu_tensors(output_tensors)
devices = nccl_util.get_tensor_device_list(input_tensors)
key = _get_comm_key_from_devices(devices)
comms = self._get_nccl_collective_communicator(key, devices)
streams = self._dev_streams_map[key]
events = self._dev_event_map[key]
# TODO(Hao): sync streams and events
self._sync_streams(devices, events, streams)
# Make the collective call
if preprocess_fn:
preprocess_fn(streams)
nccl_util.groupStart()
# TODO(Fu): how to recordStreams as there are no library functions
# We also need to make sure input tensors are not freed before their
# usages on ncclStreams finish. This can be achieved by calling
# c10::cuda::CUDACachingAllocator::recordStream, which remembers the
# usage stream (ncclStream), creates an event on the usage stream
# when GC attempts to free the input tensor, and delays GC until that
# event is done.
for i, tensor in enumerate(input_tensors):
collective_fn(tensor, output_tensors[i], comms[i], streams[i])
nccl_util.groupEnd()
if postprocess_fn:
postprocess_fn(streams)
def _point2point(
self,
tensors: list,
p2p_fn: Callable,
peer_rank: int,
peer_gpu_idx: int,
):
"""A method to encapsulate all peer-to-peer calls (i.e., send/recv).
Args:
tensors: the tensor to send or receive.
p2p_fn: the p2p function call.
peer_rank: the rank of the peer process.
peer_gpu_idx: the index of the gpu on the peer process.
"""
# check send/recv availability.
if nccl_util.get_nccl_runtime_version() < 2704:
raise RuntimeError(
"P2p send/recv requires NCCL >= 2.7.4. "
"Got '{}'.".format(nccl_util.get_nccl_runtime_version())
)
_check_gpu_tensors(tensors)
# we currently only support single device to single device send/recv.
assert len(tensors) == 1
my_gpu_idx = nccl_util.get_tensor_device(tensors[0])
comm_key = _get_comm_key_send_recv(
self.rank, my_gpu_idx, peer_rank, peer_gpu_idx
)
comms = self._get_nccl_p2p_communicator(
comm_key, my_gpu_idx, peer_rank, peer_gpu_idx
)
streams = self._dev_streams_map[comm_key]
events = self._dev_event_map[comm_key]
# TODO(Hao): sync streams and events
self._sync_streams([my_gpu_idx], events, streams)
# We have made sure that self.rank != peer_rank during API check.
peer_p2p_rank = 0 if self.rank > peer_rank else 1
for i, tensor in enumerate(tensors):
p2p_fn(tensor, comms[i], streams[i], peer_p2p_rank)
# Record the stream to avoid tensor being freed before the send/recv is completed.
torch_stream = torch.cuda.ExternalStream(streams[i].ptr)
tensor.record_stream(torch_stream)
def _flatten_for_scatter_gather(tensor_list: list, copy: bool = False):
"""Flatten the tensor for gather/scatter operations.
Args:
tensor_list: the list of tensors to be scattered/gathered.
copy: whether the copy the tensors in tensor_list into the buffer.
Returns:
The flattened tensor buffer.
"""
if not tensor_list:
raise RuntimeError("Received an empty list.")
t = tensor_list[0]
buffer_shape = [len(tensor_list)] + nccl_util.get_tensor_shape(t)
# TODO(wuxibin): cupy doesn't support bfloat16 for now,
# once it is supported, we can eliminate this if statement.
#
# Allocate using the same backend as the tensors in `tensor_list`.
# Use torch only when the tensors are torch.Tensor; otherwise fall back to CuPy.
use_torch = False
if torch_available():
try:
import torch
use_torch = isinstance(t, torch.Tensor)
except ImportError:
use_torch = False
if use_torch:
buffer = torch.empty(tuple(buffer_shape), dtype=t.dtype, device=t.device)
else:
# note we need a cupy dtype here.
dtype = nccl_util.get_cupy_tensor_dtype(t)
device = nccl_util.get_tensor_device(t)
with nccl_util.Device(device):
buffer = cupy.empty(buffer_shape, dtype=dtype)
if copy:
for i, tensor in enumerate(tensor_list):
nccl_util.copy_tensor(buffer[i], tensor)
return buffer
def _check_inputs_compatibility_for_scatter_gather(tensors, tensor_lists):
"""Check the compatibility between tensor input and tensor list input."""
if not tensors or not isinstance(tensors, list):
raise RuntimeError("The first argument 'tensors' expects a list of tensors.")
if not tensor_lists or not isinstance(tensor_lists, list):
raise RuntimeError(
"The second argument 'tensor_lists' expects a list of tensor list."
)
dtype = nccl_util.get_nccl_tensor_dtype(tensors[0])
shape = nccl_util.get_tensor_shape(tensors[0])
for i, tensor_list in enumerate(tensor_lists):
# check all tensor in `tensors` match.
dt = nccl_util.get_nccl_tensor_dtype(tensors[i])
if dt != dtype:
raise RuntimeError(
"All tensor operands to scatter/gather must "
"have the same dtype. Got '{}' and '{}'.".format(dt, dtype)
)
# Note: typically CCL libraries only requires they have the same
# number of elements; Here we make it more strict -- we require
# exact shape match.
s = nccl_util.get_tensor_shape(tensors[i])
if s != shape:
raise RuntimeError(
"All tensor operands to scatter/gather must "
"have the same shape. Got '{}' and '{}'.".format(s, shape)
)
# check all tensors in `tensor_lists` match.
for t in tensor_lists[i]:
# check dtype
dt = nccl_util.get_nccl_tensor_dtype(t)
if dt != dtype:
raise RuntimeError(
"All tensor operands to scatter/gather must "
"have the same dtype. Got '{}' and '{}'.".format(dt, dtype)
)
s = nccl_util.get_tensor_shape(t)
if s != shape:
raise RuntimeError(
"All tensor operands to scatter/gather must "
"have the same shape. Got '{}' and '{}'.".format(s, shape)
)
def _check_gpu_tensors(tensors):
"""Check all tensors are distributed on different GPUs."""
if not tensors or not isinstance(tensors, list):
raise RuntimeError("'tensors' must be a nonempty list.")
if len(tensors) > nccl_util.get_num_gpus():
raise RuntimeError(
"Tensor list cannot be larger than the number"
"of available GPUs. Got {} > {}.".format(
len(tensors), nccl_util.get_num_gpus()
)
)
t0 = tensors[0]
dt = nccl_util.get_nccl_tensor_dtype(t0)
s = nccl_util.get_tensor_shape(t0)
d = nccl_util.get_tensor_device(t0)
for i, t in enumerate(tensors):
if i == 0:
continue
# We need to check the following:
# (1) tensor is cuda (already checked during API)
# (2) tensor dtype
# (3) tensor shape match
# (4) each tensor is on a different GPU
dtype = nccl_util.get_nccl_tensor_dtype(t)
if dt != dtype:
raise RuntimeError(
"Tensors must have identical dtype. Got: '{}'.".format(dtype)
)
shape = nccl_util.get_tensor_shape(t)
if s != shape:
raise RuntimeError(
"Tensor must have identical shape. Got: '{}'.".format(shape)
)
device = nccl_util.get_tensor_device(t)
if device == d:
raise RuntimeError("Tensor must be on distinct GPUs.")
def _get_comm_key_from_devices(devices: List[int]):
"""Return a key from a list of devices for collective calls.
For example, if the tensors are on gpus 0, 1, 2, 3,
then the key would be "0,1,2,3".
Args:
devices: a list of GPU device indices
Returns:
str: a string represents the key to query the communicator cache.
"""
return ",".join([str(d) for d in devices])
def _get_comm_key_send_recv(
my_rank: int, my_gpu_idx: int, peer_rank: int, peer_gpu_idx: int
):
"""Return a key given source and destination ranks for p2p tasks.
The p2p key is in the following form:
[min_rank]_[gpu_index]:[max_rank]_[gpu_index].
Args:
my_rank: the rank of the source process.
my_gpu_idx: the source gpu index on the process.
peer_rank: the rank of the destination process.
peer_gpu_idx: the destination gpu index on the process.
Returns:
comm_key: a string key to query the communication cache.
"""
if my_rank < peer_rank:
lower_key = str(my_rank) + "_" + str(my_gpu_idx)
higher_key = str(peer_rank) + "_" + str(peer_gpu_idx)
elif my_rank > peer_rank:
lower_key = str(peer_rank) + "_" + str(peer_gpu_idx)
higher_key = str(my_rank) + "_" + str(my_gpu_idx)
else:
raise RuntimeError(
"Send and recv happens on the same process. ray.util.collective "
"does not support this case as of now. Alternatively, consider "
"doing GPU to GPU memcpy?"
)
comm_key = lower_key + ":" + higher_key
return comm_key
python/ray/util/collective/collective_group/nccl_util.py
+298
View File
@@ -0,0 +1,298 @@
"""Code to wrap some NCCL API calls."""
from typing import Any, List
import numpy
try:
import cupy
from cupy.cuda import (
Device, # noqa: F401
nccl,
)
from cupy.cuda.nccl import (
NcclCommunicator,
get_build_version,
get_version,
groupEnd, # noqa: F401
groupStart, # noqa: F401
)
except ImportError:
raise ImportError("NCCL in Ray requires Cupy being available!")
from ray.util.collective.types import ReduceOp, torch_available
NCCL_REDUCE_OP_MAP = {
ReduceOp.SUM: nccl.NCCL_SUM,
ReduceOp.PRODUCT: nccl.NCCL_PROD,
ReduceOp.MIN: nccl.NCCL_MIN,
ReduceOp.MAX: nccl.NCCL_MAX,
}
# cupy types are the same with numpy types
NUMPY_NCCL_DTYPE_MAP = {
# INT types
numpy.int_: nccl.NCCL_INT64,
numpy.uint8: nccl.NCCL_UINT8,
numpy.uint32: nccl.NCCL_UINT32,
numpy.uint64: nccl.NCCL_UINT64,
numpy.int8: nccl.NCCL_INT8,
numpy.int32: nccl.NCCL_INT32,
numpy.int64: nccl.NCCL_INT64,
# FLOAT types
numpy.half: nccl.NCCL_HALF,
numpy.float16: nccl.NCCL_FLOAT16,
numpy.float32: nccl.NCCL_FLOAT32,
numpy.float64: nccl.NCCL_FLOAT64,
numpy.double: nccl.NCCL_DOUBLE,
}
if torch_available():
import torch
import torch.utils.dlpack
TORCH_NCCL_DTYPE_MAP = {
torch.bool: nccl.NCCL_INT8,
# INT types
torch.int: nccl.NCCL_INT,
torch.uint8: nccl.NCCL_UINT8,
torch.int8: nccl.NCCL_INT8,
torch.int32: nccl.NCCL_INT32,
torch.int64: nccl.NCCL_INT64,
torch.long: nccl.NCCL_INT64,
# FLOAT types
torch.half: nccl.NCCL_HALF,
torch.float: nccl.NCCL_FLOAT,
torch.float16: nccl.NCCL_FLOAT16,
torch.float32: nccl.NCCL_FLOAT32,
torch.float64: nccl.NCCL_FLOAT64,
torch.double: nccl.NCCL_DOUBLE,
}
# Older versions of cupy don't support bfloat16.
if hasattr(nccl, "NCCL_BFLOAT16"):
TORCH_NCCL_DTYPE_MAP[torch.bfloat16] = nccl.NCCL_BFLOAT16
TORCH_NUMPY_DTYPE_MAP = {
# INT types
torch.int: numpy.int32,
torch.uint8: numpy.uint8,
torch.int8: numpy.int8,
torch.int32: numpy.int32,
torch.int64: numpy.int64,
torch.long: numpy.int64,
# FLOAT types
torch.half: numpy.half,
torch.float: numpy.float32,
torch.float16: numpy.float16,
torch.float32: numpy.float32,
torch.float64: numpy.float64,
}
def get_num_gpus():
"""Returns the number of compute-capable GPUs."""
return cupy.cuda.runtime.getDeviceCount()
def get_nccl_build_version():
return get_build_version()
def get_nccl_runtime_version():
return get_version()
def get_nccl_unique_id():
return nccl.get_unique_id()
def create_nccl_communicator(world_size: int, nccl_unique_id: bytes, rank: int):
"""Create an NCCL communicator using NCCL APIs.
Args:
world_size: the number of processes of this communicator group.
nccl_unique_id: the NCCLUniqueID for this group.
rank: the rank of this process.
Returns:
comm: an NCCL communicator.
"""
comm = NcclCommunicator(world_size, nccl_unique_id, rank)
return comm
def get_nccl_reduce_op(reduce_op: ReduceOp):
"""Map the reduce op to NCCL reduce op type.
Args:
reduce_op: ReduceOp Enum (SUM/PRODUCT/MIN/MAX).
Returns:
the mapped NCCL reduce op.
"""
if reduce_op not in NCCL_REDUCE_OP_MAP:
raise RuntimeError("NCCL does not support reduce op: '{}'.".format(reduce_op))
return NCCL_REDUCE_OP_MAP[reduce_op]
def get_nccl_tensor_dtype(tensor):
"""Return the corresponded NCCL dtype given a tensor."""
if isinstance(tensor, cupy.ndarray):
return NUMPY_NCCL_DTYPE_MAP[tensor.dtype.type]
if torch_available():
if isinstance(tensor, torch.Tensor):
return TORCH_NCCL_DTYPE_MAP[tensor.dtype]
raise ValueError(
"Unsupported tensor type. Got: {}. Supported "
"GPU tensor types are: torch.Tensor, "
"cupy.ndarray.".format(type(tensor))
)
def get_cupy_tensor_dtype(tensor):
"""Return the corresponded Cupy dtype given a tensor."""
if isinstance(tensor, cupy.ndarray):
return tensor.dtype.type
if torch_available():
if isinstance(tensor, torch.Tensor):
return TORCH_NUMPY_DTYPE_MAP[tensor.dtype]
raise ValueError(
"Unsupported tensor type. Got: {}. Supported "
"GPU tensor types are: torch.Tensor, "
"cupy.ndarray.".format(type(tensor))
)
def get_tensor_ptr(tensor):
"""Return the pointer to the underlying memory storage of a tensor."""
if isinstance(tensor, cupy.ndarray):
return tensor.data.ptr
if isinstance(tensor, numpy.ndarray):
return tensor.data
if torch_available():
if isinstance(tensor, torch.Tensor):
if not tensor.is_cuda:
raise RuntimeError(
"Torch tensor must be on GPU when using NCCL collectives."
)
return tensor.data_ptr()
raise ValueError(
"Unsupported tensor type. Got: {}. Supported "
"GPU tensor types are: torch.Tensor, "
"cupy.ndarray.".format(type(tensor))
)
def get_tensor_n_elements(tensor):
"""Return the number of elements in a tensor."""
if isinstance(tensor, cupy.ndarray) or isinstance(tensor, numpy.ndarray):
return tensor.size
if torch_available():
if isinstance(tensor, torch.Tensor):
return torch.numel(tensor)
raise ValueError(
"Unsupported tensor type. Got: {}. Supported "
"GPU tensor types are: torch.Tensor, "
"cupy.ndarray.".format(type(tensor))
)
def get_tensor_shape(tensor):
"""Return the shape of the tensor as a list."""
if isinstance(tensor, cupy.ndarray):
return list(tensor.shape)
if torch_available():
if isinstance(tensor, torch.Tensor):
return list(tensor.size())
raise ValueError(
"Unsupported tensor type. Got: {}. Supported "
"GPU tensor types are: torch.Tensor, "
"cupy.ndarray.".format(type(tensor))
)
def get_tensor_strides(tensor):
"""Return the strides of the tensor as a list."""
if isinstance(tensor, cupy.ndarray):
return [int(stride / tensor.dtype.itemsize) for stride in tensor.strides]
if torch_available():
if isinstance(tensor, torch.Tensor):
return list(tensor.stride())
raise ValueError(
"Unsupported tensor type. Got: {}. Supported "
"GPU tensor types are: torch.Tensor, "
"cupy.ndarray.".format(type(tensor))
)
def get_tensor_device(tensor):
"""Return the GPU index of a tensor."""
if isinstance(tensor, cupy.ndarray):
try:
device = tensor.device.id
except AttributeError as exec:
raise RuntimeError("The tensor is not on a valid GPU.") from exec
elif torch_available() and isinstance(tensor, torch.Tensor):
device = tensor.device.index
if not isinstance(device, int):
raise RuntimeError("The tensor is not on a valid GPU.")
else:
raise ValueError("Unsupported tensor type. Got: {}.".format(type(tensor)))
return device
def copy_tensor(dst_tensor: Any, src_tensor: Any):
"""Copy the content from src_tensor to dst_tensor.
Args:
dst_tensor: the tensor to copy from.
src_tensor: the tensor to copy to.
"""
copied = True
if isinstance(dst_tensor, cupy.ndarray) and isinstance(src_tensor, cupy.ndarray):
cupy.copyto(dst_tensor, src_tensor)
elif torch_available():
if isinstance(dst_tensor, torch.Tensor) and isinstance(
src_tensor, torch.Tensor
):
dst_tensor.copy_(src_tensor)
elif isinstance(dst_tensor, torch.Tensor) and isinstance(
src_tensor, cupy.ndarray
):
t = torch.utils.dlpack.from_dlpack(src_tensor.toDlpack())
dst_tensor.copy_(t)
elif isinstance(dst_tensor, cupy.ndarray) and isinstance(
src_tensor, torch.Tensor
):
t = cupy.fromDlpack(torch.utils.dlpack.to_dlpack(src_tensor))
cupy.copyto(dst_tensor, t)
else:
copied = False
else:
copied = False
if not copied:
raise ValueError(
"Unsupported tensor type. Got: {} and {}. Supported "
"GPU tensor types are: torch.Tensor, cupy.ndarray.".format(
type(dst_tensor), type(src_tensor)
)
)
def get_tensor_device_list(tensors: List[Any]):
"""Returns the gpu devices of the list of input tensors.
Args:
tensors: a list of tensors, each locates on a GPU.
Returns:
list: the list of GPU devices.
"""
if not isinstance(tensors, list):
raise RuntimeError(
"Expect a list of tensors each locates on a GPU device. "
"Got: '{}'.".format(type(tensors))
)
devices = [get_tensor_device(t) for t in tensors]
return devices
python/ray/util/collective/collective_group/torch_gloo_collective_group.py
+290
View File
@@ -0,0 +1,290 @@
import logging
import os
import socket
import time
from typing import TYPE_CHECKING, List, Optional, Tuple
import numpy as np
import torch
import ray
import ray.experimental.internal_kv as internal_kv
from ray._common.network_utils import find_free_port, is_ipv6
from ray.util.collective.collective_group.base_collective_group import BaseGroup
from ray.util.collective.types import (
AllGatherOptions,
AllReduceOptions,
Backend,
BarrierOptions,
BroadcastOptions,
RecvOptions,
ReduceOp,
ReduceOptions,
ReduceScatterOptions,
SendOptions,
)
if TYPE_CHECKING:
import torch
logger = logging.getLogger(__name__)
try:
import torch.distributed as dist
_TORCH_DISTRIBUTED_AVAILABLE = True
TORCH_REDUCE_OP_MAP = {
ReduceOp.SUM: dist.ReduceOp.SUM,
ReduceOp.PRODUCT: dist.ReduceOp.PRODUCT,
ReduceOp.MIN: dist.ReduceOp.MIN,
ReduceOp.MAX: dist.ReduceOp.MAX,
}
except ImportError:
_TORCH_DISTRIBUTED_AVAILABLE = False
TORCH_REDUCE_OP_MAP = None
def get_master_address_metadata_key(group_name: str):
return f"collective_group_master_address_{group_name}"
def get_address_and_port() -> Tuple[str, int]:
"""Returns the IP address and a free port on this node."""
addr = ray.util.get_node_ip_address()
port = find_free_port(socket.AF_INET6 if is_ipv6(addr) else socket.AF_INET)
return addr, port
class TorchGLOOGroup(BaseGroup):
def __init__(
self,
world_size: int,
rank: int,
group_name: str,
gloo_timeout: Optional[int] = None,
):
# Initialize the default process group only once per process.
if not dist.is_initialized():
# Rendezvous: ensure a MASTER_ADDR:MASTER_PORT is published in internal_kv.
self._rendezvous(group_name, rank, gloo_timeout)
metadata_key = get_master_address_metadata_key(group_name)
try:
metadata = internal_kv._internal_kv_get(metadata_key)
except ValueError:
raise RuntimeError(
f"TorchGLOOGroup expected metadata in internal_kv with name `{metadata_key}`. "
"TorchGLOOGroup should not be instantiated directly. "
"Use ray.experimental.collective.create_collective_group to create a group."
)
if metadata is None:
raise RuntimeError(
f"Missing rendezvous metadata for group `{group_name}` under key `{metadata_key}`."
)
metadata = metadata.decode()
master_addr, master_port = metadata.split(":")
os.environ["MASTER_ADDR"] = master_addr
os.environ["MASTER_PORT"] = master_port
dist.init_process_group(
backend="gloo", init_method="env://", world_size=world_size, rank=rank
)
# Clean up rendezvous metadata after all ranks have initialized.
# dist.init_process_group is synchronous, so all ranks have read the metadata by now.
if rank == 0:
try:
internal_kv._internal_kv_del(metadata_key)
except Exception as e:
# Ignore errors during cleanup (e.g., key already deleted)
logger.warning(
f"Failed to delete rendezvous key '{metadata_key}' during init: {e}"
)
super().__init__(world_size, rank, group_name)
# Create a subgroup for this logical group. For the default group, use WORLD.
self._is_default_group = group_name == "default"
if self._is_default_group:
self._pg = dist.group.WORLD
else:
# All ranks participate in this subgroup with global ranks [0..world_size-1].
ranks = list(range(world_size))
self._pg = dist.new_group(ranks=ranks, backend="gloo")
# Compatibility shim for legacy tests expecting a pygloo context with getTimeout().
# Store the rendezvous timeout in milliseconds, defaulting to 30000 if unspecified.
class _GlooCompatContext:
def __init__(self, timeout_ms: int):
self._timeout_ms = timeout_ms
def getTimeout(self) -> int:
return self._timeout_ms
self._gloo_context = _GlooCompatContext(
gloo_timeout if gloo_timeout is not None else 30000
)
def _rendezvous(
self, group_name: str, rank: int, gloo_timeout: Optional[int]
) -> None:
"""Rendezvous: ensure a MASTER_ADDR:MASTER_PORT is published in internal_kv.
Rank 0 publishes the address and port, other ranks wait for it.
"""
metadata_key = get_master_address_metadata_key(group_name)
if rank == 0:
addr, port = get_address_and_port()
internal_kv._internal_kv_put(metadata_key, f"{addr}:{port}")
else:
# Wait until rank 0 publishes the metadata or timeout.
deadline_s = time.time() + (gloo_timeout / 1000.0 if gloo_timeout else 30.0)
while True:
meta = internal_kv._internal_kv_get(metadata_key)
if meta is not None:
break
if time.time() > deadline_s:
raise TimeoutError(
f"Timed out waiting for GLOO rendezvous metadata for group '{group_name}'."
)
time.sleep(0.05)
def destroy_group(self):
"""GC the communicators."""
# Destroy only the subgroup for non-default groups. Allow default to be torn down explicitly.
if self._is_default_group:
# Destroy default process group to allow re-init in tests that recreate the same group.
dist.destroy_process_group()
else:
# Destroy just this subgroup.
if self._pg is not None:
dist.destroy_process_group(self._pg)
@classmethod
def backend(cls):
"""The backend of this collective group."""
return Backend.GLOO
@classmethod
def check_backend_availability(cls) -> bool:
return _TORCH_DISTRIBUTED_AVAILABLE
def _check_tensor_input(self, tensor: List["torch.Tensor"]) -> "torch.Tensor":
"""ray.util.collective wraps tensor arguments in a list.
Accept a single torch.Tensor or numpy.ndarray and unwrap/convert it.
"""
assert isinstance(tensor, list) and len(tensor) == 1
t = tensor[0]
if isinstance(t, torch.Tensor):
return t
if isinstance(t, np.ndarray):
return torch.from_numpy(t)
raise ValueError(
f"torch_gloo group only accepts torch.Tensor or numpy.ndarray, received {type(t)}"
)
def _check_tensor_list_input(
self, tensor_list: List[List["torch.Tensor"]]
) -> List["torch.Tensor"]:
"""ray.util.collective wraps tensor arguments in a list.
Accept a single list containing torch.Tensors or numpy.ndarrays and
unwrap/convert items as needed.
"""
assert isinstance(tensor_list, list) and len(tensor_list) == 1
tensor_list = tensor_list[0]
converted_tensor_list = []
for tensor in tensor_list:
if isinstance(tensor, np.ndarray):
tensor = torch.from_numpy(tensor)
converted_tensor_list.append(tensor)
elif isinstance(tensor, torch.Tensor):
converted_tensor_list.append(tensor)
else:
raise ValueError(
f"torch_gloo group only accepts torch.Tensor or numpy.ndarray types, received tensor list with value {tensor}"
)
return converted_tensor_list
def allreduce(
self,
tensor: List["torch.Tensor"],
allreduce_options: Optional[AllReduceOptions] = None,
) -> None:
if allreduce_options is None:
allreduce_options = AllReduceOptions()
tensor = self._check_tensor_input(tensor)
torch_reduce_op = TORCH_REDUCE_OP_MAP[allreduce_options.reduceOp]
dist.all_reduce(tensor, op=torch_reduce_op, group=self._pg)
def barrier(self, barrier_options=BarrierOptions()) -> None:
dist.barrier(group=self._pg)
def reduce(
self,
tensor: List["torch.Tensor"],
reduce_options: Optional[ReduceOptions] = None,
) -> None:
if reduce_options is None:
reduce_options = ReduceOptions()
t = self._check_tensor_input(tensor)
torch_reduce_op = TORCH_REDUCE_OP_MAP[reduce_options.reduceOp]
# Avoid mutating non-root ranks' user tensors to match util.collective semantics.
if self._rank == reduce_options.root_rank:
dist.reduce(
t, dst=reduce_options.root_rank, op=torch_reduce_op, group=self._pg
)
else:
tmp = t.detach().clone()
dist.reduce(
tmp, dst=reduce_options.root_rank, op=torch_reduce_op, group=self._pg
)
def allgather(
self,
tensor_list: List[List["torch.Tensor"]],
tensor: List["torch.Tensor"],
allgather_options: Optional[AllGatherOptions] = None,
) -> None:
if allgather_options is None:
allgather_options = AllGatherOptions()
tensor_list = self._check_tensor_list_input(tensor_list)
tensor = self._check_tensor_input(tensor)
dist.all_gather(tensor_list, tensor, group=self._pg)
def broadcast(
self, tensor: List["torch.Tensor"], broadcast_options=BroadcastOptions()
) -> None:
tensor = self._check_tensor_input(tensor)
dist.broadcast(tensor, src=broadcast_options.root_rank, group=self._pg)
def reducescatter(
self,
output_tensor: List["torch.Tensor"],
tensor_list: List[List["torch.Tensor"]],
reducescatter_options: Optional[ReduceScatterOptions] = None,
) -> None:
if reducescatter_options is None:
reducescatter_options = ReduceScatterOptions()
tensor_list = self._check_tensor_list_input(tensor_list)
output_tensor = self._check_tensor_input(output_tensor)
if output_tensor.shape != tensor_list[self._rank].shape:
raise ValueError(
"Output tensor has wrong shape {output_tensor.shape}, expected {tensor_list[self._rank].shape}"
)
torch_reduce_op = TORCH_REDUCE_OP_MAP[reducescatter_options.reduceOp]
# torch.distributed gloo doesn't support reducescatter. Implement a
# simple version using allreduce.
for tensor in tensor_list:
dist.all_reduce(tensor, op=torch_reduce_op, group=self._pg)
if output_tensor.data_ptr() != tensor_list[self._rank].data_ptr():
output_tensor.copy_(tensor_list[self._rank])
def send(self, tensor: List["torch.Tensor"], send_options: SendOptions) -> None:
tensor = self._check_tensor_input(tensor)
dist.send(tensor, dst=send_options.dst_rank)
def recv(self, tensor: List["torch.Tensor"], recv_options: RecvOptions) -> None:
tensor = self._check_tensor_input(tensor)
dist.recv(tensor, src=recv_options.src_rank)
python/ray/util/collective/const.py
+35
View File
@@ -0,0 +1,35 @@
"""
Constants.
Contains constants used to setup collective groups.
"""
import hashlib
import os
from enum import Enum, auto
def get_store_name(group_name: str):
"""Generate the unique name for the NCCLUniqueID store (named actor).
Args:
group_name: unique user name for the store.
Returns:
SHA256-hexlified name for the store.
"""
if not group_name:
raise ValueError("group_name is None.")
hexlified_name = hashlib.sha256(group_name.encode()).hexdigest()
return hexlified_name
class ENV(Enum):
"""ray.util.collective environment variables."""
NCCL_USE_MULTISTREAM = auto(), lambda v: (v or "True") == "True"
@property
def val(self):
"""Return the output of the lambda against the system's env value."""
_, default_fn = self.value
return default_fn(os.getenv(self.name))
python/ray/util/collective/examples/mock_internal_kv_example.py
+529
View File
@@ -0,0 +1,529 @@
import time
from typing import TYPE_CHECKING
import numpy as np
import ray
import ray.experimental.internal_kv as internal_kv
from ray.util.collective import (
allreduce,
broadcast,
create_collective_group,
init_collective_group,
)
from ray.util.collective.backend_registry import (
_global_registry,
register_collective_backend,
)
from ray.util.collective.collective_group.base_collective_group import BaseGroup
from ray.util.collective.types import (
AllGatherOptions,
AllReduceOptions,
BarrierOptions,
BroadcastOptions,
RecvOptions,
ReduceOp,
ReduceOptions,
ReduceScatterOptions,
SendOptions,
)
if TYPE_CHECKING:
pass
def _unregister_collective_backend(name: str) -> None:
"""Helper function to unregister a backend for testing purposes."""
upper_name = name.upper()
if upper_name in _global_registry._map:
del _global_registry._map[upper_name]
def get_data_key(group_name: str, rank: int, op_name: str):
return f"collective_mock_{group_name}_{op_name}_rank_{rank}"
def get_barrier_key(group_name: str, barrier_id: int):
return f"collective_mock_{group_name}_barrier_{barrier_id}"
class MockInternalKVGroup(BaseGroup):
def __init__(self, world_size: int, rank: int, group_name: str):
super().__init__(world_size, rank, group_name)
self._barrier_counter = 0
@classmethod
def backend(cls):
return "MOCK"
@classmethod
def check_backend_availability(cls) -> bool:
return True
def _check_tensor_input(self, tensor):
assert isinstance(tensor, list) and len(tensor) == 1
t = tensor[0]
if isinstance(t, np.ndarray):
return t
try:
import torch
if isinstance(t, torch.Tensor):
return t
except ImportError:
pass
raise ValueError(
f"MockInternalKVGroup only only accepts numpy.ndarray or torch.Tensor, received {type(t)}"
)
def _serialize_tensor(self, tensor):
if isinstance(tensor, np.ndarray):
return tensor.tobytes(), tensor.shape, tensor.dtype
try:
import torch
if isinstance(tensor, torch.Tensor):
return (
tensor.cpu().numpy().tobytes(),
tensor.shape,
tensor.cpu().numpy().dtype,
)
except ImportError:
pass
raise ValueError(f"Unsupported tensor type: {type(tensor)}")
def _deserialize_tensor(self, data: bytes, shape, dtype, target_tensor):
if isinstance(target_tensor, np.ndarray):
np_array = np.frombuffer(data, dtype=dtype).reshape(shape)
target_tensor[:] = np_array
else:
try:
import torch
if isinstance(target_tensor, torch.Tensor):
np_array = np.frombuffer(data, dtype=dtype).reshape(shape)
target_tensor.copy_(torch.from_numpy(np_array))
except ImportError:
pass
def broadcast(self, tensor, broadcast_options=BroadcastOptions()):
tensor = self._check_tensor_input(tensor)
root_rank = broadcast_options.root_rank
data_key = get_data_key(self._group_name, root_rank, "broadcast")
if self._rank == root_rank:
data, shape, dtype = self._serialize_tensor(tensor)
internal_kv._internal_kv_put(data_key, data)
internal_kv._internal_kv_put(f"{data_key}_shape", str(shape))
internal_kv._internal_kv_put(f"{data_key}_dtype", dtype.name)
else:
deadline_s = time.time() + 30.0
while True:
data = internal_kv._internal_kv_get(data_key)
if data is not None:
break
if time.time() > deadline_s:
raise TimeoutError(
f"Timed out waiting for broadcast data from rank {root_rank}"
)
time.sleep(0.01)
deadline_s = time.time() + 30.0
while True:
shape_data = internal_kv._internal_kv_get(f"{data_key}_shape")
dtype_data = internal_kv._internal_kv_get(f"{data_key}_dtype")
if shape_data is not None and dtype_data is not None:
break
if time.time() > deadline_s:
raise TimeoutError(
f"Timed out waiting for broadcast metadata from rank {root_rank}"
)
time.sleep(0.01)
shape_str = shape_data.decode()
shape = eval(shape_str)
dtype_name = dtype_data.decode()
dtype = np.dtype(dtype_name)
self._deserialize_tensor(data, shape, dtype, tensor)
def allreduce(self, tensor, allreduce_options=AllReduceOptions()):
tensor = self._check_tensor_input(tensor)
reduce_op = allreduce_options.reduceOp
data_key = get_data_key(self._group_name, self._rank, "allreduce")
done_key = (
f"collective_mock_{self._group_name}_allreduce_done_rank_{self._rank}"
)
data, shape, dtype = self._serialize_tensor(tensor)
internal_kv._internal_kv_put(data_key, data)
internal_kv._internal_kv_put(f"{data_key}_shape", str(shape))
internal_kv._internal_kv_put(f"{data_key}_dtype", dtype.name)
internal_kv._internal_kv_put(done_key, b"1")
deadline_s = time.time() + 30.0
while True:
all_done = True
for r in range(self._world_size):
key = f"collective_mock_{self._group_name}_allreduce_done_rank_{r}"
if internal_kv._internal_kv_get(key) is None:
all_done = False
break
if all_done:
break
if time.time() > deadline_s:
raise TimeoutError(
"Timed out waiting for allreduce data from all ranks"
)
time.sleep(0.01)
result = None
for r in range(self._world_size):
rank_data_key = get_data_key(self._group_name, r, "allreduce")
rank_data = internal_kv._internal_kv_get(rank_data_key)
rank_shape_data = internal_kv._internal_kv_get(f"{rank_data_key}_shape")
rank_dtype_data = internal_kv._internal_kv_get(f"{rank_data_key}_dtype")
rank_shape_str = rank_shape_data.decode()
rank_shape = eval(rank_shape_str)
rank_dtype_name = rank_dtype_data.decode()
rank_dtype = np.dtype(rank_dtype_name)
if isinstance(tensor, np.ndarray):
rank_tensor = np.frombuffer(rank_data, dtype=rank_dtype).reshape(
rank_shape
)
else:
import torch
rank_np = np.frombuffer(rank_data, dtype=rank_dtype).reshape(rank_shape)
rank_tensor = torch.from_numpy(rank_np)
if result is None:
result = (
rank_tensor.copy()
if isinstance(rank_tensor, np.ndarray)
else rank_tensor.clone()
)
else:
if reduce_op == ReduceOp.SUM:
result += rank_tensor
elif reduce_op == ReduceOp.PRODUCT:
result *= rank_tensor
elif reduce_op == ReduceOp.MAX:
if isinstance(result, np.ndarray):
result = np.maximum(result, rank_tensor)
else:
import torch
result = torch.maximum(result, rank_tensor)
elif reduce_op == ReduceOp.MIN:
if isinstance(result, np.ndarray):
result = np.minimum(result, rank_tensor)
else:
import torch
result = torch.minimum(result, rank_tensor)
if isinstance(tensor, np.ndarray):
tensor[:] = result
else:
import torch
if isinstance(result, np.ndarray):
tensor.copy_(torch.from_numpy(result))
else:
tensor.copy_(result)
def barrier(self, barrier_options=BarrierOptions()):
barrier_id = self._barrier_counter
barrier_key = get_barrier_key(self._group_name, barrier_id)
rank_key = f"{barrier_key}_rank_{self._rank}"
internal_kv._internal_kv_put(rank_key, b"1")
deadline_s = time.time() + 30.0
while True:
all_arrived = True
for r in range(self._world_size):
key = f"{barrier_key}_rank_{r}"
if internal_kv._internal_kv_get(key) is None:
all_arrived = False
break
if all_arrived:
break
if time.time() > deadline_s:
raise TimeoutError("Timed out waiting for barrier")
time.sleep(0.01)
self._barrier_counter += 1
def reduce(self, tensor, reduce_options=ReduceOptions()):
raise NotImplementedError("reduce is not implemented in MockInternalKVGroup")
def allgather(self, tensor_list, tensor, allgather_options=AllGatherOptions()):
raise NotImplementedError("allgather is not implemented in MockInternalKVGroup")
def reducescatter(
self, tensor, tensor_list, reducescatter_options=ReduceScatterOptions()
):
raise NotImplementedError(
"reducescatter is not implemented in MockInternalKVGroup"
)
def send(self, tensor, send_options: SendOptions):
raise NotImplementedError("send is not implemented in MockInternalKVGroup")
def recv(self, tensor, recv_options: RecvOptions):
raise NotImplementedError("recv is not implemented in MockInternalKVGroup")
def test_mock_backend_create_group():
"""Test using create_collective_group (driver-managed approach).
In this approach:
- Driver calls create_collective_group() to declare the group
- Workers only need to register the backend
- Workers do NOT call init_collective_group()
- The group is automatically initialized when workers call collective ops
"""
ray.init()
register_collective_backend("MOCK", MockInternalKVGroup)
@ray.remote
class Worker:
def __init__(self, rank):
self.rank = rank
def setup(self):
from ray.util.collective.backend_registry import register_collective_backend
register_collective_backend("MOCK", MockInternalKVGroup)
def broadcast_test(self):
if self.rank == 0:
tensor = np.array([42.0, 43.0, 44.0], dtype=np.float32)
else:
tensor = np.array([0.0, 0.0, 0.0], dtype=np.float32)
broadcast(tensor, src_rank=0)
return tensor.tolist()
def allreduce_test(self):
tensor = np.array([float(self.rank + 1)], dtype=np.float32)
allreduce(tensor, op=ReduceOp.SUM)
return tensor.item()
actors = [Worker.remote(rank=i) for i in range(3)]
create_collective_group(
actors=actors,
world_size=3,
ranks=[0, 1, 2],
backend="MOCK",
group_name="default",
)
ray.get([a.setup.remote() for a in actors])
results = ray.get([a.broadcast_test.remote() for a in actors])
expected = [[42.0, 43.0, 44.0]] * 3
if results == expected:
print("Broadcast test passed!")
else:
print(f"Broadcast test failed! Expected {expected}, got {results}")
results = ray.get([a.allreduce_test.remote() for a in actors])
if results == [6.0, 6.0, 6.0]:
print("AllReduce test passed!")
else:
print(f"AllReduce test failed! Expected [6.0, 6.0, 6.0], got {results}")
ray.shutdown()
_unregister_collective_backend("MOCK")
print("test_mock_backend_create_group completed!")
def test_mock_backend_init_group():
"""Test using init_collective_group (worker-managed approach).
In this approach:
- Workers call init_collective_group() inside their setup method
- Driver does NOT call create_collective_group()
- Each worker explicitly initializes its own group membership
"""
ray.init()
@ray.remote
class Worker:
def __init__(self, rank):
self.rank = rank
def setup(self, world_size):
from ray.util.collective.backend_registry import register_collective_backend
register_collective_backend("MOCK", MockInternalKVGroup)
init_collective_group(
world_size=world_size,
rank=self.rank,
backend="MOCK",
group_name="default",
)
def broadcast_test(self):
if self.rank == 0:
tensor = np.array([42.0, 43.0, 44.0], dtype=np.float32)
else:
tensor = np.array([0.0, 0.0, 0.0], dtype=np.float32)
broadcast(tensor, src_rank=0)
return tensor.tolist()
def allreduce_test(self):
tensor = np.array([float(self.rank + 1)], dtype=np.float32)
allreduce(tensor, op=ReduceOp.SUM)
return tensor.item()
actors = [Worker.remote(rank=i) for i in range(3)]
# Do NOT call create_collective_group here
ray.get([a.setup.remote(3) for a in actors])
results = ray.get([a.broadcast_test.remote() for a in actors])
expected = [[42.0, 43.0, 44.0]] * 3
if results == expected:
print("Broadcast test passed!")
else:
print(f"Broadcast test failed! Expected {expected}, got {results}")
results = ray.get([a.allreduce_test.remote() for a in actors])
if results == [6.0, 6.0, 6.0]:
print("AllReduce test passed!")
else:
print(f"AllReduce test failed! Expected [6.0, 6.0, 6.0], got {results}")
ray.shutdown()
_unregister_collective_backend("MOCK")
print("test_mock_backend_init_group completed!")
def test_mock_backend_worker_not_registered():
"""Test error handling when backend is not registered in worker.
This test uses create_collective_group (driver-managed approach).
The driver registers the backend, but workers do not.
When workers try to call collective ops, they should fail.
Note: We use world_size=1 to avoid "Unhandled error" messages
from multiple workers failing simultaneously.
"""
ray.init()
register_collective_backend("MOCK", MockInternalKVGroup)
@ray.remote
class Worker:
def __init__(self, rank):
self.rank = rank
def broadcast_test(self):
tensor = np.array([0.0, 0.0, 0.0], dtype=np.float32)
broadcast(tensor, src_rank=0)
return tensor.tolist()
# Use single actor to avoid multiple "Unhandled error" messages
actors = [Worker.remote(rank=0)]
create_collective_group(
actors=actors,
world_size=1,
ranks=[0],
backend="MOCK",
group_name="default",
)
test_passed = False
try:
ray.get([a.broadcast_test.remote() for a in actors])
print("ERROR: Should have raised an exception for missing registration!")
except Exception as e:
if "not registered" in str(e) or "not initialized" in str(e):
print(
"Test passed! Correctly raised error for missing worker registration."
)
test_passed = True
else:
print(f"ERROR: Unexpected error: {e}")
ray.shutdown()
_unregister_collective_backend("MOCK")
if not test_passed:
print("Test failed!")
def test_mock_backend_driver_not_registered():
"""Test error handling when backend is not registered on driver.
This test uses create_collective_group, but the driver doesn't
register the backend first, so it should fail immediately.
"""
ray.init()
@ray.remote
class Worker:
def __init__(self, rank):
self.rank = rank
actors = [Worker.remote(rank=i) for i in range(2)]
try:
create_collective_group(
actors=actors,
world_size=2,
ranks=[0, 1],
backend="MOCK",
group_name="default",
)
print("ERROR: Should have raised an exception for missing registration!")
except Exception as e:
if "not registered" in str(e):
print(
"Test passed! Correctly raised error for missing driver registration."
)
else:
print(f"ERROR: Unexpected error: {e}")
ray.shutdown()
_unregister_collective_backend("MOCK")
if __name__ == "__main__":
print("=" * 60)
print("Test 1: create_collective_group approach (driver-managed)")
print("=" * 60)
test_mock_backend_create_group()
print("\n" + "=" * 60)
print("Test 2: init_collective_group approach (worker-managed)")
print("=" * 60)
test_mock_backend_init_group()
print("\n" + "=" * 60)
print("Test 3: Error handling - worker not registered")
print("=" * 60)
test_mock_backend_worker_not_registered()
print("\n" + "=" * 60)
print("Test 4: Error handling - driver not registered")
print("=" * 60)
test_mock_backend_driver_not_registered()
python/ray/util/collective/examples/nccl_allreduce_example.py
+37
View File
@@ -0,0 +1,37 @@
import cupy as cp
import ray
import ray.util.collective as collective
@ray.remote(num_gpus=1)
class Worker:
def __init__(self):
self.send = cp.ones((4,), dtype=cp.float32)
self.recv = cp.zeros((4,), dtype=cp.float32)
def setup(self, world_size, rank):
collective.init_collective_group(world_size, rank, "nccl", "default")
return True
def compute(self):
collective.allreduce(self.send, "default")
return self.send
def destroy(self):
collective.destroy_group()
if __name__ == "__main__":
send = cp.ones((4,), dtype=cp.float32)
ray.init(num_gpus=2)
num_workers = 2
workers = []
init_rets = []
for i in range(num_workers):
w = Worker.remote()
workers.append(w)
init_rets.append(w.setup.remote(num_workers, i))
_ = ray.get(init_rets)
results = ray.get([w.compute.remote() for w in workers])
ray.shutdown()
python/ray/util/collective/examples/nccl_allreduce_example_declare_collective_group.py
+33
View File
@@ -0,0 +1,33 @@
import cupy as cp
import ray
import ray.util.collective as collective
@ray.remote(num_gpus=1)
class Worker:
def __init__(self):
self.send = cp.ones((4,), dtype=cp.float32)
def compute(self):
collective.allreduce(self.send, "177")
return self.send
if __name__ == "__main__":
ray.init(num_gpus=2)
num_workers = 2
workers = []
for i in range(num_workers):
w = Worker.remote()
workers.append(w)
_options = {
"group_name": "177",
"world_size": 2,
"ranks": [0, 1],
"backend": "nccl",
}
collective.create_collective_group(workers, **_options)
results = ray.get([w.compute.remote() for w in workers])
ray.shutdown()
python/ray/util/collective/examples/nccl_allreduce_multigpu_example.py
+43
View File
@@ -0,0 +1,43 @@
import cupy as cp
from cupy.cuda import Device
import ray
import ray.util.collective as collective
@ray.remote(num_gpus=2)
class Worker:
def __init__(self):
with Device(0):
self.send1 = cp.ones((4,), dtype=cp.float32)
with Device(1):
self.send2 = cp.ones((4,), dtype=cp.float32) * 2
self.recv = cp.zeros((4,), dtype=cp.float32)
def setup(self, world_size, rank):
collective.init_collective_group(world_size, rank, "nccl", "177")
return True
def compute(self):
collective.allreduce_multigpu([self.send1, self.send2], "177")
return [self.send1, self.send2], self.send1.device, self.send2.device
def destroy(self):
collective.destroy_collective_group("177")
if __name__ == "__main__":
ray.init(address="auto")
num_workers = 2
workers = []
init_rets = []
for i in range(num_workers):
w = Worker.remote()
workers.append(w)
init_rets.append(w.setup.remote(num_workers, i))
a = ray.get(init_rets)
results = ray.get([w.compute.remote() for w in workers])
print(results)
ray.get([w.destroy.remote() for w in workers])
ray.shutdown()
python/ray/util/collective/examples/nccl_p2p_example_multigpu.py
+53
View File
@@ -0,0 +1,53 @@
import cupy as cp
from cupy.cuda import Device
import ray
import ray.util.collective as collective
@ray.remote(num_gpus=2)
class Worker:
def __init__(self):
with Device(0):
self.send1 = cp.ones((4,), dtype=cp.float32)
with Device(1):
self.send2 = cp.ones((4,), dtype=cp.float32) * 2
with Device(0):
self.recv1 = cp.zeros((4,), dtype=cp.float32)
with Device(1):
self.recv2 = cp.zeros((4,), dtype=cp.float32)
self.rank = -1
def setup(self, world_size, rank):
self.rank = rank
collective.init_collective_group(world_size, rank, "nccl", "8")
return True
def compute(self):
if self.rank == 0:
with Device(0):
collective.send_multigpu(self.send1 * 2, 1, 1, "8")
else:
# with Device(1):
collective.recv_multigpu(self.recv2, 0, 0, "8")
return self.recv2
def destroy(self):
collective.destroy_collective_group("8")
if __name__ == "__main__":
ray.init(address="auto")
num_workers = 2
workers = []
init_rets = []
for i in range(num_workers):
w = Worker.remote()
workers.append(w)
init_rets.append(w.setup.remote(num_workers, i))
a = ray.get(init_rets)
results = ray.get([w.compute.remote() for w in workers])
print(results)
ray.get([w.destroy.remote() for w in workers])
ray.shutdown()
python/ray/util/collective/requirements.txt
+1
View File
@@ -0,0 +1 @@
cupy-cuda100
python/ray/util/collective/tests/__init__.py
View File
python/ray/util/collective/tests/conftest.py
+106
View File
@@ -0,0 +1,106 @@
"""Some fixtures for collective tests."""
import logging
import pytest
import ray
try:
from ray.util.collective.collective_group.nccl_collective_group import (
_get_comm_key_from_devices,
_get_comm_key_send_recv,
)
except Exception: # Cupy/NCCL may be unavailable on CPU-only setups
_get_comm_key_from_devices = None
_get_comm_key_send_recv = None
from ray.util.collective.const import get_store_name
logger = logging.getLogger(__name__)
logger.setLevel("INFO")
# TODO (Hao): remove this clean_up function as it sometimes crashes Ray.
def clean_up():
# If NCCL helpers are unavailable (e.g., no cupy), skip cleanup.
if _get_comm_key_from_devices is None or _get_comm_key_send_recv is None:
return
group_names = ["default", "test", "123?34!", "default2", "random"]
group_names.extend([str(i) for i in range(10)])
max_world_size = 4
all_keys = []
for name in group_names:
devices = [[0], [0, 1], [1, 0]]
for d in devices:
collective_communicator_key = _get_comm_key_from_devices(d)
all_keys.append(collective_communicator_key + "@" + name)
for i in range(max_world_size):
for j in range(max_world_size):
if i < j:
p2p_communicator_key = _get_comm_key_send_recv(i, 0, j, 0)
all_keys.append(p2p_communicator_key + "@" + name)
for group_key in all_keys:
store_name = get_store_name(group_key)
try:
actor = ray.get_actor(store_name)
except ValueError:
actor = None
if actor:
logger.debug(
"Killing actor with group_key: '{}' and store: '{}'.".format(
group_key, store_name
)
)
ray.kill(actor)
@pytest.fixture
def ray_start_single_node_2_gpus():
# Please start this fixture in a cluster with 2 GPUs.
address_info = ray.init(num_gpus=2)
yield address_info
ray.shutdown()
# Hao: this fixture is a bit tricky.
# I use a bash script to start a ray cluster on
# my own on-premise cluster before run this fixture.
@pytest.fixture
def ray_start_distributed_2_nodes_4_gpus():
# The cluster has a setup of 2 nodes, each node with 2
# GPUs. Each actor will be allocated 1 GPU.
ray.init("auto")
yield
clean_up()
ray.shutdown()
@pytest.fixture
def ray_start_distributed_multigpu_2_nodes_4_gpus():
# The cluster has a setup of 2 nodes, each node with 2
# GPUs. Each actor will be allocated 2 GPUs.
ray.init("auto")
yield
clean_up()
ray.shutdown()
@pytest.fixture
def ray_start_single_node():
address_info = ray.init(num_cpus=8)
yield address_info
ray.shutdown()
@pytest.fixture
def ray_start_distributed_2_nodes():
# The cluster has a setup of 2 nodes.
# no GPUs!
ray.init("auto")
yield
ray.shutdown()
@pytest.fixture
def shutdown_only():
yield None
ray.shutdown()
python/ray/util/collective/tests/cpu_util.py
+137
View File
@@ -0,0 +1,137 @@
import logging
import numpy as np
import torch
import ray
import ray.util.collective as col
from ray.util.collective.types import Backend, ReduceOp
logger = logging.getLogger(__name__)
@ray.remote(num_cpus=1)
class Worker:
def __init__(self):
self.buffer = None
self.list_buffer = None
def init_tensors(self):
self.buffer = np.ones((10,), dtype=np.float32)
self.list_buffer = [np.ones((10,), dtype=np.float32) for _ in range(2)]
return True
def init_group(self, world_size, rank, backend=Backend.NCCL, group_name="default"):
col.init_collective_group(world_size, rank, backend, group_name)
return True
def set_buffer(self, data):
self.buffer = data
return self.buffer
def get_buffer(self):
return self.buffer
def set_list_buffer(self, list_of_arrays, copy=False):
if copy:
copy_list = []
for tensor in list_of_arrays:
if isinstance(tensor, np.ndarray):
copy_list.append(tensor.copy())
elif isinstance(tensor, torch.Tensor):
copy_list.append(tensor.clone().detach())
self.list_buffer = copy_list
else:
self.list_buffer = list_of_arrays
return self.list_buffer
def do_allreduce(self, group_name="default", op=ReduceOp.SUM):
col.allreduce(self.buffer, group_name, op)
return self.buffer
def do_reduce(self, group_name="default", dst_rank=0, op=ReduceOp.SUM):
col.reduce(self.buffer, dst_rank, group_name, op)
return self.buffer
def do_broadcast(self, group_name="default", src_rank=0):
col.broadcast(self.buffer, src_rank, group_name)
return self.buffer
def do_allgather(self, group_name="default"):
col.allgather(self.list_buffer, self.buffer, group_name)
return self.list_buffer
def do_reducescatter(self, group_name="default", op=ReduceOp.SUM):
col.reducescatter(self.buffer, self.list_buffer, group_name, op)
return self.buffer
def do_send(self, group_name="default", dst_rank=0):
col.send(self.buffer, dst_rank, group_name)
return self.buffer
def do_recv(self, group_name="default", src_rank=0):
col.recv(self.buffer, src_rank, group_name)
return self.buffer
def destroy_group(self, group_name="default"):
col.destroy_collective_group(group_name)
return True
def report_rank(self, group_name="default"):
rank = col.get_rank(group_name)
return rank
def report_world_size(self, group_name="default"):
ws = col.get_collective_group_size(group_name)
return ws
def report_nccl_availability(self):
avail = col.nccl_available()
return avail
def report_gloo_availability(self):
avail = col.gloo_available()
return avail
def report_is_group_initialized(self, group_name="default"):
is_init = col.is_group_initialized(group_name)
return is_init
def create_collective_workers(num_workers=2, group_name="default", backend="nccl"):
actors = [None] * num_workers
for i in range(num_workers):
actor = Worker.remote()
ray.get([actor.init_tensors.remote()])
actors[i] = actor
world_size = num_workers
init_results = ray.get(
[
actor.init_group.remote(world_size, i, backend, group_name)
for i, actor in enumerate(actors)
]
)
return actors, init_results
def init_tensors_for_gather_scatter(
actors, array_size=10, dtype=np.float32, tensor_backend="numpy"
):
world_size = len(actors)
for i, a in enumerate(actors):
if tensor_backend == "numpy":
t = np.ones(array_size, dtype=dtype) * (i + 1)
elif tensor_backend == "torch":
t = torch.ones(array_size, dtype=torch.float32) * (i + 1)
else:
raise RuntimeError("Unsupported tensor backend.")
ray.get([a.set_buffer.remote(t)])
if tensor_backend == "numpy":
list_buffer = [np.ones(array_size, dtype=dtype) for _ in range(world_size)]
elif tensor_backend == "torch":
list_buffer = [
torch.ones(array_size, dtype=torch.float32) for _ in range(world_size)
]
else:
raise RuntimeError("Unsupported tensor backend.")
ray.get([a.set_list_buffer.remote(list_buffer, copy=True) for a in actors])
python/ray/util/collective/tests/distributed_cpu_tests/__init__.py
View File
python/ray/util/collective/tests/distributed_cpu_tests/test_distributed_allgather.py
+143
View File
@@ -0,0 +1,143 @@
"""Test the allgather API on a distributed Ray cluster."""
import numpy as np
import pytest
import torch
import ray
from ray.util.collective.tests.cpu_util import (
create_collective_workers,
init_tensors_for_gather_scatter,
)
from ray.util.collective.types import Backend
@pytest.mark.parametrize("backend", [Backend.GLOO])
@pytest.mark.parametrize("tensor_backend", ["numpy", "torch"])
@pytest.mark.parametrize(
"array_size", [2, 2**5, 2**10, 2**15, 2**20, [2, 2], [5, 5, 5]]
)
def test_allgather_different_array_size(
ray_start_distributed_2_nodes, array_size, tensor_backend, backend
):
world_size = 8
actors, _ = create_collective_workers(world_size, backend=backend)
init_tensors_for_gather_scatter(
actors, array_size=array_size, tensor_backend=tensor_backend
)
results = ray.get([a.do_allgather.remote() for a in actors])
for i in range(world_size):
for j in range(world_size):
if tensor_backend == "numpy":
assert (
results[i][j] == np.ones(array_size, dtype=np.float32) * (j + 1)
).all()
else:
assert (
results[i][j]
== torch.ones(array_size, dtype=torch.float32) * (j + 1)
).all()
@pytest.mark.parametrize("backend", [Backend.GLOO])
@pytest.mark.parametrize("dtype", [np.uint8, np.float16, np.float32, np.float64])
def test_allgather_different_dtype(ray_start_distributed_2_nodes, dtype, backend):
world_size = 8
actors, _ = create_collective_workers(world_size, backend=backend)
init_tensors_for_gather_scatter(actors, dtype=dtype)
results = ray.get([a.do_allgather.remote() for a in actors])
for i in range(world_size):
for j in range(world_size):
assert (results[i][j] == np.ones(10, dtype=dtype) * (j + 1)).all()
@pytest.mark.parametrize("backend", [Backend.GLOO])
@pytest.mark.parametrize("length", [0, 1, 3, 4, 7, 8])
def test_unmatched_tensor_list_length(ray_start_distributed_2_nodes, length, backend):
world_size = 8
actors, _ = create_collective_workers(world_size, backend=backend)
list_buffer = [np.ones(10, dtype=np.float32) for _ in range(length)]
ray.wait([a.set_list_buffer.remote(list_buffer, copy=True) for a in actors])
if length != world_size:
with pytest.raises(RuntimeError):
ray.get([a.do_allgather.remote() for a in actors])
else:
ray.get([a.do_allgather.remote() for a in actors])
@pytest.mark.parametrize("backend", [Backend.GLOO])
@pytest.mark.parametrize("shape", [10, 20, [4, 5], [1, 3, 5, 7]])
def test_unmatched_tensor_shape(ray_start_distributed_2_nodes, shape, backend):
world_size = 8
actors, _ = create_collective_workers(world_size, backend=backend)
init_tensors_for_gather_scatter(actors, array_size=10)
list_buffer = [np.ones(shape, dtype=np.float32) for _ in range(world_size)]
ray.get([a.set_list_buffer.remote(list_buffer, copy=True) for a in actors])
if shape != 10:
with pytest.raises(RuntimeError):
ray.get([a.do_allgather.remote() for a in actors])
else:
ray.get([a.do_allgather.remote() for a in actors])
@pytest.mark.parametrize("backend", [Backend.GLOO])
def test_allgather_torch_numpy(ray_start_distributed_2_nodes, backend):
world_size = 8
shape = [10, 10]
actors, _ = create_collective_workers(world_size, backend=backend)
# tensor is pytorch, list is numpy
for i, a in enumerate(actors):
t = torch.ones(shape, dtype=torch.float32) * (i + 1)
ray.wait([a.set_buffer.remote(t)])
list_buffer = [np.ones(shape, dtype=np.float32) for _ in range(world_size)]
ray.wait([a.set_list_buffer.remote(list_buffer, copy=True)])
results = ray.get([a.do_allgather.remote() for a in actors])
for i in range(world_size):
for j in range(world_size):
assert (results[i][j] == np.ones(shape, dtype=np.float32) * (j + 1)).all()
# tensor is numpy, list is pytorch
for i, a in enumerate(actors):
t = np.ones(shape, dtype=np.float32) * (i + 1)
ray.wait([a.set_buffer.remote(t)])
list_buffer = [
torch.ones(shape, dtype=torch.float32) for _ in range(world_size)
]
ray.wait([a.set_list_buffer.remote(list_buffer)])
results = ray.get([a.do_allgather.remote() for a in actors])
for i in range(world_size):
for j in range(world_size):
assert (
results[i][j] == torch.ones(shape, dtype=torch.float32) * (j + 1)
).all()
# some tensors in the list are pytorch, some are numpy
for i, a in enumerate(actors):
t = np.ones(shape, dtype=np.float32) * (i + 1)
ray.wait([a.set_buffer.remote(t)])
list_buffer = []
for j in range(world_size):
if j % 2 == 0:
list_buffer.append(torch.ones(shape, dtype=torch.float32))
else:
list_buffer.append(np.ones(shape, dtype=np.float32))
ray.wait([a.set_list_buffer.remote(list_buffer, copy=True)])
results = ray.get([a.do_allgather.remote() for a in actors])
for i in range(world_size):
for j in range(world_size):
if j % 2 == 0:
assert (
results[i][j] == torch.ones(shape, dtype=torch.float32) * (j + 1)
).all()
else:
assert (
results[i][j] == np.ones(shape, dtype=np.float32) * (j + 1)
).all()
if __name__ == "__main__":
import sys
import pytest
sys.exit(pytest.main(["-v", "-x", __file__]))
python/ray/util/collective/tests/distributed_cpu_tests/test_distributed_allreduce.py
+179
View File
@@ -0,0 +1,179 @@
"""Test the collective allreduice API on a distributed Ray cluster."""
import numpy as np
import pytest
import torch
import ray
from ray.util.collective.tests.cpu_util import create_collective_workers
from ray.util.collective.types import Backend, ReduceOp
@pytest.mark.parametrize("backend", [Backend.GLOO])
@pytest.mark.parametrize("group_name", ["default", "test", "123?34!"])
@pytest.mark.parametrize("world_size", [5, 6, 7, 8])
def test_allreduce_different_name(
ray_start_distributed_2_nodes, group_name, world_size, backend
):
actors, _ = create_collective_workers(
num_workers=world_size, group_name=group_name, backend=backend
)
results = ray.get([a.do_allreduce.remote(group_name) for a in actors])
assert (results[0] == np.ones((10,), dtype=np.float32) * world_size).all()
assert (results[1] == np.ones((10,), dtype=np.float32) * world_size).all()
@pytest.mark.parametrize("backend", [Backend.GLOO])
@pytest.mark.parametrize("array_size", [2, 2**5, 2**10, 2**15, 2**20])
def test_allreduce_different_array_size(
ray_start_distributed_2_nodes, array_size, backend
):
world_size = 8
actors, _ = create_collective_workers(world_size, backend=backend)
ray.wait(
[a.set_buffer.remote(np.ones(array_size, dtype=np.float32)) for a in actors]
)
results = ray.get([a.do_allreduce.remote() for a in actors])
assert (results[0] == np.ones((array_size,), dtype=np.float32) * world_size).all()
assert (results[1] == np.ones((array_size,), dtype=np.float32) * world_size).all()
@pytest.mark.parametrize("backend", [Backend.GLOO])
def test_allreduce_destroy(
ray_start_distributed_2_nodes, backend, group_name="default"
):
world_size = 8
actors, _ = create_collective_workers(world_size, backend=backend)
results = ray.get([a.do_allreduce.remote() for a in actors])
assert (results[0] == np.ones((10,), dtype=np.float32) * world_size).all()
assert (results[1] == np.ones((10,), dtype=np.float32) * world_size).all()
# destroy the group and try do work, should fail
ray.get([a.destroy_group.remote() for a in actors])
with pytest.raises(RuntimeError):
results = ray.get([a.do_allreduce.remote() for a in actors])
# reinit the same group and all reduce
ray.get(
[
actor.init_group.remote(world_size, i, backend, group_name)
for i, actor in enumerate(actors)
]
)
results = ray.get([a.do_allreduce.remote() for a in actors])
assert (
results[0] == np.ones((10,), dtype=np.float32) * world_size * world_size
).all()
assert (
results[1] == np.ones((10,), dtype=np.float32) * world_size * world_size
).all()
@pytest.mark.parametrize("backend", [Backend.GLOO])
def test_allreduce_multiple_group(ray_start_distributed_2_nodes, backend, num_groups=5):
world_size = 8
actors, _ = create_collective_workers(world_size, backend=backend)
for group_name in range(1, num_groups):
ray.get(
[
actor.init_group.remote(world_size, i, backend, str(group_name))
for i, actor in enumerate(actors)
]
)
for i in range(num_groups):
group_name = "default" if i == 0 else str(i)
results = ray.get([a.do_allreduce.remote(group_name) for a in actors])
assert (
results[0] == np.ones((10,), dtype=np.float32) * (world_size ** (i + 1))
).all()
@pytest.mark.parametrize("backend", [Backend.GLOO])
def test_allreduce_different_op(ray_start_distributed_2_nodes, backend):
world_size = 8
actors, _ = create_collective_workers(world_size, backend=backend)
# check product
ray.wait(
[
a.set_buffer.remote(np.ones(10, dtype=np.float32) * (i + 2))
for i, a in enumerate(actors)
]
)
results = ray.get([a.do_allreduce.remote(op=ReduceOp.PRODUCT) for a in actors])
product = 1
for i in range(world_size):
product = product * (i + 2)
assert (results[0] == np.ones((10,), dtype=np.float32) * product).all()
assert (results[1] == np.ones((10,), dtype=np.float32) * product).all()
# check min
ray.wait(
[
a.set_buffer.remote(np.ones(10, dtype=np.float32) * (i + 2))
for i, a in enumerate(actors)
]
)
results = ray.get([a.do_allreduce.remote(op=ReduceOp.MIN) for a in actors])
assert (results[0] == np.ones((10,), dtype=np.float32) * 2).all()
assert (results[1] == np.ones((10,), dtype=np.float32) * 2).all()
# check max
ray.wait(
[
a.set_buffer.remote(np.ones(10, dtype=np.float32) * (i + 2))
for i, a in enumerate(actors)
]
)
results = ray.get([a.do_allreduce.remote(op=ReduceOp.MAX) for a in actors])
assert (results[0] == np.ones((10,), dtype=np.float32) * 9).all()
assert (results[1] == np.ones((10,), dtype=np.float32) * 9).all()
@pytest.mark.parametrize("backend", [Backend.GLOO])
@pytest.mark.parametrize("dtype", [np.uint8, np.float16, np.float32, np.float64])
def test_allreduce_different_dtype(ray_start_distributed_2_nodes, dtype, backend):
world_size = 8
actors, _ = create_collective_workers(world_size, backend=backend)
ray.wait([a.set_buffer.remote(np.ones(10, dtype=dtype)) for a in actors])
results = ray.get([a.do_allreduce.remote() for a in actors])
assert (results[0] == np.ones((10,), dtype=dtype) * world_size).all()
assert (results[1] == np.ones((10,), dtype=dtype) * world_size).all()
@pytest.mark.parametrize("backend", [Backend.GLOO])
def test_allreduce_torch_numpy(ray_start_distributed_2_nodes, backend):
# import torch
world_size = 8
actors, _ = create_collective_workers(world_size, backend=backend)
ray.wait(
[
actors[1].set_buffer.remote(
torch.ones(
10,
)
)
]
)
results = ray.get([a.do_allreduce.remote() for a in actors])
assert (results[0] == np.ones((10,)) * world_size).all()
ray.wait(
[
actors[0].set_buffer.remote(
torch.ones(
10,
)
)
]
)
ray.wait([actors[1].set_buffer.remote(np.ones(10, dtype=np.float32))])
results = ray.get([a.do_allreduce.remote() for a in actors])
if __name__ == "__main__":
import sys
import pytest
sys.exit(pytest.main(["-v", "-x", __file__]))
python/ray/util/collective/tests/distributed_cpu_tests/test_distributed_basic_apis.py
+138
View File
@@ -0,0 +1,138 @@
"""Test the collective group APIs."""
from random import shuffle
import pytest
import ray
from ray.util.collective.tests.cpu_util import Worker, create_collective_workers
from ray.util.collective.types import Backend
@pytest.mark.parametrize("backend", [Backend.GLOO])
@pytest.mark.parametrize("world_size", [2, 3, 4])
@pytest.mark.parametrize("group_name", ["default", "test", "123?34!"])
def test_init_two_actors(
ray_start_distributed_2_nodes, world_size, group_name, backend
):
actors, results = create_collective_workers(world_size, group_name, backend=backend)
for i in range(world_size):
assert results[i]
@pytest.mark.parametrize("backend", [Backend.GLOO])
@pytest.mark.parametrize("world_size", [2, 3, 4])
def test_init_multiple_groups(ray_start_distributed_2_nodes, world_size, backend):
num_groups = 5
actors = [Worker.remote() for _ in range(world_size)]
for i in range(num_groups):
group_name = str(i)
init_results = ray.get(
[
actor.init_group.remote(
world_size, i, group_name=group_name, backend=backend
)
for i, actor in enumerate(actors)
]
)
for j in range(world_size):
assert init_results[j]
@pytest.mark.parametrize("backend", [Backend.GLOO])
@pytest.mark.parametrize("world_size", [5, 6, 7, 8])
def test_get_rank(ray_start_distributed_2_nodes, world_size, backend):
actors, _ = create_collective_workers(world_size, backend=backend)
actor0_rank = ray.get(actors[0].report_rank.remote())
assert actor0_rank == 0
actor1_rank = ray.get(actors[1].report_rank.remote())
assert actor1_rank == 1
# create a second group with a different name, and different
# orders of ranks.
new_group_name = "default2"
ranks = list(range(world_size))
shuffle(ranks)
_ = ray.get(
[
actor.init_group.remote(
world_size, ranks[i], group_name=new_group_name, backend=backend
)
for i, actor in enumerate(actors)
]
)
actor0_rank = ray.get(actors[0].report_rank.remote(new_group_name))
assert actor0_rank == ranks[0]
actor1_rank = ray.get(actors[1].report_rank.remote(new_group_name))
assert actor1_rank == ranks[1]
@pytest.mark.parametrize("backend", [Backend.GLOO])
@pytest.mark.parametrize("world_size", [5, 6, 7, 8])
def test_get_world_size(ray_start_distributed_2_nodes, world_size, backend):
actors, _ = create_collective_workers(world_size, backend=backend)
actor0_world_size = ray.get(actors[0].report_world_size.remote())
actor1_world_size = ray.get(actors[1].report_world_size.remote())
assert actor0_world_size == actor1_world_size == world_size
@pytest.mark.parametrize("backend", [Backend.GLOO])
def test_is_group_initialized(ray_start_distributed_2_nodes, backend):
world_size = 8
actors, _ = create_collective_workers(world_size, backend=backend)
# check group is_init
actor0_is_init = ray.get(actors[0].report_is_group_initialized.remote())
assert actor0_is_init
actor0_is_init = ray.get(actors[0].report_is_group_initialized.remote("random"))
assert not actor0_is_init
actor0_is_init = ray.get(actors[0].report_is_group_initialized.remote("123"))
assert not actor0_is_init
actor1_is_init = ray.get(actors[0].report_is_group_initialized.remote())
assert actor1_is_init
actor1_is_init = ray.get(actors[0].report_is_group_initialized.remote("456"))
assert not actor1_is_init
@pytest.mark.parametrize("backend", [Backend.GLOO])
def test_destroy_group(ray_start_distributed_2_nodes, backend):
world_size = 8
actors, _ = create_collective_workers(world_size, backend=backend)
# Now destroy the group at actor0
ray.wait([actors[0].destroy_group.remote()])
actor0_is_init = ray.get(actors[0].report_is_group_initialized.remote())
assert not actor0_is_init
# should go well as the group `random` does not exist at all
ray.wait([actors[0].destroy_group.remote("random")])
actor1_is_init = ray.get(actors[1].report_is_group_initialized.remote())
assert actor1_is_init
ray.wait([actors[1].destroy_group.remote("random")])
actor1_is_init = ray.get(actors[1].report_is_group_initialized.remote())
assert actor1_is_init
ray.wait([actors[1].destroy_group.remote("default")])
actor1_is_init = ray.get(actors[1].report_is_group_initialized.remote())
assert not actor1_is_init
for i in range(2, world_size):
ray.wait([actors[i].destroy_group.remote("default")])
# Now reconstruct the group using the same name
init_results = ray.get(
[
actor.init_group.remote(world_size, i, backend=backend)
for i, actor in enumerate(actors)
]
)
for i in range(world_size):
assert init_results[i]
actor0_is_init = ray.get(actors[0].report_is_group_initialized.remote())
assert actor0_is_init
actor1_is_init = ray.get(actors[0].report_is_group_initialized.remote())
assert actor1_is_init
if __name__ == "__main__":
import sys
import pytest
sys.exit(pytest.main(["-v", "-x", __file__]))
python/ray/util/collective/tests/distributed_cpu_tests/test_distributed_broadcast.py
+96
View File
@@ -0,0 +1,96 @@
"""Test the broadcast API."""
import numpy as np
import pytest
import ray
from ray.util.collective.tests.cpu_util import create_collective_workers
from ray.util.collective.types import Backend
@pytest.mark.parametrize("backend", [Backend.GLOO])
@pytest.mark.parametrize("group_name", ["default", "test", "123?34!"])
@pytest.mark.parametrize("src_rank", [0, 2, 5, 6, 7])
def test_broadcast_different_name(
ray_start_distributed_2_nodes, group_name, src_rank, backend
):
world_size = 8
actors, _ = create_collective_workers(
num_workers=world_size, group_name=group_name, backend=backend
)
ray.wait(
[
a.set_buffer.remote(np.ones((10,), dtype=np.float32) * (i + 2))
for i, a in enumerate(actors)
]
)
results = ray.get(
[
a.do_broadcast.remote(group_name=group_name, src_rank=src_rank)
for a in actors
]
)
for i in range(world_size):
assert (results[i] == np.ones((10,), dtype=np.float32) * (src_rank + 2)).all()
@pytest.mark.parametrize("backend", [Backend.GLOO])
@pytest.mark.parametrize("array_size", [2, 2**5, 2**10, 2**15, 2**20])
@pytest.mark.parametrize("src_rank", [0, 2, 5, 6, 7])
def test_broadcast_different_array_size(
ray_start_distributed_2_nodes, array_size, src_rank, backend
):
world_size = 8
actors, _ = create_collective_workers(world_size, backend=backend)
ray.wait(
[
a.set_buffer.remote(np.ones(array_size, dtype=np.float32) * (i + 2))
for i, a in enumerate(actors)
]
)
results = ray.get([a.do_broadcast.remote(src_rank=src_rank) for a in actors])
for i in range(world_size):
assert (
results[i] == np.ones((array_size,), dtype=np.float32) * (src_rank + 2)
).all()
@pytest.mark.parametrize("backend", [Backend.GLOO])
@pytest.mark.parametrize("src_rank", [0, 2, 5, 6, 7])
def test_broadcast_torch_numpy(ray_start_distributed_2_nodes, src_rank, backend):
import torch
world_size = 8
actors, _ = create_collective_workers(world_size, backend=backend)
ray.wait(
[
actors[1].set_buffer.remote(
torch.ones(
10,
)
* world_size
)
]
)
results = ray.get([a.do_broadcast.remote(src_rank=src_rank) for a in actors])
if src_rank == 0:
assert (results[0] == np.ones((10,))).all()
assert (results[1] == torch.ones((10,))).all()
else:
assert (results[0] == np.ones((10,)) * world_size).all()
assert (results[1] == torch.ones((10,)) * world_size).all()
@pytest.mark.parametrize("backend", [Backend.GLOO])
def test_broadcast_invalid_rank(ray_start_distributed_2_nodes, backend, src_rank=9):
world_size = 8
actors, _ = create_collective_workers(world_size, backend=backend)
with pytest.raises(ValueError):
_ = ray.get([a.do_broadcast.remote(src_rank=src_rank) for a in actors])
if __name__ == "__main__":
import sys
import pytest
sys.exit(pytest.main(["-v", "-x", __file__]))
python/ray/util/collective/tests/distributed_cpu_tests/test_distributed_reduce.py
+147
View File
@@ -0,0 +1,147 @@
"""Test the reduce API."""
import numpy as np
import pytest
import ray
from ray.util.collective.tests.cpu_util import create_collective_workers
from ray.util.collective.types import Backend, ReduceOp
@pytest.mark.parametrize("backend", [Backend.GLOO])
@pytest.mark.parametrize("group_name", ["default", "test", "123?34!"])
@pytest.mark.parametrize("dst_rank", [0, 2, 5, 6, 7])
def test_reduce_different_name(
ray_start_distributed_2_nodes, group_name, backend, dst_rank
):
world_size = 8
actors, _ = create_collective_workers(
num_workers=world_size, group_name=group_name, backend=backend
)
results = ray.get([a.do_reduce.remote(group_name, dst_rank) for a in actors])
for i in range(world_size):
if i == dst_rank:
assert (results[i] == np.ones((10,), dtype=np.float32) * world_size).all()
else:
assert (results[i] == np.ones((10,), dtype=np.float32)).all()
@pytest.mark.parametrize("backend", [Backend.GLOO])
@pytest.mark.parametrize("array_size", [2, 2**5, 2**10, 2**15, 2**20])
@pytest.mark.parametrize("dst_rank", [0, 2, 5, 6, 7])
def test_reduce_different_array_size(
ray_start_distributed_2_nodes, backend, array_size, dst_rank
):
world_size = 8
actors, _ = create_collective_workers(world_size, backend=backend)
ray.wait(
[a.set_buffer.remote(np.ones(array_size, dtype=np.float32)) for a in actors]
)
results = ray.get([a.do_reduce.remote(dst_rank=dst_rank) for a in actors])
for i in range(world_size):
if i == dst_rank:
assert (
results[i] == np.ones((array_size,), dtype=np.float32) * world_size
).all()
else:
assert (results[i] == np.ones((array_size,), dtype=np.float32)).all()
@pytest.mark.parametrize("backend", [Backend.GLOO])
@pytest.mark.parametrize("dst_rank", [0, 2, 5, 6, 7])
def test_reduce_different_op(ray_start_distributed_2_nodes, backend, dst_rank):
world_size = 8
actors, _ = create_collective_workers(world_size, backend=backend)
# check product
ray.wait(
[
a.set_buffer.remote(np.ones(10, dtype=np.float32) * (i + 2))
for i, a in enumerate(actors)
]
)
results = ray.get(
[a.do_reduce.remote(dst_rank=dst_rank, op=ReduceOp.PRODUCT) for a in actors]
)
product = 1
for i in range(world_size):
product = product * (i + 2)
for i in range(world_size):
if i == dst_rank:
assert (results[i] == np.ones((10,), dtype=np.float32) * product).all()
else:
assert (results[i] == np.ones((10,), dtype=np.float32) * (i + 2)).all()
# check min
ray.wait(
[
a.set_buffer.remote(np.ones(10, dtype=np.float32) * (i + 2))
for i, a in enumerate(actors)
]
)
results = ray.get(
[a.do_reduce.remote(dst_rank=dst_rank, op=ReduceOp.MIN) for a in actors]
)
for i in range(world_size):
if i == dst_rank:
assert (results[i] == np.ones((10,), dtype=np.float32) * 2).all()
else:
assert (results[i] == np.ones((10,), dtype=np.float32) * (i + 2)).all()
# check max
ray.wait(
[
a.set_buffer.remote(np.ones(10, dtype=np.float32) * (i + 2))
for i, a in enumerate(actors)
]
)
results = ray.get(
[a.do_reduce.remote(dst_rank=dst_rank, op=ReduceOp.MAX) for a in actors]
)
for i in range(world_size):
if i == dst_rank:
assert (
results[i] == np.ones((10,), dtype=np.float32) * (world_size + 1)
).all()
else:
assert (results[i] == np.ones((10,), dtype=np.float32) * (i + 2)).all()
@pytest.mark.parametrize("backend", [Backend.GLOO])
@pytest.mark.parametrize("dst_rank", [0, 2, 5, 6, 7])
def test_reduce_torch_numpy(ray_start_distributed_2_nodes, backend, dst_rank):
import torch
world_size = 8
actors, _ = create_collective_workers(world_size, backend=backend)
ray.get(
[
actors[1].set_buffer.remote(
torch.ones(
10,
)
)
]
)
results = ray.get([a.do_reduce.remote(dst_rank=dst_rank) for a in actors])
if dst_rank == 0:
assert (results[0] == np.ones((10,)) * world_size).all()
assert (results[1] == torch.ones((10,))).all()
else:
assert (results[0] == np.ones((10,))).all()
assert (results[1] == torch.ones((10,))).all()
@pytest.mark.parametrize("backend", [Backend.GLOO])
def test_reduce_invalid_rank(ray_start_distributed_2_nodes, backend, dst_rank=9):
world_size = 8
actors, _ = create_collective_workers(world_size, backend=backend)
with pytest.raises(ValueError):
_ = ray.get([a.do_reduce.remote(dst_rank=dst_rank) for a in actors])
if __name__ == "__main__":
import sys
import pytest
sys.exit(pytest.main(["-v", "-x", __file__]))
python/ray/util/collective/tests/distributed_cpu_tests/test_distributed_reducescatter.py
+131
View File
@@ -0,0 +1,131 @@
"""Test the collective reducescatter API on a distributed Ray cluster."""
import numpy as np
import pytest
import torch
import ray
from ray.util.collective.tests.cpu_util import (
create_collective_workers,
init_tensors_for_gather_scatter,
)
from ray.util.collective.types import Backend
@pytest.mark.parametrize("backend", [Backend.GLOO])
@pytest.mark.parametrize("tensor_backend", ["numpy", "torch"])
@pytest.mark.parametrize(
"array_size", [2, 2**5, 2**10, 2**15, 2**20, [2, 2], [5, 5, 5]]
)
def test_reducescatter_different_array_size(
ray_start_distributed_2_nodes, array_size, tensor_backend, backend
):
world_size = 8
actors, _ = create_collective_workers(world_size, backend=backend)
init_tensors_for_gather_scatter(
actors, array_size=array_size, tensor_backend=tensor_backend
)
results = ray.get([a.do_reducescatter.remote() for a in actors])
for i in range(world_size):
if tensor_backend == "numpy":
assert (
results[i] == np.ones(array_size, dtype=np.float32) * world_size
).all()
else:
assert (
results[i] == torch.ones(array_size, dtype=torch.float32) * world_size
).all()
@pytest.mark.parametrize("backend", [Backend.GLOO])
@pytest.mark.parametrize("dtype", [np.uint8, np.float16, np.float32, np.float64])
def test_reducescatter_different_dtype(ray_start_distributed_2_nodes, dtype, backend):
world_size = 8
actors, _ = create_collective_workers(world_size, backend=backend)
init_tensors_for_gather_scatter(actors, dtype=dtype)
results = ray.get([a.do_reducescatter.remote() for a in actors])
for i in range(world_size):
for j in range(world_size):
assert (results[i] == np.ones(10, dtype=dtype) * world_size).all()
@pytest.mark.parametrize("backend", [Backend.GLOO])
def test_reducescatter_torch_numpy(ray_start_distributed_2_nodes, backend):
world_size = 8
shape = [10, 10]
actors, _ = create_collective_workers(world_size, backend=backend)
# tensor is pytorch, list is numpy
for i, a in enumerate(actors):
t = torch.ones(shape, dtype=torch.float32) * (i + 1)
ray.wait([a.set_buffer.remote(t)])
list_buffer = [np.ones(shape, dtype=np.float32) for _ in range(world_size)]
ray.wait([a.set_list_buffer.remote(list_buffer)])
results = ray.get([a.do_reducescatter.remote() for a in actors])
for i in range(world_size):
assert (results[i] == torch.ones(shape, dtype=torch.float32) * world_size).all()
# tensor is numpy, list is pytorch
for i, a in enumerate(actors):
t = np.ones(shape, dtype=np.float32) * (i + 1)
ray.wait([a.set_buffer.remote(t)])
list_buffer = [
torch.ones(shape, dtype=torch.float32) for _ in range(world_size)
]
ray.wait([a.set_list_buffer.remote(list_buffer)])
results = ray.get([a.do_reducescatter.remote() for a in actors])
for i in range(world_size):
assert (results[i] == np.ones(shape, dtype=np.float32) * world_size).all()
# some tensors in the list are pytorch, some are numpy
for i, a in enumerate(actors):
if i % 2 == 0:
t = torch.ones(shape, dtype=torch.float32) * (i + 1)
else:
t = np.ones(shape, dtype=np.float32) * (i + 1)
ray.wait([a.set_buffer.remote(t)])
list_buffer = []
for j in range(world_size):
if j % 2 == 0:
list_buffer.append(torch.ones(shape, dtype=torch.float32))
else:
list_buffer.append(np.ones(shape, dtype=np.float32))
ray.wait([a.set_list_buffer.remote(list_buffer)])
results = ray.get([a.do_reducescatter.remote() for a in actors])
for i in range(world_size):
if i % 2 == 0:
assert (
results[i] == torch.ones(shape, dtype=torch.float32) * world_size
).all()
else:
assert (results[i] == np.ones(shape, dtype=np.float32) * world_size).all()
# mixed case
for i, a in enumerate(actors):
if i % 2 == 0:
t = torch.ones(shape, dtype=torch.float32) * (i + 1)
else:
t = np.ones(shape, dtype=np.float32) * (i + 1)
ray.wait([a.set_buffer.remote(t)])
list_buffer = []
for j in range(world_size):
if j % 2 == 0:
list_buffer.append(np.ones(shape, dtype=np.float32))
else:
list_buffer.append(torch.ones(shape, dtype=torch.float32))
ray.wait([a.set_list_buffer.remote(list_buffer)])
results = ray.get([a.do_reducescatter.remote() for a in actors])
for i in range(world_size):
if i % 2 == 0:
assert (
results[i] == torch.ones(shape, dtype=torch.float32) * world_size
).all()
else:
assert (results[i] == np.ones(shape, dtype=np.float32) * world_size).all()
if __name__ == "__main__":
import sys
import pytest
sys.exit(pytest.main(["-v", "-x", __file__]))
python/ray/util/collective/tests/distributed_cpu_tests/test_distributed_sendrecv.py
+52
View File
@@ -0,0 +1,52 @@
"""Test the send/recv API."""
import numpy as np
import pytest
import ray
from ray.util.collective.tests.cpu_util import create_collective_workers
from ray.util.collective.types import Backend
@pytest.mark.parametrize("backend", [Backend.GLOO])
@pytest.mark.parametrize("group_name", ["default", "test", "123?34!"])
@pytest.mark.parametrize("dst_rank", [0, 1, 3, 6])
@pytest.mark.parametrize("src_rank", [0, 2, 4, 7])
@pytest.mark.parametrize(
"array_size", [2**10, 2**15, 2**20, [2, 2], [5, 9, 10, 85]]
)
def test_sendrecv(
ray_start_distributed_2_nodes, group_name, array_size, src_rank, dst_rank, backend
):
if src_rank == dst_rank:
return
world_size = 8
actors, _ = create_collective_workers(
num_workers=world_size, group_name=group_name, backend=backend
)
ray.get(
[
a.set_buffer.remote(np.ones(array_size, dtype=np.float32) * (i + 1))
for i, a in enumerate(actors)
]
)
refs = []
for i in range(world_size):
refs.append(actors[i].get_buffer.remote())
refs[src_rank] = actors[src_rank].do_send.remote(group_name, dst_rank)
refs[dst_rank] = actors[dst_rank].do_recv.remote(group_name, src_rank)
results = ray.get(refs)
assert (
results[src_rank] == np.ones(array_size, dtype=np.float32) * (src_rank + 1)
).all()
assert (
results[dst_rank] == np.ones(array_size, dtype=np.float32) * (src_rank + 1)
).all()
ray.get([a.destroy_group.remote(group_name) for a in actors])
if __name__ == "__main__":
import sys
import pytest
sys.exit(pytest.main(["-v", "-x", __file__]))
python/ray/util/collective/tests/distributed_gpu_tests/__init__.py
View File
python/ray/util/collective/tests/distributed_gpu_tests/test_distributed_allgather.py
+138
View File
@@ -0,0 +1,138 @@
"""Test the allgather API on a distributed Ray cluster."""
import cupy as cp
import pytest
import torch
import ray
from ray.util.collective.tests.util import (
create_collective_workers,
init_tensors_for_gather_scatter,
)
@pytest.mark.parametrize("tensor_backend", ["cupy", "torch"])
@pytest.mark.parametrize(
"array_size", [2, 2**5, 2**10, 2**15, 2**20, [2, 2], [5, 5, 5]]
)
def test_allgather_different_array_size(
ray_start_distributed_2_nodes_4_gpus, array_size, tensor_backend
):
world_size = 4
actors, _ = create_collective_workers(world_size)
init_tensors_for_gather_scatter(
actors, array_size=array_size, tensor_backend=tensor_backend
)
results = ray.get([a.do_allgather.remote() for a in actors])
for i in range(world_size):
for j in range(world_size):
if tensor_backend == "cupy":
assert (
results[i][j] == cp.ones(array_size, dtype=cp.float32) * (j + 1)
).all()
else:
assert (
results[i][j]
== torch.ones(array_size, dtype=torch.float32).cuda() * (j + 1)
).all()
@pytest.mark.parametrize("dtype", [cp.uint8, cp.float16, cp.float32, cp.float64])
def test_allgather_different_dtype(ray_start_distributed_2_nodes_4_gpus, dtype):
world_size = 4
actors, _ = create_collective_workers(world_size)
init_tensors_for_gather_scatter(actors, dtype=dtype)
results = ray.get([a.do_allgather.remote() for a in actors])
for i in range(world_size):
for j in range(world_size):
assert (results[i][j] == cp.ones(10, dtype=dtype) * (j + 1)).all()
@pytest.mark.parametrize("length", [0, 1, 3, 4, 7, 8])
def test_unmatched_tensor_list_length(ray_start_distributed_2_nodes_4_gpus, length):
world_size = 4
actors, _ = create_collective_workers(world_size)
list_buffer = [cp.ones(10, dtype=cp.float32) for _ in range(length)]
ray.wait([a.set_list_buffer.remote(list_buffer) for a in actors])
if length != world_size:
with pytest.raises(RuntimeError):
ray.get([a.do_allgather.remote() for a in actors])
else:
ray.get([a.do_allgather.remote() for a in actors])
@pytest.mark.parametrize("shape", [10, 20, [4, 5], [1, 3, 5, 7]])
def test_unmatched_tensor_shape(ray_start_distributed_2_nodes_4_gpus, shape):
world_size = 4
actors, _ = create_collective_workers(world_size)
init_tensors_for_gather_scatter(actors, array_size=10)
list_buffer = [cp.ones(shape, dtype=cp.float32) for _ in range(world_size)]
ray.get([a.set_list_buffer.remote(list_buffer) for a in actors])
if shape != 10:
with pytest.raises(RuntimeError):
ray.get([a.do_allgather.remote() for a in actors])
else:
ray.get([a.do_allgather.remote() for a in actors])
def test_allgather_torch_cupy(ray_start_distributed_2_nodes_4_gpus):
world_size = 4
shape = [10, 10]
actors, _ = create_collective_workers(world_size)
# tensor is pytorch, list is cupy
for i, a in enumerate(actors):
t = torch.ones(shape, dtype=torch.float32).cuda() * (i + 1)
ray.wait([a.set_buffer.remote(t)])
list_buffer = [cp.ones(shape, dtype=cp.float32) for _ in range(world_size)]
ray.wait([a.set_list_buffer.remote(list_buffer)])
results = ray.get([a.do_allgather.remote() for a in actors])
for i in range(world_size):
for j in range(world_size):
assert (results[i][j] == cp.ones(shape, dtype=cp.float32) * (j + 1)).all()
# tensor is cupy, list is pytorch
for i, a in enumerate(actors):
t = cp.ones(shape, dtype=cp.float32) * (i + 1)
ray.wait([a.set_buffer.remote(t)])
list_buffer = [
torch.ones(shape, dtype=torch.float32).cuda() for _ in range(world_size)
]
ray.wait([a.set_list_buffer.remote(list_buffer)])
results = ray.get([a.do_allgather.remote() for a in actors])
for i in range(world_size):
for j in range(world_size):
assert (
results[i][j] == torch.ones(shape, dtype=torch.float32).cuda() * (j + 1)
).all()
# some tensors in the list are pytorch, some are cupy
for i, a in enumerate(actors):
t = cp.ones(shape, dtype=cp.float32) * (i + 1)
ray.wait([a.set_buffer.remote(t)])
list_buffer = []
for j in range(world_size):
if j % 2 == 0:
list_buffer.append(torch.ones(shape, dtype=torch.float32).cuda())
else:
list_buffer.append(cp.ones(shape, dtype=cp.float32))
ray.wait([a.set_list_buffer.remote(list_buffer)])
results = ray.get([a.do_allgather.remote() for a in actors])
for i in range(world_size):
for j in range(world_size):
if j % 2 == 0:
assert (
results[i][j]
== torch.ones(shape, dtype=torch.float32).cuda() * (j + 1)
).all()
else:
assert (
results[i][j] == cp.ones(shape, dtype=cp.float32) * (j + 1)
).all()
if __name__ == "__main__":
import sys
import pytest
sys.exit(pytest.main(["-v", "-x", __file__]))
python/ray/util/collective/tests/distributed_gpu_tests/test_distributed_allreduce.py
+170
View File
@@ -0,0 +1,170 @@
"""Test the collective allreduice API on a distributed Ray cluster."""
import cupy as cp
import pytest
import torch
import ray
from ray.util.collective.tests.util import create_collective_workers
from ray.util.collective.types import ReduceOp
@pytest.mark.parametrize("group_name", ["default", "test", "123?34!"])
@pytest.mark.parametrize("world_size", [2, 3, 4])
def test_allreduce_different_name(
ray_start_distributed_2_nodes_4_gpus, group_name, world_size
):
actors, _ = create_collective_workers(num_workers=world_size, group_name=group_name)
results = ray.get([a.do_allreduce.remote(group_name) for a in actors])
assert (results[0] == cp.ones((10,), dtype=cp.float32) * world_size).all()
assert (results[1] == cp.ones((10,), dtype=cp.float32) * world_size).all()
@pytest.mark.parametrize("array_size", [2, 2**5, 2**10, 2**15, 2**20])
def test_allreduce_different_array_size(
ray_start_distributed_2_nodes_4_gpus, array_size
):
world_size = 4
actors, _ = create_collective_workers(world_size)
ray.wait(
[a.set_buffer.remote(cp.ones(array_size, dtype=cp.float32)) for a in actors]
)
results = ray.get([a.do_allreduce.remote() for a in actors])
assert (results[0] == cp.ones((array_size,), dtype=cp.float32) * world_size).all()
assert (results[1] == cp.ones((array_size,), dtype=cp.float32) * world_size).all()
def test_allreduce_destroy(
ray_start_distributed_2_nodes_4_gpus, backend="nccl", group_name="default"
):
world_size = 4
actors, _ = create_collective_workers(world_size)
results = ray.get([a.do_allreduce.remote() for a in actors])
assert (results[0] == cp.ones((10,), dtype=cp.float32) * world_size).all()
assert (results[1] == cp.ones((10,), dtype=cp.float32) * world_size).all()
# destroy the group and try do work, should fail
ray.get([a.destroy_group.remote() for a in actors])
with pytest.raises(RuntimeError):
results = ray.get([a.do_allreduce.remote() for a in actors])
# reinit the same group and all reduce
ray.get(
[
actor.init_group.remote(world_size, i, backend, group_name)
for i, actor in enumerate(actors)
]
)
results = ray.get([a.do_allreduce.remote() for a in actors])
assert (
results[0] == cp.ones((10,), dtype=cp.float32) * world_size * world_size
).all()
assert (
results[1] == cp.ones((10,), dtype=cp.float32) * world_size * world_size
).all()
def test_allreduce_multiple_group(
ray_start_distributed_2_nodes_4_gpus, backend="nccl", num_groups=5
):
world_size = 4
actors, _ = create_collective_workers(world_size)
for group_name in range(1, num_groups):
ray.get(
[
actor.init_group.remote(world_size, i, backend, str(group_name))
for i, actor in enumerate(actors)
]
)
for i in range(num_groups):
group_name = "default" if i == 0 else str(i)
results = ray.get([a.do_allreduce.remote(group_name) for a in actors])
assert (
results[0] == cp.ones((10,), dtype=cp.float32) * (world_size ** (i + 1))
).all()
def test_allreduce_different_op(ray_start_distributed_2_nodes_4_gpus):
world_size = 4
actors, _ = create_collective_workers(world_size)
# check product
ray.wait(
[
a.set_buffer.remote(cp.ones(10, dtype=cp.float32) * (i + 2))
for i, a in enumerate(actors)
]
)
results = ray.get([a.do_allreduce.remote(op=ReduceOp.PRODUCT) for a in actors])
assert (results[0] == cp.ones((10,), dtype=cp.float32) * 120).all()
assert (results[1] == cp.ones((10,), dtype=cp.float32) * 120).all()
# check min
ray.wait(
[
a.set_buffer.remote(cp.ones(10, dtype=cp.float32) * (i + 2))
for i, a in enumerate(actors)
]
)
results = ray.get([a.do_allreduce.remote(op=ReduceOp.MIN) for a in actors])
assert (results[0] == cp.ones((10,), dtype=cp.float32) * 2).all()
assert (results[1] == cp.ones((10,), dtype=cp.float32) * 2).all()
# check max
ray.wait(
[
a.set_buffer.remote(cp.ones(10, dtype=cp.float32) * (i + 2))
for i, a in enumerate(actors)
]
)
results = ray.get([a.do_allreduce.remote(op=ReduceOp.MAX) for a in actors])
assert (results[0] == cp.ones((10,), dtype=cp.float32) * 5).all()
assert (results[1] == cp.ones((10,), dtype=cp.float32) * 5).all()
@pytest.mark.parametrize("dtype", [cp.uint8, cp.float16, cp.float32, cp.float64])
def test_allreduce_different_dtype(ray_start_distributed_2_nodes_4_gpus, dtype):
world_size = 4
actors, _ = create_collective_workers(world_size)
ray.wait([a.set_buffer.remote(cp.ones(10, dtype=dtype)) for a in actors])
results = ray.get([a.do_allreduce.remote() for a in actors])
assert (results[0] == cp.ones((10,), dtype=dtype) * world_size).all()
assert (results[1] == cp.ones((10,), dtype=dtype) * world_size).all()
def test_allreduce_torch_cupy(ray_start_distributed_2_nodes_4_gpus):
# import torch
world_size = 4
actors, _ = create_collective_workers(world_size)
ray.wait(
[
actors[1].set_buffer.remote(
torch.ones(
10,
).cuda()
)
]
)
results = ray.get([a.do_allreduce.remote() for a in actors])
assert (results[0] == cp.ones((10,)) * world_size).all()
ray.wait(
[
actors[0].set_buffer.remote(
torch.ones(
10,
)
)
]
)
ray.wait(
[
actors[1].set_buffer.remote(
cp.ones(
10,
)
)
]
)
with pytest.raises(RuntimeError):
results = ray.get([a.do_allreduce.remote() for a in actors])
python/ray/util/collective/tests/distributed_gpu_tests/test_distributed_basic_apis.py
+122
View File
@@ -0,0 +1,122 @@
"""Test the collective group APIs."""
from random import shuffle
import pytest
import ray
from ray.util.collective.tests.util import Worker, create_collective_workers
@pytest.mark.parametrize("world_size", [2, 3, 4])
@pytest.mark.parametrize("group_name", ["default", "test", "123?34!"])
def test_init_two_actors(ray_start_distributed_2_nodes_4_gpus, world_size, group_name):
actors, results = create_collective_workers(world_size, group_name)
for i in range(world_size):
assert results[i]
@pytest.mark.parametrize("world_size", [2, 3, 4])
def test_init_multiple_groups(ray_start_distributed_2_nodes_4_gpus, world_size):
num_groups = 1
actors = [Worker.remote() for _ in range(world_size)]
for i in range(num_groups):
group_name = str(i)
init_results = ray.get(
[
actor.init_group.remote(world_size, k, group_name=group_name)
for k, actor in enumerate(actors)
]
)
for j in range(world_size):
assert init_results[j]
@pytest.mark.parametrize("world_size", [2, 3, 4])
def test_get_rank(ray_start_distributed_2_nodes_4_gpus, world_size):
actors, _ = create_collective_workers(world_size)
actor0_rank = ray.get(actors[0].report_rank.remote())
assert actor0_rank == 0
actor1_rank = ray.get(actors[1].report_rank.remote())
assert actor1_rank == 1
# create a second group with a different name, and different
# orders of ranks.
new_group_name = "default2"
ranks = list(range(world_size))
shuffle(ranks)
ray.get(
[
actor.init_group.remote(world_size, ranks[i], group_name=new_group_name)
for i, actor in enumerate(actors)
]
)
actor0_rank = ray.get(actors[0].report_rank.remote(new_group_name))
assert actor0_rank == ranks[0]
actor1_rank = ray.get(actors[1].report_rank.remote(new_group_name))
assert actor1_rank == ranks[1]
@pytest.mark.parametrize("world_size", [2, 3, 4])
def test_get_collective_group_size(ray_start_distributed_2_nodes_4_gpus, world_size):
actors, _ = create_collective_workers(world_size)
actor0_world_size = ray.get(actors[0].report_world_size.remote())
actor1_world_size = ray.get(actors[1].report_world_size.remote())
assert actor0_world_size == actor1_world_size == world_size
def test_is_group_initialized(ray_start_distributed_2_nodes_4_gpus):
world_size = 4
actors, _ = create_collective_workers(world_size)
# check group is_init
actor0_is_init = ray.get(actors[0].report_is_group_initialized.remote())
assert actor0_is_init
actor0_is_init = ray.get(actors[0].report_is_group_initialized.remote("random"))
assert not actor0_is_init
actor0_is_init = ray.get(actors[0].report_is_group_initialized.remote("123"))
assert not actor0_is_init
actor1_is_init = ray.get(actors[0].report_is_group_initialized.remote())
assert actor1_is_init
actor1_is_init = ray.get(actors[0].report_is_group_initialized.remote("456"))
assert not actor1_is_init
def test_destroy_group(ray_start_distributed_2_nodes_4_gpus):
world_size = 4
actors, _ = create_collective_workers(world_size)
# Now destroy the group at actor0
ray.wait([actors[0].destroy_group.remote()])
actor0_is_init = ray.get(actors[0].report_is_group_initialized.remote())
assert not actor0_is_init
# should go well as the group `random` does not exist at all
ray.wait([actors[0].destroy_group.remote("random")])
actor1_is_init = ray.get(actors[1].report_is_group_initialized.remote())
assert actor1_is_init
ray.wait([actors[1].destroy_group.remote("random")])
actor1_is_init = ray.get(actors[1].report_is_group_initialized.remote())
assert actor1_is_init
ray.wait([actors[1].destroy_group.remote("default")])
actor1_is_init = ray.get(actors[1].report_is_group_initialized.remote())
assert not actor1_is_init
for i in [2, 3]:
ray.wait([actors[i].destroy_group.remote("default")])
# Now reconstruct the group using the same name
init_results = ray.get(
[actor.init_group.remote(world_size, i) for i, actor in enumerate(actors)]
)
for i in range(world_size):
assert init_results[i]
actor0_is_init = ray.get(actors[0].report_is_group_initialized.remote())
assert actor0_is_init
actor1_is_init = ray.get(actors[0].report_is_group_initialized.remote())
assert actor1_is_init
if __name__ == "__main__":
import sys
import pytest
sys.exit(pytest.main(["-v", "-x", __file__]))
python/ray/util/collective/tests/distributed_gpu_tests/test_distributed_broadcast.py
+81
View File
@@ -0,0 +1,81 @@
"""Test the broadcast API."""
import cupy as cp
import pytest
import ray
from ray.util.collective.tests.util import create_collective_workers
@pytest.mark.parametrize("group_name", ["default", "test", "123?34!"])
@pytest.mark.parametrize("src_rank", [0, 1, 2, 3])
def test_broadcast_different_name(
ray_start_distributed_2_nodes_4_gpus, group_name, src_rank
):
world_size = 4
actors, _ = create_collective_workers(num_workers=world_size, group_name=group_name)
ray.wait(
[
a.set_buffer.remote(cp.ones((10,), dtype=cp.float32) * (i + 2))
for i, a in enumerate(actors)
]
)
results = ray.get(
[
a.do_broadcast.remote(group_name=group_name, src_rank=src_rank)
for a in actors
]
)
for i in range(world_size):
assert (results[i] == cp.ones((10,), dtype=cp.float32) * (src_rank + 2)).all()
@pytest.mark.parametrize("array_size", [2, 2**5, 2**10, 2**15, 2**20])
@pytest.mark.parametrize("src_rank", [0, 1, 2, 3])
def test_broadcast_different_array_size(
ray_start_distributed_2_nodes_4_gpus, array_size, src_rank
):
world_size = 4
actors, _ = create_collective_workers(world_size)
ray.wait(
[
a.set_buffer.remote(cp.ones(array_size, dtype=cp.float32) * (i + 2))
for i, a in enumerate(actors)
]
)
results = ray.get([a.do_broadcast.remote(src_rank=src_rank) for a in actors])
for i in range(world_size):
assert (
results[i] == cp.ones((array_size,), dtype=cp.float32) * (src_rank + 2)
).all()
@pytest.mark.parametrize("src_rank", [0, 1])
def test_broadcast_torch_cupy(ray_start_distributed_2_nodes_4_gpus, src_rank):
import torch
world_size = 4
actors, _ = create_collective_workers(world_size)
ray.wait(
[
actors[1].set_buffer.remote(
torch.ones(
10,
).cuda()
* world_size
)
]
)
results = ray.get([a.do_broadcast.remote(src_rank=src_rank) for a in actors])
if src_rank == 0:
assert (results[0] == cp.ones((10,))).all()
assert (results[1] == torch.ones((10,)).cuda()).all()
else:
assert (results[0] == cp.ones((10,)) * world_size).all()
assert (results[1] == torch.ones((10,)).cuda() * world_size).all()
def test_broadcast_invalid_rank(ray_start_distributed_2_nodes_4_gpus, src_rank=3):
world_size = 2
actors, _ = create_collective_workers(world_size)
with pytest.raises(ValueError):
ray.get([a.do_broadcast.remote(src_rank=src_rank) for a in actors])
python/ray/util/collective/tests/distributed_gpu_tests/test_distributed_reduce.py
+127
View File
@@ -0,0 +1,127 @@
"""Test the reduce API."""
import cupy as cp
import pytest
import ray
from ray.util.collective.tests.util import create_collective_workers
from ray.util.collective.types import ReduceOp
@pytest.mark.parametrize("group_name", ["default", "test", "123?34!"])
@pytest.mark.parametrize("dst_rank", [0, 1, 2, 3])
def test_reduce_different_name(
ray_start_distributed_2_nodes_4_gpus, group_name, dst_rank
):
world_size = 4
actors, _ = create_collective_workers(num_workers=world_size, group_name=group_name)
results = ray.get([a.do_reduce.remote(group_name, dst_rank) for a in actors])
for i in range(world_size):
if i == dst_rank:
assert (results[i] == cp.ones((10,), dtype=cp.float32) * world_size).all()
else:
assert (results[i] == cp.ones((10,), dtype=cp.float32)).all()
@pytest.mark.parametrize("array_size", [2, 2**5, 2**10, 2**15, 2**20])
@pytest.mark.parametrize("dst_rank", [0, 1, 2, 3])
def test_reduce_different_array_size(
ray_start_distributed_2_nodes_4_gpus, array_size, dst_rank
):
world_size = 4
actors, _ = create_collective_workers(world_size)
ray.wait(
[a.set_buffer.remote(cp.ones(array_size, dtype=cp.float32)) for a in actors]
)
results = ray.get([a.do_reduce.remote(dst_rank=dst_rank) for a in actors])
for i in range(world_size):
if i == dst_rank:
assert (
results[i] == cp.ones((array_size,), dtype=cp.float32) * world_size
).all()
else:
assert (results[i] == cp.ones((array_size,), dtype=cp.float32)).all()
@pytest.mark.parametrize("dst_rank", [0, 1, 2, 3])
def test_reduce_different_op(ray_start_distributed_2_nodes_4_gpus, dst_rank):
world_size = 4
actors, _ = create_collective_workers(world_size)
# check product
ray.wait(
[
a.set_buffer.remote(cp.ones(10, dtype=cp.float32) * (i + 2))
for i, a in enumerate(actors)
]
)
results = ray.get(
[a.do_reduce.remote(dst_rank=dst_rank, op=ReduceOp.PRODUCT) for a in actors]
)
for i in range(world_size):
if i == dst_rank:
assert (results[i] == cp.ones((10,), dtype=cp.float32) * 120).all()
else:
assert (results[i] == cp.ones((10,), dtype=cp.float32) * (i + 2)).all()
# check min
ray.wait(
[
a.set_buffer.remote(cp.ones(10, dtype=cp.float32) * (i + 2))
for i, a in enumerate(actors)
]
)
results = ray.get(
[a.do_reduce.remote(dst_rank=dst_rank, op=ReduceOp.MIN) for a in actors]
)
for i in range(world_size):
if i == dst_rank:
assert (results[i] == cp.ones((10,), dtype=cp.float32) * 2).all()
else:
assert (results[i] == cp.ones((10,), dtype=cp.float32) * (i + 2)).all()
# check max
ray.wait(
[
a.set_buffer.remote(cp.ones(10, dtype=cp.float32) * (i + 2))
for i, a in enumerate(actors)
]
)
results = ray.get(
[a.do_reduce.remote(dst_rank=dst_rank, op=ReduceOp.MAX) for a in actors]
)
for i in range(world_size):
if i == dst_rank:
assert (results[i] == cp.ones((10,), dtype=cp.float32) * 5).all()
else:
assert (results[i] == cp.ones((10,), dtype=cp.float32) * (i + 2)).all()
@pytest.mark.parametrize("dst_rank", [0, 1])
def test_reduce_torch_cupy(ray_start_distributed_2_nodes_4_gpus, dst_rank):
import torch
world_size = 4
actors, _ = create_collective_workers(world_size)
ray.wait(
[
actors[1].set_buffer.remote(
torch.ones(
10,
).cuda()
)
]
)
results = ray.get([a.do_reduce.remote(dst_rank=dst_rank) for a in actors])
if dst_rank == 0:
assert (results[0] == cp.ones((10,)) * world_size).all()
assert (results[1] == torch.ones((10,)).cuda()).all()
else:
assert (results[0] == cp.ones((10,))).all()
assert (results[1] == torch.ones((10,)).cuda() * world_size).all()
def test_reduce_invalid_rank(ray_start_distributed_2_nodes_4_gpus, dst_rank=7):
world_size = 4
actors, _ = create_collective_workers(world_size)
with pytest.raises(ValueError):
ray.get([a.do_reduce.remote(dst_rank=dst_rank) for a in actors])
python/ray/util/collective/tests/distributed_gpu_tests/test_distributed_reducescatter.py
+130
View File
@@ -0,0 +1,130 @@
"""Test the collective reducescatter API on a distributed Ray cluster."""
import cupy as cp
import pytest
import torch
import ray
from ray.util.collective.tests.util import (
create_collective_workers,
init_tensors_for_gather_scatter,
)
@pytest.mark.parametrize("tensor_backend", ["cupy", "torch"])
@pytest.mark.parametrize(
"array_size", [2, 2**5, 2**10, 2**15, 2**20, [2, 2], [5, 5, 5]]
)
def test_reducescatter_different_array_size(
ray_start_distributed_2_nodes_4_gpus, array_size, tensor_backend
):
world_size = 4
actors, _ = create_collective_workers(world_size)
init_tensors_for_gather_scatter(
actors, array_size=array_size, tensor_backend=tensor_backend
)
results = ray.get([a.do_reducescatter.remote() for a in actors])
for i in range(world_size):
if tensor_backend == "cupy":
assert (
results[i] == cp.ones(array_size, dtype=cp.float32) * world_size
).all()
else:
assert (
results[i]
== torch.ones(array_size, dtype=torch.float32).cuda() * world_size
).all()
@pytest.mark.parametrize("dtype", [cp.uint8, cp.float16, cp.float32, cp.float64])
def test_reducescatter_different_dtype(ray_start_distributed_2_nodes_4_gpus, dtype):
world_size = 4
actors, _ = create_collective_workers(world_size)
init_tensors_for_gather_scatter(actors, dtype=dtype)
results = ray.get([a.do_reducescatter.remote() for a in actors])
for i in range(world_size):
for j in range(world_size):
assert (results[i] == cp.ones(10, dtype=dtype) * world_size).all()
def test_reducescatter_torch_cupy(ray_start_distributed_2_nodes_4_gpus):
world_size = 4
shape = [10, 10]
actors, _ = create_collective_workers(world_size)
# tensor is pytorch, list is cupy
for i, a in enumerate(actors):
t = torch.ones(shape, dtype=torch.float32).cuda() * (i + 1)
ray.wait([a.set_buffer.remote(t)])
list_buffer = [cp.ones(shape, dtype=cp.float32) for _ in range(world_size)]
ray.wait([a.set_list_buffer.remote(list_buffer)])
results = ray.get([a.do_reducescatter.remote() for a in actors])
for i in range(world_size):
assert (
results[i] == torch.ones(shape, dtype=torch.float32).cuda() * world_size
).all()
# tensor is cupy, list is pytorch
for i, a in enumerate(actors):
t = cp.ones(shape, dtype=cp.float32) * (i + 1)
ray.wait([a.set_buffer.remote(t)])
list_buffer = [
torch.ones(shape, dtype=torch.float32).cuda() for _ in range(world_size)
]
ray.wait([a.set_list_buffer.remote(list_buffer)])
results = ray.get([a.do_reducescatter.remote() for a in actors])
for i in range(world_size):
assert (results[i] == cp.ones(shape, dtype=cp.float32) * world_size).all()
# some tensors in the list are pytorch, some are cupy
for i, a in enumerate(actors):
if i % 2 == 0:
t = torch.ones(shape, dtype=torch.float32).cuda() * (i + 1)
else:
t = cp.ones(shape, dtype=cp.float32) * (i + 1)
ray.wait([a.set_buffer.remote(t)])
list_buffer = []
for j in range(world_size):
if j % 2 == 0:
list_buffer.append(torch.ones(shape, dtype=torch.float32).cuda())
else:
list_buffer.append(cp.ones(shape, dtype=cp.float32))
ray.wait([a.set_list_buffer.remote(list_buffer)])
results = ray.get([a.do_reducescatter.remote() for a in actors])
for i in range(world_size):
if i % 2 == 0:
assert (
results[i] == torch.ones(shape, dtype=torch.float32).cuda() * world_size
).all()
else:
assert (results[i] == cp.ones(shape, dtype=cp.float32) * world_size).all()
# mixed case
for i, a in enumerate(actors):
if i % 2 == 0:
t = torch.ones(shape, dtype=torch.float32).cuda() * (i + 1)
else:
t = cp.ones(shape, dtype=cp.float32) * (i + 1)
ray.wait([a.set_buffer.remote(t)])
list_buffer = []
for j in range(world_size):
if j % 2 == 0:
list_buffer.append(cp.ones(shape, dtype=cp.float32))
else:
list_buffer.append(torch.ones(shape, dtype=torch.float32).cuda())
ray.wait([a.set_list_buffer.remote(list_buffer)])
results = ray.get([a.do_reducescatter.remote() for a in actors])
for i in range(world_size):
if i % 2 == 0:
assert (
results[i] == torch.ones(shape, dtype=torch.float32).cuda() * world_size
).all()
else:
assert (results[i] == cp.ones(shape, dtype=cp.float32) * world_size).all()
if __name__ == "__main__":
import sys
import pytest
sys.exit(pytest.main(["-v", "-x", __file__]))
python/ray/util/collective/tests/distributed_gpu_tests/test_distributed_sendrecv.py
+40
View File
@@ -0,0 +1,40 @@
"""Test the send/recv API."""
import cupy as cp
import pytest
import ray
from ray.util.collective.tests.util import create_collective_workers
@pytest.mark.parametrize("group_name", ["default", "test", "123?34!"])
@pytest.mark.parametrize("dst_rank", [0, 1, 2, 3])
@pytest.mark.parametrize("src_rank", [0, 1, 2, 3])
@pytest.mark.parametrize(
"array_size", [2**10, 2**15, 2**20, [2, 2], [5, 9, 10, 85]]
)
def test_sendrecv(
ray_start_distributed_2_nodes_4_gpus, group_name, array_size, src_rank, dst_rank
):
if src_rank == dst_rank:
return
world_size = 4
actors, _ = create_collective_workers(num_workers=world_size, group_name=group_name)
ray.get(
[
a.set_buffer.remote(cp.ones(array_size, dtype=cp.float32) * (i + 1))
for i, a in enumerate(actors)
]
)
refs = []
for i in range(world_size):
refs.append(actors[i].get_buffer.remote())
refs[src_rank] = actors[src_rank].do_send.remote(group_name, dst_rank)
refs[dst_rank] = actors[dst_rank].do_recv.remote(group_name, src_rank)
results = ray.get(refs)
assert (
results[src_rank] == cp.ones(array_size, dtype=cp.float32) * (src_rank + 1)
).all()
assert (
results[dst_rank] == cp.ones(array_size, dtype=cp.float32) * (src_rank + 1)
).all()
ray.get([a.destroy_group.remote(group_name) for a in actors])
python/ray/util/collective/tests/distributed_multigpu_tests/__init__.py
View File
python/ray/util/collective/tests/distributed_multigpu_tests/test_distributed_multigpu_allgather.py
+87
View File
@@ -0,0 +1,87 @@
"""Test the allgather API on a distributed Ray cluster."""
import cupy as cp
import pytest
import torch
import ray
from ray.util.collective.tests.util import (
create_collective_multigpu_workers,
init_tensors_for_gather_scatter_multigpu,
)
@pytest.mark.parametrize("tensor_backend", ["cupy", "torch"])
@pytest.mark.parametrize(
"array_size", [2, 2**5, 2**10, 2**15, 2**20, [2, 2], [5, 5, 5]]
)
def test_allgather_different_array_size(
ray_start_distributed_multigpu_2_nodes_4_gpus, array_size, tensor_backend
):
world_size = 2
num_gpu_per_worker = 2
actual_world_size = world_size * num_gpu_per_worker
actors, _ = create_collective_multigpu_workers(world_size)
init_tensors_for_gather_scatter_multigpu(
actors, array_size=array_size, tensor_backend=tensor_backend
)
results = ray.get([a.do_allgather_multigpu.remote() for a in actors])
for i in range(world_size):
for j in range(num_gpu_per_worker):
for k in range(actual_world_size):
if tensor_backend == "cupy":
assert (
results[i][j][k] == cp.ones(array_size, dtype=cp.float32)
).all()
else:
assert (
results[i][j][k]
== torch.ones(array_size, dtype=torch.float32).cuda(j)
).all()
def test_allgather_torch_cupy(ray_start_distributed_multigpu_2_nodes_4_gpus):
world_size = 2
num_gpu_per_worker = 2
actual_world_size = world_size * num_gpu_per_worker
shape = [10, 10]
actors, _ = create_collective_multigpu_workers(world_size)
# tensor is pytorch, list is cupy
for i, a in enumerate(actors):
ray.get(
[a.set_buffer.remote(shape, tensor_type0="torch", tensor_type1="torch")]
)
ray.get(
[a.set_list_buffer.remote(shape, tensor_type0="cupy", tensor_type1="cupy")]
)
results = ray.get([a.do_allgather_multigpu.remote() for a in actors])
for i in range(world_size):
for j in range(num_gpu_per_worker):
for k in range(actual_world_size):
assert (results[i][j][k] == cp.ones(shape, dtype=cp.float32)).all()
# tensor is cupy, list is pytorch
for i, a in enumerate(actors):
ray.get([a.set_buffer.remote(shape, tensor_type0="cupy", tensor_type1="cupy")])
ray.get(
[
a.set_list_buffer.remote(
shape, tensor_type0="torch", tensor_type1="torch"
)
]
)
results = ray.get([a.do_allgather_multigpu.remote() for a in actors])
for i in range(world_size):
for j in range(num_gpu_per_worker):
for k in range(actual_world_size):
assert (
results[i][j][k] == torch.ones(shape, dtype=torch.float32).cuda(j)
).all()
if __name__ == "__main__":
import sys
import pytest
sys.exit(pytest.main(["-v", "-x", __file__]))
python/ray/util/collective/tests/distributed_multigpu_tests/test_distributed_multigpu_allreduce.py
+166
View File
@@ -0,0 +1,166 @@
"""Test the collective allreduice API on a distributed Ray cluster."""
import logging
import cupy as cp
import pytest
import ray
from ray.util.collective.tests.util import create_collective_multigpu_workers
from ray.util.collective.types import ReduceOp
logger = logging.getLogger(__name__)
logger.setLevel("DEBUG")
@pytest.mark.parametrize("group_name", ["default", "test", "123?34!"])
def test_allreduce_multigpu_different_name(
ray_start_distributed_multigpu_2_nodes_4_gpus, group_name
):
world_size = 2
num_gpu_per_worker = 2
actual_world_size = world_size * num_gpu_per_worker
actors, _ = create_collective_multigpu_workers(
num_workers=world_size, group_name=group_name
)
results = ray.get([a.do_allreduce_multigpu.remote(group_name) for a in actors])
assert (results[0] == cp.ones((10,), dtype=cp.float32) * actual_world_size).all()
assert (results[1] == cp.ones((10,), dtype=cp.float32) * actual_world_size).all()
@pytest.mark.parametrize("array_size", [2, 2**5, 2**10, 2**15, 2**20])
def test_allreduce_multigpu_different_array_size(
ray_start_distributed_multigpu_2_nodes_4_gpus, array_size
):
world_size = 2
num_gpu_per_worker = 2
actual_world_size = world_size * num_gpu_per_worker
actors, _ = create_collective_multigpu_workers(world_size)
ray.get([a.set_buffer.remote(array_size) for a in actors])
results = ray.get([a.do_allreduce_multigpu.remote() for a in actors])
assert (
results[0] == cp.ones((array_size,), dtype=cp.float32) * actual_world_size
).all()
assert (
results[1] == cp.ones((array_size,), dtype=cp.float32) * actual_world_size
).all()
def test_allreduce_multigpu_destroy(
ray_start_distributed_multigpu_2_nodes_4_gpus, backend="nccl", group_name="default"
):
world_size = 2
num_gpu_per_worker = 2
actual_world_size = world_size * num_gpu_per_worker
actors, _ = create_collective_multigpu_workers(world_size)
results = ray.get([a.do_allreduce_multigpu.remote() for a in actors])
assert (results[0] == cp.ones((10,), dtype=cp.float32) * actual_world_size).all()
assert (results[1] == cp.ones((10,), dtype=cp.float32) * actual_world_size).all()
# destroy the group and try do work, should fail
ray.get([a.destroy_group.remote() for a in actors])
with pytest.raises(RuntimeError):
results = ray.get([a.do_allreduce_multigpu.remote() for a in actors])
# reinit the same group and all reduce
ray.get(
[
actor.init_group.remote(world_size, i, backend, group_name)
for i, actor in enumerate(actors)
]
)
results = ray.get([a.do_allreduce_multigpu.remote() for a in actors])
assert (
results[0]
== cp.ones((10,), dtype=cp.float32) * actual_world_size * actual_world_size
).all()
assert (
results[1]
== cp.ones((10,), dtype=cp.float32) * actual_world_size * actual_world_size
).all()
def test_allreduce_multigpu_multiple_group(
ray_start_distributed_multigpu_2_nodes_4_gpus, backend="nccl", num_groups=5
):
world_size = 2
num_gpu_per_worker = 2
actual_world_size = world_size * num_gpu_per_worker
actors, _ = create_collective_multigpu_workers(world_size)
for group_name in range(1, num_groups):
ray.get(
[
actor.init_group.remote(world_size, i, backend, str(group_name))
for i, actor in enumerate(actors)
]
)
for i in range(num_groups):
group_name = "default" if i == 0 else str(i)
results = ray.get([a.do_allreduce_multigpu.remote(group_name) for a in actors])
assert (
results[0]
== cp.ones((10,), dtype=cp.float32) * (actual_world_size ** (i + 1))
).all()
def test_allreduce_multigpu_different_op(ray_start_distributed_multigpu_2_nodes_4_gpus):
world_size = 2
actors, _ = create_collective_multigpu_workers(world_size)
# check product
ray.get(actors[0].set_buffer.remote([10], value0=2, value1=3))
ray.get(actors[1].set_buffer.remote([10], value0=4, value1=5))
results = ray.get(
[a.do_allreduce_multigpu.remote(op=ReduceOp.PRODUCT) for a in actors]
)
assert (results[0] == cp.ones((10,), dtype=cp.float32) * 120).all()
assert (results[1] == cp.ones((10,), dtype=cp.float32) * 120).all()
# check min
ray.get(actors[0].set_buffer.remote([10], value0=2, value1=3))
ray.get(actors[1].set_buffer.remote([10], value0=4, value1=5))
results = ray.get([a.do_allreduce_multigpu.remote(op=ReduceOp.MIN) for a in actors])
assert (results[0] == cp.ones((10,), dtype=cp.float32) * 2).all()
assert (results[1] == cp.ones((10,), dtype=cp.float32) * 2).all()
# check max
ray.get(actors[0].set_buffer.remote([10], value0=2, value1=3))
ray.get(actors[1].set_buffer.remote([10], value0=4, value1=5))
results = ray.get([a.do_allreduce_multigpu.remote(op=ReduceOp.MAX) for a in actors])
assert (results[0] == cp.ones((10,), dtype=cp.float32) * 5).all()
assert (results[1] == cp.ones((10,), dtype=cp.float32) * 5).all()
@pytest.mark.parametrize("dtype", [cp.uint8, cp.float16, cp.float32, cp.float64])
def test_allreduce_multigpu_different_dtype(
ray_start_distributed_multigpu_2_nodes_4_gpus, dtype
):
world_size = 2
num_gpu_per_worker = 2
actual_world_size = world_size * num_gpu_per_worker
actors, _ = create_collective_multigpu_workers(world_size)
ray.get([a.set_buffer.remote([10], dtype=dtype) for a in actors])
results = ray.get([a.do_allreduce_multigpu.remote() for a in actors])
assert (results[0] == cp.ones((10,), dtype=dtype) * actual_world_size).all()
assert (results[1] == cp.ones((10,), dtype=dtype) * actual_world_size).all()
def test_allreduce_torch_cupy(ray_start_distributed_multigpu_2_nodes_4_gpus):
# import torch
world_size = 2
actual_world_size = 4
actors, _ = create_collective_multigpu_workers(world_size)
ray.get(actors[0].set_buffer.remote([10]))
ray.get(
actors[1].set_buffer.remote([10], tensor_type0="torch", tensor_type1="torch")
)
results = ray.get([a.do_allreduce_multigpu.remote() for a in actors])
assert (results[0] == cp.ones((10,)) * actual_world_size).all()
ray.get(
actors[0].set_buffer.remote([10], tensor_type0="cupy", tensor_type1="torch")
)
ray.get(
actors[1].set_buffer.remote([10], tensor_type0="torch", tensor_type1="cupy")
)
results = ray.get([a.do_allreduce_multigpu.remote() for a in actors])
assert (results[0] == cp.ones((10,)) * actual_world_size).all()
python/ray/util/collective/tests/distributed_multigpu_tests/test_distributed_multigpu_basic_apis.py
+103
View File
@@ -0,0 +1,103 @@
"""Test the collective group APIs."""
from random import shuffle
import pytest
import ray
from ray.util.collective.tests.util import create_collective_multigpu_workers
@pytest.mark.parametrize("group_name", ["default", "test", "123?34!"])
def test_init_two_actors(ray_start_distributed_multigpu_2_nodes_4_gpus, group_name):
world_size = 2
actors, results = create_collective_multigpu_workers(world_size, group_name)
for i in range(world_size):
assert results[i]
def test_report_num_gpus(ray_start_distributed_multigpu_2_nodes_4_gpus):
world_size = 2
actors, results = create_collective_multigpu_workers(world_size)
num_gpus = ray.get([actor.report_num_gpus.remote() for actor in actors])
assert num_gpus == [2, 2]
def test_get_rank(ray_start_distributed_multigpu_2_nodes_4_gpus):
world_size = 2
actors, _ = create_collective_multigpu_workers(world_size)
actor0_rank = ray.get(actors[0].report_rank.remote())
assert actor0_rank == 0
actor1_rank = ray.get(actors[1].report_rank.remote())
assert actor1_rank == 1
# create a second group with a different name, and different
# orders of ranks.
new_group_name = "default2"
ranks = list(range(world_size))
shuffle(ranks)
ray.get(
[
actor.init_group.remote(world_size, ranks[i], group_name=new_group_name)
for i, actor in enumerate(actors)
]
)
actor0_rank = ray.get(actors[0].report_rank.remote(new_group_name))
assert actor0_rank == ranks[0]
actor1_rank = ray.get(actors[1].report_rank.remote(new_group_name))
assert actor1_rank == ranks[1]
def test_is_group_initialized(ray_start_distributed_multigpu_2_nodes_4_gpus):
world_size = 2
actors, _ = create_collective_multigpu_workers(world_size)
# check group is_init
actor0_is_init = ray.get(actors[0].report_is_group_initialized.remote())
assert actor0_is_init
actor0_is_init = ray.get(actors[0].report_is_group_initialized.remote("random"))
assert not actor0_is_init
actor0_is_init = ray.get(actors[0].report_is_group_initialized.remote("123"))
assert not actor0_is_init
actor1_is_init = ray.get(actors[0].report_is_group_initialized.remote())
assert actor1_is_init
actor1_is_init = ray.get(actors[0].report_is_group_initialized.remote("456"))
assert not actor1_is_init
def test_destroy_group(ray_start_distributed_multigpu_2_nodes_4_gpus):
world_size = 2
actors, _ = create_collective_multigpu_workers(world_size)
# Now destroy the group at actor0
ray.wait([actors[0].destroy_group.remote()])
actor0_is_init = ray.get(actors[0].report_is_group_initialized.remote())
assert not actor0_is_init
# should go well as the group `random` does not exist at all
ray.wait([actors[0].destroy_group.remote("random")])
actor1_is_init = ray.get(actors[1].report_is_group_initialized.remote())
assert actor1_is_init
ray.wait([actors[1].destroy_group.remote("random")])
actor1_is_init = ray.get(actors[1].report_is_group_initialized.remote())
assert actor1_is_init
ray.wait([actors[1].destroy_group.remote("default")])
actor1_is_init = ray.get(actors[1].report_is_group_initialized.remote())
assert not actor1_is_init
# Now reconstruct the group using the same name
init_results = ray.get(
[actor.init_group.remote(world_size, i) for i, actor in enumerate(actors)]
)
for i in range(world_size):
assert init_results[i]
actor0_is_init = ray.get(actors[0].report_is_group_initialized.remote())
assert actor0_is_init
actor1_is_init = ray.get(actors[0].report_is_group_initialized.remote())
assert actor1_is_init
if __name__ == "__main__":
import sys
import pytest
sys.exit(pytest.main(["-v", "-x", __file__]))
python/ray/util/collective/tests/distributed_multigpu_tests/test_distributed_multigpu_broadcast.py
+112
View File
@@ -0,0 +1,112 @@
"""Test the broadcast API."""
import cupy as cp
import pytest
import ray
from ray.util.collective.tests.util import create_collective_multigpu_workers
@pytest.mark.parametrize("group_name", ["default", "test", "123?34!"])
@pytest.mark.parametrize("src_rank", [0, 1])
@pytest.mark.parametrize("src_gpu_index", [0, 1])
def test_broadcast_different_name(
ray_start_distributed_multigpu_2_nodes_4_gpus, group_name, src_rank, src_gpu_index
):
world_size = 2
num_gpu_per_worker = 2
actors, _ = create_collective_multigpu_workers(
num_workers=world_size, group_name=group_name
)
ray.get(actors[0].set_buffer.remote([10], value0=2, value1=3))
ray.get(actors[1].set_buffer.remote([10], value0=4, value1=5))
results = ray.get(
[
a.do_broadcast_multigpu.remote(
group_name=group_name, src_rank=src_rank, src_gpu_index=src_gpu_index
)
for a in actors
]
)
for i in range(world_size):
for j in range(num_gpu_per_worker):
val = (src_rank + 1) * 2 + src_gpu_index
assert (results[i][j] == cp.ones([10], dtype=cp.float32) * val).all()
@pytest.mark.parametrize("array_size", [2, 2**5, 2**10, 2**15, 2**20])
@pytest.mark.parametrize("src_rank", [0, 1])
@pytest.mark.parametrize("src_gpu_index", [0, 1])
def test_broadcast_different_array_size(
ray_start_distributed_multigpu_2_nodes_4_gpus, array_size, src_rank, src_gpu_index
):
world_size = 2
num_gpu_per_worker = 2
actors, _ = create_collective_multigpu_workers(world_size)
ray.get(actors[0].set_buffer.remote([array_size], value0=2, value1=3))
ray.get(actors[1].set_buffer.remote([array_size], value0=4, value1=5))
results = ray.get(
[
a.do_broadcast_multigpu.remote(
src_rank=src_rank, src_gpu_index=src_gpu_index
)
for a in actors
]
)
for i in range(world_size):
for j in range(num_gpu_per_worker):
val = (src_rank + 1) * 2 + src_gpu_index
assert (
results[i][j] == cp.ones((array_size,), dtype=cp.float32) * val
).all()
@pytest.mark.parametrize("src_rank", [0, 1])
@pytest.mark.parametrize("src_gpu_index", [0, 1])
def test_broadcast_torch_cupy(
ray_start_distributed_multigpu_2_nodes_4_gpus, src_rank, src_gpu_index
):
import torch
world_size = 2
num_gpu_per_worker = 2
actors, _ = create_collective_multigpu_workers(world_size)
ray.get(actors[0].set_buffer.remote([10], value0=2, value1=3))
ray.get(
actors[1].set_buffer.remote(
[10], value0=4, value1=5, tensor_type0="torch", tensor_type1="torch"
)
)
results = ray.get(
[
a.do_broadcast_multigpu.remote(
src_rank=src_rank, src_gpu_index=src_gpu_index
)
for a in actors
]
)
for i in range(world_size):
for j in range(num_gpu_per_worker):
val = (src_rank + 1) * 2 + src_gpu_index
if i == 0:
assert (results[i][j] == cp.ones([10], dtype=cp.float32) * val).all()
else:
assert (results[i][j] == torch.ones([10]).cuda(j) * val).all()
@pytest.mark.parametrize("src_rank", [3, 4])
@pytest.mark.parametrize("src_gpu_index", [2, 3])
def test_broadcast_invalid_rank(
ray_start_distributed_multigpu_2_nodes_4_gpus, src_rank, src_gpu_index
):
world_size = 2
actors, _ = create_collective_multigpu_workers(world_size)
with pytest.raises(ValueError):
_ = ray.get(
[
a.do_broadcast_multigpu.remote(
src_rank=src_rank, src_gpu_index=src_gpu_index
)
for a in actors
]
)
python/ray/util/collective/tests/distributed_multigpu_tests/test_distributed_multigpu_reduce.py
+194
View File
@@ -0,0 +1,194 @@
"""Test the reduce API."""
import cupy as cp
import pytest
import ray
from ray.util.collective.tests.util import create_collective_multigpu_workers
from ray.util.collective.types import ReduceOp
@pytest.mark.parametrize("group_name", ["default", "test", "123?34!"])
@pytest.mark.parametrize("dst_rank", [0, 1])
@pytest.mark.parametrize("dst_gpu_index", [0, 1])
def test_reduce_different_name(
ray_start_distributed_multigpu_2_nodes_4_gpus, group_name, dst_rank, dst_gpu_index
):
world_size = 2
num_gpu_per_worker = 2
actual_world_size = world_size * num_gpu_per_worker
actors, _ = create_collective_multigpu_workers(
num_workers=world_size, group_name=group_name
)
results = ray.get(
[
a.do_reduce_multigpu.remote(
group_name, dst_rank=dst_rank, dst_gpu_index=dst_gpu_index
)
for a in actors
]
)
for i in range(world_size):
for j in range(num_gpu_per_worker):
if i == dst_rank and j == dst_gpu_index:
assert (
results[i][j]
== cp.ones((10,), dtype=cp.float32) * actual_world_size
).all()
else:
assert (results[i][j] == cp.ones((10,), dtype=cp.float32)).all()
@pytest.mark.parametrize("array_size", [2, 2**5, 2**10, 2**15, 2**20])
@pytest.mark.parametrize("dst_rank", [0, 1])
@pytest.mark.parametrize("dst_gpu_index", [0, 1])
def test_reduce_different_array_size(
ray_start_distributed_multigpu_2_nodes_4_gpus, array_size, dst_rank, dst_gpu_index
):
world_size = 2
num_gpu_per_worker = 2
actual_world_size = world_size * num_gpu_per_worker
actors, _ = create_collective_multigpu_workers(num_workers=world_size)
ray.get(actors[0].set_buffer.remote(array_size))
ray.get(actors[1].set_buffer.remote(array_size))
results = ray.get(
[
a.do_reduce_multigpu.remote(dst_rank=dst_rank, dst_gpu_index=dst_gpu_index)
for a in actors
]
)
for i in range(world_size):
for j in range(num_gpu_per_worker):
if i == dst_rank and j == dst_gpu_index:
assert (
results[i][j]
== cp.ones((array_size,), dtype=cp.float32) * actual_world_size
).all()
else:
assert (results[i][j] == cp.ones((array_size,), dtype=cp.float32)).all()
@pytest.mark.parametrize("dst_rank", [0, 1])
@pytest.mark.parametrize("dst_gpu_index", [0, 1])
def test_reduce_different_op(
ray_start_distributed_multigpu_2_nodes_4_gpus, dst_rank, dst_gpu_index
):
world_size = 2
num_gpu_per_worker = 2
actors, _ = create_collective_multigpu_workers(world_size)
# check product
ray.get(actors[0].set_buffer.remote([10], value0=2, value1=3))
ray.get(actors[1].set_buffer.remote([10], value0=4, value1=5))
results = ray.get(
[
a.do_reduce_multigpu.remote(
dst_rank=dst_rank, dst_gpu_index=dst_gpu_index, op=ReduceOp.PRODUCT
)
for a in actors
]
)
for i in range(world_size):
for j in range(num_gpu_per_worker):
if i == dst_rank and j == dst_gpu_index:
assert (results[i][j] == cp.ones((10,), dtype=cp.float32) * 120).all()
else:
val = (i + 1) * 2 + j
assert (results[i][j] == cp.ones((10,), dtype=cp.float32) * val).all()
# check min
ray.get(actors[0].set_buffer.remote([10], value0=2, value1=3))
ray.get(actors[1].set_buffer.remote([10], value0=4, value1=5))
results = ray.get(
[
a.do_reduce_multigpu.remote(
dst_rank=dst_rank, dst_gpu_index=dst_gpu_index, op=ReduceOp.MIN
)
for a in actors
]
)
for i in range(world_size):
for j in range(num_gpu_per_worker):
if i == dst_rank and j == dst_gpu_index:
assert (results[i][j] == cp.ones((10,), dtype=cp.float32) * 2).all()
else:
val = (i + 1) * 2 + j
assert (results[i][j] == cp.ones((10,), dtype=cp.float32) * val).all()
# check max
ray.get(actors[0].set_buffer.remote([10], value0=2, value1=3))
ray.get(actors[1].set_buffer.remote([10], value0=4, value1=5))
results = ray.get(
[
a.do_reduce_multigpu.remote(
dst_rank=dst_rank, dst_gpu_index=dst_gpu_index, op=ReduceOp.MAX
)
for a in actors
]
)
for i in range(world_size):
for j in range(num_gpu_per_worker):
if i == dst_rank and j == dst_gpu_index:
assert (results[i][j] == cp.ones((10,), dtype=cp.float32) * 5).all()
else:
val = (i + 1) * 2 + j
assert (results[i][j] == cp.ones((10,), dtype=cp.float32) * val).all()
@pytest.mark.parametrize("dst_rank", [0, 1])
@pytest.mark.parametrize("dst_gpu_index", [0, 1])
def test_reduce_torch_cupy(
ray_start_distributed_multigpu_2_nodes_4_gpus, dst_rank, dst_gpu_index
):
import torch
world_size = 2
num_gpu_per_worker = 2
actors, _ = create_collective_multigpu_workers(world_size)
ray.get(actors[0].set_buffer.remote([10], value0=2, value1=3))
ray.get(
actors[1].set_buffer.remote(
[10], value0=4, value1=5, tensor_type0="torch", tensor_type1="torch"
)
)
results = ray.get(
[
a.do_reduce_multigpu.remote(dst_rank=dst_rank, dst_gpu_index=dst_gpu_index)
for a in actors
]
)
for i in range(world_size):
for j in range(num_gpu_per_worker):
val = (i + 1) * 2 + j
if dst_rank == i and dst_gpu_index == j:
if i == 0:
assert (results[i][j] == cp.ones([10], dtype=cp.float32) * 14).all()
else:
assert (results[i][j] == torch.ones([10]).cuda(j) * 14).all()
else:
if i == 0:
assert (
results[i][j] == cp.ones([10], dtype=cp.float32) * val
).all()
else:
assert (results[i][j] == torch.ones([10]).cuda(j) * val).all()
@pytest.mark.parametrize("dst_rank", [3, 4])
@pytest.mark.parametrize("dst_gpu_index", [2, 3])
def test_reduce_invalid_rank(
ray_start_distributed_multigpu_2_nodes_4_gpus, dst_rank, dst_gpu_index
):
world_size = 2
actors, _ = create_collective_multigpu_workers(world_size)
with pytest.raises(ValueError):
_ = ray.get(
[
a.do_reduce_multigpu.remote(
dst_rank=dst_rank, dst_gpu_index=dst_gpu_index
)
for a in actors
]
)
python/ray/util/collective/tests/distributed_multigpu_tests/test_distributed_multigpu_reducescatter.py
+90
View File
@@ -0,0 +1,90 @@
"""Test the collective reducescatter API on a distributed Ray cluster."""
import cupy as cp
import pytest
import torch
import ray
from ray.util.collective.tests.util import (
create_collective_multigpu_workers,
init_tensors_for_gather_scatter_multigpu,
)
@pytest.mark.parametrize("tensor_backend", ["cupy", "torch"])
@pytest.mark.parametrize(
"array_size", [2, 2**5, 2**10, 2**15, 2**20, [2, 2], [5, 5, 5]]
)
def test_reducescatter_different_array_size(
ray_start_distributed_multigpu_2_nodes_4_gpus, array_size, tensor_backend
):
world_size = 2
num_gpu_per_worker = 2
actual_world_size = world_size * num_gpu_per_worker
actors, _ = create_collective_multigpu_workers(world_size)
init_tensors_for_gather_scatter_multigpu(
actors, array_size=array_size, tensor_backend=tensor_backend
)
results = ray.get([a.do_reducescatter_multigpu.remote() for a in actors])
for i in range(world_size):
for j in range(num_gpu_per_worker):
if tensor_backend == "cupy":
assert (
results[i][j]
== cp.ones(array_size, dtype=cp.float32) * actual_world_size
).all()
else:
assert (
results[i][j]
== torch.ones(array_size, dtype=torch.float32).cuda(j)
* actual_world_size
).all()
def test_reducescatter_torch_cupy(ray_start_distributed_multigpu_2_nodes_4_gpus):
world_size = 2
num_gpu_per_worker = 2
actual_world_size = world_size * num_gpu_per_worker
shape = [10, 10]
actors, _ = create_collective_multigpu_workers(world_size)
# tensor is pytorch, list is cupy
for i, a in enumerate(actors):
ray.get(
[a.set_buffer.remote(shape, tensor_type0="torch", tensor_type1="torch")]
)
ray.get(
[a.set_list_buffer.remote(shape, tensor_type0="cupy", tensor_type1="cupy")]
)
results = ray.get([a.do_reducescatter_multigpu.remote() for a in actors])
for i in range(world_size):
for j in range(num_gpu_per_worker):
assert (
results[i][j]
== torch.ones(shape, dtype=torch.float32).cuda(j) * actual_world_size
).all()
# tensor is cupy, list is pytorch
for i, a in enumerate(actors):
ray.get([a.set_buffer.remote(shape, tensor_type0="cupy", tensor_type1="cupy")])
ray.get(
[
a.set_list_buffer.remote(
shape, tensor_type0="torch", tensor_type1="torch"
)
]
)
results = ray.get([a.do_reducescatter_multigpu.remote() for a in actors])
for i in range(world_size):
for j in range(num_gpu_per_worker):
assert (
results[i][j] == cp.ones(shape, dtype=cp.float32) * actual_world_size
).all()
if __name__ == "__main__":
import sys
import pytest
sys.exit(pytest.main(["-v", "-x", __file__]))
python/ray/util/collective/tests/distributed_multigpu_tests/test_distributed_multigpu_sendrecv.py
+54
View File
@@ -0,0 +1,54 @@
"""Test the send/recv API."""
import cupy as cp
import pytest
import ray
from ray.util.collective.tests.util import create_collective_multigpu_workers
# @pytest.mark.parametrize("group_name", ["default", "test", "123?34!"])
@pytest.mark.parametrize("dst_rank", [0, 1])
@pytest.mark.parametrize("src_rank", [0, 1])
@pytest.mark.parametrize("dst_gpu_index", [0, 1])
@pytest.mark.parametrize("src_gpu_index", [0, 1])
@pytest.mark.parametrize(
"array_size", [2**10, 2**15, 2**20, [2, 2], [5, 9, 10, 85]]
)
def test_sendrecv(
ray_start_distributed_multigpu_2_nodes_4_gpus,
array_size,
src_rank,
dst_rank,
src_gpu_index,
dst_gpu_index,
):
if src_rank == dst_rank:
return
world_size = 2
actors, _ = create_collective_multigpu_workers(num_workers=world_size)
ray.get(actors[0].set_buffer.remote(array_size, value0=2, value1=3))
ray.get(actors[1].set_buffer.remote(array_size, value0=4, value1=5))
refs = []
for i in range(world_size):
refs.append(actors[i].get_buffer.remote())
refs[src_rank][src_gpu_index] = actors[src_rank].do_send_multigpu.remote(
dst_rank=dst_rank, dst_gpu_index=dst_gpu_index, src_gpu_index=src_gpu_index
)
refs[dst_rank][dst_gpu_index] = actors[dst_rank].do_recv_multigpu.remote(
src_rank=src_rank, src_gpu_index=src_gpu_index, dst_gpu_index=dst_gpu_index
)
results = []
results_flattend = ray.get(refs[0] + refs[1])
results.append([results_flattend[0], results_flattend[1]])
results.append([results_flattend[2], results_flattend[3]])
assert (
results[src_rank][src_gpu_index]
== cp.ones(array_size, dtype=cp.float32) * ((src_rank + 1) * 2 + src_gpu_index)
).all()
assert (
results[dst_rank][dst_gpu_index]
== cp.ones(array_size, dtype=cp.float32) * ((src_rank + 1) * 2 + src_gpu_index)
).all()
ray.get([a.destroy_group.remote() for a in actors])
python/ray/util/collective/tests/single_node_cpu_tests/__init__.py
View File
python/ray/util/collective/tests/single_node_cpu_tests/test_allgather.py
+143
View File
@@ -0,0 +1,143 @@
"""Test the collective allgather API."""
import numpy as np
import pytest
import torch
import ray
from ray.util.collective.tests.cpu_util import (
create_collective_workers,
init_tensors_for_gather_scatter,
)
from ray.util.collective.types import Backend
@pytest.mark.parametrize("backend", [Backend.GLOO])
@pytest.mark.parametrize("tensor_backend", ["numpy", "torch"])
@pytest.mark.parametrize(
"array_size", [2, 2**5, 2**10, 2**15, 2**20, [2, 2], [5, 5, 5]]
)
def test_allgather_different_array_size(
ray_start_single_node, array_size, tensor_backend, backend
):
world_size = 2
actors, _ = create_collective_workers(world_size, backend=backend)
init_tensors_for_gather_scatter(
actors, array_size=array_size, tensor_backend=tensor_backend
)
results = ray.get([a.do_allgather.remote() for a in actors])
for i in range(world_size):
for j in range(world_size):
if tensor_backend == "numpy":
assert (
results[i][j] == np.ones(array_size, dtype=np.float32) * (j + 1)
).all()
else:
assert (
results[i][j]
== torch.ones(array_size, dtype=torch.float32) * (j + 1)
).all()
@pytest.mark.parametrize("backend", [Backend.GLOO])
@pytest.mark.parametrize("dtype", [np.uint8, np.float16, np.float32, np.float64])
def test_allgather_different_dtype(ray_start_single_node, dtype, backend):
world_size = 2
actors, _ = create_collective_workers(world_size, backend=backend)
init_tensors_for_gather_scatter(actors, dtype=dtype)
results = ray.get([a.do_allgather.remote() for a in actors])
for i in range(world_size):
for j in range(world_size):
assert (results[i][j] == np.ones(10, dtype=dtype) * (j + 1)).all()
@pytest.mark.parametrize("backend", [Backend.GLOO])
@pytest.mark.parametrize("length", [0, 1, 2, 3])
def test_unmatched_tensor_list_length(ray_start_single_node, length, backend):
world_size = 2
actors, _ = create_collective_workers(world_size, backend=backend)
list_buffer = [np.ones(10, dtype=np.float32) for _ in range(length)]
ray.wait([a.set_list_buffer.remote(list_buffer, copy=True) for a in actors])
if length != world_size:
with pytest.raises(RuntimeError):
ray.get([a.do_allgather.remote() for a in actors])
else:
ray.get([a.do_allgather.remote() for a in actors])
@pytest.mark.parametrize("backend", [Backend.GLOO])
@pytest.mark.parametrize("shape", [10, 20, [4, 5], [1, 3, 5, 7]])
def test_unmatched_tensor_shape(ray_start_single_node, shape, backend):
world_size = 2
actors, _ = create_collective_workers(world_size, backend=backend)
init_tensors_for_gather_scatter(actors, array_size=10)
list_buffer = [np.ones(shape, dtype=np.float32) for _ in range(world_size)]
ray.get([a.set_list_buffer.remote(list_buffer, copy=True) for a in actors])
if shape != 10:
with pytest.raises(RuntimeError):
ray.get([a.do_allgather.remote() for a in actors])
else:
ray.get([a.do_allgather.remote() for a in actors])
@pytest.mark.parametrize("backend", [Backend.GLOO])
def test_allgather_torch_numpy(ray_start_single_node, backend):
world_size = 2
shape = [10, 10]
actors, _ = create_collective_workers(world_size, backend=backend)
# tensor is pytorch, list is numpy
for i, a in enumerate(actors):
t = torch.ones(shape, dtype=torch.float32) * (i + 1)
ray.wait([a.set_buffer.remote(t)])
list_buffer = [np.ones(shape, dtype=np.float32) for _ in range(world_size)]
ray.wait([a.set_list_buffer.remote(list_buffer, copy=True)])
results = ray.get([a.do_allgather.remote() for a in actors])
for i in range(world_size):
for j in range(world_size):
assert (results[i][j] == np.ones(shape, dtype=np.float32) * (j + 1)).all()
# tensor is numpy, list is pytorch
for i, a in enumerate(actors):
t = np.ones(shape, dtype=np.float32) * (i + 1)
ray.wait([a.set_buffer.remote(t)])
list_buffer = [
torch.ones(shape, dtype=torch.float32) for _ in range(world_size)
]
ray.wait([a.set_list_buffer.remote(list_buffer, copy=True)])
results = ray.get([a.do_allgather.remote() for a in actors])
for i in range(world_size):
for j in range(world_size):
assert (
results[i][j] == torch.ones(shape, dtype=torch.float32) * (j + 1)
).all()
# some tensors in the list are pytorch, some are numpy
for i, a in enumerate(actors):
t = np.ones(shape, dtype=np.float32) * (i + 1)
ray.wait([a.set_buffer.remote(t)])
list_buffer = []
for j in range(world_size):
if j % 2 == 0:
list_buffer.append(torch.ones(shape, dtype=torch.float32))
else:
list_buffer.append(np.ones(shape, dtype=np.float32))
ray.wait([a.set_list_buffer.remote(list_buffer, copy=True)])
results = ray.get([a.do_allgather.remote() for a in actors])
for i in range(world_size):
for j in range(world_size):
if j % 2 == 0:
assert (
results[i][j] == torch.ones(shape, dtype=torch.float32) * (j + 1)
).all()
else:
assert (
results[i][j] == np.ones(shape, dtype=np.float32) * (j + 1)
).all()
if __name__ == "__main__":
import sys
import pytest
sys.exit(pytest.main(["-v", "-x", __file__]))
python/ray/util/collective/tests/single_node_cpu_tests/test_allreduce.py
+166
View File
@@ -0,0 +1,166 @@
"""Test the collective allreduice API."""
import numpy as np
import pytest
import torch
import ray
from ray.util.collective.tests.cpu_util import create_collective_workers
from ray.util.collective.types import Backend, ReduceOp
@pytest.mark.parametrize("backend", [Backend.GLOO])
@pytest.mark.parametrize("group_name", ["default", "test", "123?34!"])
def test_allreduce_different_name(ray_start_single_node, group_name, backend):
world_size = 2
actors, _ = create_collective_workers(
num_workers=world_size, group_name=group_name, backend=backend
)
results = ray.get([a.do_allreduce.remote(group_name) for a in actors])
assert (results[0] == np.ones((10,), dtype=np.float32) * world_size).all()
assert (results[1] == np.ones((10,), dtype=np.float32) * world_size).all()
@pytest.mark.parametrize("backend", [Backend.GLOO])
@pytest.mark.parametrize("array_size", [2, 2**5, 2**10, 2**15, 2**20])
def test_allreduce_different_array_size(ray_start_single_node, array_size, backend):
world_size = 2
actors, _ = create_collective_workers(world_size, backend=backend)
ray.wait(
[a.set_buffer.remote(np.ones(array_size, dtype=np.float32)) for a in actors]
)
results = ray.get([a.do_allreduce.remote() for a in actors])
assert (results[0] == np.ones((array_size,), dtype=np.float32) * world_size).all()
assert (results[1] == np.ones((array_size,), dtype=np.float32) * world_size).all()
@pytest.mark.parametrize("backend", [Backend.GLOO])
def test_allreduce_destroy(ray_start_single_node, backend, group_name="default"):
world_size = 2
actors, _ = create_collective_workers(world_size, backend=backend)
results = ray.get([a.do_allreduce.remote() for a in actors])
assert (results[0] == np.ones((10,), dtype=np.float32) * world_size).all()
assert (results[1] == np.ones((10,), dtype=np.float32) * world_size).all()
# destroy the group and try do work, should fail
ray.get([a.destroy_group.remote() for a in actors])
with pytest.raises(RuntimeError):
results = ray.get([a.do_allreduce.remote() for a in actors])
# reinit the same group and all reduce
ray.get(
[
actor.init_group.remote(world_size, i, backend, group_name)
for i, actor in enumerate(actors)
]
)
results = ray.get([a.do_allreduce.remote() for a in actors])
assert (results[0] == np.ones((10,), dtype=np.float32) * world_size * 2).all()
assert (results[1] == np.ones((10,), dtype=np.float32) * world_size * 2).all()
@pytest.mark.parametrize("backend", [Backend.GLOO])
def test_allreduce_multiple_group(ray_start_single_node, backend, num_groups=5):
world_size = 2
actors, _ = create_collective_workers(world_size, backend=backend)
for group_name in range(1, num_groups):
ray.get(
[
actor.init_group.remote(world_size, i, backend, str(group_name))
for i, actor in enumerate(actors)
]
)
for i in range(num_groups):
group_name = "default" if i == 0 else str(i)
results = ray.get([a.do_allreduce.remote(group_name) for a in actors])
assert (
results[0] == np.ones((10,), dtype=np.float32) * (world_size ** (i + 1))
).all()
@pytest.mark.parametrize("backend", [Backend.GLOO])
def test_allreduce_different_op(ray_start_single_node, backend):
world_size = 2
actors, _ = create_collective_workers(world_size, backend=backend)
# check product
ray.wait(
[
a.set_buffer.remote(np.ones(10, dtype=np.float32) * (i + 2))
for i, a in enumerate(actors)
]
)
results = ray.get([a.do_allreduce.remote(op=ReduceOp.PRODUCT) for a in actors])
assert (results[0] == np.ones((10,), dtype=np.float32) * 6).all()
assert (results[1] == np.ones((10,), dtype=np.float32) * 6).all()
# check min
ray.wait(
[
a.set_buffer.remote(np.ones(10, dtype=np.float32) * (i + 2))
for i, a in enumerate(actors)
]
)
results = ray.get([a.do_allreduce.remote(op=ReduceOp.MIN) for a in actors])
assert (results[0] == np.ones((10,), dtype=np.float32) * 2).all()
assert (results[1] == np.ones((10,), dtype=np.float32) * 2).all()
# check max
ray.wait(
[
a.set_buffer.remote(np.ones(10, dtype=np.float32) * (i + 2))
for i, a in enumerate(actors)
]
)
results = ray.get([a.do_allreduce.remote(op=ReduceOp.MAX) for a in actors])
assert (results[0] == np.ones((10,), dtype=np.float32) * 3).all()
assert (results[1] == np.ones((10,), dtype=np.float32) * 3).all()
@pytest.mark.parametrize("backend", [Backend.GLOO])
@pytest.mark.parametrize("dtype", [np.uint8, np.float16, np.float32, np.float64])
def test_allreduce_different_dtype(ray_start_single_node, dtype, backend):
world_size = 2
actors, _ = create_collective_workers(world_size, backend=backend)
ray.wait([a.set_buffer.remote(np.ones(10, dtype=dtype)) for a in actors])
results = ray.get([a.do_allreduce.remote() for a in actors])
assert (results[0] == np.ones((10,), dtype=dtype) * world_size).all()
assert (results[1] == np.ones((10,), dtype=dtype) * world_size).all()
@pytest.mark.parametrize("backend", [Backend.GLOO])
def test_allreduce_torch_numpy(ray_start_single_node, backend):
# import torch
world_size = 2
actors, _ = create_collective_workers(world_size, backend=backend)
ray.wait(
[
actors[1].set_buffer.remote(
torch.ones(
10,
)
)
]
)
results = ray.get([a.do_allreduce.remote() for a in actors])
assert (results[0] == np.ones((10,)) * world_size).all()
ray.wait(
[
actors[0].set_buffer.remote(
torch.ones(
10,
)
)
]
)
ray.wait([actors[1].set_buffer.remote(np.ones(10, dtype=np.float32))])
ray.get([a.do_allreduce.remote() for a in actors])
if __name__ == "__main__":
import sys
import pytest
sys.exit(pytest.main(["-v", "-x", __file__]))
python/ray/util/collective/tests/single_node_cpu_tests/test_basic_apis.py
+130
View File
@@ -0,0 +1,130 @@
"""Test the collective group APIs."""
import pytest
import ray
from ray.util.collective.tests.cpu_util import Worker, create_collective_workers
from ray.util.collective.types import Backend
@pytest.mark.parametrize("backend", [Backend.GLOO])
@pytest.mark.parametrize("group_name", ["default", "test", "123?34!"])
def test_init_two_actors(ray_start_single_node, group_name, backend):
world_size = 2
actors, results = create_collective_workers(world_size, group_name, backend=backend)
for i in range(world_size):
assert results[i]
@pytest.mark.parametrize("backend", [Backend.GLOO])
def test_init_multiple_groups(ray_start_single_node, backend):
world_size = 2
num_groups = 10
actors = [Worker.remote() for i in range(world_size)]
for i in range(num_groups):
group_name = str(i)
init_results = ray.get(
[
actor.init_group.remote(
world_size, k, group_name=group_name, backend=backend
)
for k, actor in enumerate(actors)
]
)
for j in range(world_size):
assert init_results[j]
@pytest.mark.parametrize("backend", [Backend.GLOO])
def test_get_rank(ray_start_single_node, backend):
world_size = 2
actors, _ = create_collective_workers(world_size, backend=backend)
actor0_rank = ray.get(actors[0].report_rank.remote())
assert actor0_rank == 0
actor1_rank = ray.get(actors[1].report_rank.remote())
assert actor1_rank == 1
# create a second group with a different name,
# and different order of ranks.
new_group_name = "default2"
ray.get(
[
actor.init_group.remote(
world_size,
world_size - 1 - i,
group_name=new_group_name,
backend=backend,
)
for i, actor in enumerate(actors)
]
)
actor0_rank = ray.get(actors[0].report_rank.remote(new_group_name))
assert actor0_rank == 1
actor1_rank = ray.get(actors[1].report_rank.remote(new_group_name))
assert actor1_rank == 0
@pytest.mark.parametrize("backend", [Backend.GLOO])
def test_get_world_size(ray_start_single_node, backend):
world_size = 2
actors, _ = create_collective_workers(world_size, backend=backend)
actor0_world_size = ray.get(actors[0].report_world_size.remote())
actor1_world_size = ray.get(actors[1].report_world_size.remote())
assert actor0_world_size == actor1_world_size == world_size
@pytest.mark.parametrize("backend", [Backend.GLOO])
def test_is_group_initialized(ray_start_single_node, backend):
world_size = 2
actors, _ = create_collective_workers(world_size, backend=backend)
# check group is_init
actor0_is_init = ray.get(actors[0].report_is_group_initialized.remote())
assert actor0_is_init
actor0_is_init = ray.get(actors[0].report_is_group_initialized.remote("random"))
assert not actor0_is_init
actor0_is_init = ray.get(actors[0].report_is_group_initialized.remote("123"))
assert not actor0_is_init
actor1_is_init = ray.get(actors[0].report_is_group_initialized.remote())
assert actor1_is_init
actor1_is_init = ray.get(actors[0].report_is_group_initialized.remote("456"))
assert not actor1_is_init
@pytest.mark.parametrize("backend", [Backend.GLOO])
def test_destroy_group(ray_start_single_node, backend):
world_size = 2
actors, _ = create_collective_workers(world_size, backend=backend)
# Now destroy the group at actor0
ray.wait([actors[0].destroy_group.remote()])
actor0_is_init = ray.get(actors[0].report_is_group_initialized.remote())
assert not actor0_is_init
# should go well as the group `random` does not exist at all
ray.wait([actors[0].destroy_group.remote("random")])
actor1_is_init = ray.get(actors[1].report_is_group_initialized.remote())
assert actor1_is_init
ray.wait([actors[1].destroy_group.remote("random")])
actor1_is_init = ray.get(actors[1].report_is_group_initialized.remote())
assert actor1_is_init
ray.wait([actors[1].destroy_group.remote("default")])
actor1_is_init = ray.get(actors[1].report_is_group_initialized.remote())
assert not actor1_is_init
# Now reconstruct the group using the same name
init_results = ray.get(
[actor.init_group.remote(world_size, i) for i, actor in enumerate(actors)]
)
for i in range(world_size):
assert init_results[i]
actor0_is_init = ray.get(actors[0].report_is_group_initialized.remote())
assert actor0_is_init
actor1_is_init = ray.get(actors[0].report_is_group_initialized.remote())
assert actor1_is_init
if __name__ == "__main__":
import sys
import pytest
sys.exit(pytest.main(["-v", "-x", __file__]))
python/ray/util/collective/tests/single_node_cpu_tests/test_broadcast.py
+94
View File
@@ -0,0 +1,94 @@
"""Test the broadcast API."""
import numpy as np
import pytest
import ray
from ray.util.collective.tests.cpu_util import create_collective_workers
from ray.util.collective.types import Backend
@pytest.mark.parametrize("backend", [Backend.GLOO])
@pytest.mark.parametrize("group_name", ["default", "test", "123?34!"])
@pytest.mark.parametrize("src_rank", [0, 1])
def test_broadcast_different_name(ray_start_single_node, group_name, src_rank, backend):
world_size = 2
actors, _ = create_collective_workers(
num_workers=world_size, group_name=group_name, backend=backend
)
ray.get(
[
a.set_buffer.remote(np.ones((10,), dtype=np.float32) * (i + 2))
for i, a in enumerate(actors)
]
)
results = ray.get(
[
a.do_broadcast.remote(group_name=group_name, src_rank=src_rank)
for a in actors
]
)
for i in range(world_size):
assert (results[i] == np.ones((10,), dtype=np.float32) * (src_rank + 2)).all()
@pytest.mark.parametrize("backend", [Backend.GLOO])
@pytest.mark.parametrize("array_size", [2, 2**5, 2**10, 2**15, 2**20])
@pytest.mark.parametrize("src_rank", [0, 1])
def test_broadcast_different_array_size(
ray_start_single_node, array_size, src_rank, backend
):
world_size = 2
actors, _ = create_collective_workers(world_size, backend=backend)
ray.get(
[
a.set_buffer.remote(np.ones(array_size, dtype=np.float32) * (i + 2))
for i, a in enumerate(actors)
]
)
results = ray.get([a.do_broadcast.remote(src_rank=src_rank) for a in actors])
for i in range(world_size):
assert (
results[i] == np.ones((array_size,), dtype=np.float32) * (src_rank + 2)
).all()
@pytest.mark.parametrize("backend", [Backend.GLOO])
@pytest.mark.parametrize("src_rank", [0, 1])
def test_broadcast_torch_numpy(ray_start_single_node, src_rank, backend):
import torch
world_size = 2
actors, _ = create_collective_workers(world_size, backend=backend)
ray.wait(
[
actors[1].set_buffer.remote(
torch.ones(
10,
)
* world_size
)
]
)
results = ray.get([a.do_broadcast.remote(src_rank=src_rank) for a in actors])
if src_rank == 0:
assert (results[0] == np.ones((10,))).all()
assert (results[1] == torch.ones((10,))).all()
else:
assert (results[0] == np.ones((10,)) * world_size).all()
assert (results[1] == torch.ones((10,)) * world_size).all()
@pytest.mark.parametrize("backend", [Backend.GLOO])
def test_broadcast_invalid_rank(ray_start_single_node, backend, src_rank=3):
world_size = 2
actors, _ = create_collective_workers(world_size, backend=backend)
with pytest.raises(ValueError):
ray.get([a.do_broadcast.remote(src_rank=src_rank) for a in actors])
if __name__ == "__main__":
import sys
import pytest
sys.exit(pytest.main(["-v", "-x", __file__]))
python/ray/util/collective/tests/single_node_cpu_tests/test_gloo_group_isolation.py
+67
View File
@@ -0,0 +1,67 @@
import time
import ray
import ray.util.collective as col
from ray.util.collective.collective_group.torch_gloo_collective_group import (
TorchGLOOGroup as GLOOGroup,
)
from ray.util.collective.types import Backend
@ray.remote
class Worker:
def __init__(self):
pass
def init_gloo_group(
self, world_size: int, rank: int, group_name: str, gloo_timeout: int = 30000
):
col.init_collective_group(
world_size, rank, Backend.GLOO, group_name, gloo_timeout
)
return True
def get_gloo_timeout(self, group_name: str) -> int:
g = col.get_group_handle(group_name)
# Check if the group is initialized correctly
assert isinstance(g, GLOOGroup)
return g._gloo_context.getTimeout()
def test_two_groups_in_one_cluster(ray_start_single_node):
name1 = "name_1"
name2 = "name_2"
time1 = 40000
time2 = 60000
w1 = Worker.remote()
ret1 = w1.init_gloo_group.remote(1, 0, name1, time1)
w2 = Worker.remote()
ret2 = w2.init_gloo_group.remote(1, 0, name2, time2)
assert ray.get(ret1)
assert ray.get(ret2)
assert ray.get(w1.get_gloo_timeout.remote(name1)) == time1
assert ray.get(w2.get_gloo_timeout.remote(name2)) == time2
def test_failure_when_initializing(shutdown_only):
# job1
ray.init()
w1 = Worker.remote()
ret1 = w1.init_gloo_group.remote(2, 0, "name_1")
ray.wait([ret1], timeout=1)
time.sleep(5)
ray.shutdown()
# job2
ray.init()
w2 = Worker.remote()
ret2 = w2.init_gloo_group.remote(1, 0, "name_1")
assert ray.get(ret2)
if __name__ == "__main__":
import sys
import pytest
sys.exit(pytest.main(["-v", "-x", __file__]))
python/ray/util/collective/tests/single_node_cpu_tests/test_reduce.py
+167
View File
@@ -0,0 +1,167 @@
"""Test the reduce API."""
import numpy as np
import pytest
import ray
from ray.util.collective.tests.cpu_util import create_collective_workers
from ray.util.collective.types import Backend, ReduceOp
@pytest.mark.parametrize("backend", [Backend.GLOO])
@pytest.mark.parametrize("group_name", ["default", "test", "123?34!"])
@pytest.mark.parametrize("dst_rank", [0, 1])
def test_reduce_different_name(ray_start_single_node, group_name, dst_rank, backend):
world_size = 2
actors, _ = create_collective_workers(
num_workers=world_size, group_name=group_name, backend=backend
)
results = ray.get([a.do_reduce.remote(group_name, dst_rank) for a in actors])
for i in range(world_size):
if i == dst_rank:
assert (results[i] == np.ones((10,), dtype=np.float32) * world_size).all()
else:
assert (results[i] == np.ones((10,), dtype=np.float32)).all()
@pytest.mark.parametrize("backend", [Backend.GLOO])
@pytest.mark.parametrize("array_size", [2, 2**5, 2**10, 2**15, 2**20])
@pytest.mark.parametrize("dst_rank", [0, 1])
def test_reduce_different_array_size(
ray_start_single_node, array_size, dst_rank, backend
):
world_size = 2
actors, _ = create_collective_workers(world_size, backend=backend)
ray.wait(
[a.set_buffer.remote(np.ones(array_size, dtype=np.float32)) for a in actors]
)
results = ray.get([a.do_reduce.remote(dst_rank=dst_rank) for a in actors])
for i in range(world_size):
if i == dst_rank:
assert (
results[i] == np.ones((array_size,), dtype=np.float32) * world_size
).all()
else:
assert (results[i] == np.ones((array_size,), dtype=np.float32)).all()
@pytest.mark.parametrize("backend", [Backend.GLOO])
@pytest.mark.parametrize("dst_rank", [0, 1])
def test_reduce_multiple_group(ray_start_single_node, dst_rank, backend, num_groups=5):
world_size = 2
actors, _ = create_collective_workers(world_size, backend=backend)
for group_name in range(1, num_groups):
ray.get(
[
actor.init_group.remote(world_size, i, backend, str(group_name))
for i, actor in enumerate(actors)
]
)
for i in range(num_groups):
group_name = "default" if i == 0 else str(i)
results = ray.get(
[
a.do_reduce.remote(dst_rank=dst_rank, group_name=group_name)
for a in actors
]
)
for j in range(world_size):
if j == dst_rank:
assert (results[j] == np.ones((10,), dtype=np.float32) * (i + 2)).all()
else:
assert (results[j] == np.ones((10,), dtype=np.float32)).all()
@pytest.mark.parametrize("backend", [Backend.GLOO])
@pytest.mark.parametrize("dst_rank", [0, 1])
def test_reduce_different_op(ray_start_single_node, dst_rank, backend):
world_size = 2
actors, _ = create_collective_workers(world_size, backend=backend)
# check product
ray.wait(
[
a.set_buffer.remote(np.ones(10, dtype=np.float32) * (i + 2))
for i, a in enumerate(actors)
]
)
results = ray.get(
[a.do_reduce.remote(dst_rank=dst_rank, op=ReduceOp.PRODUCT) for a in actors]
)
for i in range(world_size):
if i == dst_rank:
assert (results[i] == np.ones((10,), dtype=np.float32) * 6).all()
else:
assert (results[i] == np.ones((10,), dtype=np.float32) * (i + 2)).all()
# check min
ray.wait(
[
a.set_buffer.remote(np.ones(10, dtype=np.float32) * (i + 2))
for i, a in enumerate(actors)
]
)
results = ray.get(
[a.do_reduce.remote(dst_rank=dst_rank, op=ReduceOp.MIN) for a in actors]
)
for i in range(world_size):
if i == dst_rank:
assert (results[i] == np.ones((10,), dtype=np.float32) * 2).all()
else:
assert (results[i] == np.ones((10,), dtype=np.float32) * (i + 2)).all()
# check max
ray.wait(
[
a.set_buffer.remote(np.ones(10, dtype=np.float32) * (i + 2))
for i, a in enumerate(actors)
]
)
results = ray.get(
[a.do_reduce.remote(dst_rank=dst_rank, op=ReduceOp.MAX) for a in actors]
)
for i in range(world_size):
if i == dst_rank:
assert (results[i] == np.ones((10,), dtype=np.float32) * 3).all()
else:
assert (results[i] == np.ones((10,), dtype=np.float32) * (i + 2)).all()
@pytest.mark.parametrize("backend", [Backend.GLOO])
@pytest.mark.parametrize("dst_rank", [0, 1])
def test_reduce_torch_numpy(ray_start_single_node, dst_rank, backend):
import torch
world_size = 2
actors, _ = create_collective_workers(world_size, backend=backend)
ray.wait(
[
actors[1].set_buffer.remote(
torch.ones(
10,
)
)
]
)
results = ray.get([a.do_reduce.remote(dst_rank=dst_rank) for a in actors])
if dst_rank == 0:
assert (results[0] == np.ones((10,)) * world_size).all()
assert (results[1] == torch.ones((10,))).all()
else:
assert (results[0] == np.ones((10,))).all()
assert (results[1] == torch.ones((10,)) * world_size).all()
@pytest.mark.parametrize("backend", [Backend.GLOO])
def test_reduce_invalid_rank(ray_start_single_node, backend, dst_rank=3):
world_size = 2
actors, _ = create_collective_workers(world_size, backend=backend)
with pytest.raises(ValueError):
ray.get([a.do_reduce.remote(dst_rank=dst_rank) for a in actors])
if __name__ == "__main__":
import sys
import pytest
sys.exit(pytest.main(["-v", "-x", __file__]))
python/ray/util/collective/tests/single_node_cpu_tests/test_reducescatter.py
+131
View File
@@ -0,0 +1,131 @@
"""Test the collective reducescatter API."""
import numpy as np
import pytest
import torch
import ray
from ray.util.collective.tests.cpu_util import (
create_collective_workers,
init_tensors_for_gather_scatter,
)
from ray.util.collective.types import Backend
@pytest.mark.parametrize("backend", [Backend.GLOO])
@pytest.mark.parametrize("tensor_backend", ["numpy", "torch"])
@pytest.mark.parametrize(
"array_size", [2, 2**5, 2**10, 2**15, 2**20, [2, 2], [5, 5, 5]]
)
def test_reducescatter_different_array_size(
ray_start_single_node, array_size, tensor_backend, backend
):
world_size = 2
actors, _ = create_collective_workers(world_size, backend=backend)
init_tensors_for_gather_scatter(
actors, array_size=array_size, tensor_backend=tensor_backend
)
results = ray.get([a.do_reducescatter.remote() for a in actors])
for i in range(world_size):
if tensor_backend == "numpy":
assert (
results[i] == np.ones(array_size, dtype=np.float32) * world_size
).all()
else:
assert (
results[i] == torch.ones(array_size, dtype=torch.float32) * world_size
).all()
@pytest.mark.parametrize("backend", [Backend.GLOO])
@pytest.mark.parametrize("dtype", [np.uint8, np.float16, np.float32, np.float64])
def test_reducescatter_different_dtype(ray_start_single_node, dtype, backend):
world_size = 2
actors, _ = create_collective_workers(world_size, backend=backend)
init_tensors_for_gather_scatter(actors, dtype=dtype)
results = ray.get([a.do_reducescatter.remote() for a in actors])
for i in range(world_size):
for j in range(world_size):
assert (results[i] == np.ones(10, dtype=dtype) * world_size).all()
@pytest.mark.parametrize("backend", [Backend.GLOO])
def test_reducescatter_torch_numpy(ray_start_single_node, backend):
world_size = 2
shape = [10, 10]
actors, _ = create_collective_workers(world_size, backend=backend)
# tensor is pytorch, list is numpy
for i, a in enumerate(actors):
t = torch.ones(shape, dtype=torch.float32) * (i + 1)
ray.wait([a.set_buffer.remote(t)])
list_buffer = [np.ones(shape, dtype=np.float32) for _ in range(world_size)]
ray.wait([a.set_list_buffer.remote(list_buffer)])
results = ray.get([a.do_reducescatter.remote() for a in actors])
for i in range(world_size):
assert (results[i] == torch.ones(shape, dtype=torch.float32) * world_size).all()
# tensor is numpy, list is pytorch
for i, a in enumerate(actors):
t = np.ones(shape, dtype=np.float32) * (i + 1)
ray.wait([a.set_buffer.remote(t)])
list_buffer = [
torch.ones(shape, dtype=torch.float32) for _ in range(world_size)
]
ray.wait([a.set_list_buffer.remote(list_buffer)])
results = ray.get([a.do_reducescatter.remote() for a in actors])
for i in range(world_size):
assert (results[i] == np.ones(shape, dtype=np.float32) * world_size).all()
# some tensors in the list are pytorch, some are numpy
for i, a in enumerate(actors):
if i % 2 == 0:
t = torch.ones(shape, dtype=torch.float32) * (i + 1)
else:
t = np.ones(shape, dtype=np.float32) * (i + 1)
ray.wait([a.set_buffer.remote(t)])
list_buffer = []
for j in range(world_size):
if j % 2 == 0:
list_buffer.append(torch.ones(shape, dtype=torch.float32))
else:
list_buffer.append(np.ones(shape, dtype=np.float32))
ray.wait([a.set_list_buffer.remote(list_buffer)])
results = ray.get([a.do_reducescatter.remote() for a in actors])
for i in range(world_size):
if i % 2 == 0:
assert (
results[i] == torch.ones(shape, dtype=torch.float32) * world_size
).all()
else:
assert (results[i] == np.ones(shape, dtype=np.float32) * world_size).all()
# mixed case
for i, a in enumerate(actors):
if i % 2 == 0:
t = torch.ones(shape, dtype=torch.float32) * (i + 1)
else:
t = np.ones(shape, dtype=np.float32) * (i + 1)
ray.wait([a.set_buffer.remote(t)])
list_buffer = []
for j in range(world_size):
if j % 2 == 0:
list_buffer.append(np.ones(shape, dtype=np.float32))
else:
list_buffer.append(torch.ones(shape, dtype=torch.float32))
ray.wait([a.set_list_buffer.remote(list_buffer)])
results = ray.get([a.do_reducescatter.remote() for a in actors])
for i in range(world_size):
if i % 2 == 0:
assert (
results[i] == torch.ones(shape, dtype=torch.float32) * world_size
).all()
else:
assert (results[i] == np.ones(shape, dtype=np.float32) * world_size).all()
if __name__ == "__main__":
import sys
import pytest
sys.exit(pytest.main(["-v", "-x", __file__]))
python/ray/util/collective/tests/single_node_cpu_tests/test_sendrecv.py
+92
View File
@@ -0,0 +1,92 @@
"""Test the send/recv API."""
import numpy as np
import pytest
import ray
from ray.util.collective.tests.cpu_util import create_collective_workers
from ray.util.collective.types import Backend
@pytest.mark.parametrize("backend", [Backend.GLOO])
@pytest.mark.parametrize("group_name", ["default", "test", "123?34!"])
@pytest.mark.parametrize("dst_rank", [0, 1])
@pytest.mark.parametrize(
"array_size", [2, 2**5, 2**10, 2**15, 2**20, [2, 2], [5, 9, 10, 85]]
)
def test_reduce_different_name(
ray_start_single_node, group_name, array_size, dst_rank, backend
):
world_size = 2
actors, _ = create_collective_workers(
num_workers=world_size, group_name=group_name, backend=backend
)
ray.wait(
[
a.set_buffer.remote(np.ones(array_size, dtype=np.float32) * (i + 1))
for i, a in enumerate(actors)
]
)
src_rank = 1 - dst_rank
refs = []
for i, actor in enumerate(actors):
if i != dst_rank:
ref = actor.do_send.remote(group_name, dst_rank)
else:
ref = actor.do_recv.remote(group_name, src_rank)
refs.append(ref)
results = ray.get(refs)
for i in range(world_size):
assert (
results[i] == np.ones(array_size, dtype=np.float32) * (src_rank + 1)
).all()
@pytest.mark.parametrize("backend", [Backend.GLOO])
@pytest.mark.parametrize("dst_rank", [0, 1])
def test_sendrecv_torch_numpy(ray_start_single_node, dst_rank, backend):
import torch
world_size = 2
actors, _ = create_collective_workers(world_size, backend=backend)
ray.wait(
[
actors[1].set_buffer.remote(
torch.ones(
10,
)
* 2
)
]
)
src_rank = 1 - dst_rank
refs = []
for i, actor in enumerate(actors):
if i != dst_rank:
ref = actor.do_send.remote(dst_rank=dst_rank)
else:
ref = actor.do_recv.remote(src_rank=src_rank)
refs.append(ref)
results = ray.get(refs)
if dst_rank == 0:
assert (results[0] == np.ones((10,)) * 2).all()
assert (results[1] == torch.ones((10,)) * 2).all()
else:
assert (results[0] == np.ones((10,))).all()
assert (results[1] == torch.ones((10,))).all()
@pytest.mark.parametrize("backend", [Backend.GLOO])
def test_sendrecv_invalid_rank(ray_start_single_node, backend, dst_rank=3):
world_size = 2
actors, _ = create_collective_workers(world_size, backend=backend)
with pytest.raises(ValueError):
ray.get([a.do_send.remote(dst_rank=dst_rank) for a in actors])
if __name__ == "__main__":
import sys
import pytest
sys.exit(pytest.main(["-v", "-x", __file__]))
python/ray/util/collective/tests/single_node_gpu_tests/__init__.py
View File
python/ray/util/collective/tests/single_node_gpu_tests/test_allgather.py
+138
View File
@@ -0,0 +1,138 @@
"""Test the collective allgather API."""
import cupy as cp
import pytest
import torch
import ray
from ray.util.collective.tests.util import (
create_collective_workers,
init_tensors_for_gather_scatter,
)
@pytest.mark.parametrize("tensor_backend", ["cupy", "torch"])
@pytest.mark.parametrize(
"array_size", [2, 2**5, 2**10, 2**15, 2**20, [2, 2], [5, 5, 5]]
)
def test_allgather_different_array_size(
ray_start_single_node_2_gpus, array_size, tensor_backend
):
world_size = 2
actors, _ = create_collective_workers(world_size)
init_tensors_for_gather_scatter(
actors, array_size=array_size, tensor_backend=tensor_backend
)
results = ray.get([a.do_allgather.remote() for a in actors])
for i in range(world_size):
for j in range(world_size):
if tensor_backend == "cupy":
assert (
results[i][j] == cp.ones(array_size, dtype=cp.float32) * (j + 1)
).all()
else:
assert (
results[i][j]
== torch.ones(array_size, dtype=torch.float32).cuda() * (j + 1)
).all()
@pytest.mark.parametrize("dtype", [cp.uint8, cp.float16, cp.float32, cp.float64])
def test_allgather_different_dtype(ray_start_single_node_2_gpus, dtype):
world_size = 2
actors, _ = create_collective_workers(world_size)
init_tensors_for_gather_scatter(actors, dtype=dtype)
results = ray.get([a.do_allgather.remote() for a in actors])
for i in range(world_size):
for j in range(world_size):
assert (results[i][j] == cp.ones(10, dtype=dtype) * (j + 1)).all()
@pytest.mark.parametrize("length", [0, 1, 2, 3])
def test_unmatched_tensor_list_length(ray_start_single_node_2_gpus, length):
world_size = 2
actors, _ = create_collective_workers(world_size)
list_buffer = [cp.ones(10, dtype=cp.float32) for _ in range(length)]
ray.wait([a.set_list_buffer.remote(list_buffer) for a in actors])
if length != world_size:
with pytest.raises(RuntimeError):
ray.get([a.do_allgather.remote() for a in actors])
else:
ray.get([a.do_allgather.remote() for a in actors])
@pytest.mark.parametrize("shape", [10, 20, [4, 5], [1, 3, 5, 7]])
def test_unmatched_tensor_shape(ray_start_single_node_2_gpus, shape):
world_size = 2
actors, _ = create_collective_workers(world_size)
init_tensors_for_gather_scatter(actors, array_size=10)
list_buffer = [cp.ones(shape, dtype=cp.float32) for _ in range(world_size)]
ray.get([a.set_list_buffer.remote(list_buffer) for a in actors])
if shape != 10:
with pytest.raises(RuntimeError):
ray.get([a.do_allgather.remote() for a in actors])
else:
ray.get([a.do_allgather.remote() for a in actors])
def test_allgather_torch_cupy(ray_start_single_node_2_gpus):
world_size = 2
shape = [10, 10]
actors, _ = create_collective_workers(world_size)
# tensor is pytorch, list is cupy
for i, a in enumerate(actors):
t = torch.ones(shape, dtype=torch.float32).cuda() * (i + 1)
ray.wait([a.set_buffer.remote(t)])
list_buffer = [cp.ones(shape, dtype=cp.float32) for _ in range(world_size)]
ray.wait([a.set_list_buffer.remote(list_buffer)])
results = ray.get([a.do_allgather.remote() for a in actors])
for i in range(world_size):
for j in range(world_size):
assert (results[i][j] == cp.ones(shape, dtype=cp.float32) * (j + 1)).all()
# tensor is cupy, list is pytorch
for i, a in enumerate(actors):
t = cp.ones(shape, dtype=cp.float32) * (i + 1)
ray.wait([a.set_buffer.remote(t)])
list_buffer = [
torch.ones(shape, dtype=torch.float32).cuda() for _ in range(world_size)
]
ray.wait([a.set_list_buffer.remote(list_buffer)])
results = ray.get([a.do_allgather.remote() for a in actors])
for i in range(world_size):
for j in range(world_size):
assert (
results[i][j] == torch.ones(shape, dtype=torch.float32).cuda() * (j + 1)
).all()
# some tensors in the list are pytorch, some are cupy
for i, a in enumerate(actors):
t = cp.ones(shape, dtype=cp.float32) * (i + 1)
ray.wait([a.set_buffer.remote(t)])
list_buffer = []
for j in range(world_size):
if j % 2 == 0:
list_buffer.append(torch.ones(shape, dtype=torch.float32).cuda())
else:
list_buffer.append(cp.ones(shape, dtype=cp.float32))
ray.wait([a.set_list_buffer.remote(list_buffer)])
results = ray.get([a.do_allgather.remote() for a in actors])
for i in range(world_size):
for j in range(world_size):
if j % 2 == 0:
assert (
results[i][j]
== torch.ones(shape, dtype=torch.float32).cuda() * (j + 1)
).all()
else:
assert (
results[i][j] == cp.ones(shape, dtype=cp.float32) * (j + 1)
).all()
if __name__ == "__main__":
import sys
import pytest
sys.exit(pytest.main(["-v", "-x", __file__]))
python/ray/util/collective/tests/single_node_gpu_tests/test_allreduce.py
+170
View File
@@ -0,0 +1,170 @@
"""Test the collective allreduice API."""
import cupy as cp
import pytest
import torch
import ray
from ray.util.collective.tests.util import create_collective_workers
from ray.util.collective.types import ReduceOp
@pytest.mark.parametrize("group_name", ["default", "test", "123?34!"])
def test_allreduce_different_name(ray_start_single_node_2_gpus, group_name):
world_size = 2
actors, _ = create_collective_workers(num_workers=world_size, group_name=group_name)
results = ray.get([a.do_allreduce.remote(group_name) for a in actors])
assert (results[0] == cp.ones((10,), dtype=cp.float32) * world_size).all()
assert (results[1] == cp.ones((10,), dtype=cp.float32) * world_size).all()
@pytest.mark.parametrize("array_size", [2, 2**5, 2**10, 2**15, 2**20])
def test_allreduce_different_array_size(ray_start_single_node_2_gpus, array_size):
world_size = 2
actors, _ = create_collective_workers(world_size)
ray.wait(
[a.set_buffer.remote(cp.ones(array_size, dtype=cp.float32)) for a in actors]
)
results = ray.get([a.do_allreduce.remote() for a in actors])
assert (results[0] == cp.ones((array_size,), dtype=cp.float32) * world_size).all()
assert (results[1] == cp.ones((array_size,), dtype=cp.float32) * world_size).all()
def test_allreduce_destroy(
ray_start_single_node_2_gpus, backend="nccl", group_name="default"
):
world_size = 2
actors, _ = create_collective_workers(world_size)
results = ray.get([a.do_allreduce.remote() for a in actors])
assert (results[0] == cp.ones((10,), dtype=cp.float32) * world_size).all()
assert (results[1] == cp.ones((10,), dtype=cp.float32) * world_size).all()
# destroy the group and try do work, should fail
ray.get([a.destroy_group.remote() for a in actors])
with pytest.raises(RuntimeError):
results = ray.get([a.do_allreduce.remote() for a in actors])
# reinit the same group and all reduce
ray.get(
[
actor.init_group.remote(world_size, i, backend, group_name)
for i, actor in enumerate(actors)
]
)
results = ray.get([a.do_allreduce.remote() for a in actors])
assert (results[0] == cp.ones((10,), dtype=cp.float32) * world_size * 2).all()
assert (results[1] == cp.ones((10,), dtype=cp.float32) * world_size * 2).all()
def test_allreduce_multiple_group(
ray_start_single_node_2_gpus, backend="nccl", num_groups=5
):
world_size = 2
actors, _ = create_collective_workers(world_size)
for group_name in range(1, num_groups):
ray.get(
[
actor.init_group.remote(world_size, i, backend, str(group_name))
for i, actor in enumerate(actors)
]
)
for i in range(num_groups):
group_name = "default" if i == 0 else str(i)
results = ray.get([a.do_allreduce.remote(group_name) for a in actors])
assert (
results[0] == cp.ones((10,), dtype=cp.float32) * (world_size ** (i + 1))
).all()
def test_allreduce_different_op(ray_start_single_node_2_gpus):
world_size = 2
actors, _ = create_collective_workers(world_size)
# check product
ray.wait(
[
a.set_buffer.remote(cp.ones(10, dtype=cp.float32) * (i + 2))
for i, a in enumerate(actors)
]
)
results = ray.get([a.do_allreduce.remote(op=ReduceOp.PRODUCT) for a in actors])
assert (results[0] == cp.ones((10,), dtype=cp.float32) * 6).all()
assert (results[1] == cp.ones((10,), dtype=cp.float32) * 6).all()
# check min
ray.wait(
[
a.set_buffer.remote(cp.ones(10, dtype=cp.float32) * (i + 2))
for i, a in enumerate(actors)
]
)
results = ray.get([a.do_allreduce.remote(op=ReduceOp.MIN) for a in actors])
assert (results[0] == cp.ones((10,), dtype=cp.float32) * 2).all()
assert (results[1] == cp.ones((10,), dtype=cp.float32) * 2).all()
# check max
ray.wait(
[
a.set_buffer.remote(cp.ones(10, dtype=cp.float32) * (i + 2))
for i, a in enumerate(actors)
]
)
results = ray.get([a.do_allreduce.remote(op=ReduceOp.MAX) for a in actors])
assert (results[0] == cp.ones((10,), dtype=cp.float32) * 3).all()
assert (results[1] == cp.ones((10,), dtype=cp.float32) * 3).all()
@pytest.mark.parametrize("dtype", [cp.uint8, cp.float16, cp.float32, cp.float64])
def test_allreduce_different_dtype(ray_start_single_node_2_gpus, dtype):
world_size = 2
actors, _ = create_collective_workers(world_size)
ray.wait([a.set_buffer.remote(cp.ones(10, dtype=dtype)) for a in actors])
results = ray.get([a.do_allreduce.remote() for a in actors])
assert (results[0] == cp.ones((10,), dtype=dtype) * world_size).all()
assert (results[1] == cp.ones((10,), dtype=dtype) * world_size).all()
def test_allreduce_torch_cupy(ray_start_single_node_2_gpus):
# import torch
world_size = 2
actors, _ = create_collective_workers(world_size)
ray.wait(
[
actors[1].set_buffer.remote(
torch.ones(
10,
).cuda()
)
]
)
results = ray.get([a.do_allreduce.remote() for a in actors])
assert (results[0] == cp.ones((10,)) * world_size).all()
ray.wait(
[
actors[0].set_buffer.remote(
torch.ones(
10,
)
)
]
)
ray.wait(
[
actors[1].set_buffer.remote(
cp.ones(
10,
)
)
]
)
with pytest.raises(RuntimeError):
results = ray.get([a.do_allreduce.remote() for a in actors])
if __name__ == "__main__":
import sys
import pytest
sys.exit(pytest.main(["-v", "-x", __file__]))
python/ray/util/collective/tests/single_node_gpu_tests/test_basic_apis.py
+118
View File
@@ -0,0 +1,118 @@
"""Test the collective group APIs."""
import pytest
import ray
from ray.util.collective.tests.util import Worker, create_collective_workers
@pytest.mark.parametrize("group_name", ["default", "test", "123?34!"])
def test_init_two_actors(ray_start_single_node_2_gpus, group_name):
world_size = 2
actors, results = create_collective_workers(world_size, group_name)
for i in range(world_size):
assert results[i]
def test_init_multiple_groups(ray_start_single_node_2_gpus):
world_size = 2
num_groups = 10
actors = [Worker.remote() for i in range(world_size)]
for i in range(num_groups):
group_name = str(i)
init_results = ray.get(
[
actor.init_group.remote(world_size, k, group_name=group_name)
for k, actor in enumerate(actors)
]
)
for j in range(world_size):
assert init_results[j]
def test_get_rank(ray_start_single_node_2_gpus):
world_size = 2
actors, _ = create_collective_workers(world_size)
actor0_rank = ray.get(actors[0].report_rank.remote())
assert actor0_rank == 0
actor1_rank = ray.get(actors[1].report_rank.remote())
assert actor1_rank == 1
# create a second group with a different name,
# and different order of ranks.
new_group_name = "default2"
ray.get(
[
actor.init_group.remote(
world_size, world_size - 1 - i, group_name=new_group_name
)
for i, actor in enumerate(actors)
]
)
actor0_rank = ray.get(actors[0].report_rank.remote(new_group_name))
assert actor0_rank == 1
actor1_rank = ray.get(actors[1].report_rank.remote(new_group_name))
assert actor1_rank == 0
def test_get_collective_group_size(ray_start_single_node_2_gpus):
world_size = 2
actors, _ = create_collective_workers(world_size)
actor0_world_size = ray.get(actors[0].report_world_size.remote())
actor1_world_size = ray.get(actors[1].report_world_size.remote())
assert actor0_world_size == actor1_world_size == world_size
def test_is_group_initialized(ray_start_single_node_2_gpus):
world_size = 2
actors, _ = create_collective_workers(world_size)
# check group is_init
actor0_is_init = ray.get(actors[0].report_is_group_initialized.remote())
assert actor0_is_init
actor0_is_init = ray.get(actors[0].report_is_group_initialized.remote("random"))
assert not actor0_is_init
actor0_is_init = ray.get(actors[0].report_is_group_initialized.remote("123"))
assert not actor0_is_init
actor1_is_init = ray.get(actors[0].report_is_group_initialized.remote())
assert actor1_is_init
actor1_is_init = ray.get(actors[0].report_is_group_initialized.remote("456"))
assert not actor1_is_init
def test_destroy_group(ray_start_single_node_2_gpus):
world_size = 2
actors, _ = create_collective_workers(world_size)
# Now destroy the group at actor0
ray.wait([actors[0].destroy_group.remote()])
actor0_is_init = ray.get(actors[0].report_is_group_initialized.remote())
assert not actor0_is_init
# should go well as the group `random` does not exist at all
ray.wait([actors[0].destroy_group.remote("random")])
actor1_is_init = ray.get(actors[1].report_is_group_initialized.remote())
assert actor1_is_init
ray.wait([actors[1].destroy_group.remote("random")])
actor1_is_init = ray.get(actors[1].report_is_group_initialized.remote())
assert actor1_is_init
ray.wait([actors[1].destroy_group.remote("default")])
actor1_is_init = ray.get(actors[1].report_is_group_initialized.remote())
assert not actor1_is_init
# Now reconstruct the group using the same name
init_results = ray.get(
[actor.init_group.remote(world_size, i) for i, actor in enumerate(actors)]
)
for i in range(world_size):
assert init_results[i]
actor0_is_init = ray.get(actors[0].report_is_group_initialized.remote())
assert actor0_is_init
actor1_is_init = ray.get(actors[0].report_is_group_initialized.remote())
assert actor1_is_init
if __name__ == "__main__":
import sys
import pytest
sys.exit(pytest.main(["-v", "-x", __file__]))
python/ray/util/collective/tests/single_node_gpu_tests/test_broadcast.py
+79
View File
@@ -0,0 +1,79 @@
"""Test the broadcast API."""
import cupy as cp
import pytest
import ray
from ray.util.collective.tests.util import create_collective_workers
@pytest.mark.parametrize("group_name", ["default", "test", "123?34!"])
@pytest.mark.parametrize("src_rank", [0, 1])
def test_broadcast_different_name(ray_start_single_node_2_gpus, group_name, src_rank):
world_size = 2
actors, _ = create_collective_workers(num_workers=world_size, group_name=group_name)
ray.wait(
[
a.set_buffer.remote(cp.ones((10,), dtype=cp.float32) * (i + 2))
for i, a in enumerate(actors)
]
)
results = ray.get(
[
a.do_broadcast.remote(group_name=group_name, src_rank=src_rank)
for a in actors
]
)
for i in range(world_size):
assert (results[i] == cp.ones((10,), dtype=cp.float32) * (src_rank + 2)).all()
@pytest.mark.parametrize("array_size", [2, 2**5, 2**10, 2**15, 2**20])
@pytest.mark.parametrize("src_rank", [0, 1])
def test_broadcast_different_array_size(
ray_start_single_node_2_gpus, array_size, src_rank
):
world_size = 2
actors, _ = create_collective_workers(world_size)
ray.wait(
[
a.set_buffer.remote(cp.ones(array_size, dtype=cp.float32) * (i + 2))
for i, a in enumerate(actors)
]
)
results = ray.get([a.do_broadcast.remote(src_rank=src_rank) for a in actors])
for i in range(world_size):
assert (
results[i] == cp.ones((array_size,), dtype=cp.float32) * (src_rank + 2)
).all()
@pytest.mark.parametrize("src_rank", [0, 1])
def test_broadcast_torch_cupy(ray_start_single_node_2_gpus, src_rank):
import torch
world_size = 2
actors, _ = create_collective_workers(world_size)
ray.wait(
[
actors[1].set_buffer.remote(
torch.ones(
10,
).cuda()
* world_size
)
]
)
results = ray.get([a.do_broadcast.remote(src_rank=src_rank) for a in actors])
if src_rank == 0:
assert (results[0] == cp.ones((10,))).all()
assert (results[1] == torch.ones((10,)).cuda()).all()
else:
assert (results[0] == cp.ones((10,)) * world_size).all()
assert (results[1] == torch.ones((10,)).cuda() * world_size).all()
def test_broadcast_invalid_rank(ray_start_single_node_2_gpus, src_rank=3):
world_size = 2
actors, _ = create_collective_workers(world_size)
with pytest.raises(ValueError):
ray.get([a.do_broadcast.remote(src_rank=src_rank) for a in actors])
python/ray/util/collective/tests/single_node_gpu_tests/test_reduce.py
+151
View File
@@ -0,0 +1,151 @@
"""Test the reduce API."""
import cupy as cp
import pytest
import ray
from ray.util.collective.tests.util import create_collective_workers
from ray.util.collective.types import ReduceOp
@pytest.mark.parametrize("group_name", ["default", "test", "123?34!"])
@pytest.mark.parametrize("dst_rank", [0, 1])
def test_reduce_different_name(ray_start_single_node_2_gpus, group_name, dst_rank):
world_size = 2
actors, _ = create_collective_workers(num_workers=world_size, group_name=group_name)
results = ray.get([a.do_reduce.remote(group_name, dst_rank) for a in actors])
for i in range(world_size):
if i == dst_rank:
assert (results[i] == cp.ones((10,), dtype=cp.float32) * world_size).all()
else:
assert (results[i] == cp.ones((10,), dtype=cp.float32)).all()
@pytest.mark.parametrize("array_size", [2, 2**5, 2**10, 2**15, 2**20])
@pytest.mark.parametrize("dst_rank", [0, 1])
def test_reduce_different_array_size(
ray_start_single_node_2_gpus, array_size, dst_rank
):
world_size = 2
actors, _ = create_collective_workers(world_size)
ray.wait(
[a.set_buffer.remote(cp.ones(array_size, dtype=cp.float32)) for a in actors]
)
results = ray.get([a.do_reduce.remote(dst_rank=dst_rank) for a in actors])
for i in range(world_size):
if i == dst_rank:
assert (
results[i] == cp.ones((array_size,), dtype=cp.float32) * world_size
).all()
else:
assert (results[i] == cp.ones((array_size,), dtype=cp.float32)).all()
@pytest.mark.parametrize("dst_rank", [0, 1])
def test_reduce_multiple_group(ray_start_single_node_2_gpus, dst_rank, num_groups=5):
world_size = 2
actors, _ = create_collective_workers(world_size)
for group_name in range(1, num_groups):
ray.get(
[
actor.init_group.remote(world_size, i, "nccl", str(group_name))
for i, actor in enumerate(actors)
]
)
for i in range(num_groups):
group_name = "default" if i == 0 else str(i)
results = ray.get(
[
a.do_reduce.remote(dst_rank=dst_rank, group_name=group_name)
for a in actors
]
)
for j in range(world_size):
if j == dst_rank:
assert (results[j] == cp.ones((10,), dtype=cp.float32) * (i + 2)).all()
else:
assert (results[j] == cp.ones((10,), dtype=cp.float32)).all()
@pytest.mark.parametrize("dst_rank", [0, 1])
def test_reduce_different_op(ray_start_single_node_2_gpus, dst_rank):
world_size = 2
actors, _ = create_collective_workers(world_size)
# check product
ray.wait(
[
a.set_buffer.remote(cp.ones(10, dtype=cp.float32) * (i + 2))
for i, a in enumerate(actors)
]
)
results = ray.get(
[a.do_reduce.remote(dst_rank=dst_rank, op=ReduceOp.PRODUCT) for a in actors]
)
for i in range(world_size):
if i == dst_rank:
assert (results[i] == cp.ones((10,), dtype=cp.float32) * 6).all()
else:
assert (results[i] == cp.ones((10,), dtype=cp.float32) * (i + 2)).all()
# check min
ray.wait(
[
a.set_buffer.remote(cp.ones(10, dtype=cp.float32) * (i + 2))
for i, a in enumerate(actors)
]
)
results = ray.get(
[a.do_reduce.remote(dst_rank=dst_rank, op=ReduceOp.MIN) for a in actors]
)
for i in range(world_size):
if i == dst_rank:
assert (results[i] == cp.ones((10,), dtype=cp.float32) * 2).all()
else:
assert (results[i] == cp.ones((10,), dtype=cp.float32) * (i + 2)).all()
# check max
ray.wait(
[
a.set_buffer.remote(cp.ones(10, dtype=cp.float32) * (i + 2))
for i, a in enumerate(actors)
]
)
results = ray.get(
[a.do_reduce.remote(dst_rank=dst_rank, op=ReduceOp.MAX) for a in actors]
)
for i in range(world_size):
if i == dst_rank:
assert (results[i] == cp.ones((10,), dtype=cp.float32) * 3).all()
else:
assert (results[i] == cp.ones((10,), dtype=cp.float32) * (i + 2)).all()
@pytest.mark.parametrize("dst_rank", [0, 1])
def test_reduce_torch_cupy(ray_start_single_node_2_gpus, dst_rank):
import torch
world_size = 2
actors, _ = create_collective_workers(world_size)
ray.wait(
[
actors[1].set_buffer.remote(
torch.ones(
10,
).cuda()
)
]
)
results = ray.get([a.do_reduce.remote(dst_rank=dst_rank) for a in actors])
if dst_rank == 0:
assert (results[0] == cp.ones((10,)) * world_size).all()
assert (results[1] == torch.ones((10,)).cuda()).all()
else:
assert (results[0] == cp.ones((10,))).all()
assert (results[1] == torch.ones((10,)).cuda() * world_size).all()
def test_reduce_invalid_rank(ray_start_single_node_2_gpus, dst_rank=3):
world_size = 2
actors, _ = create_collective_workers(world_size)
with pytest.raises(ValueError):
ray.get([a.do_reduce.remote(dst_rank=dst_rank) for a in actors])
python/ray/util/collective/tests/single_node_gpu_tests/test_reducescatter.py
+130
View File
@@ -0,0 +1,130 @@
"""Test the collective reducescatter API."""
import cupy as cp
import pytest
import torch
import ray
from ray.util.collective.tests.util import (
create_collective_workers,
init_tensors_for_gather_scatter,
)
@pytest.mark.parametrize("tensor_backend", ["cupy", "torch"])
@pytest.mark.parametrize(
"array_size", [2, 2**5, 2**10, 2**15, 2**20, [2, 2], [5, 5, 5]]
)
def test_reducescatter_different_array_size(
ray_start_single_node_2_gpus, array_size, tensor_backend
):
world_size = 2
actors, _ = create_collective_workers(world_size)
init_tensors_for_gather_scatter(
actors, array_size=array_size, tensor_backend=tensor_backend
)
results = ray.get([a.do_reducescatter.remote() for a in actors])
for i in range(world_size):
if tensor_backend == "cupy":
assert (
results[i] == cp.ones(array_size, dtype=cp.float32) * world_size
).all()
else:
assert (
results[i]
== torch.ones(array_size, dtype=torch.float32).cuda() * world_size
).all()
@pytest.mark.parametrize("dtype", [cp.uint8, cp.float16, cp.float32, cp.float64])
def test_reducescatter_different_dtype(ray_start_single_node_2_gpus, dtype):
world_size = 2
actors, _ = create_collective_workers(world_size)
init_tensors_for_gather_scatter(actors, dtype=dtype)
results = ray.get([a.do_reducescatter.remote() for a in actors])
for i in range(world_size):
for j in range(world_size):
assert (results[i] == cp.ones(10, dtype=dtype) * world_size).all()
def test_reducescatter_torch_cupy(ray_start_single_node_2_gpus):
world_size = 2
shape = [10, 10]
actors, _ = create_collective_workers(world_size)
# tensor is pytorch, list is cupy
for i, a in enumerate(actors):
t = torch.ones(shape, dtype=torch.float32).cuda() * (i + 1)
ray.wait([a.set_buffer.remote(t)])
list_buffer = [cp.ones(shape, dtype=cp.float32) for _ in range(world_size)]
ray.wait([a.set_list_buffer.remote(list_buffer)])
results = ray.get([a.do_reducescatter.remote() for a in actors])
for i in range(world_size):
assert (
results[i] == torch.ones(shape, dtype=torch.float32).cuda() * world_size
).all()
# tensor is cupy, list is pytorch
for i, a in enumerate(actors):
t = cp.ones(shape, dtype=cp.float32) * (i + 1)
ray.wait([a.set_buffer.remote(t)])
list_buffer = [
torch.ones(shape, dtype=torch.float32).cuda() for _ in range(world_size)
]
ray.wait([a.set_list_buffer.remote(list_buffer)])
results = ray.get([a.do_reducescatter.remote() for a in actors])
for i in range(world_size):
assert (results[i] == cp.ones(shape, dtype=cp.float32) * world_size).all()
# some tensors in the list are pytorch, some are cupy
for i, a in enumerate(actors):
if i % 2 == 0:
t = torch.ones(shape, dtype=torch.float32).cuda() * (i + 1)
else:
t = cp.ones(shape, dtype=cp.float32) * (i + 1)
ray.wait([a.set_buffer.remote(t)])
list_buffer = []
for j in range(world_size):
if j % 2 == 0:
list_buffer.append(torch.ones(shape, dtype=torch.float32).cuda())
else:
list_buffer.append(cp.ones(shape, dtype=cp.float32))
ray.wait([a.set_list_buffer.remote(list_buffer)])
results = ray.get([a.do_reducescatter.remote() for a in actors])
for i in range(world_size):
if i % 2 == 0:
assert (
results[i] == torch.ones(shape, dtype=torch.float32).cuda() * world_size
).all()
else:
assert (results[i] == cp.ones(shape, dtype=cp.float32) * world_size).all()
# mixed case
for i, a in enumerate(actors):
if i % 2 == 0:
t = torch.ones(shape, dtype=torch.float32).cuda() * (i + 1)
else:
t = cp.ones(shape, dtype=cp.float32) * (i + 1)
ray.wait([a.set_buffer.remote(t)])
list_buffer = []
for j in range(world_size):
if j % 2 == 0:
list_buffer.append(cp.ones(shape, dtype=cp.float32))
else:
list_buffer.append(torch.ones(shape, dtype=torch.float32).cuda())
ray.wait([a.set_list_buffer.remote(list_buffer)])
results = ray.get([a.do_reducescatter.remote() for a in actors])
for i in range(world_size):
if i % 2 == 0:
assert (
results[i] == torch.ones(shape, dtype=torch.float32).cuda() * world_size
).all()
else:
assert (results[i] == cp.ones(shape, dtype=cp.float32) * world_size).all()
if __name__ == "__main__":
import sys
import pytest
sys.exit(pytest.main(["-v", "-x", __file__]))
python/ray/util/collective/tests/single_node_gpu_tests/test_sendrecv.py
+78
View File
@@ -0,0 +1,78 @@
"""Test the send/recv API."""
import cupy as cp
import pytest
import ray
from ray.util.collective.tests.util import create_collective_workers
@pytest.mark.parametrize("group_name", ["default", "test", "123?34!"])
@pytest.mark.parametrize("dst_rank", [0, 1])
@pytest.mark.parametrize(
"array_size", [2, 2**5, 2**10, 2**15, 2**20, [2, 2], [5, 9, 10, 85]]
)
def test_reduce_different_name(
ray_start_single_node_2_gpus, group_name, array_size, dst_rank
):
world_size = 2
actors, _ = create_collective_workers(num_workers=world_size, group_name=group_name)
ray.wait(
[
a.set_buffer.remote(cp.ones(array_size, dtype=cp.float32) * (i + 1))
for i, a in enumerate(actors)
]
)
src_rank = 1 - dst_rank
refs = []
for i, actor in enumerate(actors):
if i != dst_rank:
ref = actor.do_send.remote(group_name, dst_rank)
else:
ref = actor.do_recv.remote(group_name, src_rank)
refs.append(ref)
results = ray.get(refs)
for i in range(world_size):
assert (
results[i] == cp.ones(array_size, dtype=cp.float32) * (src_rank + 1)
).all()
@pytest.mark.parametrize("dst_rank", [0, 1])
def test_sendrecv_torch_cupy(ray_start_single_node_2_gpus, dst_rank):
import torch
world_size = 2
actors, _ = create_collective_workers(world_size)
ray.wait(
[
actors[1].set_buffer.remote(
torch.ones(
10,
).cuda()
* 2
)
]
)
src_rank = 1 - dst_rank
refs = []
for i, actor in enumerate(actors):
if i != dst_rank:
ref = actor.do_send.remote(dst_rank=dst_rank)
else:
ref = actor.do_recv.remote(src_rank=src_rank)
refs.append(ref)
results = ray.get(refs)
if dst_rank == 0:
assert (results[0] == cp.ones((10,)) * 2).all()
assert (results[1] == torch.ones((10,)).cuda() * 2).all()
else:
assert (results[0] == cp.ones((10,))).all()
assert (results[1] == torch.ones((10,)).cuda()).all()
def test_sendrecv_invalid_rank(ray_start_single_node_2_gpus, dst_rank=3):
world_size = 2
actors, _ = create_collective_workers(world_size)
with pytest.raises(ValueError):
ray.get([a.do_send.remote(dst_rank=dst_rank) for a in actors])
python/ray/util/collective/tests/util.py
+373
View File
@@ -0,0 +1,373 @@
import logging
import cupy as cp
import torch
import ray
import ray.util.collective as col
from ray.util.collective.collective_group.nccl_util import get_num_gpus
from ray.util.collective.types import Backend, ReduceOp
logger = logging.getLogger(__name__)
@ray.remote(num_gpus=1)
class Worker:
def __init__(self):
self.buffer = None
self.list_buffer = None
def init_tensors(self):
self.buffer = cp.ones((10,), dtype=cp.float32)
self.list_buffer = [cp.ones((10,), dtype=cp.float32) for _ in range(2)]
cp.cuda.Stream.null.synchronize()
return True
def init_group(self, world_size, rank, backend=Backend.NCCL, group_name="default"):
col.init_collective_group(world_size, rank, backend, group_name)
return True
def set_buffer(self, data):
self.buffer = data
return self.buffer
def get_buffer(self):
return self.buffer
def set_list_buffer(self, list_of_arrays):
self.list_buffer = list_of_arrays
return self.list_buffer
def do_allreduce(self, group_name="default", op=ReduceOp.SUM):
col.allreduce(self.buffer, group_name, op)
return self.buffer
def do_reduce(self, group_name="default", dst_rank=0, op=ReduceOp.SUM):
col.reduce(self.buffer, dst_rank, group_name, op)
return self.buffer
def do_broadcast(self, group_name="default", src_rank=0):
col.broadcast(self.buffer, src_rank, group_name)
return self.buffer
def do_allgather(self, group_name="default"):
col.allgather(self.list_buffer, self.buffer, group_name)
return self.list_buffer
def do_reducescatter(self, group_name="default", op=ReduceOp.SUM):
col.reducescatter(self.buffer, self.list_buffer, group_name, op)
return self.buffer
def do_send(self, group_name="default", dst_rank=0):
col.send(self.buffer, dst_rank, group_name)
return self.buffer
def do_recv(self, group_name="default", src_rank=0):
col.recv(self.buffer, src_rank, group_name)
return self.buffer
def destroy_group(self, group_name="default"):
col.destroy_collective_group(group_name)
return True
def report_rank(self, group_name="default"):
rank = col.get_rank(group_name)
return rank
def report_world_size(self, group_name="default"):
ws = col.get_collective_group_size(group_name)
return ws
def report_nccl_availability(self):
avail = col.nccl_available()
return avail
def report_gloo_availability(self):
avail = col.gloo_available()
return avail
def report_is_group_initialized(self, group_name="default"):
is_init = col.is_group_initialized(group_name)
return is_init
def create_collective_workers(num_workers=2, group_name="default", backend="nccl"):
actors = [None] * num_workers
for i in range(num_workers):
actor = Worker.remote()
ray.get([actor.init_tensors.remote()])
actors[i] = actor
world_size = num_workers
init_results = ray.get(
[
actor.init_group.remote(world_size, i, backend, group_name)
for i, actor in enumerate(actors)
]
)
return actors, init_results
def init_tensors_for_gather_scatter(
actors, array_size=10, dtype=cp.float32, tensor_backend="cupy"
):
world_size = len(actors)
for i, a in enumerate(actors):
if tensor_backend == "cupy":
t = cp.ones(array_size, dtype=dtype) * (i + 1)
elif tensor_backend == "torch":
t = torch.ones(array_size, dtype=torch.float32).cuda() * (i + 1)
else:
raise RuntimeError("Unsupported tensor backend.")
ray.get([a.set_buffer.remote(t)])
if tensor_backend == "cupy":
list_buffer = [cp.ones(array_size, dtype=dtype) for _ in range(world_size)]
elif tensor_backend == "torch":
list_buffer = [
torch.ones(array_size, dtype=torch.float32).cuda()
for _ in range(world_size)
]
else:
raise RuntimeError("Unsupported tensor backend.")
ray.get([a.set_list_buffer.remote(list_buffer) for a in actors])
@ray.remote(num_gpus=2)
class MultiGPUWorker:
def __init__(self):
self.buffer0 = None
self.buffer1 = None
self.list_buffer0 = None
self.list_buffer1 = None
def __del__(self):
self.buffer0 = None
self.buffer1 = None
self.list_buffer0 = None
self.list_buffer1 = None
def init_tensors(self):
with cp.cuda.Device(0):
self.buffer0 = cp.ones((10,), dtype=cp.float32)
self.list_buffer0 = [cp.ones((10,), dtype=cp.float32) for _ in range(4)]
with cp.cuda.Device(1):
self.buffer1 = cp.ones((10,), dtype=cp.float32)
self.list_buffer1 = [cp.ones((10,), dtype=cp.float32) for _ in range(4)]
cp.cuda.Stream.null.synchronize()
return True
def init_group(self, world_size, rank, backend=Backend.NCCL, group_name="default"):
col.init_collective_group(world_size, rank, backend, group_name)
return True
def set_buffer(
self,
size,
value0=1.0,
value1=1.0,
dtype=cp.float32,
tensor_type0="cupy",
tensor_type1="cupy",
):
if tensor_type0 == "cupy":
with cp.cuda.Device(0):
self.buffer0 = cp.ones(size, dtype=dtype) * value0
elif tensor_type0 == "torch":
self.buffer0 = torch.ones(size, dtype=torch.float32).cuda(0) * value0
else:
raise RuntimeError()
if tensor_type1 == "cupy":
with cp.cuda.Device(1):
self.buffer1 = cp.ones(size, dtype=dtype) * value1
elif tensor_type1 == "torch":
self.buffer1 = torch.ones(size, dtype=torch.float32).cuda(1) * value1
else:
raise RuntimeError()
cp.cuda.Device(0).synchronize()
cp.cuda.Device(1).synchronize()
# cp.cuda.Stream.null.synchronize()
return True
def set_list_buffer(
self,
size,
value0=1.0,
value1=1.0,
dtype=cp.float32,
tensor_type0="cupy",
tensor_type1="cupy",
):
if tensor_type0 == "cupy":
with cp.cuda.Device(0):
self.list_buffer0 = [
cp.ones(size, dtype=dtype) * value0 for _ in range(4)
]
elif tensor_type0 == "torch":
self.list_buffer0 = [
torch.ones(size, dtype=torch.float32).cuda(0) * value0 for _ in range(4)
]
else:
raise RuntimeError()
if tensor_type1 == "cupy":
with cp.cuda.Device(1):
self.list_buffer1 = [
cp.ones(size, dtype=dtype) * value1 for _ in range(4)
]
elif tensor_type1 == "torch":
self.list_buffer1 = [
torch.ones(size, dtype=torch.float32).cuda(1) * value1 for _ in range(4)
]
else:
raise RuntimeError()
cp.cuda.Device(0).synchronize()
cp.cuda.Device(1).synchronize()
return True
@ray.method(num_returns=2)
def get_buffer(self):
return self.buffer0, self.buffer1
def do_allreduce_multigpu(self, group_name="default", op=ReduceOp.SUM):
col.allreduce_multigpu([self.buffer0, self.buffer1], group_name, op)
cp.cuda.Device(0).synchronize()
cp.cuda.Device(1).synchronize()
return self.buffer0
def do_reduce_multigpu(
self, group_name="default", dst_rank=0, dst_gpu_index=0, op=ReduceOp.SUM
):
col.reduce_multigpu(
[self.buffer0, self.buffer1], dst_rank, dst_gpu_index, group_name, op
)
cp.cuda.Device(0).synchronize()
cp.cuda.Device(1).synchronize()
return self.buffer0, self.buffer1
def do_broadcast_multigpu(self, group_name="default", src_rank=0, src_gpu_index=0):
col.broadcast_multigpu(
[self.buffer0, self.buffer1], src_rank, src_gpu_index, group_name
)
return self.buffer0, self.buffer1
def do_allgather_multigpu(self, group_name="default"):
col.allgather_multigpu(
[self.list_buffer0, self.list_buffer1],
[self.buffer0, self.buffer1],
group_name,
)
cp.cuda.Device(0).synchronize()
cp.cuda.Device(1).synchronize()
return self.list_buffer0, self.list_buffer1
def do_reducescatter_multigpu(self, group_name="default", op=ReduceOp.SUM):
col.reducescatter_multigpu(
[self.buffer0, self.buffer1],
[self.list_buffer0, self.list_buffer1],
group_name,
op,
)
cp.cuda.Device(0).synchronize()
cp.cuda.Device(1).synchronize()
return self.buffer0, self.buffer1
def do_send_multigpu(
self, group_name="default", dst_rank=0, dst_gpu_index=0, src_gpu_index=0
):
if src_gpu_index == 0:
col.send_multigpu(self.buffer0, dst_rank, dst_gpu_index, group_name)
cp.cuda.Device(0).synchronize()
return self.buffer0
elif src_gpu_index == 1:
col.send_multigpu(self.buffer1, dst_rank, dst_gpu_index, group_name)
cp.cuda.Device(1).synchronize()
return self.buffer1
else:
raise RuntimeError()
def do_recv_multigpu(
self, group_name="default", src_rank=0, src_gpu_index=0, dst_gpu_index=0
):
if dst_gpu_index == 0:
col.recv_multigpu(self.buffer0, src_rank, src_gpu_index, group_name)
cp.cuda.Device(0).synchronize()
return self.buffer0
elif dst_gpu_index == 1:
col.recv_multigpu(self.buffer1, src_rank, src_gpu_index, group_name)
cp.cuda.Device(1).synchronize()
return self.buffer1
else:
raise RuntimeError()
def destroy_group(self, group_name="default"):
col.destroy_collective_group(group_name)
return True
def report_rank(self, group_name="default"):
rank = col.get_rank(group_name)
return rank
def report_world_size(self, group_name="default"):
ws = col.get_collective_group_size(group_name)
return ws
def report_nccl_availability(self):
avail = col.nccl_available()
return avail
def report_gloo_availability(self):
avail = col.gloo_available()
return avail
def report_is_group_initialized(self, group_name="default"):
is_init = col.is_group_initialized(group_name)
return is_init
def report_num_gpus(self):
n_gpus = get_num_gpus()
return n_gpus
def create_collective_multigpu_workers(
num_workers=2, group_name="default", backend="nccl"
):
actors = [None] * num_workers
for i in range(num_workers):
actor = MultiGPUWorker.remote()
ray.get([actor.set_buffer.remote([10])], timeout=10)
ray.get([actor.set_list_buffer.remote([10])], timeout=10)
actors[i] = actor
world_size = num_workers
init_results = ray.get(
[
actor.init_group.remote(world_size, i, backend, group_name)
for i, actor in enumerate(actors)
]
)
return actors, init_results
def init_tensors_for_gather_scatter_multigpu(
actors, array_size=10, tensor_backend="cupy"
):
for i, a in enumerate(actors):
if tensor_backend == "cupy":
ray.get([a.set_buffer.remote(array_size)])
ray.get([a.set_list_buffer.remote(array_size)])
elif tensor_backend == "torch":
ray.get(
[
a.set_buffer.remote(
array_size, tensor_type0="torch", tensor_type1="torch"
)
]
)
ray.get(
[
a.set_list_buffer.remote(
array_size, tensor_type0="torch", tensor_type1="torch"
)
]
)
else:
raise RuntimeError("Unsupported tensor backend.")
python/ray/util/collective/types.py
+122
View File
@@ -0,0 +1,122 @@
"""Types conversion between different backends."""
from dataclasses import dataclass
from datetime import timedelta
from enum import Enum
from typing import TYPE_CHECKING
_NUMPY_AVAILABLE = True
_TORCH_AVAILABLE = True
_CUPY_AVAILABLE = True
if TYPE_CHECKING:
pass
try:
import torch as th # noqa: F401
except ImportError:
_TORCH_AVAILABLE = False
try:
import cupy as cp # noqa: F401
except ImportError:
_CUPY_AVAILABLE = False
def cupy_available():
return _CUPY_AVAILABLE
def torch_available():
return _TORCH_AVAILABLE
class Backend(object):
"""A class to represent different backends."""
NCCL = "NCCL"
GLOO = "GLOO"
UNRECOGNIZED = "unrecognized"
def __new__(cls, name: str):
upper_name = name.upper()
backend = getattr(Backend, upper_name, Backend.UNRECOGNIZED)
if backend == Backend.UNRECOGNIZED:
if upper_name == "TORCH_GLOO":
return Backend.GLOO
raise ValueError(
"Unrecognized backend: '{}'. Only NCCL and GLOO are supported".format(
name
)
)
return backend
class ReduceOp(Enum):
SUM = 0
PRODUCT = 1
MIN = 2
MAX = 3
unset_timeout_ms = timedelta(milliseconds=-1)
@dataclass
class AllReduceOptions:
reduceOp = ReduceOp.SUM
timeout_ms = unset_timeout_ms
@dataclass
class BarrierOptions:
timeout_ms = unset_timeout_ms
@dataclass
class ReduceOptions:
reduceOp = ReduceOp.SUM
root_rank = 0
root_tensor = 0 # index for multi-gpu reduce operations
timeout_ms = unset_timeout_ms
@dataclass
class AllGatherOptions:
timeout_ms = unset_timeout_ms
#
# @dataclass
# class GatherOptions:
# root_rank = 0
# timeout = unset_timeout
@dataclass
class BroadcastOptions:
root_rank = 0
root_tensor = 0
timeout_ms = unset_timeout_ms
@dataclass
class ReduceScatterOptions:
reduceOp = ReduceOp.SUM
timeout_ms = unset_timeout_ms
@dataclass
class SendOptions:
dst_rank = 0
dst_gpu_index = 0
n_elements = 0
timeout_ms = unset_timeout_ms
@dataclass
class RecvOptions:
src_rank = 0
src_gpu_index = 0
n_elements = 0
unset_timeout_ms = unset_timeout_ms
python/ray/util/collective/util.py
+85
View File
@@ -0,0 +1,85 @@
"""Some utility class for Collectives."""
import asyncio
import logging
import ray
logger = logging.getLogger(__name__)
@ray.remote
class NCCLUniqueIDStore:
"""NCCLUniqueID Store as a named actor class.
Args:
name: the unique name for this named actor.
Attributes:
name: the unique name for this named actor.
nccl_id: the NCCLUniqueID held in this store.
"""
def __init__(self, name: str):
self.name = name
self.nccl_id = None
self.event = asyncio.Event()
async def set_id(self, uid: bytes):
"""
Initialize the NCCL unique ID for this store.
Args:
uid: the unique ID generated via the NCCL generate_communicator_id API.
Returns:
The NCCL unique ID set.
"""
self.nccl_id = uid
self.event.set()
return uid
async def wait_and_get_id(self):
"""Wait for the NCCL unique ID to be set and return it."""
await self.event.wait()
return self.nccl_id
def get_id(self):
"""Get the NCCL unique ID held in this store."""
if not self.nccl_id:
logger.warning(
"The NCCL ID has not been set yet for store {}.".format(self.name)
)
return self.nccl_id
@ray.remote
class Info:
"""Store the group information created via `create_collective_group`.
Note: Should be used as a NamedActor.
"""
def __init__(self):
self.ids = None
self.world_size = -1
self.rank = -1
self.backend = None
self.gloo_timeout = 30000
def set_info(self, ids, world_size, rank, backend, gloo_timeout):
"""Store collective information."""
self.ids = ids
self.world_size = world_size
self.rank = rank
self.backend = backend
self.gloo_timeout = gloo_timeout
def get_info(self):
"""Get previously stored collective information."""
return (
self.ids,
self.world_size,
self.rank,
self.backend,
self.gloo_timeout,
)
python/ray/util/common.py
+1
View File
@@ -0,0 +1 @@
INT32_MAX = (2**31) - 1
python/ray/util/dask/BUILD.bazel
+10
View File
@@ -0,0 +1,10 @@
load("@rules_python//python:defs.bzl", "py_library")
py_library(
name = "dask_lib",
srcs = glob(
["**/*.py"],
exclude = ["tests/*.py"],
),
visibility = ["__subpackages__"],
)
python/ray/util/dask/__init__.py
+75
View File
@@ -0,0 +1,75 @@
import dask
from packaging.version import Version
# Version(dask.__version__) becomes "0" during doc builds.
if Version(dask.__version__) != Version("0") and Version(dask.__version__) < Version(
"2024.11.0"
):
# Dask on Ray doesn't work if Dask version is less than 2024.11.0.
raise ImportError(
"Dask on Ray requires Dask version 2024.11.0 or later. "
"Please upgrade your Dask installation."
)
from .callbacks import (
ProgressBarCallback,
RayDaskCallback,
local_ray_callbacks,
unpack_ray_callbacks,
)
from .optimizations import dataframe_optimize
from .scheduler import (
disable_dask_on_ray,
enable_dask_on_ray,
ray_dask_get,
ray_dask_get_sync,
)
dask_persist = dask.persist
def ray_dask_persist(*args, **kwargs):
kwargs["ray_persist"] = True
return dask_persist(*args, **kwargs)
ray_dask_persist.__doc__ = dask_persist.__doc__
dask_persist_mixin = dask.base.DaskMethodsMixin.persist
def ray_dask_persist_mixin(self, **kwargs):
kwargs["ray_persist"] = True
return dask_persist_mixin(self, **kwargs)
ray_dask_persist_mixin.__doc__ = dask_persist_mixin.__doc__
# We patch dask in order to inject a kwarg into its `dask.persist()` calls,
# which the Dask-on-Ray scheduler needs.
# FIXME(Clark): Monkey patching is bad and we should try to avoid this.
def patch_dask(ray_dask_persist, ray_dask_persist_mixin):
dask.persist = ray_dask_persist
dask.base.DaskMethodsMixin.persist = ray_dask_persist_mixin
patch_dask(ray_dask_persist, ray_dask_persist_mixin)
__all__ = [
# Config
"enable_dask_on_ray",
"disable_dask_on_ray",
# Schedulers
"ray_dask_get",
"ray_dask_get_sync",
# Helpers
"ray_dask_persist",
# Callbacks
"RayDaskCallback",
"local_ray_callbacks",
"unpack_ray_callbacks",
# Optimizations
"dataframe_optimize",
"ProgressBarCallback",
]
python/ray/util/dask/callbacks.py
+310
View File
@@ -0,0 +1,310 @@
import contextlib
from collections import defaultdict, namedtuple
from datetime import datetime
from typing import Any, Dict, List, Optional
from dask.callbacks import Callback
import ray
# The names of the Ray-specific callbacks. These are the kwarg names that
# RayDaskCallback will accept on construction, and is considered the
# source-of-truth for what Ray-specific callbacks exist.
CBS = (
"ray_presubmit",
"ray_postsubmit",
"ray_pretask",
"ray_posttask",
"ray_postsubmit_all",
"ray_finish",
)
# The Ray-specific callback method names for RayDaskCallback.
CB_FIELDS = tuple("_" + field for field in CBS)
# The Ray-specific callbacks that we do _not_ wish to drop from RayCallbacks
# if not given on a RayDaskCallback instance (will be filled with None
# instead).
CBS_DONT_DROP = {"ray_pretask", "ray_posttask"}
# The Ray-specific callbacks for a single RayDaskCallback.
RayCallback = namedtuple("RayCallback", " ".join(CBS))
# The Ray-specific callbacks for one or more RayDaskCallbacks.
RayCallbacks = namedtuple("RayCallbacks", " ".join([field + "_cbs" for field in CBS]))
class RayDaskCallback(Callback):
"""
Extends Dask's `Callback` class with Ray-specific hooks. When instantiating
or subclassing this class, both the normal Dask hooks (e.g. pretask,
posttask, etc.) and the Ray-specific hooks can be provided.
See `dask.callbacks.Callback` for usage.
Caveats: Any Dask-Ray scheduler must bring the Ray-specific callbacks into
context using the `local_ray_callbacks` context manager, since the built-in
`local_callbacks` context manager provided by Dask isn't aware of this
class.
"""
# Set of active Ray-specific callbacks.
ray_active = set()
def __init__(self, **kwargs):
for cb in CBS:
cb_func = kwargs.pop(cb, None)
if cb_func is not None:
setattr(self, "_" + cb, cb_func)
super().__init__(**kwargs)
@property
def _ray_callback(self):
return RayCallback(*[getattr(self, field, None) for field in CB_FIELDS])
def __enter__(self):
self._ray_cm = add_ray_callbacks(self)
self._ray_cm.__enter__()
super().__enter__()
return self
def __exit__(self, *args):
super().__exit__(*args)
self._ray_cm.__exit__(*args)
def register(self):
type(self).ray_active.add(self._ray_callback)
super().register()
def unregister(self):
type(self).ray_active.remove(self._ray_callback)
super().unregister()
def _ray_presubmit(
self, task: Any, key: Any, deps: Dict[Any, Any]
) -> Optional[Any]:
"""Run before submitting a Ray task.
If this callback returns a non-`None` value, Ray does _not_ create
a task and uses this value as the would-be task's result value.
Args:
task: A Dask task, where the first tuple item is
the task function, and the remaining tuple items are
the task arguments, which are either the actual argument values,
or Dask keys into the deps dictionary whose
corresponding values are the argument values.
key: The Dask graph key for the given task.
deps: The dependencies of this task.
Returns:
Either None, in which case Ray submits a task, or
a non-None value, in which case Ray task doesn't submit
a task and uses this return value as the
would-be task result value.
"""
pass
def _ray_postsubmit(
self,
task: Any,
key: Any,
deps: Dict[Any, Any],
object_ref: ray.ObjectRef,
):
"""Run after submitting a Ray task.
Args:
task: A Dask task, where the first tuple item is
the task function, and the remaining tuple items are
the task arguments, which are either the actual argument values,
or Dask keys into the deps dictionary whose
corresponding values are the argument values.
key: The Dask graph key for the given task.
deps: The dependencies of this task.
object_ref: The object reference for the
return value of the Ray task.
"""
pass
def _ray_pretask(self, key: Any, object_refs: List[ray.ObjectRef]):
"""Run before executing a Dask task within a Ray task.
This method executes after Ray submits the task within a Ray
worker. The return value of this method is passed to the
_ray_posttask callback, if provided.
Args:
key: The Dask graph key for the Dask task.
object_refs: The object references
for the arguments of the Ray task.
"""
pass
def _ray_posttask(self, key: Any, result: Any, pre_state: Any):
"""Run after executing a Dask task within a Ray task.
This method executes within a Ray worker. This callback receives the
return value of the _ray_pretask callback, if provided.
Args:
key: The Dask graph key for the Dask task.
result: The task result value.
pre_state: The return value of the corresponding
_ray_pretask callback, if said callback is defined.
"""
pass
def _ray_postsubmit_all(self, object_refs: List[ray.ObjectRef], dsk: Any):
"""Run after Ray submits all tasks.
Args:
object_refs: The object references
for the output (leaf) Ray tasks of the task graph.
dsk: The Dask graph.
"""
pass
def _ray_finish(self, result: Any):
"""Run after Ray finishes executing all Ray tasks and returns the final
result.
Args:
result: The final result (output) of the Dask
computation, before any repackaging is done by
Dask collection-specific post-compute callbacks.
"""
pass
class add_ray_callbacks:
def __init__(self, *callbacks):
self.callbacks = [normalize_ray_callback(c) for c in callbacks]
RayDaskCallback.ray_active.update(self.callbacks)
def __enter__(self):
return self
def __exit__(self, *args):
for c in self.callbacks:
RayDaskCallback.ray_active.discard(c)
def normalize_ray_callback(cb):
if isinstance(cb, RayDaskCallback):
return cb._ray_callback
elif isinstance(cb, RayCallback):
return cb
else:
raise TypeError(
"Callbacks must be either 'RayDaskCallback' or 'RayCallback' namedtuple"
)
def unpack_ray_callbacks(cbs):
"""Take an iterable of callbacks, return a list of each callback."""
if cbs:
# Only drop callback methods that aren't in CBS_DONT_DROP.
return RayCallbacks(
*(
[cb for cb in cbs_ if cb or CBS[idx] in CBS_DONT_DROP] or None
for idx, cbs_ in enumerate(zip(*cbs))
)
)
else:
return RayCallbacks(*([()] * len(CBS)))
@contextlib.contextmanager
def local_ray_callbacks(callbacks=None):
"""
Allows Dask-Ray callbacks to work with nested schedulers.
Callbacks will only be used by the first started scheduler they encounter.
This means that only the outermost scheduler will use global callbacks.
"""
global_callbacks = callbacks is None
if global_callbacks:
callbacks, RayDaskCallback.ray_active = (RayDaskCallback.ray_active, set())
try:
yield callbacks or ()
finally:
if global_callbacks:
RayDaskCallback.ray_active = callbacks
class ProgressBarCallback(RayDaskCallback):
def __init__(self):
@ray.remote
class ProgressBarActor:
def __init__(self):
self._init()
def submit(self, key, deps, now):
for dep in deps.keys():
self.deps[key].add(dep)
self.submitted[key] = now
self.submission_queue.append((key, now))
def task_scheduled(self, key, now):
self.scheduled[key] = now
def finish(self, key, now):
self.finished[key] = now
def result(self):
return len(self.submitted), len(self.finished)
def report(self):
result = defaultdict(dict)
for key, finished in self.finished.items():
submitted = self.submitted[key]
scheduled = self.scheduled[key]
# deps = self.deps[key]
result[key]["execution_time"] = (
finished - scheduled
).total_seconds()
# Calculate the scheduling time.
# This is inaccurate.
# We should subtract scheduled - (last dep completed).
# But currently it is not easy because
# of how getitem is implemented in dask on ray sort.
result[key]["scheduling_time"] = (
scheduled - submitted
).total_seconds()
result["submission_order"] = self.submission_queue
return result
def ready(self):
pass
def reset(self):
self._init()
def _init(self):
self.submission_queue = []
self.submitted = defaultdict(None)
self.scheduled = defaultdict(None)
self.finished = defaultdict(None)
self.deps = defaultdict(set)
try:
self.pb = ray.get_actor("_dask_on_ray_pb")
ray.get(self.pb.reset.remote())
except ValueError:
self.pb = ProgressBarActor.options(name="_dask_on_ray_pb").remote()
ray.get(self.pb.ready.remote())
def _ray_postsubmit(self, task, key, deps, object_ref):
# Indicate the dask task is submitted.
self.pb.submit.remote(key, deps, datetime.now())
def _ray_pretask(self, key, object_refs):
self.pb.task_scheduled.remote(key, datetime.now())
def _ray_posttask(self, key, result, pre_state):
# Indicate the dask task is finished.
self.pb.finish.remote(key, datetime.now())
def _ray_finish(self, result):
print("All tasks are completed.")
python/ray/util/dask/common.py
+88
View File
@@ -0,0 +1,88 @@
import uuid
from collections import OrderedDict
from collections.abc import Iterator
from operator import getitem
from typing import Any
from dask.core import get as get_sync, quote
from dask.utils import apply
import ray
try:
from dataclasses import fields as dataclass_fields, is_dataclass
except ImportError:
# Python < 3.7
def is_dataclass(x):
return False
def dataclass_fields(x):
return []
def unpack_object_refs(*args: Any):
"""
Extract Ray object refs from a set of potentially arbitrarily nested
Python objects.
Intended use is to find all Ray object references in a set of (possibly
nested) Python objects, do something to them (get(), wait(), etc.), then
repackage them into equivalent Python objects.
Args:
*args: One or more (potentially nested) Python objects that contain
Ray object references.
Returns:
A 2-tuple of a flat list of all contained Ray object references, and a
function that, when given the corresponding flat list of concrete
values, will return a set of Python objects equivalent to that which
was given in *args, but with all Ray object references replaced with
their corresponding concrete values.
"""
object_refs = []
repack_dsk = {}
object_refs_token = uuid.uuid4().hex
def _unpack(expr):
if isinstance(expr, ray.ObjectRef):
token = expr.hex()
repack_dsk[token] = (getitem, object_refs_token, len(object_refs))
object_refs.append(expr)
return token
token = uuid.uuid4().hex
# Treat iterators like lists
typ = list if isinstance(expr, Iterator) else type(expr)
if typ in (list, tuple, set):
repack_task = (typ, [_unpack(i) for i in expr])
elif typ in (dict, OrderedDict):
repack_task = (typ, [[_unpack(k), _unpack(v)] for k, v in expr.items()])
elif is_dataclass(expr):
repack_task = (
apply,
typ,
(),
(
dict,
[
[f.name, _unpack(getattr(expr, f.name))]
for f in dataclass_fields(expr)
],
),
)
else:
return expr
repack_dsk[token] = repack_task
return token
out = uuid.uuid4().hex
repack_dsk[out] = (tuple, [_unpack(i) for i in args])
def repack(results):
dsk = repack_dsk.copy()
dsk[object_refs_token] = quote(results)
return get_sync(dsk, out)
return object_refs, repack
python/ray/util/dask/optimizations.py
+149
View File
@@ -0,0 +1,149 @@
import warnings
import dask
from dask import core
from dask.dataframe.core import _concat
from dask.highlevelgraph import HighLevelGraph
from .scheduler import MultipleReturnFunc, multiple_return_get
try:
from dask.dataframe.optimize import optimize
from dask.dataframe.shuffle import SimpleShuffleLayer, shuffle_group
except ImportError:
# SimpleShuffleLayer doesn't exist in this version of Dask.
# This is the case for dask>=2025.1.0.
SimpleShuffleLayer = None
try:
import dask_expr # noqa: F401
SimpleShuffleLayer = None
except ImportError:
pass
if SimpleShuffleLayer is not None:
class MultipleReturnSimpleShuffleLayer(SimpleShuffleLayer):
@classmethod
def clone(cls, layer: SimpleShuffleLayer):
# TODO(Clark): Probably don't need this since SimpleShuffleLayer
# implements __copy__() and the shallow clone should be enough?
return cls(
name=layer.name,
column=layer.column,
npartitions=layer.npartitions,
npartitions_input=layer.npartitions_input,
ignore_index=layer.ignore_index,
name_input=layer.name_input,
meta_input=layer.meta_input,
parts_out=layer.parts_out,
annotations=layer.annotations,
)
def __repr__(self):
return (
f"MultipleReturnSimpleShuffleLayer<name='{self.name}', "
f"npartitions={self.npartitions}>"
)
def __reduce__(self):
attrs = [
"name",
"column",
"npartitions",
"npartitions_input",
"ignore_index",
"name_input",
"meta_input",
"parts_out",
"annotations",
]
return (
MultipleReturnSimpleShuffleLayer,
tuple(getattr(self, attr) for attr in attrs),
)
def _cull(self, parts_out):
return MultipleReturnSimpleShuffleLayer(
self.name,
self.column,
self.npartitions,
self.npartitions_input,
self.ignore_index,
self.name_input,
self.meta_input,
parts_out=parts_out,
)
def _construct_graph(self):
"""Construct graph for a simple shuffle operation."""
shuffle_group_name = "group-" + self.name
shuffle_split_name = "split-" + self.name
dsk = {}
n_parts_out = len(self.parts_out)
for part_out in self.parts_out:
# TODO(Clark): Find better pattern than in-scheduler concat.
_concat_list = [
(shuffle_split_name, part_out, part_in)
for part_in in range(self.npartitions_input)
]
dsk[(self.name, part_out)] = (_concat, _concat_list, self.ignore_index)
for _, _part_out, _part_in in _concat_list:
dsk[(shuffle_split_name, _part_out, _part_in)] = (
multiple_return_get,
(shuffle_group_name, _part_in),
_part_out,
)
if (shuffle_group_name, _part_in) not in dsk:
dsk[(shuffle_group_name, _part_in)] = (
MultipleReturnFunc(
shuffle_group,
n_parts_out,
),
(self.name_input, _part_in),
self.column,
0,
self.npartitions,
self.npartitions,
self.ignore_index,
self.npartitions,
)
return dsk
def rewrite_simple_shuffle_layer(dsk, keys):
if not isinstance(dsk, HighLevelGraph):
dsk = HighLevelGraph.from_collections(id(dsk), dsk, dependencies=())
else:
dsk = dsk.copy()
layers = dsk.layers.copy()
for key, layer in layers.items():
if type(layer) is SimpleShuffleLayer:
dsk.layers[key] = MultipleReturnSimpleShuffleLayer.clone(layer)
return dsk
def dataframe_optimize(dsk, keys, **kwargs):
if not isinstance(keys, (list, set)):
keys = [keys]
keys = list(core.flatten(keys))
if not isinstance(dsk, HighLevelGraph):
dsk = HighLevelGraph.from_collections(id(dsk), dsk, dependencies=())
dsk = rewrite_simple_shuffle_layer(dsk, keys=keys)
return optimize(dsk, keys, **kwargs)
else:
def dataframe_optimize(dsk, keys, **kwargs):
warnings.warn(
"Custom dataframe shuffle optimization only works on "
"dask>=2024.11.0,<2025.1.0, you are on version "
f"{dask.__version__}."
"Doing no additional optimization aside from the default one."
)
return None

Some files were not shown because too many files have changed in this diff Show More