deltacat 0.1.18b4__py3-none-any.whl → 0.1.18b7__py3-none-any.whl

deltacat/io/memcached_object_store.py

@@ -0,0 +1,121 @@
+import logging
+from ray import cloudpickle
+from collections import defaultdict
+import time
+from deltacat.io.object_store import IObjectStore
+from typing import Any, List
+from deltacat import logs
+import uuid
+import socket
+from pymemcache.client.base import Client
+from pymemcache.client.retrying import RetryingClient
+from pymemcache.exceptions import MemcacheUnexpectedCloseError
+
+logger = logs.configure_deltacat_logger(logging.getLogger(__name__))
+
+
+class MemcachedObjectStore(IObjectStore):
+    """
+    An implementation of object store that uses Memcached.
+    """
+
+    def __init__(self, port=11212) -> None:
+        self.client_cache = {}
+        self.current_ip = None
+        self.SEPARATOR = "_"
+        self.port = port
+        super().__init__()
+
+    def put_many(self, objects: List[object], *args, **kwargs) -> List[Any]:
+        input = {}
+        result = []
+        current_ip = self._get_current_ip()
+        for obj in objects:
+            serialized = cloudpickle.dumps(obj)
+            uid = uuid.uuid4()
+            ref = self._create_ref(uid, current_ip)
+            input[uid.__str__()] = serialized
+            result.append(ref)
+
+        client = self._get_client_by_ip(current_ip)
+        if client.set_many(input, noreply=False):
+            raise RuntimeError("Unable to write few keys to cache")
+
+        return result
+
+    def put(self, obj: object, *args, **kwargs) -> Any:
+        serialized = cloudpickle.dumps(obj)
+        uid = uuid.uuid4()
+        current_ip = self._get_current_ip()
+        ref = self._create_ref(uid, current_ip)
+        client = self._get_client_by_ip(current_ip)
+
+        if client.set(uid.__str__(), serialized):
+            return ref
+        else:
+            raise RuntimeError("Unable to write to cache")
+
+    def get_many(self, refs: List[Any], *args, **kwargs) -> List[object]:
+        result = []
+        uid_per_ip = defaultdict(lambda: [])
+
+        start = time.monotonic()
+        for ref in refs:
+            uid, ip = ref.split(self.SEPARATOR)
+            uid_per_ip[ip].append(uid)
+
+        for (ip, uids) in uid_per_ip.items():
+            client = self._get_client_by_ip(ip)
+            cache_result = client.get_many(uids)
+            assert len(cache_result) == len(
+                uids
+            ), f"Not all values were returned from cache as {len(cache_result)} != {len(uids)}"
+
+            values = cache_result.values()
+            total_bytes = 0
+
+            deserialize_start = time.monotonic()
+            for serialized in values:
+                deserialized = cloudpickle.loads(serialized)
+                total_bytes += len(serialized)
+                result.append(deserialized)
+
+            deserialize_end = time.monotonic()
+            logger.debug(
+                f"The time taken to deserialize {total_bytes} bytes is: {deserialize_end - deserialize_start}",
+            )
+
+        end = time.monotonic()
+
+        logger.info(f"The total time taken to read all objects is: {end - start}")
+        return result
+
+    def get(self, ref: Any, *args, **kwargs) -> object:
+        uid, ip = ref.split(self.SEPARATOR)
+        client = self._get_client_by_ip(ip)
+        serialized = client.get(uid)
+        return cloudpickle.loads(serialized)
+
+    def _create_ref(self, uid, ip) -> str:
+        return f"{uid}{self.SEPARATOR}{ip}"
+
+    def _get_client_by_ip(self, ip_address: str):
+        if ip_address in self.client_cache:
+            return self.client_cache[ip_address]
+
+        base_client = Client((ip_address, self.port))
+        client = RetryingClient(
+            base_client,
+            attempts=3,
+            retry_delay=0.01,
+            retry_for=[MemcacheUnexpectedCloseError],
+        )
+
+        self.client_cache[ip_address] = client
+        return client
+
+    def _get_current_ip(self):
+        if self.current_ip is None:
+            self.current_ip = socket.gethostbyname(socket.gethostname())
+
+        return self.current_ip

