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

java/api/src/main/java/org/ray/api/runtime/RayRuntime.java

@@ -61,8 +61,9 @@ public interface RayRuntime {
    *
    * @param objectIds The object ids to free.
    * @param localOnly Whether only free objects for local object store or not.
+   * @param deleteCreatingTasks Whether also delete objects' creating tasks from GCS.
    */
-  void free(List<UniqueId> objectIds, boolean localOnly);
+  void free(List<UniqueId> objectIds, boolean localOnly, boolean deleteCreatingTasks);
 
   /**
    * Invoke a remote function.

java/runtime/src/main/java/org/ray/runtime/AbstractRayRuntime.java

@@ -205,8 +205,8 @@ public abstract class AbstractRayRuntime implements RayRuntime {
   }
 
   @Override
-  public void free(List<UniqueId> objectIds, boolean localOnly) {
-    rayletClient.freePlasmaObjects(objectIds, localOnly);
+  public void free(List<UniqueId> objectIds, boolean localOnly, boolean deleteCreatingTasks) {
+    rayletClient.freePlasmaObjects(objectIds, localOnly, deleteCreatingTasks);
   }
 
   private List<List<UniqueId>> splitIntoBatches(List<UniqueId> objectIds) {

java/runtime/src/main/java/org/ray/runtime/RayNativeRuntime.java

@@ -213,4 +213,22 @@ public final class RayNativeRuntime extends AbstractRayRuntime {
 
     return false;
   }
+
+  /**
+   * Query whether the raylet task exists in Gcs.
+   */
+  public boolean rayletTaskExistsInGcs(UniqueId taskId) {
+    byte[] key = ArrayUtils.addAll("RAYLET_TASK".getBytes(), taskId.getBytes());
+
+    // TODO(qwang): refactor this with `GlobalState` after this issue
+    // getting finished. https://github.com/ray-project/ray/issues/3933
+    for (RedisClient client : redisClients) {
+      if (client.exists(key)) {
+        return true;
+      }
+    }
+
+    return false;
+  }
+
 }

java/runtime/src/main/java/org/ray/runtime/raylet/MockRayletClient.java

@@ -191,7 +191,8 @@ public class MockRayletClient implements RayletClient {
   }
 
   @Override
