[serve] support running user code on same event loop (#54219)

## Why are these changes needed?

Currently we run all user code on a separate event loop (the user code
event loop) to avoid the system event loop being blocked by user code.
However if the user is confident that they are writing proper async
python code, they can choose to bypass that and run user code on the
same event loop. This will boost the performance since we avoid
cross-thread communication.

---------

Signed-off-by: Cindy Zhang <cindyzyx9@gmail.com>
Co-authored-by: akyang-anyscale <alexyang@anyscale.com>

Author: zcin

python/ray/serve/_private/constants.py

@@ -455,3 +455,10 @@ def str_to_list(s: str) -> List[str]:
 
 # Name of deployment request routing stats method implemented by user.
 REQUEST_ROUTING_STATS_METHOD = "record_routing_stats"
+
+# By default, we run user code in a separate event loop.
+# This flag can be set to 0 to run user code in the same event loop as the
+# replica's main event loop.
+RAY_SERVE_RUN_USER_CODE_IN_SEPARATE_THREAD = (
+    os.environ.get("RAY_SERVE_RUN_USER_CODE_IN_SEPARATE_THREAD", "1") == "1"
+)

python/ray/serve/_private/local_testing_mode.py

@@ -70,6 +70,7 @@ def make_local_deployment_handle(
         deployment.init_kwargs,
         deployment_id=deployment_id,
         run_sync_methods_in_threadpool=RAY_SERVE_RUN_SYNC_IN_THREADPOOL,
+        run_user_code_in_separate_thread=True,
         local_testing_mode=True,
     )
     try:

python/ray/serve/_private/replica.py

@@ -55,6 +55,7 @@
     RAY_SERVE_REPLICA_AUTOSCALING_METRIC_RECORD_PERIOD_S,
     RAY_SERVE_RUN_SYNC_IN_THREADPOOL,
     RAY_SERVE_RUN_SYNC_IN_THREADPOOL_WARNING,
+    RAY_SERVE_RUN_USER_CODE_IN_SEPARATE_THREAD,
     RECONFIGURE_METHOD,
     REQUEST_LATENCY_BUCKETS_MS,
     REQUEST_ROUTING_STATS_METHOD,
@@ -371,6 +372,7 @@ def __init__(
             init_kwargs,
             deployment_id=self._deployment_id,
             run_sync_methods_in_threadpool=RAY_SERVE_RUN_SYNC_IN_THREADPOOL,
+            run_user_code_in_separate_thread=RAY_SERVE_RUN_USER_CODE_IN_SEPARATE_THREAD,
             local_testing_mode=False,
         )
         self._semaphore = Semaphore(lambda: self.max_ongoing_requests)
@@ -1169,6 +1171,7 @@ def __init__(
         *,
         deployment_id: DeploymentID,
         run_sync_methods_in_threadpool: bool,
+        run_user_code_in_separate_thread: bool,
         local_testing_mode: bool,
     ):
         if not (inspect.isfunction(deployment_def) or inspect.isclass(deployment_def)):
@@ -1185,33 +1188,44 @@ def __init__(
         self._local_testing_mode = local_testing_mode
         self._destructor_called = False
         self._run_sync_methods_in_threadpool = run_sync_methods_in_threadpool
+        self._run_user_code_in_separate_thread = run_user_code_in_separate_thread
         self._warned_about_sync_method_change = False
         self._cached_user_method_info: Dict[str, UserMethodInfo] = {}
 
         # Will be populated in `initialize_callable`.
         self._callable = None
 
-        # All interactions with user code run on this loop to avoid blocking the
-        # replica's main event loop.
-        self._user_code_event_loop: asyncio.AbstractEventLoop = asyncio.new_event_loop()
+        if self._run_user_code_in_separate_thread:
+            # All interactions with user code run on this loop to avoid blocking the
+            # replica's main event loop.
+            self._user_code_event_loop: asyncio.AbstractEventLoop = (
+                asyncio.new_event_loop()
+            )
 
-        def _run_user_code_event_loop():
-            # Required so that calls to get the current running event loop work
-            # properly in user code.
-            asyncio.set_event_loop(self._user_code_event_loop)
-            self._user_code_event_loop.run_forever()
+            def _run_user_code_event_loop():
+                # Required so that calls to get the current running event loop work
+                # properly in user code.
+                asyncio.set_event_loop(self._user_code_event_loop)
+                self._user_code_event_loop.run_forever()
 
-        self._user_code_event_loop_thread = threading.Thread(
-            daemon=True,
-            target=_run_user_code_event_loop,
-        )
-        self._user_code_event_loop_thread.start()
+            self._user_code_event_loop_thread = threading.Thread(
+                daemon=True,
+                target=_run_user_code_event_loop,
+            )
+            self._user_code_event_loop_thread.start()
+        else:
+            self._user_code_event_loop = asyncio.get_running_loop()
+
+    @property
+    def event_loop(self) -> asyncio.AbstractEventLoop:
+        return self._user_code_event_loop
 
     def _run_user_code(f: Callable) -> Callable:
         """Decorator to run a coroutine method on the user code event loop.
 
         The method will be modified to be a sync function that returns a
-        `asyncio.Future`.
+        `asyncio.Future` if user code is running in a separate event loop.
+        Otherwise, it will return the coroutine directly.
         """
         assert inspect.iscoroutinefunction(
             f
@@ -1220,11 +1234,14 @@ def _run_user_code(f: Callable) -> Callable:
         @functools.wraps(f)
         def wrapper(self, *args, **kwargs) -> Any:
             coro = f(self, *args, **kwargs)
-            fut = asyncio.run_coroutine_threadsafe(coro, self._user_code_event_loop)
-            if self._local_testing_mode:
-                return fut
+            if self._run_user_code_in_separate_thread:
+                fut = asyncio.run_coroutine_threadsafe(coro, self._user_code_event_loop)
+                if self._local_testing_mode:
+                    return fut
 
-            return asyncio.wrap_future(fut)
+                return asyncio.wrap_future(fut)
+            else:
+                return coro
 
         return wrapper
 
@@ -1576,18 +1593,22 @@ async def call_http_entrypoint(
         # `asyncio.Event`s are not thread safe, so `call_soon_threadsafe` must be
         # used to interact with the result queue from the user callable thread.
         system_event_loop = asyncio.get_running_loop()
+        user_method_info = self.get_user_method_info(request_metadata.call_method)
 
         async def enq(item: Any):
             system_event_loop.call_soon_threadsafe(result_queue.put_nowait, item)
 
-        user_method_info = self.get_user_method_info(request_metadata.call_method)
-        call_user_method_future = self._call_http_entrypoint(
-            user_method_info, scope, receive, enq
-        )
+        if self._run_user_code_in_separate_thread:
+            call_future = self._call_http_entrypoint(
+                user_method_info, scope, receive, enq
+            )
+        else:
+            call_future = asyncio.create_task(
+                self._call_http_entrypoint(user_method_info, scope, receive, enq)
+            )
+
         first_message_peeked = False
-        async for messages in result_queue.fetch_messages_from_queue(
-            call_user_method_future
-        ):
+        async for messages in result_queue.fetch_messages_from_queue(call_future):
             # HTTP (ASGI) messages are only consumed by the proxy so batch them
             # and use vanilla pickle (we know it's safe because these messages
             # only contain primitive Python types).
@@ -1702,16 +1723,24 @@ async def call_user_generator(
         def _enqueue_thread_safe(item: Any):
             system_event_loop.call_soon_threadsafe(result_queue.put_nowait, item)
 
-        call_user_method_future = self._call_user_generator(
-            request_metadata,
-            request_args,
-            request_kwargs,
-            generator_result_callback=_enqueue_thread_safe,
-        )
+        if self._run_user_code_in_separate_thread:
+            call_future = self._call_user_generator(
+                request_metadata,
+                request_args,
+                request_kwargs,
+                generator_result_callback=_enqueue_thread_safe,
+            )
+        else:
+            call_future = asyncio.create_task(
+                self._call_user_generator(
+                    request_metadata,
+                    request_args,
+                    request_kwargs,
+                    generator_result_callback=_enqueue_thread_safe,
+                )
+            )
 
-        async for messages in result_queue.fetch_messages_from_queue(
-            call_user_method_future
-        ):
+        async for messages in result_queue.fetch_messages_from_queue(call_future):
             for msg in messages:
                 yield msg