serve.py

import asyncio
import logging
import os
import time
from pprint import pprint
from threading import Thread
from typing import Any, Dict, List

# isort: off
from unsloth import (
    FastLanguageModel,
    FastModel,
    FastVisionModel,
    is_bfloat16_supported,
)  # noqa: E402
from unsloth.chat_templates import get_chat_template  # noqa: E402

# isort: on

from fastapi import FastAPI, Request
from openai.types.chat.chat_completion import ChatCompletion
from openai.types.chat.chat_completion import Choice as ChatCompletionChoice
from openai.types.chat.chat_completion_chunk import ChatCompletionChunk
from openai.types.chat.chat_completion_chunk import Choice as ChatCompletionChunkChoice
from openai.types.chat.chat_completion_chunk import ChoiceDelta
from openai.types.chat.chat_completion_message import ChatCompletionMessage
from openai.types.chat.completion_create_params import CompletionCreateParams
from pydantic import TypeAdapter
from ray import serve
from sse_starlette import EventSourceResponse
from starlette.responses import JSONResponse
from transformers.generation.streamers import AsyncTextIteratorStreamer
from transformers.image_utils import load_image

dtype = (
    None  # None for auto detection. Float16 for Tesla T4, V100, Bfloat16 for Ampere+
)
load_in_4bit = True  # Use 4bit quantization to reduce memory usage. Can be False.
max_seq_length = 2048  # Supports RoPE Scaling interally, so choose any!
# max_seq_length = 4096 # Choose any! We auto support RoPE Scaling internally!


logger = logging.getLogger("ray.serve")

os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "expandable_segments:True"

app = FastAPI()

# middlewares = [
#     middleware
#     for middleware in ConnexionMiddleware.default_middlewares
#     if middleware is not SecurityMiddleware
# ]

# connexion_app = AsyncApp(import_name=__name__, middlewares=middlewares)

# connexion_app.add_api(
#     # "api/openai/v1/openapi/openapi.yaml",
#     "api/v1/openapi/openapi.yaml",
#     # base_path="/openai/v1",
#     base_path="/v1",
#     pythonic_params=True,
#     resolver_error=501,
# )

# # fastapi_app.mount("/api", ConnexionMiddleware(app=connexion_app, import_name=__name__))
# # app.mount("/api", ConnexionMiddleware(app=connexion_app, import_name=__name__))
# app.mount(
#     "/",
#     ConnexionMiddleware(
#         app=connexion_app,
#         import_name=__name__,
#         # middlewares=middlewares,
#     ),
# )


