chameleon/viewer/backend/models/service.py

# Copyright (c) Meta Platforms, Inc. and affiliates.
#
# This source code is licensed under the Chameleon License found in the
# LICENSE file in the root directory of this source tree.

import base64
import io
import socket
import subprocess
import time
from functools import partial

import fastapi
import PIL
import pydantic
import redis.asyncio as async_redis
import uvicorn
from fastapi import FastAPI, WebSocket, WebSocketDisconnect, WebSocketException
from websockets.exceptions import ConnectionClosedError, ConnectionClosedOK

from chameleon.viewer.backend.data_types import (
    Content,
    ContentType,
    NoOptionsForComplete,
    NoOptionsForFull,
    NoOptionsForPartial,
    NoOptionsForQueueStatus,
    WSMessageType,
    WSMultimodalMessage,
)
from chameleon.viewer.backend.models.abstract_model import (
    AbstractMultimodalGenerator,
    StreamingImage,
)
from chameleon.viewer.backend.models.chameleon_distributed import AsyncRedisCounter
from chameleon.viewer.backend.utils import get_logger

logger = get_logger(__name__)


def nvidia_smi() -> str:
    return subprocess.check_output(["nvidia-smi"], text=True)


async def await_generate_message(websocket: WebSocket) -> WSMultimodalMessage:
    while True:
        rec_message = await websocket.receive_json()
        try:
            maybe_message = WSMultimodalMessage.parse_obj(rec_message)
        except pydantic.ValidationError:
            maybe_message = None
            logger.info("Got invalid message", maybe_message)
        if maybe_message is not None:
            return maybe_message


async def async_acquire_lock(
    *,
    websocket: WebSocket,
    counter: AsyncRedisCounter,
    lock: async_redis.lock.Lock,
    interval=0.1,
    status_interval=1,
    hostname: str | None = None,
):
    start = time.time()
    await counter.add(1)
    while True:
        acquired = await lock.acquire(blocking_timeout=interval)
        if acquired:
            break
        elapsed = time.time() - start
        if elapsed > status_interval:
            n_requests = await counter.count()
            message = WSMultimodalMessage(
                message_type=WSMessageType.QUEUE_STATUS,
                content=[
                    Content(
                        content_type=ContentType.TEXT,
                        content=f"n_requests={n_requests}",
                    )
                ],
                options=NoOptionsForQueueStatus(),
                debug_info={"hostname": hostname},
            ).dict()
            await websocket.send_json(message)
            start = time.time()
    await counter.sub(1)


COORDINATOR = "coordinator"


