zenml-nightly 0.55.0.dev20240124__py3-none-any.whl → 0.72.0.dev20250116__py3-none-any.whl

zenml/integrations/mlflow/services/mlflow_deployment.py

@@ -14,6 +14,7 @@
 """Implementation of the MLflow deployment functionality."""
 
 import os
+import sys
 from typing import TYPE_CHECKING, Any, Dict, Optional, Union
 
 import numpy as np
@@ -21,6 +22,7 @@ import pandas as pd
 import requests
 from mlflow.pyfunc.backend import PyFuncBackend
 from mlflow.version import VERSION as MLFLOW_VERSION
+from pydantic import field_validator
 
 from zenml.client import Client
 from zenml.constants import DEFAULT_SERVICE_START_STOP_TIMEOUT
@@ -101,8 +103,6 @@ class MLFlowDeploymentConfig(LocalDaemonServiceConfig):
         timeout: timeout in seconds for starting and stopping the service
     """
 
-    # TODO: ServiceConfig should have additional fields such as "pipeline_run_uuid"
-    #  and "pipeline_uuid" to allow for better tracking of the service.
     model_uri: str
     model_name: str
     registry_model_name: Optional[str] = None
@@ -112,6 +112,30 @@ class MLFlowDeploymentConfig(LocalDaemonServiceConfig):
     mlserver: bool = False
     timeout: int = DEFAULT_SERVICE_START_STOP_TIMEOUT
 
+    @field_validator("mlserver")
+    @classmethod
+    def validate_mlserver_python_version(cls, mlserver: bool) -> bool:
+        """Validates the Python version if mlserver is used.
+
+        Args:
+            mlserver: set to True if the MLflow MLServer backend is used,
+                else set to False and MLflow built-in scoring server will be
+                used.
+
+        Returns:
+            the validated value
+
+        Raises:
+            ValueError: if mlserver is set to true on Python 3.12 as it is not
+                yet supported.
+        """
+        if mlserver is True and sys.version_info.minor >= 12:
+            raise ValueError(
+                "The mlserver deployment is not yet supported on Python 3.12 "
+                "or above."
+            )
+        return mlserver
+
 
 class MLFlowDeploymentService(LocalDaemonService, BaseDeploymentService):
     """MLflow deployment service used to start a local prediction server for MLflow models.
@@ -128,6 +152,7 @@ class MLFlowDeploymentService(LocalDaemonService, BaseDeploymentService):
         type="model-serving",
         flavor="mlflow",
         description="MLflow prediction service",
+        logo_url="https://public-flavor-logos.s3.eu-central-1.amazonaws.com/model_deployer/mlflow.png",
     )
 
     config: MLFlowDeploymentConfig
@@ -200,7 +225,7 @@ class MLFlowDeploymentService(LocalDaemonService, BaseDeploymentService):
             # to run the deploy the model on the local running environment
             if int(mlflow_version[0]) >= 2:
                 backend_kwargs["env_manager"] = "local"
