zenml-nightly 0.55.0.dev20240124__py3-none-any.whl → 0.72.0.dev20250116__py3-none-any.whl

zenml/zen_server/template_execution/workload_manager_interface.py

@@ -0,0 +1,92 @@
+"""Workload manager interface definition."""
+
+from abc import ABC, abstractmethod
+from typing import Dict, List, Optional
+from uuid import UUID
+
+
+class WorkloadManagerInterface(ABC):
+    """Workload manager interface."""
+
+    @abstractmethod
+    def run(
+        self,
+        workload_id: UUID,
+        image: str,
+        command: List[str],
+        arguments: List[str],
+        environment: Optional[Dict[str, str]] = None,
+        sync: bool = True,
+        timeout_in_seconds: int = 0,
+    ) -> None:
+        """Run a Docker container.
+
+        Args:
+            workload_id: Workload ID.
+            image: The Docker image to run.
+            command: The command to run in the container.
+            arguments: The arguments for the command.
+            environment: The environment to set in the container.
+            sync: If True, will wait until the container finished running before
+                returning.
+            timeout_in_seconds: Timeout in seconds to wait before cancelling
+                the container. If set to 0 the container will run until it
+                fails or finishes.
+        """
+        pass
+
+    @abstractmethod
+    def build_and_push_image(
+        self,
+        workload_id: UUID,
+        dockerfile: str,
+        image_name: str,
+        sync: bool = True,
+        timeout_in_seconds: int = 0,
+    ) -> str:
+        """Build and push a Docker image.
+
+        Args:
+            workload_id: Workload ID.
+            dockerfile: The dockerfile content to build the image.
+            image_name: The image repository and tag.
+            sync: If True, will wait until the build finished before returning.
+            timeout_in_seconds: Timeout in seconds to wait before cancelling
+                the container. If set to 0 the container will run until it
+                fails or finishes.
+
+        Returns:
+            The full image name including container registry.
+        """
+        pass
+
+    @abstractmethod
+    def delete_workload(self, workload_id: UUID) -> None:
+        """Delete a workload.
+
+        Args:
+            workload_id: Workload ID.
+        """
+        pass
+
+    @abstractmethod
+    def get_logs(self, workload_id: UUID) -> str:
+        """Get logs for a workload.
+
+        Args:
+            workload_id: Workload ID.
+
+        Returns:
+            The stored logs.
+        """
+        pass
+
+    @abstractmethod
+    def log(self, workload_id: UUID, message: str) -> None:
+        """Log a message.
+
+        Args:
+            workload_id: Workload ID.
+            message: The message to log.
+        """
+        pass

zenml/zen_server/utils.py

@@ -16,31 +16,61 @@
 import inspect
 import os
 from functools import wraps
-from typing import Any, Callable, Optional, Tuple, Type, TypeVar, cast
+from typing import (
+    TYPE_CHECKING,
+    Any,
+    Callable,
+    Dict,
+    List,
+    Optional,
+    Tuple,
+    Type,
+    TypeVar,
+    cast,
+)
 from urllib.parse import urlparse
 
 from pydantic import BaseModel, ValidationError
 
+from zenml import __version__ as zenml_version
 from zenml.config.global_config import GlobalConfiguration
 from zenml.config.server_config import ServerConfiguration
 from zenml.constants import (
+    API,
     ENV_ZENML_SERVER,
+    INFO,
+    VERSION_1,
 )
-from zenml.enums import ServerProviderType
-from zenml.exceptions import OAuthError
+from zenml.enums import StoreType
+from zenml.exceptions import IllegalOperationError, OAuthError
 from zenml.logger import get_logger
-from zenml.zen_server.deploy.deployment import ServerDeployment
-from zenml.zen_server.deploy.local.local_zen_server import (
-    LocalServerDeploymentConfig,
+from zenml.plugins.plugin_flavor_registry import PluginFlavorRegistry
+from zenml.zen_server.cache import MemoryCache
+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,
+)
 from zenml.zen_server.rbac.rbac_interface import RBACInterface
+from zenml.zen_server.template_execution.workload_manager_interface import (
+    WorkloadManagerInterface,
+)
 from zenml.zen_stores.sql_zen_store import SqlZenStore
 
+if TYPE_CHECKING:
+    from fastapi import Request
+
 logger = get_logger(__name__)
 
 _zen_store: Optional["SqlZenStore"] = None
 _rbac: Optional[RBACInterface] = None
+_feature_gate: Optional[FeatureGateInterface] = None
+_workload_manager: Optional[WorkloadManagerInterface] = None
+_plugin_flavor_registry: Optional[PluginFlavorRegistry] = None
+_memcache: Optional[MemoryCache] = None
 
 
 def zen_store() -> "SqlZenStore":
@@ -58,6 +88,19 @@ def zen_store() -> "SqlZenStore":
     return _zen_store
 
 
+def plugin_flavor_registry() -> PluginFlavorRegistry:
+    """Get the plugin flavor registry.
+
+    Returns:
+        The plugin flavor registry.
+    """
+    global _plugin_flavor_registry
+    if _plugin_flavor_registry is None:
+        _plugin_flavor_registry = PluginFlavorRegistry()
+        _plugin_flavor_registry.initialize_plugins()
+    return _plugin_flavor_registry
+
+
 def rbac() -> RBACInterface:
     """Return the initialized RBAC component.
 
@@ -86,6 +129,79 @@ def initialize_rbac() -> None:
         _rbac = implementation_class()
 
 
+def feature_gate() -> FeatureGateInterface:
+    """Return the initialized Feature Gate component.
+
+    Raises:
+        RuntimeError: If the RBAC component is not initialized.
+
+    Returns:
+        The RBAC component.
+    """
+    global _feature_gate
+    if _feature_gate is None:
+        raise RuntimeError("Feature gate component not initialized.")
+    return _feature_gate
+
+
+def initialize_feature_gate() -> None:
+    """Initialize the Feature Gate component."""
+    global _feature_gate
+
+    if (
+        feature_gate_source
+        := server_config().feature_gate_implementation_source
+    ):
+        from zenml.utils import source_utils
+
+        implementation_class = source_utils.load_and_validate_class(
+            feature_gate_source, expected_class=FeatureGateInterface
+        )
+        _feature_gate = implementation_class()
+
+
+def workload_manager() -> WorkloadManagerInterface:
+    """Return the initialized workload manager component.
+
+    Raises:
+        RuntimeError: If the workload manager component is not initialized.
+
+    Returns:
+        The workload manager component.
+    """
+    global _workload_manager
+    if _workload_manager is None:
+        raise RuntimeError("Workload manager component not initialized")
+    return _workload_manager
+
+
+def initialize_workload_manager() -> None:
+    """Initialize the workload manager component.
+
+    This does not fail if the source can't be loaded but only logs a warning.
+    """
+    global _workload_manager
+
+    if source := server_config().workload_manager_implementation_source:
+        from zenml.utils import source_utils
+
+        try:
+            workload_manager_class: Type[WorkloadManagerInterface] = (
+                source_utils.load_and_validate_class(
+                    source=source, expected_class=WorkloadManagerInterface
+                )
+            )
+        except (ModuleNotFoundError, KeyError):
+            logger.warning("Unable to load workload manager source.")
+        else:
+            _workload_manager = workload_manager_class()
+
+
+def initialize_plugins() -> None:
+    """Initialize the event plugins registry."""
+    plugin_flavor_registry().initialize_plugins()
+
+
 def initialize_zen_store() -> None:
     """Initialize the ZenML Store.
 
@@ -110,6 +226,31 @@ def initialize_zen_store() -> None:
     _zen_store = zen_store_
 
 
+def initialize_memcache(max_capacity: int, default_expiry: int) -> None:
+    """Initialize the memory cache.
+
+    Args:
+        max_capacity: The maximum capacity of the cache.
+        default_expiry: The default expiry time in seconds.
+    """
+    global _memcache
+    _memcache = MemoryCache(max_capacity, default_expiry)
+
+
+def memcache() -> MemoryCache:
+    """Return the memory cache.
+
+    Returns:
+        The memory cache.
+
+    Raises:
+        RuntimeError: If the memory cache is not initialized.
+    """
+    if _memcache is None:
+        raise RuntimeError("Memory cache not initialized")
+    return _memcache
+
+
 _server_config: Optional[ServerConfiguration] = None
 
 
@@ -125,86 +266,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
 
+    deployer = LocalServerDeployer()
+    return deployer.is_connected_to_server()
 
-def get_active_server_details() -> Tuple[str, Optional[int]]:
-    """Get the URL of the current ZenML 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.
+def show_dashboard(
+    local: bool = False,
+    ngrok_token: Optional[str] = None,
+) -> None:
+    """Show the ZenML dashboard.
 