deltacat/io/object_store.py

@@ -0,0 +1,51 @@
+from typing import List, Any
+
+
+class IObjectStore:
+    """
+    An object store interface.
+    """
+
+    def setup(self, *args, **kwargs) -> Any:
+        ...
+
+        """
+        Sets up everything needed to run the object store.
+        """
+
+    def put(self, obj: object, *args, **kwargs) -> Any:
+        """
+        Put a single object into the object store. Depending
+        on the implementation, this method can be sync or async.
+        """
+        return self.put_many([obj])[0]
+
+    def put_many(self, objects: List[object], *args, **kwargs) -> List[Any]:
+        ...
+
+        """
+        Put many objects into the object store. It would return an ordered list
+        of object references corresponding to each object in the input.
+        """
+
+    def get(self, ref: Any, *args, **kwargs) -> object:
+        """
+        Get a single object from an object store.
+        """
+        return self.get_many([ref])[0]
+
+    def get_many(self, refs: List[Any], *args, **kwargs) -> List[object]:
+        ...
+
+        """
+        Get a list of objects from the object store. Use this method to
+        avoid multiple get calls. Note that depending on implementation it may
+        or may not return ordered results.
+        """
+
+    def clear(self, *args, **kwargs) -> bool:
+        ...
+
+        """
+        Clears the object store and all the associated data in it.
+        """

deltacat/io/ray_plasma_object_store.py

@@ -0,0 +1,23 @@
+import ray
+from ray import cloudpickle
+from deltacat.io.object_store import IObjectStore
+from typing import Any, List
+
+
+class RayPlasmaObjectStore(IObjectStore):
+    """
+    An implementation of object store that uses Ray plasma object store.
+    """
+
+    def put_many(self, objects: List[object], *args, **kwargs) -> List[Any]:
+        result = []
+        for obj in objects:
+            object_ref = ray.put(obj)
+            pickled = cloudpickle.dumps(object_ref)
+            result.append(pickled)
+
+        return result
+
+    def get_many(self, refs: List[Any], *args, **kwargs) -> List[object]:
+        loaded_refs = [cloudpickle.loads(obj_id) for obj_id in refs]
+        return ray.get(loaded_refs)

deltacat/io/redis_object_store.py

@@ -0,0 +1,114 @@
+import logging
+from ray import cloudpickle
+import time
+from deltacat.io.object_store import IObjectStore
+from typing import Any, List
+from deltacat import logs
+import uuid
+import socket
+import redis
+from collections import defaultdict
+
+
+logger = logs.configure_deltacat_logger(logging.getLogger(__name__))
+
+
+class RedisObjectStore(IObjectStore):
+    """
+    An implementation of object store that uses Redis in memory DB.
+    """
+
+    def __init__(self) -> None:
+        self.client_cache = {}
+        self.current_ip = None
+        self.SEPARATOR = "_"
+        super().__init__()
+
+    def put(self, obj: object, *args, **kwargs) -> Any:
+        serialized = cloudpickle.dumps(obj)
+        uid = uuid.uuid4()
+        current_ip = self._get_current_ip()
+        ref = self._create_ref(uid, current_ip)
+
+        client = self._get_client_by_ip(current_ip)
+        if client.set(uid.__str__(), serialized):
+            return ref
+        else:
+            raise RuntimeError(f"Unable to write {ref} to cache")
+
+    def put_many(self, objects: List[object], *args, **kwargs) -> List[Any]:
+        input = {}
+        result = []
+        current_ip = self._get_current_ip()
+        for obj in objects:
+            serialized = cloudpickle.dumps(obj)
+            uid = uuid.uuid4()
+            ref = self._create_ref(uid, current_ip)
+            input[uid.__str__()] = serialized
+            result.append(ref)
+
+        client = self._get_client_by_ip(current_ip)
+
+        if client.mset(input):
+            return result
+        else:
+            raise RuntimeError("Unable to update cache")
+
+    def get_many(self, refs: List[Any], *args, **kwargs) -> List[object]:
+        result = []
+        uid_per_ip = defaultdict(lambda: [])
+
+        start = time.monotonic()
+        for ref in refs:
+            uid, ip = ref.split(self.SEPARATOR)
+            uid_per_ip[ip].append(uid)
+
+        for (ip, uids) in uid_per_ip.items():
+            client = self._get_client_by_ip(ip)
+            cache_result = client.mget(uids)
+            assert len(cache_result) == len(
+                uids
+            ), "Not all values were returned from cache"
+
+            total_bytes = 0
+
+            deserialize_start = time.monotonic()
+            for serialized in cache_result:
+                deserialized = cloudpickle.loads(serialized)
+                total_bytes += len(serialized)
+                result.append(deserialized)
+
+            deserialize_end = time.monotonic()
+            logger.debug(
+                f"The time taken to deserialize {total_bytes} bytes is: {deserialize_end - deserialize_start}",
+            )
+
+        end = time.monotonic()
+
+        logger.info(f"The total time taken to read all objects is: {end - start}")
+
+        return result
+
+    def get(self, ref: Any, *args, **kwargs) -> object:
+        uid, ip = ref.split(self.SEPARATOR)
+        client = self._get_client_by_ip(ip)
+        serialized = client.get(uid)
+        return cloudpickle.loads(serialized)
+
+    def _get_client_by_ip(self, ip_address: str):
+        if ip_address in self.client_cache:
+            return self.client_cache[ip_address]
+
+        base_client = redis.Redis(ip_address, 7777)
+
+        self.client_cache[ip_address] = base_client
+        return base_client
+
+    def _get_current_ip(self):
+        if self.current_ip is None:
+            self.current_ip = socket.gethostbyname(socket.gethostname())
+
+        return self.current_ip
+
+    def _create_ref(self, uid, ip):
+        return f"{uid}{self.SEPARATOR}{ip}"

