• 人工智能
• 大模型
• 模型推理服务
• 后端

plan_weight_upload将一份完整的运行时权重放置（runtime weight placement）转换为不可变的WeightUploadPlan。该计划向 Store 写入端提供规范的StoredWeightManifest、payload 对象位置，以及每次上传操作对应的来源证据（source evidence）。本文基于 Mooncake 仓库中 模型权重存储上传规划设计文档 与 mooncake-reshard 实现源码，深入讲解上传规划的输入约束、副本选择算法、计划内容结构、快照写入 API、恢复路径与原生 Store 依赖，帮助你理解并复现这套"一次完整权重放置 → 一个不可变上传计划"的机制。

1. 设计动机：为什么需要显式的"上传规划"层

在大模型推理/训练场景中，权重（model weight）需要从分布式运行时（多机多卡、TP/PP/EP/DP 并行切分）迁移到持久化 Store，再在任意目标放置下恢复。这个过程有两个天然风险：

1. 来源不一致：运行时权重可能处于 lease 租约期、worker 实例可能变化、多个 DP 副本的 generation 可能不一致，若直接按"当前内存地址"上传，容易读到半更新状态。
2. 布局信息丢失：Store 端保存的只是字节对象，若从 Store key 反推模型布局，会产生脆弱耦合。

Mooncake Reshard 的解法是：先由框架适配器（framework adapter）导出无地址的规范几何信息（tensor 全局形状、并行轴切分、片段几何），再由 planner 把所有绑定（binding）与放置（placement）校验、筛选并固化为一个不可变计划。Store 写入端只消费计划、执行 I/O，从不推断模型布局或并行度。

2. 输入：两份 Manifest 的契约

上传规划接受两类输入（见 上传规划实现）：

Planner 会对每个提供的绑定逐一校验其与放置的一致性（pair_manifests配对、tensor 描述符一致性校验等）。模型语义（如 tensor 如何切分、并行轴含义）保持在框架适配器中，适配器在调用本 API 之前即导出规范的 tensor 与并行轴元数据。

2.1 强制约束：上传必须来自"复制语义"的 DP 张量

实现中有一个关键限制：_collect_upload_sources会先检查是否存在has_dp_ownership(tensor)，即 DP 是否被声明为张量归属轴（每个 DP rank 只持有部分数据）。若存在这样的 DP-owned 张量，直接抛出ValueError：

"Weight Store upload requires replicated DP source tensors; DP-owned tensor snapshots are not supported"

这意味着：Store 上传的完整副本语义要求 DP 是复制轴（ReplicatedAxis），每个 DP rank 都持有完整张量；DP 作为所有权轴（每个 rank 持有一部分）的张量快照不支持直接上传，因为无法选择"一个完整副本"。这也是下面副本选择算法能成立的前提。

3. 副本选择（Replica Selection）：只存一份完整来源副本

规划器只为每次上传存储一个完整来源副本，算法分四步（对应_collect_upload_sources与complete_parallel_source_replicas，见 ownership.py）：

1. 要求复制语义的 DP 语义 + 完整逻辑覆盖：对每个 DP rank，检查该 rank 上的所有非 DP-owned 张量的片段是否无重叠地精确覆盖完整张量（boxes_exactly_cover，用全局偏移 + 局部形状做盒子覆盖校验）。只有"完整覆盖"的 DP rank 才成为候选副本。
2. 要求该副本内每个选中的来源绑定具有相同 generation：对每个候选 DP rank，收集其片段所属绑定 manifest 的 generation；若集合大小恰为 1，才把该 rank 记为"generation 一致"。全部候选都没有一致 generation 时抛错"source manifests have no complete generation-consistent DP replica"；多个候选副本 generation 不一致时抛"complete source DP replicas have inconsistent lease generations"。
3. 确定性选择最低 DP rank：selected_dp = min(generations_by_dp)，保证同一组输入永远选出同一个副本，便于复现与对账。
4. 保留每个选中片段的逻辑 TP/PP/EP 归属：选中的片段保留其rank.tp/pp/ep与parallel_tensor_owner判定结果，Store 侧不需要也不保留运行时地址。

最终 manifest 中，每个被选中的逻辑来源片段恰好对应一个 stored fragment；DP 副本不会重复 payload 对象（一个完整副本只写一份字节）。

关于筛选细节：candidates以(tensor_id, global_offset, local_shape)为几何键分组，组内用_runtime_sort_key（按dp、pp、ep、tp排名的顺序，再按worker_id、fragment_id）排序后取第一个，随后按(tensor_id, global_offset, local_shape)排序输出。

测试 test_upload_planning.py 覆盖了上述全部行为：test_upload_plan_selects_one_complete_generation_consistent_dp_replica（选中唯一完整且 generation 一致的副本）、test_upload_plan_rejects_complete_dp_replicas_at_different_generations（跨副本 generation 不一致被拒）、test_upload_plan_rejects_incomplete_dp_replica（不完整副本被拒）、test_upload_plan_rejects_dp_owned_source_tensor（DP-owned 张量被拒）、test_upload_plan_preserves_tp_pp_and_ep_fragment_ownership（TP/PP/EP 归属被保留）。

