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

zenml/integrations/kserve/steps/kserve_step_utils.py

@@ -1,293 +0,0 @@
-#  Copyright (c) ZenML GmbH 2022. 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:
-#
-#       http://www.apache.org/licenses/LICENSE-2.0
-#
-#  Unless required by applicable law or agreed to in writing, software
-#  distributed under the License is distributed on an "AS IS" BASIS,
-#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express
-#  or implied. See the License for the specific language governing
-#  permissions and limitations under the License.
-"""Utility functions used by the KServe deployer step."""
-import os
-import re
-import tempfile
-from typing import List, Optional
-
-from model_archiver.model_packaging import package_model
-from model_archiver.model_packaging_utils import ModelExportUtils
-from pydantic import BaseModel
-
-from zenml.exceptions import ValidationError
-from zenml.integrations.kserve.services.kserve_deployment import (
-    KServeDeploymentConfig,
-)
-from zenml.integrations.kserve.steps.kserve_deployer import (
-    KServeDeployerStepParameters,
-)
-from zenml.io import fileio
-from zenml.utils import io_utils
-
-
-def is_valid_model_name(model_name: str) -> bool:
-    """Checks if the model name is valid.
-
-    Args:
-        model_name: the model name to check
-
-    Returns:
-        True if the model name is valid, False otherwise.
-    """
-    pattern = re.compile("^[a-z0-9-]+$")
-    return pattern.match(model_name) is not None
-
-
-def prepare_service_config(
-    model_uri: str,
-    output_artifact_uri: str,
-    params: KServeDeployerStepParameters,
-) -> KServeDeploymentConfig:
-    """Prepare the model files for model serving.
-
-    This function ensures that the model files are in the correct format
-    and file structure required by the KServe server implementation
-    used for model serving.
-
-    Args:
-        model_uri: the URI of the model artifact being served
-        output_artifact_uri: the URI of the output artifact
-        params: the KServe deployer step parameters
-
-    Returns:
-        The URL to the model is ready for serving.
-
-    Raises:
-        RuntimeError: if the model files cannot be prepared.
-        ValidationError: if the model name is invalid.
-    """
-    served_model_uri = os.path.join(output_artifact_uri, "kserve")
-    fileio.makedirs(served_model_uri)
-
-    # TODO [ENG-773]: determine how to formalize how models are organized into
-    #   folders and sub-folders depending on the model type/format and the
-    #   KServe protocol used to serve the model.
-
-    # TODO [ENG-791]: an auto-detect built-in KServe server implementation
-    #   from the model artifact type
-
-    # TODO [ENG-792]: validate the model artifact type against the
-    #   supported built-in KServe server implementations
-    if params.service_config.predictor == "tensorflow":
-        # the TensorFlow server expects model artifacts to be
-        # stored in numbered subdirectories, each representing a model
-        # version
-        served_model_uri = os.path.join(
-            served_model_uri,
-            params.service_config.predictor,
-            params.service_config.model_name,
-        )
-        fileio.makedirs(served_model_uri)
-        io_utils.copy_dir(model_uri, os.path.join(served_model_uri, "1"))
-    elif params.service_config.predictor == "sklearn":
-        # the sklearn server expects model artifacts to be
-        # stored in a file called model.joblib
-        model_uri = os.path.join(model_uri, "model")
-        if not fileio.exists(model_uri):
-            raise RuntimeError(
-                f"Expected sklearn model artifact was not found at "
-                f"{model_uri}"
-            )
-        served_model_uri = os.path.join(
-            served_model_uri,
-            params.service_config.predictor,
-            params.service_config.model_name,
-        )
-        fileio.makedirs(served_model_uri)
-        fileio.copy(model_uri, os.path.join(served_model_uri, "model.joblib"))
-    elif not is_valid_model_name(params.service_config.model_name):
-        raise ValidationError(
-            f"Model name '{params.service_config.model_name}' is invalid. "
-            f"The model name can only include lowercase alphanumeric "
-            "characters and hyphens. Please rename your model and try again."
-        )
-    else:
-        # default treatment for all other server implementations is to
-        # simply reuse the model from the artifact store path where it
-        # is originally stored
-        served_model_uri = os.path.join(
-            served_model_uri,
-            params.service_config.predictor,
-            params.service_config.model_name,
-        )
-        fileio.makedirs(served_model_uri)
-        fileio.copy(model_uri, served_model_uri)
-
-    service_config = params.service_config.copy()
-    service_config.model_uri = served_model_uri
-    return service_config
-
-
-def prepare_torch_service_config(
-    model_uri: str,
-    output_artifact_uri: str,
-    params: KServeDeployerStepParameters,
-) -> KServeDeploymentConfig:
-    """Prepare the PyTorch model files for model serving.
-
-    This function ensures that the model files are in the correct format
-    and file structure required by the KServe server implementation
-    used for model serving.
-
-    Args:
-        model_uri: the URI of the model artifact being served
-        output_artifact_uri: the URI of the output artifact
-        params: the KServe deployer step parameters
-
-    Returns:
-        The URL to the model is ready for serving.
-
-    Raises:
-        RuntimeError: if the model files cannot be prepared.
-    """
-    deployment_folder_uri = os.path.join(output_artifact_uri, "kserve")
-    served_model_uri = os.path.join(deployment_folder_uri, "model-store")
-    config_properties_uri = os.path.join(deployment_folder_uri, "config")
-    fileio.makedirs(served_model_uri)
-    fileio.makedirs(config_properties_uri)
-
-    if params.torch_serve_parameters is None:
-        raise RuntimeError("No torch serve parameters provided")
-    else:
-        # Create a temporary folder
-        temp_dir = tempfile.mkdtemp(prefix="zenml-pytorch-temp-")
-        tmp_model_uri = os.path.join(
-            str(temp_dir), f"{params.service_config.model_name}.pt"
-        )
-
-        # Copy from artifact store to temporary file
-        fileio.copy(f"{model_uri}/checkpoint.pt", tmp_model_uri)
-
-        torch_archiver_args = TorchModelArchiver(
-            model_name=params.service_config.model_name,
-            serialized_file=tmp_model_uri,
-            model_file=params.torch_serve_parameters.model_class,
-            handler=params.torch_serve_parameters.handler,
-            export_path=temp_dir,
-            version=params.torch_serve_parameters.model_version,
-        )
-
-        manifest = ModelExportUtils.generate_manifest_json(torch_archiver_args)
-        package_model(torch_archiver_args, manifest=manifest)
-
-        # Copy from temporary file to artifact store
-        archived_model_uri = os.path.join(
-            temp_dir, f"{params.service_config.model_name}.mar"
-        )
-        if not fileio.exists(archived_model_uri):
-            raise RuntimeError(
-                f"Expected torch archived model artifact was not found at "
-                f"{archived_model_uri}"
-            )
-
-        # Copy the torch model archive artifact to the model store
-        fileio.copy(
-            archived_model_uri,
-            os.path.join(
-                served_model_uri, f"{params.service_config.model_name}.mar"
-            ),
-        )
-
-        # Get or Generate the config file
-        if params.torch_serve_parameters.torch_config:
-            # Copy the torch model config to the model store
-            fileio.copy(
-                params.torch_serve_parameters.torch_config,
-                os.path.join(config_properties_uri, "config.properties"),
-            )
-        else:
-            # Generate the config file
-            config_file_uri = generate_model_deployer_config(
-                model_name=params.service_config.model_name,
-                directory=temp_dir,
-            )
-            # Copy the torch model config to the model store
-            fileio.copy(
-                config_file_uri,
-                os.path.join(config_properties_uri, "config.properties"),
-            )
-
-    service_config = params.service_config.copy()
-    service_config.model_uri = deployment_folder_uri
-    return service_config
-
-
-class TorchModelArchiver(BaseModel):
-    """Model Archiver for PyTorch models.
-
-    Attributes:
-        model_name: Model name.
-        model_version: Model version.
-        serialized_file: Serialized model file.
-        handler: TorchServe's handler file to handle custom TorchServe inference logic.
-        extra_files: Comma separated path to extra dependency files.
-        requirements_file: Path to requirements file.
-        export_path: Path to export model.
-        runtime: Runtime of the model.
-        force: Force export of the model.
-        archive_format: Archive format.
-    """
-
-    model_name: str
-    serialized_file: str
-    model_file: str
-    handler: str
-    export_path: str
-    extra_files: Optional[List[str]] = None
-    version: Optional[str] = None
-    requirements_file: Optional[str] = None
-    runtime: Optional[str] = "python"
-    force: Optional[bool] = None
-    archive_format: Optional[str] = "default"
-
-
-def generate_model_deployer_config(
-    model_name: str,
-    directory: str,
-) -> str:
-    """Generate a model deployer config.
-
-    Args:
-        model_name: the name of the model
-        directory: the directory where the model is stored
-
-    Returns:
-        None
-    """
-    config_lines = [
-        "inference_address=http://0.0.0.0:8085",
-        "management_address=http://0.0.0.0:8085",
-        "metrics_address=http://0.0.0.0:8082",
-        "grpc_inference_port=7070",
-        "grpc_management_port=7071",
-        "enable_metrics_api=true",
-        "metrics_format=prometheus",
-        "number_of_netty_threads=4",
-        "job_queue_size=10",
-        "enable_envvars_config=true",
-        "install_py_dep_per_model=true",
-        "model_store=/mnt/models/model-store",
-    ]
-
-    with tempfile.NamedTemporaryFile(
-        suffix=".properties", mode="w+", dir=directory, delete=False
-    ) as f:
-        for line in config_lines:
-            f.write(line + "\n")
-        f.write(
-            f'model_snapshot={{"name":"startup.cfg","modelCount":1,"models":{{"{model_name}":{{"1.0":{{"defaultVersion":true,"marName":"{model_name}.mar","minWorkers":1,"maxWorkers":5,"batchSize":1,"maxBatchDelay":10,"responseTimeout":120}}}}}}}}'
-        )
-    f.close()
-    return f.name

