zenml-nightly 0.70.0.dev20241120__py3-none-any.whl → 0.70.0.dev20241125__py3-none-any.whl

zenml/zen_server/deploy/helm/templates/server-deployment.yaml

@@ -63,14 +63,14 @@ spec:
               value: "True"
             {{- end }}
             {{- if .Values.zenml.database.url }}
-            - name: ZENML_STORE_TYPE
-              value: sql
             - name: DISABLE_DATABASE_MIGRATION
               value: "True"
-            - name: ZENML_STORE_SSL_VERIFY_SERVER_CERT
-              value: {{ .Values.zenml.database.sslVerifyServerCert | default "false" | quote }}
             {{- end }}
 
+            {{- range $k, $v := include "zenml.storeEnvVariables" . | fromYaml }}
+            - name: {{ $k }}
+              value: {{ $v | quote }}
+            {{- end }}
             {{- range $k, $v := include "zenml.serverEnvVariables" . | fromYaml }}
             - name: {{ $k }}
               value: {{ $v | quote }}
@@ -104,6 +104,18 @@ spec:
             httpGet:
               path: /health
               port: http
+          lifecycle:
+            preStop:
+              exec:
+                # Give the process 15 more seconds before the SIGTERM signal is
+                # sent. This allows the endpoint removal to reach the ingress
+                # controller in time and for traffic to be routed away from the
+                # pod before it is shut down. This eliminates the number of 502
+                # errors returned to the user.
+                #
+                # See https://learnk8s.io/graceful-shutdown for more information.
+                #
+                command: ["sleep", "15"]
           resources:
             {{- toYaml .Values.resources | nindent 12 }}
       {{- with .Values.nodeSelector }}

zenml/zen_server/deploy/helm/templates/server-secret.yaml

@@ -10,23 +10,8 @@ data:
   {{- else }}
   ZENML_SERVER_JWT_SECRET_KEY: {{ $prevServerSecret.data.ZENML_SERVER_JWT_SECRET_KEY | default (randAlphaNum 32 | b64enc | quote) }}
   {{- end }}
-  {{- if .Values.zenml.database.url }}
-  ZENML_STORE_URL: {{ .Values.zenml.database.url | b64enc | quote }}
-  {{- if .Values.zenml.database.sslCa }}
-  ZENML_STORE_SSL_CA: {{ .Files.Get .Values.zenml.database.sslCa | b64enc }}
-  {{- end }}
-  {{- if .Values.zenml.database.sslCert }}
-  ZENML_STORE_SSL_CERT: {{ .Files.Get .Values.zenml.database.sslCert | b64enc }}
-  {{- end }}
-  {{- if .Values.zenml.database.sslKey }}
-  ZENML_STORE_SSL_KEY: {{ .Files.Get .Values.zenml.database.sslKey | b64enc }}
-  {{- end }}
-  {{- if .Values.zenml.database.poolSize }}
-  ZENML_STORE_POOL_SIZE: {{ .Values.zenml.database.poolSize | b64enc | quote }}
-  {{- end }}
-  {{- if .Values.zenml.database.maxOverflow }}
-  ZENML_STORE_MAX_OVERFLOW: {{ .Values.zenml.database.maxOverflow | b64enc | quote }}
-  {{- end }}
+  {{- range $k, $v := include "zenml.storeSecretEnvVariables" . | fromYaml}}
+  {{ $k }}: {{ $v | b64enc | quote }}
   {{- end }}
   {{- range $k, $v := include "zenml.secretsStoreSecretEnvVariables" . | fromYaml}}
   {{ $k }}: {{ $v | b64enc | quote }}

zenml/zen_server/routers/model_versions_endpoints.py

