# 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 asyncio import json import multiprocessing import os import random import sys import threading import time import traceback from functools import partial from typing import Any, Generator, TypeVar import redis import redis.asyncio as async_redis import torch from tokenizers import Tokenizer from chameleon.inference.image_tokenizer import ImageTokenizer from chameleon.inference.loader import load_model from chameleon.inference.vocab import VocabInfo from chameleon.viewer.backend.data_types import WSMessageType from chameleon.viewer.backend.models.abstract_model import ( DEFAULT_IMAGE_CFG_IMAGE, DEFAULT_IMAGE_CFG_TEXT, DEFAULT_MULTIMODAL_CFG_IMAGE, DEFAULT_MULTIMODAL_CFG_TEXT, AbstractMultimodalGenerator, MixedSequenceType, StreamingImage, ) from chameleon.viewer.backend.models.chameleon_local import ( ChameleonForwardMixin, ChameleonTokenizationMixin, ) from chameleon.viewer.backend.utils import get_logger logger = get_logger(__name__) START = "START" T = TypeVar("T") def find_any(queue_by_id: dict[str, list]) -> str | None: for candidate_queue_id, candidate_queue in queue_by_id.items(): if len(candidate_queue) > 0: return candidate_queue_id return None class RedisQueue: def __init__(self, redis_client: redis.Redis, name: str, interval: float = 0.1): self.redis_client = redis_client self.name = name self.interval = interval self.lock = redis.lock.Lock(redis_client, f"lock_for_{name}") def reset(self): self.redis_client.set(self.name, json.dumps({})) try: self.lock.release() except redis.lock.LockError: pass def size(self) -> int: maybe_queue_by_id = self.redis_client.get(self.name) if maybe_queue_by_id is None: return 0 else: return len(json.loads(maybe_queue_by_id)) def clear(self, queue_id: str): with self.lock: maybe_queue_by_id = self.redis_client.get(self.name) if maybe_queue_by_id is None: queue_by_id: dict[str, list] = {} else: queue_by_id: dict[str, list] = json.loads(maybe_queue_by_id) queue_by_id[queue_id] = [] self.redis_client.set(self.name, json.dumps(queue_by_id)) def put(self, queue_id: str, value: T): logger.debug( "Thread %s: Starting PUT(%s) for %s", threading.get_ident(), self.name, queue_id, ) with self.lock: maybe_queue_by_id = self.redis_client.get(self.name) if maybe_queue_by_id is None: queue_by_id: dict[str, list[T]] = {} else: queue_by_id: dict[str, list[T]] = json.loads(maybe_queue_by_id) if queue_id not in queue_by_id: queue_by_id[queue_id] = [] queue_by_id[queue_id] = [value] + queue_by_id[queue_id] self.redis_client.set(self.name, json.dumps(queue_by_id)) logger.debug( "Thread %s: Finished PUT(%s) for %s", threading.get_ident(), self.name, queue_id, ) def get(self, queue_id: str | None) -> tuple[str, T]: """ Get the next value in the queue. if queue_id is None, will get a value from any queue if queue_id is not none, will wait to get a value from a specific queue """ logger.debug( "Thread %s: Starting GET(%s) for %s", threading.get_ident(), self.name, queue_id, ) while True: with self.lock: # Initialization hasn't happened, so wait for it to happen maybe_queue_by_id = self.redis_client.get(self.name) if maybe_queue_by_id is None: continue queue_by_id: dict[str, list[T]] = json.loads(maybe_queue_by_id) if queue_id is None: queue_id = find_any(queue_by_id) # Ensure a queue_id was found or that it already existed if queue_id is not None and queue_id in queue_by_id: queue = queue_by_id[queue_id] if len(queue) == 0: continue value = queue.pop(-1) # queue is mutated and queue_by_id references it, so this works self.redis_client.set(self.name, json.dumps(queue_by_id)) logger.debug( "Thread %s: Finished GET(%s) for %s", threading.get_ident(), self.name, queue_id, ) return queue_id, value time.sleep(self.interval) class AsyncRedisQueue: def __init__( self, redis_client: async_redis.Redis, name: str, interval: float = 0.1 ) -> None: self.redis_client = redis_client self.name = name self.interval = interval self.lock = async_redis.lock.Lock(redis_client, f"lock_for_{name}") async def reset(self): await self.redis_client.set(self.name, json.dumps({})) try: await self.lock.release() except async_redis.lock.LockError: pass async def size(self) -> int: maybe_queue_by_id = await self.redis_client.get(self.name) if maybe_queue_by_id is None: return 0 else: return len(json.loads(maybe_queue_by_id)) async def clear(self, queue_id: str): logger.debug( "ASYNC Thread %s: Starting CLEAR(%s) for %s", threading.get_ident(), self.name, queue_id, ) async with self.lock: maybe_queue_by_id = await self.redis_client.get(self.name) if maybe_queue_by_id is None: queue_by_id: dict[str, list] = {} else: queue_by_id: dict[str, list] = json.loads(maybe_queue_by_id) queue_by_id[queue_id] = [] await self.redis_client.set(self.name, json.dumps(queue_by_id)) logger.debug( "ASYNC Thread %s: Finished CLEAR(%s) for %s", threading.get_ident(), self.name, queue_id, ) async def put(self, queue_id: str, value: T): logger.debug( "ASYNC Thread %s: Starting PUT(%s) for %s", threading.get_ident(), self.name, queue_id, ) async with self.lock: maybe_queue_by_id = await self.redis_client.get(self.name) if maybe_queue_by_id is None: queue_by_id: dict[str, list[T]] = {} else: queue_by_id: dict[str, list[T]] = json.loads(maybe_queue_by_id) if queue_id not in queue_by_id: queue_by_id[queue_id] = [] queue_by_id[queue_id] = [value] + queue_by_id[queue_id] await self.redis_client.set(self.name, json.dumps(queue_by_id)) logger.debug( "ASYNC Thread %s: Finished PUT(%s) for %s", threading.get_ident(), self.name, queue_id, ) async def get(self, queue_id: str | None): """ Get the next value in the queue. if queue_id is None, will get a value from any queue if queue_id is not none, will wait to get a value from a specific queue """ logger.debug( "ASYNC Thread %s: Starting GET(%s) for %s", threading.get_ident(), self.name, queue_id, ) while True: async with self.lock: maybe_queue_by_id = await self.redis_client.get(self.name) if maybe_queue_by_id is None: continue queue_by_id: dict[str, list[T]] = json.loads(maybe_queue_by_id) if queue_id is None: queue_id = find_any(queue_by_id) # Ensure a queue_id was found or that it already existed if queue_id is not None and queue_id in queue_by_id: queue: list = queue_by_id[queue_id] if len(queue) == 0: continue value = queue.pop(-1) # queue is mutated and queue_by_id references it, so this works await self.redis_client.set(self.name, json.dumps(queue_by_id)) logger.debug( "ASYNC Thread %s: Finished GET(%s) for %s", threading.get_ident(), self.name, queue_id, ) return queue_id, value await asyncio.sleep(self.interval) class AsyncRedisCounter: def __init__(self, redis_client: async_redis.Redis, name: str) -> None: self.redis_client = redis_client self.name = name self.lock = async_redis.lock.Lock(redis_client, f"lock_for_{name}") async def reset(self) -> int: try: await self.lock.release() except async_redis.lock.LockError: pass await self.redis_client.set(self.name, 0) async def add(self, n: int) -> int: async with self.lock: current_val = await self.redis_client.get(self.name) if current_val is None: current_val = 0 else: current_val = int(current_val) new_val = current_val + n await self.redis_client.set(self.name, new_val) return new_val async def sub(self, n: int) -> int: async with self.lock: current_val = await self.redis_client.get(self.name) if current_val is None: raise ValueError("Invalid sub counter when counter does not exist") current_val = int(current_val) if current_val <= 0: raise ValueError("Invalid sub counter to counter that is already zero") new_val = current_val - n await self.redis_client.set(self.name, new_val) return new_val async def count(self) -> int: value = await self.redis_client.get(self.name) if value is None: return 0 else: return int(value) def distributed_workers( model_args: dict, master_address: str, master_port: str, world_size: int, rank: int, redis_port: int, worker_queues: dict[int, multiprocessing.Queue], ) -> None: redis_client = redis.Redis("redis", redis_port) request_queue = RedisQueue(redis_client, "request") response_queue = RedisQueue(redis_client, "response") os.environ["MASTER_ADDR"] = master_address os.environ["MASTER_PORT"] = str(master_port) torch.set_default_tensor_type("torch.cuda.FloatTensor") torch.distributed.init_process_group("nccl", rank=rank, world_size=world_size) assert rank == torch.distributed.get_rank() torch.cuda.set_device(rank) is_coord = rank == 0 worker = ChameleonWorker( rank=rank, model_path=model_args["model_path"], tokenizer_path=model_args["tokenizer_path"], additional_eos_tokens=model_args["additional_eos_tokens"], ) worker_id = id(worker) logger.info("Rank %s, master_port=%s worker=%s", rank, master_port, worker_id) step = 0 while True: step += 1 redis_client.set(f"status_rank_{rank}", "Pre-coordinator sync") if is_coord: distributed_objs = [request_queue.get(None)] logger.info("Objects from queue: %s", distributed_objs) for worker_rank in range(1, world_size): worker_message = {"message": START, "src": rank, "dst": worker_rank} logger.info("Rank %s Sending: %s", rank, worker_message) worker_queues[worker_rank].put(worker_message) else: distributed_objs = [None] logger.info("Rank %s worker %s waiting for rank 0", rank, worker_id) message_from_rank_0 = worker_queues[rank].get() logger.info( "Received message from rank 0 in rank %s: %s", rank, message_from_rank_0 ) if message_from_rank_0["message"] != START: raise ValueError( f"Unexpected message from rank 0: {message_from_rank_0['message']}" ) redis_client.set(f"status_rank_{rank}", "Post-coordinator sync") try: logger.info( "Broadcast Starting: Rank %s, worker %s, step %s", rank, worker_id, step, ) redis_client.set(f"status_rank_{rank}", "Pre-torch sync") torch.distributed.broadcast_object_list(distributed_objs, src=0) redis_client.set(f"status_rank_{rank}", "Post-torch sync") logger.info( "Broadcast Complete: Rank %s, worker %s, step %s", rank, worker_id, step, ) except RuntimeError as e: logger.error( "Rank %s, worker %s, step %s, Error detected in torch broadcast: %s", rank, worker_id, step, str(e), ) raise logger.info("rank %s, objs %s", rank, distributed_objs) queue_id, data = distributed_objs[0] mode = data.pop("mode") request_id = data.pop("request_id") assert queue_id == request_id tokenized_prompt = data.pop("tokenized_prompt") try: match mode: case WSMessageType.GENERATE_TEXT: generator_fn = partial( worker._generate_text_streaming, tokenized_prompt, **data ) case WSMessageType.GENERATE_IMAGE: generator_fn = partial( worker._generate_image_streaming, tokenized_prompt, **data ) case WSMessageType.GENERATE_MULTIMODAL: generator_fn = partial( worker._generate_multimodal_streaming, tokenized_prompt, **data ) case _: logger.error( "Encountered unknown mode, crashing the program: %s", mode ) response_queue.put( queue_id, {"error": True, "final": True, "message": mode} ) raise ValueError("Unknown mode") logger.info("Rank: %s, Processing request: %s", rank, request_id) i = 0 redis_client.set(f"status_rank_{rank}", "Pre-generate") for output in generator_fn(): i += 1 if is_coord: response = {"final": False, "output": output, "error": False} logger.info( "Rank: %s, Adding to response queue: %.100s", rank, response, ) redis_client.set(f"status_rank_{rank}", f"Generate Pre Put {i}") response_queue.put(queue_id, response) redis_client.set(f"status_rank_{rank}", f"Generate Post Put {i}") else: redis_client.set(f"status_rank_{rank}", f"Generate {i}") redis_client.set(f"step_on_rank_{rank}", i) redis_client.set(f"status_rank_{rank}", "Post-generate") if is_coord: logger.info("Rank: %s, Adding final result to output queue", rank) response_queue.put(queue_id, {"final": True, "error": False}) except torch.cuda.OutOfMemoryError as e: logger.error("Encountered OOM, crashing the program: %s", e) response_queue.put( queue_id, {"error": True, "final": True, "message": str(e)} ) crash_program() except RuntimeError as e: message = str(e) if "CUDA" in message: logger.error("Encountered CUDA error, crashing the program: %s", e) response_queue.put( queue_id, {"error": True, "final": True, "message": str(e)} ) crash_program() else: logger.error( "Encountered unexpected runtime error, crashing the program: %s %s", e, traceback.format_exc(), ) response_queue.put( queue_id, {"error": True, "final": True, "message": str(e)} ) crash_program() except Exception as e: logger.error( "Encountered unexpected exception: %s %s", str(e), traceback.format_exc(), ) response_queue.put( queue_id, {"error": True, "final": True, "message": str(e)} ) crash_program() class ChameleonWorker(ChameleonForwardMixin): def __init__( self, *, rank: int, model_path: str, tokenizer_path: str, additional_eos_tokens: list[str] | None, ) -> None: self.rank = rank self.model_path = model_path self.additional_eos_tokens = additional_eos_tokens torch.set_default_device(f"cuda:{rank}") self.model = load_model(model_path, rank) self.tokenizer = Tokenizer.from_file(str(tokenizer_path)) self.vocab = VocabInfo(json.load(open(tokenizer_path))["model"]["vocab"]) logger.info( "Rank: %s, Model loaded in worker_obj: %s", rank, id(self), ) def crash_program() -> None: logger.error( "Crashing the program as instructed, likely due to distributed worker failures" ) sys.exit(1) class ChameleonDistributedGenerator(AbstractMultimodalGenerator, ChameleonTokenizationMixin): def __init__( self, *, world_size: int, model_path: str, master_port: int, tokenizer_path: str, vqgan_config_path: str, vqgan_ckpt_path: str | None = None, master_address: str = "0.0.0.0", additional_eos_tokens: list[str] | None = None, redis_port: int | None = None, ) -> None: self.master_port = master_port self.master_address = master_address self.additional_eos_tokens = additional_eos_tokens logger.info("Loading tokenizer...") tokenizer_path = tokenizer_path self.tokenizer = Tokenizer.from_file(str(tokenizer_path)) self.vocab = VocabInfo(json.load(open(tokenizer_path))["model"]["vocab"]) logger.info("Loading VQGAN...") self.image_tokenizer = ImageTokenizer(vqgan_config_path, vqgan_ckpt_path) self.redis_port = redis_port self.redis_pool = async_redis.ConnectionPool.from_url( f"redis://redis:{redis_port}" ) self.redis_client = async_redis.Redis.from_pool(self.redis_pool) self.request_queue = AsyncRedisQueue(self.redis_client, "request") self.response_queue = AsyncRedisQueue(self.redis_client, "response") self.worker_queues: dict[int, multiprocessing.Queue] = { rank: multiprocessing.Queue() for rank in range(world_size) } self.procs: list[multiprocessing.Process] = [] model_args = { "model_path": model_path, "master_address": master_address, "master_port": master_port, "tokenizer_path": tokenizer_path, "additional_eos_tokens": additional_eos_tokens, } logger.info("Launching paralle model with world_size=%s", world_size) for i in range(world_size): proc = multiprocessing.Process( target=distributed_workers, args=( model_args, master_address, master_port, world_size, i, self.redis_port, self.worker_queues, ), daemon=True, ) self.procs.append(proc) proc.start() def check_error(self, output: dict) -> None: if output["error"]: import sys print(f"check_error({output})", file=sys.stderr) self.kill_procs() logger.error( "COORDINATOR: Encountered error in managed processes, exiting: %s", output, ) crash_program() def __del__(self) -> None: self.kill_procs(error=False) def kill_procs(self, error: bool = True) -> None: if error: log_fn = logger.error else: log_fn = logger.info log_fn("Error encountered, killing worker procs: %s", self.procs) for p in self.procs: try: log_fn("Killing: %s", p) p.kill() except: log_fn("Encountered issue killing process and ignoring: %s", p) # ALLOW_ANY(get_next_output.return) async def get_next_output(self, request_id: str) -> Any: logger.info("Waiting for response for request_id=%s", request_id) queue_id, output = await self.response_queue.get(request_id) assert queue_id == request_id return output async def generate_text_streaming( self, prompt: MixedSequenceType, max_gen_tokens: int = 256, temp: float = 1.0, top_p: float = 0.8, repetition_penalty: float = 1.2, seed: int | None = None, debug: dict | None = None, ) -> Generator[str, None, None]: tokenized_prompt = self.tokens_from_inputs(prompt) request_id = f"request_{random.randint(100_000, 200_000)}" if seed is None: seed = random.randint(1, 2048) if debug is not None: debug["seed"] = seed if len(tokenized_prompt) > (4096 - 3): yield "ERROR: Your input exceeds the model's context length of 4096. Note that images consume 1024 tokens whether in input or output." return assert not isinstance(tokenized_prompt, torch.Tensor) request = { "mode": WSMessageType.GENERATE_TEXT.value, "request_id": request_id, "tokenized_prompt": tokenized_prompt, "max_gen_tokens": max_gen_tokens, "temp": temp, "top_p": top_p, "repetition_penalty": repetition_penalty, "seed": seed, } logger.info( "Sending request_id=%s: %s", request_id, request, ) await asyncio.gather( self.request_queue.clear(request_id), self.response_queue.clear(request_id), ) logger.info("Cleared request/response queue for %s", request_id) await self.request_queue.put(request_id, request) logger.info("Sent request to coordinator %s", request_id) try: while True: output = await self.get_next_output(request_id) logger.info("Received response for %s", request_id) self.check_error(output) if output["final"]: break n_outs = len(output["output"]) if n_outs != 1: logger.error( "Encountered unexpected number of %s arguments in: %s", n_outs, output["output"], ) tokens = output["output"] assert not isinstance(tokens, torch.Tensor) logger.info("output info: type=%s, value=%.20s", type(tokens), tokens) yield self.tokenizer.decode(tokens) finally: logger.info("Cleaning up queues in request_id=%s", request_id) await asyncio.gather( self.request_queue.clear(request_id), self.response_queue.clear(request_id), ) logger.info("Completed cleaning for request_id=%s", request_id) async def generate_image_streaming( self, prompt: MixedSequenceType, temp: float = 1.0, top_p: float = 0.8, cfg_image_weight: float = DEFAULT_IMAGE_CFG_IMAGE, cfg_text_weight: float = DEFAULT_IMAGE_CFG_TEXT, yield_every_n: int = 32, debug: dict | None = None, seed: int | None = None, ) -> Generator[StreamingImage, None, None]: tokenized_prompt = self.tokens_from_inputs(prompt) tokenized_prompt.append(self.vocab.begin_image) assert not isinstance(tokenized_prompt, torch.Tensor) request_id = f"request_{random.randint(100_000, 200_000)}" if seed is None: seed = random.randint(1, 2048) if debug is not None: debug["seed"] = seed if len(tokenized_prompt) > (4096 - 3 - 1024): yield "ERROR: Your input exceeds the model's context length of 4096. Note that images consume 1024 tokens whether in input or output." return request = { "mode": WSMessageType.GENERATE_IMAGE.value, "request_id": request_id, "tokenized_prompt": tokenized_prompt, "cfg_image_weight": cfg_image_weight, "cfg_text_weight": cfg_text_weight, "yield_every_n": yield_every_n, "temp": temp, "top_p": top_p, "seed": seed, } logger.info( "Sending request_id=%s: %s", request_id, request, ) await asyncio.gather( self.request_queue.clear(request_id), self.response_queue.clear(request_id), ) logger.info("Cleared request/response queue for %s", request_id) await self.request_queue.put(request_id, request) logger.info("Sent request to coordinator %s", request_id) try: while True: output = await self.get_next_output(request_id) logger.info("Received response for %s", request_id) self.check_error(output) if output["final"]: break n_outs = len(output["output"]) if n_outs != 2: logger.error( "Encountered unexpected number of %s arguments in: %s", n_outs, output["output"], ) tokens, final = output["output"] assert not isinstance(tokens, torch.Tensor) yield StreamingImage( image=self.pillow_from_bpe_tokens(torch.tensor(tokens)), final=final ) finally: logger.info("Cleaning up queues in request_id=%s", request_id) await asyncio.gather( self.request_queue.clear(request_id), self.response_queue.clear(request_id), ) logger.info("Completed cleaning for request_id=%s", request_id) async def generate_multimodal_streaming( self, prompt: MixedSequenceType, temp: float = 1.0, top_p: float = 0.8, cfg_image_weight: float = DEFAULT_MULTIMODAL_CFG_IMAGE, cfg_text_weight: float = DEFAULT_MULTIMODAL_CFG_TEXT, yield_every_n: int = 32, max_gen_tokens: int = 4096, repetition_penalty: float = 1.2, suffix_tokens: list[str] | None = None, seed: int | None = None, debug: dict | None = None, ) -> Generator[MixedSequenceType, None, None]: tokenized_prompt = self.tokens_from_inputs(prompt, suffix_tokens=suffix_tokens) assert not isinstance(tokenized_prompt, torch.Tensor) request_id = f"request_{random.randint(100_000, 200_000)}" if seed is None: seed = random.randint(1, 2048) if debug is not None: debug["seed"] = seed if len(tokenized_prompt) > (4096 - 3): yield "ERROR: Your input exceeds the model's context length of 4096. Note that images consume 1024 tokens." return request = { "mode": WSMessageType.GENERATE_MULTIMODAL.value, "request_id": request_id, "tokenized_prompt": tokenized_prompt, "cfg_image_weight": cfg_image_weight, "cfg_text_weight": cfg_text_weight, "repetition_penalty": repetition_penalty, "yield_every_n": yield_every_n, "max_gen_tokens": max_gen_tokens, "temp": temp, "top_p": top_p, "seed": seed, } logger.info( "Sending request_id=%s: %s", request_id, request, ) await asyncio.gather( self.request_queue.clear(request_id), self.response_queue.clear(request_id), ) logger.info("Cleared request/response queue for %s", request_id) await self.request_queue.put(request_id, request) logger.info("Sent request to coordinator %s", request_id) try: while True: output = await self.get_next_output(request_id) logger.info("Received response for %s", request_id) self.check_error(output) if output["final"]: break n_outs = len(output["output"]) if n_outs != 3: logger.error( "Encountered unexpected number of %s arguments in: %s", n_outs, output["output"], ) token_type, tokens, image_is_final = output["output"] assert not isinstance(tokens, torch.Tensor) match token_type: case "TEXT": yield self.tokenizer.decode(tokens) case "IMAGE": yield StreamingImage( image=self.pillow_from_bpe_tokens(torch.tensor(tokens)), final=image_is_final, ) case _: raise ValueError("Unknown token type") finally: logger.info("Cleaning up queues in request_id=%s", request_id) await self.request_queue.clear(request_id) await self.response_queue.clear(request_id)