wehub-resource-sync
145 lines
4.8 KiB
Python
145 lines
4.8 KiB
Python
from abc import ABC, abstractmethod
|
|
from typing import Any, Dict, Generic, TypeVar
|
|
|
|
from ray.experimental.channel.accelerator_context import AcceleratorContext
|
|
from ray.util.annotations import DeveloperAPI
|
|
|
|
T = TypeVar("T")
|
|
|
|
|
|
@DeveloperAPI
|
|
class DAGOperationFuture(ABC, Generic[T]):
|
|
"""
|
|
A future representing the result of a DAG operation.
|
|
|
|
This is an abstraction that is internal to each actor,
|
|
and is not exposed to the DAG caller.
|
|
"""
|
|
|
|
@abstractmethod
|
|
def wait(self):
|
|
"""
|
|
Wait for the future and return the result of the operation.
|
|
"""
|
|
raise NotImplementedError
|
|
|
|
|
|
@DeveloperAPI
|
|
class ResolvedFuture(DAGOperationFuture):
|
|
"""
|
|
A future that is already resolved. Calling `wait()` on this will
|
|
immediately return the result without blocking.
|
|
"""
|
|
|
|
def __init__(self, result: Any):
|
|
"""
|
|
Initialize a resolved future.
|
|
|
|
Args:
|
|
result: The result of the future.
|
|
"""
|
|
self._result = result
|
|
|
|
def wait(self):
|
|
"""
|
|
Wait and immediately return the result. This operation will not block.
|
|
"""
|
|
return self._result
|
|
|
|
|
|
@DeveloperAPI
|
|
class GPUFuture(DAGOperationFuture[Any]):
|
|
"""
|
|
A future for a GPU event on a CUDA stream.
|
|
|
|
This future wraps a buffer, and records an event on the given stream
|
|
when it is created. When the future is waited on, it makes the current
|
|
CUDA stream wait on the event, then returns the buffer.
|
|
|
|
The buffer must be a GPU tensor produced by an earlier operation launched
|
|
on the given stream, or it could be CPU data. Then the future guarantees
|
|
that when the wait() returns, the buffer is ready on the current stream.
|
|
|
|
The `wait()` does not block CPU.
|
|
"""
|
|
|
|
# Caching GPU futures ensures CUDA events associated with futures are properly
|
|
# destroyed instead of relying on garbage collection. The CUDA event contained
|
|
# in a GPU future is destroyed right before removing the future from the cache.
|
|
# The dictionary key is the future ID, which is the task idx of the dag operation
|
|
# that produced the future. When a future is created, it is immediately added to
|
|
# the cache. When a future has been waited on, it is removed from the cache.
|
|
# When adding a future, if its ID is already a key in the cache, the old future
|
|
# is removed. This can happen when an exception is thrown in a previous execution
|
|
# of the dag, in which case the old future is never waited on.
|
|
# Upon dag teardown, all pending futures produced by the dag are removed.
|
|
gpu_futures: Dict[int, "GPUFuture"] = {}
|
|
|
|
@staticmethod
|
|
def add_gpu_future(fut_id: int, fut: "GPUFuture") -> None:
|
|
"""
|
|
Cache the GPU future.
|
|
Args:
|
|
fut_id: GPU future ID.
|
|
fut: GPU future to be cached.
|
|
"""
|
|
if fut_id in GPUFuture.gpu_futures:
|
|
# The old future was not waited on because of an execution exception.
|
|
GPUFuture.gpu_futures.pop(fut_id).destroy_event()
|
|
GPUFuture.gpu_futures[fut_id] = fut
|
|
|
|
@staticmethod
|
|
def remove_gpu_future(fut_id: int) -> None:
|
|
"""
|
|
Remove the cached GPU future and destroy its CUDA event.
|
|
Args:
|
|
fut_id: GPU future ID.
|
|
"""
|
|
if fut_id in GPUFuture.gpu_futures:
|
|
GPUFuture.gpu_futures.pop(fut_id).destroy_event()
|
|
|
|
def __init__(self, buf: Any, fut_id: int, stream: Any = None):
|
|
"""
|
|
Initialize a GPU future on the given stream.
|
|
|
|
Args:
|
|
buf: The buffer to return when the future is resolved.
|
|
fut_id: The future ID to cache the future.
|
|
stream: The torch stream to record the event on, this event is waited
|
|
on when the future is resolved. If None, the current stream is used.
|
|
"""
|
|
if stream is None:
|
|
stream = AcceleratorContext.get().current_stream()
|
|
|
|
self._buf = buf
|
|
self._event = AcceleratorContext.get().create_event()
|
|
self._event.record(stream)
|
|
self._fut_id = fut_id
|
|
self._waited: bool = False
|
|
|
|
# Cache the GPU future such that its CUDA event is properly destroyed.
|
|
GPUFuture.add_gpu_future(fut_id, self)
|
|
|
|
def wait(self) -> Any:
|
|
"""
|
|
Wait for the future on the current CUDA stream and return the result from
|
|
the GPU operation. This operation does not block CPU.
|
|
"""
|
|
current_stream = AcceleratorContext.get().current_stream()
|
|
if not self._waited:
|
|
self._waited = True
|
|
current_stream.wait_event(self._event)
|
|
# Destroy the CUDA event after it is waited on.
|
|
GPUFuture.remove_gpu_future(self._fut_id)
|
|
|
|
return self._buf
|
|
|
|
def destroy_event(self) -> None:
|
|
"""
|
|
Destroy the CUDA event associated with this future.
|
|
"""
|
|
if self._event is None:
|
|
return
|
|
|
|
self._event = None
|