scripts/orpheus_streaming_test.py

"""
Teste REALISTA de Streaming - Simula usuários com frases variadas
Verifica se o áudio trava para algum usuário durante streaming contínuo
"""
import os
os.environ["VLLM_ATTENTION_BACKEND"] = "FLASH_ATTN"

import torch
import time
import asyncio
import random
from transformers import AutoTokenizer
from vllm import AsyncLLMEngine, AsyncEngineArgs, SamplingParams

START_TOKEN = 128259
END_TOKENS = [128009, 128260, 128261, 128257]
STOP_TOKEN = 128258
AUDIO_TOKEN_BASE = 128266

# Frases variadas - curtas, médias e longas (simulando aula de idiomas)
PHRASES_SHORT = [
    "Hello!", "Yes.", "No.", "Thank you.", "Good morning.",
    "How are you?", "I'm fine.", "See you!", "Bye!", "Nice!",
]

PHRASES_MEDIUM = [
    "I would like to order a coffee, please.",
    "The weather is really nice today, isn't it?",
    "Can you help me find the train station?",
    "I'm learning English and it's very interesting.",
    "What time does the movie start tonight?",
    "My favorite color is blue, what about yours?",
    "I work as a software engineer in the city.",
    "Do you have any recommendations for dinner?",
]

PHRASES_LONG = [
    "I've been studying English for about three years now, and I find it fascinating how much progress I've made since I started.",
    "Yesterday I went to the supermarket and bought some fruits, vegetables, and other groceries for the whole week.",
    "The conference will be held next Monday at the main auditorium, and all employees are expected to attend the presentation.",
    "Learning a new language opens up so many opportunities for travel, career advancement, and making new friends from different cultures.",
    "Could you please explain the process step by step so that I can understand it better and apply it correctly in my work?",
]

def get_random_phrase():
    """Retorna frase aleatória com distribuição realista"""
    r = random.random()
    if r < 0.4:  # 40% frases curtas
        return random.choice(PHRASES_SHORT)
    elif r < 0.8:  # 40% frases médias
        return random.choice(PHRASES_MEDIUM)
    else:  # 20% frases longas
        return random.choice(PHRASES_LONG)