zenml/integrations/kubeflow/utils.py

@@ -1,96 +0,0 @@
-#  Copyright (c) ZenML GmbH 2022. 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.
-"""KFP utilities."""
-
-
-from typing import TYPE_CHECKING
-
-from zenml.integrations.kubernetes import serialization_utils
-from zenml.integrations.kubernetes.pod_settings import KubernetesPodSettings
-from zenml.logger import get_logger
-
-logger = get_logger(__name__)
-
-if TYPE_CHECKING:
-    from kfp.dsl import ContainerOp
-
-
-def apply_pod_settings(
-    container_op: "ContainerOp",
-    settings: KubernetesPodSettings,
-) -> None:
-    """Applies Kubernetes Pod settings to a KFP container.
-
-    Args:
-        container_op: The container to which to apply the settings.
-        settings: The settings to apply.
-    """
-    from kubernetes.client.models import (
-        V1Affinity,
-        V1Toleration,
-        V1Volume,
-        V1VolumeMount,
-    )
-
-    if settings.host_ipc:
-        logger.warning(
-            "Host IPC is set to `True` but not supported in this orchestrator. "
-            "Ignoring..."
-        )
-
-    for key, value in settings.node_selectors.items():
-        container_op.add_node_selector_constraint(label_name=key, value=value)
-
-    if settings.affinity:
-        affinity: (
-            V1Affinity
-        ) = serialization_utils.deserialize_kubernetes_model(
-            settings.affinity, "V1Affinity"
-        )
-        container_op.add_affinity(affinity)
-
-    for toleration_dict in settings.tolerations:
-        toleration: (
-            V1Toleration
-        ) = serialization_utils.deserialize_kubernetes_model(
-            toleration_dict, "V1Toleration"
-        )
-        container_op.add_toleration(toleration)
-
-    if settings.volumes:
-        for v in settings.volumes:
-            volume: (
-                V1Volume
-            ) = serialization_utils.deserialize_kubernetes_model(v, "V1Volume")
-            container_op.add_volume(volume)
-
-    if settings.volume_mounts:
-        for v in settings.volume_mounts:
-            volume_mount: (
-                V1VolumeMount
-            ) = serialization_utils.deserialize_kubernetes_model(
-                v, "V1VolumeMount"
-            )
-            container_op.container.add_volume_mount(volume_mount)
-
-    resource_requests = settings.resources.get("requests") or {}
-    for name, value in resource_requests.items():
-        container_op.add_resource_request(name, value)
-
-    resource_limits = settings.resources.get("limits") or {}
-    for name, value in resource_limits.items():
-        container_op.add_resource_limit(name, value)
-
-    for name, value in settings.annotations.items():
-        container_op.add_pod_annotation(name, value)

