Improving message for access to unknown objects (#28209)

Signed-off-by: Pablo E <pabloem@apache.org>

When accessing unknown objects, Ray will kill the Python interpreter (without erroring out!). I've tried to rephrase the error to be a little easier to figure out and fix, and changed the check so that an error will be thrown rather than ending the whole process.

Author: pabloem

cpp/src/ray/runtime/abstract_ray_runtime.cc

@@ -141,7 +141,7 @@ std::vector<std::unique_ptr<::ray::TaskArg>> TransformArgs(
       auto owner_address = ray::rpc::Address{};
       if (ConfigInternal::Instance().run_mode == RunMode::CLUSTER) {
         auto &core_worker = CoreWorkerProcess::GetCoreWorker();
-        owner_address = core_worker.GetOwnerAddress(id);
+        owner_address = core_worker.GetOwnerAddressOrDie(id);
       }
       ray_arg = absl::make_unique<ray::TaskArgByReference>(id,
                                                            owner_address,

cpp/src/ray/test/cluster/cluster_mode_test.cc

@@ -73,6 +73,7 @@ TEST(RayClusterModeTest, FullTest) {
     config.redis_password_ = password;
   }
   ray::Init(config, cmd_argc, cmd_argv);
+
   /// put and get object
   auto obj = ray::Put(12345);
   auto get_result = *(ray::Get(obj));
@@ -88,9 +89,9 @@ TEST(RayClusterModeTest, FullTest) {
   EXPECT_EQ(1, task_result);
 
   /// common task with args
-  task_obj = ray::Task(Plus1).Remote(5);
-  task_result = *(ray::Get(task_obj));
-  EXPECT_EQ(6, task_result);
+  auto task_obj1 = ray::Task(Plus1).Remote(5);
+  auto task_result1 = *(ray::Get(task_obj1));
+  EXPECT_EQ(6, task_result1);
 
   ray::ActorHandle<Counter> actor = ray::Actor(RAY_FUNC(Counter::FactoryCreate))
                                         .SetMaxRestarts(1)
@@ -495,6 +496,10 @@ TEST(RayClusterModeTest, TaskWithPlacementGroup) {
 }
 
 TEST(RayClusterModeTest, NamespaceTest) {
+  if (ray::IsInitialized()) {
+    ray::Shutdown();
+  }
+  ray::Init();
   // Create a named actor in namespace `isolated_ns`.
   std::string actor_name_in_isolated_ns = "named_actor_in_isolated_ns";
   std::string isolated_ns_name = "isolated_ns";
@@ -516,11 +521,11 @@ TEST(RayClusterModeTest, NamespaceTest) {
 
   // Create a named actor in job default namespace.
   std::string actor_name_in_default_ns = "actor_name_in_default_ns";
-  actor = ray::Actor(RAY_FUNC(Counter::FactoryCreate))
-              .SetName(actor_name_in_default_ns)
-              .Remote();
-  initialized_obj = actor.Task(&Counter::Initialized).Remote();
-  EXPECT_TRUE(*initialized_obj.Get());
+  auto actor1 = ray::Actor(RAY_FUNC(Counter::FactoryCreate))
+                    .SetName(actor_name_in_default_ns)
+                    .Remote();
+  auto initialized_obj1 = actor1.Task(&Counter::Initialized).Remote();
+  EXPECT_TRUE(*initialized_obj1.Get());
   // It is visible to job default namespace.
   actor_optional = ray::GetActor<Counter>(actor_name_in_default_ns);
   EXPECT_TRUE(actor_optional);
@@ -534,6 +539,9 @@ TEST(RayClusterModeTest, GetNamespaceApiTest) {
   std::string ns = "test_get_current_namespace";
   ray::RayConfig config;
   config.ray_namespace = ns;
+  if (ray::IsInitialized()) {
+    ray::Shutdown();
+  }
   ray::Init(config, cmd_argc, cmd_argv);
   // Get namespace in driver.
   EXPECT_EQ(ray::GetNamespace(), ns);

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

@@ -6,6 +6,7 @@
 import io.ray.api.ActorHandle;
 import io.ray.api.ObjectRef;
 import io.ray.api.Ray;
+import io.ray.api.exception.RayException;
 import io.ray.api.id.ObjectId;
 import io.ray.runtime.object.NativeObjectStore;
 import io.ray.runtime.object.ObjectRefImpl;
@@ -112,9 +113,14 @@ private static void fillObjectStoreAndGet(
     if (succeed) {
       TestUtils.getRuntime().getObjectStore().getRaw(ImmutableList.of(objectId), Long.MAX_VALUE);
     } else {
-      List<Boolean> result =
-          TestUtils.getRuntime().getObjectStore().wait(ImmutableList.of(objectId), 1, 100, true);
-      Assert.assertFalse(result.get(0));
+      try {
+        List<Boolean> result =
+            TestUtils.getRuntime().getObjectStore().wait(ImmutableList.of(objectId), 0, 100, true);
+        Assert.fail(
+            "Ray did not fail when waiting for an object that does not belong in this session");
+      } catch (RayException e) {
+        // This is the expected outcome for succeed=false, as we wait for non-existent objects.
+      }
     }
   }
 
@@ -165,27 +171,7 @@ private static void sendSignal(ActorHandle<SignalActor> signal) {
   }
 
   /** Based on Python test case `test_dependency_refcounts`. */
-  public void testDependencyRefCounts() {
-    {
-      // Test that regular plasma dependency refcounts are decremented once the
-      // task finishes.
-      ActorHandle<SignalActor> signal = SignalActor.create();
-      ObjectRefImpl<TestUtils.LargeObject> largeDep =
-          (ObjectRefImpl<TestUtils.LargeObject>) Ray.put(new TestUtils.LargeObject());
-      ObjectRefImpl<Object> result =
-          (ObjectRefImpl<Object>)
-              Ray.<TestUtils.LargeObject, ActorHandle<SignalActor>, Object>task(
-                      ReferenceCountingTest::oneDep, largeDep, signal)
-                  .remote();
-      checkRefCounts(largeDep.getId(), 1, 1, result.getId(), 1, 0);
-      sendSignal(signal);
-      // Reference count should be removed once the task finishes.
-      checkRefCounts(largeDep.getId(), 1, 0, result.getId(), 1, 0);
-      del(largeDep);
-      del(result);
-      checkRefCounts(ImmutableMap.of());
-    }
-
+  public void testDependencyRefCounts1() {
     {
       // Test that inlined dependency refcounts are decremented once they are
       // inlined.
@@ -207,7 +193,9 @@ public void testDependencyRefCounts() {
       del(result);
       checkRefCounts(ImmutableMap.of());
     }
+  }
 
+  public void testDependencyRefCounts2() {
     {
       // Test that spilled plasma dependency refcounts are decremented once
       // the task finishes.
@@ -236,7 +224,9 @@ public void testDependencyRefCounts() {
       del(result);
       checkRefCounts(ImmutableMap.of());
     }
+  }
 
+  public void testDependencyRefCounts3() {
     {
       // Test that regular plasma dependency refcounts are decremented if a task
       // fails.
@@ -257,7 +247,9 @@ public void testDependencyRefCounts() {
       del(result);
       checkRefCounts(ImmutableMap.of());
     }
+  }
 
+  public void testDependencyRefCounts4() {
     {
       // Test that spilled plasma dependency refcounts are decremented if a task
       // fails.
@@ -288,6 +280,28 @@ public void testDependencyRefCounts() {
     }
   }
 
+  public void testDependencyRefCounts5() {
+    {
+      // Test that regular plasma dependency refcounts are decremented once the
+      // task finishes.
+      ActorHandle<SignalActor> signal = SignalActor.create();
+      ObjectRefImpl<TestUtils.LargeObject> largeDep =
+          (ObjectRefImpl<TestUtils.LargeObject>) Ray.put(new TestUtils.LargeObject());
+      ObjectRefImpl<Object> result =
+          (ObjectRefImpl<Object>)
+              Ray.<TestUtils.LargeObject, ActorHandle<SignalActor>, Object>task(
+                      ReferenceCountingTest::oneDep, largeDep, signal)
+                  .remote();
+      checkRefCounts(largeDep.getId(), 1, 1, result.getId(), 1, 0);
+      sendSignal(signal);
+      // Reference count should be removed once the task finishes.
+      checkRefCounts(largeDep.getId(), 1, 0, result.getId(), 1, 0);
+      del(largeDep);
+      del(result);
+      checkRefCounts(ImmutableMap.of());
+    }
+  }
+
   private static int fooBasicPinning(Object arg) {
     return 0;
   }

python/ray/_raylet.pyx

@@ -198,6 +198,10 @@ cdef int check_status(const CRayStatus& status) nogil except -1:
         raise GetTimeoutError(message)
     elif status.IsNotFound():
         raise ValueError(message)
+    elif status.IsObjectNotFound():
+        raise ValueError(message)
+    elif status.IsObjectUnknownOwner():
+        raise ValueError(message)
     else:
         raise RaySystemError(message)
 
@@ -421,6 +425,8 @@ cdef prepare_args_internal(
         c_vector[CObjectID] inlined_ids
         c_string put_arg_call_site
         c_vector[CObjectReference] inlined_refs
+        CAddress c_owner_address
+        CRayStatus op_status
 
     worker = ray._private.worker.global_worker
     put_threshold = RayConfig.instance().max_direct_call_object_size()
@@ -429,11 +435,13 @@ cdef prepare_args_internal(
     for arg in args:
         if isinstance(arg, ObjectRef):
             c_arg = (<ObjectRef>arg).native()
+            op_status = CCoreWorkerProcess.GetCoreWorker().GetOwnerAddress(
+                    c_arg, &c_owner_address)
+            check_status(op_status)
             args_vector.push_back(
                 unique_ptr[CTaskArg](new CTaskArgByReference(
                     c_arg,
-                    CCoreWorkerProcess.GetCoreWorker().GetOwnerAddress(
-                        c_arg),
+                    c_owner_address,
                     arg.call_site())))
 
         else:
@@ -1565,8 +1573,9 @@ cdef class CoreWorker:
             CTaskID c_task_id = current_task_id.native()
             c_vector[CObjectID] c_object_ids = ObjectRefsToVector(object_refs)
         with nogil:
-            check_status(CCoreWorkerProcess.GetCoreWorker().Get(
-                c_object_ids, timeout_ms, &results))
+            op_status = CCoreWorkerProcess.GetCoreWorker().Get(
+                c_object_ids, timeout_ms, &results)
+        check_status(op_status)
 
         return RayObjectsToDataMetadataPairs(results)
 
@@ -1770,8 +1779,9 @@ cdef class CoreWorker:
 
         wait_ids = ObjectRefsToVector(object_refs)
         with nogil:
-            check_status(CCoreWorkerProcess.GetCoreWorker().Wait(
-                wait_ids, num_returns, timeout_ms, &results, fetch_local))
+            op_status = CCoreWorkerProcess.GetCoreWorker().Wait(
+                wait_ids, num_returns, timeout_ms, &results, fetch_local)
+        check_status(op_status)
 
         assert len(results) == len(object_refs)
 
@@ -2322,16 +2332,20 @@ cdef class CoreWorker:
     def get_owner_address(self, ObjectRef object_ref):
         cdef:
             CObjectID c_object_id = object_ref.native()
-        return CCoreWorkerProcess.GetCoreWorker().GetOwnerAddress(
-                c_object_id).SerializeAsString()
+            CAddress c_owner_address
+        op_status = CCoreWorkerProcess.GetCoreWorker().GetOwnerAddress(
+                c_object_id, &c_owner_address)
+        check_status(op_status)
+        return c_owner_address.SerializeAsString()
 
     def serialize_object_ref(self, ObjectRef object_ref):
         cdef:
             CObjectID c_object_id = object_ref.native()
             CAddress c_owner_address = CAddress()
             c_string serialized_object_status
-        CCoreWorkerProcess.GetCoreWorker().GetOwnershipInfo(
+        op_status = CCoreWorkerProcess.GetCoreWorker().GetOwnershipInfo(
                 c_object_id, &c_owner_address, &serialized_object_status)
+        check_status(op_status)
         return (object_ref,
                 c_owner_address.SerializeAsString(),
                 serialized_object_status)

python/ray/includes/common.pxd

@@ -113,7 +113,9 @@ cdef extern from "ray/common/status.h" namespace "ray" nogil:
         c_bool IsTimedOut()
         c_bool IsInterrupted()
         c_bool ShouldExitWorker()
+        c_bool IsObjectNotFound()
         c_bool IsNotFound()
+        c_bool IsObjectUnknownOwner()
 
         c_string ToString()
         c_string CodeAsString()

python/ray/includes/libcoreworker.pxd

@@ -185,13 +185,14 @@ cdef extern from "ray/core_worker/core_worker.h" nogil:
         void PutObjectIntoPlasma(const CRayObject &object,
                                  const CObjectID &object_id)
         const CAddress &GetRpcAddress() const
-        CAddress GetOwnerAddress(const CObjectID &object_id) const
+        CRayStatus GetOwnerAddress(const CObjectID &object_id,
+                                   CAddress *owner_address) const
         c_vector[CObjectReference] GetObjectRefs(
                 const c_vector[CObjectID] &object_ids) const
 
-        void GetOwnershipInfo(const CObjectID &object_id,
-                              CAddress *owner_address,
-                              c_string *object_status)
+        CRayStatus GetOwnershipInfo(const CObjectID &object_id,
+                                    CAddress *owner_address,
+                                    c_string *object_status)
         void RegisterOwnershipInfoAndResolveFuture(
                 const CObjectID &object_id,
                 const CObjectID &outer_object_id,

python/ray/tests/BUILD

@@ -58,7 +58,6 @@ py_test_module_list(
     "test_get_locations.py",
     "test_global_state.py",
     "test_healthcheck.py",
-    "test_ray_shutdown.py",
     "test_kill_raylet_signal_log.py",
     "test_memstat.py",
     "test_protobuf_compatibility.py"
@@ -128,6 +127,7 @@ py_test_module_list(
     "test_placement_group_failover.py",
     "test_ray_init.py",
     "test_ray_init_2.py",
+    "test_ray_shutdown.py",
     "test_resource_demand_scheduler.py",
     "test_resource_metrics.py",
     "test_runtime_context.py",

python/ray/tests/test_basic.py

@@ -650,8 +650,9 @@ def h():
 
     assert ray.get(f.remote()) == (1, 2)
 
-    # Test a remote function that recursively calls itself.
 
+def test_recursive_remote_call(ray_start_shared_local_modes):
+    # Test a remote function that recursively calls itself.
     @ray.remote
     def factorial(n):
         if n == 0: