chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -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]
|
||||
Reference in New Issue
Block a user