modelexpress-client 0.7.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

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

from concurrent import futures

import grpc
import pytest
from google.protobuf import empty_pb2, wrappers_pb2

from modelexpress_rl.inference.engines.vllm.startup_probe import (
    reconcile_serving_version,
)


def _start_control_server(*, initial="default", accept_update=True):
    """Start a test server implementing vLLM's native Control methods."""
    state = {"version": initial, "updates": []}

    def update(request, _context):
        """Record and optionally apply a serving-version update."""
        state["updates"].append(request.value)
        if accept_update:
            state["version"] = request.value
        return empty_pb2.Empty()

    def get(_request, _context):
        """Return the server's current serving version."""
        return wrappers_pb2.StringValue(value=state["version"])

    server = grpc.server(futures.ThreadPoolExecutor(max_workers=2))
    server.add_generic_rpc_handlers(
        (
            grpc.method_handlers_generic_handler(
                "vllm.Control",
                {
                    "UpdateWeightVersion": grpc.unary_unary_rpc_method_handler(
                        update,
                        request_deserializer=wrappers_pb2.StringValue.FromString,
                        response_serializer=empty_pb2.Empty.SerializeToString,
                    ),
                    "GetWeightVersion": grpc.unary_unary_rpc_method_handler(
                        get,
                        request_deserializer=empty_pb2.Empty.FromString,
                        response_serializer=wrappers_pb2.StringValue.SerializeToString,
                    ),
                },
            ),
        )
    )
    port = server.add_insecure_port("127.0.0.1:0")
    server.start()
    return server, f"127.0.0.1:{port}", state


def test_reconcile_serving_version_updates_then_reads_vllm_control():
    """Update and read back the desired version through vLLM Control."""
    server, address, state = _start_control_server()
    try:
        observed = reconcile_serving_version(
            address=address,
            desired_version_uid="version-7",
        )
    finally:
        server.stop(grace=None).wait()

    assert observed == "version-7"
    assert state["updates"] == ["version-7"]


def test_reconcile_serving_version_rejects_observed_mismatch():
    """Reject a serving version that differs after reconciliation."""
    server, address, _state = _start_control_server(accept_update=False)
    try:
        with pytest.raises(RuntimeError, match="observed='default'"):
            reconcile_serving_version(
                address=address,
                desired_version_uid="version-7",
            )
    finally:
        server.stop(grace=None).wait()


def test_reconcile_without_desired_version_only_checks_control_readiness():
    """Check Control readiness without updating when no version is desired."""
    server, address, state = _start_control_server(initial="default")
    try:
        observed = reconcile_serving_version(
            address=address,
            desired_version_uid=None,
        )
    finally:
        server.stop(grace=None).wait()

    assert observed == "default"
    assert state["updates"] == []