kserve/test/e2e/modelcache/test_localmodelnamespacecache.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 asyncio
import os

import pytest
from kubernetes import client
from kubernetes.client import (
    V1ResourceRequirements,
    V1PersistentVolumeSpec,
    V1LocalVolumeSource,
    V1VolumeNodeAffinity,
    V1PersistentVolumeClaimSpec,
)
from kubernetes.client.exceptions import ApiException

from kserve import constants
from kserve.api.kserve_client import KServeClient
from kserve.models.v1alpha1_local_model_node_group import V1alpha1LocalModelNodeGroup
from kserve.models.v1alpha1_local_model_node_group_spec import (
    V1alpha1LocalModelNodeGroupSpec,
)
from kserve.models.v1alpha1_local_model_namespace_cache import (
    V1alpha1LocalModelNamespaceCache,
)
from kserve.models.v1alpha1_local_model_namespace_cache_spec import (
    V1alpha1LocalModelNamespaceCacheSpec,
)
from kserve.models.v1beta1_inference_service import V1beta1InferenceService
from kserve.models.v1beta1_inference_service_spec import V1beta1InferenceServiceSpec
from kserve.models.v1beta1_predictor_spec import V1beta1PredictorSpec
from kserve.models.v1beta1_model_spec import V1beta1ModelSpec
from kserve.models.v1beta1_model_format import V1beta1ModelFormat
from ..common.utils import KSERVE_TEST_NAMESPACE, predict_isvc


@pytest.mark.modelcache
@pytest.mark.asyncio(scope="session")
async def test_sklearn_modelnamespacecache(rest_v1_client):
    service_name = "sklearn-modelnamespacecache-worker1"
    storage_uri = "gs://kfserving-examples/models/sklearn/1.0/model"
    nodes = ["minikube-m02", "minikube-m03"]

    pv_spec = V1PersistentVolumeSpec(
        access_modes=["ReadWriteOnce"],
        storage_class_name="standard",
        capacity={"storage": "1Gi"},
        local=V1LocalVolumeSource(path="/models"),
        persistent_volume_reclaim_policy="Delete",
        node_affinity=V1VolumeNodeAffinity(
            required=client.V1NodeSelector(
                node_selector_terms=[
                    client.V1NodeSelectorTerm(
                        match_expressions=[
                            client.V1NodeSelectorRequirement(
                                key="kubernetes.io/hostname",
                                operator="In",
                                values=nodes,
                            )
                        ]
                    )
                ]
            )
        ),
    )
    pvc_spec = V1PersistentVolumeClaimSpec(
        access_modes=["ReadWriteOnce"],
        resources=V1ResourceRequirements(requests={"storage": "1Gi"}),
        storage_class_name="standard",
    )

    node_group = V1alpha1LocalModelNodeGroup(
        api_version=constants.KSERVE_V1ALPHA1,
        kind=constants.KSERVE_KIND_LOCALMODELNODEGROUP,
        metadata=client.V1ObjectMeta(
            name="sklearn-nodegroup-ns",
        ),
        spec=V1alpha1LocalModelNodeGroupSpec(
            storage_limit="1Gi",
            persistent_volume_spec=pv_spec,
            persistent_volume_claim_spec=pvc_spec,
        ),
    )

    model_cache = V1alpha1LocalModelNamespaceCache(
        api_version=constants.KSERVE_V1ALPHA1,
        kind=constants.KSERVE_KIND_LOCALMODELNAMESPACECACHE,
        metadata=client.V1ObjectMeta(
            name="sklearn-model-ns",
            namespace=KSERVE_TEST_NAMESPACE,
        ),
        spec=V1alpha1LocalModelNamespaceCacheSpec(
            model_size="10Mi",
            node_groups=[node_group.metadata.name],
            source_model_uri=storage_uri,
        ),
    )

    predictor = V1beta1PredictorSpec(
        min_replicas=1,
        model=V1beta1ModelSpec(
            model_format=V1beta1ModelFormat(
                name="sklearn",
            ),
            runtime="kserve-sklearnserver",
            resources=V1ResourceRequirements(
                requests={"cpu": "100m", "memory": "256Mi"},
                limits={"cpu": "500m", "memory": "512Mi"},
            ),
            storage_uri=storage_uri,
        ),
        node_selector={"kubernetes.io/hostname": nodes[0]},
    )

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

    kserve_client = KServeClient(
        config_file=os.environ.get("KUBECONFIG", "~/.kube/config")
    )
    kserve_client.create_local_model_node_group(node_group)
    kserve_client.create_local_model_namespace_cache(
        model_cache, namespace=KSERVE_TEST_NAMESPACE
    )
    kserve_client.wait_local_model_namespace_cache_ready(
        model_cache.metadata.name, namespace=KSERVE_TEST_NAMESPACE, nodes=nodes
    )
    kserve_client.create(isvc)
    kserve_client.wait_isvc_ready(service_name, namespace=KSERVE_TEST_NAMESPACE)
    k8s_client = kserve_client.api_instance

    # Test the model is cached on the correct nodes
    worker_node_1_cache = k8s_client.get_cluster_custom_object(
        constants.KSERVE_GROUP,
        constants.KSERVE_V1ALPHA1_VERSION,
        constants.KSERVE_PLURAL_LOCALMODELNODE,
        nodes[0],
    )
    worker_node_2_cache = k8s_client.get_cluster_custom_object(
        constants.KSERVE_GROUP,
        constants.KSERVE_V1ALPHA1_VERSION,
        constants.KSERVE_PLURAL_LOCALMODELNODE,
        nodes[1],
    )
    # Status key for namespace-scoped models is namespace/modelName
    status_key = f"{KSERVE_TEST_NAMESPACE}/{model_cache.metadata.name}"
    assert worker_node_1_cache["status"]["modelStatus"][status_key] == "ModelDownloaded"
    assert worker_node_2_cache["status"]["modelStatus"][status_key] == "ModelDownloaded"

    # Test the model is not cached on the controller node
    with pytest.raises(ApiException):
        k8s_client.get_cluster_custom_object(
            constants.KSERVE_GROUP,
            constants.KSERVE_V1ALPHA1_VERSION,
            constants.KSERVE_PLURAL_LOCALMODELNODE,
            "aks-agentpool-27350515-vmss000000",
        )

    res = await predict_isvc(rest_v1_client, service_name, "./data/iris_input.json")
    assert res["predictions"] == [1, 1]
    kserve_client.delete(service_name, KSERVE_TEST_NAMESPACE)
    # Wait for the isvc to be deleted to avoid modelcache still in use error when deleting the model cache
    await asyncio.sleep(30)
    kserve_client.delete_local_model_namespace_cache(
        model_cache.metadata.name, namespace=KSERVE_TEST_NAMESPACE
    )
    kserve_client.delete_local_model_node_group(node_group.metadata.name)