zenml-nightly 0.83.0.dev20250618__py3-none-any.whl → 0.83.0.dev20250621__py3-none-any.whl

zenml/zen_stores/migrations/versions/857843db1bcf_add_api_transaction_table.py

@@ -0,0 +1,69 @@
+"""add API transaction table [857843db1bcf].
+
+Revision ID: 857843db1bcf
+Revises: 0.83.0
+Create Date: 2025-06-13 19:47:14.165158
+
+"""
+
+import sqlalchemy as sa
+import sqlmodel
+from alembic import op
+from sqlalchemy.dialects.mysql import MEDIUMTEXT
+
+# revision identifiers, used by Alembic.
+revision = "857843db1bcf"
+down_revision = "0.83.0"
+branch_labels = None
+depends_on = None
+
+
+def upgrade() -> None:
+    """Upgrade database schema and/or data, creating a new revision."""
+    # ### commands auto generated by Alembic - please adjust! ###
+    op.create_table(
+        "api_transaction",
+        sa.Column("id", sqlmodel.sql.sqltypes.GUID(), nullable=False),
+        sa.Column("created", sa.DateTime(), nullable=False),
+        sa.Column("updated", sa.DateTime(), nullable=False),
+        sa.Column(
+            "method", sqlmodel.sql.sqltypes.AutoString(), nullable=False
+        ),
+        sa.Column("url", sa.TEXT(), nullable=False),
+        sa.Column("completed", sa.Boolean(), nullable=False),
+        sa.Column(
+            "result",
+            MEDIUMTEXT
+            if op.get_bind().dialect.name == "mysql"
+            else sa.String(length=16777215),
+            nullable=True,
+        ),
+        sa.Column("user_id", sqlmodel.sql.sqltypes.GUID(), nullable=False),
+        sa.Column("expired", sa.DateTime(), nullable=True),
+        sa.ForeignKeyConstraint(
+            ["user_id"],
+            ["user.id"],
+            name="fk_api_transaction_user_id_user",
+            ondelete="CASCADE",
+        ),
+        sa.PrimaryKeyConstraint("id"),
+    )
+    op.create_index(
+        "ix_api_transaction_completed_expired",
+        "api_transaction",
+        ["completed", "expired"],
+        unique=False,
+    )
+
+    # ### end Alembic commands ###
+
+
+def downgrade() -> None:
+    """Downgrade database schema and/or data back to the previous revision."""
+    # ### commands auto generated by Alembic - please adjust! ###
+    op.drop_index(
+        "ix_api_transaction_completed_expired",
+        table_name="api_transaction",
+    )
+    op.drop_table("api_transaction")
+    # ### end Alembic commands ###

zenml/zen_stores/rest_zen_store.py

@@ -1631,7 +1631,10 @@ class RestZenStore(BaseZenStore):
         )
 
     def get_deployment(
-        self, deployment_id: UUID, hydrate: bool = True
+        self,
+        deployment_id: UUID,
+        hydrate: bool = True,
+        step_configuration_filter: Optional[List[str]] = None,
     ) -> PipelineDeploymentResponse:
         """Get a deployment with a given ID.
 
@@ -1639,6 +1642,9 @@ class RestZenStore(BaseZenStore):
             deployment_id: ID of the deployment.
             hydrate: Flag deciding whether to hydrate the output model(s)
                 by including metadata fields in the response.
+            step_configuration_filter: List of step configurations to include in
+                the response. If not given, all step configurations will be
+                included.
 
         Returns:
             The deployment.
@@ -1647,7 +1653,10 @@ class RestZenStore(BaseZenStore):
             resource_id=deployment_id,
             route=PIPELINE_DEPLOYMENTS,
             response_model=PipelineDeploymentResponse,
-            params={"hydrate": hydrate},
+            params={
+                "hydrate": hydrate,
+                "step_configuration_filter": step_configuration_filter,
+            },
         )
 
     def list_deployments(
@@ -4194,20 +4203,6 @@ class RestZenStore(BaseZenStore):
         Returns:
             A requests session.
         """
