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

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

"""GDS loading strategy: GPUDirect Storage for direct file-to-GPU loading."""

from __future__ import annotations

import logging

from ..adapter import EngineAdapter, StrategyFailed
from .base import LoadContext, LoadStrategy, _as_load_result, register_tensors
from .context import LoadResult

logger = logging.getLogger("modelexpress.strategy_gds")


class GdsStrategy(LoadStrategy):
    """Load weights via GPUDirect Storage (direct file-to-GPU)."""

    name = "gds"
    requires = (EngineAdapter.apply_weight_iter,)

    def is_available(self, ctx: LoadContext) -> bool:
        if not super().is_available(ctx):
            return False
        if not ctx.accelerator_backend.supports_gds():
            logger.info(
                f"[Worker {ctx.global_rank}] GDS not supported on "
                f"{ctx.accelerator_backend.name}, skipping"
            )
            return False
        from ..gds_transfer import is_gds_available
        available = is_gds_available()
        if not available:
            logger.info(f"[Worker {ctx.global_rank}] GDS not available, skipping")
        return available

    def load(self, result: LoadResult, ctx: LoadContext) -> LoadResult:
        result = _as_load_result(result)
        from ..gds_loader import MxGdsLoader

        logger.info(f"[Worker {ctx.global_rank}] Attempting GDS loading...")
        gds_loader = MxGdsLoader(accelerator_backend=ctx.accelerator_backend)
        try:
            try:
                use_tqdm = getattr(ctx.load_config, "use_tqdm_on_load", True)
                revision = getattr(ctx.model_config, "revision", None)
                weights_iter = gds_loader.load_iter(
                    _get_model_path(ctx.model_config),
                    use_tqdm=use_tqdm,
                    revision=revision,
                )
            except Exception as e:
                logger.warning(
                    f"[Worker {ctx.global_rank}] GDS loading failed, falling through: {e}"
                )
                raise StrategyFailed(str(e), mutated=False) from e

            try:
                result = ctx.adapter.apply_weight_iter(result, weights_iter)
                logger.info(f"[Worker {ctx.global_rank}] GDS weight loading complete")
                result = ctx.adapter.after_weight_iter_load(result)
            except Exception as e:
                logger.warning(
                    f"[Worker {ctx.global_rank}] GDS loading failed, falling through: {e}"
                )
                raise StrategyFailed(str(e), mutated=True) from e
        finally:
            gds_loader.shutdown()

        register_tensors(result, ctx)
        return result


def _get_model_path(model_config) -> str:
    """Return the checkpoint path used by vLLM or SGLang model configs."""
    for attribute in ("model", "model_path"):
        value = getattr(model_config, attribute, None)
        if value:
            return str(value)
    raise AttributeError("model config has neither 'model' nor 'model_path'")