modelexpress-client 0.5.1

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
127
128
129
130
131
132
133
134
135
136

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

"""
Pre-mock vLLM modules so tests can run without vLLM installed.

This conftest runs before any test module is collected, injecting mock
modules into sys.modules for all vllm.* imports used by the modelexpress
package. BaseModelLoader is a real ABC so that __abstractmethods__
validation works correctly.
"""

import sys
from abc import ABC, abstractmethod
from dataclasses import dataclass, field
from unittest.mock import MagicMock

import pytest
import torch
import torch.nn as nn


@dataclass
class MockAcceleratorBackend:
    name: str = "mock"
    torch_device_type: str = "mock"
    nixl_mem_type: str = "VRAM"
    rdma_p2p: bool = False
    pool_reg: bool = False
    vmm: bool = False
    gds: bool = False
    set_device_calls: list[int] = field(default_factory=list)
    synchronize_calls: list[int | None] = field(default_factory=list)
    empty_cache_calls: int = 0

    def set_device(self, device_id: int) -> None:
        self.set_device_calls.append(device_id)

    def current_device(self) -> int:
        return 0

    def synchronize(self, device_id: int | None = None) -> None:
        self.synchronize_calls.append(device_id)

    def empty_cache(self) -> None:
        self.empty_cache_calls += 1

    def torch_device(self, device_id: int) -> torch.device:
        return torch.device("cpu")

    def is_accel_tensor(self, tensor: torch.Tensor) -> bool:
        return tensor.device.type == self.torch_device_type

    def supports_rdma_p2p(self) -> bool:
        return self.rdma_p2p

    def supports_pool_reg(self) -> bool:
        return self.pool_reg

    def supports_vmm(self) -> bool:
        return self.vmm

    def supports_gds(self) -> bool:
        return self.gds


@pytest.fixture
def mock_accelerator_backend_cls():
    return MockAcceleratorBackend


def _maybe_mock_vllm():
    """Inject vLLM mocks into sys.modules if vLLM is not installed."""
    try:
        import vllm  # noqa: F401
        return  # vLLM is available, no mocking needed
    except ImportError:
        pass

    # Real ABC so __abstractmethods__ propagates to subclasses
    class BaseModelLoader(ABC):
        def __init__(self, load_config):
            self.load_config = load_config

        @abstractmethod
        def download_model(self, model_config) -> None:
            raise NotImplementedError

        @abstractmethod
        def load_weights(self, model: nn.Module, model_config) -> None:
            raise NotImplementedError

        def load_model(self, vllm_config, model_config):
            raise NotImplementedError

    # No-op decorator that just returns the class
    def register_model_loader(load_format):
        def _wrapper(cls):
            return cls
        return _wrapper

    # Build mock module tree
    vllm_mods = {
        "vllm": MagicMock(),
        "vllm.config": MagicMock(),
        "vllm.config.load": MagicMock(),
        "vllm.model_executor": MagicMock(),
        "vllm.model_executor.model_loader": MagicMock(),
        "vllm.model_executor.model_loader.base_loader": MagicMock(),
        "vllm.model_executor.model_loader.default_loader": MagicMock(),
        "vllm.model_executor.model_loader.dummy_loader": MagicMock(),
        "vllm.model_executor.model_loader.utils": MagicMock(),
        "vllm.utils": MagicMock(),
        "vllm.utils.torch_utils": MagicMock(),
        "vllm.distributed": MagicMock(),
    }

    # Wire up real objects where behavior matters
    vllm_mods["vllm.model_executor.model_loader.base_loader"].BaseModelLoader = BaseModelLoader
    vllm_mods["vllm.model_executor.model_loader"].register_model_loader = register_model_loader
    vllm_mods["vllm.model_executor.model_loader"].BaseModelLoader = BaseModelLoader

    # set_default_torch_dtype needs to be a real context manager
    from contextlib import contextmanager

    @contextmanager
    def _set_default_torch_dtype(dtype):
        yield

    vllm_mods["vllm.utils.torch_utils"].set_default_torch_dtype = _set_default_torch_dtype

    for mod_name, mod in vllm_mods.items():
        sys.modules[mod_name] = mod


_maybe_mock_vllm()