1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
#
"""Tests for the publisher step stamp on the rendezvous blob.
The stamp answers a question a receiver otherwise has to guess at: does the shard
table I just discovered describe the step I am refitting, or the one before it? It
matters because the table carries the per-shard digests a verify gate compares
against, so a table one step behind turns correctly delivered bytes into a reported
corruption.
What is worth testing here is not that a field round-trips. It is that the stamp is
optional in *both* directions, because a fleet is upgraded one image at a time and
neither side may assume the other has it, and that a missing stamp reads as "unknown"
rather than as step 0.
"""
from __future__ import annotations
import json
from modelexpress.refit.reshard.rendezvous import (
unwrap_rendezvous_blob,
wrap_rendezvous_blob,
)
def _blob(step=None, tensors=None):
return wrap_rendezvous_blob(b"meta", "agent-r0", "h:1", tensors or [], step)
def test_the_stamp_round_trips():
assert unwrap_rendezvous_blob(_blob(step=7)).publisher_step == 7
def test_an_unstamped_blob_reads_as_unknown_not_step_zero():
"""The distinction a staleness check has to rest on. Read as 0, an unstamped
publisher looks permanently behind, and a consumer that excuses lagging
publishers would excuse every one of its shards."""
assert unwrap_rendezvous_blob(_blob()).publisher_step is None
def test_step_zero_is_carried_not_dropped():
"""0 is a legitimate step and must not collapse into "no stamp"."""
assert unwrap_rendezvous_blob(_blob(step=0)).publisher_step == 0
def test_the_stamp_is_omitted_rather_than_nulled_when_absent():
"""So an unstamped publisher emits the blob an older client would have."""
payload = json.loads(_blob().decode("utf-8"))
assert "publisher_step" not in payload
def test_the_stamp_is_the_only_added_key():
"""A wire-format change should be auditable by diffing the key sets."""
without = set(json.loads(_blob().decode("utf-8")))
with_step = set(json.loads(_blob(step=4).decode("utf-8")))
assert with_step - without == {"publisher_step"}
def test_the_first_four_fields_keep_their_positions():
"""The stamp is appended, so existing positional reads still land on the same
values. Unpacking now has to name five, which is the intended cost of one
payload type over a second unwrap function."""
payload = unwrap_rendezvous_blob(_blob(step=3))
assert payload[:4] == (b"meta", "agent-r0", "h:1", [])
def test_a_new_reader_handles_an_unstamped_publishers_blob():
"""Forward compatibility: no `publisher_step` key at all."""
payload = json.loads(_blob().decode("utf-8"))
payload.pop("publisher_step", None)
unwrapped = unwrap_rendezvous_blob(json.dumps(payload).encode("utf-8"))
assert unwrapped.publisher_step is None
def test_stamping_leaves_every_other_field_untouched():
"""What actually makes a stamped blob safe for a reader that predates the field.
Such a reader looks up the keys it knows and ignores the rest, so the property it
depends on is that the stamp is purely additive: every other key holds the same
value whether or not the stamp is present. Asserting that is honest about what
this suite can check. Running a released reader against the blob would be the
stronger test and is not something an in-process unit test can do, so it belongs
in a version-compatibility run rather than here.
"""
stamped = json.loads(_blob(step=11).decode("utf-8"))
unstamped = json.loads(_blob().decode("utf-8"))
assert {k: v for k, v in stamped.items() if k != "publisher_step"} == unstamped
def test_the_stamp_does_not_disturb_the_shard_table():
"""The stamp describes the table; it must not alter how the table decodes."""
from modelexpress.refit.reshard.rendezvous import PublishedShard, PublishedTensor
tensors = [
PublishedTensor(
name="layers.0.weight",
dtype="torch.bfloat16",
elsize=2,
full_shape=(8, 4),
shards=[
PublishedShard(
agent_name="agent-r0",
device_id=0,
addr=4096,
shard_offset=(0, 0),
shape=(4, 4),
)
],
)
]
stamped = unwrap_rendezvous_blob(_blob(step=5, tensors=tensors))
plain = unwrap_rendezvous_blob(_blob(tensors=tensors))
assert stamped[:4] == plain[:4]
assert stamped.tensors[0].shards[0].addr == 4096