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

zenml/zen_server/deploy/helm/values.yaml

@@ -10,12 +10,37 @@ zenml:
 
   replicaCount: 1
 
+  # The number of ZenML server worker threads to use. This controls the number
+  # of concurrent requests that each ZenML server pod/replica can handle at a
+  # time. If not specified, the default value is 40.
+  #
+  # NOTE: this value should be coordinated with the `zenml.database.poolSize`
+  # and `zenml.database.maxOverflow` values to ensure that the ZenML server
+  # workers do not block on database connections (i.e. the sum of the pool size
+  # and max overflow should be greater than or equal to the thread pool size).
+  #
+  # threadPoolSize: 40
+
   image:
     repository: zenmldocker/zenml-server
     pullPolicy: Always
     # Overrides the image tag whose default is the chart appVersion.
     tag:
 
+  # The URL where the ZenML server API is reachable. If not specified, the
+  # clients will use the same URL used to connect them to the ZenML server.
+  serverURL:
+
+  # The URL where the ZenML dashboard is reachable.
+  # If not specified, the `serverURL` value is used. This should be
+  # configured if the dashboard is served from a different URL than the
+  # ZenML server.
+  #
+  # This is value is used to compute the dashboard URLs during the web login
+  # authentication workflow, to print dashboard URLs in log messages when
+  # running a pipeline and for other similar tasks.
+  dashboardURL:
+
   debug: true
 
   # Flag to enable/disable the tracking process of the analytics
@@ -108,12 +133,6 @@ zenml:
     # authorization request.
     deviceAuthPollingInterval: 5
 
-    # The URL where the ZenML dashboard is hosted. Used to construct the OAuth
-    # 2.0 device authorization endpoint. If not set, a partial URL is returned
-    # to the client which is used to construct the full URL based on the
-    # server's root URL path.
-    dashboardURL:
-
     # The time in minutes that an OAuth 2.0 device is allowed to be used to
     # authenticate with the ZenML server. If not set or if
     # `zenml.auth.jwtTokenExpireMinutes` is not set, the devices are allowed to
@@ -171,13 +190,8 @@ zenml:
   # ```
   rootUrlPath:
 
-
-
   defaultProject: default
-  defaultUsername: default
-  # Use your own password here
-  defaultPassword: zenml
-
+  
   # Implicit authentication methods featured by service connectors that support
   # them are disabled by default, for security reasons. This is because they
   # allow users to authenticate to the cloud provider where ZenML is running
@@ -188,13 +202,61 @@ zenml:
   # used, which will not be persisted across pod restarts.
   # NOTE: the certificate files need to be copied in the helm chart folder and
   # the paths configured here need to be relative to the root of the helm chart.
-  database: {}
+  database:
     # url: "mysql://admin:password@zenml-mysql:3306/database"
     # sslCa: /path/to/ca.pem
     # sslCert: /path/to/client-cert.pem
     # sslKey: /path/to/client-key.pem
     # sslVerifyServerCert: True
 
+    # Connection pool settings (only relevant for MySQL databases).
+    #
+    # NOTE: these values should be coordinated with the `zenml.threadPoolSize`
+    # to ensure that the ZenML server workers do not block on database
+    # connections (i.e. the sum of the pool size and max overflow should be
+    # greater than or equal to the thread pool size).
+    #
+    # poolSize: 20
+    # maxOverflow: 20
+
+    # ZenML supports backing up the database before DB migrations are performed
+    # and restoring it in case of a DB migration failure. For more information,
+    # see the following documentation:
+    # https://docs.zenml.io/getting-started/deploying-zenml/deploy-with-helm#database-backup-and-recovery
+    #
+    # Several backup strategies are supported:
+    #
+    # disabled - no backup is performed
+    # in-memory - the database schema and data are stored in memory. This is
+    #             the fastest backup strategy, but the backup is not persisted
+    #             across pod restarts, so no manual intervention is possible
+    #             in case the automatic DB recovery fails after a failed DB
+    #             migration. Adequate memory resources should be allocated to
+    #             the ZenML server pod when using this backup strategy with
+    #             large databases.
+    #             This is the default backup strategy.
+    # dump-file - the database schema and data are dumped to a local file.
+    #             Users may optionally configure a persistent volume where
+    #             the dump file will be stored by setting the
+    #             `backupPVStorageSize` and optionally the
+    #             `backupPVStorageClass` options. If a
+    #             persistent volume is not configured, the dump file will be
+    #             stored in an emptyDir volume, which is not persisted.
+    #             NOTE: you should set the podSecurityContext.fsGroup option
+    #             if you are using a persistent volume to store the dump file.
+    # database - the database is copied to a backup database in the same
+    #            database server. This requires the `backupDatabase`
+    #            option to be set to the name of the backup database.
+    #            This backup strategy is only supported for MySQL
+    #            compatible databases and the user specified in the
+    #            database URL must have permissions to manage (create, drop, and
+    #            modify) the backup database in addition to the main
+    #            database.
+    backupStrategy: in-memory
+    # backupPVStorageClass: standard
+    # backupPVStorageSize: 1Gi
+    # backupDatabase: "zenml_backup"
+
 
   # Secrets store settings. This is used to store centralized secrets.
   secretsStore:
@@ -260,7 +322,7 @@ zenml:
     # 
     # For a list of supported authentication methods and their configuration
     # options, see the following documentation:
-    # https://docs.zenml.io/stacks-and-components/auth-management/aws-service-connector#authentication-methods
+    # https://docs.zenml.io/how-to/auth-management/aws-service-connector#authentication-methods
     # 
     # You can also use the ZenML CLI to get the list of supported authentication
     # methods and their configuration options, e.g.:
@@ -299,7 +361,8 @@ zenml:
         aws_access_key_id:
         aws_secret_access_key:
         # The AWS role ARN to use to authenticate with the AWS Secrets Manager
-        # instance. Required if the `authMethod` is set to `iam-role`.
+        # instance. Required if the `authMethod` is set to `iam-role`. Optional
+        # if the `authMethod` is set to `implicit`.
         role_arn:
         
       # The AWS region to use. This must be set to the region where the AWS
@@ -332,7 +395,7 @@ zenml:
     # 
     # For a list of supported authentication methods and their configuration
     # options, see the following documentation:
-    # https://docs.zenml.io/stacks-and-components/auth-management/gcp-service-connector#authentication-methods
+    # https://docs.zenml.io/how-to/auth-management/gcp-service-connector#authentication-methods
     # 
     # You can also use the ZenML CLI to get the list of supported authentication
     # methods and their configuration options, e.g.:
@@ -404,7 +467,7 @@ zenml:
     # 
     # For a list of supported authentication methods and their configuration
     # options, see the following documentation:
-    # https://docs.zenml.io/stacks-and-components/auth-management/azure-service-connector#authentication-methods
+    # https://docs.zenml.io/how-to/auth-management/azure-service-connector#authentication-methods
     # 
     # You can also use the ZenML CLI to get the list of supported authentication
     # methods and their configuration options, e.g.:
@@ -556,7 +619,7 @@ zenml:
     # 
     # For a list of supported authentication methods and their configuration
     # options, see the following documentation:
-    # https://docs.zenml.io/stacks-and-components/auth-management/aws-service-connector#authentication-methods
+    # https://docs.zenml.io/how-to/auth-management/aws-service-connector#authentication-methods
     # 
     # You can also use the ZenML CLI to get the list of supported authentication
     # methods and their configuration options, e.g.:
@@ -610,7 +673,7 @@ zenml:
     # 
     # For a list of supported authentication methods and their configuration
     # options, see the following documentation:
-    # https://docs.zenml.io/stacks-and-components/auth-management/gcp-service-connector#authentication-methods
+    # https://docs.zenml.io/how-to/auth-management/gcp-service-connector#authentication-methods
     # 
     # You can also use the ZenML CLI to get the list of supported authentication
     # methods and their configuration options, e.g.:
@@ -630,6 +693,8 @@ zenml:
       # - user-account - Use a GCP user account
       # - service-account - Use a GCP service account
       # - impersonation - Use the GCP service account impersonation feature
+      # - external-account - Use workload identity federation to authenticate
+      #                      with the GCP Secrets Manager API
       authMethod: service-account
 
       # The GCP Service Connector authentication configuration. This should
@@ -656,6 +721,11 @@ zenml:
         # `impersonation`.
         target_principal:
 
+        # The GCP workload identity credentials to use to authenticate with the
+        # GCP Secrets Manager instance. Required if the `authMethod` is set to
+        # `external-account`.
+        external_account_json:
+
     # Azure backup Key Vault secrets store configuration. Only relevant if the
     # `azure` backup secrets store type is configured.
     #