@serve.deployment(
    autoscaling_config={
        # https://docs.ray.io/en/latest/serve/advanced-guides/advanced-autoscaling.html#required-define-upper-and-lower-autoscaling-limits
        "max_replicas": 1,
        "min_replicas": 1,  # TOOD: set to 0
        "target_ongoing_requests": 2,  # https://docs.ray.io/en/latest/serve/advanced-guides/advanced-autoscaling.html#target-ongoing-requests-default-2
    },
    max_ongoing_requests=5,  # https://docs.ray.io/en/latest/serve/advanced-guides/advanced-autoscaling.html#max-ongoing-requests-default-5
    ray_actor_options={"num_gpus": 1},
)
@serve.ingress(app)
class ModelDeployment:
    def __init__(
        self,
        model_name: str,
    ):
        self.model_name = model_name

        model, processor = FastModel.from_pretrained(
            load_in_4bit=load_in_4bit,
            max_seq_length=max_seq_length,
            model_name=self.model_name,
        )

        # with open("chat_template.txt", "r") as f:
        #     processor.chat_template = f.read()
        #     processor.tokenizer.chat_template = processor.chat_template

        FastModel.for_inference(model)  # Enable native 2x faster inference

        self.model = model
        self.processor = processor

    def reconfigure(self, config: Dict[str, Any]):
        print("=== reconfigure ===")
        print("config:")
        print(config)
        # https://docs.ray.io/en/latest/serve/production-guide/config.html#dynamically-change-parameters-without-restarting-replicas-user-config

    @app.post("/v1/chat/completions")
    async def create_chat_completion(self, body: dict, raw_request: Request):
        """Creates a model response for the given chat conversation. Learn more in the [text generation](/docs/guides/text-generation), [vision](/docs/guides/vision), and [audio](/docs/guides/audio) guides.  Parameter support can differ depending on the model used to generate the response, particularly for newer reasoning models. Parameters that are only supported for reasoning models are noted below. For the current state of  unsupported parameters in reasoning models,  [refer to the reasoning guide](/docs/guides/reasoning).

        # noqa: E501

        :param create_chat_completion_request:
        :type create_chat_completion_request: dict | bytes

        :rtype: Union[CreateChatCompletionResponse, Tuple[CreateChatCompletionResponse, int], Tuple[CreateChatCompletionResponse, int, Dict[str, str]]
        """
        print("=== create_chat_completion ===")

        print("body:")
        pprint(body)

        ta = TypeAdapter(CompletionCreateParams)

        print("ta.validate_python...")
        pprint(ta.validate_python(body))

        max_new_tokens = body.get("max_completion_tokens", body.get("max_tokens"))
        messages = body.get("messages")
        model_name = body.get("model")
        stream = body.get("stream", False)
        temperature = body.get("temperature")
        tools = body.get("tools")

        images = []

        for message in messages:
            for content in message["content"]:
                if "type" in content and content["type"] == "image_url":
                    image_url = content["image_url"]["url"]
                    image = load_image(image_url)
                    images.append(image)

                    content["type"] = "image"
                    del content["image_url"]

        images = images if images else None

        if model_name != self.model_name:
            # adapter_path = model_name
            # self.model.load_adapter(adapter_path)

            return JSONResponse(content={"error": "Model not found"}, status_code=404)

        prompt = self.processor.apply_chat_template(
            add_generation_prompt=True,
            conversation=messages,
            # documents=documents,
            tools=tools,
            tokenize=False,  # Return string instead of token IDs
        )

        print("prompt:")
        print(prompt)

        if images:
            inputs = self.processor(text=prompt, images=images, return_tensors="pt")
        else:
            inputs = self.processor(text=prompt, return_tensors="pt")

        inputs = inputs.to(self.model.device)
        input_ids = inputs.input_ids

        class GeneratorThread(Thread):
            """Thread to generate completions in the background."""

            def __init__(self, model, **generation_kwargs):
                super().__init__()

                self.chat_completion = None
                self.generation_kwargs = generation_kwargs
                self.model = model

            def run(self):
                import torch
                import torch._dynamo.config

                try:
                    try:
                        self.generated_ids = self.model.generate(
                            **self.generation_kwargs
                        )

                    except torch._dynamo.exc.BackendCompilerFailed as e:
                        print(e)
                        print("Disabling dynamo...")

                        torch._dynamo.config.disable = True

                        self.generated_ids = self.model.generate(
                            **self.generation_kwargs
                        )

                except Exception as e:
                    print(e)
                    print("Warning: Exception in GeneratorThread")
                    self.generated_ids = []

            def join(self, timeout=None):
                super().join()

                return self.generated_ids

        decode_kwargs = dict(skip_special_tokens=True)

        streamer = (
            AsyncTextIteratorStreamer(
                self.processor,
                skip_prompt=True,
                **decode_kwargs,
            )
            if stream
            else None
        )

        generation_kwargs = dict(
            **inputs,
            max_new_tokens=max_new_tokens,
            streamer=streamer,
            temperature=temperature,
            use_cache=True,
        )

        thread = GeneratorThread(self.model, **generation_kwargs)
        thread.start()

        if stream:

            async def event_publisher():
                i = 0

                try:
                    async for new_text in streamer:
                        print("new_text:")
                        print(new_text)

                        choices: List[ChatCompletionChunkChoice] = [
                            ChatCompletionChunkChoice(
                                _request_id=None,
                                delta=ChoiceDelta(
                                    _request_id=None,
                                    content=new_text,
                                    function_call=None,
                                    refusal=None,
                                    role="assistant",
                                    tool_calls=None,
                                ),
                                finish_reason=None,
                                index=0,
                                logprobs=None,
                            )
                        ]

                        chat_completion_chunk = ChatCompletionChunk(
                            _request_id=None,
                            choices=choices,
                            created=int(time.time()),
                            id=str(i),
                            model=model_name,
                            object="chat.completion.chunk",
                            service_tier=None,
                            system_fingerprint=None,
                            usage=None,
                        )

                        yield chat_completion_chunk.model_dump_json()

                        i += 1

                except asyncio.CancelledError as e:
                    print("Disconnected from client (via refresh/close)")
                    raise e

                except Exception as e:
                    print(f"Exception: {e}")
                    raise e

            return EventSourceResponse(event_publisher())

        generated_ids = thread.join()
        input_length = input_ids.shape[1]

        batch_decoded_outputs = self.processor.batch_decode(
            generated_ids[:, input_length:],
            skip_special_tokens=True,
        )

        choices: List[ChatCompletionChoice] = []

        for i, response in enumerate(batch_decoded_outputs):
            print("response:")
            print(response)

            # try:
            # response = json.loads(response)

            #         finish_reason: str = response.get("finish_reason")
            #         tool_calls_json = response.get("tool_calls")
            #         tool_calls: List[ToolCall] = []

            #         for tool_call_json in tool_calls_json:
            #             tool_call = ToolCall(
            #                 function=FunctionToolCallArguments(
            #                     arguments=tool_call_json.get("arguments"),
            #                     name=tool_call_json.get("name"),
            #                 ),
            #                 id=tool_call_json.get("id"),
            #                 type="function",
            #             )

            #             tool_calls.append(tool_call)

            #         message: ChatMessage = ChatMessage(
            #             role="assistant",
            #             tool_calls=tool_calls,
            #         )

            #     except json.JSONDecodeError:
            #         finish_reason: str = "stop"
            #         message: ChatMessage = ChatMessage(
            #             role="assistant",
            #             content=response,
            #         )

            message = ChatCompletionMessage(
                audio=None,
                content=response,
                refusal=None,
                role="assistant",
                tool_calls=None,
            )

            choices.append(
                ChatCompletionChoice(
                    index=i,
                    finish_reason="stop",
                    logprobs=None,
                    message=message,
                )
            )

        chat_completion = ChatCompletion(
            choices=choices,
            created=int(time.time()),
            id="1",
            model=model_name,
            object="chat.completion",
            service_tier=None,
            system_fingerprint=None,
            usage=None,
        )

        return chat_completion.model_dump(mode="json")


def build_app(cli_args: Dict[str, str]) -> serve.Application:
    """Builds the Serve app based on CLI arguments."""
    return ModelDeployment.options().bind(
        cli_args.get("model_name"),
    )


# uv run serve run serve:build_app model_name="HuggingFaceTB/SmolVLM-Instruct"
# uv run serve run serve:build_app model_name="unsloth/SmolLM2-135M-Instruct-bnb-4bit"