def web_app(
    generator: AbstractMultimodalGenerator,
    debug: bool = True,
    redis_port: int | None = None,
) -> FastAPI:
    app = FastAPI(debug=debug)
    if redis_port is None:
        redis_client = None
        redis_lock = None
        queue_counter = None
    else:
        redis_client = async_redis.Redis.from_url(f"redis://redis:{redis_port}")
        redis_lock = async_redis.lock.Lock(redis_client, COORDINATOR)
        queue_counter = AsyncRedisCounter(redis_client, "count_pending")
    hostname = socket.gethostname()

    @app.get("/api/2.0/status")
    def alive() -> dict:
        return {
            "status": "alive",
            "hostname": hostname,
            "nvidia-smi": nvidia_smi(),
        }

    @app.websocket("/ws/chameleon/v2/{client_id}")
    async def websocket_chameleon_v2(*, websocket: WebSocket, client_id: str):
        logger.info("Requested client_id: %s", client_id)
        await websocket.accept()
        logger.info("Client opened %s with generator id %s", client_id, id(generator))

        try:
            while True:
                generate_message = await await_generate_message(websocket)
                logger.info("Got generate message: %s", str(generate_message)[:300])
                parsed_prompt = []
                for c in generate_message.content:
                    match c.content_type:
                        case ContentType.TEXT:
                            parsed_prompt.append(c.content)
                        case ContentType.IMAGE:
                            image_parts = c.content.split(",", 1)
                            if len(image_parts) < 2:
                                logger.error(
                                    "Encountered invalid image: %s", image_parts
                                )
                                raise WebSocketException(
                                    code=fastapi.status.WS_1008_POLICY_VIOLATION,
                                    reason=f"Invalid image: {image_parts}",
                                )
                            image_data = image_parts[1]
                            base64_image = base64.b64decode(image_data)
                            image_file = io.BytesIO(base64_image)
                            parsed_prompt.append(PIL.Image.open(image_file))
                        case _:
                            raise ValueError("Unknown content type")
                logger.info("Prompt: %s", parsed_prompt)
                partial_outputs = []
                final_contents: list[Content] = []

                match generate_message.message_type:
                    case WSMessageType.GENERATE_TEXT:
                        output_generator = generator.generate_text_streaming
                    case WSMessageType.GENERATE_IMAGE:
                        output_generator = generator.generate_image_streaming
                    case WSMessageType.GENERATE_MULTIMODAL:
                        output_generator = generator.generate_multimodal_streaming
                    case _:
                        raise WebSocketException(
                            code=fastapi.status.WS_1008_POLICY_VIOLATION,
                            reason="Unknown message type",
                        )

                logger.info(
                    "Acquiring lock for client %s generation with options: %s",
                    client_id,
                    generate_message.options,
                )
                option_args = generate_message.options.dict()
                debug_info = {"hostname": hostname}
                del option_args["message_type"]
                output_generator = partial(
                    output_generator,
                    **option_args,
                    debug=debug_info,
                )
                if redis_lock is not None:
                    await async_acquire_lock(
                        websocket=websocket,
                        lock=redis_lock,
                        hostname=hostname,
                        counter=queue_counter,
                    )
                    await redis_client.set("has_lock", client_id)

                logger.info(
                    "Starting locked generation for client %s with options: %s",
                    client_id,
                    generate_message.options,
                )
                try:
                    async for output_token in output_generator(parsed_prompt):
                        if isinstance(output_token, str):
                            content_type = ContentType.TEXT
                            content = output_token
                            message_type = WSMessageType.PARTIAL_OUTPUT
                            options = NoOptionsForPartial()
                            partial_outputs.extend(output_token)
                        elif isinstance(output_token, StreamingImage):
                            content_type = ContentType.IMAGE
                            image = output_token.image
                            img_io = io.BytesIO()
                            image.save(img_io, format="png")
                            content = (
                                "data:image/png;base64,"
                                + base64.b64encode(img_io.getvalue()).decode()
                            )
                            if output_token.final:
                                message_type = WSMessageType.FULL_OUTPUT
                                options = NoOptionsForFull()
                            else:
                                message_type = WSMessageType.PARTIAL_OUTPUT
                                options = NoOptionsForPartial()

                            if output_token.final:
                                partial_outputs.append(output_token.image)
                        else:
                            raise ValueError(f"Invalid output_token: {output_token}")

                        message_content = Content(
                            content_type=content_type, content=content
                        )
                        match content_type:
                            case ContentType.TEXT:
                                final_contents.append(message_content)
                            case ContentType.IMAGE:
                                if message_type == WSMessageType.FULL_OUTPUT:
                                    final_contents.append(message_content)
                            case _:
                                pass

                        message = WSMultimodalMessage(
                            message_type=message_type,
                            content=[message_content],
                            options=options,
                            debug_info=debug_info,
                        ).dict()
                        await websocket.send_json(message)
                finally:
                    if redis_lock is not None:
                        logger.info(
                            "Attempting release of lock for client %s generation with options: %s",
                            client_id,
                            generate_message.options,
                        )
                        owned = await redis_lock.owned()
                        if owned:
                            await redis_client.set("has_lock", "")
                            try:
                                await redis_lock.release()
                            except async_redis.lock.LockError:
                                pass

                        logger.info(
                            "Released lock for client %s generation with options: %s",
                            client_id,
                            generate_message.options,
                        )
                await websocket.send_json(
                    WSMultimodalMessage(
                        message_type=WSMessageType.COMPLETE,
                        content=final_contents,
                        options=NoOptionsForComplete(),
                        debug_info=debug_info,
                    ).dict()
                )
        except WebSocketDisconnect:
            logger.info("Client disconnected %s", client_id)
        except ConnectionClosedError:
            logger.info("Client forced a close %s", client_id)
        except ConnectionClosedOK:
            logger.info("Connection closed ok %s", client_id)
        finally:
            if redis_lock is not None:
                logger.info("Checking for client holding lock: %s", client_id)
                owned = await redis_lock.owned()
                if owned:
                    try:
                        logger.info("Attempted to release owned lock: %s", client_id)
                        await redis_lock.release()
                    except async_redis.lock.LockError:
                        pass
                    await redis_client.set("has_lock", "")

    return app


def serve(
    model: AbstractMultimodalGenerator,
    host: str,
    port: int,
    debug: bool = True,
    redis_port: int | None = None,
) -> None:
    app = web_app(model, debug=debug, redis_port=redis_port)
    # TODO: convert this to a subprocess call so enable more
    # uvicorn features like multiple workers
    uvicorn.run(app, host=host, port=port)