[Core] Fix segfault when cancel and generator is used with a high load

[rkooo567]

Why are these changes needed?

Right now, future.result() returns and rasie an exception when future.cancel() is called although the coroutine itself is actually not cancelled. I verified this from a local example.

This means the future.result can return and the task states are cleaned up (since task is cancelled) while we are running the actual coroutine & cpp code in a different thread. This causes a random sigabrt and sigsegv when there are lots of cancel & generator tasks are running.

For a generator, we are using a custom await between generator for the perf reason:

ray/python/ray/_raylet.pyx

Line 1222 in f11fa84

done = await loop.run_in_executor(

. The cpp func running in a separte thread kept running while the task is cleaned up, which causes the segfault.

The same thing can technically happen for a regular async task (if it uses thread pool). If it doesn't use a thread pool, it doesn't happen because the future.cancel is not scheduled until the coroutine is awaited (whereas for a generator, it always have await between each run because of the threadpool above).

I fixed the issue by waiting until the coroutine is finished. Note that currently it is a blocking call, but it will mostly be a short time because when the cancellation error is raised the coroutine is already cancelled.

Related issue number

Closes #39703

Checks

• I've signed off every commit(by using the -s flag, i.e., git commit -s) in this PR.
• I've run scripts/format.sh to lint the changes in this PR.
• I've included any doc changes needed for https://docs.ray.io/en/master/.
  • I've added any new APIs to the API Reference. For example, if I added a
    method in Tune, I've added it in doc/source/tune/api/ under the
    corresponding .rst file.
• I've made sure the tests are passing. Note that there might be a few flaky tests, see the recent failures at https://flakey-tests.ray.io/
• Testing Strategy
  • Unit tests
  • Release tests
  • This PR is not tested :(

[jjyao]

Right now, future.result() returns and rasie an exception when future.cancel() is called although the coroutine itself is actually not cancelled. I verified this from a local example.

Is this behavior documented by Python?

[rkooo567]

It is not documented, but a relevant issue; python/cpython#105836

[edoakes]

remove?

[edoakes]

This looks extremely sketchy and blocking the loop makes me uncomfortable.

We are already using an event to signal the completion of the coroutine inside of run_async_func_or_coro_in_event_loop:

ray/python/ray/_raylet.pyx

Line 4298 in 9d4ae4e

future.add_done_callback(lambda _: event.Notify())

I assume the current issue is that the callback added with add_done_callback is running prior to the coroutine actually completing. If that's the case, then how about we move event.Notify() into a finally block in the async_func just above?

I believe that'll accomplish the same thing without requiring the blocking wait() call.

[edoakes]

this and the type hint say asyncio.Event but a threading.Event is passed in

[rkooo567]

ah I tried with non-blocking appraoch, and it didn't work, and I changed to blocking way. Let's just go with #40122 as it makes more sense to me..

[edoakes]

why's this change needed btw?

[rkooo567]

This is to reduce the impact of the blocking call (so when things are cancelled, the coroutine is already context switched). I think with the new fix, it is not necessray

[edoakes]

@rkooo567 I just verified that my suggested change fixes the test case you've added. Without the following diff, I see the segfaults. With the following diff, the test passes consistently (without any of your other changes):

diff --git a/python/ray/_raylet.pyx b/python/ray/_raylet.pyx
index af1e777680..5449563478 100644
--- a/python/ray/_raylet.pyx
+++ b/python/ray/_raylet.pyx
@@ -4280,22 +4280,24 @@ cdef class CoreWorker:
             function_descriptor, specified_cgname)
 
         async def async_func():
-            if task_id:
-                async_task_id.set(task_id)
+            try:
+                if task_id:
+                    async_task_id.set(task_id)
 
-            if inspect.isawaitable(func_or_coro):
-                coroutine = func_or_coro
-            else:
-                coroutine = func_or_coro(*args, **kwargs)
+                if inspect.isawaitable(func_or_coro):
+                    coroutine = func_or_coro
+                else:
+                    coroutine = func_or_coro(*args, **kwargs)
 
-            return await coroutine
+                return await coroutine
+            finally:
+                event.Notify()
 
         future = asyncio.run_coroutine_threadsafe(async_func(), eventloop)
         if task_id:
             with self._task_id_to_future_lock:
                 self._task_id_to_future[task_id] = future
 
-        future.add_done_callback(lambda _: event.Notify())
         with nogil:
             (CCoreWorkerProcess.GetCoreWorker()
                 .YieldCurrentFiber(event))

[edoakes]

Made a PR to run CI against it: #40122

[rkooo567]

I assume the current issue is that the callback added with add_done_callback is running prior to the coroutine actually completing. If that's the case, then how about we move event.Notify() into a finally block in the async_func just above?

@edoakes this fix makes sense

[rkooo567]

there's a better/safer approach here; #40122