add support for deployment to K8s

.github/workflows/check-test-release.yml

@@ -92,6 +92,9 @@ jobs:
         pip install pre-commit .[tests]
     - run: pre-commit run pylint -a -v --show-diff-on-failure
       if: matrix.python != '3.7'
+    - name: Start minikube
+      if: matrix.os == 'ubuntu-latest' && matrix.python == '3.9'
+      uses: medyagh/setup-minikube@master
     - name: Run tests
       timeout-minutes: 40
       run: pytest

mlem/contrib/kubernetes/base.py

@@ -0,0 +1,219 @@
+import os
+from typing import ClassVar, List, Optional
+
+from kubernetes import client, config
+
+from mlem.config import project_config
+from mlem.core.errors import DeploymentError, EndpointNotFound, MlemError
+from mlem.core.objects import (
+    DeployState,
+    DeployStatus,
+    MlemBuilder,
+    MlemDeployment,
+    MlemEnv,
+    MlemModel,
+)
+from mlem.runtime.client import Client, HTTPClient
+from mlem.runtime.server import Server
+from mlem.ui import EMOJI_OK, echo
+
+from ..docker.base import (
+    DockerDaemon,
+    DockerImage,
+    DockerRegistry,
+    generate_docker_container_name,
+)
+from .build import build_k8s_docker
+from .context import K8sYamlBuildArgs, K8sYamlGenerator
+from .utils import create_k8s_resources, namespace_deleted, pod_is_running
+
+POD_STATE_MAPPING = {
+    "Pending": DeployStatus.STARTING,
+    "Running": DeployStatus.RUNNING,
+    "Succeeded": DeployStatus.STOPPED,
+    "Failed": DeployStatus.CRASHED,
+    "Unknown": DeployStatus.UNKNOWN,
+}
+
+
+class K8sDeploymentState(DeployState):
+    """DeployState implementation for Kubernetes deployments"""
+
+    type: ClassVar = "kubernetes"
+
+    image: Optional[DockerImage] = None
+    """Docker Image being used for Deployment"""
+    deployment_name: Optional[str] = None
+    """Name of Deployment"""
+
+
+class K8sDeployment(MlemDeployment, K8sYamlBuildArgs):
+    """MlemDeployment implementation for Kubernetes deployments"""
+
+    type: ClassVar = "kubernetes"
+    state_type: ClassVar = K8sDeploymentState
+    """Type of state for Kubernetes deployments"""
+
+    server: Optional[Server] = None
+    """Type of Server to use, with options such as FastAPI, RabbitMQ etc."""
+    registry: Optional[DockerRegistry] = DockerRegistry()
+    """Docker registry"""
+    daemon: Optional[DockerDaemon] = DockerDaemon(host="")
+    """Docker daemon"""
+    kube_config_file_path: Optional[str] = None
+    """Path for kube config file of the cluster"""
+    templates_dir: List[str] = []
+    """List of dirs where templates reside"""
+
+    def load_kube_config(self):
+        config.load_kube_config(
+            config_file=self.kube_config_file_path
+            or os.getenv("KUBECONFIG", default="~/.kube/config")
+        )
+
+    def _get_client(self, state: K8sDeploymentState) -> Client:
+        host, port = None, None
+        self.load_kube_config()
+        service = client.CoreV1Api().list_namespaced_service(self.namespace)
+        try:
+            host, port = self.service_type.get_host_and_port(
+                service, self.namespace
+            )
+        except MlemError as e:
+            raise EndpointNotFound(
+                "Couldn't determine host and port from the service deployed"
+            ) from e
+        if host is not None and port is not None:
+            return HTTPClient(host=host, port=port)
+        raise MlemError(
+            f"host and port determined are not valid, received host as {host} and port as {port}"
+        )
+
+
+class K8sEnv(MlemEnv[K8sDeployment]):
+    """MlemEnv implementation for Kubernetes Environments"""
+
+    type: ClassVar = "kubernetes"
+    deploy_type: ClassVar = K8sDeployment
+    """Type of deployment being used for the Kubernetes environment"""
+
+    registry: Optional[DockerRegistry] = None
+    """Docker registry"""
+    templates_dir: List[str] = []
+    """List of dirs where templates reside"""
+
+    def get_registry(self, meta: K8sDeployment):
+        registry = meta.registry or self.registry
+        if not registry:
+            raise MlemError(
+                "registry to be used by Docker is not set or supplied"
+            )
+        return registry
+
+    def get_image_name(self, meta: K8sDeployment):
+        return meta.image_name or generate_docker_container_name()
+
+    def get_server(self, meta: K8sDeployment):
+        return (
+            meta.server
+            or project_config(
+                meta.loc.project if meta.is_saved else None
+            ).server
+        )
+
+    def deploy(self, meta: K8sDeployment):
+        self.check_type(meta)
+        redeploy = False
+        with meta.lock_state():
+            meta.load_kube_config()
+            state: K8sDeploymentState = meta.get_state()
+            if state.image is None or meta.model_changed():
+                image_name = self.get_image_name(meta)
+                state.image = build_k8s_docker(
+                    meta=meta.get_model(),
+                    image_name=image_name,
+                    registry=self.get_registry(meta),
+                    daemon=meta.daemon,
+                    server=self.get_server(meta),
+                )
+                meta.update_model_hash(state=state)
+                redeploy = True
+
+            if (
+                state.deployment_name is None or redeploy
+            ) and state.image is not None:
+                generator = K8sYamlGenerator(
+                    namespace=meta.namespace,
+                    image_name=state.image.name,
+                    image_uri=state.image.uri,
+                    image_pull_policy=meta.image_pull_policy,
+                    port=meta.port,
+                    service_type=meta.service_type,
+                    templates_dir=meta.templates_dir or self.templates_dir,
+                )
+                create_k8s_resources(generator)
+
+                if pod_is_running(namespace=meta.namespace):
+                    deployments_list = (
+                        client.AppsV1Api().list_namespaced_deployment(
+                            namespace=meta.namespace
+                        )
+                    )
+
+                    if len(deployments_list.items) == 0:
+                        raise DeploymentError(
+                            f"Deployment {image_name} couldn't be found in {meta.namespace} namespace"
+                        )
+                    dpl_name = deployments_list.items[0].metadata.name
+                    state.deployment_name = dpl_name
+                    meta.update_state(state)
+
+                    echo(
+                        EMOJI_OK
+                        + f"Deployment {state.deployment_name} is up in {meta.namespace} namespace"
+                    )
+                else:
+                    raise DeploymentError(
+                        f"Deployment {image_name} couldn't be set-up on the Kubernetes cluster"
+                    )
+
+    def remove(self, meta: K8sDeployment):
+        self.check_type(meta)
+        with meta.lock_state():
+            meta.load_kube_config()
+            state: K8sDeploymentState = meta.get_state()
+            if state.deployment_name is not None:
+                client.CoreV1Api().delete_namespace(name=meta.namespace)
+                if namespace_deleted(meta.namespace):
+                    echo(
+                        EMOJI_OK
+                        + f"Deployment {state.deployment_name} and the corresponding service are removed from {meta.namespace} namespace"
+                    )
+                    state.deployment_name = None
+                    meta.update_state(state)
+
+    def get_status(
+        self, meta: K8sDeployment, raise_on_error=True
+    ) -> DeployStatus:
+        self.check_type(meta)
+        meta.load_kube_config()
+        state: K8sDeploymentState = meta.get_state()
+        if state.deployment_name is None:
+            return DeployStatus.NOT_DEPLOYED
+
+        pods_list = client.CoreV1Api().list_namespaced_pod(meta.namespace)
+
+        return POD_STATE_MAPPING[pods_list.items[0].status.phase]
+
+
+class K8sYamlBuilder(MlemBuilder, K8sYamlGenerator):
+    """MlemBuilder implementation for building Kubernetes manifests/yamls"""
+
+    type: ClassVar = "kubernetes"
+
+    target: str
+    """Target path for the manifest/yaml"""
+
+    def build(self, obj: MlemModel):
+        self.write(self.target)
+        echo(EMOJI_OK + f"{self.target} generated for {obj.basename}")

