[Serve LLM] Handle missing state attributes from vLLM's task-conditional init_app_state (#60812)

Signed-off-by: Kourosh Hakhamaneshi <kourosh@anyscale.com>

Author: kouroshHakha

python/ray/llm/_internal/serve/engines/vllm/vllm_engine.py

@@ -322,13 +322,17 @@ async def start(self) -> None:
                 args=args,
             )
 
-        self._oai_models = state.openai_serving_models
-        self._oai_serving_chat = state.openai_serving_chat
-        self._oai_serving_completion = state.openai_serving_completion
-        self._oai_serving_embedding = state.openai_serving_embedding
-        self._oai_serving_transcription = state.openai_serving_transcription
-        self._oai_serving_scores = state.openai_serving_scores
-        self._oai_serving_tokenization = state.openai_serving_tokenization
+        self._oai_models = getattr(state, "openai_serving_models", None)
+        self._oai_serving_chat = getattr(state, "openai_serving_chat", None)
+        self._oai_serving_completion = getattr(state, "openai_serving_completion", None)
+        self._oai_serving_embedding = getattr(state, "openai_serving_embedding", None)
+        self._oai_serving_transcription = getattr(
+            state, "openai_serving_transcription", None
+        )
+        self._oai_serving_scores = getattr(state, "openai_serving_scores", None)
+        self._oai_serving_tokenization = getattr(
+            state, "openai_serving_tokenization", None
+        )
 
         self._validate_openai_serving_models()
         self._validate_engine_client()
@@ -346,38 +350,53 @@ def _validate_openai_serving_models(self):
             self._oai_models, "load_lora_adapter"
         ), "oai_models must have a load_lora_adapter attribute"
 
-    def _validate_openai_serving_chat(self):
-        assert hasattr(
-            self._oai_serving_chat, "create_chat_completion"
-        ), "oai_serving_chat must have a create_chat_completion attribute"
+    @staticmethod
+    def _make_error(message: str) -> ErrorResponse:
+        return ErrorResponse(
+            error=ErrorInfo(message=message, type="invalid_request_error", code=400)
+        )
 
-    def _validate_openai_serving_completion(self):
-        assert hasattr(
-            self._oai_serving_completion, "create_completion"
-        ), "oai_serving_completion must have a create_completion attribute"
+    def _validate_openai_serving_chat(self) -> Optional[ErrorResponse]:
+        if self._oai_serving_chat is None:
+            return self._make_error(
+                "This model does not support the 'generate' task. "
+                "The chat completion endpoint is not available for this model."
+            )
 
-    def _validate_openai_serving_embedding(self):
-        assert hasattr(
-            self._oai_serving_embedding, "create_embedding"
-        ), "oai_serving_embedding must have a create_embedding attribute"
+    def _validate_openai_serving_completion(self) -> Optional[ErrorResponse]:
+        if self._oai_serving_completion is None:
+            return self._make_error(
+                "This model does not support the 'generate' task. "
+                "The completion endpoint is not available for this model."
+            )
 
-    def _validate_openai_serving_transcription(self):
-        assert hasattr(
-            self._oai_serving_transcription, "create_transcription"
-        ), "oai_serving_transcription must have a create_transcription attribute"
+    def _validate_openai_serving_embedding(self) -> Optional[ErrorResponse]:
+        if self._oai_serving_embedding is None:
+            return self._make_error(
+                "This model does not support the 'embed' task. "
+                "The embedding endpoint is not available for this model."
+            )
 
-    def _validate_openai_serving_scores(self):
-        assert hasattr(
-            self._oai_serving_scores, "create_score"
-        ), "oai_serving_scores must have a create_score attribute"
+    def _validate_openai_serving_transcription(self) -> Optional[ErrorResponse]:
+        if self._oai_serving_transcription is None:
+            return self._make_error(
+                "This model does not support the 'transcription' task. "
+                "The transcription endpoint is not available for this model."
+            )
 
