zenml-nightly 0.66.0.dev20240928__py3-none-any.whl → 0.68.1.dev20241102__py3-none-any.whl

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.model_dump()
-        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/__init__.py

@@ -1,32 +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.
-"""Initialization of lineage nodes."""
-
-from zenml.lineage_graph.node.artifact_node import (  # noqa
-    ArtifactNode,
-    ArtifactNodeDetails,
-)
-from zenml.lineage_graph.node.base_node import BaseNode  # noqa
-from zenml.lineage_graph.node.step_node import (  # noqa
-    StepNode,
-    StepNodeDetails,
-)
-
-__all__ = [
-    "BaseNode",
-    "ArtifactNode",
-    "StepNode",
-    "StepNodeDetails",
-    "ArtifactNodeDetails",
-]

zenml/lineage_graph/node/artifact_node.py

@@ -1,51 +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

zenml/lineage_graph/node/step_node.py

@@ -1,41 +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 step nodes."""
-
-from typing import Any, Dict, List, Tuple
-
-from zenml.enums import ExecutionStatus
-from zenml.lineage_graph.node.base_node import (
-    BaseNode,
-    BaseNodeDetails,
-)
-
-
-class StepNodeDetails(BaseNodeDetails):
-    """Captures all artifact details for the node."""
-
-    status: ExecutionStatus
-    entrypoint_name: str
-    parameters: Dict[str, Any]
-    configuration: Dict[str, Any]
-    inputs: Dict[str, Any]
-    outputs: Dict[str, Any]
-    metadata: List[Tuple[str, str, str]]  # (key, value, type)
-
-
-class StepNode(BaseNode):
-    """A class that represents a step node in a lineage graph."""
-
-    type: str = "step"
-    data: StepNodeDetails

zenml/model/model_version.py

@@ -1,38 +0,0 @@
-#  Copyright (c) ZenML GmbH 2023. 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.
-"""DEPRECATED, use `from zenml import Model` instead."""
-
-from typing import Any
-
-from zenml.logger import get_logger
-from zenml.model.model import Model
-
-logger = get_logger(__name__)
-
-
-# TODO: deprecate me
-class ModelVersion(Model):
-    """DEPRECATED, use `from zenml import Model` instead."""
-
-    def __init__(self, *args: Any, **kwargs: Any) -> None:
-        """DEPRECATED, use `from zenml import Model` instead.
-
-        Args:
-            *args: Variable length argument list.
-            **kwargs: Arbitrary keyword arguments.
-        """
-        logger.warning(
-            "`ModelVersion` is deprecated. Please use `Model` instead."
-        )
-        super().__init__(*args, **kwargs)

zenml/new/pipelines/code_archive.py

@@ -1,157 +0,0 @@
-#  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.
-"""Code archive."""
-
-import os
-from pathlib import Path
-from typing import IO, TYPE_CHECKING, Dict, Optional
-
-from zenml.logger import get_logger
-from zenml.utils import string_utils
-from zenml.utils.archivable import Archivable
-
-if TYPE_CHECKING:
-    from git.repo.base import Repo
-
-
-logger = get_logger(__name__)
-
-
-class CodeArchive(Archivable):
-    """Code archive class.
-
-    This class is used to archive user code before uploading it to the artifact
-    store. If the user code is stored in a Git repository, only files not
-    excluded by gitignores will be included in the archive.
-    """
-
-    def __init__(self, root: str) -> None:
-        """Initialize the object.
-
-        Args:
-            root: Root directory of the archive.
-        """
-        super().__init__()
-        self._root = root
-
-    @property
-    def git_repo(self) -> Optional["Repo"]:
-        """Git repository active at the code archive root.
-
-        Returns:
-            The git repository if available.
-        """
-        try:
-            # These imports fail when git is not installed on the machine
-            from git.exc import InvalidGitRepositoryError
-            from git.repo.base import Repo
-        except ImportError:
-            return None
-
-        try:
-            git_repo = Repo(path=self._root, search_parent_directories=True)
-        except InvalidGitRepositoryError:
-            return None
-
-        return git_repo
-
-    def _get_all_files(self) -> Dict[str, str]:
-        """Get all files inside the archive root.
-
-        Returns:
-            All files inside the archive root.
-        """
-        all_files = {}
-        for root, _, files in os.walk(self._root):
-            for file in files:
-                file_path = os.path.join(root, file)
-                path_in_archive = os.path.relpath(file_path, self._root)
-                all_files[path_in_archive] = file_path
-
-        return all_files
-
-    def get_files(self) -> Dict[str, str]:
-        """Gets all regular files that should be included in the archive.
-
-        Raises:
-            RuntimeError: If the code archive would not include any files.
-
-        Returns:
-            A dict {path_in_archive: path_on_filesystem} for all regular files
-            in the archive.
-        """
-        all_files = {}
-
-        if repo := self.git_repo:
-            try:
-                result = repo.git.ls_files(
-                    "--cached",
-                    "--others",
-                    "--modified",
-                    "--exclude-standard",
-                    self._root,
-                )
-            except Exception as e:
-                logger.warning(
-                    "Failed to get non-ignored files from git: %s", str(e)
-                )
-                all_files = self._get_all_files()
-            else:
-                for file in result.split():
-                    file_path = os.path.join(repo.working_dir, file)
-                    path_in_archive = os.path.relpath(file_path, self._root)
-
-                    if os.path.exists(file_path):
-                        all_files[path_in_archive] = file_path
-        else:
-            all_files = self._get_all_files()
-
-        if not all_files:
-            raise RuntimeError(
-                "The code archive to be uploaded does not contain any files. "
-                "This is probably because all files in your source root "
-                f"`{self._root}` are ignored by a .gitignore file."
-            )
-
-        # Explicitly remove .zen directories as we write an updated version
-        # to disk everytime ZenML is called. This updates the mtime of the
-        # file, which invalidates the code upload caching. The values in
-        # the .zen directory are not needed anyway as we set them as
-        # environment variables.
-        all_files = {
-            path_in_archive: file_path
-            for path_in_archive, file_path in sorted(all_files.items())
-            if ".zen" not in Path(path_in_archive).parts[:-1]
-        }
-
-        return all_files
-
-    def write_archive(
-        self, output_file: IO[bytes], use_gzip: bool = True
-    ) -> None:
-        """Writes an archive of the build context to the given file.
-
-        Args:
-            output_file: The file to write the archive to.
-            use_gzip: Whether to use `gzip` to compress the file.
-        """
-        super().write_archive(output_file=output_file, use_gzip=use_gzip)
-        archive_size = os.path.getsize(output_file.name)
-        if archive_size > 20 * 1024 * 1024:
-            logger.warning(
-                "Code archive size: `%s`. If you believe this is "
-                "unreasonably large, make sure to version your code in git and "
-                "ignore unnecessary files using a `.gitignore` file.",
-                string_utils.get_human_readable_filesize(archive_size),
-            )

zenml/new/pipelines/pipeline_decorator.py

@@ -1,118 +0,0 @@
-#  Copyright (c) ZenML GmbH 2023. 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.
-"""ZenML pipeline decorator definition."""
-
-from typing import (
-    TYPE_CHECKING,
-    Any,
-    Callable,
-    Dict,
-    Optional,
-    TypeVar,
-    Union,
-    overload,
-)
-
-from zenml.logger import get_logger
-
-if TYPE_CHECKING:
-    from zenml.config.base_settings import SettingsOrDict
-    from zenml.model.model import Model
-    from zenml.new.pipelines.pipeline import Pipeline
-    from zenml.types import HookSpecification
-
-    F = TypeVar("F", bound=Callable[..., None])
-
-logger = get_logger(__name__)
-
-
-@overload
-def pipeline(_func: "F") -> "Pipeline": ...
-
-
-@overload
-def pipeline(
-    *,
-    name: Optional[str] = None,
-    enable_cache: Optional[bool] = None,
-    enable_artifact_metadata: Optional[bool] = None,
-    enable_step_logs: Optional[bool] = None,
-    settings: Optional[Dict[str, "SettingsOrDict"]] = None,
-    extra: Optional[Dict[str, Any]] = None,
-) -> Callable[["F"], "Pipeline"]: ...
-
-
-def pipeline(
-    _func: Optional["F"] = None,
-    *,
-    name: Optional[str] = None,
-    enable_cache: Optional[bool] = None,
-    enable_artifact_metadata: Optional[bool] = None,
-    enable_step_logs: Optional[bool] = None,
-    settings: Optional[Dict[str, "SettingsOrDict"]] = None,
-    extra: Optional[Dict[str, Any]] = None,
-    on_failure: Optional["HookSpecification"] = None,
-    on_success: Optional["HookSpecification"] = None,
-    model: Optional["Model"] = None,
-    model_version: Optional["Model"] = None,  # TODO: deprecate me
-) -> Union["Pipeline", Callable[["F"], "Pipeline"]]:
-    """Decorator to create a pipeline.
-
-    Args:
-        _func: The decorated function.
-        name: The name of the pipeline. If left empty, the name of the
-            decorated function will be used as a fallback.
-        enable_cache: Whether to use caching or not.
-        enable_artifact_metadata: Whether to enable artifact metadata or not.
-        enable_step_logs: If step logs should be enabled for this pipeline.
-        settings: Settings for this pipeline.
-        extra: Extra configurations for this pipeline.
-        on_failure: Callback function in event of failure of the step. Can be a
-            function with a single argument of type `BaseException`, or a source
-            path to such a function (e.g. `module.my_function`).
-        on_success: Callback function in event of success of the step. Can be a
-            function with no arguments, or a source path to such a function
-            (e.g. `module.my_function`).
-        model: configuration of the model in the Model Control Plane.
-        model_version: DEPRECATED, please use `model` instead.
-
-    Returns:
-        A pipeline instance.
-    """
-
-    def inner_decorator(func: "F") -> "Pipeline":
-        from zenml.new.pipelines.pipeline import Pipeline
-
-        if model_version:
-            logger.warning(
-                "Pipeline decorator argument `model_version` is deprecated. Please use `model` instead."
-            )
-
-        p = Pipeline(
-            name=name or func.__name__,
-            enable_cache=enable_cache,
-            enable_artifact_metadata=enable_artifact_metadata,
-            enable_step_logs=enable_step_logs,
-            settings=settings,
-            extra=extra,
-            on_failure=on_failure,
-            on_success=on_success,
-            model=model or model_version,
-            entrypoint=func,
-        )
-
-        p.__doc__ = func.__doc__
-        return p
-
-    return inner_decorator if _func is None else inner_decorator(_func)