wehub-resource-sync
157 lines
4.5 KiB
Python
157 lines
4.5 KiB
Python
import multiprocessing.shared_memory as shm
|
|
import pickle
|
|
import sys
|
|
from dataclasses import dataclass
|
|
from typing import Any, Dict, List, Optional
|
|
|
|
import numpy
|
|
import pytest
|
|
|
|
import ray
|
|
from ray.experimental import (
|
|
CommunicatorMetadata,
|
|
TensorTransportManager,
|
|
TensorTransportMetadata,
|
|
register_tensor_transport,
|
|
)
|
|
|
|
|
|
@dataclass
|
|
class ShmTransportMetadata(TensorTransportMetadata):
|
|
shm_name: Optional[str] = None
|
|
shm_size: Optional[int] = None
|
|
|
|
|
|
@dataclass
|
|
class ShmCommunicatorMetadata(CommunicatorMetadata):
|
|
pass
|
|
|
|
|
|
class SharedMemoryTransport(TensorTransportManager):
|
|
def __init__(self):
|
|
self.shared_memory_objects: Dict[str, shm.SharedMemory] = {}
|
|
|
|
def tensor_transport_backend(self) -> str:
|
|
return "shared_memory"
|
|
|
|
@staticmethod
|
|
def is_one_sided() -> bool:
|
|
return True
|
|
|
|
@staticmethod
|
|
def can_abort_transport() -> bool:
|
|
return False
|
|
|
|
def actor_has_tensor_transport(self, actor: "ray.actor.ActorHandle") -> bool:
|
|
return True
|
|
|
|
def extract_tensor_transport_metadata(
|
|
self,
|
|
obj_id: str,
|
|
rdt_object: List[numpy.ndarray],
|
|
) -> TensorTransportMetadata:
|
|
|
|
tensor_meta = []
|
|
if rdt_object:
|
|
for tensor in rdt_object:
|
|
tensor_meta.append((tensor.shape, tensor.dtype))
|
|
|
|
serialized_rdt_object = pickle.dumps(rdt_object)
|
|
size = len(serialized_rdt_object)
|
|
# Shm name can't be as long as the obj_id, so we truncate it.
|
|
name = obj_id[:20]
|
|
shm_obj = shm.SharedMemory(name=name, create=True, size=size)
|
|
shm_obj.buf[:size] = serialized_rdt_object
|
|
self.shared_memory_objects[obj_id] = shm_obj
|
|
|
|
return ShmTransportMetadata(
|
|
tensor_meta=tensor_meta, tensor_device="cpu", shm_name=name, shm_size=size
|
|
)
|
|
|
|
def get_communicator_metadata(
|
|
self,
|
|
src_actor: "ray.actor.ActorHandle",
|
|
dst_actor: "ray.actor.ActorHandle",
|
|
backend: Optional[str] = None,
|
|
) -> CommunicatorMetadata:
|
|
return ShmCommunicatorMetadata()
|
|
|
|
def recv_multiple_tensors(
|
|
self,
|
|
obj_id: str,
|
|
tensor_transport_metadata: TensorTransportMetadata,
|
|
communicator_metadata: CommunicatorMetadata,
|
|
target_buffers: Optional[List[Any]] = None,
|
|
):
|
|
shm_name = tensor_transport_metadata.shm_name
|
|
size = tensor_transport_metadata.shm_size
|
|
shm_block = shm.SharedMemory(name=shm_name)
|
|
recv_tensors = pickle.loads(shm_block.buf[:size])
|
|
shm_block.close()
|
|
return recv_tensors
|
|
|
|
def send_multiple_tensors(
|
|
self,
|
|
tensors: List[numpy.ndarray],
|
|
tensor_transport_metadata: TensorTransportMetadata,
|
|
communicator_metadata: CommunicatorMetadata,
|
|
):
|
|
pass
|
|
|
|
def garbage_collect(
|
|
self,
|
|
obj_id: str,
|
|
tensor_transport_meta: TensorTransportMetadata,
|
|
tensors: List[numpy.ndarray],
|
|
):
|
|
self.shared_memory_objects[obj_id].close()
|
|
self.shared_memory_objects[obj_id].unlink()
|
|
del self.shared_memory_objects[obj_id]
|
|
|
|
def abort_transport(
|
|
self,
|
|
obj_id: str,
|
|
communicator_metadata: CommunicatorMetadata,
|
|
):
|
|
pass
|
|
|
|
|
|
def test_register_and_use_custom_transport(ray_start_regular):
|
|
register_tensor_transport(
|
|
"shared_memory", ["cpu"], SharedMemoryTransport, numpy.ndarray
|
|
)
|
|
|
|
@ray.remote
|
|
class Actor:
|
|
@ray.method(tensor_transport="shared_memory")
|
|
def echo(self, data):
|
|
return data
|
|
|
|
def non_rdt_echo(self, data):
|
|
return data
|
|
|
|
def sum(self, data):
|
|
return data.sum().item()
|
|
|
|
# Classes defined in test files get pickled by ref. So we need to
|
|
# explicitly pickle the transport class in this module by value.
|
|
# Note that this doesn't happen if you define the transport class on the
|
|
# driver, something with pytest convinces cloudpickle to pickle by ref.
|
|
from ray import cloudpickle
|
|
|
|
cloudpickle.register_pickle_by_value(sys.modules[SharedMemoryTransport.__module__])
|
|
|
|
actors = [Actor.remote() for _ in range(2)]
|
|
ref = actors[0].echo.remote(numpy.array([1, 2, 3]))
|
|
result = actors[1].sum.remote(ref)
|
|
assert ray.get(result) == 6
|
|
|
|
# Test that non-rdt methods that return the data type still work.
|
|
ref = actors[0].non_rdt_echo.remote(numpy.array([1, 2, 3]))
|
|
result = actors[1].sum.remote(ref)
|
|
assert ray.get(result) == 6
|
|
|
|
|
|
if __name__ == "__main__":
|
|
sys.exit(pytest.main(["-sv", __file__]))
|