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

src/ray/raylet/format/node_manager.fbs

@@ -207,6 +207,8 @@ table FreeObjectsRequest {
   // Whether keep this request with local object store
   // or send it to all the object stores.
   local_only: bool;
+  // Whether also delete objects' creating tasks from GCS.
+  delete_creating_tasks: bool;
   // List of object ids we'll delete from object store.
   object_ids: [string];
 }

src/ray/raylet/lib/java/org_ray_runtime_raylet_RayletClientImpl.cc

@@ -247,11 +247,12 @@ Java_org_ray_runtime_raylet_RayletClientImpl_nativeGenerateTaskId(
 /*
  * Class:     org_ray_runtime_raylet_RayletClientImpl
  * Method:    nativeFreePlasmaObjects
- * Signature: ([[BZ)V
+ * Signature: (J[[BZZ)V
  */
 JNIEXPORT void JNICALL
 Java_org_ray_runtime_raylet_RayletClientImpl_nativeFreePlasmaObjects(
-    JNIEnv *env, jclass, jlong client, jobjectArray objectIds, jboolean localOnly) {
+    JNIEnv *env, jclass, jlong client, jobjectArray objectIds, jboolean localOnly,
+    jboolean deleteCreatingTasks) {
   std::vector<ObjectID> object_ids;
   auto len = env->GetArrayLength(objectIds);
   for (int i = 0; i < len; i++) {
@@ -262,7 +263,7 @@ Java_org_ray_runtime_raylet_RayletClientImpl_nativeFreePlasmaObjects(
     env->DeleteLocalRef(object_id_bytes);
   }
   auto raylet_client = reinterpret_cast<RayletClient *>(client);
-  auto status = raylet_client->FreeObjects(object_ids, localOnly);
+  auto status = raylet_client->FreeObjects(object_ids, localOnly, deleteCreatingTasks);
   ThrowRayExceptionIfNotOK(env, status);
 }
 

src/ray/raylet/lib/java/org_ray_runtime_raylet_RayletClientImpl.h

@@ -91,12 +91,12 @@ Java_org_ray_runtime_raylet_RayletClientImpl_nativeGenerateTaskId(JNIEnv *, jcla
 /*
  * Class:     org_ray_runtime_raylet_RayletClientImpl
  * Method:    nativeFreePlasmaObjects
- * Signature: (J[[BZ)V
+ * Signature: (J[[BZZ)V
  */
 JNIEXPORT void JNICALL
 Java_org_ray_runtime_raylet_RayletClientImpl_nativeFreePlasmaObjects(JNIEnv *, jclass,
                                                                      jlong, jobjectArray,
-                                                                     jboolean);
+                                                                     jboolean, jboolean);
 
 /*
  * Class:     org_ray_runtime_raylet_RayletClientImpl

src/ray/raylet/node_manager.cc

@@ -711,7 +711,16 @@ void NodeManager::ProcessClientMessage(
   case protocol::MessageType::FreeObjectsInObjectStoreRequest: {
     auto message = flatbuffers::GetRoot<protocol::FreeObjectsRequest>(message_data);
     std::vector<ObjectID> object_ids = from_flatbuf<ObjectID>(*message->object_ids());
+    // Clean up objects from the object store.
     object_manager_.FreeObjects(object_ids, message->local_only());
+    if (message->delete_creating_tasks()) {
+      // Clean up their creating tasks from GCS.
+      std::vector<TaskID> creating_task_ids;
+      for (const auto &object_id : object_ids) {
+        creating_task_ids.push_back(ComputeTaskId(object_id));
+      }
+      gcs_client_->raylet_task_table().Delete(JobID::nil(), creating_task_ids);
+    }
   } break;
   case protocol::MessageType::PrepareActorCheckpointRequest: {
     ProcessPrepareActorCheckpointRequest(client, message_data);

src/ray/raylet/raylet_client.cc

@@ -349,10 +349,10 @@ ray::Status RayletClient::PushProfileEvents(const ProfileTableDataT &profile_eve
 }
 
 ray::Status RayletClient::FreeObjects(const std::vector<ray::ObjectID> &object_ids,
-                                      bool local_only) {
+                                      bool local_only, bool delete_creating_tasks) {
   flatbuffers::FlatBufferBuilder fbb;
-  auto message = ray::protocol::CreateFreeObjectsRequest(fbb, local_only,
-                                                         to_flatbuf(fbb, object_ids));
+  auto message = ray::protocol::CreateFreeObjectsRequest(
+      fbb, local_only, delete_creating_tasks, to_flatbuf(fbb, object_ids));
   fbb.Finish(message);
 
   auto status = conn_->WriteMessage(MessageType::FreeObjectsInObjectStoreRequest, &fbb);

src/ray/raylet/raylet_client.h

@@ -145,8 +145,10 @@ class RayletClient {
   /// \param object_ids A list of ObjectsIDs to be deleted.
   /// \param local_only Whether keep this request with local object store
   /// or send it to all the object stores.
+  /// \param delete_creating_tasks Whether also delete objects' creating tasks from GCS.
   /// \return ray::Status.
-  ray::Status FreeObjects(const std::vector<ray::ObjectID> &object_ids, bool local_only);
+  ray::Status FreeObjects(const std::vector<ray::ObjectID> &object_ids, bool local_only,
+                          bool deleteCreatingTasks);
 
   /// Request raylet backend to prepare a checkpoint for an actor.
   ///