worker-vllm/src/engine.py

import asyncio
import inspect
import json
import logging
import os
import time
from typing import AsyncGenerator, Optional

from dotenv import load_dotenv
from vllm import AsyncLLMEngine
from vllm.inputs import TextPrompt
from vllm.entrypoints.anthropic.protocol import AnthropicMessagesRequest, AnthropicMessagesResponse, AnthropicError, AnthropicErrorResponse
from vllm.entrypoints.anthropic.serving import AnthropicServingMessages
from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionRequest
from vllm.entrypoints.openai.chat_completion.serving import OpenAIServingChat
from vllm.entrypoints.openai.completion.protocol import CompletionRequest
from vllm.entrypoints.openai.completion.serving import OpenAIServingCompletion
from vllm.entrypoints.openai.engine.protocol import ErrorResponse
from vllm.entrypoints.openai.models.protocol import BaseModelPath, LoRAModulePath
from vllm.entrypoints.openai.models.serving import OpenAIServingModels
from vllm.entrypoints.openai.responses.protocol import ResponsesRequest, ResponsesResponse
from vllm.entrypoints.openai.responses.serving import OpenAIServingResponses
from vllm.entrypoints.serve.render.serving import OpenAIServingRender

from constants import DEFAULT_BATCH_SIZE, DEFAULT_BATCH_SIZE_GROWTH_FACTOR, DEFAULT_MAX_CONCURRENCY, DEFAULT_MIN_BATCH_SIZE
from engine_args import get_engine_args
from tokenizer import TokenizerWrapper
from utils import BatchSize, DummyRequest, JobInput, create_error_response

