kserve/test/e2e/predictor/test_huggingface_vllm_cpu.py

# Copyright 2024 The KServe Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#    http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import os

import pytest
from kubernetes import client
from kubernetes.client import V1ResourceRequirements

from kserve import (
    V1beta1PredictorSpec,
    V1beta1ModelSpec,
    V1beta1ModelFormat,
    V1beta1InferenceService,
    V1beta1InferenceServiceSpec,
    KServeClient,
)
from kserve.constants import constants
from ..common.utils import (
    KSERVE_TEST_NAMESPACE,
    generate,
    rerank,
    chat_completion_stream,
    completion_stream,
)

# Cold model loads + pinned-revision HuggingFace pulls can outrun the 600s default
# on contended CI runners. Increase the timeout for all tests in this file.
ISVC_READY_TIMEOUT_S = 900

# Knative's per-revision progress deadline (cluster default 600s) decides when a
# slow-to-start Revision is permanently marked ``RevisionFailed``. Once that
# happens, no amount of polling in ``wait_isvc_ready`` will recover. KServe
# propagates this annotation onto the Knative revision template, so raising it
# here gives CPU vLLM init (model load + bfloat16 + KV-cache build) room to
# finish.
ISVC_ANNOTATIONS = {"serving.knative.dev/progress-deadline": "20m"}


@pytest.mark.vllm
def test_huggingface_vllm_cpu_openai_chat_completions():
    service_name = "hf-qwen-chat-vllm"
    predictor = V1beta1PredictorSpec(
        min_replicas=1,
        model=V1beta1ModelSpec(
            model_format=V1beta1ModelFormat(
                name="huggingface",
            ),
            args=[
                "--model_id",
                "Qwen/Qwen2-0.5B-Instruct",
                "--model_name",
                "qwen-chat",
                "--backend",
                "vllm",
                "--max_model_len",
                "512",
                "--dtype",
                "bfloat16",
            ],
            env=[
                client.V1EnvVar(
                    name="VLLM_CPU_KVCACHE_SPACE",
                    value="1",
                ),
                client.V1EnvVar(
                    name="VLLM_ENABLE_V1_MULTIPROCESSING",
                    value="0",
                ),
            ],
            resources=V1ResourceRequirements(
                requests={"cpu": "2", "memory": "7Gi"},
                limits={"cpu": "2", "memory": "7Gi"},
            ),
        ),
    )

    isvc = V1beta1InferenceService(
        api_version=constants.KSERVE_V1BETA1,
        kind=constants.KSERVE_KIND_INFERENCESERVICE,
        metadata=client.V1ObjectMeta(
            name=service_name,
            namespace=KSERVE_TEST_NAMESPACE,
            annotations=ISVC_ANNOTATIONS,
        ),
        spec=V1beta1InferenceServiceSpec(predictor=predictor),
    )

    kserve_client = KServeClient(
        config_file=os.environ.get("KUBECONFIG", "~/.kube/config")
    )
    kserve_client.create(isvc)
    kserve_client.wait_isvc_ready(
        service_name,
        namespace=KSERVE_TEST_NAMESPACE,
        timeout_seconds=ISVC_READY_TIMEOUT_S,
    )

    res = generate(service_name, "./data/qwen_input_chat.json")
    assert res["choices"][0]["message"]["content"] == "The result of 2 + 2 is 4."

    kserve_client.delete(service_name, KSERVE_TEST_NAMESPACE)