@@ -668,7 +738,7 @@ zenml:
     # 
     # For a list of supported authentication methods and their configuration
     # options, see the following documentation:
-    # https://docs.zenml.io/stacks-and-components/auth-management/azure-service-connector#authentication-methods
+    # https://docs.zenml.io/how-to/auth-management/azure-service-connector#authentication-methods
     # 
     # You can also use the ZenML CLI to get the list of supported authentication
     # methods and their configuration options, e.g.:
@@ -737,6 +807,49 @@ zenml:
       # variable in the `zenml.secretEnvironment` variable.
       class_path: my.custom.secrets.store.MyCustomSecretsStore
 
+  # The ZenML server's secure headers configuration. This can be used to
+  # enable, disable or set custom values for security headers in the ZenML
+  # server's HTTP responses. The following values can be set for any of the
+  # supported secure headers configuration options:
+  #
+  # - `enabled`, `on`, `true` or `yes` - enables the secure header with the
+  # default value.
+  # - `disabled`, `off`, `false`, `none` or `no` - disables the secure header
+  # entirely, so that it is not set in the ZenML server's HTTP responses.
+  # - any other value - sets the secure header to the specified value.  
+  secure_headers:
+    # The `Server` HTTP header value used to identify the server. The default
+    # value is the ZenML server ID.
+    server: enabled
+    # The `Strict-Transport-Security` HTTP header value. The default value is
+    # `max-age=63072000; includeSubDomains`.
+    hsts: enabled
+    # The `X-Frame-Options` HTTP header value. The default value is `SAMEORIGIN`.
+    xfo: enabled
+    # The `X-XSS-Protection` HTTP header value. The default value is `0`.
+    # NOTE: this header is deprecated and should not be customized anymore. The
+    # `Content-Security-Policy` header should be used instead.
+    xxp: enabled
+    # The `X-Content-Type-Options` HTTP header value. The default value is
+    # `nosniff`.
+    content: enabled
+    # The `Content-Security-Policy` HTTP header value. This is by default set
+    # to a strict CSP policy that only allows content from the origins required
+    # by the ZenML dashboard.
+    # NOTE: customizing this header is discouraged, as it may cause the ZenML
+    # dashboard to malfunction.
+    csp: enabled
+    # The `Referrer-Policy` HTTP header value. The default value is
+    # `no-referrer-when-downgrade`.
+    referrer: enabled
+    # The `Cache-Control` HTTP header value. The default value is
+    # `no-store, no-cache, must-revalidate`.
+    cache: enabled
+    # The `Permissions-Policy` HTTP header value. The default value is
+    # `accelerometer=(), camera=(), geolocation=(), gyroscope=(), magnetometer=(), microphone=(), payment=(), usb=()`.
+    permissions: enabled
+
+
   # Extra environment variables to set in the ZenML server container.
   environment: {}
 
@@ -787,7 +900,7 @@ serviceAccount:
 podAnnotations: {}
 
 podSecurityContext: {}
-  # fsGroup: 2000
+  # fsGroup: 1000 # if you're using a PVC for backup, this should necessarily be set.
 
 securityContext:
   runAsNonRoot: true

zenml/zen_server/exceptions.py

@@ -20,13 +20,17 @@ from pydantic import BaseModel
 
 from zenml.exceptions import (
     AuthorizationException,
+    CredentialsNotValid,
     DoesNotExistException,
     DuplicateRunNameError,
+    EntityCreationError,
     EntityExistsError,
     IllegalOperationError,
+    MethodNotAllowedError,
     SecretExistsError,
     StackComponentExistsError,
     StackExistsError,
+    SubscriptionUpgradeRequiredError,
     ValidationError,
     ZenKeyError,
 )
@@ -38,7 +42,7 @@ if TYPE_CHECKING:
 class ErrorModel(BaseModel):
     """Base class for error responses."""
 
-    detail: Any
+    detail: Optional[Any] = None
 
 
 error_response = dict(model=ErrorModel)
