[Object Spilling] Initial Iteration of S3 adapter. (#11379)

* Finished the first iteration.

* Removed unnecessary code.

* Smartopen impl.

* Make sure tests passed.

* Addressed code review.

* Addressed code review.

* Fix issues.

* Fix issues.

python/ray/external_storage.py

@@ -1,6 +1,7 @@
 import abc
 import os
 from typing import List
+
 import ray
 
 
@@ -10,6 +11,14 @@ class ExternalStorage(metaclass=abc.ABCMeta):
     This class provides some useful functions for zero-copy object
     put/get from plasma store. Also it specifies the interface for
     object spilling.
+
+    When inheriting this class, please make sure to implement validation
+    logic inside __init__ method. When ray instance starts, it will
+    instantiating external storage to validate the config.
+
+    Raises:
+        ValueError: when given configuration for
+            the external storage is invalid.
     """
 
     def _get_objects_from_store(self, object_refs):
@@ -58,11 +67,19 @@ class NullStorage(ExternalStorage):
 
 
 class FileSystemStorage(ExternalStorage):
-    """The class for filesystem-like external storage."""
+    """The class for filesystem-like external storage.
+
+    Raises:
+        ValueError: Raises directory path to
+            spill objects doesn't exist.
+    """
 
     def __init__(self, directory_path):
         self.directory_path = directory_path
         self.prefix = "ray_spilled_object_"
+        if not os.path.exists(self.directory_path):
+            raise ValueError("The given directory path to store objects, "
+                             f"{self.directory_path}, doesn't exist.")
 
     def spill_objects(self, object_refs):
         keys = []
@@ -91,6 +108,84 @@ class FileSystemStorage(ExternalStorage):
                 self._put_object_to_store(metadata, buf_len, f, ref)
 
 
+class ExternalStorageSmartOpenImpl(ExternalStorage):
+    """The external storage class implemented by smart_open.
+    (https://github.com/RaRe-Technologies/smart_open)
+
+    Smart open supports multiple backend with the same APIs.
+
+    Args:
+        uri(str): Storage URI used for smart open.
+        prefix(str): Prefix of objects that are stored.
+        override_transport_params(dict): Overriding the default value of
+            transport_params for smart-open library.
+
+    Raises:
+        ModuleNotFoundError: If it fails to setup.
+            For example, if smart open library
+            is not downloaded, this will fail.
+    """
+
+    def __init__(self,
+                 uri: str,
+                 prefix: str = "ray_spilled_object_",
+                 override_transport_params: dict = None):
+        try:
+            from smart_open import open  # noqa
+        except ModuleNotFoundError as e:
+            raise ModuleNotFoundError(
+                "Smart open is chosen to be a object spilling "
+                "external storage, but smart_open "
+                f"is not downloaded. Original error: {e}")
+
+        self.uri = uri.strip("/")
+        self.prefix = prefix
+        self.override_transport_params = override_transport_params or {}
+        self.transport_params = {}.update(self.override_transport_params)
+
+    def spill_objects(self, object_refs):
+        keys = []
+        ray_object_pairs = self._get_objects_from_store(object_refs)
+        for ref, (buf, metadata) in zip(object_refs, ray_object_pairs):
+            key = self.prefix + ref.hex()
+            self._spill_object(key, ref, buf, metadata)
+            keys.append(key.encode())
+        return keys
+
+    def restore_spilled_objects(self, keys):
+        for k in keys:
+            key = k.decode()
+            ref = ray.ObjectRef(bytes.fromhex(key[len(self.prefix):]))
+            self._restore_spilled_object(key, ref)
+
+    def _spill_object(self, key, ref, buf, metadata):
+        from smart_open import open
+        with open(
+                self._build_uri(key), "wb",
+                transport_params=self.transport_params) as file_like:
+            metadata_len = len(metadata)
+            buf_len = len(buf)
+            file_like.write(metadata_len.to_bytes(8, byteorder="little"))
+            file_like.write(buf_len.to_bytes(8, byteorder="little"))
+            file_like.write(metadata)
+            file_like.write(memoryview(buf))
+
+    def _restore_spilled_object(self, key, ref):
+        from smart_open import open
+        with open(
+                self._build_uri(key), "rb",
+                transport_params=self.transport_params) as file_like:
+            metadata_len = int.from_bytes(
+                file_like.read(8), byteorder="little")
+            buf_len = int.from_bytes(file_like.read(8), byteorder="little")
+            metadata = file_like.read(metadata_len)
+            # read remaining data to our buffer
+            self._put_object_to_store(metadata, buf_len, file_like, ref)
+
+    def _build_uri(self, key):
+        return f"{self.uri}/{key}"
+
+
 _external_storage = NullStorage()
 
 
@@ -101,12 +196,20 @@ def setup_external_storage(config):
         storage_type = config["type"]
         if storage_type == "filesystem":
             _external_storage = FileSystemStorage(**config["params"])
+        elif storage_type == "smart_open":
+            _external_storage = ExternalStorageSmartOpenImpl(
+                **config["params"])
         else:
             raise ValueError(f"Unknown external storage type: {storage_type}")
     else:
         _external_storage = NullStorage()
 
 
+def reset_external_storage():
+    global _external_storage
+    _external_storage = NullStorage()
+
+
 def spill_objects(object_refs):
     """Spill objects to the external storage. Objects are specified
     by their object refs.