-
-        class AugmentedRetry(Retry):
-            """Augmented retry class that also retries on 429 status codes for POST requests."""
-
-            def is_retry(
-                self,
-                method: str,
-                status_code: int,
-                has_retry_after: bool = False,
-            ) -> bool:
-                if status_code == 429:
-                    return True
-                return super().is_retry(method, status_code, has_retry_after)
-
         if self._session is None:
             # We only need to initialize the session once over the lifetime
             # of the client. We can swap the token out when it expires.
@@ -4217,12 +4212,14 @@ class RestZenStore(BaseZenStore):
                 )
 
             self._session = requests.Session()
-            # Retries are triggered for idempotent HTTP methods (GET, HEAD, PUT,
-            # OPTIONS and DELETE) on specific HTTP status codes:
+            # Retries are triggered for all HTTP methods (GET, HEAD, POST, PUT,
+            # PATCH, OPTIONS and DELETE) on specific HTTP status codes:
             #
+            #     408: Request Timeout.
+            #     429: Too Many Requests.
             #     502: Bad Gateway.
             #     503: Service Unavailable.
-            #     504: Gateway Timeout.
+            #     504: Gateway Timeout
             #
             # This also handles connection level errors, if a connection attempt
             # fails due to transient issues like:
@@ -4237,12 +4234,20 @@ class RestZenStore(BaseZenStore):
             #     the timeout period.
             #     Connection Refused: If the server refuses the connection.
             #