-  public void freePlasmaObjects(List<UniqueId> objectIds, boolean localOnly) {
+  public void freePlasmaObjects(List<UniqueId> objectIds, boolean localOnly,
+                                boolean deleteCreatingTasks) {
     for (UniqueId id : objectIds) {
       store.free(id);
     }

java/runtime/src/main/java/org/ray/runtime/raylet/RayletClient.java

@@ -24,7 +24,7 @@ public interface RayletClient {
   <T> WaitResult<T> wait(List<RayObject<T>> waitFor, int numReturns, int
       timeoutMs, UniqueId currentTaskId);
 
-  void freePlasmaObjects(List<UniqueId> objectIds, boolean localOnly);
+  void freePlasmaObjects(List<UniqueId> objectIds, boolean localOnly, boolean deleteCreatingTasks);
 
   UniqueId prepareCheckpoint(UniqueId actorId);
 

java/runtime/src/main/java/org/ray/runtime/raylet/RayletClientImpl.java

@@ -123,9 +123,10 @@ public class RayletClientImpl implements RayletClient {
   }
 
   @Override
-  public void freePlasmaObjects(List<UniqueId> objectIds, boolean localOnly) {
+  public void freePlasmaObjects(List<UniqueId> objectIds, boolean localOnly,
+                                boolean deleteCreatingTasks) {
     byte[][] objectIdsArray = UniqueIdUtil.getIdBytes(objectIds);
-    nativeFreePlasmaObjects(client, objectIdsArray, localOnly);
+    nativeFreePlasmaObjects(client, objectIdsArray, localOnly, deleteCreatingTasks);
   }
 
   @Override
@@ -350,7 +351,7 @@ public class RayletClientImpl implements RayletClient {
       int taskIndex);
 
   private static native void nativeFreePlasmaObjects(long conn, byte[][] objectIds,
-      boolean localOnly) throws RayException;
+      boolean localOnly, boolean deleteCreatingTasks) throws RayException;
 
   private static native byte[] nativePrepareCheckpoint(long conn, byte[] actorId);
 

java/test/src/main/java/org/ray/api/TestUtils.java

@@ -1,15 +1,45 @@
 package org.ray.api;
 
+import java.util.function.Supplier;
 import org.ray.runtime.AbstractRayRuntime;
 import org.ray.runtime.config.RunMode;
 import org.testng.SkipException;
 
 public class TestUtils {
 
+  private static final int WAIT_INTERVAL_MS = 5;
+
   public static void skipTestUnderSingleProcess() {
     AbstractRayRuntime runtime = (AbstractRayRuntime)Ray.internal();
     if (runtime.getRayConfig().runMode == RunMode.SINGLE_PROCESS) {
       throw new SkipException("This test doesn't work under single-process mode.");
     }
   }
+
+  /**
+   * Wait until the given condition is met.
+   *
+   * @param condition A function that predicts the condition.
+   * @param timeoutMs Timeout in milliseconds.
+   * @return True if the condition is met within the timeout, false otherwise.
+   */
+  public static boolean waitForCondition(Supplier<Boolean> condition, int timeoutMs) {
+    int waitTime = 0;
+    while (true) {
+      if (condition.get()) {
+        return true;
+      }
+
+      try {
+        java.util.concurrent.TimeUnit.MILLISECONDS.sleep(WAIT_INTERVAL_MS);
+      } catch (InterruptedException e) {
+        throw new RuntimeException(e);
+      }
+      waitTime += WAIT_INTERVAL_MS;
+      if (waitTime > timeoutMs) {
+        break;
+      }
+    }
+    return false;
+  }
 }

java/test/src/main/java/org/ray/api/test/ActorTest.java

@@ -97,7 +97,7 @@ public class ActorTest extends BaseTest {
     RayObject value = Ray.call(Counter::getValue, counter);
     Assert.assertEquals(100, value.get());
     // Delete the object from the object store.
-    Ray.internal().free(ImmutableList.of(value.getId()), false);
+    Ray.internal().free(ImmutableList.of(value.getId()), false, false);
     // Wait until the object is deleted, because the above free operation is async.
     while (true) {
       GetResult<Integer> result = ((AbstractRayRuntime)

java/test/src/main/java/org/ray/api/test/PlasmaFreeTest.java

@@ -1,17 +1,16 @@
 package org.ray.api.test;
 
 import com.google.common.collect.ImmutableList;
-import java.util.ArrayList;
-import java.util.List;
 import org.ray.api.Ray;
 import org.ray.api.RayObject;
-import org.ray.api.WaitResult;
+import org.ray.api.TestUtils;
 import org.ray.api.annotation.RayRemote;
-import org.ray.api.id.UniqueId;
+import org.ray.runtime.AbstractRayRuntime;
+import org.ray.runtime.RayNativeRuntime;
+import org.ray.runtime.util.UniqueIdUtil;
 import org.testng.Assert;
 import org.testng.annotations.Test;
 
-
 public class PlasmaFreeTest extends BaseTest {
 
   @RayRemote
@@ -20,31 +19,27 @@ public class PlasmaFreeTest extends BaseTest {
   }
 
   @Test
-  public void test() {
+  public void testDeleteObjects() {
     RayObject<String> helloId = Ray.call(PlasmaFreeTest::hello);
     String helloString = helloId.get();
     Assert.assertEquals("hello", helloString);
-    List<RayObject<String>> waitFor = ImmutableList.of(helloId);
-    WaitResult<String> waitResult = Ray.wait(waitFor, 1, 2 * 1000);
-    List<RayObject<String>> readyOnes = waitResult.getReady();
-    List<RayObject<String>> unreadyOnes = waitResult.getUnready();
-    Assert.assertEquals(1, readyOnes.size());
-    Assert.assertEquals(0, unreadyOnes.size());
+    Ray.internal().free(ImmutableList.of(helloId.getId()), true, false);
 
-    List<UniqueId> freeList = new ArrayList<>();
-    freeList.add(helloId.getId());
-    Ray.internal().free(freeList, true);
-    // Flush: trigger the release function because Plasma Client has cache.
-    for (int i = 0; i < 128; i++) {
-      Ray.call(PlasmaFreeTest::hello).get();
-    }
-
-    // Check if the object has been evicted. Don't give ray.wait enough
-    // time to reconstruct the object.
-    waitResult = Ray.wait(waitFor, 1, 0);
-    readyOnes = waitResult.getReady();
-    unreadyOnes = waitResult.getUnready();
-    Assert.assertEquals(0, readyOnes.size());
-    Assert.assertEquals(1, unreadyOnes.size());
+    final boolean result = TestUtils.waitForCondition(() -> !((AbstractRayRuntime) Ray.internal())
+        .getObjectStoreProxy().get(helloId.getId(), 0).exists, 50);
+    Assert.assertTrue(result);
   }
+
+  @Test
+  public void testDeleteCreatingTasks() {
+    TestUtils.skipTestUnderSingleProcess();
+    RayObject<String> helloId = Ray.call(PlasmaFreeTest::hello);
+    Assert.assertEquals("hello", helloId.get());
+    Ray.internal().free(ImmutableList.of(helloId.getId()), true, true);
+
+    final boolean result = TestUtils.waitForCondition(() -> !((RayNativeRuntime) Ray.internal())
+          .rayletTaskExistsInGcs(UniqueIdUtil.computeTaskId(helloId.getId())), 50);
+    Assert.assertTrue(result);
+  }
+
 }