zenml-nightly 0.57.1.dev20240523__py3-none-any.whl → 0.58.0.dev20240528__py3-none-any.whl

zenml/orchestrators/utils.py

@@ -15,13 +15,14 @@
 
 import random
 from datetime import datetime
-from typing import TYPE_CHECKING, Dict, Optional
+from typing import TYPE_CHECKING, Dict, List, Optional, Set, Tuple
 from uuid import UUID
 
 from zenml.client import Client
 from zenml.config.global_config import (
     GlobalConfiguration,
 )
+from zenml.config.source import Source
 from zenml.constants import (
     ENV_ZENML_ACTIVE_STACK_ID,
     ENV_ZENML_ACTIVE_WORKSPACE_ID,
@@ -29,8 +30,16 @@ from zenml.constants import (
     PIPELINE_API_TOKEN_EXPIRES_MINUTES,
 )
 from zenml.enums import StoreType
+from zenml.exceptions import StepContextError
+from zenml.model.utils import link_artifact_config_to_model
+from zenml.models.v2.core.step_run import StepRunRequest
+from zenml.new.steps.step_context import get_step_context
 
 if TYPE_CHECKING:
+    from zenml.artifacts.external_artifact_config import (
+        ExternalArtifactConfiguration,
+    )
+    from zenml.model.model import Model
     from zenml.models import PipelineDeploymentResponse
 
 
@@ -148,3 +157,151 @@ def get_run_name(run_name_template: str) -> str:
         raise ValueError("Empty run names are not allowed.")
 
     return run_name
+
+
+def _link_pipeline_run_to_model_from_context(
+    pipeline_run_id: "UUID", model: Optional["Model"] = None
+) -> None:
+    """Links the pipeline run to the model version using artifacts data.
+
+    Args:
+        pipeline_run_id: The ID of the current pipeline run.
+        model: Model configured in the step
+    """
+    from zenml.models import ModelVersionPipelineRunRequest
+
+    if not model:
+        model_id, model_version_id = _get_model_versions_from_config()
+    else:
+        model_id, model_version_id = model.model_id, model.id
+
+    if model_id and model_version_id:
+        Client().zen_store.create_model_version_pipeline_run_link(
+            ModelVersionPipelineRunRequest(
+                user=Client().active_user.id,
+                workspace=Client().active_workspace.id,
+                pipeline_run=pipeline_run_id,
+                model=model_id,
+                model_version=model_version_id,
+            )
+        )
+
+
+def _get_model_versions_from_config() -> Tuple[Optional[UUID], Optional[UUID]]:
+    """Gets the model versions from the step model version.
+
+    Returns:
+        Tuple of (model_id, model_version_id).
+    """
+    try:
+        mc = get_step_context().model
+        return mc.model_id, mc.id
+    except StepContextError:
+        return None, None
+
+
+def _link_cached_artifacts_to_model(
+    model_from_context: Optional["Model"],
+    step_run: StepRunRequest,
+    step_source: Source,
+) -> None:
+    """Links the output artifacts of the cached step to the model version in Control Plane.
+
+    Args:
+        model_from_context: The model version of the current step.
+        step_run: The step to run.
+        step_source: The source of the step.
+    """
+    from zenml.artifacts.artifact_config import ArtifactConfig
+    from zenml.steps.base_step import BaseStep
+    from zenml.steps.utils import parse_return_type_annotations
+
+    step_instance = BaseStep.load_from_source(step_source)
+    output_annotations = parse_return_type_annotations(
+        step_instance.entrypoint
+    )
+    for output_name_, output_id in step_run.outputs.items():
+        artifact_config_ = None
+        if output_name_ in output_annotations:
+            annotation = output_annotations.get(output_name_, None)
+            if annotation and annotation.artifact_config is not None:
+                artifact_config_ = annotation.artifact_config.copy()
+        # no artifact config found or artifact was produced by `save_artifact`
+        # inside the step body, so was never in annotations
+        if artifact_config_ is None:
+            artifact_config_ = ArtifactConfig(name=output_name_)
+
+        link_artifact_config_to_model(
+            artifact_config=artifact_config_,
+            model=model_from_context,
+            artifact_version_id=output_id,
+        )
+
+
+def _link_pipeline_run_to_model_from_artifacts(
+    pipeline_run_id: UUID,
+    artifact_names: List[str],
+    external_artifacts: List["ExternalArtifactConfiguration"],
+) -> None:
+    """Links the pipeline run to the model version using artifacts data.
+
+    Args:
+        pipeline_run_id: The ID of the current pipeline run.
+        artifact_names: The name of the published output artifacts.
+        external_artifacts: The external artifacts of the step.
+    """
+    from zenml.models import ModelVersionPipelineRunRequest
+
+    models = _get_model_versions_from_artifacts(artifact_names)
+    client = Client()
+
+    # Add models from external artifacts
+    for external_artifact in external_artifacts:
+        if external_artifact.model:
+            models.add(
+                (
+                    external_artifact.model.model_id,
+                    external_artifact.model.id,
+                )
+            )
+
+    for model in models:
+        client.zen_store.create_model_version_pipeline_run_link(
+            ModelVersionPipelineRunRequest(
+                user=client.active_user.id,
+                workspace=client.active_workspace.id,
+                pipeline_run=pipeline_run_id,
+                model=model[0],
+                model_version=model[1],
+            )
+        )
+
+
+def _get_model_versions_from_artifacts(
+    artifact_names: List[str],
+) -> Set[Tuple[UUID, UUID]]:
+    """Gets the model versions from the artifacts.
+
+    Args:
+        artifact_names: The names of the published output artifacts.
+
+    Returns:
+        Set of tuples of (model_id, model_version_id).
+    """
+    models = set()
+    for artifact_name in artifact_names:
+        artifact_config = (
+            get_step_context()._get_output(artifact_name).artifact_config
+        )
+        if artifact_config is not None:
+            if (model := artifact_config._model) is not None:
+                model_version_response = model._get_or_create_model_version()
+                models.add(
+                    (
+                        model_version_response.model.id,
+                        model_version_response.id,
+                    )
+                )
+            else:
+                break
+    return models

zenml/steps/base_step.py

@@ -37,6 +37,7 @@ from typing import (
 from pydantic import BaseModel, Extra, ValidationError
 
 from zenml.client_lazy_loader import ClientLazyLoader
+from zenml.config.retry_config import StepRetryConfig
 from zenml.config.source import Source
 from zenml.constants import STEP_SOURCE_PARAMETER_NAME
 from zenml.exceptions import MissingStepParameterError, StepInterfaceError
@@ -140,6 +141,7 @@ class BaseStep(metaclass=BaseStepMeta):
         on_failure: Optional["HookSpecification"] = None,
         on_success: Optional["HookSpecification"] = None,
         model: Optional["Model"] = None,
+        retry: Optional[StepRetryConfig] = None,
         **kwargs: Any,
     ) -> None:
         """Initializes a step.
@@ -169,6 +171,7 @@ class BaseStep(metaclass=BaseStepMeta):
                 be a function with no arguments, or a source path to such a
                 function (e.g. `module.my_function`).
             model: configuration of the model version in the Model Control Plane.
+            retry: Configuration for retrying the step in case of failure.
             **kwargs: Keyword arguments passed to the step.
         """
         from zenml.config.step_configurations import PartialStepConfiguration
@@ -188,36 +191,36 @@ class BaseStep(metaclass=BaseStepMeta):
                 # We therefore disable caching unless it is explicitly enabled
                 enable_cache = False
                 logger.debug(
-                    "Step '%s': Step context required and caching not "
+                    "Step `%s`: Step context required and caching not "
                     "explicitly enabled.",
                     name,
                 )
 
         logger.debug(
-            "Step '%s': Caching %s.",
+            "Step `%s`: Caching %s.",
             name,
             "enabled" if enable_cache is not False else "disabled",
         )
         logger.debug(
-            "Step '%s': Artifact metadata %s.",
+            "Step `%s`: Artifact metadata %s.",
             name,
             "enabled" if enable_artifact_metadata is not False else "disabled",
         )
         logger.debug(
-            "Step '%s': Artifact visualization %s.",
+            "Step `%s`: Artifact visualization %s.",
             name,
             "enabled"
             if enable_artifact_visualization is not False
             else "disabled",
         )
         logger.debug(
-            "Step '%s': logs %s.",
+            "Step `%s`: logs %s.",
             name,
             "enabled" if enable_step_logs is not False else "disabled",
         )
         if model is not None:
             logger.debug(
-                "Step '%s': Is in Model context %s.",
+                "Step `%s`: Is in Model context %s.",
                 name,
                 {
                     "model": model.name,
@@ -242,6 +245,7 @@ class BaseStep(metaclass=BaseStepMeta):
             on_failure=on_failure,
             on_success=on_success,
             model=model,
+            retry=retry,
         )
         self._verify_and_apply_init_params(*args, **kwargs)
 
@@ -475,7 +479,7 @@ class BaseStep(metaclass=BaseStepMeta):
             bound_args = signature.bind_partial(*args, **kwargs)
         except TypeError as e:
             raise StepInterfaceError(
-                f"Wrong arguments when calling step '{self.name}': {e}"
+                f"Wrong arguments when calling step `{self.name}`: {e}"
             ) from e
 
         artifacts = {}
@@ -696,6 +700,7 @@ class BaseStep(metaclass=BaseStepMeta):
         on_success: Optional["HookSpecification"] = None,
         model: Optional["Model"] = None,
         merge: bool = True,
+        retry: Optional[StepRetryConfig] = None,
     ) -> T:
         """Configures the step.
 
@@ -738,6 +743,7 @@ class BaseStep(metaclass=BaseStepMeta):
                 configurations. If `False` the given configurations will
                 overwrite all existing ones. See the general description of this
                 method for an example.
+            retry: Configuration for retrying the step in case of failure.
 
         Returns:
             The step instance that this method was called on.
@@ -807,6 +813,7 @@ class BaseStep(metaclass=BaseStepMeta):
                 "failure_hook_source": failure_hook_source,
                 "success_hook_source": success_hook_source,
                 "model": model,
+                "retry": retry,
             }
         )
         config = StepConfigurationUpdate(**values)
@@ -1013,7 +1020,7 @@ To avoid this consider setting step parameters only in one place (config or code
             if output_name not in allowed_output_names:
                 raise StepInterfaceError(
                     f"Got unexpected materializers for non-existent "
-                    f"output '{output_name}' in step '{self.name}'. "
+                    f"output '{output_name}' in step `{self.name}`. "
                     f"Only materializers for the outputs "
                     f"{allowed_output_names} of this step can"
                     f" be registered."
@@ -1026,7 +1033,7 @@ To avoid this consider setting step parameters only in one place (config or code
                     ):
                         raise StepInterfaceError(
                             f"Materializer source `{source}` "
-                            f"for output '{output_name}' of step '{self.name}' "
+                            f"for output '{output_name}' of step `{self.name}` "
                             "does not resolve to a `BaseMaterializer` subclass."
                         )
 

zenml/utils/dashboard_utils.py

@@ -42,8 +42,7 @@ def is_cloud_server(server_info: ServerModel) -> bool:
     """
     return (
         "organization_id" in server_info.metadata
-        and server_info.base_url is not None
-        and "cloud.zenml.io" in server_info.base_url
+        and "cloud.zenml.io" in server_info.server_url
     )
 
 
@@ -59,7 +58,7 @@ def get_cloud_dashboard_url() -> Optional[str]:
         server_info = client.zen_store.get_store_info()
 
         if is_cloud_server(server_info):
-            return server_info.base_url
+            return server_info.dashboard_url
 
     return None
 
@@ -80,11 +79,8 @@ def get_server_dashboard_url() -> Tuple[Optional[str], bool]:
         else:
             suffix = ""
 
-        if server_info.base_url and not is_cloud_server(server_info):
-            # For cloud the base URL is set to the cloud dashboard, but in this
-            # function we want the URL of the dashboard provided by the server
-            # deployment, which for cloud is just the ZenStore URL
-            url = server_info.base_url
+        if server_info.server_url:
+            url = server_info.server_url
         else:
             url = client.zen_store.url
 

zenml/zen_server/dashboard/assets/{404-DtHkZQ2t.js → 404-DIgYOHj5.js}

@@ -1 +1 @@
-import{j as e}from"./@tanstack-GFKzItNf.js";import{h as s,r as t}from"./index-UnEiujvd.js";import{E as r}from"./EmptyState-CPJglRv0.js";import{H as a}from"./help-hEyYkC3q.js";import{L as o}from"./@react-router-CCpyfPW-.js";import"./@radix-BrEjOMvl.js";import"./@reactflow-BfgHyCG7.js";function d(){return e.jsx("div",{className:"flex min-h-screen w-full flex-col",children:e.jsx(r,{icon:e.jsx(a,{className:"h-[120px] w-[120px] fill-neutral-300"}),children:e.jsxs("div",{className:"text-center",children:[e.jsx("h1",{className:"mb-2 text-display-xs font-semibold",children:"We can't find the page you are looking for"}),e.jsx("p",{className:"text-lg text-theme-text-secondary",children:"You can try typing a different URL or we can bring you back to your Homepage."}),e.jsx("div",{className:"mt-5 flex justify-center",children:e.jsx(s,{size:"md",asChild:!0,children:e.jsx(o,{className:"w-min self-center whitespace-nowrap",to:t.home,children:e.jsx("span",{className:"px-0.5",children:"Go to Home"})})})})]})})})}export{d as default};
+import{j as e}from"./@tanstack-GFKzItNf.js";import{h as s,r as t}from"./index-DKgtNG83.js";import{E as r}from"./EmptyState-CPJglRv0.js";import{H as a}from"./help-hEyYkC3q.js";import{L as o}from"./@react-router-CCpyfPW-.js";import"./@radix-BrEjOMvl.js";import"./@reactflow-CF_wPv-n.js";function d(){return e.jsx("div",{className:"flex min-h-screen w-full flex-col",children:e.jsx(r,{icon:e.jsx(a,{className:"h-[120px] w-[120px] fill-neutral-300"}),children:e.jsxs("div",{className:"text-center",children:[e.jsx("h1",{className:"mb-2 text-display-xs font-semibold",children:"We can't find the page you are looking for"}),e.jsx("p",{className:"text-lg text-theme-text-secondary",children:"You can try typing a different URL or we can bring you back to your Homepage."}),e.jsx("div",{className:"mt-5 flex justify-center",children:e.jsx(s,{size:"md",asChild:!0,children:e.jsx(o,{className:"w-min self-center whitespace-nowrap",to:t.home,children:e.jsx("span",{className:"px-0.5",children:"Go to Home"})})})})]})})})}export{d as default};