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

zenml/models/v2/core/model.py

@@ -82,6 +82,7 @@ class ModelRequest(WorkspaceScopedRequest):
     )
     tags: Optional[List[str]] = Field(
         title="Tags associated with the model",
+        default=None,
     )
     save_models_to_registry: bool = Field(
         title="Whether to save all ModelArtifacts to Model Registry",
@@ -367,7 +368,9 @@ class ModelFilter(WorkspaceScopedTaggableFilter):
             user_filter = and_(
                 ModelSchema.user_id == UserSchema.id,
                 self.generate_name_or_id_query_conditions(
-                    value=self.user, table=UserSchema
+                    value=self.user,
+                    table=UserSchema,
+                    additional_columns=["full_name"],
                 ),
             )
             custom_filters.append(user_filter)

zenml/models/v2/core/model_version.py

@@ -325,14 +325,11 @@ class ModelVersionResponse(
     # Helper functions
     def to_model_class(
         self,
-        was_created_in_this_run: bool = False,
-        suppress_class_validation_warnings: bool = False,
+        suppress_class_validation_warnings: bool = True,
     ) -> "Model":
         """Convert response model to Model object.
 
         Args:
-            was_created_in_this_run: Whether model version was created during
-                the current run.
             suppress_class_validation_warnings: internally used to suppress
                 repeated warnings.
 
@@ -352,7 +349,6 @@ class ModelVersionResponse(
             ethics=self.model.ethics,
             tags=[t.name for t in self.tags],
             version=self.name,
-            was_created_in_this_run=was_created_in_this_run,
             suppress_class_validation_warnings=suppress_class_validation_warnings,
             model_version_id=self.id,
         )
@@ -664,7 +660,9 @@ class ModelVersionFilter(WorkspaceScopedTaggableFilter):
             user_filter = and_(
                 ModelVersionSchema.user_id == UserSchema.id,
                 self.generate_name_or_id_query_conditions(
-                    value=self.user, table=UserSchema
+                    value=self.user,
+                    table=UserSchema,
+                    additional_columns=["full_name"],
                 ),
             )
             custom_filters.append(user_filter)

zenml/models/v2/core/model_version_artifact.py

@@ -294,7 +294,9 @@ class ModelVersionArtifactFilter(WorkspaceScopedFilter):
                 == ArtifactVersionSchema.id,
                 ArtifactVersionSchema.user_id == UserSchema.id,
                 self.generate_name_or_id_query_conditions(
-                    value=self.user, table=UserSchema
+                    value=self.user,
+                    table=UserSchema,
+                    additional_columns=["full_name"],
                 ),
             )
             custom_filters.append(user_filter)

zenml/models/v2/core/model_version_pipeline_run.py

@@ -218,7 +218,9 @@ class ModelVersionPipelineRunFilter(WorkspaceScopedFilter):
                 == PipelineRunSchema.id,
                 PipelineRunSchema.user_id == UserSchema.id,
                 self.generate_name_or_id_query_conditions(
-                    value=self.user, table=UserSchema
+                    value=self.user,
+                    table=UserSchema,
+                    additional_columns=["full_name"],
                 ),
             )
             custom_filters.append(user_filter)

zenml/models/v2/core/pipeline.py

@@ -364,7 +364,9 @@ class PipelineFilter(WorkspaceScopedTaggableFilter):
             user_filter = and_(
                 PipelineSchema.user_id == UserSchema.id,
                 self.generate_name_or_id_query_conditions(
-                    value=self.user, table=UserSchema
+                    value=self.user,
+                    table=UserSchema,
+                    additional_columns=["full_name"],
                 ),
             )
             custom_filters.append(user_filter)

zenml/models/v2/core/pipeline_run.py

@@ -31,7 +31,6 @@ from zenml.config.pipeline_configurations import PipelineConfiguration
 from zenml.constants import STR_FIELD_MAX_LENGTH
 from zenml.enums import ExecutionStatus
 from zenml.models.v2.base.scoped import (
-    WorkspaceScopedFilter,
     WorkspaceScopedRequest,
     WorkspaceScopedResponse,
     WorkspaceScopedResponseBody,
@@ -575,7 +574,7 @@ class PipelineRunFilter(WorkspaceScopedTaggableFilter):
     """Model to enable advanced filtering of all Workspaces."""
 
     FILTER_EXCLUDE_FIELDS: ClassVar[List[str]] = [
-        *WorkspaceScopedFilter.FILTER_EXCLUDE_FIELDS,
+        *WorkspaceScopedTaggableFilter.FILTER_EXCLUDE_FIELDS,
         "unlisted",
         "code_repository_id",
         "build_id",
@@ -587,6 +586,7 @@ class PipelineRunFilter(WorkspaceScopedTaggableFilter):
         "stack",
         "code_repository",
         "model",
+        "stack_component",
         "pipeline_name",
         "templatable",
     ]
@@ -689,6 +689,10 @@ class PipelineRunFilter(WorkspaceScopedTaggableFilter):
         default=None,
         description="Name/ID of the model associated with the run.",
     )
+    stack_component: Optional[Union[UUID, str]] = Field(
+        default=None,
+        description="Name/ID of the stack component associated with the run.",
+    )
     templatable: Optional[bool] = Field(
         default=None, description="Whether the run is templatable."
     )
@@ -717,6 +721,8 @@ class PipelineRunFilter(WorkspaceScopedTaggableFilter):
             PipelineRunSchema,
             PipelineSchema,
             ScheduleSchema,
+            StackComponentSchema,
+            StackCompositionSchema,
             StackSchema,
             UserSchema,
         )
@@ -773,7 +779,9 @@ class PipelineRunFilter(WorkspaceScopedTaggableFilter):
             user_filter = and_(
                 PipelineRunSchema.user_id == UserSchema.id,
                 self.generate_name_or_id_query_conditions(
-                    value=self.user, table=UserSchema
+                    value=self.user,
+                    table=UserSchema,
+                    additional_columns=["full_name"],
                 ),
             )
             custom_filters.append(user_filter)
@@ -822,6 +830,19 @@ class PipelineRunFilter(WorkspaceScopedTaggableFilter):
             )
             custom_filters.append(model_filter)
 
+        if self.stack_component:
+            component_filter = and_(
+                PipelineRunSchema.deployment_id == PipelineDeploymentSchema.id,
+                PipelineDeploymentSchema.stack_id == StackSchema.id,
+                StackSchema.id == StackCompositionSchema.stack_id,
+                StackCompositionSchema.component_id == StackComponentSchema.id,
+                self.generate_name_or_id_query_conditions(
+                    value=self.stack_component,
+                    table=StackComponentSchema,
+                ),
+            )
+            custom_filters.append(component_filter)
+
         if self.pipeline_name:
             pipeline_name_filter = and_(
                 PipelineRunSchema.pipeline_id == PipelineSchema.id,

zenml/models/v2/core/run_metadata.py

@@ -21,7 +21,6 @@ from pydantic import Field, field_validator
 from zenml.constants import STR_FIELD_MAX_LENGTH, TEXT_FIELD_MAX_LENGTH
 from zenml.enums import MetadataResourceTypes
 from zenml.metadata.metadata_types import MetadataType, MetadataTypeEnum
-from zenml.model.model import Model
 from zenml.models.v2.base.scoped import (
     WorkspaceScopedFilter,
     WorkspaceScopedRequest,
@@ -239,7 +238,8 @@ class LazyRunMetadataResponse(RunMetadataResponse):
     lazy_load_artifact_name: Optional[str] = None
     lazy_load_artifact_version: Optional[str] = None
     lazy_load_metadata_name: Optional[str] = None
-    lazy_load_model: Model
+    lazy_load_model_name: str
+    lazy_load_model_version: Optional[str] = None
 
     def get_body(self) -> None:  # type: ignore[override]
         """Protects from misuse of the lazy loader.

zenml/models/v2/core/run_template.py

@@ -413,7 +413,9 @@ class RunTemplateFilter(WorkspaceScopedTaggableFilter):
             user_filter = and_(
                 RunTemplateSchema.user_id == UserSchema.id,
                 self.generate_name_or_id_query_conditions(
-                    value=self.user, table=UserSchema
+                    value=self.user,
+                    table=UserSchema,
+                    additional_columns=["full_name"],
                 ),
             )
             custom_filters.append(user_filter)

zenml/models/v2/core/stack.py

@@ -228,7 +228,7 @@ class StackResponse(
             component_dict = dict(
                 name=component.name,
                 type=str(component.type),
-                flavor=component.flavor,
+                flavor=component.flavor_name,
             )
             configuration = json.loads(
                 component.get_metadata().model_dump_json(
@@ -255,7 +255,9 @@ class StackResponse(
             Dict of analytics metadata.
         """
         metadata = super().get_analytics_metadata()
-        metadata.update({ct: c[0].flavor for ct, c in self.components.items()})
+        metadata.update(
+            {ct: c[0].flavor_name for ct, c in self.components.items()}
+        )
 
         if self.labels is not None:
             metadata.update(
@@ -381,7 +383,9 @@ class StackFilter(WorkspaceScopedFilter):
             user_filter = and_(
                 StackSchema.user_id == UserSchema.id,
                 self.generate_name_or_id_query_conditions(
-                    value=self.user, table=UserSchema
+                    value=self.user,
+                    table=UserSchema,
+                    additional_columns=["full_name"],
                 ),
             )
             custom_filters.append(user_filter)

zenml/models/v2/core/step_run.py

@@ -14,7 +14,7 @@
 """Models representing steps runs."""
 
 from datetime import datetime
-from typing import TYPE_CHECKING, Dict, List, Optional, Union
+from typing import TYPE_CHECKING, ClassVar, Dict, List, Optional, Union
 from uuid import UUID
 
 from pydantic import BaseModel, ConfigDict, Field
@@ -33,6 +33,8 @@ from zenml.models.v2.base.scoped import (
 from zenml.models.v2.core.model_version import ModelVersionResponse
 
 if TYPE_CHECKING:
+    from sqlalchemy.sql.elements import ColumnElement
+
     from zenml.models.v2.core.artifact_version import ArtifactVersionResponse
     from zenml.models.v2.core.logs import (
         LogsRequest,
@@ -156,6 +158,14 @@ class StepRunResponseBody(WorkspaceScopedResponseBody):
     """Response body for step runs."""
 
     status: ExecutionStatus = Field(title="The status of the step.")
+    start_time: Optional[datetime] = Field(
+        title="The start time of the step run.",
+        default=None,
+    )
+    end_time: Optional[datetime] = Field(
+        title="The end time of the step run.",
+        default=None,
+    )
     inputs: Dict[str, "ArtifactVersionResponse"] = Field(
         title="The input artifact versions of the step run.",
         default={},
@@ -201,16 +211,6 @@ class StepRunResponseMetadata(WorkspaceScopedResponseMetadata):
         max_length=TEXT_FIELD_MAX_LENGTH,
     )
 
-    # Timestamps
-    start_time: Optional[datetime] = Field(
-        title="The start time of the step run.",
-        default=None,
-    )
-    end_time: Optional[datetime] = Field(
-        title="The end time of the step run.",
-        default=None,
-    )
-
     # References
     logs: Optional["LogsResponse"] = Field(
         title="Logs associated with this step run.",
@@ -409,7 +409,7 @@ class StepRunResponse(
         Returns:
             the value of the property.
         """
-        return self.get_metadata().start_time
+        return self.get_body().start_time
 
     @property
     def end_time(self) -> Optional[datetime]:
@@ -418,7 +418,7 @@ class StepRunResponse(
         Returns:
             the value of the property.
         """
-        return self.get_metadata().end_time
+        return self.get_body().end_time
 
     @property
     def logs(self) -> Optional["LogsResponse"]:
@@ -490,6 +490,11 @@ class StepRunResponse(
 class StepRunFilter(WorkspaceScopedFilter):
     """Model to enable advanced filtering of step runs."""
 
+    FILTER_EXCLUDE_FIELDS: ClassVar[List[str]] = [
+        *WorkspaceScopedFilter.FILTER_EXCLUDE_FIELDS,
+        "model",
+    ]
+
     name: Optional[str] = Field(
         default=None,
         description="Name of the step run",
@@ -521,6 +526,11 @@ class StepRunFilter(WorkspaceScopedFilter):
         description="Pipeline run of this step run",
         union_mode="left_to_right",
     )
+    deployment_id: Optional[Union[UUID, str]] = Field(
+        default=None,
+        description="Deployment of this step run",
+        union_mode="left_to_right",
+    )
     original_step_run_id: Optional[Union[UUID, str]] = Field(
         default=None,
         description="Original id for this step run",
@@ -538,8 +548,42 @@ class StepRunFilter(WorkspaceScopedFilter):
     )
     model_version_id: Optional[Union[UUID, str]] = Field(
         default=None,
-        description="Model version associated with the pipeline run.",
+        description="Model version associated with the step run.",
         union_mode="left_to_right",
     )
+    model: Optional[Union[UUID, str]] = Field(
+        default=None,
+        description="Name/ID of the model associated with the step run.",
+    )
 
     model_config = ConfigDict(protected_namespaces=())
+
+    def get_custom_filters(
+        self,
+    ) -> List["ColumnElement[bool]"]:
+        """Get custom filters.
+
+        Returns:
+            A list of custom filters.
+        """
+        custom_filters = super().get_custom_filters()
+
+        from sqlmodel import and_
+
+        from zenml.zen_stores.schemas import (
+            ModelSchema,
+            ModelVersionSchema,
+            StepRunSchema,
+        )
+
+        if self.model:
+            model_filter = and_(
+                StepRunSchema.model_version_id == ModelVersionSchema.id,
+                ModelVersionSchema.model_id == ModelSchema.id,
+                self.generate_name_or_id_query_conditions(
+                    value=self.model, table=ModelSchema
+                ),
+            )
+            custom_filters.append(model_filter)
+
+        return custom_filters

zenml/models/v2/misc/auth_models.py

@@ -13,10 +13,10 @@
 #  permissions and limitations under the License.
 """Models representing OAuth2 requests and responses."""
 
-from typing import Optional
+from typing import Any, Dict, Optional
 from uuid import UUID
 
-from pydantic import BaseModel
+from pydantic import BaseModel, ConfigDict
 
 from zenml.enums import OAuthGrantTypes
 
@@ -29,6 +29,7 @@ class OAuthDeviceAuthorizationRequest(BaseModel):
     """OAuth2 device authorization grant request."""
 
     client_id: UUID
+    device_id: Optional[UUID] = None
 
 
 class OAuthDeviceVerificationRequest(BaseModel):
@@ -119,6 +120,14 @@ class OAuthTokenResponse(BaseModel):
     expires_in: Optional[int] = None
     refresh_token: Optional[str] = None
     scope: Optional[str] = None
+    cookie_name: Optional[str] = None
+    device_id: Optional[UUID] = None
+    device_metadata: Optional[Dict[str, Any]] = None
+
+    model_config = ConfigDict(
+        # Allow extra attributes to allow compatibility with different versions
+        extra="allow",
+    )
 
 
 class OAuthRedirectResponse(BaseModel):

zenml/models/v2/misc/server_models.py

@@ -52,6 +52,8 @@ class ServerModel(BaseModel):
 
     id: UUID = Field(default_factory=uuid4, title="The unique server id.")
 
+    name: Optional[str] = Field(None, title="The name of the ZenML server.")
+
     version: str = Field(
         title="The ZenML version that the server is running.",
     )
@@ -99,10 +101,6 @@ class ServerModel(BaseModel):
         {},
         title="The metadata associated with the server.",
     )
-    use_legacy_dashboard: bool = Field(
-        False,
-        title="Flag to indicate whether the server is using the legacy dashboard.",
-    )
 
     last_user_activity: Optional[datetime] = Field(
         None,

zenml/orchestrators/base_orchestrator.py

@@ -18,6 +18,10 @@ from typing import TYPE_CHECKING, Any, Dict, Iterator, Optional, Type, cast
 
 from pydantic import model_validator
 
+from zenml.constants import (
+    ENV_ZENML_PREVENT_CLIENT_SIDE_CACHING,
+    handle_bool_env_var,
+)
 from zenml.enums import ExecutionStatus, StackComponentType
 from zenml.logger import get_logger
 from zenml.metadata.metadata_types import MetadataType
@@ -162,9 +166,8 @@ class BaseOrchestrator(StackComponent, ABC):
             environment: Environment variables to set in the orchestration
                 environment. These don't need to be set if running locally.
 
-        Returns:
-            The optional return value from this method will be returned by the
-            `pipeline_instance.run()` call when someone is running a pipeline.
+        Yields:
+            Metadata for the pipeline run.
         """
 
     def run(
@@ -172,7 +175,7 @@ class BaseOrchestrator(StackComponent, ABC):
         deployment: "PipelineDeploymentResponse",
         stack: "Stack",
         placeholder_run: Optional["PipelineRunResponse"] = None,
-    ) -> Any:
+    ) -> None:
         """Runs a pipeline on a stack.
 
         Args:
@@ -185,6 +188,40 @@ class BaseOrchestrator(StackComponent, ABC):
 
         environment = get_config_environment_vars(deployment=deployment)
 
+        prevent_client_side_caching = handle_bool_env_var(
+            ENV_ZENML_PREVENT_CLIENT_SIDE_CACHING, default=False
+        )
+
+        if (
+            placeholder_run
+            and not deployment.schedule
+            and not prevent_client_side_caching
+        ):
+            from zenml.orchestrators import step_run_utils
+
+            cached_invocations = step_run_utils.create_cached_step_runs(
+                deployment=deployment,
+                pipeline_run=placeholder_run,
+                stack=stack,
+            )
+
+            for invocation_id in cached_invocations:
+                # Remove the cached step invocations from the deployment so
+                # the orchestrator does not try to run them
+                deployment.step_configurations.pop(invocation_id)
+
+            for step in deployment.step_configurations.values():
+                for invocation_id in cached_invocations:
+                    if invocation_id in step.spec.upstream_steps:
+                        step.spec.upstream_steps.remove(invocation_id)
+
+            if len(deployment.step_configurations) == 0:
+                # All steps were cached, we update the pipeline run status and
+                # don't actually use the orchestrator to run the pipeline
+                self._cleanup_run()
+                logger.info("All steps of the pipeline run were cached.")
+                return
+
         try:
             if metadata_iterator := self.prepare_or_run_pipeline(
                 deployment=deployment,

zenml/orchestrators/input_utils.py

@@ -22,18 +22,18 @@ from zenml.exceptions import InputResolutionError
 from zenml.utils import pagination_utils
 
 if TYPE_CHECKING:
-    from zenml.models import ArtifactVersionResponse
+    from zenml.models import ArtifactVersionResponse, PipelineRunResponse
 
 
 def resolve_step_inputs(
     step: "Step",
-    run_id: UUID,
+    pipeline_run: "PipelineRunResponse",
 ) -> Tuple[Dict[str, "ArtifactVersionResponse"], List[UUID]]:
     """Resolves inputs for the current step.
 
     Args:
         step: The step for which to resolve the inputs.
-        run_id: The ID of the current pipeline run.
+        pipeline_run: The current pipeline run.
 
     Raises:
         InputResolutionError: If input resolving failed due to a missing
@@ -50,7 +50,7 @@ def resolve_step_inputs(
     current_run_steps = {
         run_step.name: run_step
         for run_step in pagination_utils.depaginate(
-            Client().list_run_steps, pipeline_run_id=run_id
+            Client().list_run_steps, pipeline_run_id=pipeline_run.id
         )
     }
 
@@ -83,44 +83,60 @@ def resolve_step_inputs(
         )
 
     for name, config_ in step.config.model_artifacts_or_metadata.items():
-        issue_found = False
+        err_msg = ""
         try:
-            if config_.metadata_name is None and config_.artifact_name:
-                if artifact_ := config_.model.get_artifact(
-                    config_.artifact_name, config_.artifact_version
-                ):
-                    input_artifacts[name] = artifact_
-                else:
-                    issue_found = True
-            elif config_.artifact_name is None and config_.metadata_name:
-                # metadata values should go directly in parameters, as primitive types
-                step.config.parameters[name] = config_.model.run_metadata[
-                    config_.metadata_name
-                ].value
-            elif config_.metadata_name and config_.artifact_name:
+            context_model_version = config_._get_model_response(
+                pipeline_run=pipeline_run
+            )
+        except RuntimeError as e:
+            err_msg = str(e)
+        else:
+            if (
+                config_.artifact_name is None
+                and config_.metadata_name
+                and context_model_version.run_metadata is not None
+            ):
                 # metadata values should go directly in parameters, as primitive types
-                if artifact_ := config_.model.get_artifact(
-                    config_.artifact_name, config_.artifact_version
-                ):
-                    step.config.parameters[name] = artifact_.run_metadata[
+                step.config.parameters[name] = (
+                    context_model_version.run_metadata[
                         config_.metadata_name
                     ].value
-                else:
-                    issue_found = True
+                )
+            elif config_.artifact_name is None:
+                err_msg = (
+                    "Cannot load artifact from model version, "
+                    "no artifact name specified."
+                )
             else:
-                issue_found = True
-        except KeyError:
-            issue_found = True
-
-        if issue_found:
+                if artifact_ := context_model_version.get_artifact(
+                    config_.artifact_name, config_.artifact_version
+                ):
+                    if config_.metadata_name is None:
+                        input_artifacts[name] = artifact_
+                    elif config_.metadata_name:
+                        # metadata values should go directly in parameters, as primitive types
+                        try:
+                            step.config.parameters[name] = (
+                                artifact_.run_metadata[
+                                    config_.metadata_name
+                                ].value
+                            )
+                        except KeyError:
+                            err_msg = (
+                                f"Artifact run metadata `{config_.metadata_name}` "
+                                "could not be found in artifact "
+                                f"`{config_.artifact_name}::{config_.artifact_version}`."
+                            )
+                else:
+                    err_msg = (
+                        f"Artifact `{config_.artifact_name}::{config_.artifact_version}` "
+                        f"not found in model `{context_model_version.model.name}` "
+                        f"version `{context_model_version.name}`."
+                    )
+        if err_msg:
             raise ValueError(
-                "Cannot fetch requested information from model "
-                f"`{config_.model.name}` version "
-                f"`{config_.model.version}` given artifact "
-                f"`{config_.artifact_name}`, artifact version "
-                f"`{config_.artifact_version}`, and metadata "
-                f"key `{config_.metadata_name}` passed into "
-                f"the step `{step.config.name}`."
+                f"Failed to lazy load model version data in step `{step.config.name}`: "
+                + err_msg
             )
     for name, cll_ in step.config.client_lazy_loaders.items():
         value_ = cll_.evaluate()