class vLLMEngine:
    def __init__(self, engine = None):
        load_dotenv() # For local development
        self.engine_args = get_engine_args()

        if engine is None:
            ea = self.engine_args
            summary = {
                "model": ea.model,
                "dtype": ea.dtype,
                "quantization": ea.quantization,
                "max_model_len": ea.max_model_len,
                "tensor_parallel_size": ea.tensor_parallel_size,
                "gpu_memory_utilization": ea.gpu_memory_utilization,
            }
            if ea.tokenizer and ea.tokenizer != ea.model:
                summary["tokenizer"] = ea.tokenizer
            logging.info("Engine config: %s", summary)
            logging.debug("Full engine args: %s", ea)

            self.llm = self._initialize_llm()

            if self.engine_args.tokenizer_mode != 'mistral':
                self.tokenizer = TokenizerWrapper(self.engine_args.tokenizer or self.engine_args.model,
                                                  self.engine_args.tokenizer_revision,
                                                  self.engine_args.trust_remote_code)
            else:
                self.tokenizer = None
        else:
            self.llm = engine.llm
            self.tokenizer = engine.tokenizer
            
        self.max_concurrency = int(os.getenv("MAX_CONCURRENCY", DEFAULT_MAX_CONCURRENCY))
        self.default_batch_size = int(os.getenv("DEFAULT_BATCH_SIZE", DEFAULT_BATCH_SIZE))
        self.batch_size_growth_factor = int(os.getenv("BATCH_SIZE_GROWTH_FACTOR", DEFAULT_BATCH_SIZE_GROWTH_FACTOR))
        self.min_batch_size = int(os.getenv("MIN_BATCH_SIZE", DEFAULT_MIN_BATCH_SIZE))

    def _get_tokenizer_for_chat_template(self):
        """Get tokenizer for chat template application"""
        if self.tokenizer is not None:
            return self.tokenizer
        else:
            # For mistral models, get tokenizer from vLLM engine
            # This is a fallback - ideally chat templates should be handled by vLLM directly
            try:
                from transformers import AutoTokenizer
                tokenizer = AutoTokenizer.from_pretrained(
                    self.engine_args.tokenizer or self.engine_args.model,
                    revision=self.engine_args.tokenizer_revision or "main",
                    trust_remote_code=self.engine_args.trust_remote_code
                )
                # Create a minimal wrapper
                class MinimalTokenizerWrapper:
                    def __init__(self, tokenizer):
                        self.tokenizer = tokenizer
                        self.custom_chat_template = os.getenv("CUSTOM_CHAT_TEMPLATE")
                        self.has_chat_template = bool(self.tokenizer.chat_template) or bool(self.custom_chat_template)
                        if self.custom_chat_template and isinstance(self.custom_chat_template, str):
                            self.tokenizer.chat_template = self.custom_chat_template
                    
                    def apply_chat_template(self, input):
                        if isinstance(input, list):
                            if not self.has_chat_template:
                                raise ValueError(
                                    "Chat template does not exist for this model, you must provide a single string input instead of a list of messages"
                                )
                        elif isinstance(input, str):
                            input = [{"role": "user", "content": input}]
                        else:
                            raise ValueError("Input must be a string or a list of messages")
                        
                        return self.tokenizer.apply_chat_template(
                            input, tokenize=False, add_generation_prompt=True
                        )
                
                return MinimalTokenizerWrapper(tokenizer)
            except Exception as e:
                logging.error(f"Failed to create fallback tokenizer: {e}")
                raise e

    def dynamic_batch_size(self, current_batch_size, batch_size_growth_factor):
        return min(current_batch_size*batch_size_growth_factor, self.default_batch_size)
                           
    async def generate(self, job_input: JobInput):
        try:
            async for batch in self._generate_vllm(
                llm_input=job_input.llm_input,
                validated_sampling_params=job_input.sampling_params,
                batch_size=job_input.max_batch_size,
                stream=job_input.stream,
                apply_chat_template=job_input.apply_chat_template,
                request_id=job_input.request_id,
                batch_size_growth_factor=job_input.batch_size_growth_factor,
                min_batch_size=job_input.min_batch_size
            ):
                yield batch
        except Exception as e:
            yield {"error": create_error_response(str(e)).model_dump()}

    async def _generate_vllm(self, llm_input, validated_sampling_params, batch_size, stream, apply_chat_template, request_id, batch_size_growth_factor, min_batch_size: str) -> AsyncGenerator[dict, None]:
        if apply_chat_template or isinstance(llm_input, list):
            tokenizer_wrapper = self._get_tokenizer_for_chat_template()
            llm_input = tokenizer_wrapper.apply_chat_template(llm_input)
        results_generator = self.llm.generate(TextPrompt(prompt=llm_input), validated_sampling_params, request_id)
        n_responses, n_input_tokens, is_first_output = validated_sampling_params.n, 0, True
        last_output_texts, token_counters = ["" for _ in range(n_responses)], {"batch": 0, "total": 0}

        batch = {
            "choices": [{"tokens": []} for _ in range(n_responses)],
        }
        
        max_batch_size = batch_size or self.default_batch_size
        batch_size_growth_factor, min_batch_size = batch_size_growth_factor or self.batch_size_growth_factor, min_batch_size or self.min_batch_size
        batch_size = BatchSize(max_batch_size, min_batch_size, batch_size_growth_factor)
    

        async for request_output in results_generator:
            if is_first_output:  # Count input tokens only once
                n_input_tokens = len(request_output.prompt_token_ids)
                is_first_output = False

            for output in request_output.outputs:
                output_index = output.index
                token_counters["total"] += 1
                if stream:
                    new_output = output.text[len(last_output_texts[output_index]):]
                    batch["choices"][output_index]["tokens"].append(new_output)
                    token_counters["batch"] += 1

                    if token_counters["batch"] >= batch_size.current_batch_size:
                        batch["usage"] = {
                            "input": n_input_tokens,
                            "output": token_counters["total"],
                        }
                        yield batch
                        batch = {
                            "choices": [{"tokens": []} for _ in range(n_responses)],
                        }
                        token_counters["batch"] = 0
                        batch_size.update()

                last_output_texts[output_index] = output.text

        if not stream:
            for output_index, output in enumerate(last_output_texts):
                batch["choices"][output_index]["tokens"] = [output]
            token_counters["batch"] += 1

        if token_counters["batch"] > 0:
            batch["usage"] = {"input": n_input_tokens, "output": token_counters["total"]}
            yield batch

    def _initialize_llm(self):
        try:
            start = time.time()
            engine = AsyncLLMEngine.from_engine_args(self.engine_args)
            end = time.time()
            logging.info(f"Initialized vLLM engine in {end - start:.2f}s")
            return engine
        except Exception as e:
            logging.error("Error initializing vLLM engine: %s", e)
            raise e