@@ -29,9 +29,11 @@ from zenml.constants import (
 )
 from zenml.models import (
     ModelVersionArtifactFilter,
+    ModelVersionArtifactRequest,
     ModelVersionArtifactResponse,
     ModelVersionFilter,
     ModelVersionPipelineRunFilter,
+    ModelVersionPipelineRunRequest,
     ModelVersionPipelineRunResponse,
     ModelVersionResponse,
     ModelVersionUpdate,
@@ -198,6 +200,34 @@ model_version_artifacts_router = APIRouter(
 )
 
 
+@model_version_artifacts_router.post(
+    "",
+    responses={401: error_response, 409: error_response, 422: error_response},
+)
+@handle_exceptions
+def create_model_version_artifact_link(
+    model_version_artifact_link: ModelVersionArtifactRequest,
+    _: AuthContext = Security(authorize),
+) -> ModelVersionArtifactResponse:
+    """Create a new model version to artifact link.
+
+    Args:
+        model_version_artifact_link: The model version to artifact link to create.
+
+    Returns:
+        The created model version to artifact link.
+    """
+    model_version = zen_store().get_model_version(
+        model_version_artifact_link.model_version
+    )
+    verify_permission_for_model(model_version, action=Action.UPDATE)
+
+    mv = zen_store().create_model_version_artifact_link(
+        model_version_artifact_link
+    )
+    return mv
+
+
 @model_version_artifacts_router.get(
     "",
     response_model=Page[ModelVersionArtifactResponse],
@@ -291,6 +321,35 @@ model_version_pipeline_runs_router = APIRouter(
 )
 
 
+@model_version_pipeline_runs_router.post(
+    "",
+    responses={401: error_response, 409: error_response, 422: error_response},
+)
+@handle_exceptions
+def create_model_version_pipeline_run_link(
+    model_version_pipeline_run_link: ModelVersionPipelineRunRequest,
+    _: AuthContext = Security(authorize),
+) -> ModelVersionPipelineRunResponse:
+    """Create a new model version to pipeline run link.
+
+    Args:
+        model_version_pipeline_run_link: The model version to pipeline run link to create.
+
+    Returns:
+        - If Model Version to Pipeline Run Link already exists - returns the existing link.
+        - Otherwise, returns the newly created model version to pipeline run link.
+    """
+    model_version = zen_store().get_model_version(
+        model_version_pipeline_run_link.model_version, hydrate=False
+    )
+    verify_permission_for_model(model_version, action=Action.UPDATE)
+
+    mv = zen_store().create_model_version_pipeline_run_link(
+        model_version_pipeline_run_link
+    )
+    return mv
+
+
 @model_version_pipeline_runs_router.get(
     "",
     response_model=Page[ModelVersionPipelineRunResponse],

zenml/zen_server/routers/server_endpoints.py

@@ -22,6 +22,7 @@ from zenml.constants import (
     ACTIVATE,
     API,
     INFO,
+    LOAD_INFO,
     ONBOARDING_STATE,
     SERVER_SETTINGS,
     VERSION_1,
@@ -30,6 +31,7 @@ from zenml.enums import AuthScheme
 from zenml.exceptions import IllegalOperationError
 from zenml.models import (
     ServerActivationRequest,
+    ServerLoadInfo,
     ServerModel,
     ServerSettingsResponse,
     ServerSettingsUpdate,
@@ -71,6 +73,51 @@ def server_info() -> ServerModel:
     return zen_store().get_store_info()
 
 
+@router.get(
+    LOAD_INFO,
+    response_model=ServerLoadInfo,
+)
+@handle_exceptions
+def server_load_info(_: AuthContext = Security(authorize)) -> ServerLoadInfo:
+    """Get information about the server load.
+
+    Returns:
+        Information about the server load.
+    """
+    import threading
+
+    # Get the current number of threads
+    num_threads = len(threading.enumerate())
+
+    store = zen_store()
+
+    if store.config.driver == "sqlite":
+        # SQLite doesn't have a connection pool
+        return ServerLoadInfo(
+            threads=num_threads,
+            db_connections_total=0,
+            db_connections_active=0,
+            db_connections_overflow=0,
+        )
+
+    from sqlalchemy.pool import QueuePool
+
+    # Get the number of connections
+    pool = store.engine.pool
+    assert isinstance(pool, QueuePool)
+    idle_conn = pool.checkedin()
+    active_conn = pool.checkedout()
+    overflow_conn = max(0, pool.overflow())
+    total_conn = idle_conn + active_conn
+
+    return ServerLoadInfo(
+        threads=num_threads,
+        db_connections_total=total_conn,
+        db_connections_active=active_conn,
+        db_connections_overflow=overflow_conn,
+    )
+
+
 @router.get(
     ONBOARDING_STATE,
     responses={

zenml/zen_server/routers/workspaces_endpoints.py

@@ -20,7 +20,6 @@ from fastapi import APIRouter, Depends, Security
 
 from zenml.constants import (
     API,
-    ARTIFACTS,
     CODE_REPOSITORIES,
     GET_OR_CREATE,
     MODEL_VERSIONS,
@@ -54,10 +53,6 @@ from zenml.models import (
     ComponentResponse,
     ModelRequest,
     ModelResponse,
-    ModelVersionArtifactRequest,
-    ModelVersionArtifactResponse,
-    ModelVersionPipelineRunRequest,
-    ModelVersionPipelineRunResponse,
     ModelVersionRequest,
     ModelVersionResponse,
     Page,
@@ -1442,131 +1437,6 @@ def create_model_version(
     )
 
 
-@router.post(
-    WORKSPACES
-    + "/{workspace_name_or_id}"
-    + MODEL_VERSIONS
-    + "/{model_version_id}"
-    + ARTIFACTS,
-    response_model=ModelVersionArtifactResponse,
-    responses={401: error_response, 409: error_response, 422: error_response},
-)
-@handle_exceptions
-def create_model_version_artifact_link(
-    workspace_name_or_id: Union[str, UUID],
-    model_version_id: UUID,
-    model_version_artifact_link: ModelVersionArtifactRequest,
-    auth_context: AuthContext = Security(authorize),
-) -> ModelVersionArtifactResponse:
-    """Create a new model version to artifact link.
-
-    Args:
-        workspace_name_or_id: Name or ID of the workspace.
-        model_version_id: ID of the model version.
-        model_version_artifact_link: The model version to artifact link to create.
-        auth_context: Authentication context.
-
-    Returns:
-        The created model version to artifact link.
-
-    Raises:
-        IllegalOperationError: If the workspace or user specified in the
-            model version does not match the current workspace or authenticated
-            user.
-    """
-    workspace = zen_store().get_workspace(workspace_name_or_id)
-    if str(model_version_id) != str(model_version_artifact_link.model_version):
-        raise IllegalOperationError(
-            f"The model version id in your path `{model_version_id}` does not "
-            f"match the model version specified in the request model "
-            f"`{model_version_artifact_link.model_version}`"
-        )
-
-    if model_version_artifact_link.workspace != workspace.id:
-        raise IllegalOperationError(
-            "Creating model version to artifact links outside of the workspace scope "
-            f"of this endpoint `{workspace_name_or_id}` is "
-            f"not supported."
-        )
-    if model_version_artifact_link.user != auth_context.user.id:
-        raise IllegalOperationError(
-            "Creating model to artifact links for a user other than yourself "
-            "is not supported."
-        )
-
-    model_version = zen_store().get_model_version(model_version_id)
-    verify_permission_for_model(model_version, action=Action.UPDATE)
-
-    mv = zen_store().create_model_version_artifact_link(
-        model_version_artifact_link
-    )
-    return mv
-
-
-@router.post(
-    WORKSPACES
-    + "/{workspace_name_or_id}"
-    + MODEL_VERSIONS
-    + "/{model_version_id}"
-    + RUNS,
-    response_model=ModelVersionPipelineRunResponse,
-    responses={401: error_response, 409: error_response, 422: error_response},
-)
-@handle_exceptions
-def create_model_version_pipeline_run_link(
-    workspace_name_or_id: Union[str, UUID],
-    model_version_id: UUID,
-    model_version_pipeline_run_link: ModelVersionPipelineRunRequest,
-    auth_context: AuthContext = Security(authorize),
-) -> ModelVersionPipelineRunResponse:
-    """Create a new model version to pipeline run link.
-
-    Args:
-        workspace_name_or_id: Name or ID of the workspace.
-        model_version_id: ID of the model version.
-        model_version_pipeline_run_link: The model version to pipeline run link to create.
-        auth_context: Authentication context.
-
-    Returns:
-        - If Model Version to Pipeline Run Link already exists - returns the existing link.
-        - Otherwise, returns the newly created model version to pipeline run link.
-
-    Raises:
-        IllegalOperationError: If the workspace or user specified in the
-            model version does not match the current workspace or authenticated
-            user.
-    """
-    workspace = zen_store().get_workspace(workspace_name_or_id)
-    if str(model_version_id) != str(
-        model_version_pipeline_run_link.model_version
-    ):
-        raise IllegalOperationError(
-            f"The model version id in your path `{model_version_id}` does not "
-            f"match the model version specified in the request model "
-            f"`{model_version_pipeline_run_link.model_version}`"
-        )
-
-    if model_version_pipeline_run_link.workspace != workspace.id:
-        raise IllegalOperationError(
-            "Creating model versions outside of the workspace scope "
-            f"of this endpoint `{workspace_name_or_id}` is "
-            f"not supported."
-        )
-    if model_version_pipeline_run_link.user != auth_context.user.id:
-        raise IllegalOperationError(
-            "Creating models for a user other than yourself "
-            "is not supported."
-        )
-
-    model_version = zen_store().get_model_version(model_version_id)
-    verify_permission_for_model(model_version, action=Action.UPDATE)
-
-    mv = zen_store().create_model_version_pipeline_run_link(
-        model_version_pipeline_run_link
-    )
-    return mv
-
-
 @router.post(
     WORKSPACES + "/{workspace_name_or_id}" + SERVICES,
     response_model=ServiceResponse,

zenml/zen_server/zen_server_api.py

@@ -180,7 +180,15 @@ class RequestBodyLimit(BaseHTTPMiddleware):
         if content_length := request.headers.get("content-length"):
             if int(content_length) > self.max_bytes:
                 return Response(status_code=413)  # Request Entity Too Large
-        return await call_next(request)
+
+        try:
+            return await call_next(request)
+        except Exception:
+            logger.exception("An error occurred while processing the request")
+            return JSONResponse(
+                status_code=500,
+                content={"detail": "An unexpected error occurred."},
+            )
 
 
 class RestrictFileUploadsMiddleware(BaseHTTPMiddleware):
@@ -220,7 +228,15 @@ class RestrictFileUploadsMiddleware(BaseHTTPMiddleware):
                         "detail": "File uploads are not allowed on this endpoint."
                     },
                 )
-        return await call_next(request)
+
+        try:
+            return await call_next(request)
+        except Exception:
+            logger.exception("An error occurred while processing the request")
+            return JSONResponse(
+                status_code=500,
+                content={"detail": "An unexpected error occurred."},
+            )
 
 
 ALLOWED_FOR_FILE_UPLOAD: Set[str] = set()
@@ -252,13 +268,21 @@ async def set_secure_headers(request: Request, call_next: Any) -> Any:
     Returns:
         The response with secure headers set.
     """
+    try:
+        response = await call_next(request)
+    except Exception:
+        logger.exception("An error occurred while processing the request")
+        response = JSONResponse(
+            status_code=500,
+            content={"detail": "An unexpected error occurred."},
+        )
+
     # If the request is for the openAPI docs, don't set secure headers
     if request.url.path.startswith("/docs") or request.url.path.startswith(
         "/redoc"
     ):
-        return await call_next(request)
+        return response
 
-    response = await call_next(request)
     secure_headers().framework.fastapi(response)
     return response
 
@@ -298,7 +322,15 @@ async def track_last_user_activity(request: Request, call_next: Any) -> Any:
         zen_store()._update_last_user_activity_timestamp(
             last_user_activity=last_user_activity
         )
-    return await call_next(request)
+
+    try:
+        return await call_next(request)
+    except Exception:
+        logger.exception("An error occurred while processing the request")
+        return JSONResponse(
+            status_code=500,
+            content={"detail": "An unexpected error occurred."},
+        )
 
 
 @app.middleware("http")
@@ -330,7 +362,14 @@ async def infer_source_context(request: Request, call_next: Any) -> Any:
         )
         source_context.set(SourceContextTypes.API)
 
-    return await call_next(request)
+    try:
+        return await call_next(request)
+    except Exception:
+        logger.exception("An error occurred while processing the request")
+        return JSONResponse(
+            status_code=500,
+            content={"detail": "An unexpected error occurred."},
+        )
 
 
 @app.on_event("startup")

zenml/zen_stores/base_zen_store.py

@@ -42,6 +42,7 @@ from zenml.enums import (
     SecretsStoreType,
     StoreType,
 )
+from zenml.exceptions import IllegalOperationError
 from zenml.logger import get_logger
 from zenml.models import (
     ServerDatabaseType,
@@ -335,7 +336,7 @@ class BaseZenStore(
             # Ensure that the active stack is still valid
             try:
                 active_stack = self.get_stack(stack_id=active_stack_id)
-            except KeyError:
+            except (KeyError, IllegalOperationError):
                 logger.warning(
                     "The current %s active stack is no longer available. "
                     "Resetting the active stack to default.",

zenml/zen_stores/migrations/utils.py

@@ -273,30 +273,25 @@ class MigrationUtils(BaseModel):
                         + "\n);"
                     )
 
+                # Detect self-referential foreign keys from the table schema
+                has_self_referential_foreign_keys = False
+                for fk in table.foreign_keys:
+                    # Check if the foreign key points to the same table
+                    if fk.column.table == table:
+                        has_self_referential_foreign_keys = True
+                        break
+
                 # Store the table schema
                 store_db_info(
-                    dict(table=table.name, create_stmt=create_table_stmt)
+                    dict(
+                        table=table.name,
+                        create_stmt=create_table_stmt,
+                        self_references=has_self_referential_foreign_keys,
+                    )
                 )
 
                 # 2. extract the table data in batches
-
-                # If the table has a `created` column, we use it to sort
-                # the rows in the table starting with the oldest rows.
-                # This is to ensure that the rows are inserted in the
-                # correct order, since some tables have inner foreign key
-                # constraints.
-                if "created" in table.columns:
-                    order_by = [table.columns["created"]]
-                else:
-                    order_by = []
-                if "id" in table.columns:
-                    # If the table has an `id` column, we also use it to sort
-                    # the rows in the table, even if we already use "created"
-                    # to sort the rows. We need a unique field to sort the rows,
-                    # to break the tie between rows with the same "created"
-                    # date, otherwise the same entry might end up multiple times
-                    # in subsequent pages.
-                    order_by.append(table.columns["id"])
+                order_by = [col for col in table.primary_key]
 
                 # Fetch the number of rows in the table
                 row_count = conn.scalar(
@@ -305,7 +300,7 @@ class MigrationUtils(BaseModel):
 
                 # Fetch the data from the table in batches
                 if row_count is not None:
-                    batch_size = 50
+                    batch_size = 100
                     for i in range(0, row_count, batch_size):
                         rows = conn.execute(
                             table.select()
@@ -349,6 +344,7 @@ class MigrationUtils(BaseModel):
 
         with self.engine.begin() as connection:
             # read the DB information one JSON object at a time
+            self_references: Dict[str, bool] = {}
             for table_dump in load_db_info():
                 table_name = table_dump["table"]
                 if "create_stmt" in table_dump:
@@ -356,10 +352,22 @@ class MigrationUtils(BaseModel):
                     connection.execute(text(table_dump["create_stmt"]))
                     # Reload the database metadata after creating the table
                     metadata.reflect(bind=self.engine)
+                    self_references[table_name] = table_dump.get(
+                        "self_references", False
+                    )
 
                 if "data" in table_dump:
                     # insert the data into the database
                     table = metadata.tables[table_name]
+                    if self_references.get(table_name, False):
+                        # If the table has self-referential foreign keys, we
+                        # need to disable the foreign key checks before inserting
+                        # the rows and re-enable them afterwards. This is because
+                        # the rows need to be inserted in the correct order to
+                        # satisfy the foreign key constraints and we don't sort
+                        # the rows by creation time in the backup.
+                        connection.execute(text("SET FOREIGN_KEY_CHECKS = 0"))
+
                     for row in table_dump["data"]:
                         # Convert column values to the correct type
                         for column in table.columns:
@@ -372,10 +380,18 @@ class MigrationUtils(BaseModel):
                                     row[column.name], "utf-8"
                                 )
 
-                    # Insert the rows into the table
-                    connection.execute(
-                        table.insert().values(table_dump["data"])
-                    )
+                    # Insert the rows into the table in batches
+                    batch_size = 100
+                    for i in range(0, len(table_dump["data"]), batch_size):
+                        connection.execute(
+                            table.insert().values(
+                                table_dump["data"][i : i + batch_size]
+                            )
+                        )
+
+                    if table_dump.get("self_references", False):
+                        # Re-enable the foreign key checks after inserting the rows
+                        connection.execute(text("SET FOREIGN_KEY_CHECKS = 1"))
 
     def backup_database_to_file(self, dump_file: str) -> None:
         """Backup the database to a file.