zenml-nightly 0.71.0.dev20241213__py3-none-any.whl → 0.71.0.dev20241214__py3-none-any.whl

zenml/models/v2/core/schedule.py

@@ -279,16 +279,6 @@ class ScheduleResponse(
 class ScheduleFilter(WorkspaceScopedFilter):
     """Model to enable advanced filtering of all Users."""
 
-    workspace_id: Optional[Union[UUID, str]] = Field(
-        default=None,
-        description="Workspace scope of the schedule.",
-        union_mode="left_to_right",
-    )
-    user_id: Optional[Union[UUID, str]] = Field(
-        default=None,
-        description="User that created the schedule",
-        union_mode="left_to_right",
-    )
     pipeline_id: Optional[Union[UUID, str]] = Field(
         default=None,
         description="Pipeline that the schedule is attached to.",

zenml/models/v2/core/secret.py

@@ -15,7 +15,6 @@
 
 from datetime import datetime
 from typing import Any, ClassVar, Dict, List, Optional, Union
-from uuid import UUID
 
 from pydantic import Field, SecretStr
 
@@ -253,25 +252,12 @@ class SecretFilter(WorkspaceScopedFilter):
         default=None,
         description="Name of the secret",
     )
-
     scope: Optional[Union[SecretScope, str]] = Field(
         default=None,
         description="Scope in which to filter secrets",
         union_mode="left_to_right",
     )
 
-    workspace_id: Optional[Union[UUID, str]] = Field(
-        default=None,
-        description="Workspace of the Secret",
-        union_mode="left_to_right",
-    )
-
-    user_id: Optional[Union[UUID, str]] = Field(
-        default=None,
-        description="User that created the Secret",
-        union_mode="left_to_right",
-    )
-
     @staticmethod
     def _get_filtering_value(value: Optional[Any]) -> str:
         """Convert the value to a string that can be used for lexicographical filtering and sorting.

zenml/models/v2/core/service.py

@@ -15,19 +15,20 @@
 
 from datetime import datetime
 from typing import (
+    TYPE_CHECKING,
     Any,
     ClassVar,
     Dict,
     List,
     Optional,
     Type,
+    TypeVar,
     Union,
 )
 from uuid import UUID
 
 from pydantic import BaseModel, ConfigDict, Field
 from sqlalchemy.sql.elements import ColumnElement
-from sqlmodel import SQLModel
 
 from zenml.constants import STR_FIELD_MAX_LENGTH
 from zenml.models.v2.base.scoped import (
@@ -37,11 +38,15 @@ from zenml.models.v2.base.scoped import (
     WorkspaceScopedResponseBody,
     WorkspaceScopedResponseMetadata,
     WorkspaceScopedResponseResources,
-    WorkspaceScopedTaggableFilter,
 )
 from zenml.services.service_status import ServiceState
 from zenml.services.service_type import ServiceType
 
+if TYPE_CHECKING:
+    from zenml.zen_stores.schemas import BaseSchema
+
+    AnySchema = TypeVar("AnySchema", bound=BaseSchema)
+
 # ------------------ Request Model ------------------
 
 
@@ -376,16 +381,6 @@ class ServiceFilter(WorkspaceScopedFilter):
         description="Name of the service. Use this to filter services by "
         "their name.",
     )
-    workspace_id: Optional[Union[UUID, str]] = Field(
-        default=None,
-        description="Workspace of the service",
-        union_mode="left_to_right",
-    )
-    user_id: Optional[Union[UUID, str]] = Field(
-        default=None,
-        description="User of the service",
-        union_mode="left_to_right",
-    )
     type: Optional[str] = Field(
         default=None,
         description="Type of the service. Filter services by their type.",
@@ -457,9 +452,7 @@ class ServiceFilter(WorkspaceScopedFilter):
         "config",
     ]
     CLI_EXCLUDE_FIELDS: ClassVar[List[str]] = [
-        *WorkspaceScopedTaggableFilter.CLI_EXCLUDE_FIELDS,
-        "workspace_id",
-        "user_id",
+        *WorkspaceScopedFilter.CLI_EXCLUDE_FIELDS,
         "flavor",
         "type",
         "pipeline_step_name",
@@ -468,7 +461,7 @@ class ServiceFilter(WorkspaceScopedFilter):
     ]
 
     def generate_filter(
-        self, table: Type["SQLModel"]
+        self, table: Type["AnySchema"]
     ) -> Union["ColumnElement[bool]"]:
         """Generate the filter for the query.
 

zenml/models/v2/core/service_connector.py

@@ -801,7 +801,6 @@ class ServiceConnectorFilter(WorkspaceScopedFilter):
         default=None,
         description="The type to scope this query to.",
     )
