modelexpress-client 0.7.0

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

"""Tests for MxK8sServiceClient and the client factory."""

from __future__ import annotations

import threading
from concurrent import futures

import grpc
import pytest

from modelexpress import p2p_pb2
from modelexpress import p2p_pb2_grpc
from modelexpress.client import MxClient
from modelexpress.metadata.client_factory import create_metadata_client
from modelexpress.metadata.k8s_service_client import MxK8sServiceClient
from modelexpress.metadata.payload import worker_tensor_descriptors
from modelexpress.metadata.publish import _is_p2p_metadata_enabled
from modelexpress.metadata.source_id import compute_mx_source_id


def _base_identity() -> p2p_pb2.SourceIdentity:
    return p2p_pb2.SourceIdentity(
        mx_version="0.7.0",
        mx_source_type=p2p_pb2.MX_SOURCE_TYPE_WEIGHTS,
        model_name="deepseek-ai/DeepSeek-V3",
        backend_framework=p2p_pb2.BACKEND_FRAMEWORK_VLLM,
        tensor_parallel_size=8,
        pipeline_parallel_size=1,
        expert_parallel_size=0,
        dtype="bfloat16",
        quantization="",
        revision="abc123",
    )


# ---------------------------------------------------------------------------
# Factory selection
# ---------------------------------------------------------------------------


def test_factory_default_returns_mx_client(monkeypatch):
    monkeypatch.delenv("MX_METADATA_BACKEND", raising=False)
    assert isinstance(create_metadata_client(), MxClient)


@pytest.mark.parametrize("value", ["", "server", "redis", "kubernetes", "k8s", "crd"])
def test_factory_central_aliases_return_mx_client(monkeypatch, value):
    monkeypatch.setenv("MX_METADATA_BACKEND", value)
    assert isinstance(create_metadata_client(), MxClient)


@pytest.mark.parametrize("value", ["k8s-service", "service", "K8S-SERVICE"])
def test_factory_k8s_service_aliases_return_k8s_client(monkeypatch, value):
    monkeypatch.setenv("MX_METADATA_BACKEND", value)
    client = create_metadata_client(worker_rank=3)
    assert isinstance(client, MxK8sServiceClient)
    assert client._worker_rank == 3


def test_factory_unknown_backend_raises(monkeypatch):
    monkeypatch.setenv("MX_METADATA_BACKEND", "bogus")
    with pytest.raises(ValueError, match="Unknown MX_METADATA_BACKEND"):
        create_metadata_client()


def test_k8s_worker_rpc_retry_policy_uses_configured_values():
    client = MxK8sServiceClient(max_retries=3, backoff_seconds=0.25)

    assert client.worker_rpc_retry_policy() == (3, 0.25)


# ---------------------------------------------------------------------------
# publish_metadata / list_sources / update_status behavior
# ---------------------------------------------------------------------------


def test_publish_metadata_returns_local_source_id_and_records_rank():
    client = MxK8sServiceClient()
    identity = _base_identity()
    worker = p2p_pb2.WorkerMetadata(worker_rank=5)
    sid = client.publish_metadata(identity, worker, "worker-xyz")
    assert sid == compute_mx_source_id(identity)
    assert client._worker_rank == 5


def test_class_declares_requires_p2p_metadata():
    # metadata._is_p2p_metadata_enabled() consults this class attribute
    # to force the P2P path regardless of MX_P2P_METADATA env var.
    assert MxK8sServiceClient.REQUIRES_P2P_METADATA is True


# ---------------------------------------------------------------------------
# _is_p2p_metadata_enabled backend-driven forcing
# ---------------------------------------------------------------------------


def test_is_p2p_metadata_enabled_forced_by_k8s_service_client(monkeypatch):
    # Backend declares REQUIRES_P2P_METADATA=True, so env var is ignored.
    monkeypatch.delenv("MX_P2P_METADATA", raising=False)
    assert _is_p2p_metadata_enabled(MxK8sServiceClient()) is True

    monkeypatch.setenv("MX_P2P_METADATA", "0")
    assert _is_p2p_metadata_enabled(MxK8sServiceClient()) is True


def test_is_p2p_metadata_enabled_mx_client_defaults_on_and_honors_env_var(
    monkeypatch,
):
    monkeypatch.delenv("MX_P2P_METADATA", raising=False)
    assert _is_p2p_metadata_enabled(MxClient()) is True

    monkeypatch.setenv("MX_P2P_METADATA", "0")
    assert _is_p2p_metadata_enabled(MxClient()) is False

    monkeypatch.setenv("MX_P2P_METADATA", "1")
    assert _is_p2p_metadata_enabled(MxClient()) is True


