wehub-resource-sync
232 lines
9.8 KiB
Python
232 lines
9.8 KiB
Python
import warnings
|
|
from typing import TYPE_CHECKING, Any, Dict, List, Set, Tuple, Union
|
|
|
|
from ray.experimental.util.types import Device
|
|
|
|
if TYPE_CHECKING:
|
|
import numpy as np
|
|
import torch
|
|
|
|
|
|
_TORCH_WARNING_FILTER_ACTIVATE = True
|
|
|
|
|
|
class _SerializationContext:
|
|
def __init__(self):
|
|
# If true, then tensors found in the data to serialize are extracted
|
|
# and the caller should send them through an external transport.
|
|
self._use_external_transport: bool = False
|
|
# If _use_external_transport is True, then these are
|
|
# the tensors that should be sent or received
|
|
# out-of-band, through the external transport.
|
|
self._out_of_band_tensors: List["torch.Tensor"] = []
|
|
# During serialization, tensors sent out-of-band are replaced with
|
|
# integer placeholders. This tracks the set of placeholders seen.
|
|
self._deserialized_tensor_placeholders: Set[int] = set()
|
|
|
|
# Buffer for transferring data between tasks in the same worker process.
|
|
# The key is the channel ID, and the value is the data. We don't use a
|
|
# lock when reading/writing the buffer because a DAG node actor will only
|
|
# execute one task at a time in `do_exec_tasks`. It will not execute multiple
|
|
# Ray tasks on a single actor simultaneously.
|
|
self.intra_process_channel_buffers: Dict[str, Any] = {}
|
|
# The number of readers for each channel. When the number of readers
|
|
# reaches 0, remove the data from the buffer.
|
|
self.channel_id_to_num_readers: Dict[str, int] = {}
|
|
|
|
def set_target_device(self, device: Device) -> None:
|
|
self._target_device = device
|
|
|
|
def set_data(self, channel_id: str, value: Any, num_readers: int) -> None:
|
|
assert num_readers > 0, "num_readers must be greater than 0."
|
|
assert (
|
|
channel_id not in self.intra_process_channel_buffers
|
|
), f"Channel {channel_id} already exists in the buffer."
|
|
assert (
|
|
channel_id not in self.channel_id_to_num_readers
|
|
), f"Channel {channel_id} already exists in the channel_id_to_num_readers."
|
|
|
|
self.intra_process_channel_buffers[channel_id] = value
|
|
self.channel_id_to_num_readers[channel_id] = num_readers
|
|
|
|
def has_data(self, channel_id: str) -> bool:
|
|
return channel_id in self.intra_process_channel_buffers
|
|
|
|
def get_data(self, channel_id: str) -> Any:
|
|
assert (
|
|
channel_id in self.intra_process_channel_buffers
|
|
), f"Channel {channel_id} does not exist in the buffer."
|
|
assert (
|
|
channel_id in self.channel_id_to_num_readers
|
|
), f"Channel {channel_id} does not exist in the channel_id_to_num_readers."
|
|
|
|
self.channel_id_to_num_readers[channel_id] -= 1
|
|
if self.channel_id_to_num_readers[channel_id] == 0:
|
|
# All readers have read the data, so we can remove it.
|
|
self.channel_id_to_num_readers.pop(channel_id)
|
|
return self.intra_process_channel_buffers.pop(channel_id)
|
|
return self.intra_process_channel_buffers[channel_id]
|
|
|
|
def reset_data(self, channel_id: str) -> None:
|
|
self.intra_process_channel_buffers.pop(channel_id, None)
|
|
self.channel_id_to_num_readers.pop(channel_id, None)
|
|
|
|
def set_use_external_transport(self, use_external_transport: bool) -> None:
|
|
self._use_external_transport = use_external_transport
|
|
|
|
@property
|
|
def use_external_transport(self) -> bool:
|
|
return self._use_external_transport
|
|
|
|
def reset_out_of_band_tensors(
|
|
self, tensors: List["torch.Tensor"]
|
|
) -> Tuple[List["torch.Tensor"], Set[int]]:
|
|
"""
|
|
Return and reset the out-of-band tensors and all tensor placeholders
|
|
that were deserialized since the last call to reset.
|
|
"""
|
|
prev_tensors = self._out_of_band_tensors
|
|
deserialized_tensor_placeholders = self._deserialized_tensor_placeholders
|
|
self._out_of_band_tensors = tensors
|
|
self._deserialized_tensor_placeholders = set()
|
|
return prev_tensors, deserialized_tensor_placeholders
|
|
|
|
def serialize_tensor(
|
|
self, tensor: "torch.Tensor"
|
|
) -> Union[int, Tuple["np.ndarray", "torch.dtype", str]]:
|
|
from ray.experimental.channel import ChannelContext
|
|
|
|
ctx = ChannelContext.get_current()
|
|
if self._use_external_transport and (
|
|
ctx._torch_device is None or ctx._torch_device == tensor.device
|
|
):
|
|
# External transport is enabled and we found a tensor that matches
|
|
# our device. Add the actual tensor to a buffer. The buffer of
|
|
# tensors should later be popped by the caller and sent via
|
|
# external transport.
|
|
self._out_of_band_tensors.append(tensor)
|
|
# Return a placeholder.
|
|
return len(self._out_of_band_tensors) - 1
|
|
|
|
return self.serialize_to_numpy_or_scalar(tensor)
|
|
|
|
def serialize_to_numpy_or_scalar(
|
|
self, tensor: "torch.Tensor"
|
|
) -> Tuple[Union["np.ndarray", Any], "torch.dtype", str]:
|
|
"""
|
|
Serialize a tensor to a numpy array,
|
|
or a scalar when the tensor is 0-dim.
|
|
"""
|
|
import torch
|
|
|
|
tensor_device_type = tensor.device.type
|
|
|
|
# Transfer through Ray's shared memory store for now.
|
|
# TODO(swang): This requires two copies, one to transfer from GPU to
|
|
# CPU and another from CPU to shared memory. Ideally we should elide
|
|
# the first copy and memcpy directly from GPU to the shared memory
|
|
# buffer.
|
|
if tensor_device_type != "cpu":
|
|
tensor = tensor.to("cpu")
|
|
|
|
# Numpy does not have an equivalent dtype for all torch dtypes, so
|
|
# instead of casting directly to numpy:
|
|
# 1) for non-scalar tensors, we first use a view with a common dtype (uint8)
|
|
# and then view as numpy array.
|
|
# 2) for scalar tensors, we cannot use a uint8 view when the size differs,
|
|
# so we save the original item and type information.
|
|
if tensor.dim() > 0:
|
|
return (tensor.view(torch.uint8).numpy(), tensor.dtype, tensor_device_type)
|
|
else:
|
|
return (tensor.item(), tensor.dtype, tensor_device_type)
|
|
|
|
def deserialize_tensor(
|
|
self,
|
|
val: Union[Tuple["np.ndarray", "torch.dtype", str], int],
|
|
target_device: Device,
|
|
):
|
|
|
|
# Found a placeholder for a tensor that was serialized via accelerator.
|
|
# Replace it with the corresponding deserialized tensor.
|
|
if isinstance(val, int):
|
|
placeholder = val
|
|
self._deserialized_tensor_placeholders.add(placeholder)
|
|
assert placeholder < len(self._out_of_band_tensors), (
|
|
"placeholder",
|
|
placeholder,
|
|
"out_of_band_tensors",
|
|
self._out_of_band_tensors,
|
|
)
|
|
tensor = self._out_of_band_tensors[placeholder]
|
|
if target_device == Device.CPU:
|
|
tensor = tensor.to("cpu")
|
|
return tensor
|
|
|
|
np_array, dtype, tensor_device_type = val
|
|
return self.deserialize_from_numpy_or_scalar(
|
|
np_array, dtype, tensor_device_type, target_device
|
|
)
|
|
|
|
def deserialize_from_numpy_or_scalar(
|
|
self,
|
|
np_array: Union["np.ndarray", Any],
|
|
dtype: "torch.dtype",
|
|
tensor_device_type: str,
|
|
target_device: Device,
|
|
):
|
|
import numpy as np
|
|
import torch
|
|
|
|
if target_device == Device.DEFAULT:
|
|
target_device_type = tensor_device_type
|
|
elif target_device in [Device.GPU, Device.CUDA]:
|
|
target_device_type = "cuda"
|
|
else:
|
|
target_device_type = target_device.value
|
|
|
|
# TODO(swang): Support local P2P transfers if available.
|
|
if target_device_type != "cpu":
|
|
|
|
def convert_numpy_to_tensor(np_array):
|
|
if not isinstance(np_array, np.ndarray):
|
|
# For scalar tensors, create the 0-dim tensor.
|
|
return torch.tensor(
|
|
np_array, device=target_device_type, dtype=dtype
|
|
)
|
|
else:
|
|
# For non-scalar tensors, view as the original dtype.
|
|
# It does zero-copy convert np_array inside shared memory to
|
|
# a tensor. Since we move data to GPU immediately, it is safe.
|
|
cpu_tensor = torch.from_numpy(np_array).view(dtype)
|
|
return cpu_tensor.to(device=target_device_type)
|
|
|
|
global _TORCH_WARNING_FILTER_ACTIVATE
|
|
# filtering warning messages would be the bottleneck for
|
|
# deserializing torch tensors. Since the warning only prompts once,
|
|
# we would only deal with it for the first time.
|
|
if _TORCH_WARNING_FILTER_ACTIVATE:
|
|
with warnings.catch_warnings():
|
|
# Since np_array.is_writable is False (it is set by Ray),
|
|
# this raises a warning. Suppress it.
|
|
warnings.filterwarnings(
|
|
"ignore",
|
|
category=UserWarning,
|
|
message="The given NumPy array is not writable",
|
|
)
|
|
gpu_tensor = convert_numpy_to_tensor(np_array)
|
|
_TORCH_WARNING_FILTER_ACTIVATE = False
|
|
else:
|
|
gpu_tensor = convert_numpy_to_tensor(np_array)
|
|
|
|
return gpu_tensor
|
|
|
|
# TODO(swang): Use zero-copy from_numpy() if np_array.flags.writeable
|
|
# is True. This is safe to set when deserializing np_array if the
|
|
# upstream task has num_readers=1.
|
|
if not isinstance(np_array, np.ndarray):
|
|
# For scalar tensors, create the 0-dim tensor.
|
|
return torch.tensor(np_array, device=target_device_type, dtype=dtype)
|
|
else:
|
|
# For non-scalar tensors, view as the original dtype.
|
|
return torch.tensor(np_array, device=target_device_type).view(dtype)
|