Add hashed memcached client support

deltacat/io/memcached_object_store.py

@@ -3,13 +3,14 @@
 from collections import defaultdict
 import time
 from deltacat.io.object_store import IObjectStore
-from typing import Any, List
+from typing import Any, List, Optional
 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
+from pymemcache.client.rendezvous import RendezvousHash
 
 logger = logs.configure_deltacat_logger(logging.getLogger(__name__))
 
@@ -19,36 +20,46 @@ class MemcachedObjectStore(IObjectStore):
     An implementation of object store that uses Memcached.
     """
 
-    def __init__(self, port=11212) -> None:
+    def __init__(
+        self, storage_node_ips: Optional[List[str]] = None, port: Optional[int] = 11212
+    ) -> None:
         self.client_cache = {}
         self.current_ip = None
         self.SEPARATOR = "_"
         self.port = port
+        self.storage_node_ips = storage_node_ips
+        self.hasher = None
         super().__init__()
 
+    def initialize_hasher(self):
+        if not self.hasher and self.storage_node_ips:
+            self.hasher = RendezvousHash()
+            for n in self.storage_node_ips:
+                self.hasher.add_node(n)
+
     def put_many(self, objects: List[object], *args, **kwargs) -> List[Any]:
-        input = {}
+        input = defaultdict(dict)
         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
+            create_ref_ip = self._get_create_ref_ip(uid.__str__())
+            ref = self._create_ref(uid, create_ref_ip)
+            input[create_ref_ip][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")
+        for create_ref_ip, uid_to_object in input.items():
+            client = self._get_client_by_ip(create_ref_ip)
+            if client.set_many(uid_to_object, 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)
+        create_ref_ip = self._get_create_ref_ip(uid.__str__())
+        ref = self._create_ref(uid, create_ref_ip)
+        client = self._get_client_by_ip(create_ref_ip)
 
         if client.set(uid.__str__(), serialized):
             return ref
@@ -99,6 +110,18 @@ def get(self, ref: Any, *args, **kwargs) -> object:
     def _create_ref(self, uid, ip) -> str:
         return f"{uid}{self.SEPARATOR}{ip}"
 
+    def _get_storage_node_ip(self, key: str):
+        self.initialize_hasher()
+        storage_node_ip = self.hasher.get_node(key)
+        return storage_node_ip
+
+    def _get_create_ref_ip(self, uid: str):
+        if self.storage_node_ips:
+            create_ref_ip = self._get_storage_node_ip(uid)
+        else:
+            create_ref_ip = self._get_current_ip()
+        return create_ref_ip
+
     def _get_client_by_ip(self, ip_address: str):
         if ip_address in self.client_cache:
             return self.client_cache[ip_address]
@@ -108,7 +131,7 @@ def _get_client_by_ip(self, ip_address: str):
             base_client,
             attempts=3,
             retry_delay=0.01,
-            retry_for=[MemcacheUnexpectedCloseError],
+            retry_for=[MemcacheUnexpectedCloseError, ConnectionResetError],
         )
 
         self.client_cache[ip_address] = client

deltacat/tests/io/test_memcached_object_store.py

@@ -11,7 +11,9 @@ class TestMemcachedObjectStore(unittest.TestCase):
     def setUp(self):
         from deltacat.io.memcached_object_store import MemcachedObjectStore
 
-        self.object_store = MemcachedObjectStore()
+        self.object_store = MemcachedObjectStore(
+            storage_node_ips=["172.1.1.1", "172.2.2.2", "172.3.3.3"]
+        )
 
     @mock.patch("deltacat.io.memcached_object_store.Client")
     @mock.patch("deltacat.io.memcached_object_store.RetryingClient")
@@ -29,11 +31,10 @@ def test_put_many_sanity(
         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"])
+        result = self.object_store.put_many(["a", "b", "c"])
 
-        self.assertEqual(2, len(result))
+        self.assertEqual(3, 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")

deltacat/utils/placement.py

@@ -21,9 +21,10 @@
 
 @dataclass
 class PlacementGroupConfig:
-    def __init__(self, opts, resource):
+    def __init__(self, opts, resource, node_ips):
         self.opts = opts
         self.resource = resource
+        self.node_ips = node_ips
 
 
 class NodeGroupManager:
@@ -275,17 +276,22 @@ def _config(
     # query available resources given list of node id
     all_nodes_available_res = ray._private.state.state._available_resources_per_node()
     pg_res = {"CPU": 0, "memory": 0, "object_store_memory": 0}
+    node_ips = []
     for node_id in node_ids:
         if node_id in all_nodes_available_res:
             v = all_nodes_available_res[node_id]
             node_detail = get_node(node_id)
             pg_res["CPU"] += node_detail["resources_total"]["CPU"]
             pg_res["memory"] += v["memory"]
             pg_res["object_store_memory"] += v["object_store_memory"]
+            node_ips.append(node_detail["node_ip"])
     cluster_resources["CPU"] = int(pg_res["CPU"])
     cluster_resources["memory"] = float(pg_res["memory"])
     cluster_resources["object_store_memory"] = float(pg_res["object_store_memory"])
-    pg_config = PlacementGroupConfig(opts, cluster_resources)
+    pg_config = PlacementGroupConfig(
+        opts=opts, cluster_resources=cluster_resources, node_ips=node_ips
+    )
     logger.info(f"pg has resources:{cluster_resources}")
+    logger.debug(f"pg has node ips:{node_ips}")
 
     return pg_config