-    Returns:
-        The URL and port of the currently active server.
+    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() and gc.store is not None:
-        logger.debug("Getting URL of connected server.")
-        parsed_url = urlparse(gc.store.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])
@@ -284,6 +423,8 @@ def make_dependable(cls: Type[BaseModel]) -> Callable[..., Any]:
     """
     from fastapi import Query
 
+    from zenml.zen_server.exceptions import error_detail
+
     def init_cls_and_handle_errors(*args: Any, **kwargs: Any) -> BaseModel:
         from fastapi import HTTPException
 
@@ -291,9 +432,8 @@ def make_dependable(cls: Type[BaseModel]) -> Callable[..., Any]:
             inspect.signature(init_cls_and_handle_errors).bind(*args, **kwargs)
             return cls(*args, **kwargs)
         except ValidationError as e:
-            for error in e.errors():
-                error["loc"] = tuple(["query"] + list(error["loc"]))
-            raise HTTPException(422, detail=e.errors())
+            detail = error_detail(e, exception_type=ValueError)
+            raise HTTPException(422, detail=detail)
 
     params = {v.name: v for v in inspect.signature(cls).parameters.values()}
     query_params = getattr(cls, "API_MULTI_INPUT_PARAMS", [])
@@ -335,3 +475,163 @@ def get_ip_location(ip_address: str) -> Tuple[str, str, str]:
     except Exception:
         logger.exception(f"Could not get IP location for {ip_address}.")
         return "", "", ""
+
+
+def verify_admin_status_if_no_rbac(
+    admin_status: Optional[bool],
+    action: Optional[str] = None,
+) -> None:
+    """Validate the admin status for sensitive requests.
+
+    Only add this check in endpoints meant for admin use only.
+
+    Args:
+        admin_status: Whether the user is an admin or not. This is only used
+            if explicitly specified in the call and even if passed will be
+            ignored, if RBAC is enabled.
+        action: The action that is being performed, used for output only.
+
+    Raises:
+        IllegalOperationError: If the admin status is not valid.
+    """
+    if not server_config().rbac_enabled:
+        if not action:
+            action = "this action"
+        else:
+            action = f"`{action.strip('`')}`"
+
+        if admin_status is False:
+            raise IllegalOperationError(
+                message=f"Only admin users can perform {action} "
+                "without RBAC enabled.",
+            )
+    return
+
+
+def is_user_request(request: "Request") -> bool:
+    """Determine if the incoming request is a user request.
+
+    This function checks various aspects of the request to determine
+    if it's a user-initiated request or a system request.
+
+    Args:
+        request: The incoming FastAPI request object.
+
+    Returns:
+        True if it's a user request, False otherwise.
+    """
+    # Define system paths that should be excluded
+    system_paths: List[str] = [
+        "/health",
+        "/metrics",
+        "/system",
+        "/docs",
+        "/redoc",
+        "/openapi.json",
+    ]
+
+    user_prefix = f"{API}{VERSION_1}"
+    excluded_user_apis = [INFO]
+    # Check if this is not an excluded endpoint
+    if request.url.path in [
+        user_prefix + suffix for suffix in excluded_user_apis
+    ]:
+        return False
+
+    # Check if this is other user request
+    if request.url.path.startswith(user_prefix):
+        return True
+
+    # Exclude system paths
+    if any(request.url.path.startswith(path) for path in system_paths):
+        return False
+
+    # Exclude requests with specific headers
+    if request.headers.get("X-System-Request") == "true":
+        return False
+
+    # Exclude requests from certain user agents (e.g., monitoring tools)
+    user_agent = request.headers.get("User-Agent", "").lower()
+    system_agents = ["prometheus", "datadog", "newrelic", "pingdom"]
+    if any(agent in user_agent for agent in system_agents):
+        return False
+
+    # Check for internal IP addresses
+    client_host = request.client.host if request.client else None
+    if client_host and (
+        client_host.startswith("10.") or client_host.startswith("192.168.")
+    ):
+        return False
+
+    # Exclude OPTIONS requests (often used for CORS preflight)
+    if request.method == "OPTIONS":
+        return False
+
+    # Exclude specific query parameters that might indicate system requests
+    if request.query_params.get("system_check"):
+        return False
+
+    # If none of the above conditions are met, consider it a user request
+    return True
+
+
+def is_same_or_subdomain(source_domain: str, target_domain: str) -> bool:
+    """Check if the source domain is the same or a subdomain of the target domain.
+
+    Examples:
+        is_same_or_subdomain("example.com", "example.com") -> True
+        is_same_or_subdomain("alpha.example.com", "example.com") -> True
+        is_same_or_subdomain("alpha.example.com", ".example.com") -> True
+        is_same_or_subdomain("example.com", "alpha.example.com") -> False
+        is_same_or_subdomain("alpha.beta.example.com", "beta.example.com") -> True
+        is_same_or_subdomain("alpha.beta.example.com", "alpha.example.com") -> False
+        is_same_or_subdomain("alphabeta.gamma.example", "beta.gamma.example") -> False
+
+    Args:
+        source_domain: The source domain to check.
+        target_domain: The target domain to compare against.
+
+    Returns:
+        True if the source domain is the same or a subdomain of the target
+        domain, False otherwise.
+    """
+    import tldextract
+
+    # Extract the registered domain and suffix for both
+    src_parts = tldextract.extract(source_domain)
+    tgt_parts = tldextract.extract(target_domain)
+
+    if src_parts == tgt_parts:
+        return True  # Same domain
+
+    # Reconstruct the base domains (e.g., example.com)
+    src_base_domain = f"{src_parts.domain}.{src_parts.suffix}"
+    tgt_base_domain = f"{tgt_parts.domain}.{tgt_parts.suffix}"
+
+    if src_base_domain != tgt_base_domain:
+        return False  # Different base domains
+
+    if tgt_parts.subdomain == "":
+        return True  # Subdomain
+
+    if src_parts.subdomain.endswith(f".{tgt_parts.subdomain.lstrip('.')}"):
+        return True  # Subdomain of subdomain
+
+    return False
+
+
+def get_zenml_headers() -> Dict[str, str]:
+    """Get the ZenML specific headers to be included in requests made by the server.
+
+    Returns:
+        The ZenML specific headers.
+    """
+    config = server_config()
+    headers = {
+        "zenml-server-id": str(config.get_external_server_id()),
+        "zenml-server-version": zenml_version,
+    }
+    if config.server_url:
+        headers["zenml-server-url"] = config.server_url
+
+    return headers