Repository: mars-project/mars

Issue:
Ray storage
Implement ray storage backend for https://github.com/mars-project/mars/issues/1905


Discussion:

Resolving patch:
diff --git a/mars/storage/ray.py b/mars/storage/ray.py
new file mode 100644
--- /dev/null
+++ b/mars/storage/ray.py
@@ -0,0 +1,155 @@
+from typing import Any, Dict, List, Tuple
+from ..lib import sparse
+from ..utils import lazy_import, implements
+from .base import StorageBackend, StorageLevel, ObjectInfo
+from .core import BufferWrappedFileObject, StorageFileObject
+
+ray = lazy_import("ray")
+
+
+# TODO(fyrestone): make the SparseMatrix pickleable.
+
+def _mars_sparse_matrix_serializer(value):
+    return [value.shape, value.spmatrix]
+
+
+def _mars_sparse_matrix_deserializer(obj) -> sparse.SparseNDArray:
+    shape, spmatrix = obj
+    return sparse.matrix.SparseMatrix(spmatrix, shape=shape)
+
+
+def _register_sparse_matrix_serializer():
+    # register a custom serializer for Mars SparseMatrix
+    try:
+        ray.register_custom_serializer(
+                sparse.matrix.SparseMatrix,
+                serializer=_mars_sparse_matrix_serializer,
+                deserializer=_mars_sparse_matrix_deserializer)
+    except AttributeError:  # ray >= 1.0
+        try:
+            from ray.worker import global_worker
+
+            global_worker.check_connected()
+            context = global_worker.get_serialization_context()
+            context.register_custom_serializer(
+                    sparse.matrix.SparseMatrix,
+                    serializer=_mars_sparse_matrix_serializer,
+                    deserializer=_mars_sparse_matrix_deserializer)
+        except AttributeError:  # ray >= 1.2.0
+            ray.util.register_serializer(
+                    sparse.matrix.SparseMatrix,
+                    serializer=_mars_sparse_matrix_serializer,
+                    deserializer=_mars_sparse_matrix_deserializer)
+
+
+class RayFileLikeObject:
+    def __init__(self):
+        self._buffers = []
+        self._size = 0
+
+    def write(self, content: bytes):
+        self._buffers.append(content)
+        self._size += len(content)
+
+    def readinto(self, buffer):
+        read_bytes = 0
+        for b in self._buffers:
+            read_pos = read_bytes + len(b)
+            buffer[read_bytes:read_pos] = b
+            read_bytes = read_pos
+        return read_bytes
+
+    def close(self):
+        self._buffers.clear()
+        self._size = 0
+
+    def tell(self):
+        return self._size
+
+
+class RayFileObject(BufferWrappedFileObject):
+    def __init__(self, object_id: Any, mode: str):
+        self._object_id = object_id
+        super().__init__(mode, size=0)
+
+    def _write_init(self):
+        self._buffer = RayFileLikeObject()
+
+    def _read_init(self):
+        self._buffer = ray.get(self._object_id)
+        self._mv = memoryview(self._buffer)
+        self._size = len(self._buffer)
+
+    def write(self, content: bytes):
+        if not self._initialized:
+            self._write_init()
+            self._initialized = True
+
+        return self._buffer.write(content)
+
+    def _write_close(self):
+        worker = ray.worker.global_worker
+        metadata = ray.ray_constants.OBJECT_METADATA_TYPE_RAW
+        worker.core_worker.put_file_like_object(metadata, self._buffer.tell(), self._buffer,
+                                                self._object_id)
+
+    def _read_close(self):
+        pass
+
+
+class RayStorage(StorageBackend):
+    def __init__(self, *args, **kwargs):
+        pass
+
+    @classmethod
+    @implements(StorageBackend.setup)
+    async def setup(cls, **kwargs) -> Tuple[Dict, Dict]:
+        ray.init(ignore_reinit_error=True)
+        _register_sparse_matrix_serializer()
+        return dict(), dict()
+
+    @staticmethod
+    @implements(StorageBackend.teardown)
+    async def teardown(**kwargs):
+        pass
+
+    @property
+    @implements(StorageBackend.level)
+    def level(self) -> StorageLevel:
+        # TODO(fyrestone): return StorageLevel.MEMORY & StorageLevel.DISK
+        # if object spilling is available.
+        return StorageLevel.MEMORY
+
+    @implements(StorageBackend.get)
+    async def get(self, object_id, **kwargs) -> object:
+        return await object_id
+
+    @implements(StorageBackend.put)
+    async def put(self, obj, importance=0) -> ObjectInfo:
+        object_id = ray.put(obj)
+        # We can't get the serialized bytes length from ray.put
+        return ObjectInfo(object_id=object_id)
+
+    @implements(StorageBackend.delete)
+    async def delete(self, object_id):
+        ray.internal.free(object_id)
+
+    @implements(StorageBackend.object_info)
+    async def object_info(self, object_id) -> ObjectInfo:
+        # The performance of obtaining the object size is poor.
+        return ObjectInfo(object_id=object_id)
+
+    @implements(StorageBackend.open_writer)
+    async def open_writer(self, size=None) -> StorageFileObject:
+        new_id = ray.ObjectRef.from_random()
+        ray_writer = RayFileObject(new_id, mode='w')
+        return StorageFileObject(ray_writer, object_id=new_id)
+
+    @implements(StorageBackend.open_reader)
+    async def open_reader(self, object_id) -> StorageFileObject:
+        ray_reader = RayFileObject(object_id, mode='r')
+        return StorageFileObject(ray_reader, object_id=object_id)
+
+    @implements(StorageBackend.list)
+    async def list(self) -> List:
+        raise NotImplementedError("Ray storage does not support list")


Test patch:
diff --git a/mars/storage/tests/test_libs.py b/mars/storage/tests/test_libs.py
--- a/mars/storage/tests/test_libs.py
+++ b/mars/storage/tests/test_libs.py
@@ -32,6 +32,7 @@
 from mars.storage.plasma import PlasmaStorage
 from mars.storage.shared_memory import SharedMemoryStorage
 from mars.storage.vineyard import VineyardStorage
+from mars.storage.ray import RayStorage
 try:
     import vineyard
 except ImportError:
@@ -42,11 +43,17 @@
 except ImportError:
     cupy = None
     cudf = None
+try:
+    import ray
+except ImportError:
+    ray = None
 
 
 params = ['filesystem', 'plasma']
 if vineyard:
     params.append('vineyard')
+if ray:
+    params.append('ray')
 if sys.version_info[:2] >= (3, 8):
     params.append('shared_memory')
 
@@ -103,6 +110,14 @@ async def storage_context(request):
 
         teardown_params['object_ids'] = storage._object_ids
         await SharedMemoryStorage.teardown(**teardown_params)
+    elif request.param == 'ray':
+        params, teardown_params = await RayStorage.setup()
+        storage = RayStorage(**params)
+        assert storage.level == StorageLevel.MEMORY
+
+        yield storage
+
+        await RayStorage.teardown(**teardown_params)
 
 
 @pytest.mark.asyncio
@@ -128,7 +143,7 @@ async def test_base_operations(storage_context):
     assert info2.size == put_info2.size
 
     # FIXME: remove when list functionality is ready for vineyard.
-    if not isinstance(storage, (VineyardStorage, SharedMemoryStorage)):
+    if not isinstance(storage, (VineyardStorage, SharedMemoryStorage, RayStorage)):
         num = len(await storage.list())
         assert num == 2
         await storage.delete(info2.object_id)