modelexpress-client 0.3.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
126

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

"""Tests for ModelExpress type definitions."""

import re

from google.protobuf import __version__ as _pb_version

from modelexpress.types import TensorDescriptor, WorkerMetadata, GetMetadataResponse


class TestProtobufCompatibility:
    """Guard against generated protobuf code drifting from the installed runtime."""

    def test_p2p_pb2_gencode_matches_runtime_major_version(self):
        """Regenerate p2p_pb2.py if this fails (see pyproject.toml [dev] deps)."""
        import modelexpress.p2p_pb2 as pb2

        with open(pb2.__file__) as f:
            src = f.read()
        m = re.search(r"Protobuf Python Version: (\d+)\.", src)
        assert m, "Could not parse gencode version from p2p_pb2.py"
        gencode_major = int(m.group(1))
        runtime_major = int(_pb_version.split(".")[0])
        assert gencode_major == runtime_major, (
            f"p2p_pb2.py was generated with protobuf {gencode_major}.x "
            f"but runtime is {runtime_major}.x - regenerate with: "
            f"python -m grpc_tools.protoc -I../../modelexpress_common/proto "
            f"--python_out=modelexpress --grpc_python_out=modelexpress "
            f"../../modelexpress_common/proto/p2p.proto"
        )


class TestTensorDescriptor:
    """Tests for TensorDescriptor dataclass."""

    def test_creation(self):
        """Test basic tensor descriptor creation."""
        desc = TensorDescriptor(
            name="model.layers.0.self_attn.q_proj.weight",
            addr=0x7F8A00000000,
            size=1024 * 1024 * 1024,
            device_id=0,
            dtype="bfloat16",
        )
        assert desc.name == "model.layers.0.self_attn.q_proj.weight"
        assert desc.size == 1024 * 1024 * 1024
        assert desc.dtype == "bfloat16"

    def test_dtype_required(self):
        """Test that dtype is a required field."""
        import pytest
        with pytest.raises(TypeError):
            TensorDescriptor(
                name="test",
                addr=0,
                size=0,
                device_id=0,
            )

    def test_large_tensor(self):
        """Test with realistic large tensor values."""
        desc = TensorDescriptor(
            name="model.embed_tokens.weight",
            addr=0x7F8A00000000,
            size=32000 * 8192 * 2,
            device_id=7,
            dtype="bfloat16",
        )
        assert desc.size == 524288000


class TestWorkerMetadata:
    """Tests for WorkerMetadata dataclass."""

    def test_creation(self):
        """Test basic worker metadata creation."""
        tensors = [
            TensorDescriptor(
                name=f"layer.{i}.weight",
                addr=0x7F8A00000000 + i * 1024,
                size=1024,
                device_id=0,
                dtype="bfloat16",
            )
            for i in range(3)
        ]
        metadata = WorkerMetadata(
            worker_rank=0,
            nixl_metadata=b"test_metadata",
            tensors=tensors,
        )
        assert metadata.worker_rank == 0
        assert len(metadata.tensors) == 3
        assert metadata.nixl_metadata == b"test_metadata"


class TestGetMetadataResponse:
    """Tests for GetMetadataResponse dataclass."""

    def test_found_response(self):
        """Test response when source is found."""
        workers = [
            WorkerMetadata(
                worker_rank=i,
                nixl_metadata=b"metadata",
                tensors=[],
            )
            for i in range(4)
        ]
        response = GetMetadataResponse(
            found=True,
            workers=workers,
        )
        assert response.found is True
        assert len(response.workers) == 4

    def test_not_found_response(self):
        """Test response when source is not found."""
        response = GetMetadataResponse(
            found=False,
            workers=[],
        )
        assert response.found is False
        assert len(response.workers) == 0