-            backend = PyFuncBackend(
+            backend = PyFuncBackend(  # type: ignore[no-untyped-call]
                 config={},
                 no_conda=True,
                 workers=self.config.workers,
@@ -215,7 +240,7 @@ class MLFlowDeploymentService(LocalDaemonService, BaseDeploymentService):
                     "stack."
                 )
             experiment_tracker.configure_mlflow()
-            backend.serve(
+            backend.serve(  # type: ignore[no-untyped-call]
                 model_uri=self.config.model_uri,
                 port=self.endpoint.status.port,
                 host="localhost",
@@ -261,7 +286,7 @@ class MLFlowDeploymentService(LocalDaemonService, BaseDeploymentService):
             )
 
         if self.endpoint.prediction_url is not None:
-            if type(request) == pd.DataFrame:
+            if type(request) is pd.DataFrame:
                 response = requests.post(  # nosec
                     self.endpoint.prediction_url,
                     json={"instances": request.to_dict("records")},

zenml/integrations/mlflow/steps/mlflow_deployer.py

@@ -114,36 +114,34 @@ def mlflow_model_deployer_step(
     # Fetch the model URI from the MLflow artifact store
     model_uri = ""
     if mlflow_run_id and client.list_artifacts(mlflow_run_id, model_name):
-        model_uri = artifact_utils.get_artifact_uri(
+        model_uri = artifact_utils.get_artifact_uri(  # type: ignore[no-untyped-call]
             run_id=mlflow_run_id, artifact_path=model_name
         )
 
-    # Fetch existing services with same pipeline name, step name and model name
-    existing_services = model_deployer.find_model_server(
+    predictor_cfg = MLFlowDeploymentConfig(
+        model_name=model_name or "",
+        model_uri=model_uri,
+        workers=workers,
+        mlserver=mlserver,
         pipeline_name=pipeline_name,
         pipeline_step_name=step_name,
-        model_name=model_name,
+        timeout=timeout,
+    )
+
+    # Fetch existing services with same pipeline name, step name and model name
+    existing_services = model_deployer.find_model_server(
+        config=predictor_cfg.model_dump(),
     )
 
     # Check whether to deploy a new service
     if model_uri and deploy_decision:
-        predictor_cfg = MLFlowDeploymentConfig(
-            model_name=model_name or "",
-            model_uri=model_uri,
-            workers=workers,
-            mlserver=mlserver,
-            pipeline_name=pipeline_name,
-            run_name=run_name,
-            pipeline_run_id=run_name,
-            pipeline_step_name=step_name,
-            timeout=timeout,
-        )
         new_service = cast(
             MLFlowDeploymentService,
             model_deployer.deploy_model(
                 replace=True,
                 config=predictor_cfg,
                 timeout=timeout,
+                service_type=MLFlowDeploymentService.SERVICE_TYPE,
             ),
         )
         logger.info(
@@ -168,7 +166,7 @@ def mlflow_model_deployer_step(
                 "was negative."
             )
         logger.info(
-            f"Reusing last model server deployed by step '{step_name}' of "
+            f"Reusing last model server deployed by step `{step_name}` of "
             f"pipeline '{pipeline_name}' for model '{model_name}'..."
         )
         service = cast(MLFlowDeploymentService, existing_services[0])
@@ -228,10 +226,10 @@ def mlflow_model_registry_deployer_step(
     """
     if not registry_model_version and not registry_model_stage:
         raise ValueError(
-            "Either registry_model_version or registry_model_stage must"
-            "be provided in addition to registry_model_name to the MLflow"
-            "model registry deployer step. Since the"
-            "mlflow_model_registry_deployer_step is used in conjunction with"
+            "Either registry_model_version or registry_model_stage must "
+            "be provided in addition to registry_model_name to the MLflow "
+            "model registry deployer step. Since the "
+            "mlflow_model_registry_deployer_step is used in conjunction with "
             "the mlflow_model_registry."
         )
 
@@ -277,26 +275,25 @@ def mlflow_model_registry_deployer_step(
             f"using this step."
         )
     # fetch existing services with same pipeline name, step name and model name
+
     existing_services = (
         model_deployer.find_model_server(
-            registry_model_name=model_version.registered_model.name,
+            model_name=registry_model_name,
+            model_version=model_version.version,
         )
         if replace_existing
         else None
     )
-
     # create a config for the new model service
     metadata = model_version.metadata or ModelRegistryModelMetadata()
     predictor_cfg = MLFlowDeploymentConfig(
-        model_name=model_name or "",
+        name=model_name or None,
+        model_name=registry_model_name,
+        model_version=model_version.version,
         model_uri=model_version.model_source_uri,
-        registry_model_name=model_version.registered_model.name,
-        registry_model_version=model_version.version,
-        registry_model_stage=model_version.stage.value,
         workers=workers,
         mlserver=mlserver,
         pipeline_name=metadata.zenml_pipeline_name or "",
-        run_name=metadata.zenml_run_name or "",
         pipeline_step_name=metadata.zenml_step_name or "",
         timeout=timeout,
     )
@@ -308,6 +305,7 @@ def mlflow_model_registry_deployer_step(
             replace=True,
             config=predictor_cfg,
             timeout=timeout,
+            service_type=MLFlowDeploymentService.SERVICE_TYPE,
         ),
     )
 

zenml/integrations/mlflow/steps/mlflow_registry.py

@@ -124,7 +124,7 @@ def mlflow_register_model_step(
     if not model_source_uri and client.list_artifacts(
         mlflow_run_id, trained_model_name
     ):
-        model_source_uri = artifact_utils.get_artifact_uri(
+        model_source_uri = artifact_utils.get_artifact_uri(  # type: ignore[no-untyped-call]
             run_id=mlflow_run_id, artifact_path=trained_model_name
         )
     if not model_source_uri:
@@ -146,6 +146,8 @@ def mlflow_register_model_step(
         metadata.zenml_pipeline_run_uuid = pipeline_run_uuid
     if metadata.zenml_workspace is None:
         metadata.zenml_workspace = zenml_workspace
+    if getattr(metadata, "mlflow_run_id", None) is None:
+        setattr(metadata, "mlflow_run_id", mlflow_run_id)
 
     # Register model version
     model_version = model_registry.register_model_version(

zenml/integrations/modal/__init__.py

@@ -0,0 +1,46 @@
+#  Copyright (c) ZenML GmbH 2024. 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.
+"""Modal integration for cloud-native step execution.
+
+The Modal integration sub-module provides a step operator flavor that allows
+executing steps on Modal's cloud infrastructure.
+"""
+from typing import List, Type
+
+from zenml.integrations.constants import MODAL
+from zenml.integrations.integration import Integration
+from zenml.stack import Flavor
+
+MODAL_STEP_OPERATOR_FLAVOR = "modal"
+
+
+class ModalIntegration(Integration):
+    """Definition of Modal integration for ZenML."""
+
+    NAME = MODAL
+    REQUIREMENTS = ["modal>=0.64.49,<1"]
+
+    @classmethod
+    def flavors(cls) -> List[Type[Flavor]]:
+        """Declare the stack component flavors for the Modal integration.
+
+        Returns:
+            List of new stack component flavors.
+        """
+        from zenml.integrations.modal.flavors import ModalStepOperatorFlavor
+
+        return [ModalStepOperatorFlavor]
+
+
+ModalIntegration.check_installation()

zenml/integrations/modal/flavors/__init__.py

@@ -0,0 +1,26 @@
+#  Copyright (c) ZenML GmbH 2024. 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.
+"""Modal integration flavors."""
+
+from zenml.integrations.modal.flavors.modal_step_operator_flavor import (
+    ModalStepOperatorConfig,
+    ModalStepOperatorFlavor,
+    ModalStepOperatorSettings,
+)
+
+__all__ = [
+    "ModalStepOperatorConfig",
+    "ModalStepOperatorFlavor",
+    "ModalStepOperatorSettings",
+]

zenml/integrations/modal/flavors/modal_step_operator_flavor.py

@@ -0,0 +1,125 @@
+#  Copyright (c) ZenML GmbH 2024. 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.
+"""Modal step operator flavor."""
+
+from typing import TYPE_CHECKING, Optional, Type
+
+from zenml.config.base_settings import BaseSettings
+from zenml.integrations.modal import MODAL_STEP_OPERATOR_FLAVOR
+from zenml.step_operators import BaseStepOperatorConfig, BaseStepOperatorFlavor
+
+if TYPE_CHECKING:
+    from zenml.integrations.modal.step_operators import ModalStepOperator
+
+
+class ModalStepOperatorSettings(BaseSettings):
+    """Settings for the Modal step operator.
+
+    Specifying the region and cloud provider is only available for Enterprise
+    and Team plan customers.
+
+    Certain combinations of settings are not available. It is suggested to err
+    on the side of looser settings rather than more restrictive ones to avoid
+    pipeline execution failures. In the case of failures, however, Modal
+    provides detailed error messages that can help identify what is
+    incompatible. See more in the Modal docs at https://modal.com/docs/guide/region-selection.
+
+    Attributes:
+        gpu: The type of GPU to use for the step execution.
+        region: The region to use for the step execution.
+        cloud: The cloud provider to use for the step execution.
+    """
+
+    gpu: Optional[str] = None
+    region: Optional[str] = None
+    cloud: Optional[str] = None
+
+
+class ModalStepOperatorConfig(
+    BaseStepOperatorConfig, ModalStepOperatorSettings
+):
+    """Configuration for the Modal step operator."""
+
+    @property
+    def is_remote(self) -> bool:
+        """Checks if this stack component is running remotely.
+
+        This designation is used to determine if the stack component can be
+        used with a local ZenML database or if it requires a remote ZenML
+        server.
+
+        Returns:
+            True if this config is for a remote component, False otherwise.
+        """
+        return True
+
+
+class ModalStepOperatorFlavor(BaseStepOperatorFlavor):
+    """Modal step operator flavor."""
+
+    @property
+    def name(self) -> str:
+        """Name of the flavor.
+
+        Returns:
+            The name of the flavor.
+        """
+        return MODAL_STEP_OPERATOR_FLAVOR
+
+    @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/step_operator/modal.png"
+
+    @property
+    def config_class(self) -> Type[ModalStepOperatorConfig]:
+        """Returns `ModalStepOperatorConfig` config class.
+
+        Returns:
+            The config class.
+        """
+        return ModalStepOperatorConfig
+
+    @property
+    def implementation_class(self) -> Type["ModalStepOperator"]:
+        """Implementation class for this flavor.
+
+        Returns:
+            The implementation class.
+        """
+        from zenml.integrations.modal.step_operators import ModalStepOperator
+
+        return ModalStepOperator

zenml/integrations/modal/step_operators/__init__.py

@@ -0,0 +1,22 @@
+#  Copyright (c) ZenML GmbH 2024. 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.
+"""Modal step operator."""
+
+from zenml.integrations.modal.step_operators.modal_step_operator import (
+    ModalStepOperator,
+)
+
+__all__ = [
+    "ModalStepOperator",
+]

zenml/integrations/modal/step_operators/modal_step_operator.py

@@ -0,0 +1,242 @@
+#  Copyright (c) ZenML GmbH 2024. 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.
+"""Modal step operator implementation."""
+
+import asyncio
+from typing import TYPE_CHECKING, Dict, List, Optional, Tuple, Type, cast
+
+import modal
+from modal_proto import api_pb2
+
+from zenml.client import Client
+from zenml.config.build_configuration import BuildConfiguration
+from zenml.config.resource_settings import ByteUnit, ResourceSettings
+from zenml.enums import StackComponentType
+from zenml.integrations.modal.flavors import (
+    ModalStepOperatorConfig,
+    ModalStepOperatorSettings,
+)
+from zenml.logger import get_logger
+from zenml.stack import Stack, StackValidator
+from zenml.step_operators import BaseStepOperator
+
+if TYPE_CHECKING:
+    from zenml.config.base_settings import BaseSettings
+    from zenml.config.step_run_info import StepRunInfo
+    from zenml.models import PipelineDeploymentBase
+
+logger = get_logger(__name__)
+
+MODAL_STEP_OPERATOR_DOCKER_IMAGE_KEY = "modal_step_operator"
+
+
+def get_gpu_values(
+    settings: ModalStepOperatorSettings, resource_settings: ResourceSettings
+) -> Optional[str]:
+    """Get the GPU values for the Modal step operator.
+
+    Args:
+        settings: The Modal step operator settings.
+        resource_settings: The resource settings.
+
+    Returns:
+        The GPU string if a count is specified, otherwise the GPU type.
+    """
+    if not settings.gpu:
+        return None
+    gpu_count = resource_settings.gpu_count
+    return f"{settings.gpu}:{gpu_count}" if gpu_count else settings.gpu
+
+
+class ModalStepOperator(BaseStepOperator):
+    """Step operator to run a step on Modal.
+
+    This class defines code that can set up a Modal environment and run
+    functions in it.
+    """
+
+    @property
+    def config(self) -> ModalStepOperatorConfig:
+        """Get the Modal step operator configuration.
+
+        Returns:
+            The Modal step operator configuration.
+        """
+        return cast(ModalStepOperatorConfig, self._config)
+
+    @property
+    def settings_class(self) -> Optional[Type["BaseSettings"]]:
+        """Get the settings class for the Modal step operator.
+
+        Returns:
+            The Modal step operator settings class.
+        """
+        return ModalStepOperatorSettings
+
+    @property
+    def validator(self) -> Optional[StackValidator]:
+        """Get the stack validator for the Modal step operator.
+
+        Returns:
+            The stack validator.
+        """
+
+        def _validate_remote_components(stack: "Stack") -> Tuple[bool, str]:
+            if stack.artifact_store.config.is_local:
+                return False, (
+                    "The Modal step operator runs code remotely and "
+                    "needs to write files into the artifact store, but the "
+                    f"artifact store `{stack.artifact_store.name}` of the "
+                    "active stack is local. Please ensure that your stack "
+                    "contains a remote artifact store when using the Modal "
+                    "step operator."
+                )
+
+            container_registry = stack.container_registry
+            assert container_registry is not None
+
+            if container_registry.config.is_local:
+                return False, (
+                    "The Modal step operator runs code remotely and "
+                    "needs to push/pull Docker images, but the "
+                    f"container registry `{container_registry.name}` of the "
+                    "active stack is local. Please ensure that your stack "
+                    "contains a remote container registry when using the "
+                    "Modal step operator."
+                )
+
+            return True, ""
+
+        return StackValidator(
+            required_components={
+                StackComponentType.CONTAINER_REGISTRY,
+                StackComponentType.IMAGE_BUILDER,
+            },
+            custom_validation_function=_validate_remote_components,
+        )
+
+    def get_docker_builds(
+        self, deployment: "PipelineDeploymentBase"
+    ) -> List["BuildConfiguration"]:
+        """Get the Docker build configurations for the Modal step operator.
+
+        Args:
+            deployment: The pipeline deployment.
+
+        Returns:
+            A list of Docker build configurations.
+        """
+        builds = []
+        for step_name, step in deployment.step_configurations.items():
+            if step.config.step_operator == self.name:
+                build = BuildConfiguration(
+                    key=MODAL_STEP_OPERATOR_DOCKER_IMAGE_KEY,
+                    settings=step.config.docker_settings,
+                    step_name=step_name,
+                )
+                builds.append(build)
+
+        return builds
+
+    def launch(
+        self,
+        info: "StepRunInfo",
+        entrypoint_command: List[str],
+        environment: Dict[str, str],
+    ) -> None:
+        """Launch a step run on Modal.
+
+        Args:
+            info: The step run information.
+            entrypoint_command: The entrypoint command for the step.
+            environment: The environment variables for the step.
+
+        Raises:
+            RuntimeError: If no Docker credentials are found for the container registry.
+            ValueError: If no container registry is found in the stack.
+        """
+        settings = cast(ModalStepOperatorSettings, self.get_settings(info))
+        image_name = info.get_image(key=MODAL_STEP_OPERATOR_DOCKER_IMAGE_KEY)
+        zc = Client()
+        stack = zc.active_stack
+
+        if not stack.container_registry:
+            raise ValueError(
+                "No Container registry found in the stack. "
+                "Please add a container registry and ensure "
+                "it is correctly configured."
+            )
+
+        if docker_creds := stack.container_registry.credentials:
+            docker_username, docker_password = docker_creds
+        else:
+            raise RuntimeError(
+                "No Docker credentials found for the container registry."
+            )
+
+        my_secret = modal.secret._Secret.from_dict(
+            {
+                "REGISTRY_USERNAME": docker_username,
+                "REGISTRY_PASSWORD": docker_password,
+            }
+        )
+
+        spec = modal.image.DockerfileSpec(
+            commands=[f"FROM {image_name}"], context_files={}
+        )
+
+        zenml_image = modal.Image._from_args(
+            dockerfile_function=lambda *_, **__: spec,
+            force_build=False,
+            image_registry_config=modal.image._ImageRegistryConfig(
+                api_pb2.REGISTRY_AUTH_TYPE_STATIC_CREDS, my_secret
+            ),
+        ).env(environment)
+
+        resource_settings = info.config.resource_settings
+        gpu_values = get_gpu_values(settings, resource_settings)
+
+        app = modal.App(
+            f"zenml-{info.run_name}-{info.step_run_id}-{info.pipeline_step_name}"
+        )
+
+        async def run_sandbox() -> asyncio.Future[None]:
+            loop = asyncio.get_event_loop()
+            future = loop.create_future()
+            with modal.enable_output():
+                async with app.run():
+                    memory_mb = resource_settings.get_memory(ByteUnit.MB)
+                    memory_int = (
+                        int(memory_mb) if memory_mb is not None else None
+                    )
+                    sb = await modal.Sandbox.create.aio(
+                        "bash",
+                        "-c",
+                        " ".join(entrypoint_command),
+                        image=zenml_image,
+                        gpu=gpu_values,
+                        cpu=resource_settings.cpu_count,
+                        memory=memory_int,
+                        cloud=settings.cloud,
+                        region=settings.region,
+                        app=app,
+                        timeout=86400,  # 24h, the max Modal allows
+                    )
+
+                    await sb.wait.aio()
+
+            future.set_result(None)
+            return future
+
+        asyncio.run(run_sandbox())