deltacat/io/s3_object_store.py

@@ -0,0 +1,44 @@
+import logging
+from ray import cloudpickle
+import time
+from deltacat.io.object_store import IObjectStore
+from typing import Any, List
+from deltacat import logs
+import uuid
+from deltacat.aws import s3u as s3_utils
+
+logger = logs.configure_deltacat_logger(logging.getLogger(__name__))
+
+
+class S3ObjectStore(IObjectStore):
+    """
+    An implementation of object store that uses S3.
+    """
+
+    def __init__(self, bucket_prefix: str) -> None:
+        self.bucket = bucket_prefix
+        super().__init__()
+
+    def put_many(self, objects: List[object], *args, **kwargs) -> List[Any]:
+        result = []
+        for obj in objects:
+            serialized = cloudpickle.dumps(obj)
+            ref = uuid.uuid4()
+
+            s3_utils.upload(f"s3://{self.bucket}/{ref}", serialized)
+            result.append(ref)
+
+        return result
+
+    def get_many(self, refs: List[Any], *args, **kwargs) -> List[object]:
+        result = []
+        start = time.monotonic()
+        for ref in refs:
+            cur = s3_utils.download(f"s3://{self.bucket}/{ref}")
+            serialized = cur["Body"].read()
+            loaded = cloudpickle.loads(serialized)
+            result.append(loaded)
+        end = time.monotonic()
+
+        logger.info(f"The total time taken to read all objects is: {end - start}")
+        return result

deltacat/tests/compactor/utils/test_io.py

@@ -17,6 +17,10 @@ class TestFitInputDeltas(unittest.TestCase):
 
         super().setUpClass()
 
+    @classmethod
+    def tearDownClass(cls) -> None:
+        cls.module_patcher.stop()
+
     def test_sanity(self):
         from deltacat.compute.compactor.utils import io
 

deltacat/tests/io/test_file_object_store.py

