zenml-nightly 0.80.1.dev20250408__py3-none-any.whl → 0.80.2.dev20250410__py3-none-any.whl

zenml/integrations/kubernetes/orchestrators/kubernetes_orchestrator_entrypoint_configuration.py

@@ -13,7 +13,7 @@
 #  permissions and limitations under the License.
 """Entrypoint configuration for the Kubernetes master/orchestrator pod."""
 
-from typing import TYPE_CHECKING, List, Set
+from typing import TYPE_CHECKING, List, Optional, Set
 
 if TYPE_CHECKING:
     from uuid import UUID
@@ -21,6 +21,7 @@ if TYPE_CHECKING:
 RUN_NAME_OPTION = "run_name"
 DEPLOYMENT_ID_OPTION = "deployment_id"
 NAMESPACE_OPTION = "kubernetes_namespace"
+RUN_ID_OPTION = "run_id"
 
 
 class KubernetesOrchestratorEntrypointConfiguration:
@@ -60,6 +61,7 @@ class KubernetesOrchestratorEntrypointConfiguration:
         run_name: str,
         deployment_id: "UUID",
         kubernetes_namespace: str,
+        run_id: Optional["UUID"] = None,
     ) -> List[str]:
         """Gets all arguments that the entrypoint command should be called with.
 
@@ -67,6 +69,7 @@ class KubernetesOrchestratorEntrypointConfiguration:
             run_name: Name of the ZenML run.
             deployment_id: ID of the deployment.
             kubernetes_namespace: Name of the Kubernetes namespace.
+            run_id: Optional ID of the pipeline run. Not set for scheduled runs.
 
         Returns:
             List of entrypoint arguments.
@@ -80,4 +83,8 @@ class KubernetesOrchestratorEntrypointConfiguration:
             kubernetes_namespace,
         ]
 
+        if run_id:
+            args.append(f"--{RUN_ID_OPTION}")
+            args.append(str(run_id))
+
         return args

zenml/integrations/lightning/orchestrators/lightning_orchestrator.py

@@ -46,7 +46,7 @@ from zenml.stack import StackValidator
 from zenml.utils import code_utils, io_utils, source_utils
 
 if TYPE_CHECKING:
-    from zenml.models import PipelineDeploymentResponse
+    from zenml.models import PipelineDeploymentResponse, PipelineRunResponse
     from zenml.stack import Stack
 
 
@@ -199,6 +199,7 @@ class LightningOrchestrator(WheeledOrchestrator):
         deployment: "PipelineDeploymentResponse",
         stack: "Stack",
         environment: Dict[str, str],
+        placeholder_run: Optional["PipelineRunResponse"] = None,
     ) -> Any:
         """Creates a wheel and uploads the pipeline to Lightning.
 
@@ -223,6 +224,7 @@ class LightningOrchestrator(WheeledOrchestrator):
             stack: The stack the pipeline will run on.
             environment: Environment variables to set in the orchestration
                 environment.
+            placeholder_run: An optional placeholder run for the deployment.
 
         Raises:
             ValueError: If the schedule is not set or if the cron expression

zenml/integrations/mlflow/__init__.py

@@ -61,39 +61,12 @@ class MlflowIntegration(Integration):
         from zenml.integrations.numpy import NumpyIntegration
         from zenml.integrations.pandas import PandasIntegration
         
-
         reqs = [
             "mlflow>=2.1.1,<3",
             # TODO: remove this requirement once rapidjson is fixed
             "python-rapidjson<1.15",
-            # When you do:
-            #       pip install zenml
-            # You get all our required dependencies. However, if you follow it
-            # with:
-            #       zenml integration install mlflow
-            # This downgrades pydantic to v1 even though mlflow does not have
-            # any issues with v2. This is why we have to pin it here so a
-            # downgrade will not happen.
-            "pydantic>=2.8.0,<2.9.0",
         ]
 
-        if python_version:
-            version_minor = version.parse(python_version).minor
-        else:
-            version_minor = sys.version_info.minor
-
-        if version_minor >= 12:
-            logger.debug(
-                "The MLflow integration on Python 3.12 and above is not yet "
-                "fully supported: The extra dependencies 'mlserver' and "
-                "'mlserver-mlflow' will be skipped."
-            )
-        else:
-            reqs.extend([
-                "mlserver>=1.3.3",
-                "mlserver-mlflow>=1.3.3",
-            ])
-
         reqs.extend(NumpyIntegration.get_requirements(target_os=target_os, python_version=python_version))
         reqs.extend(PandasIntegration.get_requirements(target_os=target_os, python_version=python_version))
         return reqs

zenml/integrations/mlflow/services/mlflow_deployment.py

@@ -13,6 +13,7 @@
 #  permissions and limitations under the License.
 """Implementation of the MLflow deployment functionality."""
 
+import importlib
 import os
 import sys
 from typing import TYPE_CHECKING, Any, Dict, Optional, Union
@@ -126,14 +127,31 @@ class MLFlowDeploymentConfig(LocalDaemonServiceConfig):
             the validated value
 
         Raises:
-            ValueError: if mlserver is set to true on Python 3.12 as it is not
-                yet supported.
+            ValueError: if mlserver packages are not installed
         """
-        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."
-            )
+        if mlserver is True:
+            # For the mlserver deployment, the mlserver and
+            # mlserver-mlflow packages need to be installed separately
+            # because they rely on an older version of Pydantic.
+
+            # Check if the mlserver and mlserver-mlflow packages are installed
+            try:
+                importlib.import_module("mlserver")
+                importlib.import_module("mlserver_mlflow")
+            except ModuleNotFoundError:
+                if sys.version_info.minor >= 12:
+                    raise ValueError(
+                        "The mlserver deployment is not yet supported on "
+                        "Python 3.12 or above."
+                    )
+
+                raise ValueError(
+                    "The MLflow MLServer backend requires the `mlserver` and "
+                    "`mlserver-mlflow` packages to be installed. These "
+                    "packages are not included in the ZenML MLflow integration "
+                    "requirements. Please install them manually."
+                )
+
         return mlserver
 
 

zenml/integrations/seldon/services/seldon_deployment.py

@@ -72,7 +72,7 @@ class SeldonDeploymentConfig(ServiceConfig):
     model_metadata: Dict[str, Any] = Field(default_factory=dict)
     extra_args: Dict[str, Any] = Field(default_factory=dict)
     is_custom_deployment: Optional[bool] = False
-    spec: Optional[Dict[Any, Any]] = Field(default_factory=dict)
+    spec: Optional[Dict[Any, Any]] = Field(default_factory=dict)  # type: ignore[arg-type]
     serviceAccountName: Optional[str] = None
 
     def get_seldon_deployment_labels(self) -> Dict[str, str]:

zenml/integrations/skypilot/orchestrators/skypilot_base_vm_orchestrator.py

@@ -39,7 +39,7 @@ from zenml.orchestrators.utils import get_orchestrator_run_name
 from zenml.stack import StackValidator
 
 if TYPE_CHECKING:
-    from zenml.models import PipelineDeploymentResponse
+    from zenml.models import PipelineDeploymentResponse, PipelineRunResponse
     from zenml.stack import Stack
 
 
@@ -149,6 +149,7 @@ class SkypilotBaseOrchestrator(ContainerizedOrchestrator):
         deployment: "PipelineDeploymentResponse",
         stack: "Stack",
         environment: Dict[str, str],
+        placeholder_run: Optional["PipelineRunResponse"] = None,
     ) -> Any:
         """Runs each pipeline step in a separate Skypilot container.
 
@@ -157,6 +158,7 @@ class SkypilotBaseOrchestrator(ContainerizedOrchestrator):
             stack: The stack the pipeline will run on.
             environment: Environment variables to set in the orchestration
                 environment.
+            placeholder_run: An optional placeholder run for the deployment.
 
         Raises:
             Exception: If the pipeline run fails.

zenml/integrations/tekton/orchestrators/tekton_orchestrator.py

@@ -53,7 +53,7 @@ from zenml.utils import io_utils, yaml_utils
 
 if TYPE_CHECKING:
     from zenml.config.base_settings import BaseSettings
-    from zenml.models import PipelineDeploymentResponse
+    from zenml.models import PipelineDeploymentResponse, PipelineRunResponse
     from zenml.stack import Stack
 
 
@@ -460,6 +460,7 @@ class TektonOrchestrator(ContainerizedOrchestrator):
         deployment: "PipelineDeploymentResponse",
         stack: "Stack",
         environment: Dict[str, str],
+        placeholder_run: Optional["PipelineRunResponse"] = None,
     ) -> Any:
         """Runs the pipeline on Tekton.
 
@@ -471,6 +472,7 @@ class TektonOrchestrator(ContainerizedOrchestrator):
             stack: The stack the pipeline will run on.
             environment: Environment variables to set in the orchestration
                 environment.
+            placeholder_run: An optional placeholder run for the deployment.
 
         Raises:
             RuntimeError: If you try to run the pipelines in a notebook

zenml/integrations/vllm/services/vllm_deployment.py

@@ -150,7 +150,7 @@ class VLLMDeploymentService(LocalDaemonService, BaseDeploymentService):
 
         try:
             parser: argparse.ArgumentParser = make_arg_parser(
-                FlexibleArgumentParser()
+                FlexibleArgumentParser()  # type: ignore[no-untyped-call]
             )
             # pass in empty list to get default args
             # otherwise it will try to get the args from sys.argv

zenml/logging/step_logging.py

@@ -263,7 +263,7 @@ class StepLogsStorage:
             formatted_message = (
                 f"[{timestamp} UTC] {remove_ansi_escape_codes(text)}"
             )
-            self.buffer.append(formatted_message)
+            self.buffer.append(formatted_message.rstrip())
             self.save_to_file()
 
     @property

zenml/materializers/__init__.py

@@ -26,14 +26,14 @@ from zenml.materializers.built_in_materializer import (
 from zenml.materializers.cloudpickle_materializer import (
     CloudpickleMaterializer,
 )
+from zenml.materializers.path_materializer import PathMaterializer
+from zenml.materializers.pydantic_materializer import PydanticMaterializer
+from zenml.materializers.service_materializer import ServiceMaterializer
 from zenml.materializers.structured_string_materializer import (
     StructuredStringMaterializer,
 )
-from zenml.materializers.pydantic_materializer import PydanticMaterializer
-from zenml.materializers.service_materializer import ServiceMaterializer
 from zenml.materializers.uuid_materializer import UUIDMaterializer
 
-
 __all__ = [
     "BuiltInContainerMaterializer",
     "BuiltInMaterializer",
@@ -43,4 +43,5 @@ __all__ = [
     "PydanticMaterializer",
     "ServiceMaterializer",
     "UUIDMaterializer",
+    "PathMaterializer",
 ]

zenml/materializers/built_in_materializer.py

@@ -107,6 +107,19 @@ class BuiltInMaterializer(BaseMaterializer):
             ensure_ascii=not ZENML_MATERIALIZER_ALLOW_NON_ASCII_JSON_DUMPS,
         )
 
+    def save_visualizations(
+        self, data: Union[bool, float, int, str]
+    ) -> Dict[str, VisualizationType]:
+        """Save visualizations for the given basic type.
+
+        Args:
+            data: The data to save visualizations for.
+
+        Returns:
+            A dictionary of visualization URIs and their types.
+        """
+        return {self.data_path.replace("\\", "/"): VisualizationType.JSON}
+
     def extract_metadata(
         self, data: Union[bool, float, int, str]
     ) -> Dict[str, "MetadataType"]:
@@ -165,6 +178,17 @@ class BytesMaterializer(BaseMaterializer):
         with self.artifact_store.open(self.data_path, "wb") as file_:
             file_.write(data)
 
+    def save_visualizations(self, data: bytes) -> Dict[str, VisualizationType]:
+        """Save visualizations for the given bytes data.
+
+        Args:
+            data: The bytes data to save visualizations for.
+
+        Returns:
+            A dictionary of visualization URIs and their types.
+        """
+        return {self.data_path.replace("\\", "/"): VisualizationType.MARKDOWN}
+
 
 def _all_serializable(iterable: Iterable[Any]) -> bool:
     """For an iterable, check whether all of its elements are JSON-serializable.

zenml/materializers/path_materializer.py

@@ -0,0 +1,119 @@
+#  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 materializer for Path objects."""
+
+import os
+import shutil
+import tarfile
+from pathlib import Path
+from typing import Any, ClassVar, Tuple, Type
+
+from zenml.constants import (
+    ENV_ZENML_DISABLE_PATH_MATERIALIZER,
+    handle_bool_env_var,
+)
+from zenml.enums import ArtifactType
+from zenml.io import fileio
+from zenml.materializers.base_materializer import BaseMaterializer
+
+
+class PathMaterializer(BaseMaterializer):
+    """Materializer for Path objects.
+
+    This materializer handles `pathlib.Path` objects by storing their contents
+    in a compressed tar archive within the artifact store if it's a directory,
+    or directly copying the file if it's a single file.
+    """
+
+    ASSOCIATED_TYPES: ClassVar[Tuple[Type[Any], ...]] = (Path,)
+    ASSOCIATED_ARTIFACT_TYPE: ClassVar[ArtifactType] = ArtifactType.DATA
+    ARCHIVE_NAME: ClassVar[str] = "data.tar.gz"
+    FILE_NAME: ClassVar[str] = "file_data"
+
+    # Skip registration if the environment variable is set
+    SKIP_REGISTRATION: ClassVar[bool] = handle_bool_env_var(
+        ENV_ZENML_DISABLE_PATH_MATERIALIZER, default=False
+    )
+
+    def load(self, data_type: Type[Any]) -> Any:
+        """Copy the artifact files to a local temp directory or file.
+
+        Args:
+            data_type: Unused.
+
+        Returns:
+            Path to the local directory or file that contains the artifact.
+
+        Raises:
+            FileNotFoundError: If the artifact is not found in the artifact store.
+        """
+        # Create a temporary directory that will persist until step execution ends
+        with self.get_temporary_directory(delete_at_exit=False) as directory:
+            # Check if we're loading a file or directory by looking for the archive
+            archive_path_remote = os.path.join(self.uri, self.ARCHIVE_NAME)
+            file_path_remote = os.path.join(self.uri, self.FILE_NAME)
+
+            if fileio.exists(archive_path_remote):
+                # This is a directory artifact
+                archive_path_local = os.path.join(directory, self.ARCHIVE_NAME)
+                fileio.copy(archive_path_remote, archive_path_local)
+
+                # Extract the archive to the temporary directory
+                with tarfile.open(archive_path_local, "r:gz") as tar:
+                    tar.extractall(path=directory)
+
+                # Clean up the archive file
+                os.remove(archive_path_local)
+                return Path(directory)
+            elif fileio.exists(file_path_remote):
+                # This is a single file artifact
+                file_path_local = os.path.join(
+                    directory, os.path.basename(file_path_remote)
+                )
+                fileio.copy(file_path_remote, file_path_local)
+                return Path(file_path_local)
+            else:
+                raise FileNotFoundError(
+                    f"Could not find artifact at {archive_path_remote} or {file_path_remote}"
+                )
+
+    def save(self, data: Any) -> None:
+        """Store the directory or file in the artifact store.
+
+        Args:
+            data: Path to a local directory or file to store. Must be a Path object.
+        """
+        assert isinstance(data, Path)
+
+        if data.is_dir():
+            # Handle directory artifact
+            with self.get_temporary_directory(
+                delete_at_exit=True
+            ) as directory:
+                archive_base = os.path.join(directory, "data")
+
+                # Create tar.gz archive - automatically uses relative paths
+                shutil.make_archive(
+                    base_name=archive_base, format="gztar", root_dir=str(data)
+                )
+
+                # Copy the archive to the artifact store
+                fileio.copy(
+                    f"{archive_base}.tar.gz",
+                    os.path.join(self.uri, self.ARCHIVE_NAME),
+                )
+        else:
+            # Handle single file artifact
+            file_path_remote = os.path.join(self.uri, self.FILE_NAME)
+            fileio.copy(str(data), file_path_remote)

zenml/model/model.py

@@ -81,7 +81,7 @@ class Model(BaseModel):
     # technical attributes
     model_version_id: Optional[UUID] = None
     suppress_class_validation_warnings: bool = False
-    _model_id: UUID = PrivateAttr(None)
+    _model_id: Optional[UUID] = PrivateAttr(None)
     _number: Optional[int] = PrivateAttr(None)
     _created_model_version: bool = PrivateAttr(False)
 
@@ -130,6 +130,8 @@ class Model(BaseModel):
         """
         if self._model_id is None:
             self._get_or_create_model()
+
+        assert self._model_id is not None
         return self._model_id
 
     @property

zenml/models/__init__.py

@@ -418,7 +418,11 @@ from zenml.models.v2.misc.info_models import (
 
 # V2
 ActionResponseResources.model_rebuild()
+ActionResponseMetadata.model_rebuild()
 APIKeyResponseBody.model_rebuild()
+ArtifactResponse.model_rebuild()
+ArtifactResponseBody.model_rebuild()
+ArtifactResponseMetadata.model_rebuild()
 ArtifactVersionRequest.model_rebuild()
 ArtifactVersionResponseBody.model_rebuild()
 ArtifactVersionResponseMetadata.model_rebuild()
@@ -477,6 +481,7 @@ StepRunRequest.model_rebuild()
 StepRunResponseBody.model_rebuild()
 StepRunResponseMetadata.model_rebuild()
 StepRunResponseResources.model_rebuild()
+TagResponseBody.model_rebuild()
 TriggerExecutionResponseResources.model_rebuild()
 TriggerResponseBody.model_rebuild()
 TriggerResponseMetadata.model_rebuild()

zenml/models/v2/core/run_template.py

@@ -79,6 +79,10 @@ class RunTemplateRequest(ProjectScopedRequest):
     source_deployment_id: UUID = Field(
         title="The deployment that should be the base of the created template."
     )
+    hidden: bool = Field(
+        default=False,
+        title="Whether the run template is hidden.",
+    )
     tags: Optional[List[str]] = Field(
         default=None,
         title="Tags of the run template.",
@@ -101,6 +105,10 @@ class RunTemplateUpdate(BaseUpdate):
         title="The description of the run template.",
         max_length=TEXT_FIELD_MAX_LENGTH,
     )
+    hidden: Optional[bool] = Field(
+        default=None,
+        title="Whether the run template is hidden.",
+    )
     add_tags: Optional[List[str]] = Field(
         default=None, title="New tags to add to the run template."
     )
@@ -118,6 +126,10 @@ class RunTemplateResponseBody(ProjectScopedResponseBody):
     runnable: bool = Field(
         title="If a run can be started from the template.",
     )
+    hidden: bool = Field(
+        default=False,
+        title="Whether the run template is hidden.",
+    )
     latest_run_id: Optional[UUID] = Field(
         default=None,
         title="The ID of the latest run of the run template.",
@@ -205,6 +217,15 @@ class RunTemplateResponse(
         """
         return self.get_body().runnable
 
+    @property
+    def hidden(self) -> bool:
+        """The `hidden` property.
+
+        Returns:
+            the value of the property.
+        """
+        return self.get_body().hidden
+
     @property
     def latest_run_id(self) -> Optional[UUID]:
         """The `latest_run_id` property.
@@ -320,6 +341,7 @@ class RunTemplateFilter(ProjectScopedFilter, TaggableFilter):
         "pipeline_id",
         "pipeline",
         "stack",
+        "hidden",
     ]
     CUSTOM_SORTING_OPTIONS = [
         *ProjectScopedFilter.CUSTOM_SORTING_OPTIONS,
@@ -334,6 +356,10 @@ class RunTemplateFilter(ProjectScopedFilter, TaggableFilter):
         default=None,
         description="Name of the run template.",
     )
+    hidden: Optional[bool] = Field(
+        default=None,
+        description="Whether the run template is hidden.",
+    )
     pipeline_id: Optional[Union[UUID, str]] = Field(
         default=None,
         description="Pipeline associated with the template.",
@@ -376,7 +402,7 @@ class RunTemplateFilter(ProjectScopedFilter, TaggableFilter):
         """
         custom_filters = super().get_custom_filters(table)
 
-        from sqlmodel import and_
+        from sqlmodel import and_, col
 
         from zenml.zen_stores.schemas import (
             CodeReferenceSchema,
@@ -386,6 +412,11 @@ class RunTemplateFilter(ProjectScopedFilter, TaggableFilter):
             StackSchema,
         )
 
+        if self.hidden is not None:
+            custom_filters.append(
+                col(RunTemplateSchema.hidden).is_(self.hidden)
+            )
+
         if self.code_repository_id:
             code_repo_filter = and_(
                 RunTemplateSchema.source_deployment_id

zenml/orchestrators/base_orchestrator.py

@@ -140,6 +140,7 @@ class BaseOrchestrator(StackComponent, ABC):
         deployment: "PipelineDeploymentResponse",
         stack: "Stack",
         environment: Dict[str, str],
+        placeholder_run: Optional["PipelineRunResponse"] = None,
     ) -> Optional[Iterator[Dict[str, MetadataType]]]:
         """The method needs to be implemented by the respective orchestrator.
 
@@ -175,6 +176,7 @@ class BaseOrchestrator(StackComponent, ABC):
             stack: The stack the pipeline will run on.
             environment: Environment variables to set in the orchestration
                 environment. These don't need to be set if running locally.
+            placeholder_run: An optional placeholder run for the deployment.
 
         Yields:
             Metadata for the pipeline run.
@@ -250,6 +252,7 @@ class BaseOrchestrator(StackComponent, ABC):
                 deployment=deployment,
                 stack=stack,
                 environment=environment,
+                placeholder_run=placeholder_run,
             ):
                 for metadata_dict in metadata_iterator:
                     try: