zenml-nightly 0.57.1.dev20240522__py3-none-any.whl → 0.57.1.dev20240527__py3-none-any.whl

zenml/VERSION

@@ -1 +1 @@
-0.57.1.dev20240522
+0.57.1.dev20240527

zenml/cli/base.py

@@ -91,7 +91,7 @@ ZENML_PROJECT_TEMPLATES = dict(
     ),
     llm_finetuning=ZenMLProjectTemplateLocation(
         github_url="zenml-io/template-llm-finetuning",
-        github_tag="2024.04.05",  # Make sure it is aligned with .github/workflows/update-templates-to-examples.yml
+        github_tag="2024.05.23",  # Make sure it is aligned with .github/workflows/update-templates-to-examples.yml
     ),
 )
 
@@ -432,6 +432,13 @@ def go() -> None:
 
     zenml_tutorial_path = os.path.join(os.getcwd(), "zenml_tutorial")
 
+    if not is_jupyter_installed():
+        cli_utils.error(
+            "Jupyter Notebook or JupyterLab is not installed. "
+            "Please install the 'notebook' package with `pip` "
+            "first so you can run the tutorial notebooks."
+        )
+
     with track_handler(event=AnalyticsEvent.RUN_ZENML_GO, metadata=metadata):
         console.print(zenml_cli_privacy_message, width=80)
 
@@ -459,6 +466,7 @@ def go() -> None:
                         TUTORIAL_REPO,
                         tmp_cloned_dir,
                         branch=f"release/{zenml_version}",
+                        depth=1,  # to prevent timeouts when downloading
                     )
                 example_dir = os.path.join(
                     tmp_cloned_dir, "examples/quickstart"
@@ -483,23 +491,14 @@ def go() -> None:
         )
         input("Press ENTER to continue...")
 
-    if is_jupyter_installed():
-        try:
-            subprocess.check_call(
-                ["jupyter", "notebook"], cwd=zenml_tutorial_path
-            )
-        except subprocess.CalledProcessError as e:
-            cli_utils.error(
-                "An error occurred while launching Jupyter Notebook. "
-                "Please make sure Jupyter is properly installed and try again."
-            )
-            raise e
-    else:
+    try:
+        subprocess.check_call(["jupyter", "notebook"], cwd=zenml_tutorial_path)
+    except subprocess.CalledProcessError as e:
         cli_utils.error(
-            "Jupyter Notebook or JupyterLab is not installed. "
-            "Please install the 'notebook' package with `pip` "
-            "to run the tutorial notebooks."
+            "An error occurred while launching Jupyter Notebook. "
+            "Please make sure Jupyter is properly installed and try again."
         )
+        raise e
 
 
 def _prompt_email(event_source: AnalyticsEventSource) -> bool:

zenml/client.py

@@ -5679,7 +5679,7 @@ class Client(metaclass=ClientMetaClass):
 
     def get_model_version(
         self,
-        model_name_or_id: Union[str, UUID],
+        model_name_or_id: Optional[Union[str, UUID]] = None,
         model_version_name_or_number_or_id: Optional[
             Union[str, int, ModelStages, UUID]
         ] = None,
@@ -5702,7 +5702,17 @@ class Client(metaclass=ClientMetaClass):
         Raises:
             RuntimeError: In case method inputs don't adhere to restrictions.
             KeyError: In case no model version with the identifiers exists.
+            ValueError: In case retrieval is attempted using non UUID model version
+                identifier and no model identifier provided.
         """
+        if (
+            not is_valid_uuid(model_version_name_or_number_or_id)
+            and model_name_or_id is None
+        ):
+            raise ValueError(
+                "No model identifier provided and model version identifier "
+                f"`{model_version_name_or_number_or_id}` is not a valid UUID."
+            )
         if cll := client_lazy_loader(
             "get_model_version",
             model_name_or_id=model_name_or_id,

zenml/config/__init__.py

@@ -25,8 +25,10 @@ order to persist the configuration across sessions.
 """
 from zenml.config.docker_settings import DockerSettings
 from zenml.config.resource_settings import ResourceSettings
+from zenml.config.retry_config import StepRetryConfig
 
 __all__ = [
     "DockerSettings",
     "ResourceSettings",
+    "StepRetryConfig",
 ]

zenml/config/pipeline_configurations.py

@@ -18,6 +18,7 @@ from typing import TYPE_CHECKING, Any, Dict, Optional
 from pydantic import validator
 
 from zenml.config.constants import DOCKER_SETTINGS_KEY
+from zenml.config.retry_config import StepRetryConfig
 from zenml.config.source import Source, convert_source_validator
 from zenml.config.strict_base_model import StrictBaseModel
 from zenml.model.model import Model
@@ -43,6 +44,7 @@ class PipelineConfigurationUpdate(StrictBaseModel):
     success_hook_source: Optional[Source] = None
     model: Optional[Model] = None
     parameters: Optional[Dict[str, Any]] = None
+    retry: Optional[StepRetryConfig] = None
 
     _convert_source = convert_source_validator(
         "failure_hook_source", "success_hook_source"

zenml/config/pipeline_run_configuration.py

@@ -17,6 +17,7 @@ from typing import Any, Dict, Optional, Union
 from uuid import UUID
 
 from zenml.config.base_settings import BaseSettings
+from zenml.config.retry_config import StepRetryConfig
 from zenml.config.schedule import Schedule
 from zenml.config.step_configurations import StepConfigurationUpdate
 from zenml.config.strict_base_model import StrictBaseModel
@@ -42,3 +43,4 @@ class PipelineRunConfiguration(
     extra: Dict[str, Any] = {}
     model: Optional[Model] = None
     parameters: Optional[Dict[str, Any]] = None
+    retry: Optional[StepRetryConfig] = None

zenml/config/retry_config.py

@@ -0,0 +1,27 @@
+#  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.
+"""Retry configuration for a step."""
+
+from zenml.config.strict_base_model import StrictBaseModel
+
+
+class StepRetryConfig(StrictBaseModel):
+    """Retry configuration for a step.
+
+    Delay is an integer (specified in seconds).
+    """
+
+    max_retries: int = 1
+    delay: int = 0  # in seconds
+    backoff: int = 0

zenml/config/server_config.py

@@ -65,8 +65,16 @@ class ServerConfiguration(BaseModel):
 
     Attributes:
         deployment_type: The type of ZenML server deployment that is running.
-        base_url: The base URL of the ZenML server.
-        root_url_path: The root URL path of the ZenML server.
+        server_url: The URL where the ZenML server API is reachable. Must be
+            configured for features that involve triggering workloads from the
+            ZenML dashboard (e.g., running pipelines). If not specified, the
+            clients will use the same URL used to connect them to the ZenML
+            server.
+        dashboard_url: The URL where the ZenML dashboard is reachable.
+            If not specified, the `server_url` value is used. This should be
+            configured if the dashboard is served from a different URL than the
+            ZenML server.
+        root_url_path: The root URL path for the ZenML API and dashboard.
         auth_scheme: The authentication scheme used by the ZenML server.
         jwt_token_algorithm: The algorithm used to sign and verify JWT tokens.
         jwt_token_issuer: The issuer of the JWT tokens. If not specified, the
@@ -93,10 +101,6 @@ class ServerConfiguration(BaseModel):
             2.0 device authorization request expires.
         device_auth_polling_interval: The polling interval in seconds used to
             poll the OAuth 2.0 device authorization endpoint.
-        dashboard_url: The URL where the ZenML dashboard is hosted. Used to
-            construct the OAuth 2.0 device authorization endpoint. If not set,
-            a partial URL is returned to the client which is used to construct
-            the full URL based on the server's root URL path.
         device_expiration_minutes: The time in minutes that an OAuth 2.0 device is
             allowed to be used to authenticate with the ZenML server. If not
             set or if `jwt_token_expire_minutes` is not set, the devices are
@@ -204,7 +208,7 @@ class ServerConfiguration(BaseModel):
             of the reserved values `enabled`, `yes`, `true`, `on`, the
             `Permissions-Policy` header will be set to the default value
             (`accelerometer=(), camera=(), geolocation=(), gyroscope=(),
-              magnetometer=(), microphone=(), payment=(), usb=()`). If set to
+            magnetometer=(), microphone=(), payment=(), usb=()`). If set to
             one of the reserved values `disabled`, `no`, `none`, `false`, `off`
             or to an empty string, the `Permissions-Policy` header will not be
             included in responses.
@@ -225,7 +229,8 @@ class ServerConfiguration(BaseModel):
     """
 
     deployment_type: ServerDeploymentType = ServerDeploymentType.OTHER
-    base_url: str = ""
+    server_url: Optional[str] = None
+    dashboard_url: Optional[str] = None
     root_url_path: str = ""
     metadata: Dict[str, Any] = {}
     auth_scheme: AuthScheme = AuthScheme.OAUTH2_PASSWORD_BEARER
@@ -245,7 +250,6 @@ class ServerConfiguration(BaseModel):
     device_auth_polling_interval: int = (
         DEFAULT_ZENML_SERVER_DEVICE_AUTH_POLLING
     )
-    dashboard_url: Optional[str] = None
     device_expiration_minutes: Optional[int] = None
     trusted_device_expiration_minutes: Optional[int] = None
 

zenml/config/step_configurations.py

@@ -32,6 +32,7 @@ from zenml.artifacts.external_artifact_config import (
 from zenml.client_lazy_loader import ClientLazyLoader
 from zenml.config.base_settings import BaseSettings, SettingsOrDict
 from zenml.config.constants import DOCKER_SETTINGS_KEY, RESOURCE_SETTINGS_KEY
+from zenml.config.retry_config import StepRetryConfig
 from zenml.config.source import Source, convert_source_validator
 from zenml.config.strict_base_model import StrictBaseModel
 from zenml.logger import get_logger
@@ -137,6 +138,7 @@ class StepConfigurationUpdate(StrictBaseModel):
     failure_hook_source: Optional[Source] = None
     success_hook_source: Optional[Source] = None
     model: Optional[Model] = None
+    retry: Optional[StepRetryConfig] = None
 
     outputs: Mapping[str, PartialArtifactConfiguration] = {}
 

zenml/constants.py

@@ -148,6 +148,7 @@ ENV_ZENML_DISABLE_STEP_LOGS_STORAGE = "ZENML_DISABLE_STEP_LOGS_STORAGE"
 ENV_ZENML_PIPELINE_API_TOKEN_EXPIRES_MINUTES = (
     "ZENML_PIPELINE_API_TOKEN_EXPIRES_MINUTES"
 )
+ENV_ZENML_IGNORE_FAILURE_HOOK = "ZENML_IGNORE_FAILURE_HOOK"
 
 # ZenML Server environment variables
 ENV_ZENML_SERVER_PREFIX = "ZENML_SERVER_"

zenml/integrations/kubernetes/flavors/kubernetes_orchestrator_flavor.py

@@ -43,6 +43,7 @@ class KubernetesOrchestratorSettings(BaseSettings):
             permissions will be created.
         step_pod_service_account_name: Name of the service account to use for the
             step pods. If not provided, the default service account will be used.
+        privileged: If the container should be run in privileged mode.
         pod_settings: Pod settings to apply.
     """
 
@@ -50,6 +51,7 @@ class KubernetesOrchestratorSettings(BaseSettings):
     timeout: int = 0
     service_account_name: Optional[str] = None
     step_pod_service_account_name: Optional[str] = None
+    privileged: bool = False
     pod_settings: Optional[KubernetesPodSettings] = None
 
 

zenml/integrations/kubernetes/orchestrators/kubernetes_orchestrator.py

@@ -391,6 +391,20 @@ class KubernetesOrchestrator(ContainerizedOrchestrator):
         # Authorize pod to run Kubernetes commands inside the cluster.
         service_account_name = self._get_service_account_name(settings)
 
+        if settings.pod_settings:
+            # Remove all settings that specify on which pod to run for the
+            # orchestrator pod. These settings should only be used
+            # for the pods executing the actual steps.
+            pod_settings = settings.pod_settings.copy(
+                update={
+                    "resources": {},
+                    "node_selectors": {},
+                    "affinity": {},
+                    "tolerations": [],
+                }
+            )
+            settings = settings.copy(update={"pod_settings": pod_settings})
+
         # Schedule as CRON job if CRON schedule is given.
         if deployment.schedule:
             if not deployment.schedule.cron_expression:

zenml/integrations/kubernetes/orchestrators/manifest_utils.py

@@ -128,6 +128,9 @@ def build_pod_manifest(
     env = env.copy() if env else {}
     env.setdefault(ENV_ZENML_ENABLE_REPO_INIT_WARNINGS, "False")
 
+    security_context = k8s_client.V1SecurityContext(
+        privileged=settings.privileged
+    )
     container_spec = k8s_client.V1Container(
         name="main",
         image=image_name,
@@ -137,6 +140,7 @@ def build_pod_manifest(
             k8s_client.V1EnvVar(name=name, value=value)
             for name, value in env.items()
         ],
+        security_context=security_context,
     )
 
     pod_spec = k8s_client.V1PodSpec(

zenml/integrations/slack/alerters/slack_alerter.py

@@ -276,8 +276,6 @@ class SlackAlerter(BaseAlerter):
 
         approved = False  # will be modified by handle()
 
-        # breakpoint()
-
         @RTMClient.run_on(event="hello")  # type: ignore
         def post_initial_message(**payload: Any) -> None:
             """Post an initial message in a channel and start listening.

zenml/model/model.py

@@ -60,6 +60,8 @@ class Model(BaseModel):
         to a specific version/stage. If skipped new version will be created.
     save_models_to_registry: Whether to save all ModelArtifacts to Model Registry,
         if available in active stack.
+    model_version_id: The ID of a specific Model Version, if given - it will override
+        `name` and `version` settings. Used mostly internally.
     """
 
     name: str
@@ -73,12 +75,12 @@ class Model(BaseModel):
     tags: Optional[List[str]] = None
     version: Optional[Union[ModelStages, int, str]] = None
     save_models_to_registry: bool = True
+    model_version_id: Optional[UUID] = None
 
     suppress_class_validation_warnings: bool = False
     was_created_in_this_run: bool = False
 
     _model_id: UUID = PrivateAttr(None)
-    _id: UUID = PrivateAttr(None)
     _number: int = PrivateAttr(None)
 
     #########################
@@ -93,16 +95,21 @@ class Model(BaseModel):
                 doesn't exist and can only be read given current
                 config (you used stage name or number as
                 a version name).
+
+        Raises:
+            RuntimeError: if model version doesn't exist and
+                cannot be fetched from the Model Control Plane.
         """
-        if self._id is None:
+        if self.model_version_id is None:
             try:
-                self._get_or_create_model_version()
-            except RuntimeError:
-                logger.info(
+                mv = self._get_or_create_model_version()
+                self.model_version_id = mv.id
+            except RuntimeError as e:
+                raise RuntimeError(
                     f"Version `{self.version}` of `{self.name}` model doesn't exist "
                     "and cannot be fetched from the Model Control Plane."
-                )
-        return self._id
+                ) from e
+        return self.model_version_id
 
     @property
     def model_id(self) -> UUID:
@@ -523,61 +530,82 @@ class Model(BaseModel):
         from zenml.models import ModelRequest
 
         zenml_client = Client()
-        try:
-            model = zenml_client.zen_store.get_model(
-                model_name_or_id=self.name
+        if self.model_version_id:
+            mv = zenml_client.get_model_version(
+                model_version_name_or_number_or_id=self.model_version_id,
             )
-        except KeyError:
-            model_request = ModelRequest(
-                name=self.name,
-                license=self.license,
-                description=self.description,
-                audience=self.audience,
-                use_cases=self.use_cases,
-                limitations=self.limitations,
-                trade_offs=self.trade_offs,
-                ethics=self.ethics,
-                tags=self.tags,
-                user=zenml_client.active_user.id,
-                workspace=zenml_client.active_workspace.id,
-                save_models_to_registry=self.save_models_to_registry,
-            )
-            model_request = ModelRequest.parse_obj(model_request)
+            model = mv.model
+        else:
             try:
-                model = zenml_client.zen_store.create_model(
-                    model=model_request
-                )
-                logger.info(f"New model `{self.name}` was created implicitly.")
-            except EntityExistsError:
                 model = zenml_client.zen_store.get_model(
                     model_name_or_id=self.name
                 )
+            except KeyError:
+                model_request = ModelRequest(
+                    name=self.name,
+                    license=self.license,
+                    description=self.description,
+                    audience=self.audience,
+                    use_cases=self.use_cases,
+                    limitations=self.limitations,
+                    trade_offs=self.trade_offs,
+                    ethics=self.ethics,
+                    tags=self.tags,
+                    user=zenml_client.active_user.id,
+                    workspace=zenml_client.active_workspace.id,
+                    save_models_to_registry=self.save_models_to_registry,
+                )
+                model_request = ModelRequest.parse_obj(model_request)
+                try:
+                    model = zenml_client.zen_store.create_model(
+                        model=model_request
+                    )
+                    logger.info(
+                        f"New model `{self.name}` was created implicitly."
+                    )
+                except EntityExistsError:
+                    model = zenml_client.zen_store.get_model(
+                        model_name_or_id=self.name
+                    )
 
         self._model_id = model.id
         return model
 
-    def _get_model_version(self) -> "ModelVersionResponse":
+    def _get_model_version(
+        self, hydrate: bool = True
+    ) -> "ModelVersionResponse":
         """This method gets a model version from Model Control Plane.
 
+        Args:
+            hydrate: Flag deciding whether to hydrate the output model(s)
+                by including metadata fields in the response.
+
         Returns:
             The model version based on configuration.
         """
         from zenml.client import Client
 
         zenml_client = Client()
-        mv = zenml_client.get_model_version(
-            model_name_or_id=self.name,
-            model_version_name_or_number_or_id=self.version,
-        )
-        if not self._id:
-            self._id = mv.id
+        if self.model_version_id:
+            mv = zenml_client.get_model_version(
+                model_version_name_or_number_or_id=self.model_version_id,
+                hydrate=hydrate,
+            )
+        else:
+            mv = zenml_client.get_model_version(
+                model_name_or_id=self.name,
+                model_version_name_or_number_or_id=self.version,
+                hydrate=hydrate,
+            )
+            self.model_version_id = mv.id
 
         difference: Dict[str, Any] = {}
-        if self.description and mv.description != self.description:
-            difference["description"] = {
-                "config": self.description,
-                "db": mv.description,
-            }
+        if mv.metadata:
+            if self.description and mv.description != self.description:
+                difference["description"] = {
+                    "config": self.description,
+                    "db": mv.description,
+                }
         if self.tags:
             configured_tags = set(self.tags)
             db_tags = {t.name for t in mv.tags}
@@ -656,6 +684,7 @@ class Model(BaseModel):
                         and pipeline_mv.version is not None
                     ):
                         self.version = pipeline_mv.version
+                        self.model_version_id = pipeline_mv.model_version_id
                     else:
                         for step in context.pipeline_run.steps.values():
                             step_mv = step.config.model
@@ -666,8 +695,11 @@ class Model(BaseModel):
                                 and step_mv.version is not None
                             ):
                                 self.version = step_mv.version
+                                self.model_version_id = (
+                                    step_mv.model_version_id
+                                )
                                 break
-            if self.version:
+            if self.version or self.model_version_id:
                 model_version = self._get_model_version()
             else:
                 raise KeyError
@@ -727,7 +759,7 @@ class Model(BaseModel):
 
             logger.info(f"New model version `{self.version}` was created.")
 
-        self._id = model_version.id
+        self.model_version_id = model_version.id
         self._model_id = model_version.model.id
         self._number = model_version.number
         return model_version

zenml/models/v2/core/model_version.py

@@ -331,7 +331,7 @@ class ModelVersionResponse(
             was_created_in_this_run=was_created_in_this_run,
             suppress_class_validation_warnings=suppress_class_validation_warnings,
         )
-        mv._id = self.id
+        mv.model_version_id = self.id
 
         return mv
 

zenml/models/v2/core/pipeline_run.py

@@ -39,6 +39,7 @@ from zenml.models.v2.base.scoped import (
     WorkspaceScopedResponseMetadata,
     WorkspaceScopedResponseResources,
 )
+from zenml.models.v2.core.model_version import ModelVersionResponse
 
 if TYPE_CHECKING:
     from sqlalchemy.sql.elements import BinaryExpression, BooleanClauseList
@@ -197,6 +198,8 @@ class PipelineRunResponseMetadata(WorkspaceScopedResponseMetadata):
 class PipelineRunResponseResources(WorkspaceScopedResponseResources):
     """Class for all resource models associated with the pipeline run entity."""
 
+    model_version: Optional[ModelVersionResponse]
+
 
 class PipelineRunResponse(
     WorkspaceScopedResponse[
@@ -394,6 +397,15 @@ class PipelineRunResponse(
         """
         return self.get_metadata().orchestrator_run_id
 
+    @property
+    def model_version(self) -> Optional[ModelVersionResponse]:
+        """The `model_version` property.
+
+        Returns:
+            the value of the property.
+        """
+        return self.get_resources().model_version
+
 
 # ------------------ Filter Model ------------------
 

zenml/models/v2/core/step_run.py

@@ -30,6 +30,7 @@ from zenml.models.v2.base.scoped import (
     WorkspaceScopedResponseMetadata,
     WorkspaceScopedResponseResources,
 )
+from zenml.models.v2.core.model_version import ModelVersionResponse
 
 if TYPE_CHECKING:
     from zenml.models.v2.core.artifact_version import ArtifactVersionResponse
@@ -219,6 +220,8 @@ class StepRunResponseMetadata(WorkspaceScopedResponseMetadata):
 class StepRunResponseResources(WorkspaceScopedResponseResources):
     """Class for all resource models associated with the step run entity."""
 
+    model_version: Optional[ModelVersionResponse]
+
 
 class StepRunResponse(
     WorkspaceScopedResponse[
@@ -435,6 +438,15 @@ class StepRunResponse(
         """
         return self.get_metadata().run_metadata
 
+    @property
+    def model_version(self) -> Optional[ModelVersionResponse]:
+        """The `model_version` property.
+
+        Returns:
+            the value of the property.
+        """
+        return self.get_resources().model_version
+
 
 # ------------------ Filter Model ------------------
 

zenml/models/v2/misc/server_models.py

@@ -78,11 +78,17 @@ class ServerModel(BaseModel):
     auth_scheme: AuthScheme = Field(
         title="The authentication scheme that the server is using.",
     )
-    base_url: str = Field(
+    server_url: str = Field(
         "",
-        title="The Base URL of the server.",
+        title="The URL where the ZenML server API is reachable. If not "
+        "specified, the clients will use the same URL used to connect them to "
+        "the ZenML server.",
+    )
+    dashboard_url: str = Field(
+        "",
+        title="The URL where the ZenML dashboard is reachable. If "
+        "not specified, the `server_url` value will be used instead.",
     )
-
     analytics_enabled: bool = Field(
         default=True,  # We set a default for migrations from < 0.57.0
         title="Enable server-side analytics.",

zenml/new/pipelines/run_utils.py

@@ -17,7 +17,7 @@ from uuid import UUID
 from zenml import constants
 from zenml.client import Client
 from zenml.config.step_configurations import StepConfigurationUpdate
-from zenml.enums import ExecutionStatus
+from zenml.enums import ExecutionStatus, ModelStages
 from zenml.logger import get_logger
 from zenml.models import (
     PipelineDeploymentBase,
@@ -166,7 +166,13 @@ def _update_new_requesters(
         try:
             model._get_model_version()
             version_existed = key not in new_versions_requested
-        except KeyError:
+        except KeyError as e:
+            if model.version in ModelStages.values():
+                raise KeyError(
+                    f"Unable to get model `{model.name}` using stage "
+                    f"`{model.version}`, please check that the model "
+                    "version in given stage exists before running a pipeline."
+                ) from e
             version_existed = False
     if not version_existed:
         model.was_created_in_this_run = True