@@ -0,0 +1,86 @@
+import unittest
+from unittest import mock
+
+
+class TestFileObjectStore(unittest.TestCase):
+
+    TEST_VALUE = "test-value"
+
+    @classmethod
+    def setUpClass(cls):
+        cls.ray_mock = mock.MagicMock()
+        cls.os_mock = mock.MagicMock()
+
+        cls.module_patcher = mock.patch.dict(
+            "sys.modules", {"ray": cls.ray_mock, "os": cls.os_mock}
+        )
+        cls.module_patcher.start()
+
+        super().setUpClass()
+
+    @classmethod
+    def tearDownClass(cls) -> None:
+        cls.module_patcher.stop()
+
+    @mock.patch(
+        "deltacat.io.file_object_store.open",
+        new_callable=mock.mock_open,
+        read_data="data",
+    )
+    def test_put_many_sanity(self, mock_file):
+        from deltacat.io.file_object_store import FileObjectStore
+
+        object_store = FileObjectStore(dir_path="")
+        self.ray_mock.cloudpickle.dumps.return_value = self.TEST_VALUE
+        result = object_store.put_many(["a", "b"])
+
+        self.assertEqual(2, len(result))
+        self.assertEqual(2, mock_file.call_count)
+
+    @mock.patch(
+        "deltacat.io.file_object_store.open",
+        new_callable=mock.mock_open,
+        read_data="data",
+    )
+    def test_get_many_sanity(self, mock_file):
+        from deltacat.io.file_object_store import FileObjectStore
+
+        object_store = FileObjectStore(dir_path="")
+        self.ray_mock.cloudpickle.loads.return_value = self.TEST_VALUE
+
+        result = object_store.get_many(["test", "test"])
+
+        self.assertEqual(2, len(result))
+        self.assertEqual(2, mock_file.call_count)
+
+    @mock.patch(
+        "deltacat.io.file_object_store.open",
+        new_callable=mock.mock_open,
+        read_data="data",
+    )
+    def test_get_sanity(self, mock_file):
+        from deltacat.io.file_object_store import FileObjectStore
+
+        object_store = FileObjectStore(dir_path="")
+        self.ray_mock.cloudpickle.loads.return_value = self.TEST_VALUE
+
+        result = object_store.get("test")
+
+        self.assertEqual(self.TEST_VALUE, result)
+        self.assertEqual(1, mock_file.call_count)
+
+    @mock.patch(
+        "deltacat.io.file_object_store.open",
+        new_callable=mock.mock_open,
+        read_data="data",
+    )
+    def test_put_sanity(self, mock_file):
+        from deltacat.io.file_object_store import FileObjectStore
+
+        object_store = FileObjectStore(dir_path="")
+        self.ray_mock.cloudpickle.dumps.return_value = self.TEST_VALUE
+
+        result = object_store.put("test")
+
+        self.assertIsNotNone(result)
+        self.assertEqual(1, mock_file.call_count)

deltacat/tests/io/test_memcached_object_store.py

