zenml-nightly 0.62.0.dev20240729__py3-none-any.whl → 0.64.0.dev20240809__py3-none-any.whl

zenml/zen_server/zen_server_api.py

@@ -22,6 +22,7 @@ To run this file locally, execute:
 
 import os
 from asyncio.log import logger
+from datetime import datetime, timedelta, timezone
 from genericpath import isfile
 from typing import Any, List
 
@@ -36,7 +37,11 @@ from starlette.responses import FileResponse
 
 import zenml
 from zenml.analytics import source_context
-from zenml.constants import API, HEALTH
+from zenml.constants import (
+    API,
+    DEFAULT_ZENML_SERVER_REPORT_USER_ACTIVITY_TO_DB_SECONDS,
+    HEALTH,
+)
 from zenml.enums import AuthScheme, SourceContextTypes
 from zenml.zen_server.exceptions import error_detail
 from zenml.zen_server.routers import (
@@ -55,6 +60,7 @@ from zenml.zen_server.routers import (
     pipelines_endpoints,
     plugin_endpoints,
     run_metadata_endpoints,
+    run_templates_endpoints,
     runs_endpoints,
     schedule_endpoints,
     secrets_endpoints,
@@ -79,8 +85,10 @@ from zenml.zen_server.utils import (
     initialize_secure_headers,
     initialize_workload_manager,
     initialize_zen_store,
+    is_user_request,
     secure_headers,
     server_config,
+    zen_store,
 )
 
 if server_config().use_legacy_dashboard:
@@ -108,6 +116,12 @@ app = FastAPI(
     default_response_class=ORJSONResponse,
 )
 
+# Initialize last_user_activity
+last_user_activity: datetime = datetime.now(timezone.utc)
+last_user_activity_reported: datetime = datetime.now(timezone.utc) + timedelta(
+    seconds=-DEFAULT_ZENML_SERVER_REPORT_USER_ACTIVITY_TO_DB_SECONDS
+)
+
 
 # Customize the default request validation handler that comes with FastAPI
 # to return a JSON response that matches the ZenML API spec.
@@ -158,6 +172,44 @@ async def set_secure_headers(request: Request, call_next: Any) -> Any:
     return response
 
 
+@app.middleware("http")
+async def track_last_user_activity(request: Request, call_next: Any) -> Any:
+    """A middleware to track last user activity.
+
+    This middleware checks if the incoming request is a user request and
+    updates the last activity timestamp if it is.
+
+    Args:
+        request: The incoming request object.
+        call_next: A function that will receive the request as a parameter and
+            pass it to the corresponding path operation.
+
+    Returns:
+        The response to the request.
+    """
+    global last_user_activity
+    global last_user_activity_reported
+
+    try:
+        if is_user_request(request):
+            last_user_activity = datetime.now(timezone.utc)
+    except Exception as e:
+        logger.debug(
+            f"An unexpected error occurred while checking user activity: {e}"
+        )
+    if (
+        (
+            datetime.now(timezone.utc) - last_user_activity_reported
+        ).total_seconds()
+        > DEFAULT_ZENML_SERVER_REPORT_USER_ACTIVITY_TO_DB_SECONDS
+    ):
+        last_user_activity_reported = datetime.now(timezone.utc)
+        zen_store()._update_last_user_activity_timestamp(
+            last_user_activity=last_user_activity
+        )
+    return await call_next(request)
+
+
 @app.middleware("http")
 async def infer_source_context(request: Request, call_next: Any) -> Any:
     """A middleware to track the source of an event.
@@ -278,11 +330,11 @@ app.include_router(model_versions_endpoints.router)
 app.include_router(model_versions_endpoints.model_version_artifacts_router)
 app.include_router(model_versions_endpoints.model_version_pipeline_runs_router)
 app.include_router(pipelines_endpoints.router)
-app.include_router(pipelines_endpoints.namespace_router)
 app.include_router(pipeline_builds_endpoints.router)
 app.include_router(pipeline_deployments_endpoints.router)
 app.include_router(runs_endpoints.router)
 app.include_router(run_metadata_endpoints.router)
+app.include_router(run_templates_endpoints.router)
 app.include_router(schedule_endpoints.router)
 app.include_router(secrets_endpoints.router)
 app.include_router(secrets_endpoints.op_router)

zenml/zen_stores/base_zen_store.py

@@ -39,6 +39,7 @@ from zenml.constants import (
     DEFAULT_WORKSPACE_NAME,
     ENV_ZENML_DEFAULT_WORKSPACE_NAME,
     IS_DEBUG_ENV,
+    ZENML_PRO_CONNECTION_ISSUES_SUSPENDED_PAUSED_TENANT_HINT,
 )
 from zenml.enums import (
     SecretsStoreType,
@@ -171,9 +172,14 @@ class BaseZenStore(
             )
 
         except Exception as e:
+            zenml_pro_extra = ""
+            if ".zenml.io" in self.url:
+                zenml_pro_extra = (
+                    ZENML_PRO_CONNECTION_ISSUES_SUSPENDED_PAUSED_TENANT_HINT
+                )
             raise RuntimeError(
                 f"Error initializing {self.type.value} store with URL "
-                f"'{self.url}': {str(e)}"
+                f"'{self.url}': {str(e)}" + zenml_pro_extra
             ) from e
 
         if not skip_default_registrations:

zenml/zen_stores/migrations/versions/0.63.0_release.py

@@ -0,0 +1,23 @@
+"""Release [0.63.0].
+
+Revision ID: 0.63.0
+Revises: 7d1919bb1ef0
+Create Date: 2024-07-29 16:59:59.891939
+
+"""
+
+# revision identifiers, used by Alembic.
+revision = "0.63.0"
+down_revision = "7d1919bb1ef0"
+branch_labels = None
+depends_on = None
+
+
+def upgrade() -> None:
+    """Upgrade database schema and/or data, creating a new revision."""
+    pass
+
+
+def downgrade() -> None:
+    """Downgrade database schema and/or data back to the previous revision."""
+    pass

zenml/zen_stores/migrations/versions/0.64.0_release.py

@@ -0,0 +1,23 @@
+"""Release [0.64.0].
+
+Revision ID: 0.64.0
+Revises: 3dcc5d20e82f
+Create Date: 2024-08-08 12:25:12.058636
+
+"""
+
+# revision identifiers, used by Alembic.
+revision = "0.64.0"
+down_revision = "3dcc5d20e82f"
+branch_labels = None
+depends_on = None
+
+
+def upgrade() -> None:
+    """Upgrade database schema and/or data, creating a new revision."""
+    pass
+
+
+def downgrade() -> None:
+    """Downgrade database schema and/or data back to the previous revision."""
+    pass

zenml/zen_stores/migrations/versions/026d4577b6a0_add_code_path.py

@@ -0,0 +1,39 @@
+"""Add code path [026d4577b6a0].
+
+Revision ID: 026d4577b6a0
+Revises: 909550c7c4da
+Create Date: 2024-07-30 16:53:32.777594
+
+"""
+
+import sqlalchemy as sa
+import sqlmodel
+from alembic import op
+
+# revision identifiers, used by Alembic.
+revision = "026d4577b6a0"
+down_revision = "909550c7c4da"
+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! ###
+    with op.batch_alter_table("pipeline_deployment", schema=None) as batch_op:
+        batch_op.add_column(
+            sa.Column(
+                "code_path", sqlmodel.sql.sqltypes.AutoString(), nullable=True
+            )
+        )
+
+    # ### end Alembic commands ###
+
+
+def downgrade() -> None:
+    """Downgrade database schema and/or data back to the previous revision."""
+    # ### commands auto generated by Alembic - please adjust! ###
+    with op.batch_alter_table("pipeline_deployment", schema=None) as batch_op:
+        batch_op.drop_column("code_path")
+
+    # ### end Alembic commands ###

zenml/zen_stores/migrations/versions/3dcc5d20e82f_add_last_user_activity.py

@@ -0,0 +1,51 @@
+"""add last_user_activity [3dcc5d20e82f].
+
+Revision ID: 3dcc5d20e82f
+Revises: 909550c7c4da
+Create Date: 2024-08-07 14:49:07.623500
+
+"""
+
+from datetime import datetime, timezone
+
+import sqlalchemy as sa
+import sqlmodel
+from alembic import op
+
+# revision identifiers, used by Alembic.
+revision = "3dcc5d20e82f"
+down_revision = "026d4577b6a0"
+branch_labels = None
+depends_on = None
+
+
+def upgrade() -> None:
+    """Upgrade database schema and/or data, creating a new revision."""
+    bind = op.get_bind()
+    session = sqlmodel.Session(bind=bind)
+
+    with op.batch_alter_table("server_settings", schema=None) as batch_op:
+        batch_op.add_column(
+            sa.Column("last_user_activity", sa.DateTime(), nullable=True)
+        )
+
+    session.execute(
+        sa.text(
+            """
+            UPDATE server_settings
+            SET last_user_activity = :last_user_activity
+            """
+        ),
+        params=(dict(last_user_activity=datetime.now(timezone.utc))),
+    )
+
+    with op.batch_alter_table("server_settings", schema=None) as batch_op:
+        batch_op.alter_column(
+            "last_user_activity", existing_type=sa.DateTime(), nullable=False
+        )
+
+
+def downgrade() -> None:
+    """Downgrade database schema and/or data back to the previous revision."""
+    with op.batch_alter_table("server_settings", schema=None) as batch_op:
+        batch_op.drop_column("last_user_activity")

zenml/zen_stores/migrations/versions/7d1919bb1ef0_add_run_templates.py

@@ -0,0 +1,100 @@
+"""Add run templates [7d1919bb1ef0].
+
+Revision ID: 7d1919bb1ef0
+Revises: b59aa68fdb1f
+Create Date: 2024-07-22 11:20:00.544451
+
+"""
+
+import sqlalchemy as sa
+import sqlmodel
+from alembic import op
+from sqlalchemy.dialects import mysql
+
+# revision identifiers, used by Alembic.
+revision = "7d1919bb1ef0"
+down_revision = "b59aa68fdb1f"
+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(
+        "run_template",
+        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("name", sqlmodel.sql.sqltypes.AutoString(), nullable=False),
+        sa.Column(
+            "description",
+            sa.String(length=16777215).with_variant(mysql.MEDIUMTEXT, "mysql"),
+            nullable=True,
+        ),
+        sa.Column("user_id", sqlmodel.sql.sqltypes.GUID(), nullable=True),
+        sa.Column(
+            "workspace_id", sqlmodel.sql.sqltypes.GUID(), nullable=False
+        ),
+        sa.Column(
+            "source_deployment_id", sqlmodel.sql.sqltypes.GUID(), nullable=True
+        ),
+        sa.ForeignKeyConstraint(
+            ["source_deployment_id"],
+            ["pipeline_deployment.id"],
+            name="fk_run_template_source_deployment_id_pipeline_deployment",
+            ondelete="SET NULL",
+        ),
+        sa.ForeignKeyConstraint(
+            ["user_id"],
+            ["user.id"],
+            name="fk_run_template_user_id_user",
+            ondelete="SET NULL",
+        ),
+        sa.ForeignKeyConstraint(
+            ["workspace_id"],
+            ["workspace.id"],
+            name="fk_run_template_workspace_id_workspace",
+            ondelete="CASCADE",
+        ),
+        sa.PrimaryKeyConstraint("id"),
+        sa.UniqueConstraint(
+            "name", "workspace_id", name="unique_template_name_in_workspace"
+        ),
+    )
+    with op.batch_alter_table("pipeline_build", schema=None) as batch_op:
+        batch_op.add_column(
+            sa.Column(
+                "stack_checksum",
+                sqlmodel.sql.sqltypes.AutoString(),
+                nullable=True,
+            )
+        )
+        batch_op.drop_column("template_deployment_id")
+
+    with op.batch_alter_table("pipeline_deployment", schema=None) as batch_op:
+        batch_op.add_column(
+            sa.Column(
+                "template_id", sqlmodel.sql.sqltypes.GUID(), nullable=True
+            )
+        )
+
+    # ### end Alembic commands ###
+
+
+def downgrade() -> None:
+    """Downgrade database schema and/or data back to the previous revision."""
+    # ### commands auto generated by Alembic - please adjust! ###
+    with op.batch_alter_table("pipeline_deployment", schema=None) as batch_op:
+        batch_op.drop_column("template_id")
+
+    with op.batch_alter_table("pipeline_build", schema=None) as batch_op:
+        batch_op.add_column(
+            sa.Column(
+                "template_deployment_id", sa.CHAR(length=32), nullable=True
+            )
+        )
+        batch_op.drop_column("stack_checksum")
+
+    op.drop_table("run_template")
+    # ### end Alembic commands ###

zenml/zen_stores/migrations/versions/909550c7c4da_remove_user_hub_token.py

@@ -0,0 +1,36 @@
+"""Remove user hub token [909550c7c4da].
+
+Revision ID: 909550c7c4da
+Revises: 0.63.0
+Create Date: 2024-08-05 16:02:48.990897
+
+"""
+
+import sqlalchemy as sa
+from alembic import op
+
+# revision identifiers, used by Alembic.
+revision = "909550c7c4da"
+down_revision = "0.63.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! ###
+    with op.batch_alter_table("user", schema=None) as batch_op:
+        batch_op.drop_column("hub_token")
+
+    # ### end Alembic commands ###
+
+
+def downgrade() -> None:
+    """Downgrade database schema and/or data back to the previous revision."""
+    # ### commands auto generated by Alembic - please adjust! ###
+    with op.batch_alter_table("user", schema=None) as batch_op:
+        batch_op.add_column(
+            sa.Column("hub_token", sa.VARCHAR(), nullable=True)
+        )
+
+    # ### end Alembic commands ###

zenml/zen_stores/migrations/versions/b59aa68fdb1f_simplify_pipelines.py

@@ -0,0 +1,139 @@
+"""Simplify pipelines [b59aa68fdb1f].
+
+Revision ID: b59aa68fdb1f
+Revises: 0.62.0
+Create Date: 2024-07-04 14:00:32.830722
+
+"""
+
+from typing import Dict, Optional
+
+import sqlalchemy as sa
+import sqlmodel
+from alembic import op
+from sqlalchemy.dialects import mysql
+
+# revision identifiers, used by Alembic.
+revision = "b59aa68fdb1f"
+down_revision = "0.62.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! ###
+
+    with op.batch_alter_table("pipeline_deployment", schema=None) as batch_op:
+        batch_op.add_column(
+            sa.Column(
+                "pipeline_version_hash",
+                sqlmodel.sql.sqltypes.AutoString(),
+                nullable=True,
+            )
+        )
+        batch_op.add_column(
+            sa.Column(
+                "pipeline_spec",
+                sa.String(length=16777215).with_variant(
+                    mysql.MEDIUMTEXT, "mysql"
+                ),
+                nullable=True,
+            )
+        )
+
+    connection = op.get_bind()
+    meta = sa.MetaData()
+    meta.reflect(
+        bind=connection,
+        only=(
+            "pipeline",
+            "pipeline_run",
+            "pipeline_deployment",
+            "pipeline_build",
+            "schedule",
+        ),
+    )
+    pipeline_table = sa.Table("pipeline", meta)
+    pipeline_run_table = sa.Table("pipeline_run", meta)
+    pipeline_deployment_table = sa.Table("pipeline_deployment", meta)
+    pipeline_build_table = sa.Table("pipeline_build", meta)
+    schedule_table = sa.Table("schedule", meta)
+
+    def _migrate_pipeline_columns(
+        pipeline_id: str,
+        version_hash: Optional[str],
+        pipeline_spec: Optional[str],
+    ) -> None:
+        connection.execute(
+            sa.update(pipeline_deployment_table)
+            .where(pipeline_deployment_table.c.pipeline_id == pipeline_id)
+            .values(
+                pipeline_version_hash=version_hash, pipeline_spec=pipeline_spec
+            )
+        )
+
+    def _update_pipeline_fks(pipeline_id: str, replacement_id: str) -> None:
+        for table in [
+            pipeline_run_table,
+            pipeline_deployment_table,
+            pipeline_build_table,
+            schedule_table,
+        ]:
+            connection.execute(
+                sa.update(table)
+                .where(table.c.pipeline_id == pipeline_id)
+                .values(pipeline_id=replacement_id)
+            )
+
+    all_pipelines = connection.execute(sa.select(pipeline_table)).fetchall()
+    replacement_mapping: Dict[str, str] = {}
+
+    for pipeline in all_pipelines:
+        _migrate_pipeline_columns(
+            pipeline_id=pipeline.id,
+            version_hash=pipeline.version_hash,
+            pipeline_spec=pipeline.spec,
+        )
+
+        if replacement_id := replacement_mapping.get(pipeline.name):
+            _update_pipeline_fks(
+                pipeline_id=pipeline.id, replacement_id=replacement_id
+            )
+            connection.execute(
+                sa.delete(pipeline_table).where(
+                    pipeline_table.c.id == pipeline.id
+                )
+            )
+        else:
+            replacement_mapping[pipeline.name] = pipeline.id
+
+    with op.batch_alter_table("pipeline", schema=None) as batch_op:
+        batch_op.add_column(sa.Column("description", sa.TEXT(), nullable=True))
+        batch_op.drop_column("spec")
+        batch_op.drop_column("docstring")
+        batch_op.drop_column("version_hash")
+        batch_op.drop_column("version")
+
+    # ### end Alembic commands ###
+
+
+def downgrade() -> None:
+    """Downgrade database schema and/or data back to the previous revision."""
+    # ### commands auto generated by Alembic - please adjust! ###
+    with op.batch_alter_table("pipeline_deployment", schema=None) as batch_op:
+        batch_op.drop_column("pipeline_spec")
+        batch_op.drop_column("pipeline_version_hash")
+
+    with op.batch_alter_table("pipeline", schema=None) as batch_op:
+        batch_op.add_column(sa.Column("version", sa.VARCHAR(), nullable=False))
+        batch_op.add_column(
+            sa.Column("version_hash", sa.VARCHAR(), nullable=False)
+        )
+        batch_op.add_column(sa.Column("docstring", sa.TEXT(), nullable=True))
+        batch_op.add_column(
+            sa.Column("spec", sa.VARCHAR(length=16777215), nullable=False)
+        )
+        batch_op.drop_column("description")
+
+    # ### end Alembic commands ###