-
     name: Optional[str] = Field(
         default=None,
         description="The name to filter by",
@@ -810,16 +809,6 @@ class ServiceConnectorFilter(WorkspaceScopedFilter):
         default=None,
         description="The type of service connector to filter by",
     )
-    workspace_id: Optional[Union[UUID, str]] = Field(
-        default=None,
-        description="Workspace to filter by",
-        union_mode="left_to_right",
-    )
-    user_id: Optional[Union[UUID, str]] = Field(
-        default=None,
-        description="User to filter by",
-        union_mode="left_to_right",
-    )
     auth_method: Optional[str] = Field(
         default=None,
         title="Filter by the authentication method configured for the "

zenml/models/v2/core/stack.py

@@ -14,7 +14,17 @@
 """Models representing stacks."""
 
 import json
-from typing import TYPE_CHECKING, Any, ClassVar, Dict, List, Optional, Union
+from typing import (
+    TYPE_CHECKING,
+    Any,
+    ClassVar,
+    Dict,
+    List,
+    Optional,
+    Type,
+    TypeVar,
+    Union,
+)
 from uuid import UUID
 
 from pydantic import Field, model_validator
@@ -39,6 +49,9 @@ if TYPE_CHECKING:
     from sqlalchemy.sql.elements import ColumnElement
 
     from zenml.models.v2.core.component import ComponentResponse
+    from zenml.zen_stores.schemas import BaseSchema
+
+    AnySchema = TypeVar("AnySchema", bound=BaseSchema)
 
 
 # ------------------ Request Model ------------------
@@ -323,7 +336,6 @@ class StackFilter(WorkspaceScopedFilter):
     FILTER_EXCLUDE_FIELDS: ClassVar[List[str]] = [
         *WorkspaceScopedFilter.FILTER_EXCLUDE_FIELDS,
         "component_id",
-        "user",
         "component",
     ]
 
@@ -334,42 +346,32 @@ class StackFilter(WorkspaceScopedFilter):
     description: Optional[str] = Field(
         default=None, description="Description of the stack"
     )
-    workspace_id: Optional[Union[UUID, str]] = Field(
-        default=None,
-        description="Workspace of the stack",
-        union_mode="left_to_right",
-    )
-    user_id: Optional[Union[UUID, str]] = Field(
-        default=None,
-        description="User of the stack",
-        union_mode="left_to_right",
-    )
     component_id: Optional[Union[UUID, str]] = Field(
         default=None,
         description="Component in the stack",
         union_mode="left_to_right",
     )
-    user: Optional[Union[UUID, str]] = Field(
-        default=None,
-        description="Name/ID of the user that created the stack.",
-    )
     component: Optional[Union[UUID, str]] = Field(
         default=None, description="Name/ID of a component in the stack."
     )
 
-    def get_custom_filters(self) -> List["ColumnElement[bool]"]:
+    def get_custom_filters(
+        self, table: Type["AnySchema"]
+    ) -> List["ColumnElement[bool]"]:
         """Get custom filters.
 
+        Args:
+            table: The query table.
+
         Returns:
             A list of custom filters.
         """
-        custom_filters = super().get_custom_filters()
+        custom_filters = super().get_custom_filters(table)
 
         from zenml.zen_stores.schemas import (
             StackComponentSchema,
             StackCompositionSchema,
             StackSchema,
-            UserSchema,
         )
 
         if self.component_id:
@@ -379,17 +381,6 @@ class StackFilter(WorkspaceScopedFilter):
             )
             custom_filters.append(component_id_filter)
 
-        if self.user:
-            user_filter = and_(
-                StackSchema.user_id == UserSchema.id,
-                self.generate_name_or_id_query_conditions(
-                    value=self.user,
-                    table=UserSchema,
-                    additional_columns=["full_name"],
-                ),
-            )
-            custom_filters.append(user_filter)
-
         if self.component:
             component_filter = and_(
                 StackCompositionSchema.stack_id == StackSchema.id,

zenml/models/v2/core/step_run.py

@@ -14,7 +14,16 @@
 """Models representing steps runs."""
 
 from datetime import datetime
-from typing import TYPE_CHECKING, ClassVar, Dict, List, Optional, Union
+from typing import (
+    TYPE_CHECKING,
+    ClassVar,
+    Dict,
+    List,
+    Optional,
+    Type,
+    TypeVar,
+    Union,
+)
 from uuid import UUID
 
 from pydantic import BaseModel, ConfigDict, Field
@@ -41,6 +50,9 @@ if TYPE_CHECKING:
         LogsRequest,
         LogsResponse,
     )