def test_is_p2p_metadata_enabled_warns_on_conflicting_env(monkeypatch, caplog):
    # If the user explicitly sets MX_P2P_METADATA=0 alongside a forcing
    # backend, their setting is ignored but a warning fires.
    import logging
    monkeypatch.setenv("MX_P2P_METADATA", "0")
    with caplog.at_level(logging.WARNING, logger="modelexpress.metadata"):
        result = _is_p2p_metadata_enabled(MxK8sServiceClient())
    assert result is True
    assert any("is ignored for backend" in rec.message for rec in caplog.records)


def test_list_sources_returns_single_synthetic_ref_at_caller_rank():
    client = MxK8sServiceClient(worker_rank=2)
    identity = _base_identity()
    resp = client.list_sources(identity=identity)
    assert len(resp.instances) == 1
    inst = resp.instances[0]
    assert inst.worker_rank == 2
    assert inst.mx_source_id == compute_mx_source_id(identity)
    assert inst.model_name == identity.model_name


def test_list_sources_synthetic_ref_survives_quantized_prefilter():
    # The synthetic ref carries no accelerator (this backend learns it only from
    # GetTensorManifest). The pre-fetch accelerator gate must NOT drop it for a
    # quantized identity, or a same-family quantized k8s-service transfer would
    # be stranded before its accelerator is known. Regression for the P2P
    # quantized cross-vendor guard vs the k8s-service empty-accelerator ref.
    from modelexpress.metadata.payload import accelerators_compatible

    client = MxK8sServiceClient(worker_rank=0)
    identity = _base_identity()
    identity.quantization = "fp8"
    resp = client.list_sources(identity=identity)
    inst = resp.instances[0]
    assert inst.accelerator == ""
    # Pre-fetch gate (default defer_unknown=True) defers the unknown ref.
    assert accelerators_compatible(
        "cuda",
        inst.accelerator,
        mx_source_type=identity.mx_source_type,
        quantization=identity.quantization,
        dtype=identity.dtype,
    )


def test_list_sources_requires_identity():
    client = MxK8sServiceClient(worker_rank=0)
    with pytest.raises(ValueError, match="identity"):
        client.list_sources()


def test_list_sources_requires_known_worker_rank():
    client = MxK8sServiceClient()  # no rank, no publish_metadata yet
    with pytest.raises(RuntimeError, match="worker_rank"):
        client.list_sources(identity=_base_identity())


def test_update_status_is_noop_returning_true():
    client = MxK8sServiceClient(worker_rank=0)
    assert client.update_status("sid", "wid", 0, p2p_pb2.SOURCE_STATUS_READY) is True


def test_resolve_endpoint_substitutes_rank():
    client = MxK8sServiceClient(
        worker_rank=7,
        service_pattern="my-svc-rank-{rank}.ns.svc.cluster.local:6555",
    )
    assert client._resolve_endpoint() == "my-svc-rank-7.ns.svc.cluster.local:6555"


def test_resolve_endpoint_autoappends_port_when_pattern_has_no_port(monkeypatch):
    # Shape 2: bare hostname, client auto-computes port = base + rank.
    monkeypatch.delenv("MX_WORKER_GRPC_PORT", raising=False)
    client = MxK8sServiceClient(
        worker_rank=3,
        service_pattern="mx-sources",
    )
    assert client._resolve_endpoint() == "mx-sources:6558"  # 6555 + 3


def test_resolve_endpoint_autoappend_honors_mx_worker_grpc_port(monkeypatch):
    monkeypatch.setenv("MX_WORKER_GRPC_PORT", "9000")
    client = MxK8sServiceClient(
        worker_rank=2,
        service_pattern="mx-sources",
    )
    assert client._resolve_endpoint() == "mx-sources:9002"  # 9000 + 2


def test_resolve_endpoint_autoappend_works_with_rank_substitution(monkeypatch):
    # Pattern has {rank} in the hostname but no port; client still
    # auto-appends :base+rank.
    monkeypatch.delenv("MX_WORKER_GRPC_PORT", raising=False)
    client = MxK8sServiceClient(
        worker_rank=1,
        service_pattern="mx-sources-rank-{rank}",
    )
    assert client._resolve_endpoint() == "mx-sources-rank-1:6556"  # 6555 + 1


