[Serve] Fix serve non atomic shutdown

python/ray/serve/BUILD

@@ -672,3 +672,11 @@ py_test(
     tags = ["exclusive", "team:serve"],
     deps = [":serve_lib"],
 )
+
+py_test(
+    name = "test_endpoint_state",
+    size = "small",
+    srcs = serve_tests_srcs,
+    tags = ["exclusive", "team:serve"],
+    deps = [":serve_lib"],
+)

python/ray/serve/_private/application_state.py

@@ -177,9 +177,20 @@ def _delete_deployment(self, name):
 
     def delete(self):
         """Delete the application"""
-        logger.info(f"Deleting application '{self._name}'")
+        logger.info(
+            f"Deleting application '{self._name}'",
+            extra={"log_to_stderr": False},
+        )
         self._set_target_state(deleting=True)
 
+    def is_deleted(self) -> bool:
+        """Check whether the application is already deleted.
+
+        For an application to be considered deleted, the target state has to be set to
+        deleting and all deployments have to be deleted.
+        """
+        return self._target_state.deleting and len(self._get_live_deployments()) == 0
+
     def apply_deployment_info(
         self, deployment_name: str, deployment_info: DeploymentInfo
     ) -> None:
@@ -412,7 +423,7 @@ def update(self) -> bool:
 
         # Check if app is ready to be deleted
         if self._target_state.deleting:
-            return len(self._get_live_deployments()) == 0
+            return self.is_deleted()
         return False
 
     def get_checkpoint_data(self) -> ApplicationTargetState:
@@ -639,6 +650,16 @@ def shutdown(self) -> None:
         for app_state in self._application_states.values():
             app_state.delete()
 
+    def is_ready_for_shutdown(self) -> bool:
+        """Return whether all applications have shut down.
+
+        Iterate through all application states and check if all their applications
+        are deleted.
+        """
+        return all(
+            app_state.is_deleted() for app_state in self._application_states.values()
+        )
+
     def _save_checkpoint_func(
         self, *, writeahead_checkpoints: Optional[Dict[str, ApplicationTargetState]]
     ) -> None:

python/ray/serve/_private/client.py

@@ -20,7 +20,6 @@
     CLIENT_POLLING_INTERVAL_S,
     CLIENT_CHECK_CREATION_POLLING_INTERVAL_S,
     MAX_CACHED_HANDLES,
-    SERVE_NAMESPACE,
     SERVE_DEFAULT_APP_NAME,
 )
 from ray.serve._private.deploy_utils import get_deploy_args
@@ -95,7 +94,7 @@ def __del__(self):
     def __reduce__(self):
         raise RayServeException(("Ray Serve client cannot be serialized."))
 
-    def shutdown(self) -> None:
+    def shutdown(self, timeout_s: float = 30.0) -> None:
         """Completely shut down the connected Serve instance.
 
         Shuts down all processes and deletes all state associated with the
@@ -107,53 +106,17 @@ def shutdown(self) -> None:
             del self.handle_cache[k]
 
         if ray.is_initialized() and not self._shutdown:
-            ray.get(self._controller.shutdown.remote())
-            self._wait_for_deployments_shutdown()
-
-            ray.kill(self._controller, no_restart=True)
-
-            # Wait for the named actor entry gets removed as well.
-            started = time.time()
-            while True:
-                try:
-                    ray.get_actor(self._controller_name, namespace=SERVE_NAMESPACE)
-                    if time.time() - started > 5:
-                        logger.warning(
-                            "Waited 5s for Serve to shutdown gracefully but "
-                            "the controller is still not cleaned up. "
-                            "You can ignore this warning if you are shutting "
-                            "down the Ray cluster."
-                        )
-                        break
-                except ValueError:  # actor name is removed
-                    break
-
-            self._shutdown = True
-
-    def _wait_for_deployments_shutdown(self, timeout_s: int = 60):
-        """Waits for all deployments to be shut down and deleted.
-
-        Raises TimeoutError if this doesn't happen before timeout_s.
-        """
-        start = time.time()
-        while time.time() - start < timeout_s:
-            deployment_statuses = self.get_all_deployment_statuses()
-            if len(deployment_statuses) == 0:
-                break
-            else:
-                logger.debug(
-                    f"Waiting for shutdown, {len(deployment_statuses)} "
-                    "deployments still alive."
+            try:
+                ray.get(self._controller.graceful_shutdown.remote(), timeout=timeout_s)
+            except ray.exceptions.RayActorError:
+                # Controller has been shut down.
+                pass
+            except TimeoutError:
+                logger.warning(
+                    f"Controller failed to shut down within {timeout_s}s. "
+                    "Check controller logs for more details."
                 )
-            time.sleep(CLIENT_POLLING_INTERVAL_S)
-        else:
-            live_names = [
-                deployment_status.name for deployment_status in deployment_statuses
-            ]
-            raise TimeoutError(
-                f"Shutdown didn't complete after {timeout_s}s. "
-                f"Deployments still alive: {live_names}."
-            )
+            self._shutdown = True
 
     def _wait_for_deployment_healthy(self, name: str, timeout_s: int = -1):
         """Waits for the named deployment to enter "HEALTHY" status.

