Support multiple core workers in one process (#7623)

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

@@ -42,6 +42,10 @@ public class TestProgressListener implements IInvokedMethodListener, ITestListen
   @Override
   public void onTestFailure(ITestResult result) {
     printInfo("TEST FAILURE", getFullTestName(result));
+    Throwable throwable = result.getThrowable();
+    if (throwable != null) {
+      throwable.printStackTrace();
+    }
   }
 
   @Override

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

@@ -1,11 +1,9 @@
 package org.ray.api;
 
-import com.google.common.base.Preconditions;
 import java.io.Serializable;
 import java.util.function.Supplier;
-import org.ray.api.runtime.RayRuntime;
-import org.ray.runtime.AbstractRayRuntime;
-import org.ray.runtime.RayMultiWorkerNativeRuntime;
+import org.ray.runtime.RayRuntimeInternal;
+import org.ray.runtime.RayRuntimeProxy;
 import org.ray.runtime.config.RunMode;
 import org.testng.Assert;
 import org.testng.SkipException;
@@ -79,12 +77,13 @@ public class TestUtils {
     Assert.assertEquals(obj.get(), "hi");
   }
 
-  public static AbstractRayRuntime getRuntime() {
-    RayRuntime runtime = Ray.internal();
-    if (runtime instanceof RayMultiWorkerNativeRuntime) {
-      runtime = ((RayMultiWorkerNativeRuntime) runtime).getCurrentRuntime();
-    }
-    Preconditions.checkState(runtime instanceof AbstractRayRuntime);
-    return (AbstractRayRuntime) runtime;
+  public static RayRuntimeInternal getRuntime() {
+    return (RayRuntimeInternal) Ray.internal();
+  }
+
+  public static RayRuntimeInternal getUnderlyingRuntime() {
+    RayRuntimeProxy proxy = (RayRuntimeProxy) (java.lang.reflect.Proxy
+        .getInvocationHandler(Ray.internal()));
+    return proxy.getRuntimeObject();
   }
 }

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

@@ -134,7 +134,7 @@ public class ClassLoaderTest extends BaseTest {
         FunctionDescriptor.class, Object[].class, ActorCreationOptions.class);
     createActorMethod.setAccessible(true);
     return (RayActor<?>) createActorMethod
-        .invoke(TestUtils.getRuntime(), functionDescriptor, new Object[0], null);
+        .invoke(TestUtils.getUnderlyingRuntime(), functionDescriptor, new Object[0], null);
   }
 
   private <T> RayObject<T> callActorFunction(RayActor<?> rayActor,
@@ -143,6 +143,6 @@ public class ClassLoaderTest extends BaseTest {
         BaseActor.class, FunctionDescriptor.class, Object[].class, int.class);
     callActorFunctionMethod.setAccessible(true);
     return (RayObject<T>) callActorFunctionMethod
-        .invoke(TestUtils.getRuntime(), rayActor, functionDescriptor, args, numReturns);
+        .invoke(TestUtils.getUnderlyingRuntime(), rayActor, functionDescriptor, args, numReturns);
   }
 }

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

@@ -29,7 +29,8 @@ public class ClientExceptionTest extends BaseTest {
       try {
         TimeUnit.SECONDS.sleep(1);
         // kill raylet
-        RunManager runManager = ((RayNativeRuntime) TestUtils.getRuntime()).getRunManager();
+        RunManager runManager =
+            ((RayNativeRuntime) TestUtils.getUnderlyingRuntime()).getRunManager();
         for (Process process : runManager.getProcesses("raylet")) {
           runManager.terminateProcess("raylet", process);
         }

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

@@ -14,6 +14,7 @@ import org.ray.api.RayActor;
 import org.ray.api.RayObject;
 import org.ray.api.TestUtils;
 import org.ray.api.WaitResult;
+import org.ray.api.exception.RayException;
 import org.ray.api.id.ActorId;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
@@ -135,6 +136,105 @@ public class MultiThreadingTest extends BaseTest {
     Assert.assertEquals(actorId, actorIdTester.getId());
   }
 
+  static boolean testMissingWrapRunnable() throws InterruptedException {
+    final RayObject<Integer> fooObject = Ray.put(1);
+    final RayActor<Echo> fooActor = Ray.createActor(Echo::new);
+    final Runnable[] runnables = new Runnable[]{
+        () -> Ray.put(1),
+        () -> Ray.get(fooObject.getId()),
+        fooObject::get,
+        () -> Ray.wait(ImmutableList.of(fooObject)),
+        Ray::getRuntimeContext,
+        () -> Ray.call(MultiThreadingTest::echo, 1),
+        () -> Ray.createActor(Echo::new),
+        () -> fooActor.call(Echo::echo, 1),
+    };
+
+    // It's OK to run them in main thread.
+    for (Runnable runnable : runnables) {
+      runnable.run();
+    }
+
+    Exception[] exception = new Exception[1];
+
+    Thread thread = new Thread(Ray.wrapRunnable(() -> {
+      try {
+        // It would be OK to run them in another thread if wrapped the runnable.
+        for (Runnable runnable : runnables) {
+          runnable.run();
+        }
+      } catch (Exception ex) {
+        exception[0] = ex;
+      }
+    }));
+    thread.start();
+    thread.join();
+    if (exception[0] != null) {
+      throw new RuntimeException("Exception occurred in thread.", exception[0]);
+    }
+
+    thread = new Thread(() -> {
+      try {
+        // It wouldn't be OK to run them in another thread if not wrapped the runnable.
+        for (Runnable runnable : runnables) {
+          Assert.expectThrows(RayException.class, runnable::run);
+        }
+      } catch (Exception ex) {
+        exception[0] = ex;
+      }
+    });
+    thread.start();
+    thread.join();
+    if (exception[0] != null) {
+      throw new RuntimeException("Exception occurred in thread.", exception[0]);
+    }
+
+    Runnable[] wrappedRunnables = new Runnable[runnables.length];
+    for (int i = 0; i < runnables.length; i++) {
+      wrappedRunnables[i] = Ray.wrapRunnable(runnables[i]);
+    }
+    // It would be OK to run the wrapped runnables in the current thread.
+    for (Runnable runnable : wrappedRunnables) {
+      runnable.run();
+    }
+
+    // It would be OK to invoke Ray APIs after executing a wrapped runnable in the current thread.
+    wrappedRunnables[0].run();
+    runnables[0].run();
+
+    // Return true here to make the Ray.call returns an RayObject.
+    return true;
+  }
+
+  @Test
+  public void testMissingWrapRunnableInDriver() throws InterruptedException {
+    testMissingWrapRunnable();
+  }
+
+  @Test
+  public void testMissingWrapRunnableInWorker() {
+    Ray.call(MultiThreadingTest::testMissingWrapRunnable).get();
+  }
+
+  @Test
+  public void testGetAndSetAsyncContext() throws InterruptedException {
+    Object asyncContext = Ray.getAsyncContext();
+    Exception[] exception = new Exception[1];
+    Thread thread = new Thread(() -> {
+      try {
+        Ray.setAsyncContext(asyncContext);
+        Ray.put(1);
+      } catch (Exception ex) {
+        exception[0] = ex;
+      }
+    });
+    thread.start();
+    thread.join();
+    if (exception[0] != null) {
+      throw new RuntimeException("Exception occurred in thread.", exception[0]);
+    }
+  }
+
   private static void runTestCaseInMultipleThreads(Runnable testCase, int numRepeats) {
     ExecutorService service = Executors.newFixedThreadPool(NUM_THREADS);