chore: import upstream snapshot with attribution

python/ray/serve/_private/endpoint_state.py

@@ -0,0 +1,109 @@
+import logging
+from typing import Any, Dict, Optional
+
+from ray import cloudpickle
+from ray.serve._private.common import DeploymentID, EndpointInfo
+from ray.serve._private.constants import SERVE_LOGGER_NAME
+from ray.serve._private.long_poll import LongPollHost, LongPollNamespace
+from ray.serve._private.storage.kv_store import KVStoreBase
+
+CHECKPOINT_KEY = "serve-endpoint-state-checkpoint"
+
+logger = logging.getLogger(SERVE_LOGGER_NAME)
+
+
+class EndpointState:
+    """Manages all state for endpoints in the system.
+
+    This class is *not* thread safe, so any state-modifying methods should be
+    called with a lock held.
+    """
+
+    def __init__(self, kv_store: KVStoreBase, long_poll_host: LongPollHost):
+        self._kv_store = kv_store
+        self._long_poll_host = long_poll_host
+        self._endpoints: Dict[DeploymentID, EndpointInfo] = dict()
+
+        checkpoint = self._kv_store.get(CHECKPOINT_KEY)
+        if checkpoint is not None:
+            self._endpoints = cloudpickle.loads(checkpoint)
+
+        self._notify_route_table_changed()
+
+    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))
+
+    def _notify_route_table_changed(self):
+        self._long_poll_host.notify_changed(
+            {LongPollNamespace.ROUTE_TABLE: self._endpoints}
+        )
+
+    def _get_endpoint_for_route(self, route: str) -> Optional[DeploymentID]:
+        for endpoint, info in self._endpoints.items():
+            if info.route == route:
+                return endpoint
+
+        return None
+
+    def update_endpoint(
+        self, endpoint: DeploymentID, endpoint_info: EndpointInfo
+    ) -> None:
+        """Create or update the given endpoint.
+
+        This method is idempotent - if the endpoint already exists it will be
+        updated to match the given parameters. Calling this twice with the same
+        arguments is a no-op.
+        """
+
+        if self._endpoints.get(endpoint) == endpoint_info:
+            return
+
+        existing_route_endpoint = self._get_endpoint_for_route(endpoint_info.route)
+        if existing_route_endpoint is not None and existing_route_endpoint != endpoint:
+            logger.debug(
+                f'route_prefix "{endpoint_info.route}" is currently '
+                f'registered to deployment "{existing_route_endpoint.name}". '
+                f'Re-registering route_prefix "{endpoint_info.route}" to '
+                f'deployment "{endpoint.name}".'
+            )
+            del self._endpoints[existing_route_endpoint]
+
+        self._endpoints[endpoint] = endpoint_info
+
+        self._checkpoint()
+        self._notify_route_table_changed()
+
+    def get_endpoint_route(self, endpoint: DeploymentID) -> Optional[str]:
+        if endpoint in self._endpoints:
+            return self._endpoints[endpoint].route
+        return None
+
+    def get_endpoints(self) -> Dict[DeploymentID, Dict[str, Any]]:
+        endpoints = {}
+        for endpoint, info in self._endpoints.items():
+            endpoints[endpoint] = {
+                "route": info.route,
+            }
+        return endpoints
+
+    def delete_endpoint(self, endpoint: DeploymentID) -> None:
+        # This method must be idempotent. We should validate that the
+        # specified endpoint exists on the client.
+        if endpoint not in self._endpoints:
+            return
+
+        del self._endpoints[endpoint]
+
+        self._checkpoint()
+        self._notify_route_table_changed()