Add delete_creating_tasks option for internal.free() (#4588)

* add delete creating task objects.

* format code style

* Fix lint

* add tests add address comments.

* Refine test

* Refine java test

* Fix CI

* Refine

* Fix lint

* Fix CI

python/ray/_raylet.pyx

@@ -349,9 +349,9 @@ cdef class RayletClient:
 
         check_status(self.client.get().PushProfileEvents(profile_info))
 
-    def free_objects(self, object_ids, c_bool local_only):
+    def free_objects(self, object_ids, c_bool local_only, c_bool delete_creating_tasks):
         cdef c_vector[CObjectID] free_ids = ObjectIDsToVector(object_ids)
-        check_status(self.client.get().FreeObjects(free_ids, local_only))
+        check_status(self.client.get().FreeObjects(free_ids, local_only, delete_creating_tasks))
 
     def prepare_actor_checkpoint(self, ActorID actor_id):
         cdef CActorCheckpointID checkpoint_id

python/ray/includes/libraylet.pxd

@@ -67,7 +67,7 @@ cdef extern from "ray/raylet/raylet_client.h" nogil:
         CRayStatus PushProfileEvents(
             const GCSProfileTableDataT &profile_events)
         CRayStatus FreeObjects(const c_vector[CObjectID] &object_ids,
-                               c_bool local_only)
+                               c_bool local_only, c_bool delete_creating_tasks)
         CRayStatus PrepareActorCheckpoint(const CActorID &actor_id,
                                           CActorCheckpointID &checkpoint_id)
         CRayStatus NotifyActorResumedFromCheckpoint(

python/ray/internal/internal_api.py

@@ -8,7 +8,7 @@ from ray import profiling
 __all__ = ["free"]
 
 
-def free(object_ids, local_only=False):
+def free(object_ids, local_only=False, delete_creating_tasks=False):
     """Free a list of IDs from object stores.
 
     This function is a low-level API which should be used in restricted
@@ -25,6 +25,8 @@ def free(object_ids, local_only=False):
         object_ids (List[ObjectID]): List of object IDs to delete.
         local_only (bool): Whether only deleting the list of objects in local
             object store or all object stores.
+        delete_creating_tasks (bool): Whether also delete the object creating
+            tasks.
     """
     worker = ray.worker.get_global_worker()
 
@@ -46,4 +48,5 @@ def free(object_ids, local_only=False):
         if len(object_ids) == 0:
             return
 
-        worker.raylet_client.free_objects(object_ids, local_only)
+        worker.raylet_client.free_objects(object_ids, local_only,
+                                          delete_creating_tasks)

python/ray/tests/test_basic.py

@@ -1439,13 +1439,16 @@ def test_free_objects_multi_node(ray_start_cluster):
         assert len(l2) == 0
         return (a, b, c)
 
-    def run_one_test(actors, local_only):
+    def run_one_test(actors, local_only, delete_creating_tasks):
         (a, b, c) = create(actors)
         # The three objects should be generated on different object stores.
         assert ray.get(a) != ray.get(b)
         assert ray.get(a) != ray.get(c)
         assert ray.get(c) != ray.get(b)
-        ray.internal.free([a, b, c], local_only=local_only)
+        ray.internal.free(
+            [a, b, c],
+            local_only=local_only,
+            delete_creating_tasks=delete_creating_tasks)
         # Wait for the objects to be deleted.
         time.sleep(0.1)
         return (a, b, c)
@@ -1456,13 +1459,13 @@ def test_free_objects_multi_node(ray_start_cluster):
         ActorOnNode2.remote()
     ]
     # Case 1: run this local_only=False. All 3 objects will be deleted.
-    (a, b, c) = run_one_test(actors, False)
+    (a, b, c) = run_one_test(actors, False, False)
     (l1, l2) = ray.wait([a, b, c], timeout=0.01, num_returns=1)
     # All the objects are deleted.
     assert len(l1) == 0
     assert len(l2) == 3
     # Case 2: run this local_only=True. Only 1 object will be deleted.
-    (a, b, c) = run_one_test(actors, True)
+    (a, b, c) = run_one_test(actors, True, False)
     (l1, l2) = ray.wait([a, b, c], timeout=0.01, num_returns=3)
     # One object is deleted and 2 objects are not.
     assert len(l1) == 2
@@ -1472,6 +1475,17 @@ def test_free_objects_multi_node(ray_start_cluster):
     for object_id in l1:
         assert ray.get(object_id) != local_return
 
+    # Case3: These cases test the deleting creating tasks for the object.
+    (a, b, c) = run_one_test(actors, False, False)
+    task_table = ray.global_state.task_table()
+    for obj in [a, b, c]:
+        assert ray._raylet.compute_task_id(obj).hex() in task_table
+
+    (a, b, c) = run_one_test(actors, False, True)
+    task_table = ray.global_state.task_table()
+    for obj in [a, b, c]:
+        assert ray._raylet.compute_task_id(obj).hex() not in task_table
+
 
 def test_local_mode(shutdown_only):
     @ray.remote