-            retries = AugmentedRetry(
+            retries = Retry(
                 connect=5,
                 read=8,
                 redirect=3,
                 status=10,
-                allowed_methods=["HEAD", "GET", "PUT", "DELETE", "OPTIONS"],
+                allowed_methods=[
+                    "HEAD",
+                    "GET",
+                    "POST",
+                    "PUT",
+                    "PATCH",
+                    "DELETE",
+                    "OPTIONS",
+                ],
                 status_forcelist=[
                     408,  # Request Timeout
                     429,  # Too Many Requests
@@ -4389,14 +4394,6 @@ class RestZenStore(BaseZenStore):
         self.session.headers.update(
             {source_context.name: source_context.get().value}
         )
-        # Add a request ID to the request headers
-        request_id = str(uuid4())[:8]
-        self.session.headers.update({"X-Request-ID": request_id})
-        path = url.removeprefix(self.url)
-        start_time = time.time()
-        logger.debug(
-            f"Sending {method} request to {path} with request ID {request_id}..."
-        )
 
         # If the server replies with a credentials validation (401 Unauthorized)
         # error, we (re-)authenticate and retry the request here in the
@@ -4413,18 +4410,31 @@ class RestZenStore(BaseZenStore):
         # two times: once after initial authentication and once after
         # re-authentication.
         re_authenticated = False
+        path = url.removeprefix(self.url)
         while True:
+            # Add a request ID to the request headers
+            request_id = str(uuid4())[:8]
+            self.session.headers.update({"X-Request-ID": request_id})
+            # Add an idempotency key to the request headers to ensure that
+            # requests are idempotent.
+            self.session.headers.update({"Idempotency-Key": str(uuid4())})
+
+            start_time = time.time()
+            logger.debug(f"[{request_id}] {method} {path} started...")
+            status_code = "failed"
+
             try:
-                return self._handle_response(
-                    self.session.request(
-                        method,
-                        url,
-                        params=params if params else {},
-                        verify=self.config.verify_ssl,
-                        timeout=timeout or self.config.http_timeout,
-                        **kwargs,
-                    )
+                response = self.session.request(
+                    method,
+                    url,
+                    params=params if params else {},
+                    verify=self.config.verify_ssl,
+                    timeout=timeout or self.config.http_timeout,
+                    **kwargs,
                 )
+
+                status_code = str(response.status_code)
+                return self._handle_response(response)
             except CredentialsNotValid as e:
                 # NOTE: CredentialsNotValid is raised only when the server
                 # explicitly indicates that the credentials are not valid and
@@ -4438,7 +4448,7 @@ class RestZenStore(BaseZenStore):
                     # request again, this time with a valid API token in the
                     # header.
                     logger.debug(
-                        f"The last request with ID {request_id} was not "
+                        f"[{request_id}] The last request was not "
                         f"authenticated: {e}\n"
                         "Re-authenticating and retrying..."
                     )
@@ -4466,7 +4476,7 @@ class RestZenStore(BaseZenStore):
                     # that was rejected by the server. We attempt a
                     # re-authentication here and then retry the request.
                     logger.debug(
-                        f"The last request with ID {request_id} was authenticated "
+                        f"[{request_id}] The last request was authenticated "
                         "with an API token that was rejected by the server: "
                         f"{e}\n"
                         "Re-authenticating and retrying..."
@@ -4480,7 +4490,7 @@ class RestZenStore(BaseZenStore):
                     # The last request was made after re-authenticating but
                     # still failed. Bailing out.
                     logger.debug(
-                        f"The last request with ID {request_id} failed after "
+                        f"[{request_id}] The last request failed after "
                         "re-authenticating: {e}\n"
                         "Bailing out..."
                     )
@@ -4492,7 +4502,7 @@ class RestZenStore(BaseZenStore):
                 end_time = time.time()
                 duration = (end_time - start_time) * 1000
                 logger.debug(
-                    f"Request to {path} with request ID {request_id} took "
+                    f"[{request_id}] {status_code} {method} {path} took "
                     f"{duration:.2f}ms."
                 )
 

zenml/zen_stores/schemas/__init__.py

@@ -35,6 +35,7 @@ from zenml.zen_stores.schemas.flavor_schemas import FlavorSchema
 from zenml.zen_stores.schemas.server_settings_schemas import ServerSettingsSchema
 from zenml.zen_stores.schemas.pipeline_deployment_schemas import (
     PipelineDeploymentSchema,
+    StepConfigurationSchema,
 )
 from zenml.zen_stores.schemas.pipeline_run_schemas import PipelineRunSchema
 from zenml.zen_stores.schemas.pipeline_schemas import PipelineSchema
@@ -74,6 +75,7 @@ from zenml.zen_stores.schemas.model_schemas import (
 )
 from zenml.zen_stores.schemas.run_template_schemas import RunTemplateSchema
 from zenml.zen_stores.schemas.server_settings_schemas import ServerSettingsSchema
+from zenml.zen_stores.schemas.api_transaction_schemas import ApiTransactionSchema
 
 __all__ = [
     "ActionSchema",
@@ -91,6 +93,7 @@ __all__ = [
     "OAuthDeviceSchema",
     "PipelineBuildSchema",
     "PipelineDeploymentSchema",
+    "StepConfigurationSchema",
     "PipelineRunSchema",
     "PipelineSchema",
     "RunMetadataResourceSchema",
@@ -119,4 +122,5 @@ __all__ = [
     "ModelVersionArtifactSchema",
     "ModelVersionPipelineRunSchema",
     "ProjectSchema",
+    "ApiTransactionSchema",
 ]

zenml/zen_stores/schemas/api_transaction_schemas.py

@@ -0,0 +1,141 @@
+#  Copyright (c) ZenML GmbH 2025. 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.
+"""SQLModel implementation of idempotent API transaction tables."""
+
+from datetime import datetime, timedelta
+from typing import Any, Optional
+from uuid import UUID
+
+from sqlalchemy import TEXT, Column, String
+from sqlalchemy.dialects.mysql import MEDIUMTEXT
+from sqlmodel import Field
+
+from zenml.constants import MEDIUMTEXT_MAX_LENGTH
+from zenml.models import (
+    ApiTransactionRequest,
+    ApiTransactionResponse,
+    ApiTransactionResponseBody,
+    ApiTransactionUpdate,
+)
+from zenml.utils.time_utils import utc_now
+from zenml.zen_stores.schemas.base_schemas import BaseSchema
+from zenml.zen_stores.schemas.schema_utils import (
+    build_foreign_key_field,
+    build_index,
+)
+from zenml.zen_stores.schemas.user_schemas import UserSchema
+
+
+class ApiTransactionSchema(BaseSchema, table=True):
+    """SQL Model for API transactions."""
+
+    __tablename__ = "api_transaction"
+    __table_args__ = (
+        build_index(
+            table_name=__tablename__,
+            column_names=[
+                "completed",
+                "expired",
+            ],
+        ),
+    )
+    method: str
+    url: str = Field(sa_column=Column(TEXT, nullable=False))
+    completed: bool = Field(default=False)
+    result: Optional[str] = Field(
+        default=None,
+        sa_column=Column(
+            String(length=MEDIUMTEXT_MAX_LENGTH).with_variant(
+                MEDIUMTEXT, "mysql"
+            ),
+            nullable=True,
+        ),
+    )
+    expired: Optional[datetime] = Field(default=None, nullable=True)
+
+    user_id: UUID = build_foreign_key_field(
+        source=__tablename__,
+        target=UserSchema.__tablename__,
+        source_column="user_id",
+        target_column="id",
+        ondelete="CASCADE",
+        nullable=False,
+    )
+
+    @classmethod
+    def from_request(
+        cls, request: ApiTransactionRequest
+    ) -> "ApiTransactionSchema":
+        """Create a new API transaction from a request.
+
+        Args:
+            request: The API transaction request.
+
+        Returns:
+            The API transaction schema.
+        """
+        assert request.user is not None, "User must be set."
+        return cls(
+            id=request.transaction_id,
+            user_id=request.user,
+            method=request.method,
+            url=request.url,
+            completed=False,
+        )
+
+    def to_model(
+        self,
+        include_metadata: bool = False,
+        include_resources: bool = False,
+        **kwargs: Any,
+    ) -> ApiTransactionResponse:
+        """Convert the SQL model to a ZenML model.
+
+        Args:
+            include_metadata: Whether to include metadata in the response.
+            include_resources: Whether to include resources in the response.
+            **kwargs: Additional keyword arguments.
+
+        Returns:
+            The API transaction response.
+        """
+        response = ApiTransactionResponse(
+            id=self.id,
+            body=ApiTransactionResponseBody(
+                method=self.method,
+                url=self.url,
+                created=self.created,
+                updated=self.updated,
+                user_id=self.user_id,
+                completed=self.completed,
+            ),
+        )
+        if self.result is not None:
+            response.set_result(self.result)
+        return response
+
+    def update(self, update: ApiTransactionUpdate) -> "ApiTransactionSchema":
+        """Update the API transaction.
+
+        Args:
+            update: The API transaction update.
+
+        Returns:
+            The API transaction schema.
+        """
+        if update.result is not None:
+            self.result = update.get_result()
+        self.updated = utc_now()
+        self.expired = self.updated + timedelta(seconds=update.cache_time)
+        return self

zenml/zen_stores/schemas/pipeline_deployment_schemas.py

@@ -17,11 +17,11 @@ import json
 from typing import TYPE_CHECKING, Any, List, Optional, Sequence
 from uuid import UUID
 
-from sqlalchemy import TEXT, Column, String
+from sqlalchemy import TEXT, Column, String, UniqueConstraint
 from sqlalchemy.dialects.mysql import MEDIUMTEXT
-from sqlalchemy.orm import joinedload
+from sqlalchemy.orm import joinedload, object_session
 from sqlalchemy.sql.base import ExecutableOption
-from sqlmodel import Field, Relationship
+from sqlmodel import Field, Relationship, asc, col, select
 
 from zenml.config.pipeline_configurations import PipelineConfiguration
 from zenml.config.pipeline_spec import PipelineSpec
@@ -69,14 +69,6 @@ class PipelineDeploymentSchema(BaseSchema, table=True):
             nullable=False,
         )
     )
-    step_configurations: str = Field(
-        sa_column=Column(
-            String(length=MEDIUMTEXT_MAX_LENGTH).with_variant(
-                MEDIUMTEXT, "mysql"
-            ),
-            nullable=False,
-        )
-    )
     client_environment: str = Field(sa_column=Column(TEXT, nullable=False))
     run_name_template: str = Field(nullable=False)
     client_version: str = Field(nullable=True)
@@ -174,6 +166,46 @@ class PipelineDeploymentSchema(BaseSchema, table=True):
     step_runs: List["StepRunSchema"] = Relationship(
         sa_relationship_kwargs={"cascade": "delete"}
     )
+    step_configurations: List["StepConfigurationSchema"] = Relationship(
+        sa_relationship_kwargs={
+            "cascade": "delete",
+            "order_by": "asc(StepConfigurationSchema.index)",
+        }
+    )
+    step_count: int
+
+    def get_step_configurations(
+        self, include: Optional[List[str]] = None
+    ) -> List["StepConfigurationSchema"]:
+        """Get step configurations for the deployment.
+
+        Args:
+            include: List of step names to include. If not given, all step
+                configurations will be included.
+
+        Raises:
+            RuntimeError: If no session for the schema exists.
+
+        Returns:
+            List of step configurations.
+        """
+        if session := object_session(self):
+            query = (
+                select(StepConfigurationSchema)
+                .where(StepConfigurationSchema.deployment_id == self.id)
+                .order_by(asc(StepConfigurationSchema.index))
+            )
+
+            if include:
+                query = query.where(
+                    col(StepConfigurationSchema.name).in_(include)
+                )
+
+            return list(session.execute(query).scalars().all())
+        else:
+            raise RuntimeError(
+                "Missing DB session to fetch step configurations."
+            )
 
     @classmethod
     def get_query_options(
@@ -230,14 +262,6 @@ class PipelineDeploymentSchema(BaseSchema, table=True):
         Returns:
             The created `PipelineDeploymentSchema`.
         """
-        # Don't include the merged config in the step configurations, we
-        # reconstruct it in the `to_model` method using the pipeline
-        # configuration.
-        step_configurations = {
-            invocation_id: step.model_dump(mode="json", exclude={"config"})
-            for invocation_id, step in request.step_configurations.items()
-        }
-
         client_env = json.dumps(request.client_environment)
         if len(client_env) > TEXT_FIELD_MAX_LENGTH:
             logger.warning(
@@ -257,10 +281,7 @@ class PipelineDeploymentSchema(BaseSchema, table=True):
             code_reference_id=code_reference_id,
             run_name_template=request.run_name_template,
             pipeline_configuration=request.pipeline_configuration.model_dump_json(),
-            step_configurations=json.dumps(
-                step_configurations,
-                sort_keys=False,
-            ),
+            step_count=len(request.step_configurations),
             client_environment=client_env,
             client_version=request.client_version,
             server_version=request.server_version,
@@ -278,6 +299,7 @@ class PipelineDeploymentSchema(BaseSchema, table=True):
         include_metadata: bool = False,
         include_resources: bool = False,
         include_python_packages: bool = False,
+        step_configuration_filter: Optional[List[str]] = None,
         **kwargs: Any,
     ) -> PipelineDeploymentResponse:
         """Convert a `PipelineDeploymentSchema` to a `PipelineDeploymentResponse`.
@@ -286,6 +308,9 @@ class PipelineDeploymentSchema(BaseSchema, table=True):
             include_metadata: Whether the metadata will be filled.
             include_resources: Whether the resources will be filled.
             include_python_packages: Whether the python packages will be filled.
+            step_configuration_filter: List of step configurations to include in
+                the response. If not given, all step configurations will be
+                included.
             **kwargs: Keyword arguments to allow schema specific logic
 
 
@@ -303,10 +328,13 @@ class PipelineDeploymentSchema(BaseSchema, table=True):
             pipeline_configuration = PipelineConfiguration.model_validate_json(
                 self.pipeline_configuration
             )
-            step_configurations = json.loads(self.step_configurations)
-            for invocation_id, step in step_configurations.items():
-                step_configurations[invocation_id] = Step.from_dict(
-                    step, pipeline_configuration
+            step_configurations = {}
+            for step_configuration in self.get_step_configurations(
+                include=step_configuration_filter
+            ):
+                step_configurations[step_configuration.name] = Step.from_dict(
+                    json.loads(step_configuration.config),
+                    pipeline_configuration,
                 )
 
             client_environment = json.loads(self.client_environment)
@@ -348,3 +376,36 @@ class PipelineDeploymentSchema(BaseSchema, table=True):
             metadata=metadata,
             resources=resources,
         )
+
+
+class StepConfigurationSchema(BaseSchema, table=True):
+    """SQL Model for step configurations."""
+
+    __tablename__ = "step_configuration"
+    __table_args__ = (
+        UniqueConstraint(
+            "deployment_id",
+            "name",
+            name="unique_step_name_for_deployment",
+        ),
+    )
+
+    index: int
+    name: str
+    config: str = Field(
+        sa_column=Column(
+            String(length=MEDIUMTEXT_MAX_LENGTH).with_variant(
+                MEDIUMTEXT, "mysql"
+            ),
+            nullable=False,
+        )
+    )
+
+    deployment_id: UUID = build_foreign_key_field(
+        source=__tablename__,
+        target=PipelineDeploymentSchema.__tablename__,
+        source_column="deployment_id",
+        target_column="id",
+        ondelete="CASCADE",
+        nullable=False,
+    )

zenml/zen_stores/schemas/pipeline_run_schemas.py

@@ -389,19 +389,36 @@ class PipelineRunSchema(NamedSchema, RunMetadataInterface, table=True):
             RuntimeError: if the model creation fails.
         """
         if self.deployment is not None:
-            deployment = self.deployment.to_model(
-                include_metadata=True,
-                include_python_packages=include_python_packages,
+            config = PipelineConfiguration.model_validate_json(
+                self.deployment.pipeline_configuration
             )
+            client_environment = json.loads(self.deployment.client_environment)
 
-            config = deployment.pipeline_configuration
-            client_environment = deployment.client_environment
-
-            stack = deployment.stack
-            pipeline = deployment.pipeline
-            build = deployment.build
-            schedule = deployment.schedule
-            code_reference = deployment.code_reference
+            stack = (
+                self.deployment.stack.to_model()
+                if self.deployment.stack
+                else None
+            )
+            pipeline = (
+                self.deployment.pipeline.to_model()
+                if self.deployment.pipeline
+                else None
+            )
+            build = (
+                self.deployment.build.to_model()
+                if self.deployment.build
+                else None
+            )
+            schedule = (
+                self.deployment.schedule.to_model()
+                if self.deployment.schedule
+                else None
+            )
+            code_reference = (
+                self.deployment.code_reference.to_model()
+                if self.deployment.code_reference
+                else None
+            )
 
         elif self.pipeline_configuration is not None:
             config = PipelineConfiguration.model_validate_json(

zenml/zen_stores/schemas/step_run_schemas.py

@@ -290,10 +290,10 @@ class StepRunSchema(NamedSchema, RunMetadataInterface, table=True):
         """
         step = None
         if self.deployment is not None:
-            step_configurations = json.loads(
-                self.deployment.step_configurations
+            step_configurations = self.deployment.get_step_configurations(
+                include=[self.name]
             )
-            if self.name in step_configurations:
+            if step_configurations:
                 pipeline_configuration = (
                     PipelineConfiguration.model_validate_json(
                         self.deployment.pipeline_configuration
@@ -304,7 +304,7 @@ class StepRunSchema(NamedSchema, RunMetadataInterface, table=True):
                     inplace=True,
                 )
                 step = Step.from_dict(
-                    step_configurations[self.name],
+                    json.loads(step_configurations[0].config),
                     pipeline_configuration=pipeline_configuration,
                 )
         if not step and self.step_configuration: