wehub-resource-sync
156 lines
6.1 KiB
Python
156 lines
6.1 KiB
Python
import abc
|
|
from typing import List, Optional
|
|
|
|
import ray
|
|
from ray.air.execution.resources.request import AcquiredResources, ResourceRequest
|
|
from ray.util.annotations import DeveloperAPI
|
|
|
|
|
|
@DeveloperAPI
|
|
class ResourceManager(abc.ABC):
|
|
"""Resource manager interface.
|
|
|
|
A resource manager can be used to request resources from a Ray cluster and
|
|
allocate them to remote Ray tasks or actors.
|
|
|
|
Resources have to be requested before they can be acquired.
|
|
|
|
Resources managed by the resource manager can be in three states:
|
|
|
|
1. "Requested": The resources have been requested but are not yet available to
|
|
schedule remote Ray objects. The resource request may trigger autoscaling,
|
|
and can be cancelled if no longer needed.
|
|
2. "Ready": The requested resources are now available to schedule remote Ray
|
|
objects. They can be acquired and subsequently used remote Ray objects.
|
|
The resource request can still be cancelled if no longer needed.
|
|
3. "Acquired": The resources have been acquired by a caller to use for scheduling
|
|
remote Ray objects. Note that it is the responsibility of the caller to
|
|
schedule the Ray objects with these resources.
|
|
The associated resource request has been completed and can no longer be
|
|
cancelled. The acquired resources can be freed by the resource manager when
|
|
they are no longer used.
|
|
|
|
The flow is as follows:
|
|
|
|
.. code-block:: python
|
|
|
|
# Create resource manager
|
|
resource_manager = ResourceManager()
|
|
|
|
# Create resource request
|
|
resource_request = ResourceRequest([{"CPU": 4}])
|
|
|
|
# Pass to resource manager
|
|
resource_manager.request_resources(resource_request)
|
|
|
|
# Wait until ready
|
|
while not resource_manager.has_resources_ready(resource_request):
|
|
time.sleep(1)
|
|
|
|
# Once ready, acquire resources
|
|
acquired_resource = resource_manager.acquire_resources(resource_request)
|
|
|
|
# Bind to remote task or actor
|
|
annotated_remote_fn = acquired_resource.annotate_remote_entities(
|
|
[remote_fn])
|
|
|
|
# Run remote function. This will use the acquired resources
|
|
ray.get(annotated_remote_fn.remote())
|
|
|
|
# After using the resources, free
|
|
resource_manager.free_resources(annotated_resources)
|
|
|
|
"""
|
|
|
|
def request_resources(self, resource_request: ResourceRequest):
|
|
"""Request resources.
|
|
|
|
Depending on the backend, resources can trigger autoscaling. Requested
|
|
resources can be ready or not ready. Once they are "ready", they can
|
|
be acquired and used by remote Ray objects.
|
|
|
|
Resource requests can be cancelled anytime using ``cancel_resource_request()``.
|
|
Once acquired, the resource request is removed. Acquired resources can be
|
|
freed with ``free_resources()``.
|
|
"""
|
|
raise NotImplementedError
|
|
|
|
def cancel_resource_request(self, resource_request: ResourceRequest):
|
|
"""Cancel resource request.
|
|
|
|
Resource requests can be cancelled anytime before a resource is acquired.
|
|
Acquiring a resource will remove the associated resource request.
|
|
Acquired resources can be freed with ``free_resources()``.
|
|
"""
|
|
raise NotImplementedError
|
|
|
|
def has_resources_ready(self, resource_request: ResourceRequest) -> bool:
|
|
"""Returns True if resources for the given request are ready to be acquired."""
|
|
raise NotImplementedError
|
|
|
|
def acquire_resources(
|
|
self, resource_request: ResourceRequest
|
|
) -> Optional[AcquiredResources]:
|
|
"""Acquire resources. Returns None if resources are not ready to be acquired.
|
|
|
|
Acquiring resources will remove the associated resource request.
|
|
Acquired resources can be returned with ``free_resources()``.
|
|
"""
|
|
raise NotImplementedError
|
|
|
|
def free_resources(self, acquired_resource: AcquiredResources):
|
|
"""Free acquired resources from usage and return them to the resource manager.
|
|
|
|
Freeing resources will return the resources to the manager, but there are
|
|
no guarantees about the tasks and actors scheduled on the resources. The caller
|
|
should make sure that any references to tasks or actors scheduled on the
|
|
resources have been removed before calling ``free_resources()``.
|
|
"""
|
|
raise NotImplementedError
|
|
|
|
def get_resource_futures(self) -> List[ray.ObjectRef]:
|
|
"""Return futures for resources to await.
|
|
|
|
Depending on the backend, we use resource futures to determine availability
|
|
of resources (e.g. placement groups) or resolution of requests.
|
|
In this case, the futures can be awaited externally by the caller.
|
|
|
|
When a resource future resolved, the caller may call ``update_state()``
|
|
to force the resource manager to update its internal state immediately.
|
|
"""
|
|
return []
|
|
|
|
def update_state(self):
|
|
"""Update internal state of the resource manager.
|
|
|
|
The resource manager may have internal state that needs periodic updating.
|
|
For instance, depending on the backend, resource futures can be awaited
|
|
externally (with ``get_resource_futures()``).
|
|
|
|
If such a future resolved, the caller can instruct the resource
|
|
manager to update its internal state immediately.
|
|
"""
|
|
pass
|
|
|
|
def clear(self):
|
|
"""Reset internal state and clear all resources.
|
|
|
|
Calling this method will reset the resource manager to its initialization state.
|
|
All resources will be removed.
|
|
|
|
Clearing the state will remove tracked resources from the manager, but there are
|
|
no guarantees about the tasks and actors scheduled on the resources. The caller
|
|
should make sure that any references to tasks or actors scheduled on the
|
|
resources have been removed before calling ``clear()``.
|
|
"""
|
|
raise NotImplementedError
|
|
|
|
def __reduce__(self):
|
|
"""We disallow serialization.
|
|
|
|
Shared resource managers should live on an actor.
|
|
"""
|
|
raise ValueError(
|
|
f"Resource managers cannot be serialized. Resource manager: {str(self)}"
|
|
)
|