@@ -76,7 +80,10 @@ REST_API_EXCEPTIONS: List[Tuple[Type[Exception], int]] = [
     # 403 Forbidden
     (IllegalOperationError, 403),
     # 401 Unauthorized
+    (CredentialsNotValid, 401),
     (AuthorizationException, 401),
+    # 402 Payment required
+    (SubscriptionUpgradeRequiredError, 402),
     # 404 Not Found
     (DoesNotExistException, 404),
     (ZenKeyError, 404),
@@ -87,9 +94,12 @@ REST_API_EXCEPTIONS: List[Tuple[Type[Exception], int]] = [
     # 422 Unprocessable Entity
     (ValueError, 422),
     # 500 Internal Server Error
+    (EntityCreationError, 500),
     (RuntimeError, 500),
     # 501 Not Implemented,
     (NotImplementedError, 501),
+    # 405 Method Not Allowed
+    (MethodNotAllowedError, 405),
 ]
 
 
@@ -247,4 +257,13 @@ def exception_from_response(
 
         exception = default_exc
 
+    # There is one special case where we want to return a specific exception:
+    # 401 Unauthorized exceptions thrown directly by FastAPI in the course of
+    # authentication are interpreted as AuthorizationException, but we want to
+    # return CredentialsNotValid instead.
+    if response.status_code == 401:
+        if not isinstance(exception(), CredentialsNotValid):
+            if response.headers.get("WWW-Authenticate"):
+                return CredentialsNotValid(exc_msg)
+
     return exception(exc_msg)

zenml/zen_server/feature_gate/__init__.py

@@ -0,0 +1,13 @@
+#  Copyright (c) ZenML GmbH 2024. All Rights Reserved.
+#
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at:
+#
+#       https://www.apache.org/licenses/LICENSE-2.0
+#
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express
+#  or implied. See the License for the specific language governing
+#  permissions and limitations under the License.

zenml/zen_server/feature_gate/endpoint_utils.py

@@ -0,0 +1,61 @@
+#  Copyright (c) ZenML GmbH 2024. All Rights Reserved.
+#
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at:
+#
+#       https://www.apache.org/licenses/LICENSE-2.0
+#
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express
+#  or implied. See the License for the specific language governing
+#  permissions and limitations under the License.
+"""All endpoint utils for the feature gate implementations."""
+
+from uuid import UUID
+
+from zenml.zen_server.rbac.models import ResourceType
+from zenml.zen_server.utils import feature_gate, server_config
+
+
+def check_entitlement(resource_type: ResourceType) -> None:
+    """Queries the feature gate to see if the operation falls within the tenants entitlements.
+
+    Raises an exception if the user is not entitled to create an instance of the
+    resource. Otherwise, simply returns.
+
+    Args:
+        resource_type: The type of resource to check for.
+    """
+    if not server_config().feature_gate_enabled:
+        return
+    return feature_gate().check_entitlement(resource=resource_type)
+
+
+def report_usage(resource_type: ResourceType, resource_id: UUID) -> None:
+    """Reports the creation/usage of a feature/resource.
+
+    Args:
+        resource_type: The type of resource to report a usage for
+        resource_id: ID of the resource that was created.
+    """
+    if not server_config().feature_gate_enabled:
+        return
+    feature_gate().report_event(
+        resource=resource_type, resource_id=resource_id
+    )
+
+
+def report_decrement(resource_type: ResourceType, resource_id: UUID) -> None:
+    """Reports the deletion/deactivation of a feature/resource.
+
+    Args:
+        resource_type: The type of resource to report a decrement in count for.
+        resource_id: ID of the resource that was deleted.
+    """
+    if not server_config().feature_gate_enabled:
+        return
+    feature_gate().report_event(
+        resource=resource_type, resource_id=resource_id, is_decrement=True
+    )

zenml/zen_server/feature_gate/feature_gate_interface.py

@@ -0,0 +1,49 @@
+#  Copyright (c) ZenML GmbH 2024. All Rights Reserved.
+#
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at:
+#
+#       https://www.apache.org/licenses/LICENSE-2.0
+#
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express
+#  or implied. See the License for the specific language governing
+#  permissions and limitations under the License.
+"""Definition of the feature gate interface."""
+
+from abc import ABC, abstractmethod
+from uuid import UUID
+
+from zenml.zen_server.rbac.models import ResourceType
+
+
+class FeatureGateInterface(ABC):
+    """RBAC interface definition."""
+
+    @abstractmethod
+    def check_entitlement(self, resource: ResourceType) -> None:
+        """Checks if a user is entitled to create a resource.
+
+        Args:
+            resource: The resource the user wants to create
+
+        Raises:
+            UpgradeRequiredError in case a subscription limit is reached
+        """
+
+    @abstractmethod
+    def report_event(
+        self,
+        resource: ResourceType,
+        resource_id: UUID,
+        is_decrement: bool = False,
+    ) -> None:
+        """Reports the usage of a feature to the aggregator backend.
+
+        Args:
+            resource: The resource the user created
+            resource_id: ID of the resource that was created/deleted.
+            is_decrement: In case this event reports an actual decrement of usage
+        """

zenml/zen_server/feature_gate/zenml_cloud_feature_gate.py

@@ -0,0 +1,125 @@
+#  Copyright (c) ZenML GmbH 2024. All Rights Reserved.
+#
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at:
+#
+#       https://www.apache.org/licenses/LICENSE-2.0
+#
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express
+#  or implied. See the License for the specific language governing
+#  permissions and limitations under the License.
+"""ZenML Pro implementation of the feature gate."""
+
+from typing import Any, Dict
+from uuid import UUID
+
+from pydantic import BaseModel, Field
+
+from zenml.config.server_config import ServerConfiguration
+from zenml.exceptions import SubscriptionUpgradeRequiredError
+from zenml.logger import get_logger
+from zenml.zen_server.cloud_utils import cloud_connection
+from zenml.zen_server.feature_gate.feature_gate_interface import (
+    FeatureGateInterface,
+)
+from zenml.zen_server.rbac.models import ResourceType
+
+logger = get_logger(__name__)
+
+server_config = ServerConfiguration.get_server_config()
+
+ORGANIZATION_ID = server_config.metadata.get("organization_id", "unknown")
+
+USAGE_EVENT_ENDPOINT = "/usage-event"
+ENTITLEMENT_ENDPOINT = f"/organizations/{ORGANIZATION_ID}/entitlement"
+
+
+class RawUsageEvent(BaseModel):
+    """Model for reporting raw usage of a feature.
+
+    In case of consumables the UsageReport allows the Pricing Backend to
+    increment the usage per time-frame by 1.
+    """
+
+    organization_id: str = Field(
+        description="The organization that this usage can be attributed to.",
+    )
+    feature: ResourceType = Field(
+        description="The feature whose usage is being reported.",
+    )
+    total: int = Field(
+        description="The total amount of entities of this type."
+    )
+    metadata: Dict[str, Any] = Field(
+        default={},
+        description="Allows attaching additional metadata to events.",
+    )
+
+
+class ZenMLCloudFeatureGateInterface(FeatureGateInterface):
+    """ZenML Cloud Feature Gate implementation."""
+
+    def __init__(self) -> None:
+        """Initialize the object."""
+        self._connection = cloud_connection()
+
+    def check_entitlement(self, resource: ResourceType) -> None:
+        """Checks if a user is entitled to create a resource.
+
+        Args:
+            resource: The resource the user wants to create
+
+        Raises:
+            SubscriptionUpgradeRequiredError: in case a subscription limit is reached
+        """
+        try:
+            response = self._connection.get(
+                endpoint=ENTITLEMENT_ENDPOINT + "/" + resource, params=None
+            )
+        except SubscriptionUpgradeRequiredError:
+            raise SubscriptionUpgradeRequiredError(
+                f"Your subscription reached its `{resource}` limit. Please "
+                f"upgrade your subscription or reach out to us."
+            )
+
+        if response.status_code != 200:
+            logger.warning(
+                "Unexpected response status code from entitlement "
+                f"endpoint: {response.status_code}. Message: "
+                f"{response.json()}"
+            )
+
+    def report_event(
+        self,
+        resource: ResourceType,
+        resource_id: UUID,
+        is_decrement: bool = False,
+    ) -> None:
+        """Reports the usage of a feature to the aggregator backend.
+
+        Args:
+            resource: The resource the user created
+            resource_id: ID of the resource that was created/deleted.
+            is_decrement: In case this event reports an actual decrement of usage
+        """
+        data = RawUsageEvent(
+            organization_id=ORGANIZATION_ID,
+            feature=resource,
+            total=1 if not is_decrement else -1,
+            metadata={
+                "tenant_id": str(server_config.get_external_server_id()),
+                "resource_id": str(resource_id),
+            },
+        ).model_dump()
+        response = self._connection.post(
+            endpoint=USAGE_EVENT_ENDPOINT, data=data
+        )
+        if response.status_code != 200:
+            logger.error(
+                "Usage report not accepted by upstream backend. "
+                f"Status Code: {response.status_code}, Message: "
+                f"{response.json()}."
+            )