modelexpress-client 0.6.0

Client library for Model Express gRPC server
Documentation 
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
125

# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0

import pytest
import torch

from modelexpress.refit.reshard.rendezvous import unwrap_rendezvous_blob
from modelexpress_rl.train.engines.fsdp.publisher import (
    LocalTensorShard,
    build_fsdp_reshard_manifest,
    capture_local_shards,
)


class _Manager:
    agent_name = "trainer-r0"
    nixl_metadata = b"agent-metadata"
    listen_port = 19000


def test_capture_publishes_full_copies_and_skips_non_float():
    state_dict = {
        "w": torch.arange(8, dtype=torch.float32).reshape(2, 4),
        "b": torch.ones(4, dtype=torch.bfloat16),
        "step": torch.arange(4, dtype=torch.long),  # non-float: skipped
    }

    shards = capture_local_shards(state_dict)

    by_name = {s.name: s for s in shards}
    assert set(by_name) == {"w", "b"}
    for shard in shards:
        # Plain (non-DTensor) tensors are held in full by every rank.
        assert shard.shard_offset == (0,) * len(shard.global_shape)
        assert shard.local_shape == shard.global_shape
        assert shard.staging_tensor is None
    assert by_name["w"].global_shape == (2, 4)


def test_capture_gives_a_scalar_tensor_a_rank_zero_offset():
    state_dict = {"loss_scale": torch.tensor(1.0, dtype=torch.float32)}

    (shard,) = capture_local_shards(state_dict)

    assert shard.global_shape == ()
    assert shard.shard_offset == ()
    assert shard.local_shape == ()


def test_build_manifest_describes_the_served_buffer():
    tensor = torch.zeros(2, 4, dtype=torch.bfloat16)
    shard = LocalTensorShard(
        name="w",
        global_shape=(2, 4),
        shard_offset=(0, 0),
        local_shape=(2, 4),
        source_tensor=tensor,
    )

    blob = build_fsdp_reshard_manifest(
        manager=_Manager(), shards=[shard], metadata_endpoint="host:1234"
    )

    payload = unwrap_rendezvous_blob(blob)
    assert payload.agent_metadata == b"agent-metadata"
    (published,) = payload.tensors
    assert published.name == "w"
    assert published.dtype == "torch.bfloat16"
    assert published.elsize == 2
    assert tuple(published.full_shape) == (2, 4)
    (pshard,) = published.shards
    assert pshard.agent_name == "trainer-r0"
    assert pshard.addr == tensor.data_ptr()
    assert tuple(pshard.shard_offset) == (0, 0)
    assert tuple(pshard.shape) == (2, 4)


def test_build_manifest_groups_multiple_shards_under_one_tensor():
    top = torch.zeros(2, 4, dtype=torch.bfloat16)
    bottom = torch.zeros(2, 4, dtype=torch.bfloat16)
    shards = [
        LocalTensorShard("w", (4, 4), (0, 0), (2, 4), top),
        LocalTensorShard("w", (4, 4), (2, 0), (2, 4), bottom),
    ]

    payload = unwrap_rendezvous_blob(
        build_fsdp_reshard_manifest(
            manager=_Manager(), shards=shards, metadata_endpoint="host:1234"
        )
    )

    (published,) = payload.tensors
    offsets = sorted(tuple(s.shard_offset) for s in published.shards)
    assert offsets == [(0, 0), (2, 0)]


def test_build_manifest_rejects_non_contiguous_served_tensor():
    strided = torch.zeros(4, 4, dtype=torch.bfloat16)[:, ::2]
    shard = LocalTensorShard("w", (4, 2), (0, 0), (4, 2), strided)

    with pytest.raises(ValueError, match="contiguous"):
        build_fsdp_reshard_manifest(
            manager=_Manager(), shards=[shard], metadata_endpoint="host:1234"
        )


def test_build_manifest_rejects_an_invalid_address():
    meta = torch.zeros(2, 4, dtype=torch.bfloat16, device="meta")
    shard = LocalTensorShard("w", (2, 4), (0, 0), (2, 4), meta)

    with pytest.raises(ValueError, match="invalid address"):
        build_fsdp_reshard_manifest(
            manager=_Manager(), shards=[shard], metadata_endpoint="host:1234"
        )


def test_build_manifest_requires_host_port_endpoint():
    shard = LocalTensorShard(
        "w", (2, 4), (0, 0), (2, 4), torch.zeros(2, 4, dtype=torch.bfloat16)
    )

    with pytest.raises(ValueError, match="host:port"):
        build_fsdp_reshard_manifest(
            manager=_Manager(), shards=[shard], metadata_endpoint="no-port"
        )