mlem/contrib/kubernetes/build.py

@@ -0,0 +1,30 @@
+from typing import Optional
+
+from mlem.core.objects import MlemModel
+from mlem.runtime.server import Server
+from mlem.ui import EMOJI_BUILD, echo, set_offset
+
+from ..docker.base import DockerDaemon, DockerEnv, DockerRegistry
+from ..docker.helpers import build_model_image
+
+
+def build_k8s_docker(
+    meta: MlemModel,
+    image_name: str,
+    registry: Optional[DockerRegistry],
+    daemon: Optional[DockerDaemon],
+    server: Server,
+    platform: Optional[str] = "linux/amd64",
+    # runners usually do not support arm64 images built on Mac M1 devices
+):
+    echo(EMOJI_BUILD + f"Creating docker image {image_name}")
+    with set_offset(2):
+        return build_model_image(
+            meta,
+            image_name,
+            server,
+            DockerEnv(registry=registry, daemon=daemon),
+            tag=meta.meta_hash(),
+            force_overwrite=True,
+            platform=platform,
+        )

mlem/contrib/kubernetes/context.py

@@ -0,0 +1,55 @@
+import logging
+import os
+from enum import Enum
+from typing import ClassVar
+
+from pydantic import BaseModel
+
+from mlem.contrib.kubernetes.service import NodePortService, ServiceType
+from mlem.utils.templates import TemplateModel
+
+logger = logging.getLogger(__name__)
+
+
+class ImagePullPolicy(str, Enum):
+    always = "Always"
+    never = "Never"
+    if_not_present = "IfNotPresent"
+
+
+class K8sYamlBuildArgs(BaseModel):
+    """Class encapsulating parameters for Kubernetes manifests/yamls"""
+
+    class Config:
+        use_enum_values = True
+
+    namespace: str = "mlem"
+    """Namespace to create kubernetes resources such as pods, service in"""
+    image_name: str = "ml"
+    """Name of the docker image to be deployed"""
+    image_uri: str = "ml:latest"
+    """URI of the docker image to be deployed"""
+    image_pull_policy: ImagePullPolicy = ImagePullPolicy.always
+    """Image pull policy for the docker image to be deployed"""
+    port: int = 8080
+    """Port where the service should be available"""
+    service_type: ServiceType = NodePortService()
+    """Type of service by which endpoints of the model are exposed"""
+
+
+class K8sYamlGenerator(K8sYamlBuildArgs, TemplateModel):
+    TEMPLATE_FILE: ClassVar = "resources.yaml.j2"
+    TEMPLATE_DIR: ClassVar = os.path.dirname(__file__)
+
+    def prepare_dict(self):
+        logger.debug(
+            'Generating Resource Yaml via templates from "%s"...',
+            self.templates_dir,
+        )
+
+        logger.debug('Docker image is based on "%s".', self.image_uri)
+
+        k8s_yaml_args = self.dict()
+        k8s_yaml_args["service_type"] = self.service_type.get_string()
+        k8s_yaml_args.pop("templates_dir")
+        return k8s_yaml_args

mlem/contrib/kubernetes/resources.yaml.j2

@@ -0,0 +1,47 @@
+apiVersion: v1
+kind: Namespace
+metadata:
+  name: {{ namespace }}
+  labels:
+    name: {{ namespace }}
+
+---
+
+apiVersion: apps/v1
+kind: Deployment
+metadata:
+  name: {{ image_name }}
+  namespace: {{ namespace }}
+spec:
+  selector:
+    matchLabels:
+      app: {{ image_name }}
+  template:
+    metadata:
+      labels:
+        app: {{ image_name }}
+    spec:
+      containers:
+      - name: {{ image_name }}
+        image: {{ image_uri }}
+        imagePullPolicy: {{ image_pull_policy }}
+        ports:
+        - containerPort: {{ port }}
+
+---
+
+apiVersion: v1
+kind: Service
+metadata:
+  name: {{ image_name }}
+  namespace: {{ namespace }}
+  labels:
+    run: {{ image_name }}
+spec:
+  ports:
+  - port: {{ port }}
+    protocol: TCP
+    targetPort: {{ port }}
+  selector:
+    app: {{ image_name }}
+  type: {{ service_type }}