def test_default_service_pattern_is_bare_hostname():
    # Default is the Shape-2-friendly bare hostname. Clients that want
    # Shape 1 (rank-in-hostname) set MX_K8S_SERVICE_PATTERN explicitly.
    import os
    saved = os.environ.pop("MX_K8S_SERVICE_PATTERN", None)
    try:
        client = MxK8sServiceClient(worker_rank=0)
        assert client._service_pattern == "mx-sources"
    finally:
        if saved is not None:
            os.environ["MX_K8S_SERVICE_PATTERN"] = saved


def test_close_is_safe_noop():
    client = MxK8sServiceClient(worker_rank=0)
    client.close()
    client.close()  # idempotent


# ---------------------------------------------------------------------------
# get_metadata against a real in-process gRPC WorkerService
# ---------------------------------------------------------------------------


class _FakeWorkerServicer(p2p_pb2_grpc.WorkerServiceServicer):
    """In-process gRPC servicer that can be tuned to either succeed or
    fail with FAILED_PRECONDITION, to exercise the retry loop."""

    def __init__(
        self,
        mx_source_id: str,
        worker_rank: int,
        *,
        accelerator: str = "cuda",
        fail_first_n: int = 0,
    ):
        self._mx_source_id = mx_source_id
        self._worker_rank = worker_rank
        self._accelerator = accelerator
        self._fail_first_n = fail_first_n
        self._calls = 0
        self._lock = threading.Lock()

    @property
    def call_count(self) -> int:
        with self._lock:
            return self._calls

    def GetTensorManifest(self, request, context):
        with self._lock:
            self._calls += 1
            should_fail = self._calls <= self._fail_first_n
        if should_fail:
            context.abort(
                grpc.StatusCode.FAILED_PRECONDITION,
                "simulated stale backend",
            )
        return p2p_pb2.GetTensorManifestResponse(
            mx_source_id=self._mx_source_id,
            tensors=[
                p2p_pb2.TensorDescriptor(name="t0", size=16, device_id=0)
            ],
            metadata_endpoint="10.0.0.1:5555",
            agent_name="fake-agent",
            worker_rank=self._worker_rank,
            accelerator=self._accelerator,
        )


def _start_fake_server(servicer) -> tuple[grpc.Server, int]:
    server = grpc.server(futures.ThreadPoolExecutor(max_workers=2))
    p2p_pb2_grpc.add_WorkerServiceServicer_to_server(servicer, server)
    port = server.add_insecure_port("127.0.0.1:0")
    server.start()
    return server, port


def test_get_metadata_success_builds_synthetic_response():
    identity = _base_identity()
    sid = compute_mx_source_id(identity)
    servicer = _FakeWorkerServicer(sid, worker_rank=4)
    server, port = _start_fake_server(servicer)
    try:
        client = MxK8sServiceClient(
            worker_rank=4,
            service_pattern=f"127.0.0.1:{port}",
        )
        resp = client.get_metadata(sid, "worker-xyz")
        assert resp.found is True
        assert resp.mx_source_id == sid
        assert resp.worker.worker_rank == 4
        assert resp.worker.metadata_endpoint == "10.0.0.1:5555"
        assert resp.worker.agent_name == "fake-agent"
        assert resp.worker.accelerator == "cuda"
        assert resp.worker.status == p2p_pb2.SOURCE_STATUS_READY
        tensors = worker_tensor_descriptors(resp.worker)
        assert len(tensors) == 1
        assert tensors[0].name == "t0"
        # worker_grpc_endpoint must be populated: rdma_strategy gates
        # the P2P branch (which pulls NIXL metadata from the source's
        # listen thread) on bool(worker.worker_grpc_endpoint). An empty
        # value silently disables NIXL metadata exchange and the target
        # hits "missing nixlSerDes tag" during deserialization.
        assert resp.worker.worker_grpc_endpoint == f"127.0.0.1:{port}"
        assert servicer.call_count == 1
    finally:
        server.stop(grace=None)