python/ray/serve/_private/deployment_state.py

@@ -1279,7 +1279,10 @@ def _set_target_state_deleting(self) -> None:
         self._curr_status_info = DeploymentStatusInfo(
             self._name, DeploymentStatus.UPDATING
         )
-        logger.info(f"Deleting deployment {self._name}.")
+        logger.info(
+            f"Deleting deployment {self._name}.",
+            extra={"log_to_stderr": False},
+        )
 
     def _set_target_state(self, target_info: DeploymentInfo) -> None:
         """Set the target state for the deployment to the provided info."""
@@ -2284,6 +2287,16 @@ def shutdown(self):
         # TODO(jiaodong): Need to add some logic to prevent new replicas
         # from being created once shutdown signal is sent.
 
+    def is_ready_for_shutdown(self) -> bool:
+        """Return whether all deployments are shutdown.
+
+        Check there are no deployment states and no checkpoints.
+        """
+        return (
+            len(self._deployment_states) == 0
+            and self._kv_store.get(CHECKPOINT_KEY) is None
+        )
+
     def _save_checkpoint_func(
         self, *, writeahead_checkpoints: Optional[Dict[str, Tuple]]
     ) -> None:

python/ray/serve/_private/endpoint_state.py

@@ -34,6 +34,14 @@ def __init__(self, kv_store: KVStoreBase, long_poll_host: LongPollHost):
     def shutdown(self):
         self._kv_store.delete(CHECKPOINT_KEY)
 
+    def is_ready_for_shutdown(self) -> bool:
+        """Returns whether the endpoint checkpoint has been deleted.
+
+        Get the endpoint checkpoint from the kv store. If it is None, then it has been
+        deleted.
+        """
+        return self._kv_store.get(CHECKPOINT_KEY) is None
+
     def _checkpoint(self):
         self._kv_store.put(CHECKPOINT_KEY, cloudpickle.dumps(self._endpoints))
 

python/ray/serve/_private/http_state.py

@@ -230,6 +230,27 @@ def shutdown(self):
         self._shutting_down = True
         ray.kill(self.actor_handle, no_restart=True)
 
