通过持久化流在Jetson上部署本地AI服务

# 在 Jetson 上通过持久化流服务本地 AI 来源：https://s2.dev/blog/local-ai 随着本地 AI 越来越实用，我想自己托管模型，独立运行工作负载，不依赖任何第三方提供商，同时还想研究如何将本地模型可靠地提供给部分用户使用。NVIDIA 的 Jetson 系列是一个很好的起点，我选择了 Jetson Orin Nano Super 套件 (https://developer.nvidia.com/embedded/jetson-developer-kits)，也就是所谓的“最实惠的生成式 AI 超级计算机”！它拥有 `1024 个 CUDA 核心` 和 `32 个张量核心`，额定性能为 `67 TOPS`（每秒万亿次操作），对于我用 `Kokoro-82M` (https://github.com/hexgrad/kokoro)（一个神经文本转语音模型）构建的小型文本转语音应用来说，应该足够了。这个想法主要源于一个需求：我不想总是阅读大量文字，而是更愿意听。所以我想要一个东西：选择一些文本，挑选一个声音，然后获得一个链接，以后可以回来继续听或者分享给别人。目前这意味着把文本粘贴到一个页面里，但我最终想要一个更懒人友好的方案，那就是在同一个核心应用之上构建一个更好看的前端。除了应用本身，我还想确立一个用于本地推理的小型参考架构：一个自包含的服务层，暴露干净的 API，这样同一套设置可以支持 Web 应用、CLI 或其他服务，而无需重新改造。试一下 streamtts.dev (https://streamtts.dev/) （它就在我的 Jetson 上自托管！😉）：StreamTTS 演示：粘贴文本，选择声音，然后跟随音频流式播放 ## 不是普通的请求/响应 API (https://s2.dev/blog/local-ai#not-a-normal-requestresponse-api) 最简单的架构是这样的： ``` POST /generate 等待 返回 audio.mp3 ``` 推理比普通的 Web 请求慢。Kokoro 在这个 Jetson 上可以比实时更快地生成语音，但它仍然是 GPU 任务。一分钟的音频可能需要几秒钟的计算。冷启动的第一个句子可能更慢，因为模型堆栈需要预热。如果多个用户同时提交，一个阻塞请求就会变成一堆套接字在等待 GPU。 输出本身也是增量式的。TTS 不需要完成整个段落才让听众听到任何东西。模型可以生成一个句子，将该句子编码为 MP3，附加到某处，然后继续。如果我强迫整个过程都放在单个响应体中，就失去了这个工作负载的最佳特性。 而且我希望结果可以共享。用户应该立即被引导到一个链接，在那里他们可以“等待”模型生成所有字节。如果他们在 Jetson 还在工作时打开链接，他们应该能听到已经生成的前缀，然后跟随实时边缘。如果我们从请求-响应开始，最终会添加一堆基础设施，比如： - 队列 - 用于任务记账的数据库 - 用于存储完成文件的对象存储 - 重试逻辑 - 去重逻辑 - 清理进程 所有这些都合理。但合起来，对于一个基本承诺来说，事情就太多了： ``` 现在接受工作 稍后产生输出 让读者跟随 ``` 请求的生命周期似乎不适合这个场景。我希望推理作业在网络中断时也能无缝工作。我也不希望一个浏览器标签页被关闭就导致正在进行的生成终止。因此，输出在完成之前就应该有一个标识，读者应该能够从头开始、追赶到尾部、或者稍后回来重放相同的字节！ 总结一下，我想要： ``` 提交工作 立即获得输出流 工作者追加模型输出 客户端等待流 ``` 所有这些都可以通过持久化流干净地抽象出来。流是一个有序的记录序列，记录只是一些字节（这里是一段音频加上一点元数据）。持久化意味着每条记录都被持久保存，因此不会丢失任何东西，读者以后可以回来重放完全相同的字节。把这两者结合起来，我们就得到了一个简单但强大的构建块。向尾部追加记录，读者可以从头部开始，查找到已知的序列号，或者坐在尾部等待下一条记录到达。流存储为您提供命名的时间线： ``` 追加 记录 从 seq_num=N 读取 尾部 等待实时记录 ``` 每条记录是进度的单位。一条记录有序列号、时间戳、标头和正文。StreamTTS 不需要比这更多的结构。我们这样表示记录： ``` 标头: e: audio i: 3 d: 4210 t: "句子文本" 正文: # e = 事件类型 # i = 索引 # d = 持续时间 (ms) # t = 句子文本 ``` 输出将形如： ``` pub/casts/4LwnHZDl_vFC seq 0 元数据 seq 1 开始 seq 2 音频 句子 0 seq 3 音频 句子 1 seq 4 音频 句子 2 seq 5 eos # 流结束 ``` 架构：浏览器提交一个“播报”，Web 层在 S2 Lite 上申领流，工作者读取任务并追加音频，浏览器通过网关回读 这个流就是音频文件、实时馈送、回放日志和进度指示器。它也是 Web 服务器、GPU 工作者以及每个打开链接的浏览器之间的契约。写入者不需要知道谁在收听。读取者不需要知道写入者是否还活着。双方只同意一个命名的记录序列。 只有连接的 SSE 或 WebSocket 对于实时传输很好，但它们本身不提供持久化回放。它们将字节传递给当前连接的客户端。它们本身不会记住为晚到、断开或刷新页面的客户端准备的字节。所以如果没有人连接，WebSocket 消息就没有地方持久化发送。如果客户端断开，服务器需要其他存储来记住客户端错过的东西。如果第二个监听者在生成仍在运行时打开同一个链接，WebSocket 连接不会告诉服务器如何重放开头然后跟随实时边缘。你绝对可以通过在 SSE/WebSocket 旁边放置一个数据库或对象存储来解决这个问题。但现在实时传输和回放是两个必须一致的不同部分。有了持久化流，这种分裂可以统一！工作者只追加一次输出，实时监听者尾部跟踪流。晚到的监听者可以从 `seq_num=0` 读取，然后尾部跟踪同一个流。回放和实时播放是相同的读取路径，只是从不同的偏移量开始。 ## S2-Lite (https://s2.dev/blog/local-ai#s2-lite) S2 Lite 是一个开源自托管的、单二进制实现的 S2 持久化流 API。在这个设置中，它在 localhost 上运行，使用本地磁盘进行持久化存储，并提供带有追加、读取、尾部长轮询语义的流。 ``` s2 lite --local-root var/s2lite-data --port 4002 --no-cors ``` 我们首先创建一个 basin（命名空间），并将整个服务建模为少数几个命名的流。下面的箭头显示了哪个组件向哪个流追加数据，以及哪个组件从哪个流读取数据： 流数据流：Web 层追加任务并申领 catalog 和 cast 流，工作者读取任务并追加音频，网关尾部跟踪 cast 流 一些流在所有播报之间共享： - `jobs` 是摄入日志：每个推理请求一条记录 - `jobs/_cursor` 持有工作者提交的、在 `jobs` 中的读取偏移量 - `jobs/dead` 收集所有重试都失败的任务 - `progress/done` 为每个完成的播报收到一条收据 每个播报还有自己的两个流： - `catalog/` 是私有配方：完整文本、声音、标题、创建时间 - `pub/casts/` 是公共输出流：元数据、开始、音频...、eos ``` s2 = S2( os.environ.get("S2_ACCESS_TOKEN", "local-token"), endpoints=Endpoints( account=lite_url, basin=lite_url, ), ) config = BasinConfig( default_stream_config=StreamConfig( storage_class=StorageClass.EXPRESS, retention_policy=RETENTION_SECS, ) ) await s2.ensure_basin(basin, config=config) ``` 每个 `audio` 记录在其标头中携带句子文本和持续时间（毫秒），正文中是原始 MP3 字节。文本为浏览器提供字幕和跳转点。持续时间让播放器调度块。浏览器总是从 `seq_num=0` 开始尾部跟踪。如果流已完成，浏览器读取到 `eos` 然后停止。如果工作者仍在追加，浏览器读取现有前缀，到达尾部，然后等待下一条记录。 浏览器播放器也是围绕流的形状构建的。它不使用 Media Source Extensions 或构建一个不断增长的 MP3 文件。每个 `audio` 记录是一个完整的句子大小的 MP3 块。浏览器接收每个句子大小的 MP3 块，通过 Web Audio API (https://developer.mozilla.org/en-US/docs/Web/API/Web_Audio_API) 解码，并放置在虚拟时间线上。 ## 公平调度 (https://s2.dev/blog/local-ai#fair-scheduling) 单个 Jetson 不能像弹性推理集群那样运行 😅。假设有三个人提交文本，我不希望第一个长段落完全完成，而其他人都在等待。工作者保持多个播报活跃，并跟踪每个流相对于挂钟播放的提前量： ``` def lead(self) -> float: return self.total_ms / 1000.0 - (time.monotonic() - self.started) ``` 正向提前量意味着该流已经生成了比当前播放更靠前的音频缓冲区。负向提前量意味着监听者正在追赶实时尾部。调度循环是： ``` 在并发限制内接收任务 选择提前量最小的活跃流 为该流精确生成一个句子 追加该句子 重新计算提前量 重复 ``` 当每个活跃流都舒适地提前时，工作者会短暂休眠，而不是全力完成一个流，从而产生实时输出调度。目标是让多个公共流保持可播放。公平的单位不是请求，而是一个追加的句子。 ## 提交工作 (https://s2.dev/blog/local-ai#submitting-work) 当请求到达时，Web 进程不会加载模型。它验证文本和声音，计算一个确定性的 ID，并创建一个音频将要出现的位置。该 ID 是内容寻址的： ``` def content_id(text: str, voice: str) -> str: h = hashlib.sha256(f"{voice}\x00{text.strip()}".encode()).digest() return base64.urlsafe_b64encode(h).decode().rstrip("=")[:12] ``` 相同文本和相同声音映射到同一个流。这使重复提交变成缓存命中。写入路径是： 1. 申领 `catalog/` 使用完整配方 2. 申领 `pub/casts/` 使用一个元数据记录 3. 向 `jobs` 流追加一个任务 4. 返回 `/c/` 关键操作是申领。S2 支持带 `match_seq_num` 的条件追加。StreamTTS 使用 `match_seq_num=0`，意思是“只有在这个流为空时才追加”。 ``` payload = { "records": [{"body": json.dumps(body, separators=(",", ":"))}], "match_seq_num": 0, } ``` 如果两个人同时提交相同的文本，恰好一个请求赢得申领并将任务加入队列。另一个获得相同的链接并尾部跟踪同一个输出流。那一次追加替代了锁表、唯一性约束和去重缓存。 ## 工作者是一个持久化消费者 (https://s2.dev/blog/local-ai#the-worker-is-a-durable-consumer) 工作是唯一拥有模型并接触 GPU 的进程。它从 `jobs` 流中读取，运行 Kokoro-82M，并将音频记录追加到播报流。启动时，工作者读取上次提交的偏移量来自 `jobs/_cursor`： ``` jobs/_cursor {"offset": 123} ``` 然后它从该偏移量开始读取 `jobs`。如果没有新内容，它在尾部长轮询。微妙的部分是提交游标。StreamTTS 一次可能有多个活跃播报，它们不一定按任务顺序完成。短任务 10 可能在长任务 7 之前完成。游标只有在所有之前任务都完成时才能向前移动。工作者使用连续完成的 watermark： ``` def advance_watermark(): nonlocal committed moved = False while committed in done_above: done_above.discard(committed) committed += 1 moved = True if moved: self._commit_offset(committed) ``` 如果进程崩溃，没有特殊的恢复协议。重新启动后，工作者从上次提交的偏移量恢复。该偏移量之后的任务被重新读取。已经完成的播报通过检查其输出流是否以 `eos` 结束来跳过。未完成的播报重新运行。这就是至少一次交付加上幂等输出。对于已完成的播报，它的行为就像恰好一次，因为 `eos` 是持久化完成标记。我们也可以使用一个栅栏令牌，令牌是一个终端标记，用来标记播报已完成。 重试可能会在流中留下部分音频。因此开始记录是一个尝试边界： ``` seq 0 元数据 seq 1 开始 尝试 1 seq 2 音频 句子 0 # 工作者崩溃 seq 3 开始 尝试 2 seq 4 音频 句子 0 seq 5 音频 句子 1 seq 6 eos ``` 播放器可以将最新的开始视为可播放尝试的开始，并忽略更早的部分音频。 ## 服务读者 (https://s2.dev/blog/local-ai#serving-readers) 公共读取路径故意比内部 S2 API 窄。S2 Lite 可以写入、删除和读取任何流，但身份验证/授权留给用户自定义。因此，浏览器通过一个 StreamTTS 网关读取，该网关只允许公共播报流： ``` GET /s2/records?stream=pub/casts/&seq_num=0 ``` 网关拒绝内部流如 `jobs` 和 `catalog/*`。它还为应用提供了速率限制读取的地方。对于实时播放，同一个网关提供 SSE。S2 Lite 在内部为许多读取者共享同一个上游尾部（一个广播发送者喂养所有尾部读取者），因此网关只是将该尾部中继到浏览器。网关将一个 S2 Lite 尾部转换为给多个浏览器的 SSE 慢速客户端仍然不会对系统造成背压：每个订阅者有一个有界队列，如果队列满了，网关会丢弃该客户端而不是阻塞流。 ## 一些见解 (https://s2.dev/blog/local-ai#some-insights) 在热生成期间，`tegrastats` (https://docs.nvidia.com/jetson/archives/r36.5/DeveloperGuide/AT/JetsonLinuxDevelopmentTools/TegrastatsUtility.html) 大致如下所示： ``` GR3D_FREQ 0% VDD_IN 3295mW 空闲 GR3D_FREQ 99% VDD_IN 9911mW 生成中，[email protected] GR3D_FREQ 0% VDD_IN 6955mW 完成 ``` `GR3D_FREQ` 是 GPU 利用率。模型在生成时短暂地占用 GPU，但整个板子在这个工作负载下保持在约 `10 W` 以下，温度从未超过约 46°C。更有用的性能数据来自 `progress/done` 收据。每条收据包括 `sentences`、`audio_ms` 和 `gen_ms`，这让我可以计算 `xRT`：每秒计算产生的音频秒数。 ``` 句子 音频_ms 生成_ms xRT 声音 3 11221 2567 4.37 am_michael 2 5205 1670 3.11 af_heart 3 28224 11735 2.40 af_heart 1 917 2481 0.36 af_heart 冷启动第一个句子 ``` 预热后，生成速度大约是 `2.4x–4.4x` 实时。设备空闲后的第一个句子可能在模型重新预热时低于实时；这种冷启动行为正是调度缓冲区要隐藏的。在大约 3 倍实时速度下，同时三个实时播报是一个合理的心智模型，对这个用例来说已经足够了。 一个自托管的 AI 电台？📻 我们也可以在 *输入* 上放置一个持久化流：一个 LLM 将令牌发射到一个流中，TTS 工作者尾部跟踪该流，按自己的节奏生成语音，这样就可以建立一个有趣的电台频道！ ## 日志最大化 (https://s2.dev/blog/local-ai#logmaxxing) 这个应用的传统版本会使用几个不同的组件：用于后台任务的队列、用于状态和重试的数据库、用于存储完成 MP3 的对象存储、用于实时播放的 WebSocket 或 SSE、以及用于保留和清理的逻辑，将一个简单的流程分散到多个系统中。有了持久化流，大部分都合并为命名的日志：`jobs` 就是队列。