def test_get_metadata_populates_worker_grpc_endpoint_for_p2p_gating():
    """Regression test: worker_grpc_endpoint must equal the Service endpoint.

    rdma_strategy._receive_from_peer evaluates
    is_p2p = bool(source_worker.worker_grpc_endpoint) and only takes
    the P2P branch (which calls fetch_remote_and_wait to pull the
    source's NIXL metadata) when truthy. If this client synthesizes a
    WorkerMetadata with worker_grpc_endpoint unset, the target skips
    NIXL metadata exchange entirely and the transfer fails with
    "missing nixlSerDes tag" during deserialization.
    """
    identity = _base_identity()
    sid = compute_mx_source_id(identity)
    servicer = _FakeWorkerServicer(sid, worker_rank=0)
    server, port = _start_fake_server(servicer)
    try:
        pattern = f"mx-sources-rank-{{rank}}:{port}"
        # Override hostname resolution by pointing the pattern at localhost.
        client = MxK8sServiceClient(
            worker_rank=0,
            service_pattern=f"127.0.0.1:{port}",
        )
        resp = client.get_metadata(sid, "worker-xyz")
        assert bool(resp.worker.worker_grpc_endpoint), (
            "worker_grpc_endpoint must be non-empty so rdma_strategy "
            "takes the P2P branch"
        )
        assert resp.worker.worker_grpc_endpoint == f"127.0.0.1:{port}"
        _ = pattern  # silence unused; pattern shape documented above
    finally:
        server.stop(grace=None)


def test_get_metadata_retries_on_failed_precondition():
    identity = _base_identity()
    sid = compute_mx_source_id(identity)
    # Fail twice, succeed on the third attempt.
    servicer = _FakeWorkerServicer(sid, worker_rank=0, fail_first_n=2)
    server, port = _start_fake_server(servicer)
    try:
        client = MxK8sServiceClient(
            worker_rank=0,
            service_pattern=f"127.0.0.1:{port}",
            max_retries=5,
            backoff_seconds=0.0,
        )
        resp = client.get_metadata(sid, "worker-xyz")
        assert resp.found is True
        assert servicer.call_count == 3
    finally:
        server.stop(grace=None)


def test_get_metadata_gives_up_after_max_retries():
    identity = _base_identity()
    sid = compute_mx_source_id(identity)
    # Always fail.
    servicer = _FakeWorkerServicer(sid, worker_rank=0, fail_first_n=100)
    server, port = _start_fake_server(servicer)
    try:
        client = MxK8sServiceClient(
            worker_rank=0,
            service_pattern=f"127.0.0.1:{port}",
            max_retries=2,
            backoff_seconds=0.0,
        )
        with pytest.raises(grpc.RpcError):
            client.get_metadata(sid, "worker-xyz")
        # max_retries=2 means 3 attempts total (initial + 2 retries).
        assert servicer.call_count == 3
    finally:
        server.stop(grace=None)


def test_get_metadata_without_worker_rank_raises():
    client = MxK8sServiceClient()
    with pytest.raises(RuntimeError, match="worker_rank"):
        client.get_metadata("sid", "wid")


def test_get_metadata_rejects_source_id_mismatch():
    # Server advertises a different mx_source_id than the client asks
    # for. Client should retry up to max_retries times, then raise.
    identity = _base_identity()
    sid = compute_mx_source_id(identity)
    servicer = _FakeWorkerServicer("wrong-source-id", worker_rank=0)
    server, port = _start_fake_server(servicer)
    try:
        client = MxK8sServiceClient(
            worker_rank=0,
            service_pattern=f"127.0.0.1:{port}",
            max_retries=2,
            backoff_seconds=0.0,
        )
        with pytest.raises(RuntimeError, match="mx_source_id mismatch"):
            client.get_metadata(sid, "worker-xyz")
        # max_retries=2 -> 3 attempts total.
        assert servicer.call_count == 3
    finally:
        server.stop(grace=None)


def test_get_metadata_rejects_worker_rank_mismatch():
    # Server returns a different worker_rank than the caller's rank.
    # Typical "misconfigured Service selector" scenario.
    identity = _base_identity()
    sid = compute_mx_source_id(identity)
    servicer = _FakeWorkerServicer(sid, worker_rank=5)
    server, port = _start_fake_server(servicer)
    try:
        client = MxK8sServiceClient(
            worker_rank=3,
            service_pattern=f"127.0.0.1:{port}",
            max_retries=1,
            backoff_seconds=0.0,
        )
        with pytest.raises(RuntimeError, match="worker_rank mismatch"):
            client.get_metadata(sid, "worker-xyz")
        # max_retries=1 -> 2 attempts total.
        assert servicer.call_count == 2
    finally:
        server.stop(grace=None)