wehub-resource-sync
215 lines
8.3 KiB
Python
215 lines
8.3 KiB
Python
import time
|
|
from collections import defaultdict
|
|
from dataclasses import dataclass
|
|
from typing import Dict, List, Optional, Set
|
|
|
|
import ray
|
|
from ray.air.execution.resources.request import (
|
|
AcquiredResources,
|
|
RemoteRayEntity,
|
|
ResourceRequest,
|
|
)
|
|
from ray.air.execution.resources.resource_manager import ResourceManager
|
|
from ray.util.annotations import DeveloperAPI
|
|
from ray.util.placement_group import PlacementGroup, remove_placement_group
|
|
from ray.util.scheduling_strategies import PlacementGroupSchedulingStrategy
|
|
|
|
|
|
@DeveloperAPI
|
|
@dataclass
|
|
class PlacementGroupAcquiredResources(AcquiredResources):
|
|
placement_group: PlacementGroup
|
|
|
|
def _annotate_remote_entity(
|
|
self, entity: RemoteRayEntity, bundle: Dict[str, float], bundle_index: int
|
|
) -> RemoteRayEntity:
|
|
bundle = bundle.copy()
|
|
num_cpus = bundle.pop("CPU", 0)
|
|
num_gpus = bundle.pop("GPU", 0)
|
|
memory = bundle.pop("memory", 0.0)
|
|
|
|
return entity.options(
|
|
scheduling_strategy=PlacementGroupSchedulingStrategy(
|
|
placement_group=self.placement_group,
|
|
placement_group_bundle_index=bundle_index,
|
|
placement_group_capture_child_tasks=True,
|
|
),
|
|
num_cpus=num_cpus,
|
|
num_gpus=num_gpus,
|
|
memory=memory,
|
|
resources=bundle,
|
|
)
|
|
|
|
|
|
@DeveloperAPI
|
|
class PlacementGroupResourceManager(ResourceManager):
|
|
"""Resource manager using placement groups as the resource backend.
|
|
|
|
This manager will use placement groups to fulfill resource requests. Requesting
|
|
a resource will schedule the placement group. Acquiring a resource will
|
|
return a ``PlacementGroupAcquiredResources`` that can be used to schedule
|
|
Ray tasks and actors on the placement group. Freeing an acquired resource
|
|
will destroy the associated placement group.
|
|
|
|
Ray core does not emit events when resources are available. Instead, the
|
|
scheduling state has to be periodically updated.
|
|
|
|
Per default, placement group scheduling state is refreshed every time when
|
|
resource state is inquired, but not more often than once every ``update_interval_s``
|
|
seconds. Alternatively, staging futures can be retrieved (and awaited) with
|
|
``get_resource_futures()`` and state update can be force with ``update_state()``.
|
|
|
|
Args:
|
|
update_interval_s: Minimum interval in seconds between updating scheduling
|
|
state of placement groups.
|
|
|
|
"""
|
|
|
|
_resource_cls: AcquiredResources = PlacementGroupAcquiredResources
|
|
|
|
def __init__(self, update_interval_s: float = 0.1):
|
|
# Internally, the placement group lifecycle is like this:
|
|
# - Resources are requested with ``request_resources()``
|
|
# - A placement group is scheduled ("staged")
|
|
# - A ``PlacementGroup.ready()`` future is scheduled ("staging future")
|
|
# - We update the scheduling state when we need to
|
|
# (e.g. when ``has_resources_ready()`` is called)
|
|
# - When staging futures resolve, a placement group is moved from "staging"
|
|
# to "ready"
|
|
# - When a resource request is canceled, we remove a placement group from
|
|
# "staging". If there are not staged placement groups
|
|
# (because they are already "ready"), we remove one from "ready" instead.
|
|
# - When a resource is acquired, the pg is removed from "ready" and moved
|
|
# to "acquired"
|
|
# - When a resource is freed, the pg is removed from "acquired" and destroyed
|
|
|
|
# Mapping of placement group to request
|
|
self._pg_to_request: Dict[PlacementGroup, ResourceRequest] = {}
|
|
|
|
# PGs that are staged but not "ready", yet (i.e. not CREATED)
|
|
self._request_to_staged_pgs: Dict[
|
|
ResourceRequest, Set[PlacementGroup]
|
|
] = defaultdict(set)
|
|
|
|
# PGs that are CREATED and can be used by tasks and actors
|
|
self._request_to_ready_pgs: Dict[
|
|
ResourceRequest, Set[PlacementGroup]
|
|
] = defaultdict(set)
|
|
|
|
# Staging futures used to update internal state.
|
|
# We keep a double mapping here for better lookup efficiency.
|
|
self._staging_future_to_pg: Dict[ray.ObjectRef, PlacementGroup] = dict()
|
|
self._pg_to_staging_future: Dict[PlacementGroup, ray.ObjectRef] = dict()
|
|
|
|
# Set of acquired PGs. We keep track of these here to make sure we
|
|
# only free PGs that this manager managed.
|
|
self._acquired_pgs: Set[PlacementGroup] = set()
|
|
|
|
# Minimum time between updates of the internal state
|
|
self.update_interval_s = update_interval_s
|
|
self._last_update = time.monotonic() - self.update_interval_s - 1
|
|
|
|
def get_resource_futures(self) -> List[ray.ObjectRef]:
|
|
return list(self._staging_future_to_pg.keys())
|
|
|
|
def _maybe_update_state(self):
|
|
now = time.monotonic()
|
|
if now > self._last_update + self.update_interval_s:
|
|
self.update_state()
|
|
|
|
def update_state(self):
|
|
ready, not_ready = ray.wait(
|
|
list(self._staging_future_to_pg.keys()),
|
|
num_returns=len(self._staging_future_to_pg),
|
|
timeout=0,
|
|
)
|
|
for future in ready:
|
|
# Remove staging future
|
|
pg = self._staging_future_to_pg.pop(future)
|
|
self._pg_to_staging_future.pop(pg)
|
|
# Fetch resource request
|
|
request = self._pg_to_request[pg]
|
|
# Remove from staging, add to ready
|
|
self._request_to_staged_pgs[request].remove(pg)
|
|
self._request_to_ready_pgs[request].add(pg)
|
|
self._last_update = time.monotonic()
|
|
|
|
def request_resources(self, resource_request: ResourceRequest):
|
|
pg = resource_request.to_placement_group()
|
|
self._pg_to_request[pg] = resource_request
|
|
self._request_to_staged_pgs[resource_request].add(pg)
|
|
|
|
future = pg.ready()
|
|
self._staging_future_to_pg[future] = pg
|
|
self._pg_to_staging_future[pg] = future
|
|
|
|
def cancel_resource_request(self, resource_request: ResourceRequest):
|
|
if self._request_to_staged_pgs[resource_request]:
|
|
pg = self._request_to_staged_pgs[resource_request].pop()
|
|
|
|
# PG was staging
|
|
future = self._pg_to_staging_future.pop(pg)
|
|
self._staging_future_to_pg.pop(future)
|
|
|
|
# Cancel the pg.ready task.
|
|
# Otherwise, it will be pending node assignment forever.
|
|
ray.cancel(future)
|
|
else:
|
|
# PG might be ready
|
|
pg = self._request_to_ready_pgs[resource_request].pop()
|
|
if not pg:
|
|
raise RuntimeError(
|
|
"Cannot cancel resource request: No placement group was "
|
|
f"staged or is ready. Make sure to not cancel more resource "
|
|
f"requests than you've created. Request: {resource_request}"
|
|
)
|
|
|
|
self._pg_to_request.pop(pg)
|
|
ray.util.remove_placement_group(pg)
|
|
|
|
def has_resources_ready(self, resource_request: ResourceRequest) -> bool:
|
|
if not bool(len(self._request_to_ready_pgs[resource_request])):
|
|
# Only update state if needed
|
|
self._maybe_update_state()
|
|
|
|
return bool(len(self._request_to_ready_pgs[resource_request]))
|
|
|
|
def acquire_resources(
|
|
self, resource_request: ResourceRequest
|
|
) -> Optional[PlacementGroupAcquiredResources]:
|
|
if not self.has_resources_ready(resource_request):
|
|
return None
|
|
|
|
pg = self._request_to_ready_pgs[resource_request].pop()
|
|
self._acquired_pgs.add(pg)
|
|
|
|
return self._resource_cls(placement_group=pg, resource_request=resource_request)
|
|
|
|
def free_resources(self, acquired_resource: PlacementGroupAcquiredResources):
|
|
pg = acquired_resource.placement_group
|
|
|
|
self._acquired_pgs.remove(pg)
|
|
remove_placement_group(pg)
|
|
self._pg_to_request.pop(pg)
|
|
|
|
def clear(self):
|
|
if not ray.is_initialized():
|
|
return
|
|
|
|
for staged_pgs in self._request_to_staged_pgs.values():
|
|
for staged_pg in staged_pgs:
|
|
remove_placement_group(staged_pg)
|
|
|
|
for ready_pgs in self._request_to_ready_pgs.values():
|
|
for ready_pg in ready_pgs:
|
|
remove_placement_group(ready_pg)
|
|
|
|
for acquired_pg in self._acquired_pgs:
|
|
remove_placement_group(acquired_pg)
|
|
|
|
# Reset internal state
|
|
self.__init__(update_interval_s=self.update_interval_s)
|
|
|
|
def __del__(self):
|
|
self.clear()
|