"""Audio generation application on modal.""" from enum import Enum import itertools import json import os import select import time import traceback from typing import Optional from uuid import uuid4 import threading import queue import random import fcntl import modal import numpy as np import torch from datadog import initialize, statsd from suno_utils.audio import Audio from suno_utils.gpt.engine import GenerationConfig from suno_utils.gpt.generation_engine import align_codes, make_request, unshift_arrays_v2 from suno_utils.tasks.codec_engine import CodecEngine, Request, EMBEDDING_RATE from suno_utils.tasks.hoot import clean_text from suno_utils.worker.loader import S3Loader from suno_utils.worker.schema import HistoryPrompt, QueueItem from suno_utils.worker.utils import print_gpu_memory_usage, recursive_ls_dir from suno_utils.worker.tracing import serialize_context, distributed_trace, tracer from suno_utils.gpt.rpc_zmq import start_service_processes from suno_utils.worker.modal_model_configs import MODEL_CONFIG_DICT from suno_utils.worker.event_queue import EventType, Event, INBOUND_EVENT_QUEUE_NAME from suno_utils.worker.modal_base import get_modal_base_image_with_flash_attention from suno_utils.worker.modal_model_volume import MODEL_STORE_VOLUME_DIR, model_store_volume from suno_utils.worker.utils import ffmpeg_stream_encode, ffmpeg_stream_encode_opus_webm ############## CHANGE THESE ############## MODEL_CONFIG = MODEL_CONFIG_DICT["13b_special_8"] # Production models: # v2, # 7b_ipo # 13b_special_8, 13b_short, 13b_upload_4 # 30b_t6 # Dev models: # 30b, 30b_t6, 30b_classical, 30b_dance # 7b_dpo, 7b_ipo # v2 # 13b_special_8, 13b_upload_4, 13b_short DEPLOYMENT_TYPE = "dev" # priority, staging, dev, prod, msft ########################################## DURATION = MODEL_CONFIG.duration MODEL = MODEL_CONFIG.model # v2, v2_ft13 for prod for now COMPILE = MODEL_CONFIG.compile GPU = MODEL_CONFIG.gpu # A100, A10G MODEL_VERSION = MODEL_CONFIG.model_version IS_3B = "v2" in MODEL IS_7B = "7b" in MODEL IS_13B = "13b" in MODEL IS_30B = "30b" in MODEL if IS_3B: from suno_utils.gpt import chirp_v2 else: from suno_utils.gpt import chirp_v2_5 as chirp_v2 GPT_CKPT_PATH = MODEL_CONFIG.gpt_ckpt_path # set number of cpus. N_CPU = 8 DECODER_MAX_INPUT = 60 # 90 OOMs on A10s TOKEN_TIMEOUT_DURATION = 60 * 10 TOKENS_PER_CHUNK = 35 # controls the number of tokens sent to the decoder at a time WRITE_TO_REDIS = DEPLOYMENT_TYPE != "msft" VERBOSE_MESSAGE = DEPLOYMENT_TYPE == "dev" MOUNT_PATH = "/suno/models" WORKER_NAME = f"chirpv2_engine_{MODEL}_{DEPLOYMENT_TYPE}_{DURATION}s" if GPU == "A100" or GPU == "A100-80GB": WORKER_NAME += "_A100" elif GPU == "H100": WORKER_NAME += "_H100" elif GPU == "H200": WORKER_NAME += "_H200" APP_NAME = f"engine-{WORKER_NAME}" GPT_DDOG_SERVICE = "gpt-worker" DECODER_DDOG_SERVICE = "decoder-worker" # These queue names must correspond to the ones in streaming_api ENV_CONFIGS CHUNK_QUEUE_NAME = f"chunk-queue-{DEPLOYMENT_TYPE}" CHUNK_QUEUE_NAME_WEBM = f"chunk-queue-webm-{DEPLOYMENT_TYPE}" STREAM_KEY_QUEUE_NAME = f"stream-key-queue-{DEPLOYMENT_TYPE}" TOKEN_QUEUE_NAME = f"token-queue-{DEPLOYMENT_TYPE}" EVENTS_QUEUE_NAME = f"events-queue-{DEPLOYMENT_TYPE}" NUM_GPT_WORKER_LIMITS = { "priority": 80, "staging": 350, "dev": 3, "prod": 450, "msft": 6, } NUM_CODEC_WORKER_LIMITS = { "priority": 16, "staging": 70, "dev": 50, "prod": 2000, "msft": 300, } KEEP_WARM_DECODER = MODEL_CONFIG.keep_warm_decoder KEEP_WARM_GPT = MODEL_CONFIG.keep_warm_gpt aws_secret = modal.Secret.from_name("studio-aws") SECRETS = [ aws_secret, modal.Secret.from_name("openai-secret"), modal.Secret.from_dict( { "DD_SITE": "datadoghq.com", "DD_ENV": DEPLOYMENT_TYPE, "DD_SERVICE": WORKER_NAME, "DD_LOGS_ENABLED": "false", "DD_TRACE_ENABLED": "true" if DEPLOYMENT_TYPE == "dev" else "false", # memory leak }, ), modal.Secret.from_name("datadog-metrics"), modal.Secret.from_name("api-callback-token"), ] base_image = get_modal_base_image_with_flash_attention() class TokenSignalCode(Enum): STREAM_COMPLETE = 0 GPT_ERROR = 1 DECODER_INPUT_TIMED_OUT = 2 class ChirpWorker(S3Loader): def __init__(self, max_sequences=8, max_length_s=80): S3Loader.__init__(self) gpt_ckpt_path = chirp_v2._get_model_if_needed(GPT_CKPT_PATH, cache_dir=MOUNT_PATH) tokenizer_path = chirp_v2._get_model_if_needed(chirp_v2.TOKENIZER_PATH, cache_dir=MOUNT_PATH) # Get and print CUDA device name device_name = torch.cuda.get_device_name() print(f"Using CUDA device: {device_name}") min_batch_size = min(32, max_sequences) if MODEL_CONFIG.min_batch_size is not None: min_batch_size = MODEL_CONFIG.min_batch_size self.engine_rpc_client = start_service_processes( gpt_ckpt_path, tokenizer_path=tokenizer_path, max_sequences=max_sequences, max_length_s=max_length_s, world_size=MODEL_CONFIG.world_size, compile=COMPILE, batch_increment=MODEL_CONFIG.batch_increment, min_batch_size=min_batch_size, max_batch_size=MODEL_CONFIG.max_batch_size, ) self.tokenizer = self.engine_rpc_client.get_tokenizer() self.model_cfg = self.engine_rpc_client.get_model_cfg() @staticmethod def download_models(dir_path=MOUNT_PATH): """Use AWS CLI to download models if they don't exist.""" chirp_v2.preload_models( gpt_ckpt_path=GPT_CKPT_PATH, load_whisper=False, load_semantic=False, fetch_only=True, cache_dir=dir_path, ) recursive_ls_dir(dir_path) def code_generator(self, job: str): i = 0 while True: # this is very performance sensitive. should investigate a faster method of ipc # codes = np.zeros((100, 14)) codes = self.engine_rpc_client.get_generated_codes(job, start=i) if codes is None: return i += len(codes) if isinstance(codes, np.ndarray): yield from codes # would like to make this shorter, but slower time.sleep(4) class CodecWorker(S3Loader, CodecEngine): def __init__(self): S3Loader.__init__(self) # saves memory with dynamic batch sizes. must be set before loading the model os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "expandable_segments:True" model = chirp_v2.codec_load_model( chirp_v2._get_model_if_needed(chirp_v2.CODEC_CKPT_PATH, cache_dir=MOUNT_PATH), device="cuda", ) CodecEngine.__init__(self, model) @staticmethod def download_models(dir_path=MOUNT_PATH): from suno_utils.gpt import chirp_v2 as chirp_v2_for_codec_backwards_compat chirp_v2_for_codec_backwards_compat._get_model_if_needed( chirp_v2_for_codec_backwards_compat.CODEC_CKPT_PATH, cache_dir=dir_path ) chirp_v2_for_codec_backwards_compat.preload_codec_models( chirp_v2_for_codec_backwards_compat.CODEC_CKPT_PATH ) def download_model_wrapper_codec_a(): # this print is necessary to have modal rerun this when MODEL changes # Modal tracks referenced global variables # Change the name of the function to force a rerun print(f"Downloading Codec model {MODEL} {GPT_CKPT_PATH}") CodecWorker.download_models() def download_model_wrapper_gpt_a(): # this print is necessary to have modal rerun this when MODEL changes # Modal tracks referenced global variables # Change the name of the function to force a rerun print(f"Downloading GPT model {MODEL} {GPT_CKPT_PATH}") ChirpWorker.download_models() image = base_image.run_function( download_model_wrapper_codec_a, secrets=SECRETS, volumes={MODEL_STORE_VOLUME_DIR: model_store_volume} ).run_function( download_model_wrapper_gpt_a, secrets=SECRETS, volumes={MODEL_STORE_VOLUME_DIR: model_store_volume} ) app = modal.App(APP_NAME, image=image) # Queue used by Chirp/Decoder to stream tokens apptoken_queue = modal.Queue.from_name(TOKEN_QUEUE_NAME, create_if_missing=True) # Queue used by Decoder to send mp3 chunks under ephemeral stream_key. # This corresponds to chunk_queue_name in streaming_api, used to stream chunks. mp3_appchunk_queue = modal.Queue.from_name(CHUNK_QUEUE_NAME, create_if_missing=True) webm_appchunk_queue = modal.Queue.from_name(CHUNK_QUEUE_NAME_WEBM, create_if_missing=True) # Queue used by streaming_api to determine which ephemeral stream_key to stream from. # This corresponds to stream_key_queue_name in streaming_api. appstream_key_queue = modal.Queue.from_name(STREAM_KEY_QUEUE_NAME, create_if_missing=True) # Queue used by streaming_api to stream events like audio ready, events_queue = modal.Queue.from_name(EVENTS_QUEUE_NAME, create_if_missing=True) appinbound_event_queue = modal.Queue.from_name(INBOUND_EVENT_QUEUE_NAME, create_if_missing=True) if GPU == "A100": gpu = f"A100{':' + str(MODEL_CONFIG.world_size) if MODEL_CONFIG.world_size > 1 else ''}" elif GPU == "H100": gpu = f"H100{':' + str(MODEL_CONFIG.world_size) if MODEL_CONFIG.world_size > 1 else ''}" elif GPU == "H200": gpu = f"H200{':' + str(MODEL_CONFIG.world_size) if MODEL_CONFIG.world_size > 1 else ''}" elif GPU == "A100-80GB": gpu = f"A100-80GB{':' + str(MODEL_CONFIG.world_size) if MODEL_CONFIG.world_size > 1 else ''}" else: gpu = f"A10G{':' + str(MODEL_CONFIG.world_size) if MODEL_CONFIG.world_size > 1 else ''}" CHIRP_TIMEOUT = 660 QUEUE_ADD_TIMEOUT = 600 @app.cls( gpu=gpu, cloud=None, # "oci" if GPU == "H100" else None, secrets=SECRETS, timeout=CHIRP_TIMEOUT, # We never want to hit this timeout, it kills the worker and all jobs on it scaledown_window=400, memory=15000, retries=modal.Retries( max_retries=1, backoff_coefficient=2.0, initial_delay=5.0, ), min_containers=KEEP_WARM_GPT[DEPLOYMENT_TYPE], max_containers=NUM_GPT_WORKER_LIMITS[DEPLOYMENT_TYPE], buffer_containers=1 if DEPLOYMENT_TYPE == "prod" else 0, volumes={MODEL_STORE_VOLUME_DIR: model_store_volume}, ) @modal.concurrent(max_inputs=MODEL_CONFIG.model_concurrency) class ChirpV2Stub: # TODO: this should be renamed to sth else def __init__(self): import torch torch.set_num_threads(N_CPU) num_gpus = torch.cuda.device_count() cuda_device = torch.cuda.current_device() print(f"Found {num_gpus} GPUs. Using GPU {cuda_device}.") options = {"statsd_host": "127.0.0.1", "statsd_port": 8125} initialize(**options) chirp_v2.preload_models( gpt_ckpt_path=GPT_CKPT_PATH, load_gpt=False, load_semantic=False, load_whisper=False, cache_dir=MOUNT_PATH, ) self.worker = ChirpWorker( max_sequences=MODEL_CONFIG.max_sequences, max_length_s=MODEL_CONFIG.duration, ) def _get_gconf_kwargs(self, item: QueueItem): text = item.prompt_text or "" text_tags = item.metadata.get("tags", None) if text_tags == "random": text_tags = None custom_config = {} if MODEL == "v2": cfg_coef = 1.25 # condition on non-english, higher cfg text # note this uses fastext, coupled with usage of whisper now detected_lang = chirp_v2._get_text_lang(text) if (detected_lang is not None and not detected_lang == "en") and text: cfg_coef = 1.7 print(f"detected non-english lyrics, increase cfg, lang is {detected_lang}") # condition on genre, rap elif text_tags and ("rap" in text_tags or "hip-hop" in text_tags): cfg_coef = 1.4 custom_config = dict( cfg_coef=cfg_coef, cfg_coef_tags=1.9, temp_semantic=0.9, temp_coarse=0.85, min_text_offset=128, top_k_semantic=1000, top_k_coarse=1000, top_p_semantic=None, top_p_coarse=None, eos_pad_duration_s=0.24, max_tag_len=128, ) elif MODEL == "v2_ft13" or MODEL == "v2_dev": custom_config = dict( cfg_coef=1.25, cfg_coef_tags=1.9, temp_semantic=0.9, temp_coarse=0.85, min_text_offset=128, top_k_semantic=1000, top_k_coarse=1000, top_p_semantic=None, top_p_coarse=None, use_whisper=False, eos_pad_duration_s=0.24, max_tag_len=128, ) elif MODEL == "7b_dpo": custom_config = dict( cfg_coef=1.1, # collect the data for now cfg_coef_neg_tags=-2, cfg_coef_tags_max_steps=None, # collect the data for now min_eos_p=0.1, use_whisper=False, max_tag_len=128, ) elif MODEL == "7b": custom_config = dict( cfg_coef=1.25, # increase the cfg a bit for now cfg_coef_neg_tags=-2, use_whisper=False, max_tag_len=128, ) elif MODEL == "7b_ipo" or MODEL == "7b_special": custom_config = dict( cfg_coef=1.0, # no text cfg cfg_coef_tags=2.0, cfg_coef_neg_tags=-2, min_eos_p=0.1, # this is okay use_whisper=False, max_tag_len=128, ) elif MODEL == "7b_fast": custom_config = dict( cfg_coef=1.0, # for dpo-ed without cfg cfg_coef_tags=2.0, cfg_coef_neg_tags=-1.0, cfg_coef_tags_max_steps=250, # for dpo-ed without cfg min_text_offset=0, min_eos_p=0.1, use_whisper=False, n_repeat_tags=3, n_repeat_neg_tags=3, max_tag_len=128, ) elif MODEL == "13b_ft_1": # this is effectively a bot model custom_config = dict( cfg_coef=1.0, # turn off the text cfg -- use it in exp cfg_coef_tags=2.0, cfg_coef_neg_tags=-1.0, cfg_coef_tags_max_steps=50, # reduce cost further min_eos_p=0.1, # this is okay use_whisper=False, min_text_offset=0, n_repeat_tags=1, n_repeat_neg_tags=1, ) elif ( MODEL == "13b_special_8" or MODEL == "13b_special_31" or MODEL == "13b_special_test" or MODEL == "13b_special_test_2" or MODEL == "13b_special_test_3" or MODEL == "13b_short" or MODEL == "13b_tech" ): custom_config = dict( cfg_coef=1.0, # no text cfg cfg_coef_tags=2.0, cfg_coef_neg_tags=-1.0, min_eos_p=0.1, # this is okay use_whisper=False, min_text_offset=0, max_tag_len=128, n_repeat_tags=1, min_p_semantic=0.005, min_p_coarse=0.005, ) elif MODEL == "13b_upload_4": custom_config = dict( cfg_coef=1.0, # no text cfg cfg_coef_tags=2.5, cfg_coef_neg_tags=-1.0, min_eos_p=0.1, # this is okay use_whisper=False, min_text_offset=0, max_tag_len=128, n_repeat_tags=2, n_repeat_neg_tags=2, min_p_semantic=0.005, min_p_coarse=0.005, ) elif ( MODEL == "30b" or MODEL == "30b_t2" or MODEL == "30b_t3" or MODEL == "30b_t5" or MODEL == "30b_t6" or MODEL == "30b_t6_test" or MODEL == "30b_classical" or MODEL == "30b_dance" ): custom_config = dict( cfg_coef=1.3 if (MODEL == "30b" and DEPLOYMENT_TYPE == "dev") else 1.0, # for dpo-ed without cfg cfg_coef_tags=2.0, cfg_coef_neg_tags=-1.0, cfg_coef_tags_max_steps=250, # for dpo-ed without cfg min_text_offset=0, min_eos_p=0.1, use_whisper=False, max_tag_len=128, n_repeat_tags=3, n_repeat_neg_tags=3, min_p_semantic=0.005, min_p_coarse=0.005, ) # increase artist / cover text cfg and max steps if item.is_artist_condition or item.is_cover_condition: custom_config["cfg_coef"] = 1.2 custom_config["cfg_coef_tags_max_steps"] = 25 * 30 custom_config["cfg_coef_max_steps"] = 25 * 30 else: raise ValueError(f"Unknown model {MODEL}") input_negative_tags = item.metadata.get("negative_tags") if isinstance(input_negative_tags, str) and len(input_negative_tags) > 0: # if negative tags are provided, we want to repeat them more times custom_config["n_repeat_neg_tags"] = 3 # Note that experiment config will overwrite default config if special_config := item.metadata.get("model_config"): # has specific config instruction print(f"ChirpStub: load specific config: {special_config}, {item.id}") # update the custom config # note this will overwrite existing parameters custom_config.update(**special_config) # Note forced infer config will overwrite experiment config if (special_config := item.metadata.get("forced_infer_config")) and DEPLOYMENT_TYPE == "dev": # has specific config instruction print(f"item.metadata: {item.metadata}") print(f"ChirpStub: load forced infer config: {special_config}, {item.id}") # update the custom config # note this will overwrite existing parameters custom_config.update(**{k: v for k, v in special_config.items() if v is not None}) # Filter kwargs that are valid properties of GenerationConfig valid_gconf_props = set(vars(GenerationConfig()).keys()) filtered_config = {k: v for k, v in custom_config.items() if k in valid_gconf_props} # Print removed keys for debugging removed_keys = set(custom_config.keys()) - set(filtered_config.keys()) if removed_keys: print(f"ChirpStub: Removed invalid GenerationConfig properties: {removed_keys}") custom_config = filtered_config return custom_config @tracer.wrap(service=GPT_DDOG_SERVICE) def _make_gconf(self, item: QueueItem, history: HistoryPrompt) -> GenerationConfig: # build a config text = item.prompt_text or "" text_tags = item.metadata.get("tags", None) if text_tags == "random": text_tags = None text_neg_tags = None text_start_control_tags = None text_end_control_tags = None history_audio = history.prompt_audio if IS_3B: # v2 if "instrumental" not in (text_tags or "") and len(text.split()) > 5: text_neg_tags = "instrumental noise" else: # v3 and onwards is_instrumental = False text_neg_tags = "repetitive, loop," if IS_13B: text_neg_tags += " noisy, distorted," # cleaned_text will remove all meta tags like: [chorus] cleaned_text = clean_text(text).strip() if len(cleaned_text) == 0: # there is no asr-able text # very likely this should be instrumental # consistent with the FE update is_instrumental = True text_tags = "instrumental, " + (text_tags or "") + ", instrumental" if MODEL == "13b_upload_4": # we want to trick the audio upload worker # to generate a bit longer audio text = "[instrumental] \n [instrumental] \n ... \n [instrumental] \n [instrumental]" print(f"Request {item.id} hack text input to : {text}") else: text = "[instrumental, instrumental, instrumental]" + text text_neg_tags += " vocal, sing, voices, chant, hum, speech" if item.is_one_box_generation or history_audio is None: # can be None, or empty "" # for one box, we want to try to end the songs print("ChirpStub: Pad start/end to one-box generation or history is None.") # from v3 for 7b we will use the `{}` annotation # note that for instrumental, we don't want vocals, and we don't pad control {end} text_start_control_tags = "{start:0} " if IS_30B else "{start} " if not is_instrumental and (MODEL == "13b_short" or MODEL == "7b_fast"): text_start_control_tags = "{vocals:start} " text_end_control_tags = "{end:vocals:end} " if not is_instrumental and IS_7B: text_start_control_tags = "{start;vocals:start} " if not is_instrumental and IS_30B: text_start_control_tags = ( "{start:0;vocals:start} " if random.random() > 0.5 else "{start:0;vocals:intro} " ) if history_audio is None: text_end_control_tags = "{end} " # we will try to end one box using end tag if item.is_one_box_generation: text = text + (" \n [end]" if not is_instrumental else "") # special infilling control tag overwrite if IS_30B and item.is_infill: infill_start_s = item.metadata.get("infill_start_s", 0) infill_end_s = item.metadata.get("infill_end_s", None) # these are intro/outro durations control_tags = [] if (infill_duration := item.metadata.get("infill_dur_s")) and item.metadata.get( "task" ) == "infill": replaced_duration = max(round(infill_duration), 1) if infill_start_s == 0: control_tags.append("start:0") history_context_duration = ( history.prompt_audio.shape[0] // 25 if isinstance(history.prompt_audio, np.ndarray) else 0 ) future_context_duration = ( history.future_audio.shape[0] // 25 if isinstance(history.future_audio, np.ndarray) else 0 ) total_duration = int( history_context_duration + replaced_duration + future_context_duration ) control_tags.append("duration:" + str(total_duration)) elif item.metadata.get("task") == "infill_intro": replaced_duration = max(int(infill_end_s - infill_start_s), 1) intro_duration = item.metadata.get("infill_dur_s", 10) future_context_duration = ( history.future_audio.shape[0] // 25 if isinstance(history.future_audio, np.ndarray) else 20 ) total_duration = int(future_context_duration + replaced_duration + intro_duration) control_tags.append("start:0") control_tags.append("duration:" + str(total_duration)) elif item.metadata.get("task") == "infill_outro": replaced_duration = max(int(infill_end_s - infill_start_s), 1) outro_duration = item.metadata.get("infill_dur_s", 10) history_context_duration = ( history.prompt_audio.shape[0] // 25 if isinstance(history.prompt_audio, np.ndarray) else 20 ) total_duration = int(history_context_duration + replaced_duration + outro_duration) control_tags.append("duration:" + str(total_duration)) if control_tags: text_start_control_tags = "{" + ";".join(control_tags) + "} " print( f"ChirpStub: item {item.id} infilling control tags: {text_start_control_tags}." ) # load the other inference parameters here inference_parameters = self._get_gconf_kwargs(item) # overwrite the negative tags from prompt input_negative_tags = item.metadata.get("negative_tags") if isinstance(input_negative_tags, str) and len(input_negative_tags) > 0: text_negative_tags = input_negative_tags # if negative tags are provided, we want to repeat them more times inference_parameters["n_repeat_neg_tags"] = 3 else: text_negative_tags = text_neg_tags if item.metadata.get("negative_tags"): # use gpt prompt print( f"ChirpStub: load specific negative_tags: {item.metadata.get('negative_tags')}, {item.id}" ) # set the generation max duration if history_audio is not None: if MODEL == "13b_upload_4": # will dynamically change based on input duration max_gen_duration_s = MODEL_CONFIG.duration - history_audio.shape[0] // 25 - 1 print( f"Upload extend clip: {item.id} customized max generation duration is {max_gen_duration_s}." ) else: max_gen_duration_s = MODEL_CONFIG.duration - MODEL_CONFIG.max_history_duration_s else: max_gen_duration_s = MODEL_CONFIG.duration # for bot generation, we want to limit the generation duration, but keep it random :) if item.is_bot_generation and MODEL == "13b_ft_1": # this is a beta distribution, mean is roughply 10 / (10 + 10) = 0.5 ~ 120 sec max_gen_duration_s = 10 + int(random.betavariate(10, 10) * 230) print(f"Bot generation: {item.id}, max duration: {max_gen_duration_s}.") # make the configuration cfg = chirp_v2.GenerationConfig( text=text, text_tags=text_tags, text_neg_tags=text_negative_tags, text_start_control_tags=text_start_control_tags, text_end_control_tags=text_end_control_tags, n_batch=1, max_gen_duration_s=max_gen_duration_s, rep_penality=0, history_arr=history_audio, history_text=history.prompt_lyrics, future_arr=history.future_audio, cover_arr=history.cover_audio, artist_arr=history.artist_audio, history_do_asr=False, stream=True, allow_eos=True, **inference_parameters, # pass in the other configuration parameters ) if VERBOSE_MESSAGE: print(f"ChirpStub: Model {item.model_name}: gconf: {cfg}") return cfg @modal.method() @distributed_trace("gpt_generate", GPT_DDOG_SERVICE, env_name=DEPLOYMENT_TYPE) def generate(self, queue_item: str, history: Optional[HistoryPrompt] = None): tracer.current_span().set_tag("model", MODEL) tracer.current_span().set_tag("duration", DURATION) tracer.current_span().set_tag("deployment_type", DEPLOYMENT_TYPE) start_time = time.time() item = QueueItem(**json.loads(queue_item)) item_id = item.id ids = [item_id] if history is None: history = HistoryPrompt() try: if DEPLOYMENT_TYPE == "dev": # these are mostly for debugging but not practical for prod os.environ["CUDA_LAUNCH_BLOCKING"] = "1" os.environ["TORCH_CUDA_ALLOC_SYNC"] = "1" if VERBOSE_MESSAGE: print_gpu_memory_usage(self.__class__.__name__) torch.cuda.reset_max_memory_allocated() request_prep_time = time.time() print( f"ChirpStub: Adding job to queue: {item_id}, " f"modalID {modal.current_input_id()}, {request_prep_time - start_time:.2f}s elapsed." ) device_name = torch.cuda.get_device_name() dd_tags = [ f"model:{MODEL}", f"env:{DEPLOYMENT_TYPE}", f"env_name:{WORKER_NAME}", f"modal_cloud_provider:{os.environ.get('MODAL_CLOUD_PROVIDER', 'unknown')}", f"modal_environment:{os.environ.get('MODAL_ENVIRONMENT', 'unknown')}", f"modal_image_id:{os.environ.get('MODAL_IMAGE_ID', 'unknown')}", f"modal_region:{os.environ.get('MODAL_REGION', 'unknown')}", f"gen_type:{item.metadata.get('type', 'unknown')}", f"gen_task:{item.metadata.get('task', 'unknown')}", # this is not always set for stuff like basic gens f"cuda_device:{device_name}", f"custom_lyrics:{not bool(item.metadata.get('gpt_description_prompt'))}", # f"modal_task_id:{os.environ.get('MODAL_TASK_ID', 'unknown')}", # this increases costs ] # # this watches for events from clients # inbound_event_queue = InboundEventQueue( # appinbound_event_queue, # item_id, # timeout=TOKEN_TIMEOUT_DURATION, # ) # inbound_event_queue.start() gconf = self._make_gconf(item, history) max_history_duration_s = MODEL_CONFIG.max_history_duration_s # for infilling, we don't want to limit the history duration # they can be as long as possible... if item.is_infill: max_history_duration_s = None gconf = chirp_v2.prep_gconf( gconf, max_history_duration_s=max_history_duration_s, cfg=self.worker.model_cfg, ) # Each unique Chirp request is assigned a unique ephemeral stream_key. # Replace item in the stream_key_queue. # streaming_api streams from the first item in the stream_key_queue. stream_key = "stream-key-" + str(uuid4()) appstream_key_queue.get(block=False, partition=item_id) appstream_key_queue.put(stream_key, partition=item_id, partition_ttl=TOKEN_TIMEOUT_DURATION) request = make_request(item_id, gconf, self.worker.model_cfg, self.worker.tokenizer) with tracer.trace("add_request_to_worker"): start_add_time = time.time() print( f"Adding request to worker: {item_id}, {start_add_time - start_time:.2f}s elapsed." ) result = self.worker.engine_rpc_client.add_request(request) if result == "Job already exists": raise ValueError(f"ChirpStub: Job already exists for {item_id}") while self.worker.engine_rpc_client.get_job_state(request.id) is None: time.sleep(0.2) end_add_time = time.time() statsd.distribution( "engine.add_request_duration_seconds", end_add_time - start_add_time, tags=dd_tags ) statsd.distribution( "engine.free_caches", self.worker.engine_rpc_client.get_free_caches(), tags=dd_tags ) statsd.distribution( "engine.used_caches", self.worker.engine_rpc_client.get_used_caches(), tags=dd_tags ) statsd.gauge( "engine.evictions_last_minute", self.worker.engine_rpc_client.get_evictions_last_minute(), tags=dd_tags, ) num_jobs = self.worker.engine_rpc_client.get_num_jobs() num_streams = self.worker.engine_rpc_client.get_num_streams() avg_streams_per_job = num_streams / num_jobs if num_jobs > 0 else 0 print( f"Number of jobs: {num_jobs}, Number of streams: {num_streams}, Avg streams per job: {avg_streams_per_job}" ) statsd.distribution("engine.num_jobs", num_jobs, tags=dd_tags) statsd.distribution("engine.num_streams", num_streams, tags=dd_tags) statsd.distribution("engine.streams_per_job", avg_streams_per_job, tags=dd_tags) utilization = torch.cuda.utilization() print(f"CUDA Utilization: {utilization}") statsd.distribution("cuda.utilization", utilization, tags=dd_tags) print( f"{item_id}: Time till added request to worker: {end_add_time - start_time:.2f} seconds" ) statsd.distribution("engine.time_before_add", start_add_time - start_time, tags=dd_tags) statsd.distribution("engine.time_after_add", end_add_time - start_time, tags=dd_tags) code_stream = self.worker.code_generator(item_id) aligned_code_stream = align_codes(code_stream, self.worker.model_cfg) # if we want to include history in the output include_history_s = item.metadata.get("include_history_s", 0) include_future_s = item.metadata.get("include_future_s", 0) if include_history_s > 0 and gconf.history_arr is not None: history_codes = gconf.history_arr[-int(25 * include_history_s) :, 1:].astype(np.int32) print( f"ChirpStub: Including {include_history_s} seconds of history. {history_codes.shape} tokens." ) aligned_code_stream = itertools.chain( torch.from_numpy(history_codes), aligned_code_stream, ) if include_future_s > 0 and gconf.future_arr is not None: future_codes = gconf.future_arr[: int(25 * include_future_s), 1:].astype(np.int32) print( f"ChirpStub: Including {include_future_s} seconds of future. {future_codes.shape} tokens." ) aligned_code_stream = itertools.chain( aligned_code_stream, torch.from_numpy(future_codes), ) codes = [] stream_size = 0 start_decoder = False with tracer.trace("generate_tokens") as span: span_str = serialize_context(span.context) for code in aligned_code_stream: # if not inbound_event_queue.local_queue.empty(): # event = inbound_event_queue.local_queue.get() # print(f"ChirpStub: Received event {event}") # if event.event_type is EventType.CANCEL: # print(f"ChirpStub: Canceling generation for {item_id}") # # Note this doesnt remove the job from the engine immediately # # it will be removed when rest of this function finishes # break codes.append(code.numpy().astype(np.int16)) if len(codes) >= TOKENS_PER_CHUNK: with tracer.trace("put_tokens"): tracer.current_span().set_tag("stream_size", stream_size) tracer.current_span().set_tag("n_codes", len(codes)) apptoken_queue.put( codes, partition=stream_key, partition_ttl=TOKEN_TIMEOUT_DURATION, block=False, timeout=QUEUE_ADD_TIMEOUT, ) stream_size += len(codes) codes = [] if not start_decoder: start_decoder = True # start the decoder with sth in the dict DecoderStub.generate.spawn( item.json(), stream_key=stream_key, parent_context=span_str ) print( f"ChirpStub: Start the decoder: {item_id}, {time.time() - start_time:.2f}s elapsed." ) statsd.distribution( "engine.time_to_start_decoder", time.time() - start_time, tags=dd_tags, ) if time.time() - start_time > CHIRP_TIMEOUT - 120: # manually catch timeouts to prevent modal from killing the worker print( f"ChirpStub: Timeout for {item_id}, {time.time() - start_time:.2f}s elapsed." ) raise TimeoutError( f"ChirpStub: Timeout for {item_id}, {time.time() - start_time:.2f}s elapsed." ) print(f"{item_id}: Time till sending last chunk: {time.time() - start_time:.2f} seconds") # if the gen is really short, we need to send the last chunk, and still start the job with tracer.trace("send_last_chunk", resource="token_queue_operation"): apptoken_queue.put( codes + [TokenSignalCode.STREAM_COMPLETE], partition=stream_key, partition_ttl=TOKEN_TIMEOUT_DURATION, block=False, timeout=QUEUE_ADD_TIMEOUT, ) stream_size += len(codes) print(f"ChirpStub: Item id: {item_id}, total stream size: {stream_size}") print(f"{item_id}: Time till sent last chunk: {time.time() - start_time:.2f} seconds") statsd.distribution( "engine.time_after_send_last_chunk", time.time() - start_time, tags=dd_tags ) if not start_decoder: start_decoder = True DecoderStub.generate.spawn( item.json(), stream_key=stream_key, parent_context=span_str ) start_write_npz_time = time.time() with tracer.trace("write_npz"): job_state = self.worker.engine_rpc_client.get_job_state(request.id) print(f"{item_id}: Job state: {job_state}") codes = self.worker.engine_rpc_client.get_generated_codes(request.id) raw_arrays = unshift_arrays_v2(codes, self.worker.model_cfg, job_state["eos_step"]) raw_arrays = raw_arrays.cpu().numpy() self.worker._write_npz( item, raw_arrays, APP_NAME, MODEL_VERSION, generated_arr=raw_arrays, history_arr=gconf.history_arr, future_arr=gconf.future_arr, pre_history_arr=history.pre_history_arr, post_future_arr=history.post_future_arr, cover_arr=history.cover_audio, artist_arr=history.artist_audio, history_lyrics=history.prompt_lyrics, fuller_array=job_state["full_array"], ) statsd.distribution("engine.job_itl", job_state["itl"], tags=dd_tags) statsd.distribution("engine.job_ttl", job_state["ttl"], tags=dd_tags) statsd.distribution("engine.job_ttft", job_state["ttft"], tags=dd_tags) print(f"{item_id}: Time writing npz: {time.time() - start_write_npz_time:.2f} seconds") print(f"{item_id}: Time till removing job: {time.time() - start_time:.2f} seconds") statsd.distribution("engine.time_before_remove_job", time.time() - start_time, tags=dd_tags) self.worker.engine_rpc_client.remove_job(request.id) if VERBOSE_MESSAGE: max_mem = torch.cuda.max_memory_allocated() / 1e9 print(f"ChirpStub: Max memory allocated: {max_mem:.2f} GB") except Exception as e: error_message = str(e) # or e.args[0] if args attribute is used if "Job already exists" in error_message: print(f"ChirpStub: Job {item_id} already exists, skipping.") return finish_time = time.time() apptoken_queue.put( [TokenSignalCode.GPT_ERROR], partition=item_id, partition_ttl=TOKEN_TIMEOUT_DURATION, block=False, ) item.notify_progress( { "id": item_id, "model": item.model_name, "n_audios": len(ids), "ok": 0, "gen_duration": finish_time - start_time, "error_type": "generation_failure", "error_message": "Failed generation.", } ) # TODO: such workers should be stopped ASAP. find a proper modal way to do this if "CUDA error: an illegal memory access was encountered" in error_message: print("ChirpStub: Caught a CUDA memory error, this worker is busted. Sleep to death.") modal.experimental.stop_fetching_inputs() if "stream has been closed" in error_message: print("ChirpStub: Caught a stream closed error, this worker is busted. Sleep to death.") modal.experimental.stop_fetching_inputs() self.worker.engine_rpc_client.remove_job(request.id) print(f"ChirpStub: item_id: {item_id} failed with error: {e}") traceback.print_exc() raise e @modal.exit() def cleanup_processes(self): print("ChirpStub: Cleaning up processes") self.worker.engine_rpc_client.terminate_processes() def fill_local_token_queue(local_token_queue: queue.Queue, partition: str): # Ensure the queue has not expired. if apptoken_queue.len(partition=partition) == 0: # Fail the decoding job for this partition. local_token_queue.put([TokenSignalCode.DECODER_INPUT_TIMED_OUT]) return # Stream tokens into local queue. for token_batch in apptoken_queue.iterate(partition=partition, item_poll_timeout=60): local_token_queue.put(token_batch) if ( isinstance(token_batch[-1], TokenSignalCode) and token_batch[-1] == TokenSignalCode.STREAM_COMPLETE ): return # print(f"Put {len(token_batch)} tokens into queue, {tokens.qsize()} tokens in queue {partition}.") # If we exited the loop without STREAM_COMPLETE, Chirp lagged out and stopped generating tokens. local_token_queue.put([TokenSignalCode.DECODER_INPUT_TIMED_OUT]) # TODO: when the GPT worker is spinning up -- they should trigger a multi-warm up call to the decoder # So the decoder scales up before the tokens are generated @app.cls( cpu=12, gpu="A10G", secrets=SECRETS, timeout=600, # TODO: with current issue this is too long scaledown_window=400, memory=32000, retries=modal.Retries( max_retries=1, backoff_coefficient=2.0, initial_delay=5.0, ), min_containers=KEEP_WARM_DECODER[DEPLOYMENT_TYPE], max_containers=NUM_CODEC_WORKER_LIMITS[DEPLOYMENT_TYPE], buffer_containers=1 if DEPLOYMENT_TYPE == "prod" else 0, volumes={MODEL_STORE_VOLUME_DIR: model_store_volume}, ) @modal.concurrent(max_inputs=DECODER_MAX_INPUT) class DecoderStub: def __init__(self): import torch torch.set_num_threads(N_CPU) num_gpus = torch.cuda.device_count() device_name = torch.cuda.get_device_name() print(f"Found {num_gpus} GPUs. Using CUDA device: {device_name}") options = {"statsd_host": "127.0.0.1", "statsd_port": 8125} initialize(**options) self.worker = CodecWorker() self.worker.start() self.modal_f_video_generator = modal.Cls.lookup( f"videos-v2-{'dev' if DEPLOYMENT_TYPE == 'dev' else 'prod'}", "DummyV0Stub", )().write_video def print_info(self, queue_item_id: str, message: str): print(f"DecoderStub ({queue_item_id}): {message}") @modal.method() @distributed_trace("decoder_generate", DECODER_DDOG_SERVICE, env_name=DEPLOYMENT_TYPE) def generate(self, item_json: str, stream_key: Optional[str] = None): start_time = time.time() if VERBOSE_MESSAGE: print_gpu_memory_usage(self.__class__.__name__) torch.cuda.reset_max_memory_allocated() item = QueueItem(**json.loads(item_json)) item_id = item.id device_name = torch.cuda.get_device_name() dd_tags = [ f"model:{MODEL}", f"env:{DEPLOYMENT_TYPE}", f"env_name:{WORKER_NAME}", f"modal_cloud_provider:{os.environ.get('MODAL_CLOUD_PROVIDER', 'unknown')}", f"modal_environment:{os.environ.get('MODAL_ENVIRONMENT', 'unknown')}", f"modal_image_id:{os.environ.get('MODAL_IMAGE_ID', 'unknown')}", f"modal_region:{os.environ.get('MODAL_REGION', 'unknown')}", f"gen_type:{item.metadata.get('type', 'unknown')}", f"gen_task:{item.metadata.get('task', 'unknown')}", # this is not always set for stuff like basic gens f"cuda_device:{device_name}", f"custom_lyrics:{not bool(item.metadata.get('gpt_description_prompt'))}", ] token_queue = queue.Queue() # If specified, use stream_key to read tokens from the token queue. # Otherwise, fall back to the item id. partition = stream_key or item_id loader_thread = threading.Thread( target=fill_local_token_queue, args=(token_queue, partition), daemon=True, ) loader_thread.start() ids = [item_id] first_log_time = None error = None attempts = 2 self.print_info(item_id, "Starting decoder job") statsd.gauge("decoder.num_jobs", len(self.worker.active_jobs()), tags=dd_tags) try: received_tokens = [] while attempts > 0: error = None attempts -= 1 audios: list[Audio] = [] audio_idx = 0 first_ffmpeg_read_time_mp3 = None first_ffmpeg_read_time_webm = None with ( ffmpeg_stream_encode() as mp3_proc, ffmpeg_stream_encode_opus_webm() as webm_proc, ): # Increase pipe buffer size (e.g., to 1 MB) # Should fix broken pipe error PIPE_BUF_SIZE = 1024 * 1024 # 1 MB fcntl.fcntl(mp3_proc.stdin.fileno(), fcntl.F_SETPIPE_SZ, PIPE_BUF_SIZE) fcntl.fcntl(mp3_proc.stdout.fileno(), fcntl.F_SETPIPE_SZ, PIPE_BUF_SIZE) fcntl.fcntl(webm_proc.stdin.fileno(), fcntl.F_SETPIPE_SZ, PIPE_BUF_SIZE) fcntl.fcntl(webm_proc.stdout.fileno(), fcntl.F_SETPIPE_SZ, PIPE_BUF_SIZE) os.set_blocking(mp3_proc.stdout.fileno(), False) os.set_blocking(webm_proc.stdout.fileno(), False) os.set_blocking(mp3_proc.stdin.fileno(), False) os.set_blocking(webm_proc.stdin.fileno(), False) mp3_poll = select.poll() mp3_poll.register(mp3_proc.stdout, select.POLLIN) webm_poll = select.poll() webm_poll.register(webm_proc.stdout, select.POLLIN) mp3_local_chunk = b"" webm_local_chunk = b"" partially_written_chunk_mp3 = b"" partially_written_chunk_webm = b"" # not currently necessary to input tokens list since its empty, but could be filled in the future with tracer.trace("add_request"): job = self.worker.add_request(Request(item_id)) # add existing tokens if this is a retry for token in received_tokens: job.add_token(token) while True: # check for new decoded audio if ( audio_idx >= len(job.generated_audios) and len(partially_written_chunk_mp3) == 0 and len(partially_written_chunk_webm) == 0 and job.decoding_completed # we close both together, so checking one is enough and not mp3_proc.stdin.closed ): mp3_proc.stdin.close() webm_proc.stdin.close() # if there is audio to write, and we have space in the ffmpeg buffers, write it to ffmpeg while ( audio_idx < len(job.generated_audios) or len(partially_written_chunk_mp3) > 0 or len(partially_written_chunk_webm) > 0 ): made_progress = False with tracer.trace("write_audio_to_ffmpeg"): # grab new audios until we can fill the partially written chunks # or run out of new audios while ( audio_idx < len(job.generated_audios) and len(partially_written_chunk_mp3) == 0 and len(partially_written_chunk_webm) == 0 ): if audio_idx == 0: statsd.distribution( "decoder.first_generated_audio_time", time.time() - start_time, tags=dd_tags, ) audio = job.generated_audios[audio_idx] audio_idx += 1 audios.append(audio) partially_written_chunk_mp3 = np.ascontiguousarray( audio.array_float.T ).tobytes() partially_written_chunk_webm = partially_written_chunk_mp3 # now we are either out of chunks or have a non-empty partially written chunk if ( len(partially_written_chunk_mp3) == 0 and len(partially_written_chunk_webm) == 0 ): break try: with tracer.trace("write_audio_to_ffmpeg_mp3"): try: write_len = mp3_proc.stdin.write(partially_written_chunk_mp3) if write_len is not None: partially_written_chunk_mp3 = ( partially_written_chunk_mp3[write_len:] ) made_progress |= write_len > 0 except BlockingIOError: # no space in the buffer pass with tracer.trace("write_audio_to_ffmpeg_webm"): try: write_len = webm_proc.stdin.write( partially_written_chunk_webm ) if write_len is not None: partially_written_chunk_webm = ( partially_written_chunk_webm[write_len:] ) made_progress |= write_len > 0 except BlockingIOError: # no space in the buffer pass except BrokenPipeError: self.print_info( item_id, f"Broken pipe, skipping. Wrote {len(audios)} chunks.", ) error = "broken_pipe" break if not made_progress: break mp3_done = False webm_done = False # read audio from ffmpeg if mp3_poll.poll(1): with tracer.trace("read_audio_from_ffmpeg_mp3"): # modal write has 256KiB limit mp3_read_res = mp3_proc.stdout.read(64 * 1024) if mp3_read_res: mp3_local_chunk += mp3_read_res else: # if poll() returns an event but we read 0 bytes, the pipe is closed # at the other end mp3_done = True if first_ffmpeg_read_time_mp3 is None: first_ffmpeg_read_time_mp3 = time.time() statsd.distribution( "decoder.first_ffmpeg_read_time", first_ffmpeg_read_time_mp3 - start_time, tags=dd_tags + ["format:mp3"], ) if first_log_time is None: self.print_info( item_id, f"First audio chunk reading mp3: {time.time() - start_time:.2f}s", ) if webm_poll.poll(1): with tracer.trace("read_audio_from_ffmpeg_webm"): # modal write has 256KiB limit webm_read_res = webm_proc.stdout.read(64 * 1024) if webm_read_res: webm_local_chunk += webm_read_res else: webm_done = True if first_ffmpeg_read_time_webm is None: first_ffmpeg_read_time_webm = time.time() statsd.distribution( "decoder.first_ffmpeg_read_time", first_ffmpeg_read_time_webm - start_time, tags=dd_tags + ["format:webm"], ) if first_log_time is None: self.print_info( item_id, f"First audio chunk reading webm: {time.time() - start_time:.2f}s", ) decoded_duration = sum([audio.duration_s for audio in audios]) should_buffer = first_log_time is None and decoded_duration < 0.6 if ( len(mp3_local_chunk) > 1024 * 16 and len(webm_local_chunk) > 1024 * 16 and not should_buffer ): with tracer.trace("write_audio_to_modal"): if WRITE_TO_REDIS: try: mp3_appchunk_queue.put( mp3_local_chunk, partition=partition, partition_ttl=TOKEN_TIMEOUT_DURATION, block=False, timeout=QUEUE_ADD_TIMEOUT, ) webm_appchunk_queue.put( webm_local_chunk, partition=partition, partition_ttl=TOKEN_TIMEOUT_DURATION, block=False, timeout=QUEUE_ADD_TIMEOUT, ) except Exception: self.print_info(item_id, "Failed to write chunk to Modal.") pass mp3_local_chunk = b"" webm_local_chunk = b"" if first_log_time is None: first_log_time = time.time() elapsed_time = first_log_time - start_time self.print_info( item_id, f"First audio chunk writing: {elapsed_time:.2f}s elapsed since request received.", ) statsd.distribution( "decoder.first_write_time_seconds", elapsed_time, tags=dd_tags, ) item.notify_progress( { "id": item.id, "type": "streaming", }, ) events_queue.put( { "type": "gen_streaming", "data": {}, }, partition=item.id, partition_ttl=60, ) if ( gen_request_start_time := item.metadata.get( "gen_request_start_time", None ) ) is not None: first_audio_write_time_millis = int(time.time() * 1000) elapsed_millis = ( first_audio_write_time_millis - gen_request_start_time ) if elapsed_millis > 0: print( f"DecoderStub ({item.id}): Start audio write time from gen request: {elapsed_millis // 1000}s" ) if elapsed_millis > 600_000: print( f"DecoderStub ({item.id}): Start audio write time from gen request took >10m: {elapsed_millis // 1000}s" ) statsd.distribution( "decoder.audio_write_start_millis.distribution", elapsed_millis, tags=dd_tags, ) if error is not None: break # pull new input tokens from thread queue while True: try: token_batch = token_queue.get_nowait() except queue.Empty: break are_all_codes = all( not isinstance(token, TokenSignalCode) for token in token_batch ) if are_all_codes: # if all are codes we convert them to torch in a batch for speed with tracer.trace("add_tokens_to_job_batch"): tokens_batch = np.stack(token_batch) for token in tokens_batch: job.add_token(token) received_tokens.append(token) else: for token in token_batch: if not isinstance(token, TokenSignalCode): with tracer.trace("add_token_to_job"): if not isinstance(token, np.ndarray): raise ValueError( f"Expected np.ndarray, got {type(token)}" ) # print(f"Adding token to job {item_id}") token = torch.from_numpy(token) job.add_token(token) received_tokens.append(token) else: if token == TokenSignalCode.GPT_ERROR: error = "gpt_error" attempts = 0 elif ( token == TokenSignalCode.DECODER_INPUT_TIMED_OUT and not job.request.input_tokens_finished ): error = "decoder_input_timed_out" attempts = 0 elif token == TokenSignalCode.STREAM_COMPLETE: self.print_info(item_id, "Received EOS token.") else: raise ValueError("This shouldn't happen.") job.request.input_tokens_finished = True if mp3_done and webm_done: break time.sleep(0.1) if error is None: break elif attempts == 0: self.print_info( item_id, f"Failed to process after few attempts, error: {error}", ) else: self.print_info( item_id, f"Retrying due to {error}, {attempts} attempts left.", ) time.sleep(5) mp3_proc.stdin.close() mp3_proc.stdout.close() mp3_proc.terminate() webm_proc.stdin.close() webm_proc.stdout.close() webm_proc.terminate() statsd.distribution( "decoder.step_batch_size", np.mean(self.worker.batch_size_history), tags=dd_tags ) # try to empty the generated audios still? while audio_idx < len(job.generated_audios): audio = job.generated_audios[audio_idx] audio_idx += 1 audios.append(audio) self.print_info(item_id, "Decoder job complete") if VERBOSE_MESSAGE: max_mem = torch.cuda.max_memory_allocated() / 1e9 self.print_info(item_id, f"Max memory allocated: {max_mem:.2f} GB") loop_end_time = time.time() self.worker.remove_job(job.request.id) if len(audios) == 0: self.print_info(item_id, f"No audio generated. Recived tokens: {len(received_tokens)}") error = "decoder_error" if error is None: audio = Audio.concatenate(audios) if audio.duration_s != len(job.tokens) / EMBEDDING_RATE: print( f"Warning: Audio duration {audio.duration_s} does not match expected duration {len(job.tokens) / EMBEDDING_RATE} for item_id {item_id}" ) if WRITE_TO_REDIS: with tracer.trace("write_audio_to_modal"): try: mp3_appchunk_queue.put( mp3_local_chunk, partition=partition, partition_ttl=TOKEN_TIMEOUT_DURATION, block=False, timeout=QUEUE_ADD_TIMEOUT, ) webm_appchunk_queue.put( webm_local_chunk, partition=partition, partition_ttl=TOKEN_TIMEOUT_DURATION, block=False, timeout=QUEUE_ADD_TIMEOUT, ) # write empty bytes to modal to signal end of stream mp3_appchunk_queue.put( b"", partition=partition, partition_ttl=TOKEN_TIMEOUT_DURATION, block=False, timeout=QUEUE_ADD_TIMEOUT, ) webm_appchunk_queue.put( b"", partition=partition, partition_ttl=TOKEN_TIMEOUT_DURATION, block=False, timeout=QUEUE_ADD_TIMEOUT, ) if ( gen_request_start_time := item.metadata.get( "gen_request_start_time", None ) ) is not None: first_audio_write_time_millis = int(time.time() * 1000) elapsed_millis = first_audio_write_time_millis - gen_request_start_time if elapsed_millis > 0: print( f"DecoderStub ({item.id}): Finish audio write time from gen request: {elapsed_millis // 1000}s" ) if elapsed_millis > 600_000: print( f"DecoderStub ({item.id}): Finish audio write time from gen request took >10m: {elapsed_millis // 1000}s" ) statsd.distribution( "decoder.audio_write_finish_millis.distribution", elapsed_millis, tags=dd_tags, ) except Exception: self.print_info(item_id, "Failed to write chunk to Modal.") pass if first_log_time is not None: realtime_factor = audio.duration_s / (loop_end_time - first_log_time) self.print_info( item_id, f"Finished writing chunks. Realtime factor: {realtime_factor:.2f}x", ) statsd.increment("decoder.complete", 1, tags=dd_tags) statsd.distribution("decoder.duration_seconds", audio.duration_s, tags=dd_tags) statsd.distribution("decoder.realtime_factor", realtime_factor, tags=dd_tags) with tracer.trace("write_audio_to_s3"): self.worker._write_audio_only(item, audio) self.print_info(item_id, "Finished writing to S3.") else: audio = None raise Exception(f"Failed to generate audio for item_id {item_id}") item.notify_progress( { "id": item_id, "model": item.model_name, "n_audios": len(ids), "ok": error is None, "gen_duration": time.time() - start_time, "ids": ids, "durations": [round(audio.duration_s, 3) if audio is not None else 0], } ) video_time = time.time() # Start the video jobs. self.modal_f_video_generator.spawn(item.json(), f"image_{item.id}.png") self.print_info( item_id, f"Spawning videos {time.time() - video_time}", ) return item_id except Exception as e: self.print_info(item_id, f"Exception occurs: {e}") finish_time = time.time() item.notify_progress( { "id": item_id, "model": item.model_name, "n_audios": len(ids), "ok": 0, "gen_duration": finish_time - start_time, "error_type": "generation_failure", "error_message": "Failed audio generation.", } ) traceback.print_exc() raise e def test_cancel_job(): import time # Create a job item_id = str(uuid4()) job_data = json.dumps( dict( id=item_id, prompt_text="This is a test job to be cancelled." * 20, metadata={"tags": "test"}, ) ) model = ChirpV2Stub() job = model.generate.spawn(job_data) # Cancel the job appinbound_event_queue.put( Event(item_id, EventType.CANCEL), partition=item_id, partition_ttl=TOKEN_TIMEOUT_DURATION ) time.sleep(1000) def test_stress_test(): import random random.seed(1) inputs = [] genres = ["edm", "rap", "rock", "pop", "country", "jazz", "classical", "metal", "blues"] for i in range(9000): uid = str(uuid4()) inputs.append( json.dumps( dict( id=uid, prompt_text=f"""Hello world {i}, """ * random.randint(1, 200), metadata={"tags": f"{genres[i % len(genres)]}"}, ), ), ) model = ChirpV2Stub() for input in inputs: model.generate.spawn(input) time.sleep(20000000) @app.local_entrypoint() def main(): # test_cancel_job() test_stress_test()