diff --git a/docs/book/component-guide/model-deployers/README.md b/docs/book/component-guide/model-deployers/README.md index cd5150f0811..c5281dbed7e 100644 --- a/docs/book/component-guide/model-deployers/README.md +++ b/docs/book/component-guide/model-deployers/README.md @@ -63,6 +63,7 @@ integrations: | [Hugging Face](huggingface.md) | `huggingface` | `huggingface` | Deploys ML model on Hugging Face Inference Endpoints | | [Databricks](databricks.md) | `databricks` | `databricks` | Deploying models to Databricks Inference Endpoints with Databricks | | [vLLM](vllm.md) | `vllm` | `vllm` | Deploys LLM using vLLM locally | +| [vLLM](vllm.md) | `vllm-kubernetes` | `vllm` | Deploys LLM using vLLM on Kubernetes | | [Custom Implementation](custom.md) | _custom_ | | Extend the Artifact Store abstraction and provide your own implementation | {% hint style="info" %} diff --git a/docs/book/component-guide/model-deployers/vllm.md b/docs/book/component-guide/model-deployers/vllm.md index 62fc47f52a3..cb413e719c7 100644 --- a/docs/book/component-guide/model-deployers/vllm.md +++ b/docs/book/component-guide/model-deployers/vllm.md @@ -73,4 +73,133 @@ Within the `VLLMDeploymentService` you can configure: * `dtype`: Data type for model weights and activations. Allowed choices: ['auto', 'half', 'float16', 'bfloat16', 'float', 'float32'] * `revision`: The specific model version to use. It can be a branch name, a tag name, or a commit id. If unspecified, will use the default version. +## Deploying on Kubernetes + +The `vllm-kubernetes` flavor deploys the vLLM OpenAI-compatible server as a long-lived Deployment and Service on a Kubernetes cluster, instead of running it as a local daemon process. It builds on the same [model deployer abstraction](README.md) as [Seldon Core](seldon.md), with the Kubernetes cluster reached through a `kubectl` context, an in-cluster configuration, or a [Service Connector](https://docs.zenml.io/how-to/infrastructure-deployment/auth-management/). + +Only `deploy` and `delete` are supported for this flavor. Stopping and starting a deployed server raise `NotImplementedError`, delete it and deploy again instead. + +### Installation + +The Kubernetes flavor talks to the cluster through the `kubernetes` package and to the deployed server through the `openai` client. It does not need the `vllm` package itself. If you only need the Kubernetes flavor, you can skip the `vllm` install: + +```bash +pip install kubernetes openai +``` + +`zenml integration install vllm -y` also works and additionally installs `vllm`, which is only needed for the local flavor. + +### Registering the model deployer + +```bash +zenml model-deployer register --flavor=vllm-kubernetes \ + --kubernetes_namespace= \ + --hf_token= +``` + +You can configure: + +* `kubernetes_context`: name of a `kubectl` context to deploy to. Ignored if the component is linked to a Kubernetes service connector. +* `kubernetes_namespace`: namespace vLLM deployments are provisioned into, defaults to `zenml-vllm`. Can be overridden per deployment. +* `incluster`: if `True`, connects using the configuration of the cluster the client itself runs in. Requires the client to run inside a Kubernetes pod, and ignores `kubernetes_context` when set. +* `default_image`: image used when a deployment doesn't specify its own, defaults to `vllm/vllm-openai:v0.25.1`. +* `default_service_type`: Kubernetes Service type used when a deployment doesn't specify its own, defaults to `LoadBalancer`. +* `hf_token`: Hugging Face access token used to download gated models, used as a fallback for deployments that configure neither their own token nor an existing secret. + +### Using a service connector + +As with Seldon Core, the recommended way to authenticate to the target cluster is [a Service Connector](https://docs.zenml.io/how-to/infrastructure-deployment/auth-management/) instead of `kubernetes_context`. Pick the connector for your cloud provider (AWS, GCP, Azure) or the generic [Kubernetes Service Connector](https://docs.zenml.io/how-to/infrastructure-deployment/auth-management/kubernetes-service-connector), then connect it to the model deployer: + +```sh +zenml service-connector register --type kubernetes --auto-configure + +zenml model-deployer register --flavor=vllm-kubernetes +zenml model-deployer connect --connector +``` + +We can now use the model deployer in our stack: + +```bash +zenml stack update --model-deployer= +``` + +### Deploy an LLM + +The deploy entry point is the same `deploy_model` method used by every model deployer. Here is an example step that deploys a gated Hugging Face model with a GPU: + +```python +from typing import Annotated +from zenml import pipeline, step +from zenml.integrations.vllm.model_deployers import KubernetesVLLMModelDeployer +from zenml.integrations.vllm.services import ( + VLLMKubernetesDeploymentService, + VLLMKubernetesServiceConfig, +) + + +@step +def deploy_vllm_kubernetes( + model: str, + timeout: int = 1200, +) -> Annotated[VLLMKubernetesDeploymentService, "vllm_deployment"]: + model_deployer = KubernetesVLLMModelDeployer.get_active_model_deployer() + + service_config = VLLMKubernetesServiceConfig( + model=model, + resources={"limits": {"nvidia.com/gpu": "1"}}, + hf_token="", + ) + + service = model_deployer.deploy_model( + config=service_config, + service_type=VLLMKubernetesDeploymentService.SERVICE_TYPE, + timeout=timeout, + ) + return service + + +@pipeline +def deploy_vllm_kubernetes_pipeline(model: str, timeout: int = 1200): + deploy_vllm_kubernetes(model=model, timeout=timeout) +``` + +`VLLMKubernetesServiceConfig` accepts the same engine arguments as the local `VLLMDeploymentService` (`model`, `tokenizer`, `served_model_name`, `trust_remote_code`, `tokenizer_mode`, `dtype`, `revision`), plus: + +* `namespace`: Kubernetes namespace to deploy into. Falls back to the model deployer's `kubernetes_namespace` when unset. +* `image`: container image running the vLLM OpenAI server. Falls back to the model deployer's `default_image` when unset. +* `port`: port the vLLM OpenAI server listens on, defaults to `8000`. +* `service_type`: Kubernetes Service type used to expose the server. Falls back to the model deployer's `default_service_type` when unset. +* `replicas`: number of pod replicas for the Deployment, defaults to `1`. +* `resources`: native Kubernetes resource requests and limits for the container, e.g. `{"limits": {"nvidia.com/gpu": "1"}}`. A GPU limit is automatically mirrored into requests, since Kubernetes requires the two to match for extended resources. +* `pod_settings`: a `KubernetesPodSettings` object for node selectors, tolerations, affinity, and other pod configuration not modeled as a dedicated field. +* `hf_token`: Hugging Face access token injected into the container through a managed secret. +* `existing_hf_secret`: name of an existing Kubernetes secret to use instead of `hf_token`. Must contain the token under a `token` key. +* `env`: additional environment variables set on the container. +* `extra_serve_args`: additional CLI args appended to the vLLM OpenAI server args, for engine options not modeled as dedicated fields. +* `shm_size`: size limit of the `/dev/shm` volume mounted into the container, defaults to `2Gi`. + +### Reaching the deployed server + +A deployment that doesn't set its own `service_type` is exposed through the model deployer's `default_service_type`, which is `LoadBalancer` unless the deployer was registered with a different one, for example `--default_service_type=ClusterIP`. + +With a `LoadBalancer` Service, point an OpenAI client at the external address assigned to the Service. With a `ClusterIP` Service, which only has an address inside the cluster, forward the Service port instead to reach it from your local machine: + +```bash +kubectl port-forward -n svc/ 8000:8000 +``` + +and then point an OpenAI client at `http://localhost:8000/v1`. + +### Custom images + +`image` (per deployment) and `default_image` (on the model deployer) both accept any vLLM OpenAI-compatible image, and override the pinned default `vllm/vllm-openai:v0.25.1`. For clusters without a GPU, use a CPU-capable image such as `vllm/vllm-openai-cpu`. + +### Teardown + +```bash +zenml model-deployer models delete +``` + +This deletes the Deployment, Service, and the managed secret if one was created for `hf_token`. The namespace itself is never deleted. +
ZenML Scarf
diff --git a/src/zenml/integrations/kubernetes/manifest_utils.py b/src/zenml/integrations/kubernetes/manifest_utils.py index 796a9dd0bc2..3de6567bbc0 100644 --- a/src/zenml/integrations/kubernetes/manifest_utils.py +++ b/src/zenml/integrations/kubernetes/manifest_utils.py @@ -225,7 +225,8 @@ def add_pod_settings( for container in pod_spec.containers: assert isinstance(container, k8s_client.V1Container) - container._resources = settings.resources + if settings.resources: + container._resources = settings.resources if settings.volume_mounts: if container.volume_mounts: container.volume_mounts.extend(settings.volume_mounts) diff --git a/src/zenml/integrations/vllm/__init__.py b/src/zenml/integrations/vllm/__init__.py index 2c6a3a25bc1..77dd2d2d893 100644 --- a/src/zenml/integrations/vllm/__init__.py +++ b/src/zenml/integrations/vllm/__init__.py @@ -19,6 +19,7 @@ from zenml.integrations.constants import VLLM VLLM_MODEL_DEPLOYER = "vllm" +VLLM_KUBERNETES_MODEL_DEPLOYER_FLAVOR = "vllm-kubernetes" logger = get_logger(__name__) @@ -28,13 +29,28 @@ class VLLMIntegration(Integration): NAME = VLLM - REQUIREMENTS = ["vllm>=0.6.0,<0.7.0", "openai>=1.0.0"] + REQUIREMENTS = [ + "vllm>=0.6.0,<0.7.0", + "openai>=1.0.0", + "kubernetes>=21.7,<26", + ] @classmethod def activate(cls) -> None: """Activates the integration.""" from zenml.integrations.vllm import services + try: + from zenml.integrations.vllm.services import ( # noqa: F401 + vllm_kubernetes_deployment, + ) + except ImportError: + logger.debug( + "Could not import the vLLM Kubernetes deployment service. " + "The `vllm-kubernetes` model deployer flavor will not be " + "usable until the `kubernetes` package is installed." + ) + @classmethod def flavors(cls) -> List[Type[Flavor]]: """Declare the stack component flavors for the vLLM integration. @@ -42,8 +58,11 @@ def flavors(cls) -> List[Type[Flavor]]: Returns: List of stack component flavors for this integration. """ - from zenml.integrations.vllm.flavors import VLLMModelDeployerFlavor + from zenml.integrations.vllm.flavors import ( + KubernetesVLLMModelDeployerFlavor, + VLLMModelDeployerFlavor, + ) - return [VLLMModelDeployerFlavor] + return [VLLMModelDeployerFlavor, KubernetesVLLMModelDeployerFlavor] diff --git a/src/zenml/integrations/vllm/flavors/__init__.py b/src/zenml/integrations/vllm/flavors/__init__.py index 1a45c9da09a..76c52ea4ee3 100644 --- a/src/zenml/integrations/vllm/flavors/__init__.py +++ b/src/zenml/integrations/vllm/flavors/__init__.py @@ -13,9 +13,18 @@ # permissions and limitations under the License. """vLLM integration flavors.""" +from zenml.integrations.vllm.flavors.vllm_kubernetes_model_deployer_flavor import ( # noqa + KubernetesVLLMModelDeployerConfig, + KubernetesVLLMModelDeployerFlavor, +) from zenml.integrations.vllm.flavors.vllm_model_deployer_flavor import ( # noqa VLLMModelDeployerConfig, VLLMModelDeployerFlavor, ) -__all__ = ["VLLMModelDeployerConfig", "VLLMModelDeployerFlavor"] +__all__ = [ + "KubernetesVLLMModelDeployerConfig", + "KubernetesVLLMModelDeployerFlavor", + "VLLMModelDeployerConfig", + "VLLMModelDeployerFlavor", +] diff --git a/src/zenml/integrations/vllm/flavors/vllm_kubernetes_model_deployer_flavor.py b/src/zenml/integrations/vllm/flavors/vllm_kubernetes_model_deployer_flavor.py new file mode 100644 index 00000000000..0f15730df55 --- /dev/null +++ b/src/zenml/integrations/vllm/flavors/vllm_kubernetes_model_deployer_flavor.py @@ -0,0 +1,164 @@ +# Copyright (c) ZenML GmbH 2026. All Rights Reserved. +# +# 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: +# +# https://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. +"""vLLM Kubernetes model deployer flavor.""" + +from typing import TYPE_CHECKING, Optional, Type + +from pydantic import Field + +from zenml.constants import KUBERNETES_CLUSTER_RESOURCE_TYPE +from zenml.enums import KubernetesServiceType +from zenml.integrations.vllm import VLLM_KUBERNETES_MODEL_DEPLOYER_FLAVOR +from zenml.model_deployers.base_model_deployer import ( + BaseModelDeployerConfig, + BaseModelDeployerFlavor, +) +from zenml.models import ServiceConnectorRequirements +from zenml.utils.secret_utils import SecretField + +if TYPE_CHECKING: + from zenml.integrations.vllm.model_deployers import ( + KubernetesVLLMModelDeployer, + ) + +DEFAULT_VLLM_IMAGE = "vllm/vllm-openai:v0.25.1" + + +class KubernetesVLLMModelDeployerConfig(BaseModelDeployerConfig): + """Configuration for the Kubernetes vLLM model deployer. + + Attributes: + kubernetes_context: name of a Kubernetes context to deploy vLLM servers to. If the stack component is linked to a Kubernetes service connector, this field is ignored. + kubernetes_namespace: default Kubernetes namespace for vLLM deployments. Can be overridden per-deployment on the service config. + incluster: if `True`, connect using the Kubernetes configuration of the cluster the client itself is running in. Requires the client to run in a Kubernetes pod. If set, `kubernetes_context` is ignored. + default_image: default vLLM server image used when a deployment doesn't specify one. + default_service_type: default Kubernetes Service type used when a deployment doesn't specify one. + hf_token: Hugging Face access token used to download gated models when a deployment doesn't configure its own token or existing secret. + """ + + kubernetes_context: Optional[str] = Field( + default=None, + description="Name of a Kubernetes context to deploy vLLM servers " + "to. Ignored when the stack component is linked to a Kubernetes " + "service connector", + ) + kubernetes_namespace: str = Field( + default="zenml-vllm", + description="Kubernetes namespace in which vLLM deployment " + "servers are provisioned and managed by ZenML", + ) + incluster: bool = Field( + default=False, + description="Connect using the Kubernetes configuration of the " + "cluster the client itself is running in. Requires the client to " + "run inside a Kubernetes pod. When set, `kubernetes_context` is " + "ignored", + ) + default_image: str = Field( + default=DEFAULT_VLLM_IMAGE, + description="Container image used to run the vLLM OpenAI-" + "compatible server when a deployment doesn't specify its own " + "image", + ) + default_service_type: KubernetesServiceType = Field( + default=KubernetesServiceType.LOAD_BALANCER, + description="Kubernetes Service type used to expose a vLLM " + "deployment when a deployment doesn't specify its own service " + "type", + ) + hf_token: Optional[str] = SecretField( + default=None, + description="Hugging Face access token used to download gated " + "models. Used as a fallback for deployments that configure " + "neither their own token nor an existing Kubernetes secret", + ) + + +class KubernetesVLLMModelDeployerFlavor(BaseModelDeployerFlavor): + """Kubernetes vLLM model deployer flavor.""" + + @property + def name(self) -> str: + """Name of the flavor. + + Returns: + The name of the flavor. + """ + return VLLM_KUBERNETES_MODEL_DEPLOYER_FLAVOR + + @property + def service_connector_requirements( + self, + ) -> Optional[ServiceConnectorRequirements]: + """Service connector resource requirements for service connectors. + + Specifies resource requirements that are used to filter the available + service connector types that are compatible with this flavor. + + Returns: + Requirements for compatible service connectors, if a service + connector is required for this flavor. + """ + return ServiceConnectorRequirements( + resource_type=KUBERNETES_CLUSTER_RESOURCE_TYPE, + ) + + @property + def docs_url(self) -> Optional[str]: + """A url to point at docs explaining this flavor. + + Returns: + A flavor docs url. + """ + return self.generate_default_docs_url() + + @property + def sdk_docs_url(self) -> Optional[str]: + """A url to point at SDK docs explaining this flavor. + + Returns: + A flavor SDK docs url. + """ + return self.generate_default_sdk_docs_url() + + @property + def logo_url(self) -> str: + """A url to represent the flavor in the dashboard. + + Returns: + The flavor logo. + """ + return "https://public-flavor-logos.s3.eu-central-1.amazonaws.com/model_deployer/vllm.png" + + @property + def config_class(self) -> Type[KubernetesVLLMModelDeployerConfig]: + """Returns `KubernetesVLLMModelDeployerConfig` config class. + + Returns: + The config class. + """ + return KubernetesVLLMModelDeployerConfig + + @property + def implementation_class(self) -> Type["KubernetesVLLMModelDeployer"]: + """Implementation class for this flavor. + + Returns: + The implementation class. + """ + from zenml.integrations.vllm.model_deployers import ( + KubernetesVLLMModelDeployer, + ) + + return KubernetesVLLMModelDeployer diff --git a/src/zenml/integrations/vllm/model_deployers/__init__.py b/src/zenml/integrations/vllm/model_deployers/__init__.py index cb7edfe68ce..591d7b2e9d5 100644 --- a/src/zenml/integrations/vllm/model_deployers/__init__.py +++ b/src/zenml/integrations/vllm/model_deployers/__init__.py @@ -17,3 +17,15 @@ ) __all__ = ["VLLMModelDeployer"] + +# The Kubernetes vLLM model deployer depends on the `kubernetes` package. +# Guard the import so that the local vLLM model deployer keeps working in +# environments where `kubernetes` is not installed. +try: + from zenml.integrations.vllm.model_deployers.vllm_kubernetes_model_deployer import ( # noqa + KubernetesVLLMModelDeployer, + ) + + __all__.append("KubernetesVLLMModelDeployer") +except ImportError: + pass diff --git a/src/zenml/integrations/vllm/model_deployers/vllm_kubernetes_model_deployer.py b/src/zenml/integrations/vllm/model_deployers/vllm_kubernetes_model_deployer.py new file mode 100644 index 00000000000..aea3b948cbd --- /dev/null +++ b/src/zenml/integrations/vllm/model_deployers/vllm_kubernetes_model_deployer.py @@ -0,0 +1,283 @@ +# Copyright (c) ZenML GmbH 2026. All Rights Reserved. +# +# 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: +# +# https://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. +"""Implementation of the Kubernetes vLLM Model Deployer.""" + +import os +from typing import ClassVar, Dict, Optional, Type, cast +from uuid import UUID + +from kubernetes import client as k8s_client + +from zenml.integrations.kubernetes import kube_utils +from zenml.integrations.kubernetes.k8s_applier import KubernetesApplier +from zenml.integrations.vllm.flavors.vllm_kubernetes_model_deployer_flavor import ( + KubernetesVLLMModelDeployerConfig, + KubernetesVLLMModelDeployerFlavor, +) +from zenml.integrations.vllm.services.vllm_deployment import ( + VLLM_HEALTHCHECK_URL_PATH, + VLLM_PREDICTION_URL_PATH, +) +from zenml.integrations.vllm.services.vllm_kubernetes_deployment import ( + VLLMKubernetesDeploymentService, + VLLMKubernetesServiceConfig, +) +from zenml.logger import get_logger +from zenml.model_deployers import BaseModelDeployer, BaseModelDeployerFlavor +from zenml.services.service import BaseService, ServiceConfig + +logger = get_logger(__name__) + +DEFAULT_VLLM_KUBERNETES_DEPLOYMENT_START_STOP_TIMEOUT = 300 + + +class KubernetesVLLMModelDeployer(BaseModelDeployer): + """Kubernetes vLLM model deployer stack component implementation.""" + + NAME: ClassVar[str] = "vLLM Kubernetes" + FLAVOR: ClassVar[Type[BaseModelDeployerFlavor]] = ( + KubernetesVLLMModelDeployerFlavor + ) + + _k8s_client: Optional[k8s_client.ApiClient] = None + _applier: Optional[KubernetesApplier] = None + + @property + def config(self) -> KubernetesVLLMModelDeployerConfig: + """Returns the `KubernetesVLLMModelDeployerConfig` config. + + Returns: + The configuration. + """ + return cast(KubernetesVLLMModelDeployerConfig, self._config) + + def get_kube_client( + self, incluster: Optional[bool] = None + ) -> k8s_client.ApiClient: + """Getter for the Kubernetes API client. + + Args: + incluster: Whether to use the in-cluster config or not. Overrides + the `incluster` setting in the config. + + Raises: + RuntimeError: if the Kubernetes connector behaves unexpectedly. + + Returns: + The Kubernetes API client. + """ + if incluster is None: + incluster = self.config.incluster + + if incluster: + try: + kube_utils.load_kube_config( + incluster=incluster, + ) + self._k8s_client = k8s_client.ApiClient() + return self._k8s_client + except Exception as e: + if self.connector: + message = ( + "Falling back to using the linked service connector " + "configuration." + ) + elif self.config.kubernetes_context: + message = ( + f"Falling back to using the configured " + f"'{self.config.kubernetes_context}' kubernetes context." + ) + else: + raise RuntimeError( + f"The model deployer failed to load the in-cluster " + f"Kubernetes configuration and there is no service " + f"connector or kubernetes_context to fall back to: {e}" + ) from e + + logger.debug( + f"Could not load the in-cluster Kubernetes configuration: " + f"{e}. {message}" + ) + + if self._k8s_client and not self.connector_has_expired(): + return self._k8s_client + + connector = self.get_connector() + if connector: + client = connector.connect() + if not isinstance(client, k8s_client.ApiClient): + raise RuntimeError( + f"Expected a k8s_client.ApiClient while trying to use the " + f"linked connector, but got {type(client)}." + ) + self._k8s_client = client + else: + kube_utils.load_kube_config( + context=self.config.kubernetes_context, + ) + self._k8s_client = k8s_client.ApiClient() + + return self._k8s_client + + @property + def k8s_applier(self) -> KubernetesApplier: + """Get or create the Kubernetes Applier instance. + + Returns: + Kubernetes Applier instance. + """ + if self.connector_has_expired(): + self._applier = None + + if not self._applier: + self._applier = KubernetesApplier( + api_client=self.get_kube_client() + ) + return self._applier + + @staticmethod + def get_model_server_info( # type: ignore[override] + service_instance: "VLLMKubernetesDeploymentService", + ) -> Dict[str, Optional[str]]: + """Return implementation specific information on the model server. + + Args: + service_instance: vLLM Kubernetes deployment service object. + + Returns: + A dictionary containing the model server information. + """ + base_url = service_instance.service_url + return { + "PREDICTION_URL": os.path.join(base_url, VLLM_PREDICTION_URL_PATH) + if base_url + else None, + "HEALTH_CHECK_URL": os.path.join( + base_url, VLLM_HEALTHCHECK_URL_PATH + ) + if base_url + else None, + "MODEL": service_instance.config.model, + "IMAGE": service_instance.image, + "NAMESPACE": service_instance.namespace, + "DEPLOYMENT_NAME": service_instance.deployment_name, + "SERVICE_TYPE": str(service_instance.service_type), + } + + def perform_deploy_model( + self, + id: UUID, + config: ServiceConfig, + timeout: int = DEFAULT_VLLM_KUBERNETES_DEPLOYMENT_START_STOP_TIMEOUT, + ) -> BaseService: + """Create a new vLLM Kubernetes deployment service. + + Args: + id: the UUID of the model server to deploy. + config: the configuration of the model to be deployed with vLLM + on Kubernetes. + timeout: the timeout in seconds to wait for the vLLM server to + be provisioned and successfully started. If set to 0, the + method will return immediately after the vLLM server is + provisioned, without waiting for it to fully start. + + Returns: + The ZenML vLLM Kubernetes deployment service object that can be + used to interact with the remote vLLM server. + """ + config = cast(VLLMKubernetesServiceConfig, config) + + config.namespace = config.namespace or self.config.kubernetes_namespace + config.image = config.image or self.config.default_image + config.service_type = ( + config.service_type or self.config.default_service_type + ) + if not config.hf_token and not config.existing_hf_secret: + config.hf_token = self.config.hf_token + + service = VLLMKubernetesDeploymentService(uuid=id, config=config) + logger.info( + f"Creating a new vLLM Kubernetes deployment service: {service}" + ) + service.start(timeout=timeout) + + return service + + def perform_stop_model( + self, + service: BaseService, + timeout: int = DEFAULT_VLLM_KUBERNETES_DEPLOYMENT_START_STOP_TIMEOUT, + force: bool = False, + ) -> BaseService: + """Stop a vLLM Kubernetes model server. + + Args: + service: The service to stop. + timeout: timeout in seconds to wait for the service to stop. + force: if True, force the service to stop. + + Raises: + NotImplementedError: stopping vLLM Kubernetes model servers is + not supported. + + Returns: + The stopped service. + """ + raise NotImplementedError( + "Stopping vLLM Kubernetes model servers is not implemented. Try " + "deleting the vLLM Kubernetes model server instead." + ) + + def perform_start_model( + self, + service: BaseService, + timeout: int = DEFAULT_VLLM_KUBERNETES_DEPLOYMENT_START_STOP_TIMEOUT, + ) -> BaseService: + """Start a vLLM Kubernetes model deployment server. + + Args: + service: The service to start. + timeout: timeout in seconds to wait for the service to become + active. If set to 0, the method will return immediately + after provisioning the service, without waiting for it to + become active. + + Raises: + NotImplementedError: since we don't support starting vLLM + Kubernetes model servers. + + Returns: + The started service. + """ + raise NotImplementedError( + "Starting vLLM Kubernetes model servers is not implemented" + ) + + def perform_delete_model( + self, + service: BaseService, + timeout: int = DEFAULT_VLLM_KUBERNETES_DEPLOYMENT_START_STOP_TIMEOUT, + force: bool = False, + ) -> None: + """Delete a vLLM Kubernetes model deployment server. + + Args: + service: The service to delete. + timeout: timeout in seconds to wait for the service to stop. If + set to 0, the method will return immediately after + deprovisioning the service, without waiting for it to stop. + force: if True, force the service to stop. + """ + service = cast(VLLMKubernetesDeploymentService, service) + service.stop(timeout=timeout, force=force) diff --git a/src/zenml/integrations/vllm/services/__init__.py b/src/zenml/integrations/vllm/services/__init__.py index 67b4158c937..192096e13c5 100644 --- a/src/zenml/integrations/vllm/services/__init__.py +++ b/src/zenml/integrations/vllm/services/__init__.py @@ -15,5 +15,27 @@ from zenml.integrations.vllm.services.vllm_deployment import ( # noqa VLLMDeploymentService, + VLLMEngineArgs, VLLMServiceConfig, ) + +__all__ = [ + "VLLMDeploymentService", + "VLLMEngineArgs", + "VLLMServiceConfig", +] + +# The Kubernetes vLLM deployment service depends on the `kubernetes` +# package. Guard the import so that the local vLLM deployment service keeps +# working in environments where `kubernetes` is not installed. +try: + from zenml.integrations.vllm.services.vllm_kubernetes_deployment import ( # noqa + VLLMKubernetesDeploymentService, + VLLMKubernetesServiceConfig, + ) + + __all__.extend( + ["VLLMKubernetesDeploymentService", "VLLMKubernetesServiceConfig"] + ) +except ImportError: + pass diff --git a/src/zenml/integrations/vllm/services/vllm_deployment.py b/src/zenml/integrations/vllm/services/vllm_deployment.py index 0ff4617f20e..c0cf065d6ab 100644 --- a/src/zenml/integrations/vllm/services/vllm_deployment.py +++ b/src/zenml/integrations/vllm/services/vllm_deployment.py @@ -17,6 +17,8 @@ import os from typing import Any, List, Optional, Union +from pydantic import BaseModel + from zenml.constants import DEFAULT_LOCAL_SERVICE_IP_ADDRESS from zenml.logger import get_logger from zenml.models.v2.misc.service import ServiceType @@ -71,13 +73,10 @@ def prediction_url(self) -> Optional[str]: return os.path.join(uri, self.config.prediction_url_path) -class VLLMServiceConfig(LocalDaemonServiceConfig): - """vLLM service configurations.""" +class VLLMEngineArgs(BaseModel): + """vLLM engine arguments shared across deployment targets.""" model: str - port: int - host: Optional[str] = None - blocking: bool = True # If unspecified, model name or path will be used. tokenizer: Optional[str] = None served_model_name: Optional[Union[str, List[str]]] = None @@ -92,6 +91,14 @@ class VLLMServiceConfig(LocalDaemonServiceConfig): revision: Optional[str] = None +class VLLMServiceConfig(VLLMEngineArgs, LocalDaemonServiceConfig): + """vLLM service configurations.""" + + port: int + host: Optional[str] = None + blocking: bool = True + + class VLLMDeploymentService(LocalDaemonService, BaseDeploymentService): """vLLM Inference Server Deployment Service.""" diff --git a/src/zenml/integrations/vllm/services/vllm_kubernetes_deployment.py b/src/zenml/integrations/vllm/services/vllm_kubernetes_deployment.py new file mode 100644 index 00000000000..946de667a41 --- /dev/null +++ b/src/zenml/integrations/vllm/services/vllm_kubernetes_deployment.py @@ -0,0 +1,945 @@ +# Copyright (c) ZenML GmbH 2025. All Rights Reserved. +# +# 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: +# +# https://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. +"""Implementation of the vLLM Kubernetes Inference Server Service.""" + +import os +from typing import ( + TYPE_CHECKING, + Any, + Dict, + Generator, + List, + Optional, + Tuple, + cast, +) + +from kubernetes import client as k8s_client +from pydantic import Field + +from zenml.enums import KubernetesServiceType, ServiceState +from zenml.integrations.kubernetes import kube_utils +from zenml.integrations.kubernetes.k8s_applier import ResourceInventoryItem +from zenml.integrations.kubernetes.manifest_utils import add_pod_settings +from zenml.integrations.kubernetes.pod_settings import KubernetesPodSettings +from zenml.integrations.kubernetes.serialization_utils import ( + normalize_resource_to_dict, +) +from zenml.integrations.vllm.flavors.vllm_kubernetes_model_deployer_flavor import ( + DEFAULT_VLLM_IMAGE, +) +from zenml.integrations.vllm.services.vllm_deployment import ( + VLLM_HEALTHCHECK_URL_PATH, + VLLM_PREDICTION_URL_PATH, + VLLMEngineArgs, +) +from zenml.logger import get_logger +from zenml.models.v2.misc.service import ServiceType +from zenml.services.service import BaseDeploymentService, ServiceConfig +from zenml.services.service_status import ServiceStatus + +if TYPE_CHECKING: + from zenml.integrations.vllm.model_deployers.vllm_kubernetes_model_deployer import ( + KubernetesVLLMModelDeployer, + ) + +logger = get_logger(__name__) + + +VLLM_FIELD_MANAGER = "zenml-vllm-deployer" +VLLM_CONTAINER_NAME = "vllm" +VLLM_HF_TOKEN_ENV_VAR = "HF_TOKEN" +VLLM_HF_TOKEN_SECRET_KEY = "token" +VLLM_HEALTH_PROBE_PATH = f"/{VLLM_HEALTHCHECK_URL_PATH}" +SHM_VOLUME_NAME = "shm" +SHM_VOLUME_MOUNT_PATH = "/dev/shm" +POD_FAILURE_REASONS = {"CrashLoopBackOff", "ImagePullBackOff", "ErrImagePull"} + + +def _build_container_args( + engine_args: VLLMEngineArgs, + port: int, + extra_serve_args: List[str], +) -> List[str]: + """Build vLLM OpenAI server CLI args from engine arguments. + + Args: + engine_args: vLLM engine arguments. + port: Port the server listens on inside the container. + extra_serve_args: Additional CLI args appended verbatim. + + Returns: + CLI args for the vLLM OpenAI API server entrypoint. + """ + args = [ + "--model", + engine_args.model, + "--port", + str(port), + "--host", + "0.0.0.0", + ] + if engine_args.tokenizer: + args += ["--tokenizer", engine_args.tokenizer] + if engine_args.served_model_name: + served_model_names = ( + engine_args.served_model_name + if isinstance(engine_args.served_model_name, list) + else [engine_args.served_model_name] + ) + args += ["--served-model-name", *served_model_names] + if engine_args.trust_remote_code: + args.append("--trust-remote-code") + if engine_args.tokenizer_mode: + args += ["--tokenizer-mode", engine_args.tokenizer_mode] + if engine_args.dtype: + args += ["--dtype", engine_args.dtype] + if engine_args.revision: + args += ["--revision", engine_args.revision] + args += extra_serve_args + return args + + +def _build_probe( + port: int, + initial_delay_seconds: int, + period_seconds: int, + timeout_seconds: int, + failure_threshold: int, +) -> k8s_client.V1Probe: + """Build an HTTP GET probe against the vLLM health endpoint. + + Args: + port: Port the health endpoint listens on. + initial_delay_seconds: Delay before the first probe. + period_seconds: Interval between probes. + timeout_seconds: Timeout for each probe. + failure_threshold: Consecutive failures before the probe fails. + + Returns: + The configured probe. + """ + return k8s_client.V1Probe( + http_get=k8s_client.V1HTTPGetAction( + path=VLLM_HEALTH_PROBE_PATH, port=port + ), + initial_delay_seconds=initial_delay_seconds, + period_seconds=period_seconds, + timeout_seconds=timeout_seconds, + failure_threshold=failure_threshold, + ) + + +def _resources_with_mirrored_gpu_requests( + resources: Dict[str, Dict[str, str]], +) -> Dict[str, Dict[str, str]]: + """Mirror a configured GPU limit into the resource requests. + + Kubernetes requires requests to equal limits for extended resources + like `nvidia.com/gpu`, so a config that only sets a GPU limit would + otherwise leave the pod unschedulable. + + Args: + resources: Native Kubernetes resource requests and limits. + + Returns: + The resources with the GPU limit mirrored into requests. + """ + gpu_limit = resources.get("limits", {}).get("nvidia.com/gpu") + if not gpu_limit: + return resources + + requests = dict(resources.get("requests", {})) + requests.setdefault("nvidia.com/gpu", gpu_limit) + merged = dict(resources) + merged["requests"] = requests + return merged + + +def _pod_container_name(pod: k8s_client.V1Pod) -> str: + """Name of a pod's primary container, or the default vLLM container name. + + Args: + pod: Pod to inspect. + + Returns: + The container name. + """ + return ( + pod.spec.containers[0].name + if pod.spec and pod.spec.containers + else VLLM_CONTAINER_NAME + ) + + +def build_vllm_deployment_manifest( + name: str, + namespace: str, + image: str, + replicas: int, + port: int, + engine_args: VLLMEngineArgs, + extra_serve_args: List[str], + labels: Dict[str, str], + resources: Dict[str, Dict[str, str]], + shm_size: Optional[str], + env: Dict[str, str], + hf_secret_name: Optional[str], + pod_settings: Optional[KubernetesPodSettings], +) -> Dict[str, Any]: + """Build the manifest for the vLLM Deployment. + + Args: + name: Name of the Deployment. + namespace: Namespace to deploy into. + image: Container image to run. + replicas: Number of pod replicas. + port: Port the vLLM server listens on. + engine_args: vLLM engine arguments used to build the container args. + extra_serve_args: Additional CLI args appended to the container args. + labels: Labels applied to the Deployment and its pod template. + resources: Native Kubernetes resource requests and limits. + shm_size: Size limit of the `/dev/shm` volume. No volume is added + if unset. + env: Environment variables set on the container. + hf_secret_name: Name of the secret holding the Hugging Face token. + No `HF_TOKEN` variable is added if unset. + pod_settings: Additional pod configuration merged into the pod spec. + + Returns: + The Deployment manifest. + """ + env_vars = [ + k8s_client.V1EnvVar(name=key, value=value) + for key, value in env.items() + ] + if hf_secret_name: + env_vars.append( + k8s_client.V1EnvVar( + name=VLLM_HF_TOKEN_ENV_VAR, + value_from=k8s_client.V1EnvVarSource( + secret_key_ref=k8s_client.V1SecretKeySelector( + name=hf_secret_name, + key=VLLM_HF_TOKEN_SECRET_KEY, + ) + ), + ) + ) + + volumes = [] + volume_mounts = [] + if shm_size: + volumes.append( + k8s_client.V1Volume( + name=SHM_VOLUME_NAME, + empty_dir=k8s_client.V1EmptyDirVolumeSource( + medium="Memory", size_limit=shm_size + ), + ) + ) + volume_mounts.append( + k8s_client.V1VolumeMount( + name=SHM_VOLUME_NAME, mount_path=SHM_VOLUME_MOUNT_PATH + ) + ) + + container = k8s_client.V1Container( + name=VLLM_CONTAINER_NAME, + image=image, + args=_build_container_args( + engine_args=engine_args, + port=port, + extra_serve_args=extra_serve_args, + ), + ports=[k8s_client.V1ContainerPort(container_port=port)], + env=env_vars, + volume_mounts=volume_mounts, + readiness_probe=_build_probe( + port=port, + initial_delay_seconds=10, + period_seconds=10, + timeout_seconds=5, + failure_threshold=60, + ), + liveness_probe=_build_probe( + port=port, + initial_delay_seconds=600, + period_seconds=10, + timeout_seconds=5, + failure_threshold=3, + ), + ) + + pod_labels = dict(labels) + pod_metadata = k8s_client.V1ObjectMeta(labels=pod_labels) + pod_spec = k8s_client.V1PodSpec(containers=[container], volumes=volumes) + + merged_resources = dict(resources) + if pod_settings: + # The service config's resources take precedence over + # `pod_settings.resources` per top-level key. The merged result, + # with GPU limits mirrored into requests, is applied to the + # container once below, after `add_pod_settings` has applied the + # rest of the pod settings. + for key in ("requests", "limits"): + if pod_settings.resources.get(key): + merged_resources[key] = { + **merged_resources.get(key, {}), + **pod_settings.resources[key], + } + + add_pod_settings( + pod_spec=pod_spec, + settings=pod_settings, + substitutions={"{{ image }}": image}, + ) + if pod_settings.labels: + pod_labels.update(pod_settings.labels) + if pod_settings.annotations: + pod_metadata.annotations = pod_settings.annotations + + mirrored_resources = _resources_with_mirrored_gpu_requests( + merged_resources + ) + if mirrored_resources: + container.resources = k8s_client.V1ResourceRequirements( + requests=mirrored_resources.get("requests"), + limits=mirrored_resources.get("limits"), + ) + + deployment = k8s_client.V1Deployment( + api_version="apps/v1", + kind="Deployment", + metadata=k8s_client.V1ObjectMeta( + name=name, namespace=namespace, labels=labels + ), + spec=k8s_client.V1DeploymentSpec( + replicas=replicas, + selector=k8s_client.V1LabelSelector(match_labels=labels), + template=k8s_client.V1PodTemplateSpec( + metadata=pod_metadata, spec=pod_spec + ), + ), + ) + + return cast( + Dict[str, Any], + k8s_client.ApiClient().sanitize_for_serialization(deployment), + ) + + +def build_vllm_service_manifest( + name: str, + namespace: str, + port: int, + service_type: KubernetesServiceType, + selector_labels: Dict[str, str], + labels: Dict[str, str], +) -> Dict[str, Any]: + """Build the manifest for the vLLM Service. + + Args: + name: Name of the Service. + namespace: Namespace to deploy into. + port: Port exposed by the Service and forwarded to the pod. + service_type: Kubernetes Service type. + selector_labels: Labels used to select the target pods. + labels: Labels applied to the Service itself. + + Returns: + The Service manifest. + """ + service = k8s_client.V1Service( + api_version="v1", + kind="Service", + metadata=k8s_client.V1ObjectMeta( + name=name, namespace=namespace, labels=labels + ), + spec=k8s_client.V1ServiceSpec( + type=service_type, + selector=selector_labels, + ports=[k8s_client.V1ServicePort(port=port, target_port=port)], + ), + ) + return cast( + Dict[str, Any], + k8s_client.ApiClient().sanitize_for_serialization(service), + ) + + +class VLLMKubernetesServiceConfig(VLLMEngineArgs, ServiceConfig): + """vLLM Kubernetes service configuration. + + Attributes: + namespace: Kubernetes namespace to deploy into. Resolved from the + deployer configuration when unset. + image: Container image running the vLLM OpenAI server. Resolved + from the deployer configuration when unset. + port: Port the vLLM OpenAI server listens on inside the container, + and the port exposed by the Service. + service_type: Type of Kubernetes Service used to expose the vLLM + server. Resolved from the deployer configuration when unset. + replicas: Number of pod replicas for the Deployment. + resources: Native Kubernetes resource requests and limits for the + container. GPU limits are mirrored into requests automatically, + since Kubernetes requires the two to match for extended + resources. + pod_settings: Additional pod configuration merged into the + generated pod spec, for volumes, affinity, tolerations, and + other settings not modeled as dedicated fields. + hf_token: Hugging Face access token injected into the container as + `HF_TOKEN` through a managed secret. + existing_hf_secret: Name of an existing Kubernetes secret to use + instead of `hf_token`. Must contain the token under a `token` + key. + env: Additional environment variables set on the vLLM container. + extra_serve_args: Additional CLI args appended to the vLLM OpenAI + server args, for engine options not modeled as dedicated + fields. + shm_size: Size limit of the `/dev/shm` volume mounted into the + container. No volume is mounted if unset. + """ + + namespace: Optional[str] = Field( + default=None, + description="Kubernetes namespace to deploy into.", + ) + image: Optional[str] = Field( + default=None, + description="Container image running the vLLM OpenAI server.", + ) + port: int = Field( + default=8000, + description="Port the vLLM OpenAI server listens on.", + ) + service_type: Optional[KubernetesServiceType] = Field( + default=None, + description="Type of Kubernetes Service used to expose the vLLM server.", + ) + replicas: int = Field( + default=1, + description="Number of pod replicas for the Deployment.", + ) + resources: Dict[str, Dict[str, str]] = Field( + default_factory=dict, + description="Native Kubernetes resource requests and limits for the container.", + ) + pod_settings: Optional[KubernetesPodSettings] = Field( + default=None, + description="Additional pod configuration merged into the generated pod spec.", + ) + hf_token: Optional[str] = Field( + default=None, + description="Hugging Face access token injected into the container through a managed secret.", + ) + existing_hf_secret: Optional[str] = Field( + default=None, + description="Name of an existing Kubernetes secret to use instead of `hf_token`.", + ) + env: Dict[str, str] = Field( + default_factory=dict, + description="Additional environment variables set on the vLLM container.", + ) + extra_serve_args: List[str] = Field( + default_factory=list, + description="Additional CLI args appended to the vLLM OpenAI server args.", + ) + shm_size: Optional[str] = Field( + default="2Gi", + description="Size limit of the `/dev/shm` volume mounted into the container.", + ) + + +class VLLMKubernetesServiceStatus(ServiceStatus): + """vLLM Kubernetes service status.""" + + +class VLLMKubernetesDeploymentService(BaseDeploymentService): + """A service that represents a vLLM inference server running on Kubernetes. + + Attributes: + config: service configuration. + status: service status. + """ + + SERVICE_TYPE = ServiceType( + name="vllm-kubernetes-deployment", + type="model-serving", + flavor="vllm-kubernetes", + description="vLLM Inference prediction service running on Kubernetes", + ) + + config: VLLMKubernetesServiceConfig + status: VLLMKubernetesServiceStatus = Field( + default_factory=lambda: VLLMKubernetesServiceStatus() + ) + + def _get_deployer(self) -> "KubernetesVLLMModelDeployer": + """Get the active Kubernetes vLLM model deployer. + + Returns: + The active Kubernetes vLLM model deployer. + """ + from zenml.integrations.vllm.model_deployers.vllm_kubernetes_model_deployer import ( + KubernetesVLLMModelDeployer, + ) + + return cast( + "KubernetesVLLMModelDeployer", + KubernetesVLLMModelDeployer.get_active_model_deployer(), + ) + + @property + def namespace(self) -> str: + """Kubernetes namespace this service is deployed to. + + Returns: + The configured namespace, or the deployer's configured + `kubernetes_namespace` if unset. + """ + return ( + self.config.namespace + or self._get_deployer().config.kubernetes_namespace + ) + + @property + def image(self) -> str: + """Container image running the vLLM OpenAI server. + + Returns: + The configured image, the deployer's configured `default_image` + if unset, or `DEFAULT_VLLM_IMAGE` if no deployer is resolvable. + """ + if self.config.image: + return self.config.image + try: + return self._get_deployer().config.default_image + except TypeError: + return DEFAULT_VLLM_IMAGE + + @property + def service_type(self) -> KubernetesServiceType: + """Kubernetes Service type used to expose the vLLM server. + + Returns: + The configured Service type, the deployer's configured + `default_service_type` if unset, or + `KubernetesServiceType.CLUSTER_IP` if no deployer is resolvable. + """ + if self.config.service_type: + return self.config.service_type + try: + return self._get_deployer().config.default_service_type + except TypeError: + return KubernetesServiceType.CLUSTER_IP + + @property + def deployment_name(self) -> str: + """Name of the Kubernetes Deployment for this service. + + Returns: + The sanitized Deployment name. + """ + return kube_utils.sanitize_label(f"vllm-{self.uuid}") + + @property + def service_name(self) -> str: + """Name of the Kubernetes Service for this service. + + Returns: + The sanitized Service name. + """ + return self.deployment_name + + @property + def hf_secret_name(self) -> str: + """Name of the managed secret holding the Hugging Face token. + + Returns: + The sanitized secret name. + """ + return kube_utils.sanitize_label(f"vllm-hf-{self.uuid}") + + @property + def _labels(self) -> Dict[str, str]: + """Labels identifying the Kubernetes resources owned by this service. + + Returns: + The label dictionary. + """ + return { + "managed-by": "zenml", + "zenml-service-uuid": str(self.uuid), + } + + def _resolve_hf_secret_name(self) -> Optional[str]: + """Resolve the name of the secret holding the Hugging Face token. + + Returns: + The secret name, or `None` if neither `hf_token` nor + `existing_hf_secret` is configured. + """ + if self.config.hf_token: + return self.hf_secret_name + return self.config.existing_hf_secret + + def provision(self) -> None: + """Provision or update the vLLM Kubernetes Deployment and Service.""" + deployer = self._get_deployer() + core_api = k8s_client.CoreV1Api(deployer.get_kube_client()) + + kube_utils.create_namespace( + core_api=core_api, namespace=self.namespace + ) + + if self.config.hf_token: + kube_utils.create_or_update_secret( + core_api=core_api, + namespace=self.namespace, + secret_name=self.hf_secret_name, + data={VLLM_HF_TOKEN_SECRET_KEY: self.config.hf_token}, + ) + + labels = self._labels + manifests = [ + build_vllm_deployment_manifest( + name=self.deployment_name, + namespace=self.namespace, + image=self.image, + replicas=self.config.replicas, + port=self.config.port, + engine_args=self.config, + extra_serve_args=self.config.extra_serve_args, + labels=labels, + resources=self.config.resources, + shm_size=self.config.shm_size, + env=self.config.env, + hf_secret_name=self._resolve_hf_secret_name(), + pod_settings=self.config.pod_settings, + ), + build_vllm_service_manifest( + name=self.service_name, + namespace=self.namespace, + port=self.config.port, + service_type=self.service_type, + selector_labels=labels, + labels=labels, + ), + ] + + deployer.k8s_applier.provision( + manifests, + default_namespace=self.namespace, + field_manager=VLLM_FIELD_MANAGER, + force=True, + ) + + def _list_pods( + self, core_api: k8s_client.CoreV1Api + ) -> List[k8s_client.V1Pod]: + """List the pods belonging to this service's Deployment. + + Args: + core_api: Kubernetes Core V1 API client. + + Returns: + The pods matching the service's label selector. + """ + label_selector = ",".join( + f"{key}={value}" for key, value in self._labels.items() + ) + pod_list = core_api.list_namespaced_pod( + namespace=self.namespace, label_selector=label_selector + ) + return cast(List[k8s_client.V1Pod], pod_list.items) + + def _get_pod_failure_message( + self, pods: List[k8s_client.V1Pod] + ) -> Optional[str]: + """Look for pod-level failure conditions among the deployment's pods. + + Args: + pods: The deployment's pods. + + Returns: + A failure message if a pod is crash-looping or failing to pull + its image. `None` otherwise. + """ + for pod in pods: + container_name = _pod_container_name(pod) + container_status = kube_utils.get_container_status( + pod, container_name + ) + if ( + container_status + and container_status.waiting + and container_status.waiting.reason in POD_FAILURE_REASONS + ): + return kube_utils.get_pod_failure_details(pod, container_name) + + return None + + def _get_pod_pending_message( + self, pods: List[k8s_client.V1Pod] + ) -> Optional[str]: + """Look for pod-level pending details among the deployment's pods. + + Args: + pods: The deployment's pods. + + Returns: + Pending details, such as an `Unschedulable` condition, if a pod + is not yet running. `None` otherwise. + """ + for pod in pods: + container_name = _pod_container_name(pod) + pending_message = kube_utils.get_pod_pending_details( + pod, container_name + ) + if pending_message: + return pending_message + + return None + + def check_status(self) -> Tuple[ServiceState, str]: + """Check the current operational state of the vLLM Kubernetes Deployment. + + Returns: + The operational state of the Deployment and a message providing + additional information about that state. + """ + deployer = self._get_deployer() + applier = deployer.k8s_applier + + deployment = applier.get_resource( + name=self.deployment_name, + namespace=self.namespace, + kind="Deployment", + api_version="apps/v1", + ) + if deployment is None: + return (ServiceState.INACTIVE, "") + + deployment_dict = normalize_resource_to_dict(deployment) + status = deployment_dict.get("status") or {} + spec = deployment_dict.get("spec") or {} + desired_replicas = spec.get("replicas", 0) + available_replicas = status.get("availableReplicas", 0) + + conditions = { + condition.get("type"): condition + for condition in status.get("conditions") or [] + } + progressing = conditions.get("Progressing", {}) + if progressing.get("reason") == "ProgressDeadlineExceeded": + return ( + ServiceState.ERROR, + progressing.get( + "message", + f"Deployment '{self.deployment_name}' rollout stuck.", + ), + ) + + core_api = k8s_client.CoreV1Api(deployer.get_kube_client()) + pods = self._list_pods(core_api) + pod_failure_message = self._get_pod_failure_message(pods) + if pod_failure_message: + return (ServiceState.ERROR, pod_failure_message) + + if desired_replicas and available_replicas >= desired_replicas: + return ( + ServiceState.ACTIVE, + f"vLLM deployment '{self.deployment_name}' is available", + ) + + # An unschedulable pod may still be scheduled once the cluster + # autoscaler adds a node, so this stays PENDING_STARTUP instead of + # ERROR. ProgressDeadlineExceeded above is the signal for a rollout + # that is no longer transient. + pod_pending_message = self._get_pod_pending_message(pods) + return ( + ServiceState.PENDING_STARTUP, + pod_pending_message + or f"vLLM deployment '{self.deployment_name}' is starting up", + ) + + def deprovision(self, force: bool = False) -> None: + """Deprovision the vLLM Kubernetes Deployment, Service, and secret. + + Args: + force: if True, the remote deployment instance will be + forcefully deprovisioned. + """ + deployer = self._get_deployer() + + inventory = [ + ResourceInventoryItem( + api_version="v1", + kind="Service", + namespace=self.namespace, + name=self.service_name, + ), + ResourceInventoryItem( + api_version="apps/v1", + kind="Deployment", + namespace=self.namespace, + name=self.deployment_name, + ), + ] + if self.config.hf_token: + inventory.append( + ResourceInventoryItem( + api_version="v1", + kind="Secret", + namespace=self.namespace, + name=self.hf_secret_name, + ) + ) + + deployer.k8s_applier.delete_from_inventory( + inventory=inventory, propagation_policy="Foreground" + ) + + def get_logs( + self, + follow: bool = False, + tail: Optional[int] = None, + ) -> Generator[str, bool, None]: + """Get the logs of the vLLM Kubernetes deployment. + + Args: + follow: if True, the logs will be streamed as they are written. + tail: only retrieve the last NUM lines of log output. + + Yields: + A generator that can be accessed to get the service logs. + """ + deployer = self._get_deployer() + core_api = k8s_client.CoreV1Api(deployer.get_kube_client()) + + pods = self._list_pods(core_api) + if not pods: + return + + pod_name = pods[0].metadata.name + + if follow: + response = core_api.read_namespaced_pod_log( + name=pod_name, + namespace=self.namespace, + follow=True, + tail_lines=tail, + _preload_content=False, + ) + for line in response: + yield ( + line.decode("utf-8").rstrip() + if isinstance(line, bytes) + else str(line).rstrip() + ) + else: + response = core_api.read_namespaced_pod_log( + name=pod_name, + namespace=self.namespace, + tail_lines=tail, + _preload_content=False, + ) + logs = response.data.decode("utf-8") + for line in logs.split("\n"): + if line: + yield line + + def _resolve_service_url(self) -> Optional[str]: + """Resolve the base URL of the underlying Kubernetes Service. + + Returns: + The base URL, or `None` if the Service resource does not exist + yet or has no resolvable address. + """ + deployer = self._get_deployer() + service = deployer.k8s_applier.get_resource( + name=self.service_name, + namespace=self.namespace, + kind="Service", + api_version="v1", + ) + if service is None: + return None + + core_api = k8s_client.CoreV1Api(deployer.get_kube_client()) + return kube_utils.build_service_url( + core_api=core_api, service=service, namespace=self.namespace + ) + + @property + def service_url(self) -> Optional[str]: + """Base URL of the underlying Kubernetes Service. + + Returns: + The base URL, or `None` if the service is not running or the + Service resource has no resolvable address. + """ + if not self.is_running: + return None + return self._resolve_service_url() + + @property + def prediction_url(self) -> Optional[str]: + """The prediction URI exposed by the vLLM Kubernetes service. + + Returns: + The prediction URI, or `None` if the service is not yet ready. + """ + base_url = self.service_url + if base_url is None: + return None + return os.path.join(base_url, VLLM_PREDICTION_URL_PATH) + + @property + def healthcheck_url(self) -> Optional[str]: + """The healthcheck URI exposed by the vLLM Kubernetes service. + + Returns: + The healthcheck URI, or `None` if the service is not yet ready. + """ + base_url = self.service_url + if base_url is None: + return None + return os.path.join(base_url, VLLM_HEALTHCHECK_URL_PATH) + + def predict(self, data: Any) -> Any: + """Make a prediction using the service. + + Args: + data: Prompt to run inference on. + + Raises: + Exception: if the service is not running. + + Returns: + The prediction result. + """ + prediction_url = self.prediction_url + if prediction_url is None: + raise Exception( + "vLLM Inference service is not running. " + "Please start the service before making predictions." + ) + + from openai import OpenAI + + client = OpenAI(api_key="EMPTY", base_url=prediction_url) + models = client.models.list() + model = models.data[0].id + return client.completions.create(model=model, prompt=data) diff --git a/tests/unit/integrations/vllm/__init__.py b/tests/unit/integrations/vllm/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/integrations/vllm/test_vllm_kubernetes_flavor.py b/tests/unit/integrations/vllm/test_vllm_kubernetes_flavor.py new file mode 100644 index 00000000000..f3415974275 --- /dev/null +++ b/tests/unit/integrations/vllm/test_vllm_kubernetes_flavor.py @@ -0,0 +1,64 @@ +# Copyright (c) ZenML GmbH 2026. All Rights Reserved. +# +# 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: +# +# https://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. +"""Unit tests for the vLLM Kubernetes model deployer flavor.""" + +from zenml.constants import KUBERNETES_CLUSTER_RESOURCE_TYPE +from zenml.enums import KubernetesServiceType +from zenml.integrations.vllm.flavors.vllm_kubernetes_model_deployer_flavor import ( + KubernetesVLLMModelDeployerConfig, + KubernetesVLLMModelDeployerFlavor, +) +from zenml.integrations.vllm.model_deployers.vllm_kubernetes_model_deployer import ( + KubernetesVLLMModelDeployer, +) + + +def test_flavor_name(): + """The flavor name is vllm-kubernetes.""" + flavor = KubernetesVLLMModelDeployerFlavor() + assert flavor.name == "vllm-kubernetes" + + +def test_flavor_config_and_implementation_classes(): + """The flavor exposes the Kubernetes vLLM config and deployer classes.""" + flavor = KubernetesVLLMModelDeployerFlavor() + assert flavor.config_class is KubernetesVLLMModelDeployerConfig + assert flavor.implementation_class is KubernetesVLLMModelDeployer + + +def test_flavor_service_connector_requirements_target_kubernetes_cluster(): + """The flavor requires a kubernetes-cluster service connector resource.""" + flavor = KubernetesVLLMModelDeployerFlavor() + requirements = flavor.service_connector_requirements + + assert requirements is not None + assert requirements.resource_type == KUBERNETES_CLUSTER_RESOURCE_TYPE + + +def test_config_defaults(): + """The config defaults match the documented Kubernetes vLLM defaults.""" + config = KubernetesVLLMModelDeployerConfig() + + assert config.kubernetes_namespace == "zenml-vllm" + assert config.default_image == "vllm/vllm-openai:v0.25.1" + assert config.incluster is False + assert config.default_service_type == KubernetesServiceType.LOAD_BALANCER + assert config.kubernetes_context is None + assert config.hf_token is None + + +def test_config_is_not_local(): + """The Kubernetes vLLM deployer is not a local stack component.""" + config = KubernetesVLLMModelDeployerConfig() + assert config.is_local is False diff --git a/tests/unit/integrations/vllm/test_vllm_kubernetes_manifests.py b/tests/unit/integrations/vllm/test_vllm_kubernetes_manifests.py new file mode 100644 index 00000000000..3854d53df66 --- /dev/null +++ b/tests/unit/integrations/vllm/test_vllm_kubernetes_manifests.py @@ -0,0 +1,303 @@ +# Copyright (c) ZenML GmbH 2026. All Rights Reserved. +# +# 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: +# +# https://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. +"""Unit tests for the vLLM Kubernetes manifest builders.""" + +import json + +from zenml.enums import KubernetesServiceType +from zenml.integrations.vllm.services.vllm_deployment import VLLMEngineArgs +from zenml.integrations.vllm.services.vllm_kubernetes_deployment import ( + build_vllm_deployment_manifest, + build_vllm_service_manifest, +) + +LABELS = {"managed-by": "zenml", "zenml-service-uuid": "abc"} + + +def _build_deployment(**overrides): + """Build a Deployment manifest with sensible defaults. + + Args: + overrides: Keyword arguments overriding the defaults. + + Returns: + The Deployment manifest. + """ + kwargs = dict( + name="vllm-test", + namespace="zenml-vllm", + image="vllm/vllm-openai:v0.25.1", + replicas=1, + port=8000, + engine_args=VLLMEngineArgs(model="facebook/opt-125m"), + extra_serve_args=[], + labels=LABELS, + resources={}, + shm_size=None, + env={}, + hf_secret_name=None, + pod_settings=None, + ) + kwargs.update(overrides) + return build_vllm_deployment_manifest(**kwargs) + + +def _container(manifest): + """Extract the vLLM container spec from a Deployment manifest. + + Args: + manifest: A Deployment manifest. + + Returns: + The container spec. + """ + return manifest["spec"]["template"]["spec"]["containers"][0] + + +def test_container_args_from_engine_fields(): + """Engine fields translate to their corresponding CLI args.""" + engine_args = VLLMEngineArgs( + model="facebook/opt-125m", + tokenizer="facebook/opt-125m-tok", + served_model_name=["opt", "opt-alias"], + trust_remote_code=True, + tokenizer_mode="slow", + dtype="float16", + revision="main", + ) + manifest = _build_deployment( + engine_args=engine_args, extra_serve_args=["--foo", "bar"] + ) + args = _container(manifest)["args"] + + assert args == [ + "--model", + "facebook/opt-125m", + "--port", + "8000", + "--host", + "0.0.0.0", + "--tokenizer", + "facebook/opt-125m-tok", + "--served-model-name", + "opt", + "opt-alias", + "--trust-remote-code", + "--tokenizer-mode", + "slow", + "--dtype", + "float16", + "--revision", + "main", + "--foo", + "bar", + ] + + +def test_container_args_minimal_engine_fields(): + """Only the required model field and engine defaults produce args.""" + manifest = _build_deployment( + engine_args=VLLMEngineArgs(model="facebook/opt-125m"), port=9000 + ) + args = _container(manifest)["args"] + + assert args == [ + "--model", + "facebook/opt-125m", + "--port", + "9000", + "--host", + "0.0.0.0", + "--tokenizer-mode", + "auto", + "--dtype", + "auto", + ] + + +def test_served_model_name_as_single_string(): + """A single served_model_name is passed as one CLI value.""" + engine_args = VLLMEngineArgs( + model="facebook/opt-125m", served_model_name="opt" + ) + manifest = _build_deployment(engine_args=engine_args) + args = _container(manifest)["args"] + + assert args[args.index("--served-model-name") + 1] == "opt" + + +def test_probes_use_health_path_and_configured_port(): + """Readiness and liveness probes hit /health on the configured port.""" + manifest = _build_deployment(port=8888) + container = _container(manifest) + + for probe_key in ("readinessProbe", "livenessProbe"): + probe = container[probe_key] + assert probe["httpGet"]["path"] == "/health" + assert probe["httpGet"]["port"] == 8888 + + +def test_readiness_and_liveness_probes_differ_in_thresholds(): + """The liveness probe tolerates a much longer startup than readiness.""" + manifest = _build_deployment() + container = _container(manifest) + + assert container["readinessProbe"]["initialDelaySeconds"] == 10 + assert container["readinessProbe"]["failureThreshold"] == 60 + assert container["livenessProbe"]["initialDelaySeconds"] == 600 + assert container["livenessProbe"]["failureThreshold"] == 3 + + +def test_gpu_limit_mirrored_into_requests(): + """A configured GPU limit is mirrored into requests.""" + manifest = _build_deployment( + resources={"limits": {"nvidia.com/gpu": "1", "cpu": "2"}} + ) + resources = _container(manifest)["resources"] + + assert resources["limits"] == {"nvidia.com/gpu": "1", "cpu": "2"} + assert resources["requests"] == {"nvidia.com/gpu": "1"} + + +def test_existing_gpu_request_is_not_overwritten(): + """An explicit GPU request is left untouched, not replaced by the limit.""" + manifest = _build_deployment( + resources={ + "limits": {"nvidia.com/gpu": "2"}, + "requests": {"nvidia.com/gpu": "1"}, + } + ) + resources = _container(manifest)["resources"] + + assert resources["requests"]["nvidia.com/gpu"] == "1" + + +def test_no_gpu_limit_leaves_resources_unmirrored(): + """Resources without a GPU limit are passed through unchanged.""" + manifest = _build_deployment(resources={"limits": {"cpu": "2"}}) + resources = _container(manifest)["resources"] + + assert "requests" not in resources + + +def test_shm_volume_added_with_memory_medium_and_size_limit(): + """A configured shm_size adds a Memory-backed /dev/shm emptyDir volume.""" + manifest = _build_deployment(shm_size="4Gi") + pod_spec = manifest["spec"]["template"]["spec"] + + assert pod_spec["volumes"] == [ + {"name": "shm", "emptyDir": {"medium": "Memory", "sizeLimit": "4Gi"}} + ] + assert _container(manifest)["volumeMounts"] == [ + {"name": "shm", "mountPath": "/dev/shm"} + ] + + +def test_no_shm_volume_when_shm_size_unset(): + """No volume is added when shm_size is unset.""" + manifest = _build_deployment(shm_size=None) + pod_spec = manifest["spec"]["template"]["spec"] + + assert pod_spec["volumes"] == [] + assert _container(manifest)["volumeMounts"] == [] + + +def test_hf_token_secret_key_ref_added_when_secret_name_set(): + """An HF_TOKEN env var referencing the secret is added when configured.""" + manifest = _build_deployment(hf_secret_name="vllm-hf-abc") + env = _container(manifest)["env"] + hf_token_env = next(item for item in env if item["name"] == "HF_TOKEN") + + assert hf_token_env["valueFrom"]["secretKeyRef"] == { + "name": "vllm-hf-abc", + "key": "token", + } + + +def test_hf_token_env_absent_when_secret_name_unset(): + """No HF_TOKEN env var is added when no secret name is passed.""" + manifest = _build_deployment(hf_secret_name=None) + env = _container(manifest)["env"] + + assert not any(item["name"] == "HF_TOKEN" for item in env) + + +def test_custom_env_vars_included_alongside_hf_token(): + """Custom env vars and the HF_TOKEN secret ref coexist.""" + manifest = _build_deployment( + env={"FOO": "BAR"}, hf_secret_name="vllm-hf-abc" + ) + env = _container(manifest)["env"] + + assert {"name": "FOO", "value": "BAR"} in env + assert any(item["name"] == "HF_TOKEN" for item in env) + + +def test_deployment_manifest_is_json_serializable(): + """The Deployment manifest can be JSON-serialized as-is.""" + manifest = _build_deployment( + resources={"limits": {"nvidia.com/gpu": "1"}}, + shm_size="2Gi", + env={"FOO": "BAR"}, + hf_secret_name="vllm-hf-abc", + ) + + json.dumps(manifest) + + +def test_service_manifest_selector_matches_deployment_pod_labels(): + """The Service selector matches the Deployment's pod template labels.""" + deployment_manifest = _build_deployment() + service_manifest = build_vllm_service_manifest( + name="vllm-test", + namespace="zenml-vllm", + port=8000, + service_type=KubernetesServiceType.CLUSTER_IP, + selector_labels=LABELS, + labels=LABELS, + ) + + pod_labels = deployment_manifest["spec"]["template"]["metadata"]["labels"] + assert service_manifest["spec"]["selector"] == pod_labels + + +def test_service_manifest_exposes_configured_port_and_type(): + """The Service exposes the configured port and Service type.""" + service_manifest = build_vllm_service_manifest( + name="vllm-test", + namespace="zenml-vllm", + port=8000, + service_type=KubernetesServiceType.LOAD_BALANCER, + selector_labels=LABELS, + labels=LABELS, + ) + + assert service_manifest["spec"]["type"] == "LoadBalancer" + assert service_manifest["spec"]["ports"] == [ + {"port": 8000, "targetPort": 8000} + ] + + +def test_service_manifest_is_json_serializable(): + """The Service manifest can be JSON-serialized as-is.""" + service_manifest = build_vllm_service_manifest( + name="vllm-test", + namespace="zenml-vllm", + port=8000, + service_type=KubernetesServiceType.CLUSTER_IP, + selector_labels=LABELS, + labels=LABELS, + ) + + json.dumps(service_manifest) diff --git a/tests/unit/integrations/vllm/test_vllm_kubernetes_service.py b/tests/unit/integrations/vllm/test_vllm_kubernetes_service.py new file mode 100644 index 00000000000..e713ddce4ad --- /dev/null +++ b/tests/unit/integrations/vllm/test_vllm_kubernetes_service.py @@ -0,0 +1,545 @@ +# Copyright (c) ZenML GmbH 2026. All Rights Reserved. +# +# 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: +# +# https://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. +"""Unit tests for the vLLM Kubernetes deployment service.""" + +from unittest.mock import MagicMock, patch +from uuid import uuid4 + +from kubernetes import client as k8s_client + +from zenml.enums import KubernetesServiceType, ServiceState +from zenml.integrations.vllm.flavors.vllm_kubernetes_model_deployer_flavor import ( + DEFAULT_VLLM_IMAGE, +) +from zenml.integrations.vllm.model_deployers.vllm_kubernetes_model_deployer import ( + KubernetesVLLMModelDeployer, +) +from zenml.integrations.vllm.services.vllm_kubernetes_deployment import ( + VLLMKubernetesDeploymentService, + VLLMKubernetesServiceConfig, +) + + +def _build_service(**config_overrides): + """Build a service with a minimal Kubernetes config. + + Args: + config_overrides: Keyword arguments overriding the config defaults. + + Returns: + The service. + """ + config_kwargs = dict( + name="vllm-test", + model="facebook/opt-125m", + namespace="zenml-vllm", + port=8000, + ) + config_kwargs.update(config_overrides) + return VLLMKubernetesDeploymentService( + uuid=uuid4(), config=VLLMKubernetesServiceConfig(**config_kwargs) + ) + + +def _mock_deployer(deployment=None): + """Build a deployer mock whose applier returns the given Deployment. + + Args: + deployment: Deployment resource returned by `get_resource`. + + Returns: + The deployer mock. + """ + deployer = MagicMock() + deployer.k8s_applier.get_resource.return_value = deployment + return deployer + + +def test_check_status_inactive_when_deployment_missing(): + """A missing Deployment maps to INACTIVE.""" + service = _build_service() + deployer = _mock_deployer(deployment=None) + + with patch.object( + KubernetesVLLMModelDeployer, + "get_active_model_deployer", + return_value=deployer, + ): + state, message = service.check_status() + + assert state == ServiceState.INACTIVE + assert message == "" + + +def test_check_status_active_when_available_replicas_meet_desired(): + """Available replicas meeting the desired count maps to ACTIVE.""" + service = _build_service() + deployment = { + "spec": {"replicas": 2}, + "status": {"availableReplicas": 2, "conditions": []}, + } + deployer = _mock_deployer(deployment=deployment) + + with ( + patch.object( + KubernetesVLLMModelDeployer, + "get_active_model_deployer", + return_value=deployer, + ), + patch.object(service, "_list_pods", return_value=[]), + patch.object(service, "_get_pod_failure_message", return_value=None), + ): + state, _ = service.check_status() + + assert state == ServiceState.ACTIVE + + +def test_check_status_error_on_progress_deadline_exceeded(): + """A stuck rollout maps to ERROR with the condition's message.""" + service = _build_service() + deployment = { + "spec": {"replicas": 2}, + "status": { + "availableReplicas": 0, + "conditions": [ + { + "type": "Progressing", + "reason": "ProgressDeadlineExceeded", + "message": "deployment exceeded its progress deadline", + } + ], + }, + } + deployer = _mock_deployer(deployment=deployment) + + with patch.object( + KubernetesVLLMModelDeployer, + "get_active_model_deployer", + return_value=deployer, + ): + state, message = service.check_status() + + assert state == ServiceState.ERROR + assert message == "deployment exceeded its progress deadline" + + +def test_check_status_pending_startup_when_replicas_not_yet_available(): + """Fewer available replicas than desired maps to PENDING_STARTUP.""" + service = _build_service() + deployment = { + "spec": {"replicas": 2}, + "status": {"availableReplicas": 1, "conditions": []}, + } + deployer = _mock_deployer(deployment=deployment) + + with ( + patch.object( + KubernetesVLLMModelDeployer, + "get_active_model_deployer", + return_value=deployer, + ), + patch.object(service, "_list_pods", return_value=[]), + patch.object(service, "_get_pod_failure_message", return_value=None), + patch.object(service, "_get_pod_pending_message", return_value=None), + ): + state, _ = service.check_status() + + assert state == ServiceState.PENDING_STARTUP + + +def test_check_status_error_on_crash_loop_backoff(): + """A crash-looping pod maps to ERROR with the failure details.""" + service = _build_service() + deployment = { + "spec": {"replicas": 1}, + "status": {"availableReplicas": 0, "conditions": []}, + } + deployer = _mock_deployer(deployment=deployment) + + with ( + patch.object( + KubernetesVLLMModelDeployer, + "get_active_model_deployer", + return_value=deployer, + ), + patch.object(service, "_list_pods", return_value=[]), + patch.object( + service, + "_get_pod_failure_message", + return_value="container waiting reason: CrashLoopBackOff", + ), + ): + state, message = service.check_status() + + assert state == ServiceState.ERROR + assert message == "container waiting reason: CrashLoopBackOff" + + +def test_check_status_pending_startup_on_unschedulable_pod(): + """An unschedulable pod maps to PENDING_STARTUP, not ERROR. + + An unschedulable pod is indistinguishable at observation time from the + cluster autoscaler being about to add a node that fits it. The + unschedulable detail is still surfaced in the status message. + """ + service = _build_service() + deployment = { + "spec": {"replicas": 1}, + "status": {"availableReplicas": 0, "conditions": []}, + } + deployer = _mock_deployer(deployment=deployment) + pending_message = ( + "pod `vllm-test-0`: pod condition: PodScheduled, Unschedulable, " + "message=0/3 nodes are available: 3 Insufficient nvidia.com/gpu" + ) + + with ( + patch.object( + KubernetesVLLMModelDeployer, + "get_active_model_deployer", + return_value=deployer, + ), + patch.object(service, "_list_pods", return_value=[]), + patch.object(service, "_get_pod_failure_message", return_value=None), + patch.object( + service, + "_get_pod_pending_message", + return_value=pending_message, + ), + ): + state, message = service.check_status() + + assert state == ServiceState.PENDING_STARTUP + assert message == pending_message + + +def test_list_pods_uses_label_selector_and_namespace(): + """Pods are listed with the service's namespace and label selector.""" + service = _build_service() + core_api = MagicMock() + core_api.list_namespaced_pod.return_value = k8s_client.V1PodList(items=[]) + + service._list_pods(core_api) + + core_api.list_namespaced_pod.assert_called_once_with( + namespace="zenml-vllm", + label_selector=f"managed-by=zenml,zenml-service-uuid={service.uuid}", + ) + + +def test_get_pod_failure_message_detects_crash_loop_backoff(): + """A crash-looping container is surfaced as a failure message.""" + service = _build_service() + pod = k8s_client.V1Pod( + spec=k8s_client.V1PodSpec( + containers=[k8s_client.V1Container(name="vllm", image="img")] + ), + status=k8s_client.V1PodStatus( + container_statuses=[ + k8s_client.V1ContainerStatus( + name="vllm", + image="img", + image_id="", + ready=False, + restart_count=1, + state=k8s_client.V1ContainerState( + waiting=k8s_client.V1ContainerStateWaiting( + reason="CrashLoopBackOff", + ) + ), + ) + ] + ), + ) + + message = service._get_pod_failure_message([pod]) + + assert message is not None + assert "CrashLoopBackOff" in message + + +def test_get_pod_failure_message_none_when_no_pods_failing(): + """No failure message is returned when no pod is crash-looping.""" + service = _build_service() + pod = k8s_client.V1Pod( + spec=k8s_client.V1PodSpec( + containers=[k8s_client.V1Container(name="vllm", image="img")] + ), + status=k8s_client.V1PodStatus(container_statuses=[]), + ) + + assert service._get_pod_failure_message([pod]) is None + + +def test_get_pod_pending_message_detects_unschedulable_condition(): + """An unschedulable pod condition is surfaced as a pending message.""" + service = _build_service() + pod = k8s_client.V1Pod( + status=k8s_client.V1PodStatus( + conditions=[ + k8s_client.V1PodCondition( + type="PodScheduled", + status="False", + reason="Unschedulable", + message="0/3 nodes are available", + ) + ] + ), + ) + + message = service._get_pod_pending_message([pod]) + + assert message is not None + assert "Unschedulable" in message + + +def test_get_pod_pending_message_none_when_no_pods_pending(): + """No pending message is returned when no pod has a pending condition.""" + service = _build_service() + + assert service._get_pod_pending_message([]) is None + + +def test_deployment_and_service_names_are_deterministic_and_sanitized(): + """Resource names are derived from the service UUID and sanitized.""" + service_uuid = uuid4() + service = _build_service() + service.uuid = service_uuid + + assert service.deployment_name == f"vllm-{service_uuid}" + assert service.service_name == f"vllm-{service_uuid}" + assert service.hf_secret_name == f"vllm-hf-{service_uuid}" + + +def test_namespace_falls_back_to_deployer_namespace_when_unset(): + """Reading namespace without a configured value falls back to the deployer's namespace.""" + service = VLLMKubernetesDeploymentService( + uuid=uuid4(), + config=VLLMKubernetesServiceConfig( + name="vllm-test", model="facebook/opt-125m", port=8000 + ), + ) + deployer = MagicMock() + deployer.config.kubernetes_namespace = "deployer-namespace" + + with patch.object( + KubernetesVLLMModelDeployer, + "get_active_model_deployer", + return_value=deployer, + ): + assert service.namespace == "deployer-namespace" + + +def test_namespace_uses_configured_value_when_set(): + """Reading namespace with a configured value returns it as-is.""" + service = _build_service() + + assert service.namespace == "zenml-vllm" + + +def test_service_type_falls_back_to_deployer_service_type_when_unset(): + """Reading service_type without a configured value falls back to the deployer's default_service_type.""" + service = _build_service() + deployer = MagicMock() + deployer.config.default_service_type = KubernetesServiceType.NODE_PORT + + with patch.object( + KubernetesVLLMModelDeployer, + "get_active_model_deployer", + return_value=deployer, + ): + assert service.service_type == KubernetesServiceType.NODE_PORT + + +def test_service_type_falls_back_to_cluster_ip_when_no_deployer_resolvable(): + """Reading service_type without a configured value or resolvable deployer falls back to CLUSTER_IP.""" + service = _build_service() + + with patch.object( + KubernetesVLLMModelDeployer, + "get_active_model_deployer", + side_effect=TypeError("no active model deployer"), + ): + assert service.service_type == KubernetesServiceType.CLUSTER_IP + + +def test_service_type_uses_configured_value_when_set(): + """Reading service_type with a configured value returns it as-is.""" + service = _build_service(service_type=KubernetesServiceType.LOAD_BALANCER) + + assert service.service_type == KubernetesServiceType.LOAD_BALANCER + + +def test_image_falls_back_to_deployer_default_image_when_unset(): + """Reading image without a configured value falls back to the deployer's default_image.""" + service = _build_service() + deployer = MagicMock() + deployer.config.default_image = "vllm/deployer-image:latest" + + with patch.object( + KubernetesVLLMModelDeployer, + "get_active_model_deployer", + return_value=deployer, + ): + assert service.image == "vllm/deployer-image:latest" + + +def test_image_falls_back_to_default_constant_when_no_deployer_resolvable(): + """Reading image without a configured value or resolvable deployer falls back to DEFAULT_VLLM_IMAGE.""" + service = _build_service() + + with patch.object( + KubernetesVLLMModelDeployer, + "get_active_model_deployer", + side_effect=TypeError("no active model deployer"), + ): + assert service.image == DEFAULT_VLLM_IMAGE + + +def test_image_uses_configured_value_when_set(): + """Reading image with a configured value returns it as-is.""" + service = _build_service(image="vllm/custom-image:latest") + + assert service.image == "vllm/custom-image:latest" + + +def test_service_url_resolves_base_url_when_running(): + """The base URL is resolved when the service is running.""" + service = _build_service() + + with ( + patch.object( + VLLMKubernetesDeploymentService, + "check_status", + return_value=(ServiceState.ACTIVE, ""), + ), + patch.object( + service, + "_resolve_service_url", + return_value="http://1.2.3.4:8000", + ), + ): + assert service.service_url == "http://1.2.3.4:8000" + + +def test_service_url_none_when_service_not_running(): + """No base URL is resolved while the Deployment is not active.""" + service = _build_service() + + with patch.object( + VLLMKubernetesDeploymentService, + "check_status", + return_value=(ServiceState.INACTIVE, ""), + ): + assert service.service_url is None + + +def test_prediction_url_composes_base_url_with_prediction_path(): + """The prediction URL appends the OpenAI-compatible path to the base URL.""" + service = _build_service() + + with ( + patch.object( + VLLMKubernetesDeploymentService, + "check_status", + return_value=(ServiceState.ACTIVE, ""), + ), + patch.object( + service, + "_resolve_service_url", + return_value="http://1.2.3.4:8000", + ), + ): + assert service.prediction_url == "http://1.2.3.4:8000/v1" + + +def test_healthcheck_url_composes_base_url_with_health_path(): + """The healthcheck URL appends the health path to the base URL.""" + service = _build_service() + + with ( + patch.object( + VLLMKubernetesDeploymentService, + "check_status", + return_value=(ServiceState.ACTIVE, ""), + ), + patch.object( + service, + "_resolve_service_url", + return_value="http://1.2.3.4:8000", + ), + ): + assert service.healthcheck_url == "http://1.2.3.4:8000/health" + + +def test_prediction_url_none_when_service_not_running(): + """No prediction URL is exposed while the Deployment is not active.""" + service = _build_service() + + with patch.object( + VLLMKubernetesDeploymentService, + "check_status", + return_value=(ServiceState.INACTIVE, ""), + ): + assert service.prediction_url is None + + +def test_prediction_url_none_when_service_url_not_resolvable(): + """No prediction URL is exposed when the Service has no resolvable address.""" + service = _build_service() + + with ( + patch.object( + VLLMKubernetesDeploymentService, + "check_status", + return_value=(ServiceState.ACTIVE, ""), + ), + patch.object(service, "_resolve_service_url", return_value=None), + ): + assert service.prediction_url is None + + +def test_resolve_service_url_none_when_service_resource_missing(): + """No base URL is resolved when the Service resource does not exist.""" + service = _build_service() + deployer = _mock_deployer(deployment=None) + + with patch.object( + KubernetesVLLMModelDeployer, + "get_active_model_deployer", + return_value=deployer, + ): + assert service._resolve_service_url() is None + + +def test_resolve_service_url_delegates_to_build_service_url(): + """The base URL is built via kube_utils.build_service_url.""" + service = _build_service() + deployer = _mock_deployer(deployment={"spec": {"type": "ClusterIP"}}) + + with ( + patch.object( + KubernetesVLLMModelDeployer, + "get_active_model_deployer", + return_value=deployer, + ), + patch( + "zenml.integrations.vllm.services.vllm_kubernetes_deployment." + "kube_utils.build_service_url", + return_value="http://10.0.0.1:8000", + ) as build_service_url_mock, + ): + assert service._resolve_service_url() == "http://10.0.0.1:8000" + + build_service_url_mock.assert_called_once() diff --git a/tests/unit/integrations/vllm/test_vllm_service_configs.py b/tests/unit/integrations/vllm/test_vllm_service_configs.py new file mode 100644 index 00000000000..332c21e46c8 --- /dev/null +++ b/tests/unit/integrations/vllm/test_vllm_service_configs.py @@ -0,0 +1,86 @@ +# Copyright (c) ZenML GmbH 2026. All Rights Reserved. +# +# 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: +# +# https://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. +"""Unit tests for the vLLM service configs. + +`VLLMEngineArgs` was extracted as a shared mixin for the local +`VLLMServiceConfig` and the Kubernetes `VLLMKubernetesServiceConfig`. These +tests pin down the field defaults on both configs so the extraction did not +change either one's behavior. +""" + +from zenml.integrations.vllm.services.vllm_deployment import VLLMServiceConfig +from zenml.integrations.vllm.services.vllm_kubernetes_deployment import ( + VLLMKubernetesServiceConfig, +) + + +def test_local_service_config_engine_defaults(): + """VLLMServiceConfig keeps its engine field defaults after the mixin extraction.""" + config = VLLMServiceConfig( + name="vllm-test", model="facebook/opt-125m", port=8000 + ) + + assert config.model == "facebook/opt-125m" + assert config.tokenizer is None + assert config.served_model_name is None + assert config.trust_remote_code is False + assert config.tokenizer_mode == "auto" + assert config.dtype == "auto" + assert config.revision is None + + +def test_local_service_config_service_defaults(): + """VLLMServiceConfig keeps its own non-engine field defaults.""" + config = VLLMServiceConfig( + name="vllm-test", model="facebook/opt-125m", port=8000 + ) + + assert config.port == 8000 + assert config.host is None + assert config.blocking is True + + +def test_kubernetes_service_config_engine_defaults(): + """VLLMKubernetesServiceConfig inherits the same engine field defaults.""" + config = VLLMKubernetesServiceConfig( + name="vllm-test", model="facebook/opt-125m" + ) + + assert config.model == "facebook/opt-125m" + assert config.tokenizer is None + assert config.served_model_name is None + assert config.trust_remote_code is False + assert config.tokenizer_mode == "auto" + assert config.dtype == "auto" + assert config.revision is None + + +def test_kubernetes_service_config_defaults(): + """VLLMKubernetesServiceConfig field defaults match the plan.""" + config = VLLMKubernetesServiceConfig( + name="vllm-test", model="facebook/opt-125m" + ) + + assert config.port == 8000 + assert config.replicas == 1 + assert config.shm_size == "2Gi" + assert config.namespace is None + assert config.image is None + assert config.service_type is None + assert config.resources == {} + assert config.pod_settings is None + assert config.hf_token is None + assert config.existing_hf_secret is None + assert config.env == {} + assert config.extra_serve_args == []