async def main():
    print("=" * 70)
    print("TESTE REALISTA DE STREAMING - FRASES VARIADAS")
    print("=" * 70)

    print("\n[1] Carregando modelo...")
    tokenizer = AutoTokenizer.from_pretrained("canopylabs/orpheus-3b-0.1-ft")

    engine_args = AsyncEngineArgs(
        model="canopylabs/orpheus-3b-0.1-ft",
        dtype="bfloat16",
        max_model_len=4096,
        gpu_memory_utilization=0.95,
        max_num_seqs=256,
        enable_chunked_prefill=True,
        enable_prefix_caching=True,
        enforce_eager=False,
    )
    engine = AsyncLLMEngine.from_engine_args(engine_args)

    sampling_params = SamplingParams(
        temperature=0.2,
        top_p=0.9,
        max_tokens=4096,
        stop_token_ids=[STOP_TOKEN],
        repetition_penalty=1.1,
    )

    def format_prompt(text, voice="tara"):
        adapted_prompt = f"{voice}: {text}"
        prompt_tokens = tokenizer(adapted_prompt, return_tensors="pt")
        start_token = torch.tensor([[START_TOKEN]], dtype=torch.int64)
        end_tokens = torch.tensor([END_TOKENS], dtype=torch.int64)
        all_input_ids = torch.cat([start_token, prompt_tokens.input_ids, end_tokens], dim=1)
        return tokenizer.decode(all_input_ids[0])

    async def simulate_streaming(text, user_id):
        """
        Simula streaming real - verifica se cada chunk chega a tempo
        Retorna True se streaming foi fluido, False se travou
        """
        prompt_string = format_prompt(text)
        start = time.time()

        ttff = None  # Time to first frame
        audio_tokens = 0
        last_frame_time = None
        frame_gaps = []  # Tempo entre frames
        max_gap = 0
        frames_generated = 0

        # Cada frame de 7 tokens = ~23ms de áudio
        # Se o gap entre frames > 23ms, o áudio vai travar
        FRAME_DURATION = 0.023  # 23ms por frame

        async for output in engine.generate(prompt_string, sampling_params, f"user_{user_id}"):
            current_audio = sum(1 for t in output.outputs[0].token_ids if t >= AUDIO_TOKEN_BASE)
            current_frames = current_audio // 7

            # Novo frame gerado?
            if current_frames > frames_generated:
                now = time.time()

                if frames_generated == 0:
                    ttff = now - start
                    last_frame_time = now
                else:
                    gap = now - last_frame_time
                    # Normalizar pelo número de frames gerados de uma vez
                    new_frames = current_frames - frames_generated
                    gap_per_frame = gap / new_frames
                    frame_gaps.append(gap_per_frame)
                    max_gap = max(max_gap, gap_per_frame)
                    last_frame_time = now

                frames_generated = current_frames

        total_time = time.time() - start
        audio_duration = frames_generated * FRAME_DURATION

        # Streaming travou se algum gap foi maior que o dobro da duração do frame
        # (dando margem para buffering)
        BUFFER_MARGIN = 3.0  # 3x a duração do frame como margem
        streaming_ok = max_gap < (FRAME_DURATION * BUFFER_MARGIN) if frame_gaps else True

        avg_gap = sum(frame_gaps) / len(frame_gaps) if frame_gaps else 0

        return {
            'user_id': user_id,
            'text_len': len(text),
            'ttff': ttff or 0,
            'total_time': total_time,
            'audio_duration': audio_duration,
            'frames': frames_generated,
            'avg_gap': avg_gap,
            'max_gap': max_gap,
            'streaming_ok': streaming_ok,
            'rtf': total_time / audio_duration if audio_duration > 0 else 999,
        }

    print("\n[2] Warmup...")
    await simulate_streaming("Hello world.", 0)

    print("\n[3] Testando streaming realista...")
    print("=" * 70)

    results_summary = []

    for num_users in [32, 64, 100, 128, 150, 200, 256]:
        print(f"\n>>> {num_users} USUÁRIOS COM FRASES VARIADAS <<<")

        # Gerar frases aleatórias para cada usuário
        user_phrases = [get_random_phrase() for _ in range(num_users)]

        # Mostrar distribuição
        short = sum(1 for p in user_phrases if len(p) < 20)
        medium = sum(1 for p in user_phrases if 20 <= len(p) < 80)
        long = sum(1 for p in user_phrases if len(p) >= 80)
        print(f"    Distribuição: {short} curtas, {medium} médias, {long} longas")

        try:
            start_batch = time.time()
            tasks = [simulate_streaming(phrase, i) for i, phrase in enumerate(user_phrases)]
            results = await asyncio.gather(*tasks, return_exceptions=True)
            batch_time = time.time() - start_batch

            errors = [r for r in results if isinstance(r, Exception)]
            successful = [r for r in results if not isinstance(r, Exception)]

            if errors:
                print(f"    ERROS: {len(errors)}")

            if successful:
                streaming_ok = [r for r in successful if r['streaming_ok']]
                streaming_bad = [r for r in successful if not r['streaming_ok']]

                max_gaps = [r['max_gap'] * 1000 for r in successful]
                ttffs = [r['ttff'] * 1000 for r in successful]
                rtfs = [r['rtf'] for r in successful]

                print(f"    Streaming OK: {len(streaming_ok)}/{len(successful)}")
                print(f"    TTFF: min={min(ttffs):.0f}ms, max={max(ttffs):.0f}ms, avg={sum(ttffs)/len(ttffs):.0f}ms")
                print(f"    Max Gap: min={min(max_gaps):.0f}ms, max={max(max_gaps):.0f}ms, avg={sum(max_gaps)/len(max_gaps):.0f}ms")
                print(f"    RTF: min={min(rtfs):.3f}, max={max(rtfs):.3f}, avg={sum(rtfs)/len(rtfs):.3f}")

                if streaming_bad:
                    print(f"    ⚠️  STREAMING TRAVOU para {len(streaming_bad)} usuários!")
                    worst = max(streaming_bad, key=lambda x: x['max_gap'])
                    print(f"       Pior caso: user_{worst['user_id']} com gap de {worst['max_gap']*1000:.0f}ms")
                else:
                    print(f"    ✓ Streaming fluido para TODOS os usuários!")

                results_summary.append({
                    'users': num_users,
                    'ok': len(streaming_ok),
                    'bad': len(streaming_bad),
                    'max_gap_worst': max(max_gaps),
                    'ttff_max': max(ttffs),
                    'rtf_max': max(rtfs),
                })

        except Exception as e:
            print(f"    CRASH: {type(e).__name__}: {str(e)[:80]}")
            break

    print("\n" + "=" * 70)
    print("RESUMO - TESTE DE STREAMING REALISTA")
    print("=" * 70)

    print("\n| Users | Stream OK | Travou | Max Gap | TTFF Max | RTF Max |")
    print("|-------|-----------|--------|---------|----------|---------|")
    for r in results_summary:
        status = "✓" if r['bad'] == 0 else "⚠️"
        print(f"| {r['users']:5} | {r['ok']:9} | {r['bad']:6} | {r['max_gap_worst']:6.0f}ms | {r['ttff_max']:7.0f}ms | {r['rtf_max']:7.3f} |")

    # Encontrar limite
    all_ok = [r for r in results_summary if r['bad'] == 0]
    if all_ok:
        max_ok = max(all_ok, key=lambda x: x['users'])
        print(f"\n>>> MÁXIMO SEM TRAVAMENTO: {max_ok['users']} usuários <<<")

    some_bad = [r for r in results_summary if r['bad'] > 0]
    if some_bad:
        first_bad = min(some_bad, key=lambda x: x['users'])
        print(f">>> PRIMEIRO TRAVAMENTO: {first_bad['users']} usuários ({first_bad['bad']} travaram) <<<")

    print("=" * 70)

if __name__ == "__main__":
    asyncio.run(main())