4. 计划内容（Plan Contents）：WeightUploadPlan的结构

WeightUploadPlan是不可变数据类（定义见 contracts.py），包含：

• manifest：StoredWeightManifest，其group、manifest、payload key 全部不可变（对象创建后不可修改）；
• operations：每个 stored fragment 对应一个UploadOperation；
• source_placement_id与source_placement_digest：来源放置身份与摘要，用于上传时校验来源未漂移；
• transaction_group_id与control_key：上传事务组及其决策（decision）控制键。

4.1UploadOperation：一份来源证据 + 目标几何

每个UploadOperation携带：

构造期做了大量契约校验：来源放置与绑定的 placement_fragment_id 必须一致；来源片段与目标片段的 tensor_id/global_offset/local_shape/nbytes 必须完全一致；target.object_offset必须为 0；source_snapshot.lease_generation必须等于source_generation。

4.2 Key 布局与 fragment 摘要

plan_weight_upload生成的 key 布局（namespace默认"default"、key_prefix默认"weights"，均校验非空）：

{key_prefix}/{namespace}/{resource_id}/{revision}/{weight_generation} ├── /manifest # StoredWeightManifest 本体 ├── /payload/{transaction_id}/{fragment_id} # 每个片段的字节对象 └── /transactions/{transaction_id} └── /decision # 事务决策控制键

其中：

• transaction_id为uuid4().hex；
• 每个fragment_id是对"{tensor_id}|{global_offset}|{local_shape}"做SHA-256 摘要并截取前 24 个十六进制字符（_fragment_digest），因此几何相同的片段在任何机器、任何代际下得到同一 fragment_id——这是幂等与去重的基础；
• 整个StoredWeightManifest还有manifest_digest：对to_json()（sort_keys=True、紧凑分隔符）的字节流做 SHA-256，供manifest_identity使用；
• WeightUploadPlan的__post_init__还会验证：transaction group 必须属于 manifest group；control_key必须等于{transaction_group_id}/decision；operations 的目标集合必须与 manifest fragments 完全一致；payload object_key 必须属于本事务。

关于 payload 写入端：文档指出 "The payload writer rebinds this evidence to a fresh runtime manifest and acquires the framework allocation guard before Store I/O"。即上传时WeightUploadService.upload会先做一系列新鲜度校验（validate_manifest_pair、revision/generation/placement_id/digest 匹配、instance_id 陈旧检测、lease/generation 陈旧检测、same_runtime_snapshot比对），必要时通过acquire_weight_binding_token获取新的绑定 token 与 allocation guards，再执行batch_put_from批量写入；失败时以FAILED_DRAINED终态释放 token（见 upload.py）。

5. 执行边界（Execution Boundary）：Store 写入端只消费计划

• Store writer 负责：payload 写入、注册（registration）、事务提交/中止（commit/abort）、Store 到运行时读取（Store-to-runtime reads）。
• Store writer 不负责：从 Store key 推断模型布局或并行度（模型语义在框架适配器侧）。

这种边界隔离带来两个好处：一是 Store 端逻辑与框架无关，可复用于多种框架适配器；二是 key 布局、事务结构可以自由演进而不影响模型语义。

事务侧由WeightUploadTransaction（transaction.py）管理：require_writable（检查/认领事务可写性）、commit（把决策写入control_key并发布 manifest）、abort_upload（清理 payload）。注意 writer 的abort()在_commit_decision_may_exist为真时不允许 abort，提示"retry commit instead"——避免在决策边界模糊时误清理已可能提交的数据。

6. 快照 API（Snapshot API）：begin_weight_snapshot

快照入口在 entrypoint.py 与 store.py：

MooncakeDistributedStore.begin_weight_snapshot(descriptor, adapter) # -> WeightStoreWriter

调用流程（WeightStoreWriter，见 writer.py）：

1. 构造时，adapter.export_source(snapshot)一次性导出完整来源放置与活体绑定（WeightSnapshotSource），并校验resource_id/revision/weight_generation与WeightSnapshotDescriptor一致；
2. 随即调用weight_store.plan_upload(...)生成内部WeightUploadPlan；
3. 调用方逐 tensor 写入：writer.write_tensor(tensor_id, tensor)——适配器的resolve_fragment_ids(tensor_id, tensor, source)把框架 tensor 映射为规范 placement fragment ID，并拒绝空解析、重复解析、跨快照片段、跨 tensor 片段、重复提交；
4. 某个参与者（participant）的必需片段全部提交后，自动_flush_participant触发该参与者的实际上传（WeightUploadService.upload）；
5. 全部片段就绪后，writer.commit()才发布StoredWeightManifest——commit 前任何缺失片段都会 abort 并抛错（"Weight snapshot is missing required fragments"）；
6. 支持with上下文管理器：正常退出自动 commit，异常退出自动 abort（除非 commit decision 可能已存在）。