zenml/lineage_graph/lineage_graph.py

@@ -1,244 +0,0 @@
-#  Copyright (c) ZenML GmbH 2022. 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.
-"""Class for lineage graph generation."""
-
-from typing import TYPE_CHECKING, List, Optional, Tuple, Union
-
-from pydantic import BaseModel
-
-from zenml.enums import ExecutionStatus
-from zenml.lineage_graph.edge import Edge
-from zenml.lineage_graph.node import (
-    ArtifactNode,
-    ArtifactNodeDetails,
-    StepNode,
-    StepNodeDetails,
-)
-from zenml.lineage_graph.node.artifact_node import ArtifactNodeStatus
-
-if TYPE_CHECKING:
-    from zenml.models import (
-        ArtifactVersionResponse,
-        PipelineRunResponse,
-        StepRunResponse,
-    )
-
-
-ARTIFACT_PREFIX = "artifact_"
-STEP_PREFIX = "step_"
-
-
-class LineageGraph(BaseModel):
-    """A lineage graph representation of a PipelineRunResponseModel."""
-
-    nodes: List[Union[StepNode, ArtifactNode]] = []
-    edges: List[Edge] = []
-    root_step_id: Optional[str] = None
-    run_metadata: List[Tuple[str, str, str]] = []
-
-    def generate_run_nodes_and_edges(self, run: "PipelineRunResponse") -> None:
-        """Initializes a lineage graph from a pipeline run.
-
-        Args:
-            run: The PipelineRunResponseModel to generate the lineage graph for.
-        """
-        self.run_metadata = [
-            (m.key, str(m.value), str(m.type))
-            for m in run.run_metadata.values()
-        ]
-
-        for step in run.steps.values():
-            self.generate_step_nodes_and_edges(step)
-
-        self.add_external_artifacts(run)
-        self.add_direct_edges(run)
-
-    def generate_step_nodes_and_edges(self, step: "StepRunResponse") -> None:
-        """Generates the nodes and edges for a step and its artifacts.
-
-        Args:
-            step: The step to generate the nodes and edges for.
-        """
-        step_id = STEP_PREFIX + str(step.id)
-
-        # Set a root step if it doesn't exist yet
-        if self.root_step_id is None:
-            self.root_step_id = step_id
-
-        # Add the step node
-        self.add_step_node(step, step_id)
-
-        # Add nodes and edges for all output artifacts
-        for artifact_name, artifact_version in step.outputs.items():
-            artifact_version_id = ARTIFACT_PREFIX + str(artifact_version.id)
-            if step.status == ExecutionStatus.CACHED:
-                artifact_status = ArtifactNodeStatus.CACHED
-            elif step.status == ExecutionStatus.COMPLETED:
-                artifact_status = ArtifactNodeStatus.CREATED
-            else:
-                artifact_status = ArtifactNodeStatus.UNKNOWN
-            self.add_artifact_node(
-                artifact=artifact_version,
-                id=artifact_version_id,
-                name=artifact_name,
-                step_id=str(step_id),
-                status=artifact_status,
-            )
-            self.add_edge(step_id, artifact_version_id)
-
-        # Add nodes and edges for all input artifacts
-        for artifact_name, artifact_version in step.inputs.items():
-            artifact_version_id = ARTIFACT_PREFIX + str(artifact_version.id)
-            self.add_edge(artifact_version_id, step_id)
-
-    def add_external_artifacts(self, run: "PipelineRunResponse") -> None:
-        """Adds all external artifacts to the lineage graph.
-
-        Args:
-            run: The pipeline run to add external artifacts for.
-        """
-        nodes_ids = {node.id for node in self.nodes}
-        for step in run.steps.values():
-            for artifact_name, artifact_version in step.inputs.items():
-                artifact_version_id = ARTIFACT_PREFIX + str(
-                    artifact_version.id
-                )
-                if artifact_version_id not in nodes_ids:
-                    self.add_artifact_node(
-                        artifact=artifact_version,
-                        id=artifact_version_id,
-                        name=artifact_name,
-                        step_id=str(artifact_version.producer_step_run_id),
-                        status=ArtifactNodeStatus.EXTERNAL,
-                    )
-
-    def add_direct_edges(self, run: "PipelineRunResponse") -> None:
-        """Add all direct edges between nodes generated by `after=...`.
-
-        Args:
-            run: The pipeline run to add direct edges for.
-        """
-        for step in run.steps.values():
-            step_id = STEP_PREFIX + str(step.id)
-            for parent_step_id_uuid in step.parent_step_ids:
-                parent_step_id = STEP_PREFIX + str(parent_step_id_uuid)
-                if not self.has_artifact_link(step_id, parent_step_id):
-                    self.add_edge(parent_step_id, step_id)
-
-    def has_artifact_link(self, step_id: str, parent_step_id: str) -> bool:
-        """Checks if a step has an artifact link to a parent step.
-
-        This is the case for all parent steps that were not specified via
-        `after=...`.
-
-        Args:
-            step_id: The node ID of the step to check.
-            parent_step_id: T node ID of the parent step to check.
-
-        Returns:
-            True if the steps are linked via an artifact, False otherwise.
-        """
-        parent_outputs, child_inputs = set(), set()
-        for edge in self.edges:
-            if edge.source == parent_step_id:
-                parent_outputs.add(edge.target)
-            if edge.target == step_id:
-                child_inputs.add(edge.source)
-        return bool(parent_outputs.intersection(child_inputs))
-
-    def add_step_node(
-        self,
-        step: "StepRunResponse",
-        id: str,
-    ) -> None:
-        """Adds a step node to the lineage graph.
-
-        Args:
-            step: The step to add a node for.
-            id: The id of the step node.
-        """
-        step_config = step.config.dict()
-        if step_config:
-            step_config = {
-                key: value
-                for key, value in step_config.items()
-                if key not in ["inputs", "outputs", "parameters"] and value
-            }
-        self.nodes.append(
-            StepNode(
-                id=id,
-                data=StepNodeDetails(
-                    execution_id=str(step.id),
-                    name=step.name,  # redundant for consistency
-                    status=step.status,
-                    entrypoint_name=step.config.name,  # redundant for consistency
-                    parameters=step.config.parameters,
-                    configuration=step_config,
-                    inputs={k: v.uri for k, v in step.inputs.items()},
-                    outputs={k: v.uri for k, v in step.outputs.items()},
-                    metadata=[
-                        (m.key, str(m.value), str(m.type))
-                        for m in step.run_metadata.values()
-                    ],
-                ),
-            )
-        )
-
-    def add_artifact_node(
-        self,
-        artifact: "ArtifactVersionResponse",
-        id: str,
-        name: str,
-        step_id: str,
-        status: ArtifactNodeStatus,
-    ) -> None:
-        """Adds an artifact node to the lineage graph.
-
-        Args:
-            artifact: The artifact to add a node for.
-            id: The id of the artifact node.
-            name: The input or output name of the artifact.
-            step_id: The id of the step that produced the artifact.
-            status: The status of the step that produced the artifact.
-        """
-        node = ArtifactNode(
-            id=id,
-            data=ArtifactNodeDetails(
-                execution_id=str(artifact.id),
-                name=name,
-                status=status,
-                is_cached=status == ArtifactNodeStatus.CACHED,
-                artifact_type=artifact.type,
-                artifact_data_type=artifact.data_type.import_path,
-                parent_step_id=step_id,
-                producer_step_id=str(artifact.producer_step_run_id),
-                uri=artifact.uri,
-                metadata=[
-                    (m.key, str(m.value), str(m.type))
-                    for m in artifact.run_metadata.values()
-                ],
-            ),
-        )
-        self.nodes.append(node)
-
-    def add_edge(self, source: str, target: str) -> None:
-        """Adds an edge to the lineage graph.
-
-        Args:
-            source: The source node id.
-            target: The target node id.
-        """
-        self.edges.append(
-            Edge(id=source + "_" + target, source=source, target=target)
-        )

zenml/lineage_graph/node/artifact_node.py

@@ -1,52 +0,0 @@
-#  Copyright (c) ZenML GmbH 2022. 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.
-"""Class for all lineage artifact nodes."""
-
-
-from typing import List, Optional, Tuple
-
-from zenml.lineage_graph.node.base_node import (
-    BaseNode,
-    BaseNodeDetails,
-)
-from zenml.utils.enum_utils import StrEnum
-
-
-class ArtifactNodeStatus(StrEnum):
-    """Enum that represents the status of an artifact."""
-
-    CACHED = "cached"
-    CREATED = "created"
-    EXTERNAL = "external"
-    UNKNOWN = "unknown"
-
-
-class ArtifactNodeDetails(BaseNodeDetails):
-    """Captures all artifact details for the node."""
-
-    status: ArtifactNodeStatus
-    is_cached: bool
-    artifact_type: str
-    artifact_data_type: str
-    parent_step_id: str
-    producer_step_id: Optional[str]
-    uri: str
-    metadata: List[Tuple[str, str, str]]  # (key, value, type)
-
-
-class ArtifactNode(BaseNode):
-    """A class that represents an artifact node in a lineage graph."""
-
-    type: str = "artifact"
-    data: ArtifactNodeDetails