zenml-nightly 0.66.0.dev20240927__py3-none-any.whl → 0.68.1.dev20241102__py3-none-any.whl

zenml/zen_server/template_execution/utils.py

@@ -3,7 +3,7 @@
 import copy
 import hashlib
 import sys
-from typing import Any, Dict, List, Optional, Set, Tuple
+from typing import Any, Dict, List, Optional
 from uuid import UUID
 
 from fastapi import BackgroundTasks
@@ -22,11 +22,9 @@ from zenml.constants import (
     ENV_ZENML_ACTIVE_WORKSPACE_ID,
 )
 from zenml.enums import ExecutionStatus, StackComponentType, StoreType
-from zenml.integrations.utils import get_integration_for_module
 from zenml.logger import get_logger
 from zenml.models import (
     CodeReferenceRequest,
-    ComponentResponse,
     FlavorFilter,
     PipelineDeploymentRequest,
     PipelineDeploymentResponse,
@@ -35,15 +33,15 @@ from zenml.models import (
     RunTemplateResponse,
     StackResponse,
 )
-from zenml.new.pipelines.build_utils import compute_stack_checksum
-from zenml.new.pipelines.run_utils import (
+from zenml.pipelines.build_utils import compute_stack_checksum
+from zenml.pipelines.run_utils import (
     create_placeholder_run,
     get_default_run_name,
     validate_run_config_is_runnable_from_server,
     validate_stack_is_runnable_from_server,
 )
 from zenml.stack.flavor import Flavor
-from zenml.utils import dict_utils, settings_utils
+from zenml.utils import dict_utils, requirements_utils, settings_utils
 from zenml.zen_server.auth import AuthContext
 from zenml.zen_server.template_execution.runner_entrypoint_configuration import (
     RunnerEntrypointConfiguration,
@@ -151,8 +149,8 @@ def run_template(
     assert placeholder_run
 
     def _task() -> None:
-        pypi_requirements, apt_packages = get_requirements_for_stack(
-            stack=stack
+        pypi_requirements, apt_packages = (
+            requirements_utils.get_requirements_for_stack(stack=stack)
         )
 
         if build.python_version:
@@ -248,7 +246,7 @@ def ensure_async_orchestrator(
     """
     orchestrator = stack.components[StackComponentType.ORCHESTRATOR][0]
     flavors = zen_store().list_flavors(
-        FlavorFilter(name=orchestrator.flavor, type=orchestrator.type)
+        FlavorFilter(name=orchestrator.flavor_name, type=orchestrator.type)
     )
     flavor = Flavor.from_model(flavors[0])
 
@@ -266,59 +264,6 @@ def ensure_async_orchestrator(
         )
 
 
-def get_requirements_for_stack(
-    stack: StackResponse,
-) -> Tuple[List[str], List[str]]:
-    """Get requirements for a stack model.
-
-    Args:
-        stack: The stack for which to get the requirements.
-
-    Returns:
-        Tuple of PyPI and APT requirements of the stack.
-    """
-    pypi_requirements: Set[str] = set()
-    apt_packages: Set[str] = set()
-
-    for component_list in stack.components.values():
-        assert len(component_list) == 1
-        component = component_list[0]
-        (
-            component_pypi_requirements,
-            component_apt_packages,
-        ) = get_requirements_for_component(component=component)
-        pypi_requirements = pypi_requirements.union(
-            component_pypi_requirements
-        )
-        apt_packages = apt_packages.union(component_apt_packages)
-
-    return sorted(pypi_requirements), sorted(apt_packages)
-
-
-def get_requirements_for_component(
-    component: ComponentResponse,
-) -> Tuple[List[str], List[str]]:
-    """Get requirements for a component model.
-
-    Args:
-        component: The component for which to get the requirements.
-
-    Returns:
-        Tuple of PyPI and APT requirements of the component.
-    """
-    flavors = zen_store().list_flavors(
-        FlavorFilter(name=component.flavor, type=component.type)
-    )
-    assert len(flavors) == 1
-    flavor_source = flavors[0].source
-    integration = get_integration_for_module(module_name=flavor_source)
-
-    if integration:
-        return integration.get_requirements(), integration.APT_PACKAGES
-    else:
-        return [], []
-
-
 def generate_image_hash(dockerfile: str) -> str:
     """Generate a hash of the Dockerfile.
 
@@ -512,7 +457,7 @@ def get_pipeline_run_analytics_metadata(
     own_stack = stack_creator and stack_creator.id == deployment.user.id
 
     stack_metadata = {
-        component_type.value: component_list[0].flavor
+        component_type.value: component_list[0].flavor_name
         for component_type, component_list in stack.components.items()
     }
 

zenml/zen_server/utils.py

@@ -39,14 +39,14 @@ from zenml.constants import (
     INFO,
     VERSION_1,
 )
-from zenml.enums import ServerProviderType
+from zenml.enums import StoreType
 from zenml.exceptions import IllegalOperationError, OAuthError
 from zenml.logger import get_logger
 from zenml.plugins.plugin_flavor_registry import PluginFlavorRegistry
-from zenml.zen_server.deploy.deployment import ServerDeployment
-from zenml.zen_server.deploy.local.local_zen_server import (
-    LocalServerDeploymentConfig,
+from zenml.zen_server.deploy.deployment import (
+    LocalServerDeployment,
 )
+from zenml.zen_server.deploy.exceptions import ServerDeploymentNotFoundError
 from zenml.zen_server.exceptions import http_exception_from_error
 from zenml.zen_server.feature_gate.feature_gate_interface import (
     FeatureGateInterface,
@@ -237,86 +237,84 @@ 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["LocalServerDeployment"]:
+    """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()
+    from zenml.zen_server.deploy.deployer import LocalServerDeployer
 
-    if not servers:
+    deployer = LocalServerDeployer()
+    try:
+        return deployer.get_server()
+    except ServerDeploymentNotFoundError:
         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
+def connected_to_local_server() -> bool:
+    """Check if the client is connected to a local server.
 
+    Returns:
+        True if the client is connected to a local server, False otherwise.
+    """
+    from zenml.zen_server.deploy.deployer import LocalServerDeployer
 
-def get_active_server_details() -> Tuple[str, Optional[int]]:
-    """Get the URL of the current ZenML Server.
+    deployer = LocalServerDeployer()
+    return deployer.is_connected_to_server()
 
-    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.
 
-    Returns:
-        The URL and port of the currently active server.
+def show_dashboard(
+    local: bool = False,
+    ngrok_token: Optional[str] = None,
+) -> None:
+    """Show the ZenML dashboard.
+
+    Args:
+        local: Whether to show the dashboard for the local server or the
+            one for the active server.
+        ngrok_token: An ngrok auth token to use for exposing the ZenML
+            dashboard on a public domain. Primarily used for accessing the
+            dashboard in Colab.
 
     Raises:
-        RuntimeError: If no server is active.
+        RuntimeError: If no server is connected.
     """
-    # Check for connected servers first
-    gc = GlobalConfiguration()
-    if not gc.uses_default_store():
-        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.")
-
-    if server and server.status and server.status.url:
-        if isinstance(server.config, LocalServerDeploymentConfig):
-            return server.status.url, server.config.port
-        return server.status.url, None
-
-    raise RuntimeError(
-        "ZenML is not connected to any server right now. Please use "
-        "`zenml connect` to connect to a server or spin up a new local server "
-        "via `zenml up`."
-    )
+    from zenml.utils.dashboard_utils import show_dashboard
+    from zenml.utils.networking_utils import get_or_create_ngrok_tunnel
+
+    url: Optional[str] = None
+    if not local:
+        gc = GlobalConfiguration()
+        if gc.store_configuration.type == StoreType.REST:
+            url = gc.store_configuration.url
+
+    if not url:
+        # Else, check for local servers
+        server = get_local_server()
+        if server and server.status and server.status.url:
+            url = server.status.url
+
+    if not url:
+        raise RuntimeError(
+            "ZenML is not connected to any server right now. Please use "
+            "`zenml login` to connect to a server or spin up a new local server "
+            "via `zenml login --local`."
+        )
+
+    if ngrok_token:
+        parsed_url = urlparse(url)
+
+        ngrok_url = get_or_create_ngrok_tunnel(
+            ngrok_token=ngrok_token, port=parsed_url.port or 80
+        )
+        logger.debug(f"Tunneling dashboard from {url} to {ngrok_url}.")
+        url = ngrok_url
+
+    show_dashboard(url)
 
 
 F = TypeVar("F", bound=Callable[..., Any])

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

@@ -24,11 +24,9 @@ from typing import (
     Type,
     Union,
 )
-from urllib.parse import urlparse
 from uuid import UUID
 
 from pydantic import BaseModel, ConfigDict, model_validator
-from requests import ConnectionError
 
 import zenml
 from zenml.config.global_config import GlobalConfiguration
@@ -39,13 +37,11 @@ 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,
     StoreType,
 )
-from zenml.exceptions import AuthorizationException
 from zenml.logger import get_logger
 from zenml.models import (
     ServerDatabaseType,
@@ -57,9 +53,6 @@ from zenml.models import (
     WorkspaceResponse,
 )
 from zenml.utils.pydantic_utils import before_validator_handler
-from zenml.zen_stores.secrets_stores.sql_secrets_store import (
-    SqlSecretsStoreConfiguration,
-)
 from zenml.zen_stores.zen_store_interface import ZenStoreInterface
 
 logger = get_logger(__name__)
@@ -136,51 +129,10 @@ class BaseZenStore(
                 stack and user in the store will be skipped.
             **kwargs: Additional keyword arguments to pass to the Pydantic
                 constructor.
-
-        Raises:
-            RuntimeError: If the store cannot be initialized.
-            AuthorizationException: If the store cannot be initialized due to
-                authentication errors.
         """
         super().__init__(**kwargs)
 
-        try:
-            self._initialize()
-
-        # Handle cases where the ZenML server is not available
-        except ConnectionError as e:
-            error_message = (
-                "Cannot connect to the ZenML database because the ZenML server "
-                f"at {self.url} is not running."
-            )
-            if urlparse(self.url).hostname in ["localhost", "127.0.0.1"]:
-                recommendation = (
-                    "Please run `zenml down` and `zenml up` to restart the "
-                    "server."
-                )
-            else:
-                recommendation = (
-                    "Please run `zenml disconnect` and `zenml connect --url "
-                    f"{self.url}` to reconnect to the server."
-                )
-            raise RuntimeError(f"{error_message}\n{recommendation}") from e
-
-        except AuthorizationException as e:
-            raise AuthorizationException(
-                f"Authorization failed for store at '{self.url}'. Please check "
-                f"your credentials: {str(e)}"
-            )
-
-        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)}" + zenml_pro_extra
-            ) from e
+        self._initialize()
 
         if not skip_default_registrations:
             logger.debug("Initializing database")
@@ -293,6 +245,9 @@ class BaseZenStore(
         Returns:
             The default store configuration.
         """
+        from zenml.zen_stores.secrets_stores.sql_secrets_store import (
+            SqlSecretsStoreConfiguration,
+        )
         from zenml.zen_stores.sql_zen_store import SqlZenStoreConfiguration
 
         config = SqlZenStoreConfiguration(
@@ -426,7 +381,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 +394,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.67.0_release.py

@@ -0,0 +1,23 @@
+"""Release [0.67.0].
+
+Revision ID: 0.67.0
+Revises: 0.66.0
+Create Date: 2024-09-25 10:39:36.711715
+
+"""
+
+# revision identifiers, used by Alembic.
+revision = "0.67.0"
+down_revision = "0.66.0"
+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.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/0.68.1_release.py

@@ -0,0 +1,23 @@
+"""Release [0.68.1].
+
+Revision ID: 0.68.1
+Revises: 0.68.0
+Create Date: 2024-10-28 10:34:17.307663
+
+"""
+
+# revision identifiers, used by Alembic.
+revision = "0.68.1"
+down_revision = "0.68.0"
+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