Add gcs object manager (#8298)

python/ray/gcs_utils.py

@@ -14,6 +14,7 @@ from ray.core.generated.gcs_pb2 import (
     TablePubsub,
     TaskTableData,
     ResourceTableData,
+    ObjectLocationInfo,
 )
 
 __all__ = [
@@ -33,6 +34,7 @@ __all__ = [
     "TaskTableData",
     "ResourceTableData",
     "construct_error_message",
+    "ObjectLocationInfo",
 ]
 
 FUNCTION_PREFIX = "RemoteFunction:"

python/ray/includes/global_state_accessor.pxd

@@ -1,6 +1,10 @@
 from libcpp.string cimport string as c_string
 from libcpp cimport bool as c_bool
 from libcpp.vector cimport vector as c_vector
+from libcpp.memory cimport unique_ptr
+from ray.includes.unique_ids cimport (
+    CObjectID
+)
 
 cdef extern from "ray/gcs/gcs_client/global_state_accessor.h" nogil:
     cdef cppclass CGlobalStateAccessor "ray::gcs::GlobalStateAccessor":
@@ -11,3 +15,5 @@ cdef extern from "ray/gcs/gcs_client/global_state_accessor.h" nogil:
         void Disconnect()
         c_vector[c_string] GetAllJobInfo()
         c_vector[c_string] GetAllProfileInfo()
+        c_vector[c_string] GetAllObjectInfo()
+        unique_ptr[c_string] GetObjectInfo(const CObjectID &object_id)

python/ray/includes/global_state_accessor.pxi

@@ -1,7 +1,13 @@
+from ray.includes.unique_ids cimport (
+    CObjectID
+)
+
 from ray.includes.global_state_accessor cimport (
     CGlobalStateAccessor,
 )
 
+from libcpp.string cimport string as c_string
+
 cdef class GlobalStateAccessor:
     """Cython wrapper class of C++ `ray::gcs::GlobalStateAccessor`."""
     cdef:
@@ -25,3 +31,12 @@ cdef class GlobalStateAccessor:
 
     def get_profile_table(self):
         return self.inner.get().GetAllProfileInfo()
+
+    def get_object_table(self):
+        return self.inner.get().GetAllObjectInfo()
+
+    def get_object_info(self, object_id):
+        object_info = self.inner.get().GetObjectInfo(CObjectID.FromBinary(object_id.binary()))
+        if object_info:
+            return c_string(object_info.get().data(), object_info.get().size())
+        return None

python/ray/state.py

@@ -10,8 +10,7 @@ from ray import (
     gcs_utils,
     services,
 )
-from ray.utils import (decode, binary_to_object_id, binary_to_hex,
-                       hex_to_binary)
+from ray.utils import (decode, binary_to_hex, hex_to_binary)
 
 from ray._raylet import GlobalStateAccessor
 
@@ -256,38 +255,6 @@ class GlobalState:
             result.extend(list(client.scan_iter(match=pattern)))
         return result
 
-    def _object_table(self, object_id):
-        """Fetch and parse the object table information for a single object ID.
-
-        Args:
-            object_id: An object ID to get information about.
-
-        Returns:
-            A dictionary with information about the object ID in question.
-        """
-        # Allow the argument to be either an ObjectID or a hex string.
-        if not isinstance(object_id, ray.ObjectID):
-            object_id = ray.ObjectID(hex_to_binary(object_id))
-
-        # Return information about a single object ID.
-        message = self._execute_command(object_id, "RAY.TABLE_LOOKUP",
-                                        gcs_utils.TablePrefix.Value("OBJECT"),
-                                        "", object_id.binary())
-        if message is None:
-            return {}
-        gcs_entry = gcs_utils.GcsEntry.FromString(message)
-
-        assert len(gcs_entry.entries) > 0
-
-        entry = gcs_utils.ObjectTableData.FromString(gcs_entry.entries[0])
-
-        object_info = {
-            "DataSize": entry.object_size,
-            "Manager": entry.manager,
-        }
-
-        return object_info
-
     def object_table(self, object_id=None):
         """Fetch and parse the object table info for one or more object IDs.
 
@@ -299,23 +266,42 @@ class GlobalState:
             Information from the object table.
         """
         self._check_connected()
-        if object_id is not None:
-            # Return information about a single object ID.
-            return self._object_table(object_id)
-        else:
-            # Return the entire object table.
-            object_keys = self._keys(gcs_utils.TablePrefix_OBJECT_string + "*")
-            object_ids_binary = {
-                key[len(gcs_utils.TablePrefix_OBJECT_string):]
-                for key in object_keys
-            }
 
+        if object_id is not None:
+            object_id = ray.ObjectID(hex_to_binary(object_id))
+            object_info = self.global_state_accessor.get_object_info(object_id)
+            if object_info is None:
+                return {}
+            else:
+                object_location_info = gcs_utils.ObjectLocationInfo.FromString(
+                    object_info)
+                return self._gen_object_info(object_location_info)
+        else:
+            object_table = self.global_state_accessor.get_object_table()
             results = {}
-            for object_id_binary in object_ids_binary:
-                results[binary_to_object_id(object_id_binary)] = (
-                    self._object_table(binary_to_object_id(object_id_binary)))
+            for i in range(len(object_table)):
+                object_location_info = gcs_utils.ObjectLocationInfo.FromString(
+                    object_table[i])
+                results[binary_to_hex(object_location_info.object_id)] = \
+                    self._gen_object_info(object_location_info)
             return results
 
+    def _gen_object_info(self, object_location_info):
+        """Parse object location info.
+        Returns:
+            Information from object.
+        """
+        locations = []
+        for location in object_location_info.locations:
+            locations.append(ray.utils.binary_to_hex(location.manager))
+
+        object_info = {
+            "ObjectID": ray.utils.binary_to_hex(
+                object_location_info.object_id),
+            "Locations": locations,
+        }
+        return object_info
+
     def _actor_table(self, actor_id):
         """Fetch and parse the actor table information for a single actor ID.
 

python/ray/tests/test_advanced_3.py

@@ -511,26 +511,6 @@ def test_put_pins_object(ray_start_object_store_memory):
         ray.get(y_id)
 
 
-@pytest.mark.parametrize(
-    "ray_start_object_store_memory", [150 * 1024 * 1024], indirect=True)
-def test_redis_lru_with_set(ray_start_object_store_memory):
-    x = np.zeros(8 * 10**7, dtype=np.uint8)
-    x_id = ray.put(x, weakref=True)
-
-    # Remove the object from the object table to simulate Redis LRU eviction.
-    removed = False
-    start_time = time.time()
-    while time.time() < start_time + 10:
-        if ray.state.state.redis_clients[0].delete(b"OBJECT" +
-                                                   x_id.binary()) == 1:
-            removed = True
-            break
-    assert removed
-
-    # Now evict the object from the object store.
-    ray.put(x)  # This should not crash.
-
-
 def test_decorated_function(ray_start_regular):
     def function_invocation_decorator(f):
         def new_f(args, kwargs):