@pytest.mark.vllm
def test_huggingface_vllm_cpu_text_completion_streaming():
    service_name = "hf-qwen-cmpl-stream-vllm"
    predictor = V1beta1PredictorSpec(
        min_replicas=1,
        model=V1beta1ModelSpec(
            model_format=V1beta1ModelFormat(
                name="huggingface",
            ),
            args=[
                "--model_id",
                "Qwen/Qwen2-0.5B",
                "--model_name",
                "qwen-cmpl-stream",
                "--backend",
                "vllm",
                "--max_model_len",
                "512",
                "--dtype",
                "bfloat16",
            ],
            env=[
                client.V1EnvVar(
                    name="VLLM_CPU_KVCACHE_SPACE",
                    value="1",
                ),
                client.V1EnvVar(
                    name="VLLM_ENABLE_V1_MULTIPROCESSING",
                    value="0",
                ),
            ],
            resources=V1ResourceRequirements(
                requests={"cpu": "2", "memory": "7Gi"},
                limits={"cpu": "2", "memory": "7Gi"},
            ),
        ),
    )

    isvc = V1beta1InferenceService(
        api_version=constants.KSERVE_V1BETA1,
        kind=constants.KSERVE_KIND_INFERENCESERVICE,
        metadata=client.V1ObjectMeta(
            name=service_name,
            namespace=KSERVE_TEST_NAMESPACE,
            annotations=ISVC_ANNOTATIONS,
        ),
        spec=V1beta1InferenceServiceSpec(predictor=predictor),
    )

    kserve_client = KServeClient(
        config_file=os.environ.get("KUBECONFIG", "~/.kube/config")
    )
    kserve_client.create(isvc)
    kserve_client.wait_isvc_ready(
        service_name,
        namespace=KSERVE_TEST_NAMESPACE,
        timeout_seconds=ISVC_READY_TIMEOUT_S,
    )

    full_response, _ = completion_stream(
        service_name, "./data/qwen_input_cmpl_stream.json"
    )
    assert full_response.strip() == "The result of 2 + 2 is 4."

    kserve_client.delete(service_name, KSERVE_TEST_NAMESPACE)


@pytest.mark.vllm
def test_huggingface_vllm_cpu_openai_completions():
    service_name = "hf-qwen-cmpl-vllm"
    predictor = V1beta1PredictorSpec(
        min_replicas=1,
        model=V1beta1ModelSpec(
            model_format=V1beta1ModelFormat(
                name="huggingface",
            ),
            args=[
                "--model_id",
                "Qwen/Qwen2-0.5B",
                "--model_name",
                "qwen-cmpl",
                "--backend",
                "vllm",
                "--max_model_len",
                "512",
                "--dtype",
                "bfloat16",
            ],
            env=[
                client.V1EnvVar(
                    name="VLLM_CPU_KVCACHE_SPACE",
                    value="1",
                ),
                client.V1EnvVar(
                    name="VLLM_ENABLE_V1_MULTIPROCESSING",
                    value="0",
                ),
            ],
            resources=V1ResourceRequirements(
                requests={"cpu": "2", "memory": "7Gi"},
                limits={"cpu": "2", "memory": "7Gi"},
            ),
        ),
    )

    isvc = V1beta1InferenceService(
        api_version=constants.KSERVE_V1BETA1,
        kind=constants.KSERVE_KIND_INFERENCESERVICE,
        metadata=client.V1ObjectMeta(
            name=service_name,
            namespace=KSERVE_TEST_NAMESPACE,
            annotations=ISVC_ANNOTATIONS,
        ),
        spec=V1beta1InferenceServiceSpec(predictor=predictor),
    )

    kserve_client = KServeClient(
        config_file=os.environ.get("KUBECONFIG", "~/.kube/config")
    )
    kserve_client.create(isvc)
    kserve_client.wait_isvc_ready(
        service_name,
        namespace=KSERVE_TEST_NAMESPACE,
        timeout_seconds=ISVC_READY_TIMEOUT_S,
    )
    res = generate(service_name, "./data/qwen_input_cmpl.json", chat_completions=False)
    assert res["choices"][0]["text"].strip() == "The result of 2 + 2 is 4."

    kserve_client.delete(service_name, KSERVE_TEST_NAMESPACE)