@@ -0,0 +1,158 @@
+import unittest
+from unittest import mock
+
+
+@mock.patch("deltacat.io.memcached_object_store.cloudpickle")
+@mock.patch("deltacat.io.memcached_object_store.socket")
+class TestMemcachedObjectStore(unittest.TestCase):
+
+    TEST_VALUE = "test-value"
+
+    def setUp(self):
+        from deltacat.io.memcached_object_store import MemcachedObjectStore
+
+        self.object_store = MemcachedObjectStore()
+
+    @mock.patch("deltacat.io.memcached_object_store.Client")
+    @mock.patch("deltacat.io.memcached_object_store.RetryingClient")
+    def test_put_many_sanity(
+        self,
+        mock_retrying_client,
+        mock_client,
+        mock_socket,
+        mock_cloudpickle,
+    ):
+        mock_cloudpickle.dumps.return_value = self.TEST_VALUE
+        mock_cloudpickle.loads.return_value = self.TEST_VALUE
+        mock_socket.gethostbyname.return_value = "0.0.0.0"
+        mock_socket.gethostname.return_value = "test-host"
+        mock_retrying_client.return_value = mock_client.return_value
+        mock_client.return_value.set_many.return_value = []
+
+        result = self.object_store.put_many(["a", "b"])
+
+        self.assertEqual(2, len(result))
+        self.assertRegex(result[0], ".*_.*")
+        self.assertEqual(1, mock_client.return_value.set_many.call_count)
+
+    @mock.patch("deltacat.io.memcached_object_store.Client")
+    @mock.patch("deltacat.io.memcached_object_store.RetryingClient")
+    def test_put_many_when_cache_fails(
+        self, mock_retrying_client, mock_client, mock_socket, mock_cloudpickle
+    ):
+        mock_cloudpickle.dumps.return_value = self.TEST_VALUE
+        mock_cloudpickle.loads.return_value = self.TEST_VALUE
+        mock_socket.gethostbyname.return_value = "0.0.0.0"
+        mock_socket.gethostname.return_value = "test-host"
+        mock_retrying_client.return_value = mock_client.return_value
+        mock_client.return_value.set_many.return_value = ["abcd"]
+
+        with self.assertRaises(RuntimeError):
+            self.object_store.put_many(["a", "b"])
+
+        self.assertEqual(1, mock_client.return_value.set_many.call_count)
+
+    @mock.patch("deltacat.io.memcached_object_store.Client")
+    @mock.patch("deltacat.io.memcached_object_store.RetryingClient")
+    def test_get_many_sanity(
+        self, mock_retrying_client, mock_client, mock_socket, mock_cloudpickle
+    ):
+        mock_cloudpickle.dumps.return_value = self.TEST_VALUE
+        mock_cloudpickle.loads.return_value = self.TEST_VALUE
+        mock_socket.gethostbyname.return_value = "0.0.0.0"
+        mock_socket.gethostname.return_value = "test-host"
+        mock_client.return_value.get_many.return_value = {
+            "key1": "value1",
+            "key2": "value2",
+        }
+        mock_retrying_client.return_value = mock_client.return_value
+
+        result = self.object_store.get_many(["test_ip", "test_ip"])
+
+        self.assertEqual(2, len(result))
+        self.assertEqual(1, mock_client.return_value.get_many.call_count)
+
+    @mock.patch("deltacat.io.memcached_object_store.Client")
+    @mock.patch("deltacat.io.memcached_object_store.RetryingClient")
+    def test_get_many_when_cache_expired(
+        self, mock_retrying_client, mock_client, mock_socket, mock_cloudpickle
+    ):
+        mock_cloudpickle.dumps.return_value = self.TEST_VALUE
+        mock_cloudpickle.loads.return_value = self.TEST_VALUE
+        mock_socket.gethostbyname.return_value = "0.0.0.0"
+        mock_socket.gethostname.return_value = "test-host"
+        mock_client.return_value.get_many.return_value = {"key1": "value1"}
+        mock_retrying_client.return_value = mock_client.return_value
+
+        with self.assertRaises(AssertionError):
+            self.object_store.get_many(["test_ip", "test_ip"])
+
+        self.assertEqual(1, mock_client.return_value.get_many.call_count)
+
+    @mock.patch("deltacat.io.memcached_object_store.Client")
+    @mock.patch("deltacat.io.memcached_object_store.RetryingClient")
+    def test_get_sanity(
+        self, mock_retrying_client, mock_client, mock_socket, mock_cloudpickle
+    ):
+        mock_cloudpickle.dumps.return_value = self.TEST_VALUE
+        mock_cloudpickle.loads.return_value = self.TEST_VALUE
+        mock_socket.gethostbyname.return_value = "0.0.0.0"
+        mock_socket.gethostname.return_value = "test-host"
+        mock_client.return_value.get.return_value = self.TEST_VALUE
+        mock_retrying_client.return_value = mock_client.return_value
+
+        result = self.object_store.get("test_ip")
+
+        self.assertEqual(self.TEST_VALUE, result)
+        self.assertEqual(1, mock_client.return_value.get.call_count)
+
+    @mock.patch("deltacat.io.memcached_object_store.Client")
+    @mock.patch("deltacat.io.memcached_object_store.RetryingClient")
+    def test_get_when_cache_fails(
+        self, mock_retrying_client, mock_client, mock_socket, mock_cloudpickle
+    ):
+        mock_cloudpickle.dumps.return_value = self.TEST_VALUE
+        mock_cloudpickle.loads.return_value = self.TEST_VALUE
+        mock_socket.gethostbyname.return_value = "0.0.0.0"
+        mock_socket.gethostname.return_value = "test-host"
+        mock_client.return_value.get.side_effect = RuntimeError()
+        mock_retrying_client.return_value = mock_client.return_value
+
+        with self.assertRaises(RuntimeError):
+            self.object_store.get("test_ip")
+
+        self.assertEqual(1, mock_client.return_value.get.call_count)
+
+    @mock.patch("deltacat.io.memcached_object_store.Client")
+    @mock.patch("deltacat.io.memcached_object_store.RetryingClient")
+    def test_put_sanity(
+        self, mock_retrying_client, mock_client, mock_socket, mock_cloudpickle
+    ):
+        mock_cloudpickle.dumps.return_value = self.TEST_VALUE
+        mock_cloudpickle.loads.return_value = self.TEST_VALUE
+        mock_socket.gethostbyname.return_value = "0.0.0.0"
+        mock_socket.gethostname.return_value = "test-host"
+        mock_retrying_client.return_value = mock_client.return_value
+        mock_client.return_value.set.return_value = True
+
+        result = self.object_store.put("test")
+
+        self.assertIsNotNone(result)
+        self.assertEqual(1, mock_client.return_value.set.call_count)
+
+    @mock.patch("deltacat.io.memcached_object_store.Client")
+    @mock.patch("deltacat.io.memcached_object_store.RetryingClient")
+    def test_put_when_cache_fails(
+        self, mock_retrying_client, mock_client, mock_socket, mock_cloudpickle
+    ):
+        mock_cloudpickle.dumps.return_value = self.TEST_VALUE
+        mock_cloudpickle.loads.return_value = self.TEST_VALUE
+        mock_socket.gethostbyname.return_value = "0.0.0.0"
+        mock_socket.gethostname.return_value = "test-host"
+        mock_retrying_client.return_value = mock_client.return_value
+        mock_client.return_value.set.return_value = False
+
+        with self.assertRaises(RuntimeError):
+            self.object_store.put("test_ip")
+
+        self.assertEqual(1, mock_client.return_value.set.call_count)

