[Java] Add fetchLocal parameter in Ray.wait() (#13604)

java/api/src/main/java/io/ray/api/Ray.java

@@ -87,6 +87,24 @@ public final class Ray extends RayCall {
     return internal().get(objectList);
   }
 
+  /**
+   * Wait for a list of RayObjects to be available, until specified number of objects are ready, or
+   * specified timeout has passed.
+   *
+   * @param waitList A list of object references to wait for.
+   * @param numReturns The number of objects that should be returned.
+   * @param timeoutMs The maximum time in milliseconds to wait before returning.
+   * @param fetchLocal If true, wait for the object to be downloaded onto the local node before
+   *     returning it as ready. If false, ray.wait() will not trigger fetching of objects to the
+   *     local node and will return immediately once the object is available anywhere in the
+   *     cluster.
+   * @return Two lists, one containing locally available objects, one containing the rest.
+   */
+  public static <T> WaitResult<T> wait(
+      List<ObjectRef<T>> waitList, int numReturns, int timeoutMs, boolean fetchLocal) {
+    return internal().wait(waitList, numReturns, timeoutMs, fetchLocal);
+  }
+
   /**
    * Wait for a list of RayObjects to be locally available, until specified number of objects are
    * ready, or specified timeout has passed.
@@ -97,30 +115,29 @@ public final class Ray extends RayCall {
    * @return Two lists, one containing locally available objects, one containing the rest.
    */
   public static <T> WaitResult<T> wait(List<ObjectRef<T>> waitList, int numReturns, int timeoutMs) {
-    return internal().wait(waitList, numReturns, timeoutMs);
+    return wait(waitList, numReturns, timeoutMs, true);
   }
 
   /**
-   * A convenient helper method for Ray.wait. It will wait infinitely until specified number of
-   * objects are locally available.
+   * Wait for a list of RayObjects to be locally available, until specified number of objects are
+   * ready.
    *
    * @param waitList A list of object references to wait for.
    * @param numReturns The number of objects that should be returned.
    * @return Two lists, one containing locally available objects, one containing the rest.
    */
   public static <T> WaitResult<T> wait(List<ObjectRef<T>> waitList, int numReturns) {
-    return internal().wait(waitList, numReturns, Integer.MAX_VALUE);
+    return wait(waitList, numReturns, Integer.MAX_VALUE);
   }
 
   /**
-   * A convenient helper method for Ray.wait. It will wait infinitely until all objects are locally
-   * available.
+   * Wait for a list of RayObjects to be locally available.
    *
    * @param waitList A list of object references to wait for.
    * @return Two lists, one containing locally available objects, one containing the rest.
    */
   public static <T> WaitResult<T> wait(List<ObjectRef<T>> waitList) {
-    return internal().wait(waitList, waitList.size(), Integer.MAX_VALUE);
+    return wait(waitList, waitList.size());
   }
 
   /**

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

@@ -53,15 +53,20 @@ public interface RayRuntime {
   <T> List<T> get(List<ObjectRef<T>> objectRefs);
 
   /**
-   * Wait for a list of RayObjects to be locally available, until specified number of objects are
-   * ready, or specified timeout has passed.
+   * Wait for a list of RayObjects to be available, until specified number of objects are ready, or
+   * specified timeout has passed.
    *
    * @param waitList A list of ObjectRef to wait for.
    * @param numReturns The number of objects that should be returned.
    * @param timeoutMs The maximum time in milliseconds to wait before returning.
+   * @param fetchLocal If true, wait for the object to be downloaded onto the local node before
+   *     returning it as ready. If false, ray.wait() will not trigger fetching of objects to the
+   *     local node and will return immediately once the object is available anywhere in the
+   *     cluster.
    * @return Two lists, one containing locally available objects, one containing the rest.
    */
-  <T> WaitResult<T> wait(List<ObjectRef<T>> waitList, int numReturns, int timeoutMs);
+  <T> WaitResult<T> wait(
+      List<ObjectRef<T>> waitList, int numReturns, int timeoutMs, boolean fetchLocal);
 
   /**
    * Free a list of objects from Plasma Store.

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

@@ -105,8 +105,9 @@ public abstract class AbstractRayRuntime implements RayRuntimeInternal {
   }
 
   @Override
-  public <T> WaitResult<T> wait(List<ObjectRef<T>> waitList, int numReturns, int timeoutMs) {
-    return objectStore.wait(waitList, numReturns, timeoutMs);
+  public <T> WaitResult<T> wait(
+      List<ObjectRef<T>> waitList, int numReturns, int timeoutMs, boolean fetchLocal) {
+    return objectStore.wait(waitList, numReturns, timeoutMs, fetchLocal);
   }
 
   @Override

java/runtime/src/main/java/io/ray/runtime/object/LocalModeObjectStore.java

@@ -60,7 +60,8 @@ public class LocalModeObjectStore extends ObjectStore {
   }
 
   @Override
-  public List<Boolean> wait(List<ObjectId> objectIds, int numObjects, long timeoutMs) {
+  public List<Boolean> wait(
+      List<ObjectId> objectIds, int numObjects, long timeoutMs, boolean fetchLocal) {
     waitInternal(objectIds, numObjects, timeoutMs);
     return objectIds.stream().map(pool::containsKey).collect(Collectors.toList());
   }

java/runtime/src/main/java/io/ray/runtime/object/NativeObjectStore.java

@@ -45,8 +45,9 @@ public class NativeObjectStore extends ObjectStore {
   }
 
   @Override
-  public List<Boolean> wait(List<ObjectId> objectIds, int numObjects, long timeoutMs) {
-    return nativeWait(toBinaryList(objectIds), numObjects, timeoutMs);
+  public List<Boolean> wait(
+      List<ObjectId> objectIds, int numObjects, long timeoutMs, boolean fetchLocal) {
+    return nativeWait(toBinaryList(objectIds), numObjects, timeoutMs, fetchLocal);
   }
 
   @Override
@@ -113,7 +114,7 @@ public class NativeObjectStore extends ObjectStore {
   private static native List<NativeRayObject> nativeGet(List<byte[]> ids, long timeoutMs);
 
   private static native List<Boolean> nativeWait(
-      List<byte[]> objectIds, int numObjects, long timeoutMs);
+      List<byte[]> objectIds, int numObjects, long timeoutMs, boolean fetchLocal);
 
   private static native void nativeDelete(List<byte[]> objectIds, boolean localOnly);
 

java/runtime/src/main/java/io/ray/runtime/object/ObjectStore.java

@@ -117,25 +117,36 @@ public abstract class ObjectStore {
   }
 
   /**
-   * Wait for a list of objects to appear in the object store.
+   * Wait for a list of RayObjects to be available, until specified number of objects are ready, or
+   * specified timeout has passed.
    *
    * @param objectIds IDs of the objects to wait for.
    * @param numObjects Number of objects that should appear.
    * @param timeoutMs Timeout in milliseconds, wait infinitely if it's negative.
+   * @param fetchLocal If true, wait for the object to be downloaded onto the local node before
+   *     returning it as ready. If false, ray.wait() will not trigger fetching of objects to the
+   *     local node and will return immediately once the object is available anywhere in the
+   *     cluster.
    * @return A bitset that indicates each object has appeared or not.
    */
-  public abstract List<Boolean> wait(List<ObjectId> objectIds, int numObjects, long timeoutMs);
+  public abstract List<Boolean> wait(
+      List<ObjectId> objectIds, int numObjects, long timeoutMs, boolean fetchLocal);
 
   /**
-   * Wait for a list of RayObjects to be locally available, until specified number of objects are
-   * ready, or specified timeout has passed.
+   * Wait for a list of RayObjects to be available, until specified number of objects are ready, or
+   * specified timeout has passed.
    *
    * @param waitList A list of object references to wait for.
    * @param numReturns The number of objects that should be returned.
    * @param timeoutMs The maximum time in milliseconds to wait before returning.
+   * @param fetchLocal If true, wait for the object to be downloaded onto the local node before
+   *     returning it as ready. If false, ray.wait() will not trigger fetching of objects to the
+   *     local node and will return immediately once the object is available anywhere in the
+   *     cluster.
    * @return Two lists, one containing locally available objects, one containing the rest.
    */
-  public <T> WaitResult<T> wait(List<ObjectRef<T>> waitList, int numReturns, int timeoutMs) {
+  public <T> WaitResult<T> wait(
+      List<ObjectRef<T>> waitList, int numReturns, int timeoutMs, boolean fetchLocal) {
     Preconditions.checkNotNull(waitList);
     if (waitList.isEmpty()) {
       return new WaitResult<>(Collections.emptyList(), Collections.emptyList());
@@ -144,7 +155,7 @@ public abstract class ObjectStore {
     List<ObjectId> ids =
         waitList.stream().map(ref -> ((ObjectRefImpl<?>) ref).getId()).collect(Collectors.toList());
 
-    List<Boolean> ready = wait(ids, numReturns, timeoutMs);
+    List<Boolean> ready = wait(ids, numReturns, timeoutMs, fetchLocal);
     List<ObjectRef<T>> readyList = new ArrayList<>();
     List<ObjectRef<T>> unreadyList = new ArrayList<>();
 

java/test.sh

@@ -41,6 +41,15 @@ bazel build //java:gen_maven_deps
 echo "Build test jar."
 bazel build //java:all_tests_deploy.jar
 
+java/generate_jni_header_files.sh
+
+if ! git diff --exit-code -- java src/ray/core_worker/lib/java; then
+  echo "Files are changed after build. Common cases are:"
+  echo "    * Java native methods doesn't match JNI files. You need to either update Java code or JNI code."
+  echo "    * pom_template.xml and pom.xml doesn't match. You need to either update pom_template.xml or pom.xml."
+  exit 1
+fi
+
 # Enable multi-worker feature in Java test
 TEST_ARGS=(-Dray.job.num-java-workers-per-process=10)
 

java/test/pom.xml

@@ -117,41 +117,6 @@
           </execution>
         </executions>
       </plugin>
-
-      <plugin>
-        <groupId>com.diffplug.spotless</groupId>
-        <artifactId>spotless-maven-plugin</artifactId>
-        <version>2.6.1</version>
-        <configuration>
-          <!-- optional: limit format enforcement to just the files changed by this feature branch -->
-<!--          <ratchetFrom>origin/main</ratchetFrom>-->
-          <formats>
-            <!-- you can define as many formats as you want, each is independent -->
-            <format>
-              <!-- define the files to apply to -->
-              <includes>
-                <include>.java</include>
-              </includes>
-              <excludes>
-              </excludes>
-              <!-- define the steps to apply to those files -->
-              <trimTrailingWhitespace/>
-              <endWithNewline/>
-              <indent>
-                <tabs>true</tabs>
-                <spacesPerTab>4</spacesPerTab>
-              </indent>
-            </format>
-          </formats>
-          <!-- define a language-specific format -->
-          <java>
-            <googleJavaFormat>
-              <version>1.7</version> <!-- optional -->
-              <style>GOOGLE</style>  <!-- or AOSP (optional) -->
-            </googleJavaFormat>
-          </java>
-        </configuration>
-      </plugin>
     </plugins>
   </build>
 </project>

java/test/src/main/java/io/ray/test/PlasmaFreeTest.java

@@ -25,7 +25,7 @@ public class PlasmaFreeTest extends BaseTest {
             () ->
                 !TestUtils.getRuntime()
                     .getObjectStore()
-                    .wait(ImmutableList.of(((ObjectRefImpl<String>) helloId).getId()), 1, 0)
+                    .wait(ImmutableList.of(((ObjectRefImpl<String>) helloId).getId()), 1, 0, true)
                     .get(0),
             50);
     if (TestUtils.isSingleProcessMode()) {

java/test/src/main/java/io/ray/test/ReferenceCountingTest.java

@@ -119,7 +119,7 @@ public class ReferenceCountingTest extends BaseTest {
       TestUtils.getRuntime().getObjectStore().getRaw(ImmutableList.of(objectId), Long.MAX_VALUE);
     } else {
       List<Boolean> result =
-          TestUtils.getRuntime().getObjectStore().wait(ImmutableList.of(objectId), 1, 100);
+          TestUtils.getRuntime().getObjectStore().wait(ImmutableList.of(objectId), 1, 100, true);
       Assert.assertFalse(result.get(0));
     }
   }

src/ray/core_worker/lib/java/io_ray_runtime_RayNativeRuntime.h

@@ -25,7 +25,7 @@ extern "C" {
  * Class:     io_ray_runtime_RayNativeRuntime
  * Method:    nativeInitialize
  * Signature:
- * (ILjava/lang/String;ILjava/lang/String;Ljava/lang/String;Ljava/lang/String;[BLio/ray/runtime/gcs/GcsClientOptions;ILjava/lang/String;Ljava/util/Map;)V
+ * (ILjava/lang/String;ILjava/lang/String;Ljava/lang/String;Ljava/lang/String;[BLio/ray/runtime/gcs/GcsClientOptions;ILjava/lang/String;Ljava/util/Map;[B)V
  */
 JNIEXPORT void JNICALL Java_io_ray_runtime_RayNativeRuntime_nativeInitialize(
     JNIEnv *, jclass, jint, jstring, jint, jstring, jstring, jstring, jbyteArray, jobject,

src/ray/core_worker/lib/java/io_ray_runtime_object_NativeObjectStore.h

@@ -52,7 +52,7 @@ JNIEXPORT jobject JNICALL Java_io_ray_runtime_object_NativeObjectStore_nativeGet
 /*
  * Class:     io_ray_runtime_object_NativeObjectStore
  * Method:    nativeWait
- * Signature: (Ljava/util/List;IJ)Ljava/util/List;
+ * Signature: (Ljava/util/List;IJZ)Ljava/util/List;
  */
 JNIEXPORT jobject JNICALL Java_io_ray_runtime_object_NativeObjectStore_nativeWait(
     JNIEnv *, jclass, jobject, jint, jlong, jboolean);
@@ -68,7 +68,7 @@ JNIEXPORT void JNICALL Java_io_ray_runtime_object_NativeObjectStore_nativeDelete
 /*
  * Class:     io_ray_runtime_object_NativeObjectStore
  * Method:    nativeAddLocalReference
- * Signature: ([B)V
+ * Signature: ([B[B)V
  */
 JNIEXPORT void JNICALL
 Java_io_ray_runtime_object_NativeObjectStore_nativeAddLocalReference(JNIEnv *, jclass,
@@ -78,7 +78,7 @@ Java_io_ray_runtime_object_NativeObjectStore_nativeAddLocalReference(JNIEnv *, j
 /*
  * Class:     io_ray_runtime_object_NativeObjectStore
  * Method:    nativeRemoveLocalReference
- * Signature: ([B)V
+ * Signature: ([B[B)V
  */
 JNIEXPORT void JNICALL
 Java_io_ray_runtime_object_NativeObjectStore_nativeRemoveLocalReference(JNIEnv *, jclass,

src/ray/core_worker/lib/java/io_ray_runtime_task_NativeTaskExecutor.h

@@ -21,25 +21,6 @@
 #ifdef __cplusplus
 extern "C" {
 #endif
-#undef io_ray_runtime_task_NativeTaskExecutor_NUM_ACTOR_CHECKPOINTS_TO_KEEP
-#define io_ray_runtime_task_NativeTaskExecutor_NUM_ACTOR_CHECKPOINTS_TO_KEEP 20L
-/*
- * Class:     io_ray_runtime_task_NativeTaskExecutor
- * Method:    nativePrepareCheckpoint
- * Signature: ()[B
- */
-JNIEXPORT jbyteArray JNICALL
-Java_io_ray_runtime_task_NativeTaskExecutor_nativePrepareCheckpoint(JNIEnv *, jclass);
-
-/*
- * Class:     io_ray_runtime_task_NativeTaskExecutor
- * Method:    nativeNotifyActorResumedFromCheckpoint
- * Signature: ([B)V
- */
-JNIEXPORT void JNICALL
-Java_io_ray_runtime_task_NativeTaskExecutor_nativeNotifyActorResumedFromCheckpoint(
-    JNIEnv *, jclass, jbyteArray);
-
 #ifdef __cplusplus
 }
 #endif

src/ray/core_worker/lib/java/io_ray_runtime_task_NativeTaskSubmitter.h

@@ -74,13 +74,13 @@ Java_io_ray_runtime_task_NativeTaskSubmitter_nativeRemovePlacementGroup(JNIEnv *
 /*
  * Class:     io_ray_runtime_task_NativeTaskSubmitter
  * Method:    nativeWaitPlacementGroupReady
- * Signature: (J)Z
+ * Signature: ([BI)Z
  */
 JNIEXPORT jboolean JNICALL
-Java_io_ray_runtime_task_NativeTaskSubmitter__nativeWaitPlacementGroupReady(JNIEnv *,
-                                                                            jclass,
-                                                                            jbyteArray,
-                                                                            jint);
+Java_io_ray_runtime_task_NativeTaskSubmitter_nativeWaitPlacementGroupReady(JNIEnv *,
+                                                                           jclass,
+                                                                           jbyteArray,
+                                                                           jint);
 
 #ifdef __cplusplus
 }