Add placement group java api (#9611)

* add part code

* add part code

* add part code

* fix code style

* fix review comment

* fix review comment

* add part code

* add part code

* add part code

* add part code

* fix review comment

* fix review comment

* fix code style

* fix review comment

* fix lint error

* fix lint error

Co-authored-by: 灵洵 <fengbin.ffb@antfin.com>

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

@@ -2,11 +2,14 @@ package io.ray.api;
 
 import io.ray.api.id.ObjectId;
 import io.ray.api.id.UniqueId;
+import io.ray.api.placementgroup.PlacementGroup;
+import io.ray.api.placementgroup.PlacementStrategy;
 import io.ray.api.runtime.RayRuntime;
 import io.ray.api.runtime.RayRuntimeFactory;
 import io.ray.api.runtimecontext.RuntimeContext;
 import java.util.ArrayList;
 import java.util.List;
+import java.util.Map;
 import java.util.Optional;
 import java.util.concurrent.Callable;
 
@@ -240,4 +243,22 @@ public final class Ray extends RayCall {
   public static RuntimeContext getRuntimeContext() {
     return runtime.getRuntimeContext();
   }
+
+  /**
+   * Create a placement group.
+   * A placement group is used to place actors according to a specific strategy
+   * and resource constraints.
+   * It will sends a request to GCS to preallocate the specified resources, which is asynchronous.
+   * If the specified resource cannot be allocated, it will wait for the resource
+   * to be updated and rescheduled.
+   * This function only works when gcs actor manager is turned on.
+   *
+   * @param bundles Preallocated resource list.
+   * @param strategy Actor placement strategy.
+   * @return A handle to the created placement group.
+   */
+  public static PlacementGroup createPlacementGroup(List<Map<String, Double>> bundles,
+      PlacementStrategy strategy) {
+    return runtime.createPlacementGroup(bundles, strategy);
+  }
 }

java/api/src/main/java/io/ray/api/call/BaseActorCreator.java

@@ -2,6 +2,7 @@ package io.ray.api.call;
 
 import io.ray.api.Ray;
 import io.ray.api.options.ActorCreationOptions;
+import io.ray.api.placementgroup.PlacementGroup;
 import java.util.Map;
 
 /**
@@ -85,7 +86,6 @@ public class BaseActorCreator<T extends BaseActorCreator> {
   }
 
   /**
-   * /**
    * Set the max number of concurrent calls to allow for this actor.
    * <p>
    * The max concurrency defaults to 1 for threaded execution.
@@ -100,6 +100,19 @@ public class BaseActorCreator<T extends BaseActorCreator> {
     return self();
   }
 
+  /**
+   * Set the placement group to place this actor in.
+   *
+   * @param group The placement group of the actor.
+   * @param bundleIndex The index of the bundle to place this actor in.
+   * @return self
+   * @see ActorCreationOptions.Builder#setPlacementGroup(PlacementGroup, int)
+   */
+  public T setPlacementGroup(PlacementGroup group, int bundleIndex) {
+    builder.setPlacementGroup(group, bundleIndex);
+    return self();
+  }
+
   @SuppressWarnings("unchecked")
   private T self() {
     return (T) this;

java/api/src/main/java/io/ray/api/options/ActorCreationOptions.java

@@ -1,6 +1,7 @@
 package io.ray.api.options;
 
 import io.ray.api.Ray;
+import io.ray.api.placementgroup.PlacementGroup;
 import java.util.HashMap;
 import java.util.Map;
 
@@ -13,15 +14,20 @@ public class ActorCreationOptions extends BaseTaskOptions {
   public final int maxRestarts;
   public final String jvmOptions;
   public final int maxConcurrency;
+  public final PlacementGroup group;
+  public final int bundleIndex;
 
   private ActorCreationOptions(boolean global, String name, Map<String, Double> resources,
-                               int maxRestarts, String jvmOptions, int maxConcurrency) {
+                               int maxRestarts, String jvmOptions, int maxConcurrency,
+                               PlacementGroup group, int bundleIndex) {
     super(resources);
     this.global = global;
     this.name = name;
     this.maxRestarts = maxRestarts;
     this.jvmOptions = jvmOptions;
     this.maxConcurrency = maxConcurrency;
+    this.group = group;
+    this.bundleIndex = bundleIndex;
   }
 
   /**
@@ -34,6 +40,8 @@ public class ActorCreationOptions extends BaseTaskOptions {
     private int maxRestarts = 0;
     private String jvmOptions = null;
     private int maxConcurrency = 1;
+    private PlacementGroup group;
+    private int bundleIndex;
 
     /**
      * Set the actor name of a named actor.
@@ -135,9 +143,22 @@ public class ActorCreationOptions extends BaseTaskOptions {
       return this;
     }
 
+    /**
+     * Set the placement group to place this actor in.
+     *
+     * @param group The placement group of the actor.
+     * @param bundleIndex The index of the bundle to place this actor in.
+     * @return self
+     */
+    public Builder setPlacementGroup(PlacementGroup group, int bundleIndex) {
+      this.group = group;
+      this.bundleIndex = bundleIndex;
+      return this;
+    }
+
     public ActorCreationOptions build() {
       return new ActorCreationOptions(
-          global, name, resources, maxRestarts, jvmOptions, maxConcurrency);
+          global, name, resources, maxRestarts, jvmOptions, maxConcurrency, group, bundleIndex);
     }
   }
 

java/api/src/main/java/io/ray/api/placementgroup/PlacementGroup.java

@@ -0,0 +1,10 @@
+package io.ray.api.placementgroup;
+
+/**
+ * A placement group is used to place interdependent actors according to a specific strategy
+ * {@link PlacementStrategy}.
+ * When a placement group is created, the corresponding actor slots and resources are preallocated.
+ * A placement group consists of one or more bundles plus a specific placement strategy.
+ */
+public interface PlacementGroup {
+}

java/api/src/main/java/io/ray/api/placementgroup/PlacementStrategy.java

@@ -0,0 +1,25 @@
+package io.ray.api.placementgroup;
+
+/**
+ * The actor placement strategy.
+ */
+public enum PlacementStrategy {
+  /**
+   * Packs Bundles close together inside nodes as tight as possible.
+   */
+  PACK(0),
+  /**
+   * Places Bundles across distinct nodes as even as possible.
+   */
+  SPREAD(1);
+
+  private int value = 0;
+
+  PlacementStrategy(int value) {
+    this.value = value;
+  }
+
+  public int value() {
+    return this.value;
+  }
+}

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

@@ -14,8 +14,11 @@ import io.ray.api.id.ObjectId;
 import io.ray.api.id.UniqueId;
 import io.ray.api.options.ActorCreationOptions;
 import io.ray.api.options.CallOptions;
+import io.ray.api.placementgroup.PlacementGroup;
+import io.ray.api.placementgroup.PlacementStrategy;
 import io.ray.api.runtimecontext.RuntimeContext;
 import java.util.List;
+import java.util.Map;
 import java.util.Optional;
 import java.util.concurrent.Callable;
 
@@ -169,6 +172,9 @@ public interface RayRuntime {
   PyActorHandle createActor(PyActorClass pyActorClass, Object[] args,
                             ActorCreationOptions options);
 
+  PlacementGroup createPlacementGroup(List<Map<String, Double>> bundles,
+      PlacementStrategy strategy);
+
   RuntimeContext getRuntimeContext();
 
   Object getAsyncContext();

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

@@ -17,6 +17,8 @@ import io.ray.api.id.ActorId;
 import io.ray.api.id.ObjectId;
 import io.ray.api.options.ActorCreationOptions;
 import io.ray.api.options.CallOptions;
+import io.ray.api.placementgroup.PlacementGroup;
+import io.ray.api.placementgroup.PlacementStrategy;
 import io.ray.api.runtimecontext.RuntimeContext;
 import io.ray.runtime.config.RayConfig;
 import io.ray.runtime.context.RuntimeContextImpl;
@@ -35,6 +37,7 @@ import io.ray.runtime.task.FunctionArg;
 import io.ray.runtime.task.TaskExecutor;
 import io.ray.runtime.task.TaskSubmitter;
 import java.util.List;
+import java.util.Map;
 import java.util.Optional;
 import java.util.concurrent.Callable;
 import org.slf4j.Logger;
@@ -155,6 +158,12 @@ public abstract class AbstractRayRuntime implements RayRuntimeInternal {
     return (PyActorHandle) createActorImpl(functionDescriptor, args, options);
   }
 
+  @Override
+  public PlacementGroup createPlacementGroup(List<Map<String, Double>> bundles,
+      PlacementStrategy strategy) {
+    return taskSubmitter.createPlacementGroup(bundles, strategy);
+  }
+
   @SuppressWarnings("unchecked")
   @Override
   public <T extends BaseActorHandle> T getActorHandle(ActorId actorId) {

java/runtime/src/main/java/io/ray/runtime/placementgroup/PlacementGroupId.java

@@ -0,0 +1,58 @@
+package io.ray.runtime.placementgroup;
+
+import io.ray.api.id.BaseId;
+import java.io.Serializable;
+import java.nio.ByteBuffer;
+import java.util.Arrays;
+import java.util.Random;
+
+/**
+ * Represents the id of a placement group.
+ */
+public class PlacementGroupId extends BaseId implements Serializable {
+
+  public static final int LENGTH = 16;
+
+  public static final PlacementGroupId NIL = nil();
+
+  private PlacementGroupId(byte[] id) {
+    super(id);
+  }
+
+  /**
+   * Creates a PlacementGroupId from the given ByteBuffer.
+   */
+  public static PlacementGroupId fromByteBuffer(ByteBuffer bb) {
+    return new PlacementGroupId(byteBuffer2Bytes(bb));
+  }
+
+  /**
+   * Create a PlacementGroupId instance according to the given bytes.
+   */
+  public static PlacementGroupId fromBytes(byte[] bytes) {
+    return new PlacementGroupId(bytes);
+  }
+
+  /**
+   * Generate a nil PlacementGroupId.
+   */
+  private static PlacementGroupId nil() {
+    byte[] b = new byte[LENGTH];
+    Arrays.fill(b, (byte) 0xFF);
+    return new PlacementGroupId(b);
+  }
+
+  /**
+   * Generate an PlacementGroupId with random value. Used for local mode and test only.
+   */
+  public static PlacementGroupId fromRandom() {
+    byte[] b = new byte[LENGTH];
+    new Random().nextBytes(b);
+    return new PlacementGroupId(b);
+  }
+
+  @Override
+  public int size() {
+    return LENGTH;
+  }
+}

java/runtime/src/main/java/io/ray/runtime/placementgroup/PlacementGroupImpl.java

@@ -0,0 +1,28 @@
+package io.ray.runtime.placementgroup;
+
+import io.ray.api.placementgroup.PlacementGroup;
+
+/**
+ * The default implementation of `PlacementGroup` interface.
+ */
+public class PlacementGroupImpl implements PlacementGroup {
+
+  private PlacementGroupId id;
+  private int bundleCount = 0;
+
+  public PlacementGroupImpl() {
+  }
+
+  public PlacementGroupImpl(PlacementGroupId id, int bundleCount) {
+    this.id = id;
+    this.bundleCount = bundleCount;
+  }
+
+  public PlacementGroupId getId() {
+    return id;
+  }
+
+  public int getBundleCount() {
+    return bundleCount;
+  }
+}

java/runtime/src/main/java/io/ray/runtime/task/LocalModeTaskSubmitter.java

@@ -12,6 +12,8 @@ import io.ray.api.id.TaskId;
 import io.ray.api.id.UniqueId;
 import io.ray.api.options.ActorCreationOptions;
 import io.ray.api.options.CallOptions;
+import io.ray.api.placementgroup.PlacementGroup;
+import io.ray.api.placementgroup.PlacementStrategy;
 import io.ray.runtime.RayRuntimeInternal;
 import io.ray.runtime.actor.LocalModeActorHandle;
 import io.ray.runtime.context.LocalModeWorkerContext;
@@ -27,6 +29,7 @@ import io.ray.runtime.generated.Common.TaskSpec;
 import io.ray.runtime.generated.Common.TaskType;
 import io.ray.runtime.object.LocalModeObjectStore;
 import io.ray.runtime.object.NativeRayObject;
+import io.ray.runtime.placementgroup.PlacementGroupImpl;
 import java.nio.ByteBuffer;
 import java.util.ArrayList;
 import java.util.Collections;
@@ -102,7 +105,7 @@ public class LocalModeTaskSubmitter implements TaskSubmitter {
     // Check whether task arguments are ready.
     for (TaskArg arg : taskSpec.getArgsList()) {
       ByteString idByteString = arg.getObjectRef().getObjectId();
-      if (idByteString  != ByteString.EMPTY) {
+      if (idByteString != ByteString.EMPTY) {
         ObjectId id = new ObjectId(idByteString.toByteArray());
         if (!objectStore.isObjectReady(id)) {
           unreadyObjects.add(id);
@@ -209,6 +212,12 @@ public class LocalModeTaskSubmitter implements TaskSubmitter {
     }
   }
 
+  @Override
+  public PlacementGroup createPlacementGroup(List<Map<String, Double>> bundles,
+      PlacementStrategy strategy) {
+    return new PlacementGroupImpl();
+  }
+
   @Override
   public BaseActorHandle getActor(ActorId actorId) {
     return actorHandles.get(actorId).copy();

java/runtime/src/main/java/io/ray/runtime/task/NativeTaskSubmitter.java

@@ -8,9 +8,14 @@ import io.ray.api.id.ActorId;
 import io.ray.api.id.ObjectId;
 import io.ray.api.options.ActorCreationOptions;
 import io.ray.api.options.CallOptions;
+import io.ray.api.placementgroup.PlacementGroup;
+import io.ray.api.placementgroup.PlacementStrategy;
 import io.ray.runtime.actor.NativeActorHandle;
 import io.ray.runtime.functionmanager.FunctionDescriptor;
+import io.ray.runtime.placementgroup.PlacementGroupId;
+import io.ray.runtime.placementgroup.PlacementGroupImpl;
 import java.util.List;
+import java.util.Map;
 import java.util.Optional;
 import java.util.stream.Collectors;
 import org.apache.commons.lang3.StringUtils;
@@ -34,11 +39,20 @@ public class NativeTaskSubmitter implements TaskSubmitter {
   @Override
   public BaseActorHandle createActor(FunctionDescriptor functionDescriptor, List<FunctionArg> args,
                                      ActorCreationOptions options) throws IllegalArgumentException {
-    if (options != null && StringUtils.isNotBlank(options.name)) {
-      Optional<BaseActorHandle> actor =
-          options.global ? Ray.getGlobalActor(options.name) : Ray.getActor(options.name);
-      Preconditions.checkArgument(!actor.isPresent(),
-          String.format("Actor of name %s exists", options.name));
+    if (options != null) {
+      if (options.group != null) {
+        PlacementGroupImpl group = (PlacementGroupImpl)options.group;
+        Preconditions.checkArgument(options.bundleIndex >= 0
+                && options.bundleIndex < group.getBundleCount(),
+            String.format("Bundle index %s is invalid", options.bundleIndex));
+      }
+
+      if (StringUtils.isNotBlank(options.name)) {
+        Optional<BaseActorHandle> actor =
+            options.global ? Ray.getGlobalActor(options.name) : Ray.getActor(options.name);
+        Preconditions.checkArgument(!actor.isPresent(),
+            String.format("Actor of name %s exists", options.name));
+      }
     }
     byte[] actorId = nativeCreateActor(functionDescriptor, functionDescriptor.hashCode(), args,
         options);
@@ -63,6 +77,13 @@ public class NativeTaskSubmitter implements TaskSubmitter {
     return returnIds.stream().map(ObjectId::new).collect(Collectors.toList());
   }
 
+  @Override
+  public PlacementGroup createPlacementGroup(List<Map<String, Double>> bundles,
+      PlacementStrategy strategy) {
+    byte[] bytes = nativeCreatePlacementGroup(bundles, strategy.value());
+    return new PlacementGroupImpl(PlacementGroupId.fromBytes(bytes), bundles.size());
+  }
+
   private static native List<byte[]> nativeSubmitTask(FunctionDescriptor functionDescriptor,
       int functionDescriptorHash, List<FunctionArg> args, int numReturns, CallOptions callOptions);
 
@@ -73,4 +94,7 @@ public class NativeTaskSubmitter implements TaskSubmitter {
   private static native List<byte[]> nativeSubmitActorTask(byte[] actorId,
       FunctionDescriptor functionDescriptor, int functionDescriptorHash, List<FunctionArg> args,
       int numReturns, CallOptions callOptions);
+
+  private static native byte[] nativeCreatePlacementGroup(List<Map<String, Double>> bundles,
+      int strategy);
 }

java/runtime/src/main/java/io/ray/runtime/task/TaskSubmitter.java

@@ -5,8 +5,11 @@ import io.ray.api.id.ActorId;
 import io.ray.api.id.ObjectId;
 import io.ray.api.options.ActorCreationOptions;
 import io.ray.api.options.CallOptions;
+import io.ray.api.placementgroup.PlacementGroup;
+import io.ray.api.placementgroup.PlacementStrategy;
 import io.ray.runtime.functionmanager.FunctionDescriptor;
 import java.util.List;
+import java.util.Map;
 
 /**
  * A set of methods to submit tasks and create actors.
@@ -47,6 +50,15 @@ public interface TaskSubmitter {
   List<ObjectId> submitActorTask(BaseActorHandle actor, FunctionDescriptor functionDescriptor,
                                  List<FunctionArg> args, int numReturns, CallOptions options);
 
+  /**
+   * Create a placement group.
+   * @param bundles Preallocated resource list.
+   * @param strategy Actor placement strategy.
+   * @return A handle to the created placement group.
+   */
+  PlacementGroup createPlacementGroup(List<Map<String, Double>> bundles,
+      PlacementStrategy strategy);
+
   BaseActorHandle getActor(ActorId actorId);
 
 }

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

@@ -0,0 +1,50 @@
+package io.ray.test;
+
+import io.ray.api.ActorHandle;
+import io.ray.api.Ray;
+import io.ray.api.id.ActorId;
+import io.ray.api.placementgroup.PlacementGroup;
+import io.ray.api.placementgroup.PlacementStrategy;
+import java.util.ArrayList;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+import org.testng.Assert;
+import org.testng.annotations.Test;
+
+@Test
+public class PlacementGroupTest extends BaseTest {
+
+  public static class Counter {
+
+    private int value;
+
+    public Counter(int initValue) {
+      this.value = initValue;
+    }
+
+    public int getValue() {
+      return value;
+    }
+  }
+
+  // TODO(ffbin): Currently Java doesn't support multi-node tests.
+  // This test just creates a placement group with one bundle.
+  // It's not comprehensive to test all placement group test cases.
+  public void testCreateAndCallActor() {
+    List<Map<String, Double>> bundles = new ArrayList<>();
+    Map<String, Double> bundle = new HashMap<>();
+    bundle.put("CPU", 1.0);
+    bundles.add(bundle);
+    PlacementStrategy strategy = PlacementStrategy.PACK;
+    PlacementGroup placementGroup = Ray.createPlacementGroup(bundles, strategy);
+
+    // Test creating an actor from a constructor.
+    ActorHandle<Counter> actor = Ray.actor(Counter::new, 1)
+        .setPlacementGroup(placementGroup, 0).remote();
+    Assert.assertNotEquals(actor.getId(), ActorId.NIL);
+
+    // Test calling an actor.
+    Assert.assertEquals(Integer.valueOf(1), actor.task(Counter::getValue).remote().get());
+  }
+}