用 SageMaker AI 和 vLLM 构建实时语音应用：从"等音频传完"到"边说边转"

# 假设你已经有一个 vLLM 兼容的语音模型（如 Whisper-large-v3 的 vLLM 适配版）
# 以下使用 vLLM 官方镜像作为基础

# 登录 ECR
aws ecr get-login-password --region us-east-1 | docker login --username AWS --password-stdin 763104351884.dkr.ecr.us-east-1.amazonaws.com

# 拉取 vLLM 基础镜像
docker pull vllm/vllm-openai:latest

# 如果需要自定义 entrypoint 或挂载语音预处理逻辑，可在此基础上构建
# 这里直接使用官方镜像，模型权重从 S3 加载

import sagemaker
from sagemaker.model import Model

role = sagemaker.get_execution_role()
sagemaker_session = sagemaker.session.Session()

# 模型权重已上传到 S3
model_s3_path = "s3://my-voice-models/whisper-large-v3-vllm/"

model = Model(
    image_uri="vllm/vllm-openai:latest",  # 替换为你推到 ECR 的镜像 URI
    model_data=model_s3_path,
    role=role,
    sagemaker_session=sagemaker_session,
    env={
        # vLLM 启动参数
        "MODEL_NAME": "/opt/ml/model",  # SageMaker 标准模型挂载路径
        "VLLM_SERVE_FROM_S3": "0",
        # 开启 streaming
        "ENABLE_STREAMING": "1",
        # GPU 并行与 batch 配置
        "TENSOR_PARALLEL_SIZE": "1",
        "MAX_NUM_SEQS": "32",  # 同时服务 32 路语音流
        "MAX_NUM_BATCHED_TOKENS": "4096",
    },
)

# 部署到 ml.g5.12xlarge（4xA10G），根据模型大小调整实例类型
predictor = model.deploy(
    initial_instance_count=1,
    instance_type="ml.g5.12xlarge",
    endpoint_name="voice-streaming-vllm",
)

import json
import requests

endpoint_url = "https://runtime.sagemaker.us-east-1.amazonaws.com/endpoints/voice-streaming-vllm/invocations"

# 模拟：将音频分 chunk 流式推送
# 实际场景中，chunk 来自麦克风实时采集
audio_chunks = [
    b"<binary_audio_chunk_1>",  # 前 100ms 音频
    b"<binary_audio_chunk_2>",  # 接下来 100ms
    b"<binary_audio_chunk_3>",
    # ... 持续推送
]

# 使用 SageMaker InvokeEndpointWithResponseStream（流式响应 API）
import boto3

sm_runtime = boto3.client("sagemaker-runtime")

# 发送第一个 chunk，开启流式会话
# 实际生产中需配合 WebSocket 或自定义流式协议
response = sm_runtime.invoke_endpoint(
    EndpointName="voice-streaming-vllm",
    ContentType="application/json",
    Body=json.dumps({
        "model": "whisper-large-v3",
        "audio": "<base64_encoded_chunk>",
        "stream": True,
        "language": "zh",
    }),
)

# 解析流式返回
result = json.loads(response["Body"].read())
print(result["text"])  # 增量转录片段

# ---- 更完整的流式调用（SageMaker Streaming API）----
# SageMaker 2024 年新增 InvokeEndpointWithResponseStream
stream_response = sm_runtime.invoke_endpoint_with_response_stream(
    EndpointName="voice-streaming-vllm",
    ContentType="application/json",
    Body=json.dumps({
        "model": "whisper-large-v3",
        "audio": "<base64_encoded_full_or_chunk>",
        "stream": True,
    }),
)

for event in stream_response["Body"]:
    chunk = json.loads(event["PayloadFragment"])
    print(chunk.get("text", ""), end="", flush=True)
# 输出效果：文字逐片段出现，像打字机一样

import pyaudio
import base64
import json
import threading
import queue

CHUNK_DURATION_MS = 100  # 每 100ms 采集一个 chunk
SAMPLE_RATE = 16000
CHUNK_SIZE = int(SAMPLE_RATE * CHUNK_DURATION_MS / 1000)

audio_queue = queue.Queue()

# 音频采集线程
def capture_audio():
    pa = pyaudio.PyAudio()
    stream = pa.open(
        format=pyaudio.paInt16,
        channels=1,
        rate=SAMPLE_RATE,
        input=True,
        frames_per_buffer=CHUNK_SIZE,
    )
    while True:
        data = stream.read(CHUNK_SIZE, exception_on_overflow=False)
        audio_queue.put(data)

threading.Thread(target=capture_audio, daemon=True).start()

# 推理推送线程——实际生产中应使用 WebSocket 持久连接
# 此处简化为逐 chunk 调用
import boto3
sm_runtime = boto3.client("sagemaker-runtime")

full_transcript = ""

while True:
    chunk_data = audio_queue.get()
    encoded = base64.b64encode(chunk_data).decode("utf-8")

    resp = sm_runtime.invoke_endpoint_with_response_stream(
        EndpointName="voice-streaming-vllm",
        ContentType="application/json",
        Body=json.dumps({
            "model": "whisper-large-v3",
            "audio_chunk": encoded,
            "stream": True,
            "is_final": False,  # 标记是否为最后一个 chunk
        }),
    )

    for event in resp["Body"]:
        fragment = json.loads(event["PayloadFragment"])
        if "text" in fragment:
            full_transcript += fragment["text"]
            print(f"[实时] {full_transcript}")