class OpenAIvLLMEngine(vLLMEngine):
    def __init__(self, vllm_engine):
        super().__init__(vllm_engine)
        self.served_model_name = os.getenv("OPENAI_SERVED_MODEL_NAME_OVERRIDE") or self.engine_args.served_model_name or self.engine_args.model
        self.response_role = os.getenv("OPENAI_RESPONSE_ROLE") or "assistant"
        self.lora_adapters = self._load_lora_adapters()

        # Always defer OpenAI engine initialization to the first request.
        # asyncio.run() creates a temporary event loop that gets closed, but async
        # components (tokenizer pool, serving engines) bind futures to that loop.
        # When Runpod's serverless handler runs in its own event loop, those futures
        # are "attached to a different loop" causing RuntimeError.
        # This affects all configurations, not just LoRA.
        self._engines_initialized = False
        if self.lora_adapters:
            logging.info(f"LoRA mode: {len(self.lora_adapters)} adapter(s) will load on first request")
            for adapter in self.lora_adapters:
                logging.info(f"  - {adapter.name}: {adapter.path}")
        else:
            logging.info("OpenAI engines will initialize on first request")

        # Handle both integer and boolean string values for RAW_OPENAI_OUTPUT
        raw_output_env = os.getenv("RAW_OPENAI_OUTPUT", "1")
        if raw_output_env.lower() in ('true', 'false'):
            self.raw_openai_output = raw_output_env.lower() == 'true'
        else:
            self.raw_openai_output = bool(int(raw_output_env))

    def _load_lora_adapters(self):
        lora_modules_env = os.getenv("LORA_MODULES", "")
        if not lora_modules_env:
            return []

        try:
            parsed = json.loads(lora_modules_env)
        except json.JSONDecodeError as e:
            logging.error(
                "LORA_MODULES could not be parsed as JSON: %s — no LoRA adapters loaded. Value: %r",
                e, lora_modules_env,
            )
            return []

        # Accept a single adapter dict as well as an array
        if isinstance(parsed, dict):
            parsed = [parsed]

        if not isinstance(parsed, list):
            logging.error(
                "LORA_MODULES must be a JSON array of adapter objects, got %s — no LoRA adapters loaded.",
                type(parsed).__name__,
            )
            return []

        adapters = []
        for i, adapter in enumerate(parsed):
            try:
                adapters.append(LoRAModulePath(**adapter))
                logging.info("Loaded LoRA adapter config [%d]: %s", i, adapter)
            except Exception as e:
                logging.error(
                    "Failed to parse LoRA adapter at index %d: %s. Config: %r",
                    i, e, adapter,
                )

        if parsed and not adapters:
            logging.error(
                "LORA_MODULES specified %d adapter(s) but none could be loaded — "
                "OpenAI model name lookups for LoRA adapters will fail.",
                len(parsed),
            )

        return adapters

    async def _ensure_engines_initialized(self):
        """Initialize engines on first request to avoid event loop mismatch.

        In Runpod Serverless, the startup code runs outside the handler's event
        loop. Deferring initialization to the first request ensures all async
        components (tokenizer pool, serving engines, LoRA state) are created in
        the correct event loop context.
        """
        if not self._engines_initialized:
            logging.info("Initializing OpenAI serving engines...")
            await self._initialize_engines()
            self._engines_initialized = True
            logging.info("OpenAI serving engines initialized successfully")

    async def _initialize_engines(self):
        self.model_config = self.llm.model_config
        self.base_model_paths = [
            BaseModelPath(name=self.served_model_name, model_path=self.engine_args.model)
        ]

        self.serving_models = OpenAIServingModels(
            engine_client=self.llm,
            base_model_paths=self.base_model_paths,
            lora_modules=self.lora_adapters,
        )
        await self.serving_models.init_static_loras()

        # Get chat template from vLLM tokenizer if available
        chat_template = None
        if self.tokenizer and hasattr(self.tokenizer, 'tokenizer'):
            chat_template = self.tokenizer.tokenizer.chat_template

        self.openai_serving_render = OpenAIServingRender(
            model_config=self.llm.model_config,
            renderer=self.llm.renderer,
            model_registry=self.serving_models.registry,
            request_logger=None,
            chat_template=chat_template,
            chat_template_content_format="auto",
            trust_request_chat_template=os.getenv('TRUST_REQUEST_CHAT_TEMPLATE', 'false').lower() == 'true',
            enable_auto_tools=os.getenv('ENABLE_AUTO_TOOL_CHOICE', 'false').lower() == 'true',
            exclude_tools_when_tool_choice_none=os.getenv('EXCLUDE_TOOLS_WHEN_TOOL_CHOICE_NONE', 'false').lower() == 'true',
            tool_parser=os.getenv('TOOL_CALL_PARSER', "") or None,
            reasoning_parser=os.getenv('REASONING_PARSER', "") or None,
            log_error_stack=os.getenv('LOG_ERROR_STACK', 'false').lower() == 'true',
        )

        self.chat_engine = OpenAIServingChat(
            engine_client=self.llm,
            models=self.serving_models,
            response_role=self.response_role,
            openai_serving_render=self.openai_serving_render,
            request_logger=None,
            chat_template=chat_template,
            chat_template_content_format="auto",
            trust_request_chat_template=os.getenv('TRUST_REQUEST_CHAT_TEMPLATE', 'false').lower() == 'true',
            return_tokens_as_token_ids=os.getenv('RETURN_TOKENS_AS_TOKEN_IDS', 'false').lower() == 'true',
            reasoning_parser=os.getenv('REASONING_PARSER', "") or "",
            enable_auto_tools=os.getenv('ENABLE_AUTO_TOOL_CHOICE', 'false').lower() == 'true',
            exclude_tools_when_tool_choice_none=os.getenv('EXCLUDE_TOOLS_WHEN_TOOL_CHOICE_NONE', 'false').lower() == 'true',
            tool_parser=os.getenv('TOOL_CALL_PARSER', "") or None,
            enable_prompt_tokens_details=os.getenv('ENABLE_PROMPT_TOKENS_DETAILS', 'false').lower() == 'true',
            enable_force_include_usage=os.getenv('ENABLE_FORCE_INCLUDE_USAGE', 'false').lower() == 'true',
            enable_log_outputs=os.getenv('ENABLE_LOG_OUTPUTS', 'false').lower() == 'true',
        )
        self.completion_engine = OpenAIServingCompletion(
            engine_client=self.llm,
            models=self.serving_models,
            openai_serving_render=self.openai_serving_render,
            request_logger=None,
            return_tokens_as_token_ids=os.getenv('RETURN_TOKENS_AS_TOKEN_IDS', 'false').lower() == 'true',
            enable_prompt_tokens_details=os.getenv('ENABLE_PROMPT_TOKENS_DETAILS', 'false').lower() == 'true',
            enable_force_include_usage=os.getenv('ENABLE_FORCE_INCLUDE_USAGE', 'false').lower() == 'true',
        )
        self.responses_engine = OpenAIServingResponses(
            engine_client=self.llm,
            models=self.serving_models,
            openai_serving_render=self.openai_serving_render,
            request_logger=None,
            chat_template=chat_template,
            chat_template_content_format="auto",
            return_tokens_as_token_ids=os.getenv('RETURN_TOKENS_AS_TOKEN_IDS', 'false').lower() == 'true',
            reasoning_parser=os.getenv('REASONING_PARSER', "") or "",
            enable_auto_tools=os.getenv('ENABLE_AUTO_TOOL_CHOICE', 'false').lower() == 'true',
            tool_parser=os.getenv('TOOL_CALL_PARSER', "") or None,
            tool_server=None,
            enable_prompt_tokens_details=os.getenv('ENABLE_PROMPT_TOKENS_DETAILS', 'false').lower() == 'true',
            enable_force_include_usage=os.getenv('ENABLE_FORCE_INCLUDE_USAGE', 'false').lower() == 'true',
            enable_log_outputs=os.getenv('ENABLE_LOG_OUTPUTS', 'false').lower() == 'true',
        )
        self.messages_engine = AnthropicServingMessages(
            engine_client=self.llm,
            models=self.serving_models,
            response_role=self.response_role,
            openai_serving_render=self.openai_serving_render,
            request_logger=None,
            chat_template=chat_template,
            chat_template_content_format="auto",
            return_tokens_as_token_ids=os.getenv('RETURN_TOKENS_AS_TOKEN_IDS', 'false').lower() == 'true',
            reasoning_parser=os.getenv('REASONING_PARSER', "") or "",
            enable_auto_tools=os.getenv('ENABLE_AUTO_TOOL_CHOICE', 'false').lower() == 'true',
            tool_parser=os.getenv('TOOL_CALL_PARSER', "") or None,
            enable_prompt_tokens_details=os.getenv('ENABLE_PROMPT_TOKENS_DETAILS', 'false').lower() == 'true',
            enable_force_include_usage=os.getenv('ENABLE_FORCE_INCLUDE_USAGE', 'false').lower() == 'true',
        )

        warmup = getattr(self.chat_engine, 'warmup', None)
        if callable(warmup):
            result = warmup()
            if inspect.isawaitable(result):
                await result

    async def generate(self, openai_request: JobInput):
        # Ensure engines are ready (no-op if already initialized at startup)
        await self._ensure_engines_initialized()

        if openai_request.openai_route == "/v1/models":
            yield await self._handle_model_request()
        elif openai_request.openai_route in ["/v1/chat/completions", "/v1/completions"]:
            async for response in self._handle_chat_or_completion_request(openai_request):
                yield response
        elif openai_request.openai_route == "/v1/responses":
            async for response in self._handle_responses_request(openai_request):
                yield response
        elif openai_request.openai_route == "/v1/messages":
            async for response in self._handle_messages_request(openai_request):
                yield response
        else:
            yield create_error_response("Invalid route").model_dump()
    
    async def _handle_model_request(self):
        models = await self.serving_models.show_available_models()
        return models.model_dump()
    
    async def _handle_chat_or_completion_request(self, openai_request: JobInput):
        if openai_request.openai_route == "/v1/chat/completions":
            request_class = ChatCompletionRequest
            generator_function = self.chat_engine.create_chat_completion
        elif openai_request.openai_route == "/v1/completions":
            request_class = CompletionRequest
            generator_function = self.completion_engine.create_completion
        
        try:
            request = request_class(
                **openai_request.openai_input
            )
        except Exception as e:
            yield create_error_response(str(e)).model_dump()
            return
        
        dummy_request = DummyRequest()
        response_generator = await generator_function(request, raw_request=dummy_request)

        if not openai_request.openai_input.get("stream") or isinstance(response_generator, ErrorResponse):
            yield response_generator.model_dump()
        else:
            batch = []
            batch_token_counter = 0
            batch_size = BatchSize(self.default_batch_size, self.min_batch_size, self.batch_size_growth_factor)
        
            async for chunk_str in response_generator:
                if "data" in chunk_str:
                    if self.raw_openai_output:
                        data = chunk_str
                    elif "[DONE]" in chunk_str:
                        continue
                    else:
                        data = json.loads(chunk_str.removeprefix("data: ").rstrip("\n\n")) if not self.raw_openai_output else chunk_str
                    batch.append(data)
                    batch_token_counter += 1
                    if batch_token_counter >= batch_size.current_batch_size:
                        if self.raw_openai_output:
                            batch = "".join(batch)
                        yield batch
                        batch = []
                        batch_token_counter = 0
                        batch_size.update()
            if batch:
                if self.raw_openai_output:
                    batch = "".join(batch)
                yield batch

    async def _handle_responses_request(self, openai_request: JobInput):
        request_id = getattr(openai_request, "request_id", "unknown")

        try:
            request = ResponsesRequest(**openai_request.openai_input)
        except Exception as e:
            logging.error(
                "Invalid ResponsesRequest JSON: %s",
                e,
                extra={"request_id": request_id}
            )
            yield create_error_response(
                "Invalid request format",
                err_type="BadRequestError"
            ).model_dump()
            return

        dummy_request = DummyRequest()
        try:
            response = await self.responses_engine.create_responses(request, raw_request=dummy_request)
        except Exception as e:
            logging.error(
                "Failed to create Responses: %s",
                e,
                extra={"request_id": request_id},
                exc_info=True
            )
            yield create_error_response(
                "Internal server error during response generation",
                err_type="InternalServerError"
            ).model_dump()
            return

        if isinstance(response, (ErrorResponse, ResponsesResponse)):
            yield response.model_dump()
            return

        try:
            async for event in response:
                if not hasattr(event, "type"):
                    continue
                event_type = getattr(event, "type", "unknown")
                yield f"event: {event_type}\ndata: {event.model_dump_json(indent=None)}\n\n"
        except Exception as e:
            logging.error(
                "Error processing responses stream: %s",
                e,
                extra={"request_id": request_id},
                exc_info=True
            )
            error_payload = create_error_response(
                "Streaming response failed",
                err_type="InternalServerError"
            ).model_dump_json()
            yield f"event: error\ndata: {error_payload}\n\n"

    async def _handle_messages_request(self, openai_request: JobInput):
        request_id = getattr(openai_request, "request_id", "unknown")

        try:
            request = AnthropicMessagesRequest(**openai_request.openai_input)
        except Exception as e:
            logging.error(
                "Invalid AnthropicMessagesRequest: %s",
                e,
                extra={"request_id": request_id}
            )
            yield AnthropicErrorResponse(
                error=AnthropicError(
                    type="invalid_request_error",
                    message="Invalid request format"
                )
            ).model_dump()
            return

        dummy_request = DummyRequest()

        try:
            response = await self.messages_engine.create_messages(request, raw_request=dummy_request)
        except Exception as e:
            logging.error(
                "Failed to create messages: %s",
                e,
                extra={"request_id": request_id},
                exc_info=True
            )
            yield AnthropicErrorResponse(
                error=AnthropicError(
                    type="internal_error",
                    message="Failed to generate messages"
                )
            ).model_dump()
            return

        if isinstance(response, ErrorResponse):
            error_type = getattr(response, "type", "internal_error")
            error_message = getattr(response, "message", "Unknown error")
            yield AnthropicErrorResponse(
                error=AnthropicError(type=error_type, message=error_message)
            ).model_dump()
            return

        if isinstance(response, AnthropicMessagesResponse):
            yield response.model_dump(exclude_none=True)
            return

        try:
            async for chunk in response:
                yield chunk
        except Exception as e:
            logging.error(
                "Error streaming messages: %s",
                e,
                extra={"request_id": request_id},
                exc_info=True
            )
            error_payload = AnthropicErrorResponse(
                error=AnthropicError(
                    type="internal_error",
                    message="Error while streaming messages"
                )
            ).model_dump_json()
            yield f"event: error\ndata: {error_payload}\n\n"