-    def _validate_openai_serving_tokenization(self):
-        assert hasattr(
-            self._oai_serving_tokenization, "create_tokenize"
-        ), "oai_serving_tokenization must have a create_tokenize attribute"
-        assert hasattr(
-            self._oai_serving_tokenization, "create_detokenize"
-        ), "oai_serving_tokenization must have a create_detokenize attribute"
+    def _validate_openai_serving_scores(self) -> Optional[ErrorResponse]:
+        if self._oai_serving_scores is None:
+            return self._make_error(
+                "This model does not support the 'score' task. "
+                "The score endpoint is not available for this model."
+            )
+
+    def _validate_openai_serving_tokenization(self) -> Optional[ErrorResponse]:
+        if self._oai_serving_tokenization is None:
+            return self._make_error(
+                "This model does not support the 'tokenization' task. "
+                "The tokenization endpoint is not available for this model."
+            )
 
     def _validate_engine_client(self):
         assert hasattr(
@@ -486,7 +505,9 @@ async def chat(
         request: ChatCompletionRequest,
         raw_request_info: Optional[RawRequestInfo] = None,
     ) -> AsyncGenerator[Union[str, ChatCompletionResponse, ErrorResponse], None]:
-        self._validate_openai_serving_chat()
+        if error := self._validate_openai_serving_chat():
+            yield error
+            return
 
         raw_request: Optional[Request] = RawRequestInfo.to_starlette_request_optional(
             raw_request_info
@@ -514,7 +535,9 @@ async def completions(
         request: CompletionRequest,
         raw_request_info: Optional[RawRequestInfo] = None,
     ) -> AsyncGenerator[Union[str, CompletionResponse, ErrorResponse], None]:
-        self._validate_openai_serving_completion()
+        if error := self._validate_openai_serving_completion():
+            yield error
+            return
 
         raw_request: Optional[Request] = RawRequestInfo.to_starlette_request_optional(
             raw_request_info
@@ -544,7 +567,9 @@ async def embeddings(
         request: EmbeddingRequest,
         raw_request_info: Optional[RawRequestInfo] = None,
     ) -> AsyncGenerator[Union[EmbeddingResponse, ErrorResponse], None]:
-        self._validate_openai_serving_embedding()
+        if error := self._validate_openai_serving_embedding():
+            yield error
+            return
 
         raw_request: Optional[Request] = RawRequestInfo.to_starlette_request_optional(
             raw_request_info
@@ -566,7 +591,9 @@ async def transcriptions(
         request: TranscriptionRequest,
         raw_request_info: Optional[RawRequestInfo] = None,
     ) -> AsyncGenerator[Union[str, TranscriptionResponse, ErrorResponse], None]:
-        self._validate_openai_serving_transcription()
+        if error := self._validate_openai_serving_transcription():
+            yield error
+            return
 
         # Extract audio data from the request file
         audio_data = await request.file.read()
@@ -600,7 +627,9 @@ async def score(
         request: ScoreRequest,
         raw_request_info: Optional[RawRequestInfo] = None,
     ) -> AsyncGenerator[Union[ScoreResponse, ErrorResponse], None]:
-        self._validate_openai_serving_scores()
+        if error := self._validate_openai_serving_scores():
+            yield error
+            return
 
         raw_request: Optional[Request] = RawRequestInfo.to_starlette_request_optional(
             raw_request_info
@@ -620,7 +649,9 @@ async def tokenize(
         request: TokenizeRequest,
         raw_request_info: Optional[RawRequestInfo] = None,
     ) -> AsyncGenerator[Union[TokenizeResponse, ErrorResponse], None]:
-        self._validate_openai_serving_tokenization()
+        if error := self._validate_openai_serving_tokenization():
+            yield error
+            return
 
         raw_request: Optional[Request] = RawRequestInfo.to_starlette_request_optional(
             raw_request_info
@@ -640,7 +671,9 @@ async def detokenize(
         request: DetokenizeRequest,
         raw_request_info: Optional[RawRequestInfo] = None,
     ) -> AsyncGenerator[Union[DetokenizeResponse, ErrorResponse], None]:
-        self._validate_openai_serving_tokenization()
+        if error := self._validate_openai_serving_tokenization():
+            yield error
+            return
 
         raw_request: Optional[Request] = RawRequestInfo.to_starlette_request_optional(
             raw_request_info