+    from zenml.zen_stores.schemas import BaseSchema
+
+    AnySchema = TypeVar("AnySchema", bound=BaseSchema)
 
 
 class StepRunInputResponse(ArtifactVersionResponse):
@@ -553,16 +565,6 @@ class StepRunFilter(WorkspaceScopedFilter):
         description="Original id for this step run",
         union_mode="left_to_right",
     )
-    user_id: Optional[Union[UUID, str]] = Field(
-        default=None,
-        description="User that produced this step run",
-        union_mode="left_to_right",
-    )
-    workspace_id: Optional[Union[UUID, str]] = Field(
-        default=None,
-        description="Workspace of this step run",
-        union_mode="left_to_right",
-    )
     model_version_id: Optional[Union[UUID, str]] = Field(
         default=None,
         description="Model version associated with the step run.",
@@ -576,18 +578,20 @@ class StepRunFilter(WorkspaceScopedFilter):
         default=None,
         description="The run_metadata to filter the step runs by.",
     )
-
     model_config = ConfigDict(protected_namespaces=())
 
     def get_custom_filters(
-        self,
+        self, table: Type["AnySchema"]
     ) -> List["ColumnElement[bool]"]:
         """Get custom filters.
 
+        Args:
+            table: The query table.
+
         Returns:
             A list of custom filters.
         """
-        custom_filters = super().get_custom_filters()
+        custom_filters = super().get_custom_filters(table)
 
         from sqlmodel import and_
 

zenml/models/v2/core/trigger.py

@@ -13,7 +13,17 @@
 #  permissions and limitations under the License.
 """Collection of all models concerning triggers."""
 
-from typing import TYPE_CHECKING, Any, ClassVar, Dict, List, Optional, Union
+from typing import (
+    TYPE_CHECKING,
+    Any,
+    ClassVar,
+    Dict,
+    List,
+    Optional,
+    Type,
+    TypeVar,
+    Union,
+)
 from uuid import UUID
 
 from pydantic import Field, model_validator
@@ -39,6 +49,9 @@ if TYPE_CHECKING:
         ActionResponse,
     )
     from zenml.models.v2.core.event_source import EventSourceResponse
+    from zenml.zen_stores.schemas import BaseSchema
+
+    AnySchema = TypeVar("AnySchema", bound=BaseSchema)
 
 
 # ------------------ Request Model ------------------
@@ -358,10 +371,13 @@ class TriggerFilter(WorkspaceScopedFilter):
     )
 
     def get_custom_filters(
-        self,
+        self, table: Type["AnySchema"]
     ) -> List["ColumnElement[bool]"]:
         """Get custom filters.
 
+        Args:
+            table: The query table.
+
         Returns:
             A list of custom filters.
         """
@@ -373,7 +389,7 @@ class TriggerFilter(WorkspaceScopedFilter):
             TriggerSchema,
         )
 
-        custom_filters = super().get_custom_filters()
+        custom_filters = super().get_custom_filters(table)
 
         if self.event_source_flavor:
             event_source_flavor_filter = and_(

zenml/orchestrators/step_launcher.py

@@ -179,12 +179,10 @@ class StepLauncher:
                         pipeline_run_id=pipeline_run.id,
                         pipeline_run_metadata=pipeline_run_metadata,
                     )
