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

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

"""Content-addressed source ID computation.

Mirrors the Rust ``compute_mx_source_id`` in
``modelexpress_server/src/source_identity.rs``. When the central server
is in the loop, it computes the ID and returns it to clients; in
decentralized / K8s-Service-routed deployments where there is no central
server, clients compute the ID locally. The two implementations MUST
produce identical IDs for identical identities, or the handshake
validation that ties ``mx_source_id`` to the content of
``SourceIdentity`` breaks.
"""

from __future__ import annotations

import hashlib
import json

from .. import p2p_pb2

_OPTIONAL_SOURCE_ID_FIELDS = (
    "backend_framework_version",
    "torch_version",
    "cuda_version",
    "triton_version",
    "gpu_arch",
    "compile_config_digest",
)


def compute_mx_source_id(identity: p2p_pb2.SourceIdentity) -> str:
    """Compute the 16-char hex ``mx_source_id`` for a ``SourceIdentity``.

    All string fields are lowercased and ``extra_parameters`` is sorted
    by key so semantically-identical identities produce identical IDs
    regardless of case or insertion order.
    """
    canonical = _canonical_json(identity)
    digest = hashlib.sha256(canonical.encode()).hexdigest()
    return digest[:16]


def _canonical_json(identity: p2p_pb2.SourceIdentity) -> str:
    # Normalize extra_parameters deterministically:
    # 1. sort by original key (ASCII/byte order, matches Rust's String::cmp)
    # 2. lowercase both keys and values
    # 3. on case-colliding keys, keep the first value we see (stable because
    #    the sort already picks a canonical survivor)
    # Rust's HashMap iteration order is non-deterministic, so this explicit
    # sort-then-dedup is what keeps cross-language hashes stable when users
    # pass case-colliding keys.
    sorted_extra: dict[str, str] = {}
    for k, v in sorted(identity.extra_parameters.items()):
        lk = k.lower()
        if lk not in sorted_extra:
            sorted_extra[lk] = v.lower()
    payload = {
        "mx_version": identity.mx_version.lower(),
        "mx_source_type": identity.mx_source_type,
        "model_name": identity.model_name.lower(),
        "backend_framework": identity.backend_framework,
        "tensor_parallel_size": identity.tensor_parallel_size,
        "pipeline_parallel_size": identity.pipeline_parallel_size,
        "expert_parallel_size": identity.expert_parallel_size,
        "dtype": identity.dtype.lower(),
        "quantization": identity.quantization.lower(),
        "extra_parameters": sorted_extra,
        "revision": identity.revision.lower(),
    }
    # Include optional compatibility fields only when set. That preserves
    # existing weight source IDs while letting artifact/cache sources split IDs
    # by framework, compiler, CUDA, and GPU compatibility constraints.
    for field in _OPTIONAL_SOURCE_ID_FIELDS:
        value = getattr(identity, field)
        if value:
            payload[field] = value.lower()
    # Rust's serde_json::json! stores keys in a BTreeMap, which sorts
    # them alphabetically on serialization. Match that ordering with
    # sort_keys=True so both sides produce the exact same bytes for
    # SHA256 to hash.
    return json.dumps(payload, separators=(",", ":"), sort_keys=True)