+    def is_ready_for_shutdown(self) -> bool:
+        """Return whether the HTTP proxy actor is shutdown.
+
+        For an HTTP proxy actor to be considered shutdown, it must be marked as
+        _shutting_down and the actor must be dead. If the actor is dead, the health
+        check will return RayActorError.
+        """
+        if not self._shutting_down:
+            return False
+
+        try:
+            ray.get(self._actor_handle.check_health.remote(), timeout=0.001)
+        except ray.exceptions.RayActorError:
+            # The actor is dead, so it's ready for shutdown.
+            return True
+        except ray.exceptions.GetTimeoutError:
+            # The actor is still alive, so it's not ready for shutdown.
+            return False
+
+        return False
+
 
 class HTTPState:
     """Manages all state for HTTP proxies in the system.
@@ -269,6 +290,17 @@ def shutdown(self) -> None:
         for proxy_state in self._proxy_states.values():
             proxy_state.shutdown()
 
+    def is_ready_for_shutdown(self) -> bool:
+        """Return whether all proxies are shutdown.
+
+        Iterate through all proxy states and check if all their proxy actors
+        are shutdown.
+        """
+        return all(
+            proxy_state.is_ready_for_shutdown()
+            for proxy_state in self._proxy_states.values()
+        )
+
     def get_config(self):
         return self._config
 

python/ray/serve/controller.py

@@ -187,6 +187,9 @@ async def __init__(
             worker_id=ray.get_runtime_context().get_worker_id(),
             log_file_path=get_component_logger_file_path(),
         )
+        self._shutting_down = False
+        self._shutdown = asyncio.Event()
+        self._shutdown_start_time = None
 
         run_background_task(self.run_control_loop())
 
@@ -307,6 +310,12 @@ async def run_control_loop(self) -> None:
         recovering_timeout = RECOVERING_LONG_POLL_BROADCAST_TIMEOUT_S
         start_time = time.time()
         while True:
+            if self._shutting_down:
+                try:
+                    self.shutdown()
+                except Exception:
+                    logger.exception("Exception during shutdown.")
+
             if (
                 not self.done_recovering_event.is_set()
                 and time.time() - start_time > recovering_timeout
@@ -439,15 +448,84 @@ def get_root_url(self):
                 )
         return http_config.root_url
 
+    def config_checkpoint_deleted(self) -> bool:
+        """Returns whether the config checkpoint has been deleted.
+
+        Get the config checkpoint from the kv store. If it is None, then it has been
+        deleted.
+        """
+        return self.kv_store.get(CONFIG_CHECKPOINT_KEY) is None
+
     def shutdown(self):
-        """Shuts down the serve instance completely."""
+        """Shuts down the serve instance completely.
+
+        This method will only be triggered when `self._shutting_down` is true. It
+        deletes the kv store for config checkpoints, sets application state to deleting,
+        delete all deployments, and shuts down all HTTP proxies. Once all these
+        resources are released, it then kills the controller actor.
+        """
+        if not self._shutting_down:
+            return
+
+        if self._shutdown_start_time is None:
+            self._shutdown_start_time = time.time()
+
+        logger.info("Controller shutdown started!", extra={"log_to_stderr": False})
         self.kv_store.delete(CONFIG_CHECKPOINT_KEY)
         self.application_state_manager.shutdown()
         self.deployment_state_manager.shutdown()
         self.endpoint_state.shutdown()
         if self.http_state:
             self.http_state.shutdown()
 
+        config_checkpoint_deleted = self.config_checkpoint_deleted()
+        application_is_shutdown = self.application_state_manager.is_ready_for_shutdown()
+        deployment_is_shutdown = self.deployment_state_manager.is_ready_for_shutdown()
+        endpoint_is_shutdown = self.endpoint_state.is_ready_for_shutdown()
+        http_state_is_shutdown = (
+            self.http_state is None or self.http_state.is_ready_for_shutdown()
+        )
+        if (
+            config_checkpoint_deleted
+            and application_is_shutdown
+            and deployment_is_shutdown
+            and endpoint_is_shutdown
+            and http_state_is_shutdown
+        ):
+            logger.warning(
+                "All resources have shut down, shutting down controller!",
+                extra={"log_to_stderr": False},
+            )
+            _controller_actor = ray.get_runtime_context().current_actor
+            self._shutdown.set()
+            ray.kill(_controller_actor, no_restart=True)
+        elif time.time() - self._shutdown_start_time > 10:
+            if not config_checkpoint_deleted:
+                logger.warning(
+                    f"{CONFIG_CHECKPOINT_KEY} not yet deleted",
+                    extra={"log_to_stderr": False},
+                )
+            if not application_is_shutdown:
+                logger.warning(
+                    "application not yet shutdown",
+                    extra={"log_to_stderr": False},
+                )
+            if not deployment_is_shutdown:
+                logger.warning(
+                    "deployment not yet shutdown",
+                    extra={"log_to_stderr": False},
+                )
+            if not endpoint_is_shutdown:
+                logger.warning(
+                    "endpoint not yet shutdown",
+                    extra={"log_to_stderr": False},
+                )
+            if not http_state_is_shutdown:
+                logger.warning(
+                    "http_state not yet shutdown",
+                    extra={"log_to_stderr": False},
+                )
+
     def deploy(
         self,
         name: str,
@@ -835,6 +913,20 @@ def record_multiplexed_replica_info(self, info: MultiplexedReplicaInfo):
         """
         self.deployment_state_manager.record_multiplexed_replica_info(info)
 
+    async def graceful_shutdown(self, wait: bool = True):
+        """Set the shutting down flag on controller to signal shutdown in
+        run_control_loop().
+
+        This is used to signal to the controller that it should proceed with shutdown
+        process, so it can shut down gracefully. It also waits until the shutdown
+        event is triggered if wait is true.
+        """
+        self._shutting_down = True
+        if not wait:
+            return
+
+        await self._shutdown.wait()
+
 
 @ray.remote(num_cpus=0, max_calls=1)
 def deploy_serve_application(
@@ -943,7 +1035,7 @@ def __init__(
         http_proxy_port: int = 8000,
     ):
         try:
-            self._controller = ray.get_actor(controller_name, namespace="serve")
+            self._controller = ray.get_actor(controller_name, namespace=SERVE_NAMESPACE)
         except ValueError:
             self._controller = None
         if self._controller is None:
@@ -956,7 +1048,7 @@ def __init__(
                 max_restarts=-1,
                 max_task_retries=-1,
                 resources={HEAD_NODE_RESOURCE_NAME: 0.001},
-                namespace="serve",
+                namespace=SERVE_NAMESPACE,
                 max_concurrency=CONTROLLER_MAX_CONCURRENCY,
             ).remote(
                 controller_name,