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

zenml/zen_stores/sql_zen_store.py

@@ -55,7 +55,7 @@ from pydantic import (
     field_validator,
     model_validator,
 )
-from sqlalchemy import asc, case, desc, func
+from sqlalchemy import func
 from sqlalchemy.engine import URL, Engine, make_url
 from sqlalchemy.exc import (
     ArgumentError,
@@ -71,6 +71,7 @@ from sqlmodel import (
     col,
     create_engine,
     delete,
+    desc,
     or_,
     select,
 )
@@ -100,7 +101,6 @@ from zenml.constants import (
     ENV_ZENML_SERVER,
     FINISHED_ONBOARDING_SURVEY_KEY,
     MAX_RETRIES_FOR_VERSIONED_ENTITY_CREATION,
-    SORT_PIPELINES_BY_LATEST_RUN_KEY,
     SQL_STORE_BACKUP_DIRECTORY_NAME,
     TEXT_FIELD_MAX_LENGTH,
     handle_bool_env_var,
@@ -117,7 +117,6 @@ from zenml.enums import (
     OnboardingStep,
     SecretScope,
     SecretsStoreType,
-    SorterOps,
     StackComponentType,
     StackDeploymentProvider,
     StepRunInputArtifactType,
@@ -298,7 +297,11 @@ from zenml.utils.networking_utils import (
     replace_localhost_with_internal_hostname,
 )
 from zenml.utils.pydantic_utils import before_validator_handler
-from zenml.utils.string_utils import random_str, validate_name
+from zenml.utils.string_utils import (
+    format_name_template,
+    random_str,
+    validate_name,
+)
 from zenml.zen_stores import template_utils
 from zenml.zen_stores.base_zen_store import (
     BaseZenStore,
@@ -4358,69 +4361,14 @@ class SqlZenStore(BaseZenStore):
         Returns:
             A list of all pipelines matching the filter criteria.
         """
-        query: Union[Select[Any], SelectOfScalar[Any]] = select(PipelineSchema)
-        _custom_conversion: Optional[Callable[[Any], PipelineResponse]] = None
-
-        column, operand = pipeline_filter_model.sorting_params
-        if column == SORT_PIPELINES_BY_LATEST_RUN_KEY:
-            with Session(self.engine) as session:
-                max_date_subquery = (
-                    # If no run exists for the pipeline yet, we use the pipeline
-                    # creation date as a fallback, otherwise newly created
-                    # pipeline would always be at the top/bottom
-                    select(
-                        PipelineSchema.id,
-                        case(
-                            (
-                                func.max(PipelineRunSchema.created).is_(None),
-                                PipelineSchema.created,
-                            ),
-                            else_=func.max(PipelineRunSchema.created),
-                        ).label("run_or_created"),
-                    )
-                    .outerjoin(
-                        PipelineRunSchema,
-                        PipelineSchema.id == PipelineRunSchema.pipeline_id,  # type: ignore[arg-type]
-                    )
-                    .group_by(col(PipelineSchema.id))
-                    .subquery()
-                )
-
-                if operand == SorterOps.DESCENDING:
-                    sort_clause = desc
-                else:
-                    sort_clause = asc
-
-                query = (
-                    # We need to include the subquery in the select here to
-                    # make this query work with the distinct statement. This
-                    # result will be removed in the custom conversion function
-                    # applied later
-                    select(PipelineSchema, max_date_subquery.c.run_or_created)
-                    .where(PipelineSchema.id == max_date_subquery.c.id)
-                    .order_by(sort_clause(max_date_subquery.c.run_or_created))
-                    # We always add the `id` column as a tiebreaker to ensure a
-                    # stable, repeatable order of items, otherwise subsequent
-                    # pages might contain the same items.
-                    .order_by(col(PipelineSchema.id))
-                )
-
-            def _custom_conversion(row: Any) -> PipelineResponse:
-                return cast(
-                    PipelineResponse,
-                    row[0].to_model(
-                        include_metadata=hydrate, include_resources=True
-                    ),
-                )
-
         with Session(self.engine) as session:
+            query = select(PipelineSchema)
             return self.filter_and_paginate(
                 session=session,
                 query=query,
                 table=PipelineSchema,
                 filter_model=pipeline_filter_model,
                 hydrate=hydrate,
-                custom_schema_to_model_conversion=_custom_conversion,
             )
 
     def count_pipelines(self, filter_model: Optional[PipelineFilter]) -> int:
@@ -5211,6 +5159,20 @@ class SqlZenStore(BaseZenStore):
                         "already exists."
                     )
 
+            if model_version_id := self._get_or_create_model_version_for_run(
+                new_run
+            ):
+                new_run.model_version_id = model_version_id
+                session.add(new_run)
+                session.commit()
+
+                self.create_model_version_pipeline_run_link(
+                    ModelVersionPipelineRunRequest(
+                        model_version=model_version_id, pipeline_run=new_run.id
+                    )
+                )
+                session.refresh(new_run)
+
             return new_run.to_model(
                 include_metadata=True, include_resources=True
             )
@@ -8167,25 +8129,17 @@ class SqlZenStore(BaseZenStore):
                     f"with ID '{step_run.pipeline_run_id}' found."
                 )
 
-            # Check if the step name already exists in the pipeline run
-            existing_step_run = session.exec(
-                select(StepRunSchema)
-                .where(StepRunSchema.name == step_run.name)
-                .where(
-                    StepRunSchema.pipeline_run_id == step_run.pipeline_run_id
-                )
-            ).first()
-            if existing_step_run is not None:
+            step_schema = StepRunSchema.from_request(step_run)
+            session.add(step_schema)
+            try:
+                session.commit()
+            except IntegrityError:
                 raise EntityExistsError(
                     f"Unable to create step `{step_run.name}`: A step with "
                     f"this name already exists in the pipeline run with ID "
                     f"'{step_run.pipeline_run_id}'."
                 )
 
-            # Create the step
-            step_schema = StepRunSchema.from_request(step_run)
-            session.add(step_schema)
-
             # Add logs entry for the step if exists
             if step_run.logs is not None:
                 log_entry = LogsSchema(
@@ -8281,6 +8235,21 @@ class SqlZenStore(BaseZenStore):
             session.commit()
             session.refresh(step_schema)
 
+            if model_version_id := self._get_or_create_model_version_for_run(
+                step_schema
+            ):
+                step_schema.model_version_id = model_version_id
+                session.add(step_schema)
+                session.commit()
+
+                self.create_model_version_pipeline_run_link(
+                    ModelVersionPipelineRunRequest(
+                        model_version=model_version_id,
+                        pipeline_run=step_schema.pipeline_run_id,
+                    )
+                )
+                session.refresh(step_schema)
+
             return step_schema.to_model(
                 include_metadata=True, include_resources=True
             )
@@ -10283,6 +10252,22 @@ class SqlZenStore(BaseZenStore):
 
     # ----------------------------- Model Versions -----------------------------
 
+    def _get_or_create_model(
+        self, model_request: ModelRequest
+    ) -> Tuple[bool, ModelResponse]:
+        """Get or create a model.
+
+        Args:
+            model_request: The model request.
+
+        Returns:
+            A boolean whether the model was created or not, and the model.
+        """
+        try:
+            return True, self.create_model(model_request)
+        except EntityExistsError:
+            return False, self.get_model(model_request.name)
+
     def _get_next_numeric_version_for_model(
         self, session: Session, model_id: UUID
     ) -> int:
@@ -10307,55 +10292,276 @@ class SqlZenStore(BaseZenStore):
         else:
             return int(current_max_version) + 1
 
-    def _model_version_exists(self, model_id: UUID, version: str) -> bool:
+    def _model_version_exists(
+        self,
+        model_id: UUID,
+        version: Optional[str] = None,
+        producer_run_id: Optional[UUID] = None,
+    ) -> bool:
         """Check if a model version with a certain version exists.
 
         Args:
             model_id: The model ID of the version.
             version: The version name.
+            producer_run_id: The producer run ID. If given, checks if a numeric
+                version for the producer run exists.
 
         Returns:
-            If a model version with the given version name exists.
+            If a model version for the given arguments exists.
         """
+        query = select(ModelVersionSchema.id).where(
+            ModelVersionSchema.model_id == model_id
+        )
+
+        if version:
+            query = query.where(ModelVersionSchema.name == version)
+
+        if producer_run_id:
+            query = query.where(
+                ModelVersionSchema.producer_run_id_if_numeric
+                == producer_run_id,
+            )
+
         with Session(self.engine) as session:
-            return (
-                session.exec(
-                    select(ModelVersionSchema.id)
-                    .where(ModelVersionSchema.model_id == model_id)
-                    .where(ModelVersionSchema.name == version)
-                ).first()
-                is not None
+            return session.exec(query).first() is not None
+
+    def _get_model_version(
+        self,
+        model_id: UUID,
+        version_name: Optional[str] = None,
+        producer_run_id: Optional[UUID] = None,
+    ) -> ModelVersionResponse:
+        """Get a model version.
+
+        Args:
+            model_id: The ID of the model.
+            version_name: The name of the model version.
+            producer_run_id: The ID of the producer pipeline run. If this is
+                set, only numeric versions created as part of the pipeline run
+                will be returned.
+
+        Raises:
+            ValueError: If no version name or producer run ID was provided.
+            KeyError: If no model version was found.
+
+        Returns:
+            The model version.
+        """
+        query = select(ModelVersionSchema).where(
+            ModelVersionSchema.model_id == model_id
+        )
+
+        if version_name:
+            if version_name.isnumeric():
+                query = query.where(
+                    ModelVersionSchema.number == int(version_name)
+                )
+                error_text = (
+                    f"No version with number {version_name} found "
+                    f"for model {model_id}."
+                )
+            elif version_name in ModelStages.values():
+                if version_name == ModelStages.LATEST:
+                    query = query.order_by(
+                        desc(col(ModelVersionSchema.number))
+                    ).limit(1)
+                else:
+                    query = query.where(
+                        ModelVersionSchema.stage == version_name
+                    )
+                error_text = (
+                    f"No {version_name} stage version found for "
+                    f"model {model_id}."
+                )
+            else:
+                query = query.where(ModelVersionSchema.name == version_name)
+                error_text = (
+                    f"No {version_name} version found for model {model_id}."
+                )
+
+        elif producer_run_id:
+            query = query.where(
+                ModelVersionSchema.producer_run_id_if_numeric
+                == producer_run_id,
+            )
+            error_text = (
+                f"No numeric model version found for model {model_id} "
+                f"and producer run {producer_run_id}."
+            )
+        else:
+            raise ValueError(
+                "Version name or producer run id need to be specified."
             )
 
-    @track_decorator(AnalyticsEvent.CREATED_MODEL_VERSION)
-    def create_model_version(
-        self, model_version: ModelVersionRequest
+        with Session(self.engine) as session:
+            schema = session.exec(query).one_or_none()
+
+            if not schema:
+                raise KeyError(error_text)
+
+            return schema.to_model(
+                include_metadata=True, include_resources=True
+            )
+
+    def _get_or_create_model_version(
+        self,
+        model_version_request: ModelVersionRequest,
+        producer_run_id: Optional[UUID] = None,
+    ) -> Tuple[bool, ModelVersionResponse]:
+        """Get or create a model version.
+
+        Args:
+            model_version_request: The model version request.
+            producer_run_id: ID of the producer pipeline run.
+
+        Raises:
+            EntityCreationError: If the model version creation failed.
+
+        Returns:
+            A boolean whether the model version was created or not, and the
+            model version.
+        """
+        try:
+            model_version = self._create_model_version(
+                model_version=model_version_request,
+                producer_run_id=producer_run_id,
+            )
+            track(event=AnalyticsEvent.CREATED_MODEL_VERSION)
+            return True, model_version
+        except EntityCreationError:
+            # Need to explicitly re-raise this here as otherwise the catching
+            # of the RuntimeError would include this
+            raise
+        except RuntimeError:
+            return False, self._get_model_version(
+                model_id=model_version_request.model,
+                producer_run_id=producer_run_id,
+            )
+        except EntityExistsError:
+            return False, self._get_model_version(
+                model_id=model_version_request.model,
+                version_name=model_version_request.name,
+            )
+
+    def _get_or_create_model_version_for_run(
+        self, pipeline_or_step_run: Union[PipelineRunSchema, StepRunSchema]
+    ) -> Optional[UUID]:
+        """Get or create a model version for a pipeline or step run.
+
+        Args:
+            pipeline_or_step_run: The pipeline or step run for which to create
+                the model version.
+
+        Returns:
+            The model version.
+        """
+        if isinstance(pipeline_or_step_run, PipelineRunSchema):
+            producer_run_id = pipeline_or_step_run.id
+            pipeline_run = pipeline_or_step_run.to_model(include_metadata=True)
+            configured_model = pipeline_run.config.model
+            substitutions = pipeline_run.config.substitutions
+        else:
+            producer_run_id = pipeline_or_step_run.pipeline_run_id
+            step_run = pipeline_or_step_run.to_model(include_metadata=True)
+            configured_model = step_run.config.model
+            substitutions = step_run.config.substitutions
+
+        if not configured_model:
+            return None
+
+        model_request = ModelRequest(
+            name=format_name_template(
+                configured_model.name, substitutions=substitutions
+            ),
+            license=configured_model.license,
+            description=configured_model.description,
+            audience=configured_model.audience,
+            use_cases=configured_model.use_cases,
+            limitations=configured_model.limitations,
+            trade_offs=configured_model.trade_offs,
+            ethics=configured_model.ethics,
+            save_models_to_registry=configured_model.save_models_to_registry,
+            user=pipeline_or_step_run.user_id,
+            workspace=pipeline_or_step_run.workspace_id,
+        )
+
+        _, model_response = self._get_or_create_model(
+            model_request=model_request
+        )
+
+        version_name = None
+        if configured_model.version is not None:
+            version_name = format_name_template(
+                str(configured_model.version), substitutions=substitutions
+            )
+
+            # If the model version was specified to be a numeric version or
+            # stage we don't try to create it (which will fail because it is not
+            # allowed) but try to fetch it immediately
+            if (
+                version_name.isnumeric()
+                or version_name in ModelStages.values()
+            ):
+                return self._get_model_version(
+                    model_id=model_response.id, version_name=version_name
+                ).id
+
+        model_version_request = ModelVersionRequest(
+            model=model_response.id,
+            name=version_name,
+            description=configured_model.description,
+            tags=configured_model.tags,
+            user=pipeline_or_step_run.user_id,
+            workspace=pipeline_or_step_run.workspace_id,
+        )
+
+        _, model_version_response = self._get_or_create_model_version(
+            model_version_request=model_version_request,
+            producer_run_id=producer_run_id,
+        )
+        return model_version_response.id
+
+    def _create_model_version(
+        self,
+        model_version: ModelVersionRequest,
+        producer_run_id: Optional[UUID] = None,
     ) -> ModelVersionResponse:
         """Creates a new model version.
 
         Args:
             model_version: the Model Version to be created.
+            producer_run_id: ID of the pipeline run that produced this model
+                version.
 
         Returns:
             The newly created model version.
 
         Raises:
-            ValueError: If `number` is not None during model version creation.
+            ValueError: If the requested version name is invalid.
             EntityExistsError: If a model version with the given name already
                 exists.
             EntityCreationError: If the model version creation failed.
+            RuntimeError: If an auto-incremented model version already exists
+                for the producer run.
         """
-        if model_version.number is not None:
-            raise ValueError(
-                "`number` field  must be None during model version creation."
-            )
+        has_custom_name = False
+        if model_version.name:
+            has_custom_name = True
+            validate_name(model_version)
 
-        model = self.get_model(model_version.model)
+            if model_version.name.isnumeric():
+                raise ValueError(
+                    "Can't create model version with custom numeric model "
+                    "version name."
+                )
 
-        has_custom_name = model_version.name is not None
-        if has_custom_name:
-            validate_name(model_version)
+            if str(model_version.name).lower() in ModelStages.values():
+                raise ValueError(
+                    "Can't create model version with a name that is used as a "
+                    f"model version stage ({ModelStages.values()})."
+                )
 
+        model = self.get_model(model_version.model)
         model_version_id = None
 
         remaining_tries = MAX_RETRIES_FOR_VERSIONED_ENTITY_CREATION
@@ -10363,17 +10569,19 @@ class SqlZenStore(BaseZenStore):
             remaining_tries -= 1
             try:
                 with Session(self.engine) as session:
-                    model_version.number = (
+                    model_version_number = (
                         self._get_next_numeric_version_for_model(
                             session=session,
                             model_id=model.id,
                         )
                     )
                     if not has_custom_name:
-                        model_version.name = str(model_version.number)
+                        model_version.name = str(model_version_number)
 
                     model_version_schema = ModelVersionSchema.from_request(
-                        model_version
+                        model_version,
+                        model_version_number=model_version_number,
+                        producer_run_id=producer_run_id,
                     )
                     session.add(model_version_schema)
                     session.commit()
@@ -10394,6 +10602,13 @@ class SqlZenStore(BaseZenStore):
                         f"{model_version.name}): A model with the "
                         "same name and version already exists."
                     )
+                elif producer_run_id and self._model_version_exists(
+                    model_id=model.id, producer_run_id=producer_run_id
+                ):
+                    raise RuntimeError(
+                        "Auto-incremented model version already exists for "
+                        f"producer run {producer_run_id}."
+                    )
                 elif remaining_tries == 0:
                     raise EntityCreationError(
                         f"Failed to create version for model "
@@ -10412,10 +10627,9 @@ class SqlZenStore(BaseZenStore):
                     )
                     logger.debug(
                         "Failed to create model version %s "
-                        "(version %s) due to an integrity error. "
+                        "due to an integrity error. "
                         "Retrying in %f seconds.",
                         model.name,
-                        model_version.number,
                         sleep_duration,
                     )
                     time.sleep(sleep_duration)
@@ -10430,6 +10644,20 @@ class SqlZenStore(BaseZenStore):
 
         return self.get_model_version(model_version_id)
 
+    @track_decorator(AnalyticsEvent.CREATED_MODEL_VERSION)
+    def create_model_version(
+        self, model_version: ModelVersionRequest
+    ) -> ModelVersionResponse:
+        """Creates a new model version.
+
+        Args:
+            model_version: the Model Version to be created.
+
+        Returns:
+            The newly created model version.
+        """
+        return self._create_model_version(model_version=model_version)
+
     def get_model_version(
         self, model_version_id: UUID, hydrate: bool = True
     ) -> ModelVersionResponse:

{zenml_nightly-0.71.0.dev20241212.dist-info → zenml_nightly-0.71.0.dev20241214.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: zenml-nightly
-Version: 0.71.0.dev20241212
+Version: 0.71.0.dev20241214
 Summary: ZenML: Write production-ready ML code.
 Home-page: https://zenml.io
 License: Apache-2.0