[Java] Refine tests and fix single-process mode (#4265)

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

@@ -28,4 +28,9 @@ public class RayDevRuntime extends AbstractRayRuntime {
   public MockObjectStore getObjectStore() {
     return store;
   }
+
+  @Override
+  public Worker getWorker() {
+    return ((MockRayletClient) rayletClient).getCurrentWorker();
+  }
 }

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

@@ -4,7 +4,6 @@ import com.google.common.base.Preconditions;
 import org.ray.api.RuntimeContext;
 import org.ray.api.id.UniqueId;
 import org.ray.runtime.config.RunMode;
-import org.ray.runtime.config.WorkerMode;
 import org.ray.runtime.task.TaskSpec;
 
 public class RuntimeContextImpl implements RuntimeContext {
@@ -22,8 +21,10 @@ public class RuntimeContextImpl implements RuntimeContext {
 
   @Override
   public UniqueId getCurrentActorId() {
-    Preconditions.checkState(runtime.rayConfig.workerMode == WorkerMode.WORKER);
-    return runtime.getWorker().getCurrentActorId();
+    Worker worker = runtime.getWorker();
+    Preconditions.checkState(worker != null && !worker.getCurrentActorId().isNil(),
+        "This method should only be called from an actor.");
+    return worker.getCurrentActorId();
   }
 
   @Override

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

@@ -79,7 +79,6 @@ public class Worker {
    * Execute a task.
    */
   public void execute(TaskSpec spec) {
-    LOGGER.info("Executing task {}", spec.taskId);
     LOGGER.debug("Executing task {}", spec);
     UniqueId returnId = spec.returnIds[0];
     ClassLoader oldLoader = Thread.currentThread().getContextClassLoader();
@@ -123,7 +122,7 @@ public class Worker {
         maybeLoadCheckpoint(result, returnId);
         currentActor = result;
       }
-      LOGGER.info("Finished executing task {}", spec.taskId);
+      LOGGER.debug("Finished executing task {}", spec.taskId);
     } catch (Exception e) {
       LOGGER.error("Error executing task " + spec, e);
       if (!spec.isActorCreationTask()) {

java/runtime/src/main/java/org/ray/runtime/objectstore/MockObjectStore.java

@@ -95,11 +95,10 @@ public class MockObjectStore implements ObjectStoreLink {
     ArrayList<ObjectStoreData> rets = new ArrayList<>();
     for (byte[] id : objectIds) {
       try {
-        Constructor<ObjectStoreData> constructor = ObjectStoreData.class.getConstructor(
-                byte[].class, byte[].class);
+        Constructor<?> constructor = ObjectStoreData.class.getDeclaredConstructors()[0];
         constructor.setAccessible(true);
-        rets.add(constructor.newInstance(metadata.get(new UniqueId(id)),
-                data.get(new UniqueId(id))));
+        rets.add((ObjectStoreData) constructor.newInstance(data.get(new UniqueId(id)),
+                metadata.get(new UniqueId(id))));
       } catch (Exception e) {
         throw new RuntimeException(e);
       }

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

@@ -39,6 +39,7 @@ public class MockRayletClient implements RayletClient {
   private final ExecutorService exec;
   private final Deque<Worker> idleWorkers;
   private final Map<UniqueId, Worker> actorWorkers;
+  private final ThreadLocal<Worker> currentWorker;
 
   public MockRayletClient(RayDevRuntime runtime, int numberThreads) {
     this.runtime = runtime;
@@ -48,6 +49,7 @@ public class MockRayletClient implements RayletClient {
     exec = Executors.newFixedThreadPool(numberThreads);
     idleWorkers = new LinkedList<>();
     actorWorkers = new HashMap<>();
+    currentWorker = new ThreadLocal<>();
   }
 
   public synchronized void onObjectPut(UniqueId id) {
@@ -60,22 +62,28 @@ public class MockRayletClient implements RayletClient {
     }
   }
 
+  public Worker getCurrentWorker() {
+    return currentWorker.get();
+  }
+
   /**
    * Get a worker from the worker pool to run the given task.
    */
   private Worker getWorker(TaskSpec task) {
-    if (task.isActorTask()) {
-      return actorWorkers.get(task.actorId);
-    }
     Worker worker;
-    if (idleWorkers.size() > 0) {
-      worker = idleWorkers.pop();
+    if (task.isActorTask()) {
+      worker = actorWorkers.get(task.actorId);
     } else {
-      worker = new Worker(runtime);
-    }
-    if (task.isActorCreationTask()) {
-      actorWorkers.put(task.actorCreationId, worker);
+      if (idleWorkers.size() > 0) {
+        worker = idleWorkers.pop();
+      } else {
+        worker = new Worker(runtime);
+      }
+      if (task.isActorCreationTask()) {
+        actorWorkers.put(task.actorCreationId, worker);
+      }
     }
+    currentWorker.set(worker);
     return worker;
   }
 
@@ -83,6 +91,7 @@ public class MockRayletClient implements RayletClient {
    * Return the worker to the worker pool.
    */
   private void returnWorker(Worker worker) {
+    currentWorker.remove();
     idleWorkers.push(worker);
   }
 
@@ -105,9 +114,7 @@ public class MockRayletClient implements RayletClient {
                     new byte[]{}, new byte[]{});
           }
         } finally {
-          if (!task.isActorCreationTask() && !task.isActorTask()) {
-            returnWorker(worker);
-          }
+          returnWorker(worker);
         }
       });
     } else {

java/runtime/src/main/resources/ray.default.conf

@@ -100,7 +100,7 @@ ray {
   // ----------------------------
   dev-runtime {
     // Number of threads that you process tasks
-    execution-parallelism: 5
+    execution-parallelism: 10
   }
 
 }

java/test.sh

@@ -2,36 +2,25 @@
 
 # Cause the script to exit if a single command fails.
 set -e
-
 # Show explicitly which commands are currently running.
 set -x
 
 ROOT_DIR=$(cd "$(dirname "${BASH_SOURCE:-$0}")"; pwd)
-$ROOT_DIR/../build.sh -l java
 
 pushd $ROOT_DIR/../java
+echo "Compiling Java code."
 mvn clean install -Dmaven.test.skip
-check_style=$(mvn checkstyle:check)
-echo "${check_style}"
-[[ ${check_style} =~ "BUILD FAILURE" ]] && exit 1
 
-# test raylet
-mvn test | tee mvn_test
-if [ `grep -c "BUILD FAILURE" mvn_test` -eq '0' ]; then
-    rm mvn_test
-    echo "Tests passed under CLUSTER mode!"
-else
-    rm mvn_test
-    exit 1
-fi
-# test raylet under SINGLE_PROCESS mode
-mvn test -Dray.run-mode=SINGLE_PROCESS | tee dev_mvn_test
-if [ `grep -c "BUILD FAILURE" dev_mvn_test` -eq '0' ]; then
-    rm dev_mvn_test
-    echo "Tests passed under SINGLE_PROCESS mode!"
-else
-    rm dev_mvn_test
-    exit 1
-fi
+echo "Checking code format."
+mvn checkstyle:check
+
+echo "Running tests under cluster mode."
+ENABLE_MULTI_LANGUAGE_TESTS=1 mvn test
+
+echo "Running tests under single-process mode."
+mvn test -Dray.run-mode=SINGLE_PROCESS
+
+set +x
+set +e
 
 popd

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

@@ -9,7 +9,7 @@ public class TestUtils {
   public static void skipTestUnderSingleProcess() {
     AbstractRayRuntime runtime = (AbstractRayRuntime)Ray.internal();
     if (runtime.getRayConfig().runMode == RunMode.SINGLE_PROCESS) {
-      throw new SkipException("Skip case.");
+      throw new SkipException("This test doesn't work under single-process mode.");
     }
   }
 }

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

@@ -44,13 +44,9 @@ public class ActorReconstructionTest extends BaseTest {
     }
   }
 
-  @Override
-  public void beforeEachCase() {
-    TestUtils.skipTestUnderSingleProcess();
-  }
-
   @Test
   public void testActorReconstruction() throws InterruptedException, IOException {
+    TestUtils.skipTestUnderSingleProcess();
     ActorCreationOptions options = new ActorCreationOptions(new HashMap<>(), 1);
     RayActor<Counter> actor = Ray.createActor(Counter::new, options);
     // Call increase 3 times.
@@ -130,6 +126,8 @@ public class ActorReconstructionTest extends BaseTest {
 
   @Test
   public void testActorCheckpointing() throws IOException, InterruptedException {
+    TestUtils.skipTestUnderSingleProcess();
+
     ActorCreationOptions options = new ActorCreationOptions(new HashMap<>(), 1);
     RayActor<CheckpointableCounter> actor = Ray.createActor(CheckpointableCounter::new, options);
     // Call increase 3 times.
@@ -138,8 +136,6 @@ public class ActorReconstructionTest extends BaseTest {
     }
     // Assert that the actor wasn't resumed from a checkpoint.
     Assert.assertFalse(Ray.call(CheckpointableCounter::wasResumedFromCheckpoint, actor).get());
-
-    // Kill the actor process.
     int pid = Ray.call(CheckpointableCounter::getPid, actor).get();
     Runtime.getRuntime().exec("kill -9 " + pid);
     // Wait for the actor to be killed.

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

@@ -5,6 +5,7 @@ import java.util.concurrent.TimeUnit;
 import org.ray.api.Ray;
 import org.ray.api.RayActor;
 import org.ray.api.RayObject;
+import org.ray.api.TestUtils;
 import org.ray.api.annotation.RayRemote;
 import org.ray.api.exception.UnreconstructableException;
 import org.ray.api.id.UniqueId;
@@ -90,6 +91,7 @@ public class ActorTest extends BaseTest {
 
   @Test
   public void testUnreconstructableActorObject() throws InterruptedException {
+    TestUtils.skipTestUnderSingleProcess();
     RayActor<Counter> counter = Ray.createActor(Counter::new, 100);
     // Call an actor method.
     RayObject value = Ray.call(Counter::getValue, counter);

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

@@ -1,27 +1,31 @@
 package org.ray.api.test;
 
+import java.lang.reflect.Method;
 import org.ray.api.Ray;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
 import org.testng.annotations.AfterMethod;
 import org.testng.annotations.BeforeMethod;
 
 public class BaseTest {
 
+  private static final Logger LOGGER = LoggerFactory.getLogger(BaseTest.class);
+
   @BeforeMethod
-  public void setUp() {
+  public void setUpBase(Method method) {
+    LOGGER.info("===== Running test: "
+        + method.getDeclaringClass().getName() + "." + method.getName());
     System.setProperty("ray.home", "../..");
     System.setProperty("ray.resources", "CPU:4,RES-A:4");
-    beforeInitRay();
     Ray.init();
-    beforeEachCase();
   }
 
   @AfterMethod
-  public void tearDown() {
+  public void tearDownBase() {
     // TODO(qwang): This is double check to check that the socket file is removed actually.
     // We could not enable this until `systemInfo` enabled.
     //File rayletSocketFIle = new File(Ray.systemInfo().rayletSocketName());
     Ray.shutdown();
-    afterShutdownRay();
 
     //remove raylet socket file
     //rayletSocketFIle.delete();
@@ -31,15 +35,4 @@ public class BaseTest {
     System.clearProperty("ray.resources");
   }
 
-  protected void beforeInitRay() {
-
-  }
-
-  protected void afterShutdownRay() {
-
-  }
-
-  protected void beforeEachCase() {
-
-  }
 }

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

@@ -17,13 +17,9 @@ public class ClientExceptionTest extends BaseTest {
 
   private static final Logger LOGGER = LoggerFactory.getLogger(ClientExceptionTest.class);
 
-  @Override
-  public void beforeEachCase() {
-    TestUtils.skipTestUnderSingleProcess();
-  }
-
   @Test
   public void testWaitAndCrash() {
+    TestUtils.skipTestUnderSingleProcess();
     UniqueId randomId = UniqueId.randomId();
     RayObject<String> notExisting = new RayObjectImpl(randomId);
 

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

@@ -55,30 +55,29 @@ public class FailureTest extends BaseTest {
     }
   }
 
-  @Override
-  public void beforeEachCase() {
-    TestUtils.skipTestUnderSingleProcess();
-  }
-
   @Test
   public void testNormalTaskFailure() {
+    TestUtils.skipTestUnderSingleProcess();
     assertTaskFailedWithRayTaskException(Ray.call(FailureTest::badFunc));
   }
 
   @Test
   public void testActorCreationFailure() {
+    TestUtils.skipTestUnderSingleProcess();
     RayActor<BadActor> actor = Ray.createActor(BadActor::new, true);
     assertTaskFailedWithRayTaskException(Ray.call(BadActor::badMethod, actor));
   }
 
   @Test
   public void testActorTaskFailure() {
+    TestUtils.skipTestUnderSingleProcess();
     RayActor<BadActor> actor = Ray.createActor(BadActor::new, false);
     assertTaskFailedWithRayTaskException(Ray.call(BadActor::badMethod, actor));
   }
 
   @Test
   public void testWorkerProcessDying() {
+    TestUtils.skipTestUnderSingleProcess();
     try {
       Ray.call(FailureTest::badFunc2).get();
       Assert.fail("This line shouldn't be reached.");
@@ -90,6 +89,7 @@ public class FailureTest extends BaseTest {
 
   @Test
   public void testActorProcessDying() {
+    TestUtils.skipTestUnderSingleProcess();
     RayActor<BadActor> actor = Ray.createActor(BadActor::new, false);
     try {
       Ray.call(BadActor::badMethod2, actor).get();

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

@@ -2,7 +2,7 @@ package org.ray.api.test;
 
 import com.google.common.collect.ImmutableList;
 import java.io.File;
-import java.lang.ProcessBuilder.Redirect;
+import java.lang.reflect.Method;
 import java.util.List;
 import java.util.concurrent.TimeUnit;
 import org.ray.api.Ray;
@@ -33,13 +33,13 @@ public class MultiLanguageClusterTest {
 
   /**
    * Execute an external command.
+   *
    * @return Whether the command succeeded.
    */
   private boolean executeCommand(List<String> command, int waitTimeoutSeconds) {
     try {
       LOGGER.info("Executing command: {}", String.join(" ", command));
-      Process process = new ProcessBuilder(command).redirectOutput(Redirect.INHERIT)
-          .redirectError(Redirect.INHERIT).start();
+      Process process = new ProcessBuilder(command).inheritIO().start();
       process.waitFor(waitTimeoutSeconds, TimeUnit.SECONDS);
       return process.exitValue() == 0;
     } catch (Exception e) {
@@ -48,11 +48,12 @@ public class MultiLanguageClusterTest {
   }
 
   @BeforeMethod
-  public void setUp() {
-    // Check whether 'ray' command is installed.
-    boolean rayCommandExists = executeCommand(ImmutableList.of("which", "ray"), 5);
-    if (!rayCommandExists) {
-      throw new SkipException("Skipping test, because ray command doesn't exist.");
+  public void setUp(Method method) {
+    String testName = method.getName();
+    if (!"1".equals(System.getenv("ENABLE_MULTI_LANGUAGE_TESTS"))) {
+      LOGGER.info("Skip " + testName +
+          " because env variable ENABLE_MULTI_LANGUAGE_TESTS isn't set");
+      throw new SkipException("Skip test.");
     }
 
     // Delete existing socket files.
@@ -64,15 +65,20 @@ public class MultiLanguageClusterTest {
     }
 
     // Start ray cluster.
+    String testDir = System.getProperty("user.dir");
+    String workerOptions = String.format("-Dray.home=%s/../../", testDir);
+    workerOptions +=
+        " -classpath " + String.format("%s/../../build/java/*:%s/target/*", testDir, testDir);
     final List<String> startCommand = ImmutableList.of(
         "ray",
         "start",
         "--head",
         "--redis-port=6379",
-        "--include-java",
         String.format("--plasma-store-socket-name=%s", PLASMA_STORE_SOCKET_NAME),
         String.format("--raylet-socket-name=%s", RAYLET_SOCKET_NAME),
-        "--java-worker-options=-classpath ../../build/java/*:../../java/test/target/*"
+        "--load-code-from-local",
+        "--include-java",
+        "--java-worker-options=" + workerOptions
     );
     if (!executeCommand(startCommand, 10)) {
       throw new RuntimeException("Couldn't start ray cluster.");

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

@@ -12,6 +12,7 @@ import java.util.concurrent.TimeUnit;
 import org.ray.api.Ray;
 import org.ray.api.RayActor;
 import org.ray.api.RayObject;
+import org.ray.api.TestUtils;
 import org.ray.api.WaitResult;
 import org.ray.api.annotation.RayRemote;
 import org.testng.Assert;
@@ -73,11 +74,15 @@ public class MultiThreadingTest extends BaseTest {
 
   @Test
   public void testInDriver() {
+    // TODO(hchen): Fix this test under single-process mode.
+    TestUtils.skipTestUnderSingleProcess();
     testMultiThreading();
   }
 
   @Test
   public void testInWorker() {
+    // Single-process mode doesn't have real workers.
+    TestUtils.skipTestUnderSingleProcess();
     RayObject<String> obj = Ray.call(MultiThreadingTest::testMultiThreading);
     Assert.assertEquals("ok", obj.get());
   }

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

@@ -4,6 +4,7 @@ import org.apache.arrow.plasma.PlasmaClient;
 import org.apache.arrow.plasma.exceptions.DuplicateObjectException;
 
 import org.ray.api.Ray;
+import org.ray.api.TestUtils;
 import org.ray.api.id.UniqueId;
 import org.ray.runtime.AbstractRayRuntime;
 import org.testng.Assert;
@@ -13,6 +14,7 @@ public class PlasmaStoreTest extends BaseTest {
 
   @Test
   public void testPutWithDuplicateId() {
+    TestUtils.skipTestUnderSingleProcess();
     UniqueId objectId = UniqueId.randomId();
     AbstractRayRuntime runtime = (AbstractRayRuntime) Ray.internal();
     PlasmaClient store = new PlasmaClient(runtime.getRayConfig().objectStoreSocketName, "", 0);

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

@@ -4,18 +4,20 @@ import org.ray.api.Ray;
 import org.ray.api.RayObject;
 import org.ray.api.annotation.RayRemote;
 import org.testng.Assert;
+import org.testng.annotations.AfterClass;
+import org.testng.annotations.BeforeClass;
 import org.testng.annotations.Test;
 
 public class RedisPasswordTest extends BaseTest {
 
-  @Override
-  public void beforeInitRay() {
+  @BeforeClass
+  public void setUp() {
     System.setProperty("ray.redis.head-password", "12345678");
     System.setProperty("ray.redis.password", "12345678");
   }
 
-  @Override
-  public void afterShutdownRay() {
+  @AfterClass
+  public void tearDown() {
     System.clearProperty("ray.redis.head-password");
     System.clearProperty("ray.redis.password");
   }

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

@@ -25,18 +25,15 @@ public class ResourcesManagementTest extends BaseTest {
 
   @RayRemote
   public static class Echo {
+
     public Integer echo(Integer number) {
       return number;
     }
   }
 
-  @Override
-  public void beforeEachCase() {
-    TestUtils.skipTestUnderSingleProcess();
-  }
-
   @Test
   public void testMethods() {
+    TestUtils.skipTestUnderSingleProcess();
     CallOptions callOptions1 = new CallOptions(ImmutableMap.of("CPU", 4.0, "GPU", 0.0));
 
     // This is a case that can satisfy required resources.
@@ -57,6 +54,7 @@ public class ResourcesManagementTest extends BaseTest {
 
   @Test
   public void testActors() {
+    TestUtils.skipTestUnderSingleProcess();
 
     ActorCreationOptions actorCreationOptions1 =
         new ActorCreationOptions(ImmutableMap.of("CPU", 2.0, "GPU", 0.0));

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

@@ -5,6 +5,8 @@ import org.ray.api.RayActor;
 import org.ray.api.annotation.RayRemote;
 import org.ray.api.id.UniqueId;
 import org.testng.Assert;
+import org.testng.annotations.AfterClass;
+import org.testng.annotations.BeforeClass;
 import org.testng.annotations.Test;
 
 public class RuntimeContextTest extends BaseTest {
@@ -14,13 +16,20 @@ public class RuntimeContextTest extends BaseTest {
   private static String RAYLET_SOCKET_NAME = "/tmp/ray/test/raylet_socket";
   private static String OBJECT_STORE_SOCKET_NAME = "/tmp/ray/test/object_store_socket";
 
-  @Override
-  public void beforeInitRay() {
+  @BeforeClass
+  public void setUp() {
     System.setProperty("ray.driver.id", DRIVER_ID.toString());
     System.setProperty("ray.raylet.socket-name", RAYLET_SOCKET_NAME);
     System.setProperty("ray.object-store.socket-name", OBJECT_STORE_SOCKET_NAME);
   }
 
+  @AfterClass
+  public void tearDown() {
+    System.clearProperty("ray.driver.id");
+    System.clearProperty("ray.raylet.socket-name");
+    System.clearProperty("ray.object-store.socket-name");
+  }
+
   @Test
   public void testRuntimeContextInDriver() {
     Assert.assertEquals(DRIVER_ID, Ray.getRuntimeContext().getCurrentDriverId());

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

@@ -17,13 +17,9 @@ public class StressTest extends BaseTest {
     return x;
   }
 
-  @Override
-  public void beforeEachCase() {
-    TestUtils.skipTestUnderSingleProcess();
-  }
-
   @Test
   public void testSubmittingTasks() {
+    TestUtils.skipTestUnderSingleProcess();
     for (int numIterations : ImmutableList.of(1, 10, 100, 1000)) {
       int numTasks = 1000 / numIterations;
       for (int i = 0; i < numIterations; i++) {
@@ -40,6 +36,7 @@ public class StressTest extends BaseTest {
 
   @Test
   public void testDependency() {
+    TestUtils.skipTestUnderSingleProcess();
     RayObject<Integer> x = Ray.call(StressTest::echo, 1);
     for (int i = 0; i < 1000; i++) {
       x = Ray.call(StressTest::echo, x);
@@ -77,6 +74,7 @@ public class StressTest extends BaseTest {
 
   @Test
   public void testSubmittingManyTasksToOneActor() {
+    TestUtils.skipTestUnderSingleProcess();
     RayActor<Actor> actor = Ray.createActor(Actor::new);
     List<UniqueId> objectIds = new ArrayList<>();
     for (int i = 0; i < 10; i++) {
@@ -90,6 +88,7 @@ public class StressTest extends BaseTest {
 
   @Test
   public void testPuttingAndGettingManyObjects() {
+    TestUtils.skipTestUnderSingleProcess();
     Integer objectToPut = 1;
     List<RayObject<Integer>> objects = new ArrayList<>();
     for (int i = 0; i < 100_000; i++) {