-
-                pipeline_model_version, pipeline_run = (
-                    step_run_utils.prepare_pipeline_run_model_version(
-                        pipeline_run
-                    )
-                )
+                    if model_version := pipeline_run.model_version:
+                        step_run_utils.log_model_version_dashboard_url(
+                            model_version=model_version
+                        )
 
                 request_factory = step_run_utils.StepRunRequestFactory(
                     deployment=self._deployment,
@@ -209,12 +207,10 @@ class StepLauncher:
                     step_run = Client().zen_store.create_run_step(
                         step_run_request
                     )
-
-                    step_model_version, step_run = (
-                        step_run_utils.prepare_step_run_model_version(
-                            step_run=step_run, pipeline_run=pipeline_run
+                    if model_version := step_run.model_version:
+                        step_run_utils.log_model_version_dashboard_url(
+                            model_version=model_version
                         )
-                    )
 
                 if not step_run.status.is_finished:
                     logger.info(f"Step `{self._step_name}` has started.")
@@ -289,8 +285,8 @@ class StepLauncher:
                         f"Using cached version of step `{self._step_name}`."
                     )
                     if (
-                        model_version := step_model_version
-                        or pipeline_model_version
+                        model_version := step_run.model_version
+                        or pipeline_run.model_version
                     ):
                         step_run_utils.link_output_artifacts_to_model_version(
                             artifacts=step_run.outputs,

zenml/orchestrators/step_run_utils.py

@@ -14,7 +14,7 @@
 """Utilities for creating step runs."""
 
 from datetime import datetime
-from typing import TYPE_CHECKING, Dict, List, Optional, Set, Tuple
+from typing import Dict, List, Optional, Set, Tuple
 
 from zenml.client import Client
 from zenml.config.step_configurations import Step
@@ -24,21 +24,13 @@ from zenml.logger import get_logger
 from zenml.model.utils import link_artifact_version_to_model_version
 from zenml.models import (
     ArtifactVersionResponse,
-    ModelVersionPipelineRunRequest,
     ModelVersionResponse,
     PipelineDeploymentResponse,
     PipelineRunResponse,
-    PipelineRunUpdate,
     StepRunRequest,
-    StepRunResponse,
-    StepRunUpdate,
 )
 from zenml.orchestrators import cache_utils, input_utils, utils
 from zenml.stack import Stack
-from zenml.utils import pagination_utils, string_utils
-
-if TYPE_CHECKING:
-    from zenml.model.model import Model
 
 logger = get_logger(__name__)
 
@@ -293,10 +285,6 @@ def create_cached_step_runs(
         deployment=deployment, pipeline_run=pipeline_run, stack=stack
     )
 
-    pipeline_model_version, pipeline_run = prepare_pipeline_run_model_version(
-        pipeline_run=pipeline_run
-    )
-
     while (
         cache_candidates := find_cacheable_invocation_candidates(
             deployment=deployment,
@@ -311,7 +299,9 @@ def create_cached_step_runs(
 
             # Make sure the request factory has the most up to date pipeline
             # run to avoid hydration calls
-            request_factory.pipeline_run = pipeline_run
+            request_factory.pipeline_run = Client().get_pipeline_run(
+                pipeline_run.id
+            )
             try:
                 step_run_request = request_factory.create_request(
                     invocation_id
@@ -336,15 +326,10 @@ def create_cached_step_runs(
 
             step_run = Client().zen_store.create_run_step(step_run_request)
 
-            # Refresh the pipeline run here to make sure we have the latest
-            # state
-            pipeline_run = Client().get_pipeline_run(pipeline_run.id)
-
-            step_model_version, step_run = prepare_step_run_model_version(
-                step_run=step_run, pipeline_run=pipeline_run
-            )
-
-            if model_version := step_model_version or pipeline_model_version:
+            if (
+                model_version := step_run.model_version
+                or pipeline_run.model_version
+            ):
                 link_output_artifacts_to_model_version(
                     artifacts=step_run.outputs,
                     model_version=model_version,
@@ -356,169 +341,6 @@ def create_cached_step_runs(
     return cached_invocations
 
 
-def get_or_create_model_version_for_pipeline_run(
-    model: "Model",
-    pipeline_run: PipelineRunResponse,
-    substitutions: Dict[str, str],
-) -> Tuple[ModelVersionResponse, bool]:
-    """Get or create a model version as part of a pipeline run.
-
-    Args:
-        model: The model to get or create.
-        pipeline_run: The pipeline run for which the model should be created.
-        substitutions: Substitutions to apply to the model version name.
-
-    Returns:
-        The model version and a boolean indicating whether it was newly created
-        or not.
-    """
-    # Copy the model before modifying it so we don't accidently modify
-    # configurations in which the model object is potentially referenced
-    model = model.model_copy()
-
-    if model.model_version_id:
-        return model._get_model_version(), False
-    elif model.version:
-        if isinstance(model.version, str):
-            model.version = string_utils.format_name_template(
-                model.version,
-                substitutions=substitutions,
-            )
-        model.name = string_utils.format_name_template(
-            model.name,
-            substitutions=substitutions,
-        )
-
-        return (
-            model._get_or_create_model_version(),
-            model._created_model_version,
-        )
-
-    # The model version should be created as part of this run
-    # -> We first check if it was already created as part of this run, and if
-    # not we do create it. If this is running in two parallel steps, we might
-    # run into issues that this will create two versions. Ideally, all model
-    # versions required for a pipeline run and its steps could be created
-    # server-side at run creation time before the first step starts.
-    if model_version := get_model_version_created_by_pipeline_run(
-        model_name=model.name, pipeline_run=pipeline_run
-    ):
-        return model_version, False
-    else:
-        return model._get_or_create_model_version(), True
-
-
-def get_model_version_created_by_pipeline_run(
-    model_name: str, pipeline_run: PipelineRunResponse
-) -> Optional[ModelVersionResponse]:
-    """Get a model version that was created by a specific pipeline run.
-
-    This function does not refresh the pipeline run, so it will only try to
-    fetch the model version from existing steps if they're already part of the
-    response.
-
-    Args:
-        model_name: The model name for which to get the version.
-        pipeline_run: The pipeline run for which to get the version.
-
-    Returns:
-        A model version with the given name created by the run, or None if such
-        a model version does not exist.
-    """
-    if pipeline_run.config.model and pipeline_run.model_version:
-        if (
-            pipeline_run.config.model.name == model_name
-            and pipeline_run.config.model.version is None
-        ):
-            return pipeline_run.model_version
-
-    # We fetch a list of hydrated step runs here in order to avoid hydration
-    # calls for each step separately.
-    candidate_step_runs = pagination_utils.depaginate(
-        Client().list_run_steps,
-        pipeline_run_id=pipeline_run.id,
-        model=model_name,
-        hydrate=True,
-    )
-    for step_run in candidate_step_runs:
-        if step_run.config.model and step_run.model_version:
-            if (
-                step_run.config.model.name == model_name
-                and step_run.config.model.version is None
-            ):
-                return step_run.model_version
-
-    return None
-
-
-def prepare_pipeline_run_model_version(
-    pipeline_run: PipelineRunResponse,
-) -> Tuple[Optional[ModelVersionResponse], PipelineRunResponse]:
-    """Prepare the model version for a pipeline run.
-
-    Args:
-        pipeline_run: The pipeline run for which to prepare the model version.
-
-    Returns:
-        The prepared model version and the updated pipeline run.
-    """
-    model_version = None
-
-    if pipeline_run.model_version:
-        model_version = pipeline_run.model_version
-    elif config_model := pipeline_run.config.model:
-        model_version, _ = get_or_create_model_version_for_pipeline_run(
-            model=config_model,
-            pipeline_run=pipeline_run,
-            substitutions=pipeline_run.config.substitutions,
-        )
-        pipeline_run = Client().zen_store.update_run(
-            run_id=pipeline_run.id,
-            run_update=PipelineRunUpdate(model_version_id=model_version.id),
-        )
-        link_pipeline_run_to_model_version(
-            pipeline_run=pipeline_run, model_version=model_version
-        )
-        log_model_version_dashboard_url(model_version)
-
-    return model_version, pipeline_run
-
-
-def prepare_step_run_model_version(
-    step_run: StepRunResponse, pipeline_run: PipelineRunResponse
-) -> Tuple[Optional[ModelVersionResponse], StepRunResponse]:
-    """Prepare the model version for a step run.
-
-    Args:
-        step_run: The step run for which to prepare the model version.
-        pipeline_run: The pipeline run of the step.
-
-    Returns:
-        The prepared model version and the updated step run.
-    """
-    model_version = None
-
-    if step_run.model_version:
-        model_version = step_run.model_version
-    elif config_model := step_run.config.model:
-        model_version, created = get_or_create_model_version_for_pipeline_run(
-            model=config_model,
-            pipeline_run=pipeline_run,
-            substitutions=step_run.config.substitutions,
-        )
-        step_run = Client().zen_store.update_run_step(
-            step_run_id=step_run.id,
-            step_run_update=StepRunUpdate(model_version_id=model_version.id),
-        )
-        link_pipeline_run_to_model_version(
-            pipeline_run=pipeline_run, model_version=model_version
-        )
-        if created:
-            log_model_version_dashboard_url(model_version)
-
-    return model_version, step_run
-
-
 def log_model_version_dashboard_url(
     model_version: ModelVersionResponse,
 ) -> None:
@@ -546,24 +368,6 @@ def log_model_version_dashboard_url(
         )
 
 
-def link_pipeline_run_to_model_version(
-    pipeline_run: PipelineRunResponse, model_version: ModelVersionResponse
-) -> None:
-    """Link a pipeline run to a model version.
-
-    Args:
-        pipeline_run: The pipeline run to link.
-        model_version: The model version to link.
-    """
-    client = Client()
-    client.zen_store.create_model_version_pipeline_run_link(
-        ModelVersionPipelineRunRequest(
-            pipeline_run=pipeline_run.id,
-            model_version=model_version.id,
-        )
-    )
-
-
 def link_output_artifacts_to_model_version(
     artifacts: Dict[str, List[ArtifactVersionResponse]],
     model_version: ModelVersionResponse,