@pytest.mark.vllm
def test_huggingface_vllm_openai_chat_completions_streaming():
    service_name = "hf-qwen-chat-stream-vllm"
    predictor = V1beta1PredictorSpec(
        min_replicas=1,
        model=V1beta1ModelSpec(
            model_format=V1beta1ModelFormat(
                name="huggingface",
            ),
            args=[
                "--model_id",
                "Qwen/Qwen2-0.5B-Instruct",
                "--model_name",
                "qwen-chat-stream",
                "--backend",
                "vllm",
                "--max_model_len",
                "512",
                "--dtype",
                "bfloat16",
            ],
            env=[
                client.V1EnvVar(
                    name="VLLM_CPU_KVCACHE_SPACE",
                    value="1",
                ),
                client.V1EnvVar(
                    name="VLLM_ENABLE_V1_MULTIPROCESSING",
                    value="0",
                ),
            ],
            resources=V1ResourceRequirements(
                requests={"cpu": "2", "memory": "7Gi"},
                limits={"cpu": "2", "memory": "7Gi"},
            ),
        ),
    )

    isvc = V1beta1InferenceService(
        api_version=constants.KSERVE_V1BETA1,
        kind=constants.KSERVE_KIND_INFERENCESERVICE,
        metadata=client.V1ObjectMeta(
            name=service_name,
            namespace=KSERVE_TEST_NAMESPACE,
            annotations=ISVC_ANNOTATIONS,
        ),
        spec=V1beta1InferenceServiceSpec(predictor=predictor),
    )

    kserve_client = KServeClient(
        config_file=os.environ.get("KUBECONFIG", "~/.kube/config")
    )
    kserve_client.create(isvc)
    kserve_client.wait_isvc_ready(
        service_name,
        namespace=KSERVE_TEST_NAMESPACE,
        timeout_seconds=ISVC_READY_TIMEOUT_S,
    )

    full_response, _ = chat_completion_stream(
        service_name, "./data/qwen_input_chat_stream.json"
    )
    assert full_response.strip() == "The result of 2 + 2 is 4."

    kserve_client.delete(service_name, KSERVE_TEST_NAMESPACE)


@pytest.mark.vllm
def test_huggingface_vllm_cpu_rerank():
    service_name = "bge-reranker-base"
    predictor = V1beta1PredictorSpec(
        min_replicas=1,
        model=V1beta1ModelSpec(
            model_format=V1beta1ModelFormat(
                name="huggingface",
            ),
            args=[
                "--model_id",
                "BAAI/bge-reranker-base",
                "--backend",
                "vllm",
                "--model_revision",
                "2cfc18c9415c912f9d8155881c133215df768a70",
                "--tokenizer_revision",
                "2cfc18c9415c912f9d8155881c133215df768a70",
                "--max-model-len",
                "100",
                "--dtype",
                "bfloat16",
                "--enforce-eager",
            ],
            env=[
                client.V1EnvVar(
                    name="VLLM_CPU_KVCACHE_SPACE",
                    value="1",
                ),
                client.V1EnvVar(
                    name="VLLM_ENABLE_V1_MULTIPROCESSING",
                    value="0",
                ),
            ],
            resources=V1ResourceRequirements(
                requests={"cpu": "2", "memory": "6Gi"},
                limits={"cpu": "2", "memory": "6Gi"},
            ),
        ),
    )

    isvc = V1beta1InferenceService(
        api_version=constants.KSERVE_V1BETA1,
        kind=constants.KSERVE_KIND_INFERENCESERVICE,
        metadata=client.V1ObjectMeta(
            name=service_name,
            namespace=KSERVE_TEST_NAMESPACE,
            annotations=ISVC_ANNOTATIONS,
        ),
        spec=V1beta1InferenceServiceSpec(predictor=predictor),
    )

    kserve_client = KServeClient(
        config_file=os.environ.get("KUBECONFIG", "~/.kube/config")
    )
    kserve_client.create(isvc)
    kserve_client.wait_isvc_ready(
        service_name,
        namespace=KSERVE_TEST_NAMESPACE,
        timeout_seconds=ISVC_READY_TIMEOUT_S,
    )

    res = rerank(service_name, "./data/bge-reranker-base.json")
    assert res["results"][0]["index"] == 1
    assert res["results"][0]["relevance_score"] == pytest.approx(1.0, rel=1e-2)
    assert res["results"][0]["document"]["text"] == "The capital of France is Paris."
    assert res["results"][1]["index"] == 0
    assert res["results"][1]["relevance_score"] == pytest.approx(0.0, abs=1e-2)
    assert res["results"][1]["document"]["text"] == "The capital of Brazil is Brasilia."

    kserve_client.delete(service_name, KSERVE_TEST_NAMESPACE)