Basic reference counting & pinning (#6554)

python/ray/tests/test_reference_counting.py

@@ -4,10 +4,12 @@ from __future__ import division
 from __future__ import print_function
 
 import os
+import json
 import copy
 import tempfile
 import numpy as np
 import time
+import pytest
 import logging
 import uuid
 
@@ -27,7 +29,7 @@ def _check_refcounts(expected):
         assert submitted == actual[object_id]["submitted"]
 
 
-def check_refcounts(expected, timeout=1):
+def check_refcounts(expected, timeout=10):
     start = time.time()
     while True:
         try:
@@ -156,7 +158,99 @@ def test_dependency_refcounts(ray_start_regular):
     check_refcounts({})
 
 
+def test_basic_pinning(shutdown_only):
+    ray.init(object_store_memory=100 * 1024 * 1024)
+
+    @ray.remote
+    def f(array):
+        return np.sum(array)
+
+    @ray.remote
+    class Actor(object):
+        def __init__(self):
+            # Hold a long-lived reference to a ray.put object's ID. The object
+            # should not be garbage collected while the actor is alive because
+            # the object is pinned by the raylet.
+            self.large_object = ray.put(
+                np.zeros(25 * 1024 * 1024, dtype=np.uint8))
+
+        def get_large_object(self):
+            return ray.get(self.large_object)
+
+    actor = Actor.remote()
+
+    # Fill up the object store with short-lived objects. These should be
+    # evicted before the long-lived object whose reference is held by
+    # the actor.
+    for batch in range(10):
+        intermediate_result = f.remote(
+            np.zeros(10 * 1024 * 1024, dtype=np.uint8))
+        ray.get(intermediate_result)
+
+    # The ray.get below would fail with only LRU eviction, as the object
+    # that was ray.put by the actor would have been evicted.
+    ray.get(actor.get_large_object.remote())
+
+
+def test_pending_task_dependency_pinning(shutdown_only):
+    ray.init(object_store_memory=100 * 1024 * 1024, use_pickle=True)
+
+    @ray.remote
+    def pending(input1, input2):
+        return
+
+    @ray.remote
+    def slow(dep):
+        pass
+
+    # The object that is ray.put here will go out of scope immediately, so if
+    # pending task dependencies aren't considered, it will be evicted before
+    # the ray.get below due to the subsequent ray.puts that fill up the object
+    # store.
+    np_array = np.zeros(40 * 1024 * 1024, dtype=np.uint8)
+    random_id = ray.ObjectID.from_random()
+    oid = pending.remote(np_array, slow.remote(random_id))
+
+    for _ in range(2):
+        ray.put(np_array)
+
+    ray.worker.global_worker.put_object(None, object_id=random_id)
+    ray.get(oid)
+
+
+def test_feature_flag(shutdown_only):
+    ray.init(
+        object_store_memory=100 * 1024 * 1024,
+        _internal_config=json.dumps({
+            "object_pinning_enabled": 0
+        }))
+
+    @ray.remote
+    def f(array):
+        return np.sum(array)
+
+    @ray.remote
+    class Actor(object):
+        def __init__(self):
+            self.large_object = ray.put(
+                np.zeros(25 * 1024 * 1024, dtype=np.uint8))
+
+        def get_large_object(self):
+            return ray.get(self.large_object)
+
+    actor = Actor.remote()
+
+    for batch in range(10):
+        intermediate_result = f.remote(
+            np.zeros(10 * 1024 * 1024, dtype=np.uint8))
+        ray.get(intermediate_result)
+
+    # The ray.get below fails with only LRU eviction, as the object
+    # that was ray.put by the actor should have been evicted.
+    with pytest.raises(ray.exceptions.RayTimeoutError):
+        ray.get(actor.get_large_object.remote(), timeout=1)
+
+
 if __name__ == "__main__":
-    import pytest
     import sys
     sys.exit(pytest.main(["-v", __file__]))