zenml-nightly 0.67.0.dev20241002__py3-none-any.whl → 0.68.0.dev20241027__py3-none-any.whl

zenml/zen_server/utils.py

@@ -39,7 +39,6 @@ from zenml.constants import (
     INFO,
     VERSION_1,
 )
-from zenml.enums import ServerProviderType
 from zenml.exceptions import IllegalOperationError, OAuthError
 from zenml.logger import get_logger
 from zenml.plugins.plugin_flavor_registry import PluginFlavorRegistry
@@ -237,45 +236,23 @@ def server_config() -> ServerConfiguration:
     return _server_config
 
 
-def get_active_deployment(local: bool = False) -> Optional["ServerDeployment"]:
-    """Get the active local or remote server deployment.
+def get_local_server() -> Optional["ServerDeployment"]:
+    """Get the active local server.
 
-    Call this function to retrieve the local or remote server deployment that
-    was last provisioned on this machine.
-
-    Args:
-        local: Whether to return the local active deployment or the remote one.
+    Call this function to retrieve the local server deployed on this machine.
 
     Returns:
-        The local or remote active server deployment or None, if no deployment
-        was found.
+        The local server deployment or None, if no local server deployment was
+        found.
     """
     from zenml.zen_server.deploy.deployer import ServerDeployer
 
     deployer = ServerDeployer()
-    if local:
-        servers = deployer.list_servers(provider_type=ServerProviderType.LOCAL)
-        if not servers:
-            servers = deployer.list_servers(
-                provider_type=ServerProviderType.DOCKER
-            )
-    else:
-        servers = deployer.list_servers()
-
+    servers = deployer.list_servers()
     if not servers:
         return None
 
-    for server in servers:
-        if server.config.provider in [
-            ServerProviderType.LOCAL,
-            ServerProviderType.DOCKER,
-        ]:
-            if local:
-                return server
-        elif not local:
-            return server
-
-    return None
+    return servers[0]
 
 
 def get_active_server_details() -> Tuple[str, Optional[int]]:
@@ -284,8 +261,7 @@ def get_active_server_details() -> Tuple[str, Optional[int]]:
     When multiple servers are present, the following precedence is used to
     determine which server to use:
     - If the client is connected to a server, that server has precedence.
-    - If no server is connected, a server that was deployed remotely has
-        precedence over a server that was deployed locally.
+    - If no server is connected, the local server is used, if present.
 
     Returns:
         The URL and port of the currently active server.
@@ -299,14 +275,8 @@ def get_active_server_details() -> Tuple[str, Optional[int]]:
         logger.debug("Getting URL of connected server.")
         parsed_url = urlparse(gc.store_configuration.url)
         return f"{parsed_url.scheme}://{parsed_url.hostname}", parsed_url.port
-    # Else, check for deployed servers
-    server = get_active_deployment(local=False)
-    if server:
-        logger.debug("Getting URL of remote server.")
-    else:
-        server = get_active_deployment(local=True)
-        logger.debug("Getting URL of local server.")
-
+    # Else, check for local servers
+    server = get_local_server()
     if server and server.status and server.status.url:
         if isinstance(server.config, LocalServerDeploymentConfig):
             return server.status.url, server.config.port

zenml/zen_server/zen_server_api.py

@@ -24,7 +24,7 @@ import os
 from asyncio.log import logger
 from datetime import datetime, timedelta, timezone
 from genericpath import isfile
-from typing import Any, List
+from typing import Any, List, Set
 
 from anyio import to_thread
 from fastapi import FastAPI, HTTPException, Request
@@ -32,8 +32,18 @@ from fastapi.exceptions import RequestValidationError
 from fastapi.responses import ORJSONResponse
 from fastapi.staticfiles import StaticFiles
 from fastapi.templating import Jinja2Templates
+from starlette.middleware.base import (
+    BaseHTTPMiddleware,
+    RequestResponseEndpoint,
+)
 from starlette.middleware.cors import CORSMiddleware
-from starlette.responses import FileResponse
+from starlette.responses import (
+    FileResponse,
+    JSONResponse,
+    RedirectResponse,
+    Response,
+)
+from starlette.types import ASGIApp
 
 import zenml
 from zenml.analytics import source_context
@@ -43,6 +53,7 @@ from zenml.constants import (
     HEALTH,
 )
 from zenml.enums import AuthScheme, SourceContextTypes
+from zenml.models import ServerDeploymentType
 from zenml.zen_server.exceptions import error_detail
 from zenml.zen_server.routers import (
     actions_endpoints,
@@ -93,10 +104,7 @@ from zenml.zen_server.utils import (
     zen_store,
 )
 
-if server_config().use_legacy_dashboard:
-    DASHBOARD_DIRECTORY = "dashboard_legacy"
-else:
-    DASHBOARD_DIRECTORY = "dashboard"
+DASHBOARD_DIRECTORY = "dashboard"
 
 
 def relative_path(rel: str) -> str:
@@ -143,6 +151,79 @@ def validation_exception_handler(
     return ORJSONResponse(error_detail(exc, ValueError), status_code=422)
 
 
+class RequestBodyLimit(BaseHTTPMiddleware):
+    """Limits the size of the request body."""
+
+    def __init__(self, app: ASGIApp, max_bytes: int) -> None:
+        """Limits the size of the request body.
+
+        Args:
+            app: The FastAPI app.
+            max_bytes: The maximum size of the request body.
+        """
+        super().__init__(app)
+        self.max_bytes = max_bytes
+
+    async def dispatch(
+        self, request: Request, call_next: RequestResponseEndpoint
+    ) -> Response:
+        """Limits the size of the request body.
+
+        Args:
+            request: The incoming request.
+            call_next: The next function to be called.
+
+        Returns:
+            The response to the request.
+        """
+        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)
+
+
+class RestrictFileUploadsMiddleware(BaseHTTPMiddleware):
+    """Restrict file uploads to certain paths."""
+
+    def __init__(self, app: FastAPI, allowed_paths: Set[str]):
+        """Restrict file uploads to certain paths.
+
+        Args:
+            app: The FastAPI app.
+            allowed_paths: The allowed paths.
+        """
+        super().__init__(app)
+        self.allowed_paths = allowed_paths
+
+    async def dispatch(
+        self, request: Request, call_next: RequestResponseEndpoint
+    ) -> Response:
+        """Restrict file uploads to certain paths.
+
+        Args:
+            request: The incoming request.
+            call_next: The next function to be called.
+
+        Returns:
+            The response to the request.
+        """
+        if request.method == "POST":
+            content_type = request.headers.get("content-type", "")
+            if (
+                "multipart/form-data" in content_type
+                and request.url.path not in self.allowed_paths
+            ):
+                return JSONResponse(
+                    status_code=403,
+                    content={
+                        "detail": "File uploads are not allowed on this endpoint."
+                    },
+                )
+        return await call_next(request)
+
+
+ALLOWED_FOR_FILE_UPLOAD: Set[str] = set()
+
 app.add_middleware(
     CORSMiddleware,
     allow_origins=server_config().cors_allow_origins,
@@ -151,6 +232,13 @@ app.add_middleware(
     allow_headers=["*"],
 )
 
+app.add_middleware(
+    RequestBodyLimit, max_bytes=server_config().max_request_body_size_in_bytes
+)
+app.add_middleware(
+    RestrictFileUploadsMiddleware, allowed_paths=ALLOWED_FOR_FILE_UPLOAD
+)
+
 
 @app.middleware("http")
 async def set_secure_headers(request: Request, call_next: Any) -> Any:
@@ -261,17 +349,14 @@ def initialize() -> None:
     initialize_secure_headers()
 
 
-if server_config().use_legacy_dashboard:
-    app.mount(
-        "/static",
-        StaticFiles(
-            directory=relative_path(
-                os.path.join(DASHBOARD_DIRECTORY, "static")
-            ),
-            check_dir=False,
-        ),
-    )
-else:
+DASHBOARD_REDIRECT_URL = None
+if (
+    server_config().dashboard_url
+    and server_config().deployment_type == ServerDeploymentType.CLOUD
+):
+    DASHBOARD_REDIRECT_URL = server_config().dashboard_url
+
+if not DASHBOARD_REDIRECT_URL:
     app.mount(
         "/assets",
         StaticFiles(
@@ -311,6 +396,9 @@ async def dashboard(request: Request) -> Any:
     Raises:
         HTTPException: If the dashboard files are not included.
     """
+    if DASHBOARD_REDIRECT_URL:
+        return RedirectResponse(url=DASHBOARD_REDIRECT_URL)
+
     if not os.path.isfile(
         os.path.join(relative_path(DASHBOARD_DIRECTORY), "index.html")
     ):
@@ -419,6 +507,8 @@ async def catch_all(request: Request, file_path: str) -> Any:
     Returns:
         The ZenML dashboard.
     """
+    if DASHBOARD_REDIRECT_URL:
+        return RedirectResponse(url=DASHBOARD_REDIRECT_URL)
     # some static files need to be served directly from the root dashboard
     # directory
     if file_path and file_path in root_static_files:

zenml/zen_stores/base_zen_store.py

@@ -426,7 +426,6 @@ class BaseZenStore(
         secrets_store_type = SecretsStoreType.NONE
         if isinstance(self, SqlZenStore) and self.config.secrets_store:
             secrets_store_type = self.config.secrets_store.type
-        use_legacy_dashboard = server_config.use_legacy_dashboard
         return ServerModel(
             id=GlobalConfiguration().user_id,
             active=True,
@@ -440,7 +439,6 @@ class BaseZenStore(
             dashboard_url=server_config.dashboard_url or "",
             analytics_enabled=GlobalConfiguration().analytics_opt_in,
             metadata=metadata,
-            use_legacy_dashboard=use_legacy_dashboard,
         )
 
     def is_local_store(self) -> bool:

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

@@ -0,0 +1,23 @@
+"""Release [0.68.0].
+
+Revision ID: 0.68.0
+Revises: 1d8f30c54477
+Create Date: 2024-10-24 13:14:56.912159
+
+"""
+
+# revision identifiers, used by Alembic.
+revision = "0.68.0"
+down_revision = "1d8f30c54477"
+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/1d8f30c54477_migrate_to_new_.py

@@ -0,0 +1,124 @@
+"""migrate to new ModelVersionDataLazyLoader structures [1d8f30c54477].
+
+Revision ID: 1d8f30c54477
+Revises: 0.67.0
+Create Date: 2024-10-08 15:55:34.727543
+
+"""
+
+import json
+
+import sqlalchemy as sa
+import sqlmodel
+from alembic import op
+
+# revision identifiers, used by Alembic.
+revision = "1d8f30c54477"
+down_revision = "0.67.0"
+branch_labels = None
+depends_on = None
+
+
+update_query_pd = sa.text(
+    "UPDATE pipeline_deployment SET step_configurations = :data WHERE id = :id_"
+)
+
+update_query_sr = sa.text(
+    "UPDATE step_run SET step_configuration = :data WHERE id = :id_"
+)
+
+
+def upgrade() -> None:
+    """Upgrade database schema and/or data, creating a new revision."""
+    bind = op.get_bind()
+    session = sqlmodel.Session(bind=bind)
+    # update pipeline_deployment
+    rows_pd = session.execute(
+        sa.text(
+            """
+             SELECT id, step_configurations 
+             FROM pipeline_deployment 
+             WHERE step_configurations IS NOT NULL
+             """
+        )
+    )
+    for id_, data in rows_pd:
+        try:
+            data_dict = json.loads(data)
+        except json.JSONDecodeError:
+            continue
+        has_changes = False
+        for k in data_dict:
+            if config := data_dict[k].get("config", {}):
+                if model_artifacts_or_metadata := config.get(
+                    "model_artifacts_or_metadata", {}
+                ):
+                    for (
+                        inp_name,
+                        inp_config,
+                    ) in model_artifacts_or_metadata.items():
+                        if model := inp_config.get("model", {}):
+                            del data_dict[k]["config"][
+                                "model_artifacts_or_metadata"
+                            ][inp_name]["model"]
+                            data_dict[k]["config"][
+                                "model_artifacts_or_metadata"
+                            ][inp_name]["model_name"] = model["name"]
+                            data_dict[k]["config"][
+                                "model_artifacts_or_metadata"
+                            ][inp_name]["model_version"] = str(
+                                model["version"]
+                            )
+                            has_changes = True
+        if has_changes:
+            data = json.dumps(data_dict)
+            session.execute(update_query_pd, params=(dict(data=data, id_=id_)))
+
+    # update step_run
+    rows_sr = session.execute(
+        sa.text(
+            """
+             SELECT id, step_configuration 
+             FROM step_run 
+             WHERE step_configuration IS NOT NULL
+             """
+        )
+    )
+    for id_, data in rows_sr:
+        try:
+            data_dict = json.loads(data)
+        except json.JSONDecodeError:
+            continue
+        has_changes = False
+        if config := data_dict.get("config", {}):
+            if model_artifacts_or_metadata := config.get(
+                "model_artifacts_or_metadata", {}
+            ):
+                for (
+                    inp_name,
+                    inp_config,
+                ) in model_artifacts_or_metadata.items():
+                    if model := inp_config.get("model", {}):
+                        del data_dict["config"]["model_artifacts_or_metadata"][
+                            inp_name
+                        ]["model"]
+                        data_dict["config"]["model_artifacts_or_metadata"][
+                            inp_name
+                        ]["model_name"] = model["name"]
+                        data_dict["config"]["model_artifacts_or_metadata"][
+                            inp_name
+                        ]["model_version"] = str(model["version"])
+                        has_changes = True
+        if has_changes:
+            data = json.dumps(data_dict)
+            session.execute(update_query_sr, params=(dict(data=data, id_=id_)))
+    session.commit()
+
+
+def downgrade() -> None:
+    """Downgrade database schema and/or data back to the previous revision.
+
+    Raises:
+        NotImplementedError: Downgrade is not supported for this migration.
+    """
+    raise NotImplementedError("Downgrade is not supported for this migration.")

zenml/zen_stores/schemas/pipeline_run_schemas.py

@@ -412,10 +412,7 @@ class PipelineRunSchema(NamedSchema, table=True):
         Returns:
             The updated `PipelineRunSchema`.
         """
-        if (
-            self.orchestrator_run_id
-            or self.status != ExecutionStatus.INITIALIZING
-        ):
+        if not self.is_placeholder_run():
             raise RuntimeError(
                 f"Unable to replace pipeline run {self.id} which is not a "
                 "placeholder run."
@@ -439,3 +436,14 @@ class PipelineRunSchema(NamedSchema, table=True):
         self.updated = datetime.utcnow()
 
         return self
+
+    def is_placeholder_run(self) -> bool:
+        """Whether the pipeline run is a placeholder run.
+
+        Returns:
+            Whether the pipeline run is a placeholder run.
+        """
+        return (
+            self.orchestrator_run_id is None
+            and self.status == ExecutionStatus.INITIALIZING
+        )

zenml/zen_stores/schemas/step_run_schemas.py

@@ -262,6 +262,8 @@ class StepRunSchema(NamedSchema, table=True):
         body = StepRunResponseBody(
             user=self.user.to_model() if self.user else None,
             status=ExecutionStatus(self.status),
+            start_time=self.start_time,
+            end_time=self.end_time,
             inputs=input_artifacts,
             outputs=output_artifacts,
             created=self.created,
@@ -278,8 +280,6 @@ class StepRunSchema(NamedSchema, table=True):
                 code_hash=self.code_hash,
                 docstring=self.docstring,
                 source_code=self.source_code,
-                start_time=self.start_time,
-                end_time=self.end_time,
                 logs=self.logs.to_model() if self.logs else None,
                 deployment_id=self.deployment_id,
                 pipeline_run_id=self.pipeline_run_id,

zenml/zen_stores/sql_zen_store.py

@@ -973,7 +973,6 @@ class SqlZenStore(BaseZenStore):
             ValueError: if the filtered page number is out of bounds.
             RuntimeError: if the schema does not have a `to_model` method.
         """
-        query = query.distinct()
         query = filter_model.apply_filter(query=query, table=table)
         query = query.distinct()
 
@@ -4108,7 +4107,8 @@ class SqlZenStore(BaseZenStore):
         Returns:
             A list of all pipelines matching the filter criteria.
         """
-        query = select(PipelineSchema)
+        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:
@@ -4141,7 +4141,12 @@ class SqlZenStore(BaseZenStore):
                     sort_clause = asc
 
                 query = (
-                    query.where(PipelineSchema.id == max_date_subquery.c.id)
+                    # 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
@@ -4149,6 +4154,14 @@ class SqlZenStore(BaseZenStore):
                     .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:
             return self.filter_and_paginate(
                 session=session,
@@ -4156,6 +4169,7 @@ class SqlZenStore(BaseZenStore):
                 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:
@@ -8093,7 +8107,9 @@ class SqlZenStore(BaseZenStore):
             session.commit()
             session.refresh(existing_step_run)
 
-            return existing_step_run.to_model(include_metadata=True)
+            return existing_step_run.to_model(
+                include_metadata=True, include_resources=True
+            )
 
     @staticmethod
     def _set_run_step_parent_step(
@@ -8304,6 +8320,14 @@ class SqlZenStore(BaseZenStore):
             num_steps=num_steps,
         )
 
+        if pipeline_run.is_placeholder_run() and not new_status.is_finished:
+            # If the pipeline run is a placeholder run, no step has been started
+            # for the run yet. This means the orchestrator hasn't started
+            # running yet, and this method is most likely being called as
+            # part of the creation of some cached steps. In this case, we don't
+            # update the status unless the run is finished.
+            return
+
         if new_status != pipeline_run.status:
             run_update = PipelineRunUpdate(status=new_status)
             if new_status in {
@@ -9181,7 +9205,7 @@ class SqlZenStore(BaseZenStore):
                 # Running inside server with external auth
                 return False
 
-            if config.auto_activate or config.use_legacy_dashboard:
+            if config.auto_activate:
                 return True
 
         else:

{zenml_nightly-0.67.0.dev20241002.dist-info → zenml_nightly-0.68.0.dev20241027.dist-info}/METADATA

@@ -1,20 +1,19 @@
 Metadata-Version: 2.1
 Name: zenml-nightly
-Version: 0.67.0.dev20241002
+Version: 0.68.0.dev20241027
 Summary: ZenML: Write production-ready ML code.
 Home-page: https://zenml.io
 License: Apache-2.0
 Keywords: machine learning,production,pipeline,mlops,devops
 Author: ZenML GmbH
 Author-email: info@zenml.io
-Requires-Python: >=3.8,<3.13
+Requires-Python: >=3.9,<3.13
 Classifier: Development Status :: 4 - Beta
 Classifier: Intended Audience :: Developers
 Classifier: Intended Audience :: Science/Research
 Classifier: Intended Audience :: System Administrators
 Classifier: License :: OSI Approved :: Apache Software License
 Classifier: Programming Language :: Python :: 3
-Classifier: Programming Language :: Python :: 3.8
 Classifier: Programming Language :: Python :: 3.9
 Classifier: Programming Language :: Python :: 3.10
 Classifier: Programming Language :: Python :: 3.11
@@ -31,7 +30,6 @@ Provides-Extra: connectors-gcp
 Provides-Extra: connectors-kubernetes
 Provides-Extra: dev
 Provides-Extra: gcsfs
-Provides-Extra: mlstacks
 Provides-Extra: s3fs
 Provides-Extra: sagemaker
 Provides-Extra: secrets-aws
@@ -85,7 +83,6 @@ Requires-Dist: mkdocs (>=1.2.3,<2.0.0) ; extra == "dev"
 Requires-Dist: mkdocs-awesome-pages-plugin (>=2.6.1,<3.0.0) ; extra == "dev"
 Requires-Dist: mkdocs-material (>=8.1.7,<9.0.0) ; extra == "dev"
 Requires-Dist: mkdocstrings (>=0.17.0,<0.18.0) ; extra == "dev"
-Requires-Dist: mlstacks (==0.10.0) ; extra == "mlstacks"
 Requires-Dist: mypy (==1.7.1) ; extra == "dev"
 Requires-Dist: orjson (>=3.10.0,<3.11.0) ; extra == "server"
 Requires-Dist: packaging (>=24.1)
@@ -105,7 +102,6 @@ Requires-Dist: pytest-rerunfailures (>=13.0) ; extra == "dev"
 Requires-Dist: pytest-split (>=0.8.1,<0.9.0) ; extra == "dev"
 Requires-Dist: python-dateutil (>=2.8.1,<3.0.0)
 Requires-Dist: python-multipart (>=0.0.9,<0.1.0) ; extra == "server"
-Requires-Dist: python-terraform (>=0.10.1,<0.11.0) ; extra == "terraform"
 Requires-Dist: pyyaml (>=6.0.1)
 Requires-Dist: pyyaml-include (<2.0) ; extra == "templates"
 Requires-Dist: rich[jupyter] (>=12.0.0)
@@ -231,7 +227,7 @@ Thank you for your support! 🌟
 ## 🤸 Quickstart
 [![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/zenml-io/zenml/blob/main/examples/quickstart/quickstart.ipynb)
 
-[Install ZenML](https://docs.zenml.io/getting-started/installation) via [PyPI](https://pypi.org/project/zenml/). Python 3.8 - 3.11 is required:
+[Install ZenML](https://docs.zenml.io/getting-started/installation) via [PyPI](https://pypi.org/project/zenml/). Python 3.9 - 3.12 is required:
 
 ```bash
 pip install "zenml[server]" notebook
@@ -472,7 +468,7 @@ the Apache License Version 2.0.
     <a href="https://github.com/zenml-io/zenml-projects">Projects Showcase</a>
     <br />
     <br />
-    🎉 Version 0.67.0 is out. Check out the release notes
+    🎉 Version 0.68.0 is out. Check out the release notes
     <a href="https://github.com/zenml-io/zenml/releases">here</a>.
     <br />
     🖥️ Download our VS Code Extension <a href="https://marketplace.visualstudio.com/items?itemName=ZenML.zenml-vscode">here</a>.