mirror of
https://github.com/wassname/ray.git
synced 2026-07-18 12:40:56 +08:00
Distributed ref counting for serialized ObjectIDs (#6945)
* Skeleton plus a unit test for simple borrower case * First unit test passes - forward an ID and task returns with 1 submitted task pending on the inner ID * Invariant for contained_in * Unit test passes for testing task return without creating a borrower * Wrap ref count functionality in test case * Fix bad delete * Unit test and fix for borrowers creating more borrowers * Unit test and fix for simple borrowing, but owner sends call after borrower's ref count goes to 0 * Refactor: - keep a sentinel ref count for task argument IDs - keep contained_in_borrowed in addition to contained_in_owned * Unit test for nested IDs passes * Refactor so that an object ID can only be contained in 1 borrowed ID at a time * Add check * Fix * Unit test (passes) to test nesting object IDs but no borrowers created * Unit test for nested objects from different owners passes, refactor to unset contained_in when popping refs * Unit tests for borrowers receiving an ObjectID from multiple sources, skip adding ownership info if we already have it to handle duplicate refs * Unit test for returning object ID passes * More unit tests for returning object IDs pass * Add serialized ID tests * fix serialization issue * remove swap * It builds! * debugging and some fixes: - register handler for WaitForRefRemoved - don't create a python reference for arg IDs - pass in client factory into ReferenceCounter - fix bad decrement in PopBorrowerRefs * Fix accounting for serialized IDs: - don't decrement for IDs on dependency resolution, wait until task finished - add object IDs that were inlined when building the arguments to the task spec, pin these on the task executor until task finishes * mu_ -> mutex_ * lint * fix build * clear outer_object_id * add direct call type check * Fix test for direct call IDs and return IDs for actor calls * Fix CoreWorkerClient.Addr() * Remove unneeded lock * Remove unnecessary ObjectID refs * Fix worker holding serialized refs test * Fix hex IDs * fix * fix tests * fix tests * refactor and cleanups * lint * Put inlined Ids in task args and some cleanup * Add back gc.collect() line for test case * Refactor and fixes: - store inlined IDs in RayObject - allow storing objects with inlined IDs in memory store - pin objects that were promoted to plasma * oops * make sure worker ID is set in address, pass in rpc::Address to CoreWorkerClient * todos * cleanups and test builds * Fix tests * Add feature flag * cleanups * address comments and some cleanups * cleanup * fix recursive test * Comments for tests * Turn off ref counting by default * Skip tests * Fix some bugs for test_array.py, java build * Don't include nested objects in the ref count when the feature flag is off * C++ feature flag does not work... * Remove * Turn on python tests and add a warning when plasma objects are evicted before being pinned * Fix build and remove irrelevant test * Fix for java * Revert "Fix build and remove irrelevant test" This reverts commit 056cca9b263ed05b0f9ab2250907338edcbca2d5. * Fix ray.internal.free * Fixes and skip some flaky tests * fix java build * fix windows build * Add IDs contained in owned objects * Update src/ray/protobuf/core_worker.proto Co-Authored-By: Edward Oakes <ed.nmi.oakes@gmail.com> * Update src/ray/core_worker/reference_count.cc Co-Authored-By: Edward Oakes <ed.nmi.oakes@gmail.com> * Update src/ray/protobuf/core_worker.proto Co-Authored-By: Edward Oakes <ed.nmi.oakes@gmail.com> * Update src/ray/protobuf/core_worker.proto Co-Authored-By: Edward Oakes <ed.nmi.oakes@gmail.com> * Update src/ray/core_worker/reference_count.h Co-Authored-By: Edward Oakes <ed.nmi.oakes@gmail.com> * Update src/ray/core_worker/reference_count.h Co-Authored-By: Edward Oakes <ed.nmi.oakes@gmail.com> * Update src/ray/core_worker/reference_count.cc Co-Authored-By: Edward Oakes <ed.nmi.oakes@gmail.com> * Apply suggestions from code review Co-Authored-By: Edward Oakes <ed.nmi.oakes@gmail.com> * update * Try to fix ::test_direct_call_serialized_id_eviction Co-authored-by: Edward Oakes <ed.nmi.oakes@gmail.com>
This commit is contained in:
Stephanie Wang
co-authored by
Edward Oakes
Edward Oakes
parent
4a12243336
commit
f76ce836b2
@@ -8,6 +8,7 @@ import time
|
||||
import pytest
|
||||
import logging
|
||||
import uuid
|
||||
import gc
|
||||
|
||||
import ray
|
||||
import ray.cluster_utils
|
||||
@@ -18,7 +19,13 @@ logger = logging.getLogger(__name__)
|
||||
|
||||
@pytest.fixture
|
||||
def one_worker_100MiB(request):
|
||||
yield ray.init(num_cpus=1, object_store_memory=100 * 1024 * 1024)
|
||||
config = json.dumps({
|
||||
"distributed_ref_counting_enabled": 1,
|
||||
})
|
||||
yield ray.init(
|
||||
num_cpus=1,
|
||||
object_store_memory=100 * 1024 * 1024,
|
||||
_internal_config=config)
|
||||
ray.shutdown()
|
||||
|
||||
|
||||
@@ -266,7 +273,6 @@ def test_feature_flag(shutdown_only):
|
||||
# Remote function takes serialized reference and doesn't hold onto it after
|
||||
# finishing. Referenced object shouldn't be evicted while the task is pending
|
||||
# and should be evicted after it returns.
|
||||
@pytest.mark.skip("Serialized ObjectID reference counting not implemented.")
|
||||
def test_basic_serialized_reference(one_worker_100MiB):
|
||||
@ray.remote
|
||||
def pending(ref, dep):
|
||||
@@ -286,6 +292,9 @@ def test_basic_serialized_reference(one_worker_100MiB):
|
||||
# Remove the local reference.
|
||||
array_oid_bytes = array_oid.binary()
|
||||
del array_oid
|
||||
# Needed due to Python GC issue in cloudpickle.
|
||||
# https://github.com/cloudpipe/cloudpickle/issues/343
|
||||
gc.collect()
|
||||
|
||||
# Check that the remote reference pins the object.
|
||||
_fill_object_store_and_get(array_oid_bytes)
|
||||
@@ -301,7 +310,8 @@ def test_basic_serialized_reference(one_worker_100MiB):
|
||||
# Call a recursive chain of tasks that pass a serialized reference to the end
|
||||
# of the chain. The reference should still exist while the final task in the
|
||||
# chain is running and should be removed once it finishes.
|
||||
@pytest.mark.skip("Serialized ObjectID reference counting not implemented.")
|
||||
@pytest.mark.skip("Memory not freed due to Python GC issue in cloudpickle "
|
||||
"(https://github.com/cloudpipe/cloudpickle/issues/343).")
|
||||
def test_recursive_serialized_reference(one_worker_100MiB):
|
||||
@ray.remote
|
||||
def recursive(ref, dep, max_depth, depth=0):
|
||||
@@ -325,7 +335,7 @@ def test_recursive_serialized_reference(one_worker_100MiB):
|
||||
del array_oid
|
||||
|
||||
tail_oid = head_oid
|
||||
for _ in range(max_depth - 1):
|
||||
for _ in range(max_depth):
|
||||
tail_oid = ray.get(tail_oid)
|
||||
|
||||
# Check that the remote reference pins the object.
|
||||
@@ -333,7 +343,7 @@ def test_recursive_serialized_reference(one_worker_100MiB):
|
||||
|
||||
# Fulfill the dependency, causing the tail task to finish.
|
||||
ray.worker.global_worker.put_object(None, object_id=random_oid)
|
||||
ray.get(tail_oid)
|
||||
assert ray.get(tail_oid) is None
|
||||
|
||||
# Reference should be gone, check that array gets evicted.
|
||||
_fill_object_store_and_get(array_oid_bytes, succeed=False)
|
||||
@@ -342,7 +352,6 @@ def test_recursive_serialized_reference(one_worker_100MiB):
|
||||
# Test that a passed reference held by an actor after the method finishes
|
||||
# is kept until the reference is removed from the actor. Also tests giving
|
||||
# the actor a duplicate reference to the same object ID.
|
||||
@pytest.mark.skip("Serialized ObjectID reference counting not implemented.")
|
||||
def test_actor_holding_serialized_reference(one_worker_100MiB):
|
||||
@ray.remote
|
||||
class GreedyActor(object):
|
||||
@@ -376,6 +385,9 @@ def test_actor_holding_serialized_reference(one_worker_100MiB):
|
||||
# Remove the local reference.
|
||||
array_oid_bytes = array_oid.binary()
|
||||
del array_oid
|
||||
# Needed due to Python GC issue in cloudpickle.
|
||||
# https://github.com/cloudpipe/cloudpickle/issues/343
|
||||
gc.collect()
|
||||
|
||||
# Test that the remote references still pin the object.
|
||||
_fill_object_store_and_get(array_oid_bytes)
|
||||
@@ -392,7 +404,8 @@ def test_actor_holding_serialized_reference(one_worker_100MiB):
|
||||
# Test that a passed reference held by an actor after a task finishes
|
||||
# is kept until the reference is removed from the worker. Also tests giving
|
||||
# the worker a duplicate reference to the same object ID.
|
||||
@pytest.mark.skip("Serialized ObjectID reference counting not implemented.")
|
||||
@pytest.mark.skip("Memory not freed due to Python GC issue in cloudpickle "
|
||||
"(https://github.com/cloudpipe/cloudpickle/issues/343).")
|
||||
def test_worker_holding_serialized_reference(one_worker_100MiB):
|
||||
@ray.remote
|
||||
def child(dep1, dep2):
|
||||
@@ -428,7 +441,6 @@ def test_worker_holding_serialized_reference(one_worker_100MiB):
|
||||
|
||||
|
||||
# Test that an object containing object IDs within it pins the inner IDs.
|
||||
@pytest.mark.skip("Serialized ObjectID reference counting not implemented.")
|
||||
def test_basic_nested_ids(one_worker_100MiB):
|
||||
inner_oid = ray.put(np.zeros(40 * 1024 * 1024, dtype=np.uint8))
|
||||
outer_oid = ray.put([inner_oid])
|
||||
@@ -436,6 +448,9 @@ def test_basic_nested_ids(one_worker_100MiB):
|
||||
# Remove the local reference to the inner object.
|
||||
inner_oid_bytes = inner_oid.binary()
|
||||
del inner_oid
|
||||
# Needed due to Python GC issue in cloudpickle.
|
||||
# https://github.com/cloudpipe/cloudpickle/issues/343
|
||||
gc.collect()
|
||||
|
||||
# Check that the outer reference pins the inner object.
|
||||
_fill_object_store_and_get(inner_oid_bytes)
|
||||
@@ -447,7 +462,8 @@ def test_basic_nested_ids(one_worker_100MiB):
|
||||
|
||||
# Test that an object containing object IDs within it pins the inner IDs
|
||||
# recursively and for submitted tasks.
|
||||
@pytest.mark.skip("Serialized ObjectID reference counting not implemented.")
|
||||
@pytest.mark.skip("Memory not freed due to Python GC issue in cloudpickle "
|
||||
"(https://github.com/cloudpipe/cloudpickle/issues/343).")
|
||||
def test_recursively_nest_ids(one_worker_100MiB):
|
||||
@ray.remote
|
||||
def recursive(ref, dep, max_depth, depth=0):
|
||||
@@ -474,7 +490,7 @@ def test_recursively_nest_ids(one_worker_100MiB):
|
||||
del array_oid, nested_oid
|
||||
|
||||
tail_oid = head_oid
|
||||
for _ in range(max_depth - 1):
|
||||
for _ in range(max_depth):
|
||||
tail_oid = ray.get(tail_oid)
|
||||
|
||||
# Check that the remote reference pins the object.
|
||||
@@ -490,7 +506,8 @@ def test_recursively_nest_ids(one_worker_100MiB):
|
||||
|
||||
# Test that serialized objectIDs returned from remote tasks are pinned until
|
||||
# they go out of scope on the caller side.
|
||||
@pytest.mark.skip("Serialized ObjectID reference counting not implemented.")
|
||||
@pytest.mark.skip("Memory not freed due to Python GC issue in cloudpickle "
|
||||
"(https://github.com/cloudpipe/cloudpickle/issues/343).")
|
||||
def test_return_object_id(one_worker_100MiB):
|
||||
@ray.remote
|
||||
def put():
|
||||
@@ -519,7 +536,8 @@ def test_return_object_id(one_worker_100MiB):
|
||||
|
||||
# Test that serialized objectIDs returned from remote tasks are pinned if
|
||||
# passed into another remote task by the caller.
|
||||
@pytest.mark.skip("Serialized ObjectID reference counting not implemented.")
|
||||
@pytest.mark.skip("Memory not freed due to Python GC issue in cloudpickle "
|
||||
"(https://github.com/cloudpipe/cloudpickle/issues/343).")
|
||||
def test_pass_returned_object_id(one_worker_100MiB):
|
||||
@ray.remote
|
||||
def put():
|
||||
@@ -555,7 +573,8 @@ def test_pass_returned_object_id(one_worker_100MiB):
|
||||
# returned by another task to the end of the chain. The reference should still
|
||||
# exist while the final task in the chain is running and should be removed once
|
||||
# it finishes.
|
||||
@pytest.mark.skip("Serialized ObjectID reference counting not implemented.")
|
||||
@pytest.mark.skip("Memory not freed due to Python GC issue in cloudpickle "
|
||||
"(https://github.com/cloudpipe/cloudpickle/issues/343).")
|
||||
def test_recursively_pass_returned_object_id(one_worker_100MiB):
|
||||
@ray.remote
|
||||
def put():
|
||||
@@ -583,7 +602,7 @@ def test_recursively_pass_returned_object_id(one_worker_100MiB):
|
||||
del outer_oid
|
||||
|
||||
tail_oid = head_oid
|
||||
for _ in range(max_depth - 1):
|
||||
for _ in range(max_depth):
|
||||
tail_oid = ray.get(tail_oid)
|
||||
|
||||
# Check that the remote reference pins the object.
|
||||
|
||||
Reference in New Issue
Block a user