# 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)