[ID Refactor] Shorten the length of JobID to 4 bytes (#5110)

* WIP

* Fix

* Add jobid test

* Fix

* Add python part

* Fix

* Fix tes

* Remove TODOs

* Fix C++ tests

* Lint

* Fix

* Fix exporting functions in multiple ray.init

* Fix java test

* Fix lint

* Fix linting

* Address comments.

* FIx

* Address and fix linting

* Refine and fix

* Fix

* address

* Address comments.

* Fix linting

* Fix

* Address

* Address comments.

* Address

* Address

* Fix

* Fix

* Fix

* Fix lint

* Fix

* Fix linting

* Address comments.

* Fix linting

* Address comments.

* Fix linting

* address comments.

* Fix

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

@@ -76,8 +76,6 @@ public abstract class AbstractRayRuntime implements RayRuntime {
     this.rayConfig = rayConfig;
     functionManager = new FunctionManager(rayConfig.jobResourcePath);
     worker = new Worker(this);
-    workerContext = new WorkerContext(rayConfig.workerMode,
-        rayConfig.jobId, rayConfig.runMode);
     runtimeContext = new RuntimeContextImpl(this);
   }
 

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

@@ -1,5 +1,7 @@
 package org.ray.runtime;
 
+import java.util.concurrent.atomic.AtomicInteger;
+import org.ray.api.id.JobId;
 import org.ray.runtime.config.RayConfig;
 import org.ray.runtime.objectstore.MockObjectStore;
 import org.ray.runtime.objectstore.ObjectStoreProxy;
@@ -13,9 +15,16 @@ public class RayDevRuntime extends AbstractRayRuntime {
 
   private MockObjectStore store;
 
+  private AtomicInteger jobCounter = new AtomicInteger(0);
+
   @Override
   public void start() {
     store = new MockObjectStore(this);
+    if (rayConfig.getJobId().isNil()) {
+      rayConfig.setJobId(nextJobId());
+    }
+    workerContext = new WorkerContext(rayConfig.workerMode,
+        rayConfig.getJobId(), rayConfig.runMode);
     objectStoreProxy = new ObjectStoreProxy(this, null);
     rayletClient = new MockRayletClient(this, rayConfig.numberExecThreadsForDevRuntime);
   }
@@ -33,4 +42,8 @@ public class RayDevRuntime extends AbstractRayRuntime {
   public Worker getWorker() {
     return ((MockRayletClient) rayletClient).getCurrentWorker();
   }
+
+  private JobId nextJobId() {
+    return JobId.fromInt(jobCounter.getAndIncrement());
+  }
 }

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

@@ -11,6 +11,7 @@ import java.nio.file.Paths;
 import java.nio.file.StandardCopyOption;
 import java.util.HashMap;
 import java.util.Map;
+import org.ray.api.id.JobId;
 import org.ray.runtime.config.RayConfig;
 import org.ray.runtime.config.WorkerMode;
 import org.ray.runtime.gcs.GcsClient;
@@ -94,6 +95,12 @@ public final class RayNativeRuntime extends AbstractRayRuntime {
 
     gcsClient = new GcsClient(rayConfig.getRedisAddress(), rayConfig.redisPassword);
 
+    if (rayConfig.getJobId() == JobId.NIL) {
+      rayConfig.setJobId(gcsClient.nextJobId());
+    }
+
+    workerContext = new WorkerContext(rayConfig.workerMode,
+        rayConfig.getJobId(), rayConfig.runMode);
     // TODO(qwang): Get object_store_socket_name and raylet_socket_name from Redis.
     objectStoreProxy = new ObjectStoreProxy(this, rayConfig.objectStoreSocketName);
 

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

@@ -2,6 +2,7 @@ package org.ray.runtime;
 
 import com.google.common.base.Preconditions;
 import java.util.List;
+import org.ray.api.id.JobId;
 import org.ray.api.id.UniqueId;
 import org.ray.api.runtimecontext.NodeInfo;
 import org.ray.api.runtimecontext.RuntimeContext;
@@ -17,7 +18,7 @@ public class RuntimeContextImpl implements RuntimeContext {
   }
 
   @Override
-  public UniqueId getCurrentJobId() {
+  public JobId getCurrentJobId() {
     return runtime.getWorkerContext().getCurrentJobId();
   }
 

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

@@ -1,11 +1,13 @@
 package org.ray.runtime;
 
 import com.google.common.base.Preconditions;
+import org.ray.api.id.JobId;
 import org.ray.api.id.TaskId;
 import org.ray.api.id.UniqueId;
 import org.ray.runtime.config.RunMode;
 import org.ray.runtime.config.WorkerMode;
 import org.ray.runtime.task.TaskSpec;
+import org.ray.runtime.util.IdUtil;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
@@ -29,7 +31,7 @@ public class WorkerContext {
 
   private ThreadLocal<TaskSpec> currentTask;
 
-  private UniqueId currentJobId;
+  private JobId currentJobId;
 
   private ClassLoader currentClassLoader;
 
@@ -43,7 +45,7 @@ public class WorkerContext {
    */
   private RunMode runMode;
 
-  public WorkerContext(WorkerMode workerMode, UniqueId jobId, RunMode runMode) {
+  public WorkerContext(WorkerMode workerMode, JobId jobId, RunMode runMode) {
     mainThreadId = Thread.currentThread().getId();
     taskIndex = ThreadLocal.withInitial(() -> 0);
     putIndex = ThreadLocal.withInitial(() -> 0);
@@ -52,15 +54,13 @@ public class WorkerContext {
     currentTask = ThreadLocal.withInitial(() -> null);
     currentClassLoader = null;
     if (workerMode == WorkerMode.DRIVER) {
-      // TODO(qwang): Assign the driver id to worker id
-      // once we treat driver id as a special worker id.
-      workerId = jobId;
+      workerId = IdUtil.computeDriverId(jobId);
       currentTaskId.set(TaskId.randomId());
       currentJobId = jobId;
     } else {
       workerId = UniqueId.randomId();
       this.currentTaskId.set(TaskId.NIL);
-      this.currentJobId = UniqueId.NIL;
+      this.currentJobId = JobId.NIL;
     }
   }
 
@@ -119,7 +119,7 @@ public class WorkerContext {
   /**
    * The ID of the current job.
    */
-  public UniqueId getCurrentJobId() {
+  public JobId getCurrentJobId() {
     return currentJobId;
   }
 

java/runtime/src/main/java/org/ray/runtime/config/RayConfig.java

@@ -10,7 +10,7 @@ import java.io.File;
 import java.util.ArrayList;
 import java.util.List;
 import java.util.Map;
-import org.ray.api.id.UniqueId;
+import org.ray.api.id.JobId;
 import org.ray.runtime.util.NetworkUtil;
 import org.ray.runtime.util.ResourceUtil;
 import org.ray.runtime.util.StringUtil;
@@ -32,7 +32,7 @@ public class RayConfig {
   public final WorkerMode workerMode;
   public final RunMode runMode;
   public final Map<String, Double> resources;
-  public final UniqueId jobId;
+  private JobId jobId;
   public final String logDir;
   public final boolean redirectOutput;
   public final List<String> libraryPath;
@@ -108,9 +108,9 @@ public class RayConfig {
     // Job id.
     String jobId = config.getString("ray.job.id");
     if (!jobId.isEmpty()) {
-      this.jobId = UniqueId.fromHexString(jobId);
+      this.jobId = JobId.fromHexString(jobId);
     } else {
-      this.jobId = UniqueId.randomId();
+      this.jobId = JobId.NIL;
     }
     // Log dir.
     logDir = removeTrailingSlash(config.getString("ray.log-dir"));
@@ -198,6 +198,14 @@ public class RayConfig {
     return redisPort;
   }
 
+  public void setJobId(JobId jobId) {
+    this.jobId = jobId;
+  }
+
+  public JobId getJobId() {
+    return this.jobId;
+  }
+
   @Override
   public String toString() {
     return "RayConfig{"

java/runtime/src/main/java/org/ray/runtime/functionmanager/FunctionManager.java

@@ -24,7 +24,7 @@ import org.apache.commons.lang3.tuple.ImmutablePair;
 import org.apache.commons.lang3.tuple.Pair;
 import org.objectweb.asm.Type;
 import org.ray.api.function.RayFunc;
-import org.ray.api.id.UniqueId;
+import org.ray.api.id.JobId;
 import org.ray.runtime.util.LambdaUtils;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
@@ -48,7 +48,7 @@ public class FunctionManager {
   /**
    * Mapping from the job id to the functions that belong to this job.
    */
-  private Map<UniqueId, JobFunctionTable> jobFunctionTables = new HashMap<>();
+  private Map<JobId, JobFunctionTable> jobFunctionTables = new HashMap<>();
 
   /**
    * The resource path which we can load the job's jar resources.
@@ -72,7 +72,7 @@ public class FunctionManager {
    * @param func The lambda.
    * @return A RayFunction object.
    */
-  public RayFunction getFunction(UniqueId jobId, RayFunc func) {
+  public RayFunction getFunction(JobId jobId, RayFunc func) {
     JavaFunctionDescriptor functionDescriptor = RAY_FUNC_CACHE.get().get(func.getClass());
     if (functionDescriptor == null) {
       SerializedLambda serializedLambda = LambdaUtils.getSerializedLambda(func);
@@ -92,7 +92,7 @@ public class FunctionManager {
    * @param functionDescriptor The function descriptor.
    * @return A RayFunction object.
    */
-  public RayFunction getFunction(UniqueId jobId, JavaFunctionDescriptor functionDescriptor) {
+  public RayFunction getFunction(JobId jobId, JavaFunctionDescriptor functionDescriptor) {
     JobFunctionTable jobFunctionTable = jobFunctionTables.get(jobId);
     if (jobFunctionTable == null) {
       ClassLoader classLoader;

java/runtime/src/main/java/org/ray/runtime/gcs/GcsClient.java

@@ -11,6 +11,7 @@ import java.util.stream.Collectors;
 import org.apache.commons.lang3.ArrayUtils;
 import org.ray.api.Checkpointable.Checkpoint;
 import org.ray.api.id.BaseId;
+import org.ray.api.id.JobId;
 import org.ray.api.id.TaskId;
 import org.ray.api.id.UniqueId;
 import org.ray.api.runtimecontext.NodeInfo;
@@ -164,6 +165,11 @@ public class GcsClient {
     return checkpoints;
   }
 
+  public JobId nextJobId() {
+    int jobCounter = (int) primary.incr("JobCounter".getBytes());
+    return JobId.fromInt(jobCounter);
+  }
+
   private RedisClient getShardClient(BaseId key) {
     return shards.get((int) Long.remainderUnsigned(IdUtil.murmurHashCode(key),
         shards.size()));

java/runtime/src/main/java/org/ray/runtime/gcs/RedisClient.java

@@ -107,4 +107,9 @@ public class RedisClient {
     }
   }
 
+  public long incr(byte[] key) {
+    try (Jedis jedis = jedisPool.getResource()) {
+      return jedis.incr(key).intValue();
+    }
+  }
 }

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

@@ -17,6 +17,7 @@ import java.util.concurrent.Executors;
 import org.apache.commons.lang3.NotImplementedException;
 import org.ray.api.RayObject;
 import org.ray.api.WaitResult;
+import org.ray.api.id.JobId;
 import org.ray.api.id.ObjectId;
 import org.ray.api.id.TaskId;
 import org.ray.api.id.UniqueId;
@@ -164,7 +165,7 @@ public class MockRayletClient implements RayletClient {
   }
 
   @Override
-  public TaskId generateTaskId(UniqueId jobId, TaskId parentTaskId, int taskIndex) {
+  public TaskId generateTaskId(JobId jobId, TaskId parentTaskId, int taskIndex) {
     return TaskId.randomId();
   }
 

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

@@ -3,6 +3,7 @@ package org.ray.runtime.raylet;
 import java.util.List;
 import org.ray.api.RayObject;
 import org.ray.api.WaitResult;
+import org.ray.api.id.JobId;
 import org.ray.api.id.ObjectId;
 import org.ray.api.id.TaskId;
 import org.ray.api.id.UniqueId;
@@ -21,7 +22,7 @@ public interface RayletClient {
 
   void notifyUnblocked(TaskId currentTaskId);
 
-  TaskId generateTaskId(UniqueId jobId, TaskId parentTaskId, int taskIndex);
+  TaskId generateTaskId(JobId jobId, TaskId parentTaskId, int taskIndex);
 
   <T> WaitResult<T> wait(List<RayObject<T>> waitFor, int numReturns, int
       timeoutMs, TaskId currentTaskId);

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

@@ -15,6 +15,7 @@ import java.util.stream.Collectors;
 import org.ray.api.RayObject;
 import org.ray.api.WaitResult;
 import org.ray.api.exception.RayException;
+import org.ray.api.id.JobId;
 import org.ray.api.id.ObjectId;
 import org.ray.api.id.TaskId;
 import org.ray.api.id.UniqueId;
@@ -39,7 +40,7 @@ public class RayletClientImpl implements RayletClient {
 
   // TODO(qwang): JobId parameter can be removed once we embed jobId in driverId.
   public RayletClientImpl(String schedulerSockName, UniqueId clientId,
-      boolean isWorker, UniqueId jobId) {
+                          boolean isWorker, JobId jobId) {
     client = nativeInit(schedulerSockName, clientId.getBytes(),
         isWorker, jobId.getBytes());
   }
@@ -107,7 +108,7 @@ public class RayletClientImpl implements RayletClient {
   }
 
   @Override
-  public TaskId generateTaskId(UniqueId jobId, TaskId parentTaskId, int taskIndex) {
+  public TaskId generateTaskId(JobId jobId, TaskId parentTaskId, int taskIndex) {
     byte[] bytes = nativeGenerateTaskId(jobId.getBytes(), parentTaskId.getBytes(), taskIndex);
     return new TaskId(bytes);
   }
@@ -146,7 +147,7 @@ public class RayletClientImpl implements RayletClient {
     }
 
     // Parse common fields.
-    UniqueId jobId = UniqueId.fromByteBuffer(taskSpec.getJobId().asReadOnlyByteBuffer());
+    JobId jobId = JobId.fromByteBuffer(taskSpec.getJobId().asReadOnlyByteBuffer());
     TaskId taskId = TaskId.fromByteBuffer(taskSpec.getTaskId().asReadOnlyByteBuffer());
     TaskId parentTaskId = TaskId.fromByteBuffer(taskSpec.getParentTaskId().asReadOnlyByteBuffer());
     int parentCounter = (int) taskSpec.getParentCounter();

java/runtime/src/main/java/org/ray/runtime/runner/RunManager.java

@@ -189,6 +189,8 @@ public class RunManager {
         client.auth(rayConfig.headRedisPassword);
       }
       client.set("UseRaylet", "1");
+      // Set job counter to compute job id.
+      client.set("JobCounter", "0");
       // Register the number of Redis shards in the primary shard, so that clients
       // know how many redis shards to expect under RedisShards.
       client.set("NumRedisShards", Integer.toString(rayConfig.numberRedisShards));

java/runtime/src/main/java/org/ray/runtime/task/TaskSpec.java

@@ -5,6 +5,7 @@ import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.List;
 import java.util.Map;
+import org.ray.api.id.JobId;
 import org.ray.api.id.ObjectId;
 import org.ray.api.id.TaskId;
 import org.ray.api.id.UniqueId;
@@ -19,7 +20,7 @@ import org.ray.runtime.util.IdUtil;
 public class TaskSpec {
 
   // ID of the job that created this task.
-  public final UniqueId jobId;
+  public final JobId jobId;
 
   // Task ID of the task.
   public final TaskId taskId;
@@ -81,7 +82,7 @@ public class TaskSpec {
   }
 
   public TaskSpec(
-      UniqueId jobId,
+      JobId jobId,
       TaskId taskId,
       TaskId parentTaskId,
       int parentCounter,

java/runtime/src/main/java/org/ray/runtime/util/IdUtil.java

@@ -3,13 +3,14 @@ package org.ray.runtime.util;
 import com.google.common.base.Preconditions;
 import java.nio.ByteBuffer;
 import java.nio.ByteOrder;
+import java.util.Arrays;
 import java.util.List;
 import org.ray.api.id.BaseId;
+import org.ray.api.id.JobId;
 import org.ray.api.id.ObjectId;
 import org.ray.api.id.TaskId;
 import org.ray.api.id.UniqueId;
 
-
 /**
  * Helper method for different Ids.
  * Note: any changes to these methods must be synced with C++ helper functions
@@ -153,6 +154,18 @@ public class IdUtil {
   }
 
 
+  /**
+   * Compute the driver id from the given job.
+   */
+  public static UniqueId computeDriverId(JobId jobId) {
+    byte[] bytes = new byte[UniqueId.LENGTH];
+    System.arraycopy(jobId.getBytes(), 0, bytes, 0, jobId.size());
+    Arrays.fill(bytes, jobId.size(), UniqueId.LENGTH, (byte)0xFF);
+    ByteBuffer wbb = ByteBuffer.wrap(bytes);
+    wbb.order(ByteOrder.LITTLE_ENDIAN);
+    return new UniqueId(bytes);
+  }
+
   /**
    * Compute the murmur hash code of this ID.
    */

java/runtime/src/test/java/org/ray/runtime/functionmanager/FunctionManagerTest.java

@@ -4,6 +4,7 @@ import java.io.File;
 import java.nio.file.Files;
 import java.nio.file.Paths;
 import java.util.Map;
+import java.util.Random;
 import javax.tools.JavaCompiler;
 import javax.tools.ToolProvider;
 import org.apache.commons.io.FileUtils;
@@ -12,6 +13,7 @@ import org.apache.commons.lang3.tuple.Pair;
 import org.ray.api.annotation.RayRemote;
 import org.ray.api.function.RayFunc0;
 import org.ray.api.function.RayFunc1;
+import org.ray.api.id.JobId;
 import org.ray.api.id.UniqueId;
 import org.ray.runtime.functionmanager.FunctionManager.JobFunctionTable;
 import org.testng.Assert;
@@ -64,19 +66,19 @@ public class FunctionManagerTest {
   public void testGetFunctionFromRayFunc() {
     final FunctionManager functionManager = new FunctionManager(null);
     // Test normal function.
-    RayFunction func = functionManager.getFunction(UniqueId.NIL, fooFunc);
+    RayFunction func = functionManager.getFunction(JobId.NIL, fooFunc);
     Assert.assertFalse(func.isConstructor());
     Assert.assertEquals(func.getFunctionDescriptor(), fooDescriptor);
     Assert.assertNotNull(func.getRayRemoteAnnotation());
 
     // Test actor method
-    func = functionManager.getFunction(UniqueId.NIL, barFunc);
+    func = functionManager.getFunction(JobId.NIL, barFunc);
     Assert.assertFalse(func.isConstructor());
     Assert.assertEquals(func.getFunctionDescriptor(), barDescriptor);
     Assert.assertNull(func.getRayRemoteAnnotation());
 
     // Test actor constructor
-    func = functionManager.getFunction(UniqueId.NIL, barConstructor);
+    func = functionManager.getFunction(JobId.NIL, barConstructor);
     Assert.assertTrue(func.isConstructor());
     Assert.assertEquals(func.getFunctionDescriptor(), barConstructorDescriptor);
     Assert.assertNotNull(func.getRayRemoteAnnotation());
@@ -86,19 +88,19 @@ public class FunctionManagerTest {
   public void testGetFunctionFromFunctionDescriptor() {
     final FunctionManager functionManager = new FunctionManager(null);
     // Test normal function.
-    RayFunction func = functionManager.getFunction(UniqueId.NIL, fooDescriptor);
+    RayFunction func = functionManager.getFunction(JobId.NIL, fooDescriptor);
     Assert.assertFalse(func.isConstructor());
     Assert.assertEquals(func.getFunctionDescriptor(), fooDescriptor);
     Assert.assertNotNull(func.getRayRemoteAnnotation());
 
     // Test actor method
-    func = functionManager.getFunction(UniqueId.NIL, barDescriptor);
+    func = functionManager.getFunction(JobId.NIL, barDescriptor);
     Assert.assertFalse(func.isConstructor());
     Assert.assertEquals(func.getFunctionDescriptor(), barDescriptor);
     Assert.assertNull(func.getRayRemoteAnnotation());
 
     // Test actor constructor
-    func = functionManager.getFunction(UniqueId.NIL, barConstructorDescriptor);
+    func = functionManager.getFunction(JobId.NIL, barConstructorDescriptor);
     Assert.assertTrue(func.isConstructor());
     Assert.assertEquals(func.getFunctionDescriptor(), barConstructorDescriptor);
     Assert.assertNotNull(func.getRayRemoteAnnotation());
@@ -119,7 +121,7 @@ public class FunctionManagerTest {
 
   @Test
   public void testGetFunctionFromLocalResource() throws Exception {
-    UniqueId jobId = UniqueId.randomId();
+    JobId jobId = JobId.fromInt(1);
     final String resourcePath = FileUtils.getTempDirectoryPath() + "/ray_test_resources";
     final String jobResourcePath = resourcePath + "/" + jobId.toString();
     File jobResourceDir = new File(jobResourcePath);