WeightSnapshotAdapter协议（snapshot.py）定义了三个方法：export_source、resolve_fragment_ids、source_allocation_guards（返回框架侧 lifetime guards，供上传前绑定）。该协议正是"模型语义留在框架适配器"这一设计落地的载体。

6.1 存储策略归属

文档明确：writer 拥有其不可变 payload 与元数据对象的存储策略；逐 tensor 的复制（replication）、分片（partition）、upsert 参数在本 API 之外。因此每次提交的快照都有且只有一个显式 manifest 契约（one explicit manifest contract），不存在"同一快照被不同策略部分写入"的歧义。

7. 恢复路径（Restore）：从StoredWeightManifest出发

恢复从StoredWeightManifest开始（加载实现见 load.py）：

1. load_manifest(manifest_key)：store.get(manifest_key)读取 JSON →StoredWeightManifest.from_json严格反序列化（拒绝重复字段、非有限数值、字段缺失）→ 校验manifest.manifest_key与请求 key 一致；
2. loader 依据每个 stored fragment 的tensor_id、全局偏移、局部形状、object key、object offset、字节长度重建来源几何；
3. 结合目标放置与运行时绑定，经plan_stored_transfer_to_target_placement+bind_logical_transfer_plan生成WeightLoadPlan（其校验保证 transfer 的 resource/revision/generation 与 manifest 一致、来源片段与 manifest fragments 一致）；
4. 最终通过get_into_ranges把对象范围读入目标缓冲（StoreBackend.get_into_ranges支持prepare_get_into_ranges_snapshot+get_into_ranges_from_snapshot的快照读路径，见 backend.py）。

这条路径不使用 legacyTensorMetadata头，也不使用 legacy 并行 tensor 重建元数据——恢复所需的一切几何信息都内联在StoredWeightManifest中。

7.1 持久化边界：什么进 Store，什么留在运行时

• 进 Store（manifest 内）：逻辑 tensor 描述符（TensorDescriptor：global_shape、dtype、itemsize、layer_id、expert_id、layout_fingerprint、shard_dims、parallel_axes）、片段几何（fragment_id、global_offset、local_shape、nbytes、aliases）、payload key、快照身份（namespace、resource_id、revision、weight_generation、group_id、manifest_key、created_at、digest）；
• 留在活体绑定/守卫路径（不进 Store）：运行时地址（address）、分配所有者（allocation owners）、lease、worker 实例。这些是易失的、随进程生命周期变化的，Store 保存它们没有意义且会引入陈旧地址风险。

8. 原生 Store 要求（Native Store Requirement）：ReplicateConfig.group_ids

原生 Store 写入端使用group 语义，确保 payload、manifest、事务控制对象分别落在其声明 group 中。这依赖 Mooncake wheel 的ReplicateConfig暴露group_ids字段（即文档所述 PR #3000 引入的 API）。

在 structured_object_store.py 中可以看到：

• 第 559-581 行：getattr(config, "group_ids", None)读取group_ids，支持字符串或字符串列表；若列表长度 > 1 则抛错要求"恰好一个 group"，然后以该 group_id 生成写配置；
• 第 6594-6620 行附近：group_ids被纳入已知配置键，并提供按物理 key 数量对齐 group 的辅助逻辑（normalize_to_list、空列表时置None走非分组路径）。

适配器在 Store I/O 开始前，若检测到旧版绑定（无group_ids的ReplicateConfig），会直接拒绝（"The adapter rejects an older binding before Store I/O starts"）。上传时 payload 写配置通过config_factory([plan.manifest.group_id] * len(batch), "payload")统一挂到 manifest group 下（upload.py），从而保证 payload 与 manifest、事务对象同组隔离、可统一管理。

9. 小结：一次上传规划的生命周期

把全文串起来，plan_weight_upload的完整生命周期是：

1. 框架适配器导出WeightPlacementManifest+WeightRuntimeBindingManifest（模型语义在此终结）；
2. planner 校验所有绑定，剔除 DP-owned 张量，选出完整、generation 一致、最低 DP rank的单一来源副本，保留 TP/PP/EP 逻辑归属；
3. 生成不可变WeightUploadPlan：规范 manifest、逐片段UploadOperation（含来源证据）、来源身份/摘要、事务组与控制键；
4. Store writer 按write_tensor→ 参与者级 flush →commit()的顺序，在 allocation guard 保护下执行batch_put_from，最后在control_key写入决策并发布 manifest；
5. 恢复端从StoredWeightManifest反序列化几何，经plan_stored_transfer_to_target_placement+get_into_ranges读回目标运行时。

整个设计的关键取舍可以概括为三句话：逻辑几何进 Store、运行时地址留内存；每快照一份显式 manifest 契约；Store 端永远不做布局推断。这套机制使权重上传具备可复现性（确定性副本选择）、一致性（generation/lease/instance 三重新鲜度校验）与可审计性（digest + 决策控制键），是 Mooncake Reshard 权重存储链路的核心底座。

• 人工智能
• 大模型
• 模型推理服务
• 后端