deltacat/tests/io/test_ray_plasma_object_store.py

@@ -0,0 +1,54 @@
+import unittest
+from unittest import mock
+
+
+class TestRayPlasmaObjectStore(unittest.TestCase):
+
+    TEST_VALUE = "test-value"
+
+    @classmethod
+    def setUpClass(cls):
+        from deltacat.io.ray_plasma_object_store import RayPlasmaObjectStore
+
+        cls.object_store = RayPlasmaObjectStore()
+
+        super().setUpClass()
+
+    @mock.patch("deltacat.io.ray_plasma_object_store.ray")
+    @mock.patch("deltacat.io.ray_plasma_object_store.cloudpickle")
+    def test_put_many_sanity(self, mock_cloudpickle, mock_ray):
+        mock_ray.put.return_value = "c"
+        mock_cloudpickle.dumps.return_value = self.TEST_VALUE
+        result = self.object_store.put_many(["a", "b"])
+
+        self.assertEqual(2, len(result))
+
+    @mock.patch("deltacat.io.ray_plasma_object_store.ray")
+    @mock.patch("deltacat.io.ray_plasma_object_store.cloudpickle")
+    def test_get_many_sanity(self, mock_cloudpickle, mock_ray):
+        mock_ray.get.return_value = ["a", "b"]
+        mock_cloudpickle.loads.return_value = self.TEST_VALUE
+
+        result = self.object_store.get_many(["test", "test"])
+
+        self.assertEqual(2, len(result))
+
+    @mock.patch("deltacat.io.ray_plasma_object_store.ray")
+    @mock.patch("deltacat.io.ray_plasma_object_store.cloudpickle")
+    def test_get_sanity(self, mock_cloudpickle, mock_ray):
+        mock_ray.get.return_value = [self.TEST_VALUE]
+        mock_cloudpickle.loads.return_value = self.TEST_VALUE
+
+        result = self.object_store.get("test")
+
+        self.assertEqual(self.TEST_VALUE, result)
+
+    @mock.patch("deltacat.io.ray_plasma_object_store.ray")
+    @mock.patch("deltacat.io.ray_plasma_object_store.cloudpickle")
+    def test_put_sanity(self, mock_cloudpickle, mock_ray):
+        mock_ray.put.return_value = "c"
+        mock_cloudpickle.dumps.return_value = self.TEST_VALUE
+
+        result = self.object_store.put("test")
+
+        self.assertEqual(self.TEST_VALUE, result)