[xray] Track ray.get calls as task dependencies (#2362)

java/runtime-common/src/main/java/org/ray/core/WorkerContext.java

@@ -2,6 +2,7 @@ package org.ray.core;
 
 import org.ray.api.UniqueID;
 import org.ray.core.model.RayParameters;
+import org.ray.core.model.WorkerMode;
 import org.ray.spi.model.TaskSpec;
 
 public class WorkerContext {
@@ -35,7 +36,11 @@ public class WorkerContext {
 
     TaskSpec dummy = new TaskSpec();
     dummy.parentTaskId = UniqueID.nil;
-    dummy.taskId = UniqueID.nil;
+    if (params.worker_mode == WorkerMode.DRIVER) {
+      dummy.taskId = UniqueID.randomId();
+    } else {
+      dummy.taskId = UniqueID.nil;
+    }
     dummy.actorId = UniqueID.nil;
     dummy.driverId = params.driver_id;
     prepare(dummy, null);

java/runtime-native/src/main/java/org/ray/core/impl/RayNativeRuntime.java

@@ -109,6 +109,7 @@ public class RayNativeRuntime extends RayRuntime {
           WorkerContext.currentWorkerId(),
           UniqueID.nil,
           isWorker,
+          WorkerContext.currentTask().taskId,
           0
       );
 
@@ -237,4 +238,4 @@ public class RayNativeRuntime extends RayRuntime {
       throw new TaskExecutionException(log, e);
     }
   }
-}
+}

java/runtime-native/src/main/java/org/ray/spi/impl/DefaultLocalSchedulerClient.java

@@ -26,13 +26,13 @@ public class DefaultLocalSchedulerClient implements LocalSchedulerLink {
   private long client = 0;
 
   public DefaultLocalSchedulerClient(String schedulerSockName, UniqueID clientId, UniqueID actorId,
-                                     boolean isWorker, long numGpus) {
+                                     boolean isWorker, UniqueID driverId, long numGpus) {
     client = _init(schedulerSockName, clientId.getBytes(), actorId.getBytes(), isWorker,
-        numGpus);
+        driverId.getBytes(), numGpus);
   }
 
   private static native long _init(String localSchedulerSocket, byte[] workerId, byte[] actorId,
-                                   boolean isWorker, long numGpus);
+                                   boolean isWorker, byte[] driverTaskId, long numGpus);
 
   private static native byte[] _computePutId(long client, byte[] taskId, int putIndex);