From 7fdcbce3e7c6b9b48e8a988b7410c2164b9f71a2 Mon Sep 17 00:00:00 2001 From: Stefan Nica Date: Wed, 13 Dec 2023 18:17:31 +0100 Subject: [PATCH 1/7] Add service connector support to AWS secrets store --- .../service_connectors/service_connector.py | 35 +++---- .../secrets_stores/aws_secrets_store.py | 98 +++++++++++++++---- 2 files changed, 97 insertions(+), 36 deletions(-) diff --git a/src/zenml/service_connectors/service_connector.py b/src/zenml/service_connectors/service_connector.py index 9c8842851a2..4a601c473d0 100644 --- a/src/zenml/service_connectors/service_connector.py +++ b/src/zenml/service_connectors/service_connector.py @@ -831,9 +831,6 @@ def has_expired(self) -> bool: True if the connector has expired, False otherwise. """ if not self.expires_at: - logger.info( - "connector authentication credentials have no expiration time." - ) return False expires_at = self.expires_at.replace(tzinfo=timezone.utc) @@ -967,6 +964,7 @@ def validate_runtime_args( def connect( self, + verify: bool = True, **kwargs: Any, ) -> Any: """Authenticate and connect to a resource. @@ -982,6 +980,8 @@ def connect( main service connector. Args: + verify: Whether to verify that the connector can access the + configured resource before connecting to it. kwargs: Additional implementation specific keyword arguments to use to configure the client. @@ -993,22 +993,23 @@ def connect( AuthorizationException: If the connector's authentication credentials have expired. """ - resource_type, resource_id = self.validate_runtime_args( - resource_type=self.resource_type, - resource_id=self.resource_id, - require_resource_type=True, - require_resource_id=True, - ) - - if self.has_expired(): - raise AuthorizationException( - "the connector's authentication credentials have expired." + if verify: + resource_type, resource_id = self.validate_runtime_args( + resource_type=self.resource_type, + resource_id=self.resource_id, + require_resource_type=True, + require_resource_id=True, ) - self._verify( - resource_type=resource_type, - resource_id=resource_id, - ) + if self.has_expired(): + raise AuthorizationException( + "the connector's authentication credentials have expired." + ) + + self._verify( + resource_type=resource_type, + resource_id=resource_id, + ) return self._connect_to_resource( **kwargs, diff --git a/src/zenml/zen_stores/secrets_stores/aws_secrets_store.py b/src/zenml/zen_stores/secrets_stores/aws_secrets_store.py index fa68b817ebb..a26bf4072c8 100644 --- a/src/zenml/zen_stores/secrets_stores/aws_secrets_store.py +++ b/src/zenml/zen_stores/secrets_stores/aws_secrets_store.py @@ -26,12 +26,13 @@ List, Optional, Type, + cast, ) -from uuid import UUID +from uuid import UUID, uuid4 import boto3 from botocore.exceptions import ClientError -from pydantic import SecretStr +from pydantic import BaseModel, Field, root_validator from zenml.analytics.enums import AnalyticsEvent from zenml.analytics.utils import track_decorator @@ -43,6 +44,12 @@ SecretsStoreType, ) from zenml.exceptions import EntityExistsError +from zenml.integrations.aws.service_connectors.aws_service_connector import ( + AWS_CONNECTOR_TYPE, + AWS_RESOURCE_TYPE, + AWSAuthenticationMethods, + AWSServiceConnector, +) from zenml.logger import get_logger from zenml.models import ( Page, @@ -50,6 +57,10 @@ SecretRequestModel, SecretResponseModel, SecretUpdateModel, + ServiceConnectorRequest, +) +from zenml.service_connectors.service_connector_registry import ( + service_connector_registry, ) from zenml.zen_stores.secrets_stores.base_secrets_store import ( BaseSecretsStore, @@ -61,6 +72,10 @@ AWS_ZENML_SECRET_NAME_PREFIX = "zenml" +class AWSSecretsStoreConnector(BaseModel): + """AWS secrets store connector configuration.""" + + class AWSSecretsStoreConfiguration(SecretsStoreConfiguration): """AWS secrets store configuration. @@ -86,18 +101,40 @@ class AWSSecretsStoreConfiguration(SecretsStoreConfiguration): """ type: SecretsStoreType = SecretsStoreType.AWS - region_name: str - aws_access_key_id: Optional[SecretStr] = None - aws_secret_access_key: Optional[SecretStr] = None - aws_session_token: Optional[SecretStr] = None + + auth_method: AWSAuthenticationMethods = AWSAuthenticationMethods.SECRET_KEY + auth_config: Dict[str, Any] = Field(default_factory=dict) + list_page_size: int = 100 secret_list_refresh_timeout: int = 0 + @root_validator(pre=True) + def populate_config(cls, values: Dict[str, Any]) -> Dict[str, Any]: + """Populate the connector configuration from legacy attributes. + + Args: + values: Dict representing user-specified runtime settings. + + Returns: + Validated settings. + + Raises: + ValueError: If the connector attribute is not set. + """ + if "auth_method" not in values or "auth_config" not in values: + values["auth_config"] = dict( + aws_access_key_id=values.get("aws_access_key_id"), + aws_secret_access_key=values.get("aws_secret_access_key"), + region=values.get("region_name"), + ) + + return values + class Config: """Pydantic configuration class.""" # Forbid extra attributes set in the class. - extra = "forbid" + extra = "allow" class AWSSecretsStore(BaseSecretsStore): @@ -159,6 +196,7 @@ class AWSSecretsStore(BaseSecretsStore): ] = AWSSecretsStoreConfiguration _client: Optional[Any] = None + _connector: Optional[AWSServiceConnector] = None @property def client(self) -> Any: @@ -167,21 +205,43 @@ def client(self) -> Any: Returns: The AWS Secrets Manager client. """ + if self._connector is not None: + # If the client connector expires, we'll try to get a new one. + if self._connector.has_expired(): + self._connector = None + self._client = None + + if self._connector is None: + # Initialize a base AWS service connector with the credentials from + # the configuration. + request = ServiceConnectorRequest( + name="secrets-store", + connector_type=AWS_CONNECTOR_TYPE, + resource_types=[AWS_RESOURCE_TYPE], + user=uuid4(), + workspace=uuid4(), + auth_method=self.config.auth_method, + configuration=self.config.auth_config, + ) + base_connector = service_connector_registry.instantiate_connector( + model=request + ) + self._connector = cast( + AWSServiceConnector, base_connector.get_connector_client() + ) + if self._client is None: # Initialize the AWS Secrets Manager client with the - # credentials from the configuration, if provided. - self._client = boto3.client( + # credentials from the connector. + session = self._connector.connect( + # Don't verify again because we already did that when we + # initialized the connector. + verify=False + ) + assert isinstance(session, boto3.Session) + self._client = session.client( "secretsmanager", - region_name=self.config.region_name, - aws_access_key_id=self.config.aws_access_key_id.get_secret_value() - if self.config.aws_access_key_id - else None, - aws_secret_access_key=self.config.aws_secret_access_key.get_secret_value() - if self.config.aws_secret_access_key - else None, - aws_session_token=self.config.aws_session_token.get_secret_value() - if self.config.aws_session_token - else None, + region_name=self._connector.config.region, ) return self._client From d4a47d6cfb88e2997a6fea92bcf6bb9be7776327 Mon Sep 17 00:00:00 2001 From: Stefan Nica Date: Thu, 14 Dec 2023 19:23:53 +0100 Subject: [PATCH 2/7] Add GCP and Azure service connectors support to secrets stores --- src/zenml/integrations/aws/__init__.py | 4 + .../flavors/aws_container_registry_flavor.py | 10 +- .../flavors/sagemaker_orchestrator_flavor.py | 7 +- .../flavors/sagemaker_step_operator_flavor.py | 4 +- .../aws_service_connector.py | 9 +- src/zenml/integrations/azure/__init__.py | 5 + .../flavors/azure_artifact_store_flavor.py | 11 +- .../azure_service_connector.py | 2 +- src/zenml/integrations/gcp/__init__.py | 4 + .../gcp/flavors/gcp_artifact_store_flavor.py | 4 +- .../gcp/flavors/gcp_image_builder_flavor.py | 10 +- .../gcp/flavors/vertex_orchestrator_flavor.py | 7 +- .../flavors/vertex_step_operator_flavor.py | 7 +- .../gcp_service_connector.py | 9 +- .../secrets_stores/aws_secrets_store.py | 136 ++++++---------- .../secrets_stores/azure_secrets_store.py | 141 ++++++++-------- .../secrets_stores/gcp_secrets_store.py | 111 ++++++++++--- .../service_connector_secrets_store.py | 153 ++++++++++++++++++ 18 files changed, 420 insertions(+), 214 deletions(-) create mode 100644 src/zenml/zen_stores/secrets_stores/service_connector_secrets_store.py diff --git a/src/zenml/integrations/aws/__init__.py b/src/zenml/integrations/aws/__init__.py index 311bb68e639..87525add512 100644 --- a/src/zenml/integrations/aws/__init__.py +++ b/src/zenml/integrations/aws/__init__.py @@ -29,6 +29,10 @@ AWS_SAGEMAKER_STEP_OPERATOR_FLAVOR = "sagemaker" AWS_SAGEMAKER_ORCHESTRATOR_FLAVOR = "sagemaker" +# Service connector constants +AWS_CONNECTOR_TYPE = "aws" +AWS_RESOURCE_TYPE = "aws-generic" +S3_RESOURCE_TYPE = "s3-bucket" class AWSIntegration(Integration): """Definition of AWS integration for ZenML.""" diff --git a/src/zenml/integrations/aws/flavors/aws_container_registry_flavor.py b/src/zenml/integrations/aws/flavors/aws_container_registry_flavor.py index 3bc914a7e1c..117a8d3c21e 100644 --- a/src/zenml/integrations/aws/flavors/aws_container_registry_flavor.py +++ b/src/zenml/integrations/aws/flavors/aws_container_registry_flavor.py @@ -17,11 +17,15 @@ from pydantic import validator +from zenml.constants import DOCKER_REGISTRY_RESOURCE_TYPE from zenml.container_registries.base_container_registry import ( BaseContainerRegistryConfig, BaseContainerRegistryFlavor, ) -from zenml.integrations.aws import AWS_CONTAINER_REGISTRY_FLAVOR +from zenml.integrations.aws import ( + AWS_CONNECTOR_TYPE, + AWS_CONTAINER_REGISTRY_FLAVOR, +) from zenml.models import ServiceConnectorRequirements if TYPE_CHECKING: @@ -81,8 +85,8 @@ def service_connector_requirements( connector is required for this flavor. """ return ServiceConnectorRequirements( - connector_type="aws", - resource_type="docker-registry", + connector_type=AWS_CONNECTOR_TYPE, + resource_type=DOCKER_REGISTRY_RESOURCE_TYPE, resource_id_attr="uri", ) diff --git a/src/zenml/integrations/aws/flavors/sagemaker_orchestrator_flavor.py b/src/zenml/integrations/aws/flavors/sagemaker_orchestrator_flavor.py index c6ad36647fe..d35f3831ee0 100644 --- a/src/zenml/integrations/aws/flavors/sagemaker_orchestrator_flavor.py +++ b/src/zenml/integrations/aws/flavors/sagemaker_orchestrator_flavor.py @@ -16,7 +16,10 @@ from typing import TYPE_CHECKING, Any, Dict, Optional, Type, Union from zenml.config.base_settings import BaseSettings -from zenml.integrations.aws import AWS_SAGEMAKER_STEP_OPERATOR_FLAVOR +from zenml.integrations.aws import ( + AWS_RESOURCE_TYPE, + AWS_SAGEMAKER_STEP_OPERATOR_FLAVOR, +) from zenml.models import ServiceConnectorRequirements from zenml.orchestrators import BaseOrchestratorConfig from zenml.orchestrators.base_orchestrator import BaseOrchestratorFlavor @@ -167,7 +170,7 @@ def service_connector_requirements( Requirements for compatible service connectors, if a service connector is required for this flavor. """ - return ServiceConnectorRequirements(resource_type="aws-generic") + return ServiceConnectorRequirements(resource_type=AWS_RESOURCE_TYPE) @property def docs_url(self) -> Optional[str]: diff --git a/src/zenml/integrations/aws/flavors/sagemaker_step_operator_flavor.py b/src/zenml/integrations/aws/flavors/sagemaker_step_operator_flavor.py index 9d1a0e22351..c063892ccc5 100644 --- a/src/zenml/integrations/aws/flavors/sagemaker_step_operator_flavor.py +++ b/src/zenml/integrations/aws/flavors/sagemaker_step_operator_flavor.py @@ -16,9 +16,9 @@ from typing import TYPE_CHECKING, Any, Dict, Optional, Type, Union from zenml.config.base_settings import BaseSettings -from zenml.integrations.aws import AWS_SAGEMAKER_STEP_OPERATOR_FLAVOR -from zenml.integrations.aws.service_connectors.aws_service_connector import ( +from zenml.integrations.aws import ( AWS_RESOURCE_TYPE, + AWS_SAGEMAKER_STEP_OPERATOR_FLAVOR, ) from zenml.models import ServiceConnectorRequirements from zenml.step_operators.base_step_operator import ( diff --git a/src/zenml/integrations/aws/service_connectors/aws_service_connector.py b/src/zenml/integrations/aws/service_connectors/aws_service_connector.py index ada229acc69..fd17fb3ec00 100644 --- a/src/zenml/integrations/aws/service_connectors/aws_service_connector.py +++ b/src/zenml/integrations/aws/service_connectors/aws_service_connector.py @@ -42,6 +42,11 @@ KUBERNETES_CLUSTER_RESOURCE_TYPE, ) from zenml.exceptions import AuthorizationException +from zenml.integrations.aws import ( + AWS_CONNECTOR_TYPE, + AWS_RESOURCE_TYPE, + S3_RESOURCE_TYPE, +) from zenml.logger import get_logger from zenml.models import ( AuthenticationMethodModel, @@ -61,10 +66,6 @@ logger = get_logger(__name__) - -AWS_CONNECTOR_TYPE = "aws" -AWS_RESOURCE_TYPE = "aws-generic" -S3_RESOURCE_TYPE = "s3-bucket" EKS_KUBE_API_TOKEN_EXPIRATION = 60 DEFAULT_IAM_ROLE_TOKEN_EXPIRATION = 3600 # 1 hour DEFAULT_STS_TOKEN_EXPIRATION = 43200 # 12 hours diff --git a/src/zenml/integrations/azure/__init__.py b/src/zenml/integrations/azure/__init__.py index 6f5d07df58d..f7292963e0e 100644 --- a/src/zenml/integrations/azure/__init__.py +++ b/src/zenml/integrations/azure/__init__.py @@ -29,6 +29,11 @@ AZURE_SECRETS_MANAGER_FLAVOR = "azure" AZUREML_STEP_OPERATOR_FLAVOR = "azureml" +# Service connector constants +AZURE_CONNECTOR_TYPE = "azure" +AZURE_RESOURCE_TYPE = "azure-generic" +BLOB_RESOURCE_TYPE = "blob-container" + class AzureIntegration(Integration): """Definition of Azure integration for ZenML.""" diff --git a/src/zenml/integrations/azure/flavors/azure_artifact_store_flavor.py b/src/zenml/integrations/azure/flavors/azure_artifact_store_flavor.py index b86d813e4a4..4bf209f1441 100644 --- a/src/zenml/integrations/azure/flavors/azure_artifact_store_flavor.py +++ b/src/zenml/integrations/azure/flavors/azure_artifact_store_flavor.py @@ -19,7 +19,11 @@ BaseArtifactStoreConfig, BaseArtifactStoreFlavor, ) -from zenml.integrations.azure import AZURE_ARTIFACT_STORE_FLAVOR +from zenml.integrations.azure import ( + AZURE_ARTIFACT_STORE_FLAVOR, + AZURE_CONNECTOR_TYPE, + BLOB_RESOURCE_TYPE, +) from zenml.models import ServiceConnectorRequirements from zenml.stack.authentication_mixin import AuthenticationConfigMixin @@ -27,11 +31,6 @@ from zenml.integrations.azure.artifact_stores import AzureArtifactStore -AZURE_CONNECTOR_TYPE = "azure" -AZURE_RESOURCE_TYPE = "azure-generic" -BLOB_RESOURCE_TYPE = "blob-container" - - class AzureArtifactStoreConfig( BaseArtifactStoreConfig, AuthenticationConfigMixin ): diff --git a/src/zenml/integrations/azure/service_connectors/azure_service_connector.py b/src/zenml/integrations/azure/service_connectors/azure_service_connector.py index 377c7d0bbc5..4d0d10bb9b1 100644 --- a/src/zenml/integrations/azure/service_connectors/azure_service_connector.py +++ b/src/zenml/integrations/azure/service_connectors/azure_service_connector.py @@ -34,7 +34,7 @@ KUBERNETES_CLUSTER_RESOURCE_TYPE, ) from zenml.exceptions import AuthorizationException -from zenml.integrations.azure.flavors.azure_artifact_store_flavor import ( +from zenml.integrations.azure import ( AZURE_CONNECTOR_TYPE, AZURE_RESOURCE_TYPE, BLOB_RESOURCE_TYPE, diff --git a/src/zenml/integrations/gcp/__init__.py b/src/zenml/integrations/gcp/__init__.py index 3067cd41620..474ff313bde 100644 --- a/src/zenml/integrations/gcp/__init__.py +++ b/src/zenml/integrations/gcp/__init__.py @@ -37,6 +37,10 @@ GCP_VERTEX_ORCHESTRATOR_FLAVOR = "vertex" GCP_VERTEX_STEP_OPERATOR_FLAVOR = "vertex" +# Service connector constants +GCP_CONNECTOR_TYPE = "gcp" +GCP_RESOURCE_TYPE = "gcp-generic" +GCS_RESOURCE_TYPE = "gcs-bucket" class GcpIntegration(Integration): """Definition of Google Cloud Platform integration for ZenML.""" diff --git a/src/zenml/integrations/gcp/flavors/gcp_artifact_store_flavor.py b/src/zenml/integrations/gcp/flavors/gcp_artifact_store_flavor.py index cfd14223340..1721193c263 100644 --- a/src/zenml/integrations/gcp/flavors/gcp_artifact_store_flavor.py +++ b/src/zenml/integrations/gcp/flavors/gcp_artifact_store_flavor.py @@ -19,7 +19,7 @@ BaseArtifactStoreConfig, BaseArtifactStoreFlavor, ) -from zenml.integrations.gcp import GCP_ARTIFACT_STORE_FLAVOR +from zenml.integrations.gcp import GCP_ARTIFACT_STORE_FLAVOR, GCS_RESOURCE_TYPE from zenml.models import ServiceConnectorRequirements from zenml.stack.authentication_mixin import AuthenticationConfigMixin @@ -64,7 +64,7 @@ def service_connector_requirements( connector is required for this flavor. """ return ServiceConnectorRequirements( - resource_type="gcs-bucket", + resource_type=GCS_RESOURCE_TYPE, resource_id_attr="path", ) diff --git a/src/zenml/integrations/gcp/flavors/gcp_image_builder_flavor.py b/src/zenml/integrations/gcp/flavors/gcp_image_builder_flavor.py index 0af13db002f..3614ced37da 100644 --- a/src/zenml/integrations/gcp/flavors/gcp_image_builder_flavor.py +++ b/src/zenml/integrations/gcp/flavors/gcp_image_builder_flavor.py @@ -18,7 +18,11 @@ from pydantic import PositiveInt from zenml.image_builders import BaseImageBuilderConfig, BaseImageBuilderFlavor -from zenml.integrations.gcp import GCP_IMAGE_BUILDER_FLAVOR +from zenml.integrations.gcp import ( + GCP_CONNECTOR_TYPE, + GCP_IMAGE_BUILDER_FLAVOR, + GCP_RESOURCE_TYPE, +) from zenml.integrations.gcp.google_credentials_mixin import ( GoogleCredentialsConfigMixin, ) @@ -82,8 +86,8 @@ def service_connector_requirements( connector is required for this flavor. """ return ServiceConnectorRequirements( - connector_type="gcp", - resource_type="gcp-generic", + connector_type=GCP_CONNECTOR_TYPE, + resource_type=GCP_RESOURCE_TYPE, ) @property diff --git a/src/zenml/integrations/gcp/flavors/vertex_orchestrator_flavor.py b/src/zenml/integrations/gcp/flavors/vertex_orchestrator_flavor.py index 798f821e4d1..17236d81c87 100644 --- a/src/zenml/integrations/gcp/flavors/vertex_orchestrator_flavor.py +++ b/src/zenml/integrations/gcp/flavors/vertex_orchestrator_flavor.py @@ -16,7 +16,10 @@ from typing import TYPE_CHECKING, Dict, Optional, Tuple, Type from zenml.config.base_settings import BaseSettings -from zenml.integrations.gcp import GCP_VERTEX_ORCHESTRATOR_FLAVOR +from zenml.integrations.gcp import ( + GCP_RESOURCE_TYPE, + GCP_VERTEX_ORCHESTRATOR_FLAVOR, +) from zenml.integrations.gcp.google_credentials_mixin import ( GoogleCredentialsConfigMixin, ) @@ -173,7 +176,7 @@ def service_connector_requirements( connector is required for this flavor. """ return ServiceConnectorRequirements( - resource_type="gcp-generic", + resource_type=GCP_RESOURCE_TYPE, ) @property diff --git a/src/zenml/integrations/gcp/flavors/vertex_step_operator_flavor.py b/src/zenml/integrations/gcp/flavors/vertex_step_operator_flavor.py index df866155e3d..eb2a12f1d17 100644 --- a/src/zenml/integrations/gcp/flavors/vertex_step_operator_flavor.py +++ b/src/zenml/integrations/gcp/flavors/vertex_step_operator_flavor.py @@ -16,7 +16,10 @@ from typing import TYPE_CHECKING, Optional, Type from zenml.config.base_settings import BaseSettings -from zenml.integrations.gcp import GCP_VERTEX_STEP_OPERATOR_FLAVOR +from zenml.integrations.gcp import ( + GCP_RESOURCE_TYPE, + GCP_VERTEX_STEP_OPERATOR_FLAVOR, +) from zenml.integrations.gcp.google_credentials_mixin import ( GoogleCredentialsConfigMixin, ) @@ -116,7 +119,7 @@ def service_connector_requirements( connector is required for this flavor. """ return ServiceConnectorRequirements( - resource_type="gcp-generic", + resource_type=GCP_RESOURCE_TYPE, ) @property diff --git a/src/zenml/integrations/gcp/service_connectors/gcp_service_connector.py b/src/zenml/integrations/gcp/service_connectors/gcp_service_connector.py index 6211482a6cb..917d137d6a1 100644 --- a/src/zenml/integrations/gcp/service_connectors/gcp_service_connector.py +++ b/src/zenml/integrations/gcp/service_connectors/gcp_service_connector.py @@ -43,6 +43,11 @@ KUBERNETES_CLUSTER_RESOURCE_TYPE, ) from zenml.exceptions import AuthorizationException +from zenml.integrations.gcp import ( + GCP_CONNECTOR_TYPE, + GCP_RESOURCE_TYPE, + GCS_RESOURCE_TYPE, +) from zenml.logger import get_logger from zenml.models import ( AuthenticationMethodModel, @@ -62,10 +67,6 @@ logger = get_logger(__name__) - -GCP_CONNECTOR_TYPE = "gcp" -GCP_RESOURCE_TYPE = "gcp-generic" -GCS_RESOURCE_TYPE = "gcs-bucket" GKE_KUBE_API_TOKEN_EXPIRATION = 60 DEFAULT_IMPERSONATE_TOKEN_EXPIRATION = 3600 # 1 hour diff --git a/src/zenml/zen_stores/secrets_stores/aws_secrets_store.py b/src/zenml/zen_stores/secrets_stores/aws_secrets_store.py index a26bf4072c8..95b07ba6235 100644 --- a/src/zenml/zen_stores/secrets_stores/aws_secrets_store.py +++ b/src/zenml/zen_stores/secrets_stores/aws_secrets_store.py @@ -26,17 +26,15 @@ List, Optional, Type, - cast, ) -from uuid import UUID, uuid4 +from uuid import UUID import boto3 from botocore.exceptions import ClientError -from pydantic import BaseModel, Field, root_validator +from pydantic import root_validator from zenml.analytics.enums import AnalyticsEvent from zenml.analytics.utils import track_decorator -from zenml.config.secrets_store_config import SecretsStoreConfiguration from zenml.enums import ( GenericFilterOps, LogicalOperators, @@ -44,11 +42,12 @@ SecretsStoreType, ) from zenml.exceptions import EntityExistsError -from zenml.integrations.aws.service_connectors.aws_service_connector import ( +from zenml.integrations.aws import ( AWS_CONNECTOR_TYPE, AWS_RESOURCE_TYPE, +) +from zenml.integrations.aws.service_connectors.aws_service_connector import ( AWSAuthenticationMethods, - AWSServiceConnector, ) from zenml.logger import get_logger from zenml.models import ( @@ -57,13 +56,10 @@ SecretRequestModel, SecretResponseModel, SecretUpdateModel, - ServiceConnectorRequest, ) -from zenml.service_connectors.service_connector_registry import ( - service_connector_registry, -) -from zenml.zen_stores.secrets_stores.base_secrets_store import ( - BaseSecretsStore, +from zenml.zen_stores.secrets_stores.service_connector_secrets_store import ( + ServiceConnectorSecretsStore, + ServiceConnectorSecretsStoreConfiguration, ) logger = get_logger(__name__) @@ -72,20 +68,11 @@ AWS_ZENML_SECRET_NAME_PREFIX = "zenml" -class AWSSecretsStoreConnector(BaseModel): - """AWS secrets store connector configuration.""" - - -class AWSSecretsStoreConfiguration(SecretsStoreConfiguration): +class AWSSecretsStoreConfiguration(ServiceConnectorSecretsStoreConfiguration): """AWS secrets store configuration. Attributes: type: The type of the store. - region_name: The AWS region name to use. - aws_access_key_id: The AWS access key ID to use to authenticate. - aws_secret_access_key: The AWS secret access key to use to - authenticate. - aws_session_token: The AWS session token to use to authenticate. list_page_size: The number of secrets to fetch per page when listing secrets. secret_list_refresh_timeout: The number of seconds to wait after @@ -101,13 +88,27 @@ class AWSSecretsStoreConfiguration(SecretsStoreConfiguration): """ type: SecretsStoreType = SecretsStoreType.AWS - - auth_method: AWSAuthenticationMethods = AWSAuthenticationMethods.SECRET_KEY - auth_config: Dict[str, Any] = Field(default_factory=dict) - list_page_size: int = 100 secret_list_refresh_timeout: int = 0 + @property + def region(self) -> str: + """The AWS region to use. + + Returns: + The AWS region to use. + + Raises: + ValueError: If the region is not configured. + """ + region = self.auth_config.get("region") + if region: + return str(region) + + raise ValueError( + "AWS `region` must be specified in the auth_config." + ) + @root_validator(pre=True) def populate_config(cls, values: Dict[str, Any]) -> Dict[str, Any]: """Populate the connector configuration from legacy attributes. @@ -122,6 +123,7 @@ def populate_config(cls, values: Dict[str, Any]) -> Dict[str, Any]: ValueError: If the connector attribute is not set. """ if "auth_method" not in values or "auth_config" not in values: + values["auth_method"] = AWSAuthenticationMethods.SECRET_KEY values["auth_config"] = dict( aws_access_key_id=values.get("aws_access_key_id"), aws_secret_access_key=values.get("aws_secret_access_key"), @@ -137,7 +139,7 @@ class Config: extra = "allow" -class AWSSecretsStore(BaseSecretsStore): +class AWSSecretsStore(ServiceConnectorSecretsStore): """Secrets store implementation that uses the AWS Secrets Manager API. This secrets store implementation uses the AWS Secrets Manager API to @@ -181,69 +183,15 @@ class AWSSecretsStore(BaseSecretsStore): workspace) does not update the secret's `updated` timestamp. This is a limitation of the AWS Secrets Manager API (updating AWS tags does not update the secret's `updated` timestamp). - - - Attributes: - config: The configuration of the AWS secrets store. - TYPE: The type of the store. - CONFIG_TYPE: The type of the store configuration. """ config: AWSSecretsStoreConfiguration TYPE: ClassVar[SecretsStoreType] = SecretsStoreType.AWS CONFIG_TYPE: ClassVar[ - Type[SecretsStoreConfiguration] + Type[ServiceConnectorSecretsStoreConfiguration] ] = AWSSecretsStoreConfiguration - - _client: Optional[Any] = None - _connector: Optional[AWSServiceConnector] = None - - @property - def client(self) -> Any: - """Initialize and return the AWS Secrets Manager client. - - Returns: - The AWS Secrets Manager client. - """ - if self._connector is not None: - # If the client connector expires, we'll try to get a new one. - if self._connector.has_expired(): - self._connector = None - self._client = None - - if self._connector is None: - # Initialize a base AWS service connector with the credentials from - # the configuration. - request = ServiceConnectorRequest( - name="secrets-store", - connector_type=AWS_CONNECTOR_TYPE, - resource_types=[AWS_RESOURCE_TYPE], - user=uuid4(), - workspace=uuid4(), - auth_method=self.config.auth_method, - configuration=self.config.auth_config, - ) - base_connector = service_connector_registry.instantiate_connector( - model=request - ) - self._connector = cast( - AWSServiceConnector, base_connector.get_connector_client() - ) - - if self._client is None: - # Initialize the AWS Secrets Manager client with the - # credentials from the connector. - session = self._connector.connect( - # Don't verify again because we already did that when we - # initialized the connector. - verify=False - ) - assert isinstance(session, boto3.Session) - self._client = session.client( - "secretsmanager", - region_name=self._connector.config.region, - ) - return self._client + SERVICE_CONNECTOR_TYPE: ClassVar[str] = AWS_CONNECTOR_TYPE + SERVICE_CONNECTOR_RESOURCE_TYPE: ClassVar[str] = AWS_RESOURCE_TYPE # ==================================== # Secrets Store interface implementation @@ -253,13 +201,21 @@ def client(self) -> Any: # Initialization and configuration # -------------------------------- - def _initialize(self) -> None: - """Initialize the AWS secrets store.""" - logger.debug("Initializing AWSSecretsStore") + def _initialize_client_from_connector(self, client: Any) -> Any: + """Initialize the GCP Secrets Manager client from the service connector client. - # Initialize the AWS client early, just to catch any configuration or - # authentication errors early, before the Secrets Store is used. - _ = self.client + Args: + client: The authenticated client object returned by the service + connector. + + Returns: + The GCP Secrets Manager client. + """ + assert isinstance(client, boto3.Session) + return client.client( + "secretsmanager", + region_name=self.config.region, + ) # ------ # Secrets diff --git a/src/zenml/zen_stores/secrets_stores/azure_secrets_store.py b/src/zenml/zen_stores/secrets_stores/azure_secrets_store.py index 1fd6baddfcb..eb379be7cb7 100644 --- a/src/zenml/zen_stores/secrets_stores/azure_secrets_store.py +++ b/src/zenml/zen_stores/secrets_stores/azure_secrets_store.py @@ -20,30 +20,34 @@ import uuid from datetime import datetime from typing import ( + Any, ClassVar, Dict, List, Optional, Type, - Union, + cast, ) from uuid import UUID +from azure.core.credentials import TokenCredential from azure.core.exceptions import HttpResponseError, ResourceNotFoundError -from azure.identity import ( - ClientSecretCredential, - DefaultAzureCredential, -) from azure.keyvault.secrets import SecretClient -from pydantic import SecretStr +from pydantic import root_validator from zenml.analytics.enums import AnalyticsEvent from zenml.analytics.utils import track_decorator -from zenml.config.secrets_store_config import SecretsStoreConfiguration from zenml.enums import ( SecretsStoreType, ) from zenml.exceptions import EntityExistsError +from zenml.integrations.azure import ( + AZURE_CONNECTOR_TYPE, + AZURE_RESOURCE_TYPE, +) +from zenml.integrations.azure.service_connectors.azure_service_connector import ( + AzureAuthenticationMethods, +) from zenml.logger import get_logger from zenml.models import ( Page, @@ -52,8 +56,9 @@ SecretResponseModel, SecretUpdateModel, ) -from zenml.zen_stores.secrets_stores.base_secrets_store import ( - BaseSecretsStore, +from zenml.zen_stores.secrets_stores.service_connector_secrets_store import ( + ServiceConnectorSecretsStore, + ServiceConnectorSecretsStoreConfiguration, ) logger = get_logger(__name__) @@ -64,39 +69,53 @@ ZENML_AZURE_SECRET_UPDATED_KEY = "zenml_secret_updated" -class AzureSecretsStoreConfiguration(SecretsStoreConfiguration): +class AzureSecretsStoreConfiguration( + ServiceConnectorSecretsStoreConfiguration +): """Azure secrets store configuration. Attributes: type: The type of the store. key_vault_name: Name of the Azure Key Vault that this secrets store will use to store secrets. - azure_client_id: The client ID of the Azure application service - principal that will be used to access the Azure Key Vault. If not - set, the default Azure credential chain will be used. - azure_client_secret: The client secret of the Azure application - service principal that will be used to access the Azure Key Vault. - If not set, the default Azure credential chain will be used. - azure_tenant_id: The tenant ID of the Azure application service - principal that will be used to access the Azure Key Vault. If not - set, the default Azure credential chain will be used. """ type: SecretsStoreType = SecretsStoreType.AZURE - key_vault_name: str - azure_client_id: Optional[SecretStr] = None - azure_client_secret: Optional[SecretStr] = None - azure_tenant_id: Optional[SecretStr] = None + + @root_validator(pre=True) + def populate_config(cls, values: Dict[str, Any]) -> Dict[str, Any]: + """Populate the connector configuration from legacy attributes. + + Args: + values: Dict representing user-specified runtime settings. + + Returns: + Validated settings. + + Raises: + ValueError: If the connector attribute is not set. + """ + if "auth_method" not in values or "auth_config" not in values: + values[ + "auth_method" + ] = AzureAuthenticationMethods.SERVICE_PRINCIPAL + values["auth_config"] = dict( + client_id=values.get("azure_client_id"), + client_secret=values.get("azure_client_secret"), + tenant_id=values.get("azure_tenant_id"), + ) + + return values class Config: """Pydantic configuration class.""" # Forbid extra attributes set in the class. - extra = "forbid" + extra = "allow" -class AzureSecretsStore(BaseSecretsStore): +class AzureSecretsStore(ServiceConnectorSecretsStore): """Secrets store implementation that uses the Azure Key Vault API. This secrets store implementation uses the Azure Key Vault API to @@ -132,20 +151,15 @@ class AzureSecretsStore(BaseSecretsStore): fetch all versions for every secret to get the created_on timestamp during a list operation. So instead we manage the `created` and `updated` timestamps ourselves and save them as tags in the Azure Key Vault secret. - - Attributes: - config: The configuration of the Azure secrets store. - TYPE: The type of the store. - CONFIG_TYPE: The type of the store configuration. """ config: AzureSecretsStoreConfiguration TYPE: ClassVar[SecretsStoreType] = SecretsStoreType.AZURE CONFIG_TYPE: ClassVar[ - Type[SecretsStoreConfiguration] + Type[ServiceConnectorSecretsStoreConfiguration] ] = AzureSecretsStoreConfiguration - - _client: Optional[SecretClient] = None + SERVICE_CONNECTOR_TYPE: ClassVar[str] = AZURE_CONNECTOR_TYPE + SERVICE_CONNECTOR_RESOURCE_TYPE: ClassVar[str] = AZURE_RESOURCE_TYPE @property def client(self) -> SecretClient: @@ -154,39 +168,7 @@ def client(self) -> SecretClient: Returns: The Azure Key Vault client. """ - if self._client is None: - azure_logger = logging.getLogger("azure") - - # Suppress the INFO logging level of the Azure SDK if the - # ZenML logging level is WARNING or lower. - if logger.level <= logging.WARNING: - azure_logger.setLevel(logging.WARNING) - else: - azure_logger.setLevel(logging.INFO) - - # Initialize the Azure Key Vault client with the - # credentials from the configuration, if provided. - vault_url = f"https://{self.config.key_vault_name}.vault.azure.net" - credential: Union[ - ClientSecretCredential, - DefaultAzureCredential, - ] - if ( - self.config.azure_client_id - and self.config.azure_tenant_id - and self.config.azure_client_secret - ): - credential = ClientSecretCredential( - tenant_id=self.config.azure_tenant_id.get_secret_value(), - client_id=self.config.azure_client_id.get_secret_value(), - client_secret=self.config.azure_client_secret.get_secret_value(), - ) - else: - credential = DefaultAzureCredential() - self._client = SecretClient( - vault_url=vault_url, credential=credential - ) - return self._client + return cast(SecretClient, super().client) # ==================================== # Secrets Store interface implementation @@ -196,13 +178,28 @@ def client(self) -> SecretClient: # Initialization and configuration # -------------------------------- - def _initialize(self) -> None: - """Initialize the Azure secrets store.""" - logger.debug("Initializing AzureSecretsStore") + def _initialize_client_from_connector(self, client: Any) -> Any: + """Initialize the Azure Key Vault client from the service connector client. + + Args: + client: The authenticated client object returned by the service + connector. + + Returns: + The Azure Key Vault client. + """ + assert isinstance(client, TokenCredential) + azure_logger = logging.getLogger("azure") + + # Suppress the INFO logging level of the Azure SDK if the + # ZenML logging level is WARNING or lower. + if logger.level <= logging.WARNING: + azure_logger.setLevel(logging.WARNING) + else: + azure_logger.setLevel(logging.INFO) - # Initialize the Azure client early, just to catch any configuration or - # authentication errors early, before the Secrets Store is used. - _ = self.client + vault_url = f"https://{self.config.key_vault_name}.vault.azure.net" + return SecretClient(vault_url=vault_url, credential=client) # ------ # Secrets diff --git a/src/zenml/zen_stores/secrets_stores/gcp_secrets_store.py b/src/zenml/zen_stores/secrets_stores/gcp_secrets_store.py index 632e82b9974..22559e08020 100644 --- a/src/zenml/zen_stores/secrets_stores/gcp_secrets_store.py +++ b/src/zenml/zen_stores/secrets_stores/gcp_secrets_store.py @@ -16,6 +16,7 @@ import json import math +import os import re import uuid from datetime import datetime @@ -28,19 +29,28 @@ Tuple, Type, Union, + cast, ) from uuid import UUID from google.api_core import exceptions as google_exceptions from google.cloud.secretmanager import SecretManagerServiceClient +from google.oauth2 import service_account as gcp_service_account +from pydantic import root_validator from zenml.analytics.enums import AnalyticsEvent from zenml.analytics.utils import track_decorator -from zenml.config.secrets_store_config import SecretsStoreConfiguration from zenml.enums import ( SecretsStoreType, ) from zenml.exceptions import EntityExistsError +from zenml.integrations.gcp import ( + GCP_CONNECTOR_TYPE, + GCP_RESOURCE_TYPE, +) +from zenml.integrations.gcp.service_connectors.gcp_service_connector import ( + GCPAuthenticationMethods, +) from zenml.logger import get_logger from zenml.models import ( Page, @@ -49,8 +59,9 @@ SecretResponseModel, SecretUpdateModel, ) -from zenml.zen_stores.secrets_stores.base_secrets_store import ( - BaseSecretsStore, +from zenml.zen_stores.secrets_stores.service_connector_secrets_store import ( + ServiceConnectorSecretsStore, + ServiceConnectorSecretsStoreConfiguration, ) logger = get_logger(__name__) @@ -65,53 +76,111 @@ ZENML_GCP_SECRET_UPDATED_KEY = "zenml-secret-updated" -class GCPSecretsStoreConfiguration(SecretsStoreConfiguration): +class GCPSecretsStoreConfiguration(ServiceConnectorSecretsStoreConfiguration): """GCP secrets store configuration. Attributes: type: The type of the store. - project_id: The GCP project ID where the secrets are stored. - """ type: SecretsStoreType = SecretsStoreType.GCP - project_id: str + + @property + def project_id(self) -> str: + """Get the GCP project ID. + + Returns: + The GCP project ID. + + Raises: + ValueError: If the project ID is not set. + """ + project_id = self.auth_config.get("project_id") + if project_id: + return str(project_id) + + raise ValueError("GCP `project_id` must be specified in auth_config.") + + @root_validator(pre=True) + def populate_config(cls, values: Dict[str, Any]) -> Dict[str, Any]: + """Populate the connector configuration from legacy attributes. + + Args: + values: Dict representing user-specified runtime settings. + + Returns: + Validated settings. + + Raises: + ValueError: If the connector attribute is not set. + """ + if "auth_method" not in values or "auth_config" not in values: + values["auth_method"] = GCPAuthenticationMethods.SERVICE_ACCOUNT + values["auth_config"] = dict( + project_id=values.get("project_id"), + ) + if os.environ.get("GOOGLE_APPLICATION_CREDENTIALS"): + # Load the service account credentials from the file + with open(os.environ["GOOGLE_APPLICATION_CREDENTIALS"]) as f: + values["auth_config"]["service_account_json"] = f.read() + + return values class Config: """Pydantic configuration class.""" # Forbid extra attributes set in the class. - extra = "forbid" + extra = "allow" -class GCPSecretsStore(BaseSecretsStore): +class GCPSecretsStore(ServiceConnectorSecretsStore): """Secrets store implementation that uses the GCP Secrets Manager API.""" config: GCPSecretsStoreConfiguration TYPE: ClassVar[SecretsStoreType] = SecretsStoreType.GCP CONFIG_TYPE: ClassVar[ - Type[SecretsStoreConfiguration] + Type[ServiceConnectorSecretsStoreConfiguration] ] = GCPSecretsStoreConfiguration + SERVICE_CONNECTOR_TYPE: ClassVar[str] = GCP_CONNECTOR_TYPE + SERVICE_CONNECTOR_RESOURCE_TYPE: ClassVar[str] = GCP_RESOURCE_TYPE _client: Optional[SecretManagerServiceClient] = None - def _initialize(self) -> None: - """Initialize the GCP secrets store.""" - logger.debug("Initializing GCPSecretsStore") - - # Initialize the GCP client. - _ = self.client - @property - def client(self) -> Any: + def client(self) -> SecretManagerServiceClient: """Initialize and return the GCP Secrets Manager client. + Returns: + The GCP Secrets Manager client instance. + """ + return cast(SecretManagerServiceClient, super().client) + + # ==================================== + # Secrets Store interface implementation + # ==================================== + + # -------------------------------- + # Initialization and configuration + # -------------------------------- + + def _initialize_client_from_connector(self, client: Any) -> Any: + """Initialize the GCP Secrets Manager client from the service connector client. + + Args: + client: The authenticated client object returned by the service + connector. + Returns: The GCP Secrets Manager client. """ - if self._client is None: - self._client = SecretManagerServiceClient() - return self._client + assert isinstance(client, gcp_service_account.Credentials) + return SecretManagerServiceClient( + project=self.config.project_id, credentials=client + ) + + # ------ + # Secrets + # ------ @property def parent_name(self) -> str: diff --git a/src/zenml/zen_stores/secrets_stores/service_connector_secrets_store.py b/src/zenml/zen_stores/secrets_stores/service_connector_secrets_store.py new file mode 100644 index 00000000000..14134bcc292 --- /dev/null +++ b/src/zenml/zen_stores/secrets_stores/service_connector_secrets_store.py @@ -0,0 +1,153 @@ +# Copyright (c) ZenML GmbH 2023. 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. +"""Base secrets store class used for all secrets stores that use a service connector.""" + + +from abc import abstractmethod +from threading import Lock +from typing import ( + Any, + ClassVar, + Dict, + Optional, + Type, +) +from uuid import uuid4 + +from pydantic import Field + +from zenml.config.secrets_store_config import SecretsStoreConfiguration +from zenml.logger import get_logger +from zenml.models import ( + ServiceConnectorRequest, +) +from zenml.service_connectors.service_connector import ServiceConnector +from zenml.service_connectors.service_connector_registry import ( + service_connector_registry, +) +from zenml.zen_stores.secrets_stores.base_secrets_store import ( + BaseSecretsStore, +) + +logger = get_logger(__name__) + + +class ServiceConnectorSecretsStoreConfiguration(SecretsStoreConfiguration): + """Base configuration for secrets stores that use a service connector. + + Attributes: + auth_method: The service connector authentication method to use. + auth_config: The service connector authentication configuration. + """ + + auth_method: str + auth_config: Dict[str, Any] = Field(default_factory=dict) + + +class ServiceConnectorSecretsStore(BaseSecretsStore): + """Base secrets store class for service connector-based secrets stores. + + All secrets store implementations that use a Service Connector to + authenticate and connect to the secrets store back-end should inherit from + this class and: + + * implement the `_initialize_client_from_connector` method + * use a configuration class that inherits from + `ServiceConnectorSecretsStoreConfiguration` + * set the `SERVICE_CONNECTOR_TYPE` to the service connector type used + to connect to the secrets store back-end + * set the `SERVICE_CONNECTOR_RESOURCE_TYPE` to the resource type used + to connect to the secrets store back-end + """ + + config: ServiceConnectorSecretsStoreConfiguration + CONFIG_TYPE: ClassVar[Type[ServiceConnectorSecretsStoreConfiguration]] + SERVICE_CONNECTOR_TYPE: ClassVar[str] + SERVICE_CONNECTOR_RESOURCE_TYPE: ClassVar[str] + + _connector: Optional[ServiceConnector] = None + _client: Optional[Any] = None + _lock: Lock = Lock() + + def _initialize(self) -> None: + """Initialize the secrets store.""" + # Initialize the client early, just to catch any configuration or + # authentication errors early, before the Secrets Store is used. + _ = self.client + + def _get_client(self) -> Any: + """Initialize and return the secrets store API client. + + Returns: + The secrets store API client object. + """ + if self._connector is not None: + # If the client connector expires, we'll try to get a new one. + if self._connector.has_expired(): + self._connector = None + self._client = None + + if self._connector is None: + # Initialize a base AWS service connector with the credentials from + # the configuration. + request = ServiceConnectorRequest( + name="secrets-store", + connector_type=self.SERVICE_CONNECTOR_TYPE, + resource_types=[self.SERVICE_CONNECTOR_RESOURCE_TYPE], + user=uuid4(), # Use a fake user ID + workspace=uuid4(), # Use a fake workspace ID + auth_method=self.config.auth_method, + configuration=self.config.auth_config, + ) + base_connector = service_connector_registry.instantiate_connector( + model=request + ) + self._connector = base_connector.get_connector_client() + + if self._client is None: + # Use the connector to get a client object. + client = self._connector.connect( + # Don't verify again because we already did that when we + # initialized the connector. + verify=False + ) + + self._client = self._initialize_client_from_connector(client) + return self._client + + @property + def client(self) -> Any: + """Get the secrets store API client. + + Returns: + The secrets store API client instance. + """ + # Multiple API calls can be made to this secrets store in the context + # of different threads. We want to make sure that we only initialize + # the client once, and then reuse it. We have to use a lock to treat + # this method as a critical section. + with self._lock: + return self._get_client() + + @abstractmethod + def _initialize_client_from_connector(self, client: Any) -> Any: + """Initialize the client from the service connector. + + Args: + client: The authenticated client object returned by the service + connector. + + Returns: + The initialized client instance. + """ From 19e84e8bc85a7bacdd5363bfecac6a9004ecfc67 Mon Sep 17 00:00:00 2001 From: Stefan Nica Date: Fri, 15 Dec 2023 19:43:23 +0100 Subject: [PATCH 3/7] Add helm chart support and update docs --- .../zenml-self-hosted/deploy-with-docker.md | 37 +++- .../zenml-self-hosted/deploy-with-helm.md | 79 ++++---- .../deploy/helm/templates/server-db-job.yaml | 10 ++ .../helm/templates/server-deployment.yaml | 10 ++ .../deploy/helm/templates/server-secret.yaml | 9 + src/zenml/zen_server/deploy/helm/values.yaml | 168 +++++++++++++++++- .../secrets_stores/aws_secrets_store.py | 18 +- .../secrets_stores/azure_secrets_store.py | 14 +- .../secrets_stores/gcp_secrets_store.py | 20 ++- .../service_connector_secrets_store.py | 2 +- 10 files changed, 308 insertions(+), 59 deletions(-) diff --git a/docs/book/deploying-zenml/zenml-self-hosted/deploy-with-docker.md b/docs/book/deploying-zenml/zenml-self-hosted/deploy-with-docker.md index fd03604ffbc..cd0f5385620 100644 --- a/docs/book/deploying-zenml/zenml-self-hosted/deploy-with-docker.md +++ b/docs/book/deploying-zenml/zenml-self-hosted/deploy-with-docker.md @@ -75,17 +75,32 @@ The SQL database is used as the default secret store. You only need to configure These configuration options are only relevant if you're using the AWS Secrets Manager as the secrets store backend. * **ZENML\_SECRETS\_STORE\_TYPE:** Set this to `aws` in order to set this type of secret store. -* **ZENML\_SECRETS\_STORE\_REGION\_NAME**: The AWS region to use. This must be set to the region where the AWS Secrets Manager service that you want to use is located. -* **ZENML\_SECRETS\_STORE\_AWS\_ACCESS\_KEY\_ID**: The AWS access key ID to use for authentication. This must be set to a valid AWS access key ID that has access to the AWS Secrets Manager service that you want to use. If you are using an IAM role attached to an EKS cluster to authenticate, you can omit this variable. NOTE: this is the same as setting the `AWS_ACCESS_KEY_ID` environment variable. -* **ZENML\_SECRETS\_STORE\_AWS\_SECRET\_ACCESS\_KEY**: The AWS secret access key to use for authentication. This must be set to a valid AWS secret access key that has access to the AWS Secrets Manager service that you want to use. If you are using an IAM role attached to an EKS cluster to authenticate, you can omit this variable. NOTE: this is the same as setting the `AWS_SECRET_ACCESS_KEY` environment variable. -* **ZENML\_SECRETS\_STORE\_AWS\_SESSION\_TOKEN**: Optional AWS session token to use for authentication. NOTE: this is the same as setting the `AWS_SESSION_TOKEN` environment variable. + +The AWS Secrets Store uses the ZenML AWS Service Connector under the hood to authenticate with the AWS Secrets Manager API. This means that you can use any of the [authentication methods supported by the AWS Service Connector](https://docs.zenml.io/stacks-and-components/auth-management/aws-service-connector#authentication-methods) to authenticate with the AWS Secrets Manager API. The following configuration options are supported: + +* **ZENML\_SECRETS\_STORE\_AUTH\_METHOD**: The AWS Service Connector authentication method to use (e.g. `secret-key` or `iam-role`). +* **ZENML\_SECRETS\_STORE\_AUTH\_CONFIG**: The AWS Service Connector configuration, in JSON format (e.g. `{"role_arn": "arn:aws:iam::123456789012:role/MyRole"}`). * **ZENML\_SECRETS\_STORE\_SECRET\_LIST\_REFRESH\_TIMEOUT**: AWS' [Secrets Manager](https://aws.amazon.com/secrets-manager) has a known issue where it does not immediately reflect new and updated secrets in the `list_secrets` results. To work around this issue, you can set this refresh timeout value to a non-zero value to get the ZenML server to wait after creating or updating an AWS secret until the changes are reflected in the secrets returned by `list_secrets` or the number of seconds specified by this value has elapsed. Defaults to `0` (disabled). Should not be set to a high value as it may cause thread starvation in the ZenML server on high load. + +> **Note:** The remaining configuration options are deprecated and may be removed in a future release. Instead, you should set the `ZENML_SECRETS_STORE_AUTH_METHOD` and `ZENML_SECRETS_STORE_AUTH_CONFIG` variables to use the AWS Service Connector authentication method. + +* **ZENML\_SECRETS\_STORE\_REGION\_NAME**: The AWS region to use. This must be set to the region where the AWS Secrets Manager service that you want to use is located. +* **ZENML\_SECRETS\_STORE\_AWS\_ACCESS\_KEY\_ID**: The AWS access key ID to use for authentication. This must be set to a valid AWS access key ID that has access to the AWS Secrets Manager service that you want to use. If you are using an IAM role attached to an EKS cluster to authenticate, you can omit this variable. +* **ZENML\_SECRETS\_STORE\_AWS\_SECRET\_ACCESS\_KEY**: The AWS secret access key to use for authentication. This must be set to a valid AWS secret access key that has access to the AWS Secrets Manager service that you want to use. If you are using an IAM role attached to an EKS cluster to authenticate, you can omit this variable. {% endtab %} {% tab title="GCP" %} These configuration options are only relevant if you're using the GCP Secrets Manager as the secrets store backend. * **ZENML\_SECRETS\_STORE\_TYPE:** Set this to `gcp` in order to set this type of secret store. + +The GCP Secrets Store uses the ZenML GCP Service Connector under the hood to authenticate with the GCP Secrets Manager API. This means that you can use any of the [authentication methods supported by the GCP Service Connector](https://docs.zenml.io/stacks-and-components/auth-management/gcp-service-connector#authentication-methods) to authenticate with the GCP Secrets Manager API. The following configuration options are supported: + +* **ZENML\_SECRETS\_STORE\_AUTH\_METHOD**: The GCP Service Connector authentication method to use (e.g. `service-account`). +* **ZENML\_SECRETS\_STORE\_AUTH\_CONFIG**: The GCP Service Connector configuration, in JSON format (e.g. `{"project_id": "my-project", "service_account_json": { ... }}`). + +> **Note:** The remaining configuration options are deprecated and may be removed in a future release. Instead, you should set the `ZENML_SECRETS_STORE_AUTH_METHOD` and `ZENML_SECRETS_STORE_AUTH_CONFIG` variables to use the GCP Service Connector authentication method. + * **ZENML\_SECRETS\_STORE\_PROJECT\_ID**: The GCP project ID to use. This must be set to the project ID where the GCP Secrets Manager service that you want to use is located. * **GOOGLE\_APPLICATION\_CREDENTIALS**: The path to the GCP service account credentials file to use for authentication. This must be set to a valid GCP service account credentials file that has access to the GCP Secrets Manager service that you want to use. If you are using a GCP service account attached to a GKE cluster to authenticate, you can omit this variable. NOTE: the path to the credentials file must be mounted into the container. {% endtab %} @@ -95,9 +110,17 @@ These configuration options are only relevant if you're using Azure Key Vault as * **ZENML\_SECRETS\_STORE\_TYPE:** Set this to `azure` in order to set this type of secret store. * **ZENML\_SECRETS\_STORE\_KEY\_VAULT\_NAME**: The name of the Azure Key Vault. This must be set to point to the Azure Key Vault instance that you want to use. -* **ZENML\_SECRETS\_STORE\_AZURE\_CLIENT\_ID**: The Azure application service principal client ID to use to authenticate with the Azure Key Vault API. If you are running the ZenML server hosted in Azure and are using a managed identity to access the Azure Key Vault service, you can omit this variable. NOTE: this is the same as setting the `AZURE_CLIENT_ID` environment variable. -* **ZENML\_SECRETS\_STORE\_AZURE\_CLIENT\_SECRET**: The Azure application service principal client secret to use to authenticate with the Azure Key Vault API. If you are running the ZenML server hosted in Azure and are using a managed identity to access the Azure Key Vault service, you can omit this variable. NOTE: this is the same as setting the `AZURE_CLIENT_SECRET` environment variable. -* **ZENML\_SECRETS\_STORE\_AZURE\_TENANT\_ID**: The Azure application service principal tenant ID to use to authenticate with the Azure Key Vault API. If you are running the ZenML server hosted in Azure and are using a managed identity to access the Azure Key Vault service, you can omit this variable. NOTE: this is the same as setting the `AZURE_TENANT_ID` environment variable. + +The Azure Secrets Store uses the ZenML Azure Service Connector under the hood to authenticate with the Azure Key Vault API. This means that you can use any of the [authentication methods supported by the Azure Service Connector](https://docs.zenml.io/stacks-and-components/auth-management/azure-service-connector#authentication-methods) to authenticate with the Azure Key Vault API. The following configuration options are supported: + +* **ZENML\_SECRETS\_STORE\_AUTH\_METHOD**: The Azure Service Connector authentication method to use (e.g. `service-account`). +* **ZENML\_SECRETS\_STORE\_AUTH\_CONFIG**: The Azure Service Connector configuration, in JSON format (e.g. `{"tenant_id": "my-tenant-id", "client_id": "my-client-id", "client_secret": "my-client-secret"}`). + +> **Note:** The remaining configuration options are deprecated and may be removed in a future release. Instead, you should set the `ZENML_SECRETS_STORE_AUTH_METHOD` and `ZENML_SECRETS_STORE_AUTH_CONFIG` variables to use the Azure Service Connector authentication method. + +* **ZENML\_SECRETS\_STORE\_AZURE\_CLIENT\_ID**: The Azure application service principal client ID to use to authenticate with the Azure Key Vault API. If you are running the ZenML server hosted in Azure and are using a managed identity to access the Azure Key Vault service, you can omit this variable. +* **ZENML\_SECRETS\_STORE\_AZURE\_CLIENT\_SECRET**: The Azure application service principal client secret to use to authenticate with the Azure Key Vault API. If you are running the ZenML server hosted in Azure and are using a managed identity to access the Azure Key Vault service, you can omit this variable. +* **ZENML\_SECRETS\_STORE\_AZURE\_TENANT\_ID**: The Azure application service principal tenant ID to use to authenticate with the Azure Key Vault API. If you are running the ZenML server hosted in Azure and are using a managed identity to access the Azure Key Vault service, you can omit this variable. {% endtab %} {% tab title="Hashicorp" %} diff --git a/docs/book/deploying-zenml/zenml-self-hosted/deploy-with-helm.md b/docs/book/deploying-zenml/zenml-self-hosted/deploy-with-helm.md index 93272c16425..2614bd8fcb5 100644 --- a/docs/book/deploying-zenml/zenml-self-hosted/deploy-with-helm.md +++ b/docs/book/deploying-zenml/zenml-self-hosted/deploy-with-helm.md @@ -312,7 +312,10 @@ This method requires you to configure a DNS service like AWS Route 53 or Google {% tab title="AWS" %} #### Using the AWS Secrets Manager as a secrets store backend -Unless explicitly disabled or configured otherwise, the ZenML server will use the SQL database as a secrets store backend. If you want to use the AWS Secrets Manager instead, you need to configure it in the Helm values. Depending on where you deploy your ZenML server and how your Kubernetes cluster is configured, you may also need to provide the AWS credentials needed to access the AWS Secrets Manager API: +Unless explicitly disabled or configured otherwise, the ZenML server will use the SQL database as a secrets store backend. If you want to use the AWS Secrets Manager instead, you need to configure it in the Helm values. Depending on where you deploy your ZenML server and how your Kubernetes cluster is configured, you may also need to provide the AWS credentials needed to access the AWS Secrets Manager API. + +The AWS Secrets Store uses the ZenML AWS Service Connector under the hood to authenticate with the AWS Secrets Manager API. This means that you can use any of the [authentication methods supported by the AWS Service Connector](https://docs.zenml.io/stacks-and-components/auth-management/aws-service-connector#authentication-methods) to authenticate with the AWS Secrets Manager API: + ```yaml zenml: @@ -331,24 +334,28 @@ Unless explicitly disabled or configured otherwise, the ZenML server will use th # Configuration for the AWS Secrets Manager secrets store aws: - # The AWS region to use. This must be set to the region where the AWS - # Secrets Manager service that you want to use is located. - region_name: us-east-1 - - # The AWS credentials to use to authenticate with the AWS Secrets - # Manager instance. You can omit these if you are running the ZenML server - # in an AWS EKS cluster that has an IAM role attached to it that has - # permissions to access the AWS Secrets Manager instance. - aws_access_key_id: - aws_secret_access_key: - aws_session_token: + # The AWS Service Connector authentication method to use. + authMethod: secret-key + + # The AWS Service Connector configuration. + authConfig: + # The AWS region to use. This must be set to the region where the AWS + # Secrets Manager service that you want to use is located. + region: us-east-1 + + # The AWS credentials to use to authenticate with the AWS Secrets + aws_access_key_id: + aws_secret_access_key: ``` {% endtab %} {% tab title="GCP" %} #### Using the GCP Secrets Manager as a secrets store backend -Unless explicitly disabled or configured otherwise, the ZenML server will use the SQL database as a secrets store backend. If you want to use the GCP Secrets Manager instead, you need to configure it in the Helm values. Depending on where you deploy your ZenML server and how your Kubernetes cluster is configured, you may also need to provide the GCP credentials needed to access the GCP Secrets Manager API: +Unless explicitly disabled or configured otherwise, the ZenML server will use the SQL database as a secrets store backend. If you want to use the GCP Secrets Manager instead, you need to configure it in the Helm values. Depending on where you deploy your ZenML server and how your Kubernetes cluster is configured, you may also need to provide the GCP credentials needed to access the GCP Secrets Manager API. + +The GCP Secrets Store uses the ZenML GCP Service Connector under the hood to authenticate with the GCP Secrets Manager API. This means that you can use any of the [authentication methods supported by the GCP Service Connector](https://docs.zenml.io/stacks-and-components/auth-management/gcp-service-connector#authentication-methods) to authenticate with the GCP Secrets Manager API: + ```yaml zenml: @@ -367,18 +374,19 @@ Unless explicitly disabled or configured otherwise, the ZenML server will use th # Configuration for the GCP Secrets Manager secrets store gcp: - # The GCP project ID to use. This must be set to the project ID where the - # GCP Secrets Manager service that you want to use is located. - project_id: my-gcp-project + # The GCP Service Connector authentication method to use. + authMethod: service-account - # Path to the GCP credentials file to use to authenticate with the GCP Secrets - # Manager instance. You can omit this if you are running the ZenML server - # in a GCP GKE cluster that uses workload identity to authenticate with - # GCP services without the need for credentials. - # NOTE: the credentials file needs to be copied in the helm chart folder - # and the path configured here needs to be relative to the root of the - # helm chart. - google_application_credentials: cloud-credentials.json + # The GCP Service Connector configuration. + authConfig: + + # The GCP project ID to use. This must be set to the project ID where the + # GCP Secrets Manager service that you want to use is located. + project_id: my-gcp-project + + # GCP credentials JSON to use to authenticate with the GCP Secrets + # Manager instance. + google_application_credentials: '{...}' serviceAccount: @@ -393,7 +401,10 @@ Unless explicitly disabled or configured otherwise, the ZenML server will use th {% tab title="Azure" %} #### Using the Azure Key Vault as a secrets store backend -Unless explicitly disabled or configured otherwise, the ZenML server will use the SQL database as a secrets store backend. If you want to use the Azure Key Vault service instead, you need to configure it in the Helm values. Depending on where you deploy your ZenML server and how your Kubernetes cluster is configured, you may also need to provide the Azure credentials needed to access the Azure Key Vault API: +Unless explicitly disabled or configured otherwise, the ZenML server will use the SQL database as a secrets store backend. If you want to use the Azure Key Vault service instead, you need to configure it in the Helm values. Depending on where you deploy your ZenML server and how your Kubernetes cluster is configured, you may also need to provide the Azure credentials needed to access the Azure Key Vault API. + +The Azure Secrets Store uses the ZenML Azure Service Connector under the hood to authenticate with the Azure Key Vault API. This means that you can use any of the [authentication methods supported by the Azure Service Connector](https://docs.zenml.io/stacks-and-components/auth-management/azure-service-connector#authentication-methods) to authenticate with the Azure Key Vault API: + ```yaml zenml: @@ -416,13 +427,17 @@ Unless explicitly disabled or configured otherwise, the ZenML server will use th # Key Vault instance that you want to use. key_vault_name: - # The Azure application service principal credentials to use to - # authenticate with the Azure Key Vault API. You can omit these if you are - # running the ZenML server hosted in Azure and are using a managed - # identity to access the Azure Key Vault service. - azure_client_id: - azure_client_secret: - azure_tenant_id: + # The Azure Service Connector authentication method to use. + authMethod: service-principal + + # The Azure Service Connector configuration. + authConfig: + + # The Azure application service principal credentials to use to + # authenticate with the Azure Key Vault API. + client_id: + client_secret: + tenant_id: ``` {% endtab %} diff --git a/src/zenml/zen_server/deploy/helm/templates/server-db-job.yaml b/src/zenml/zen_server/deploy/helm/templates/server-db-job.yaml index 99138fc71cc..5d47721425a 100644 --- a/src/zenml/zen_server/deploy/helm/templates/server-db-job.yaml +++ b/src/zenml/zen_server/deploy/helm/templates/server-db-job.yaml @@ -83,18 +83,28 @@ spec: - name: ZENML_SECRETS_STORE_TYPE value: {{ .Values.zenml.secretsStore.type | quote }} {{- if eq .Values.zenml.secretsStore.type "aws" }} + - name: ZENML_SECRETS_STORE_AUTH_METHOD + value: {{ .Values.zenml.secretsStore.aws.authMethod | quote }} + {{- if .Values.zenml.secretsStore.aws.region_name }} - name: ZENML_SECRETS_STORE_REGION_NAME value: {{ .Values.zenml.secretsStore.aws.region_name | quote }} + {{- endif }} - name: ZENML_SECRETS_STORE_SECRET_LIST_REFRESH_TIMEOUT value: {{ .Values.zenml.secretsStore.aws.secret_list_refresh_timeout | quote }} {{- else if eq .Values.zenml.secretsStore.type "gcp" }} + - name: ZENML_SECRETS_STORE_AUTH_METHOD + value: {{ .Values.zenml.secretsStore.gcp.authMethod | quote }} + {{- if .Values.zenml.secretsStore.gcp.project_id }} - name: ZENML_SECRETS_STORE_PROJECT_ID value: {{ .Values.zenml.secretsStore.gcp.project_id | quote }} + {{- endif }} {{- if .Values.zenml.secretsStore.gcp.google_application_credentials }} - name: GOOGLE_APPLICATION_CREDENTIALS value: /gcp-credentials/credentials.json {{- end }} {{- else if eq .Values.zenml.secretsStore.type "azure" }} + - name: ZENML_SECRETS_STORE_AUTH_METHOD + value: {{ .Values.zenml.secretsStore.azure.authMethod | quote }} - name: ZENML_SECRETS_STORE_KEY_VAULT_NAME value: {{ .Values.zenml.secretsStore.azure.key_vault_name | quote }} {{- else if eq .Values.zenml.secretsStore.type "hashicorp" }} diff --git a/src/zenml/zen_server/deploy/helm/templates/server-deployment.yaml b/src/zenml/zen_server/deploy/helm/templates/server-deployment.yaml index 380ba137277..d16dca8726a 100644 --- a/src/zenml/zen_server/deploy/helm/templates/server-deployment.yaml +++ b/src/zenml/zen_server/deploy/helm/templates/server-deployment.yaml @@ -161,18 +161,28 @@ spec: - name: ZENML_SECRETS_STORE_TYPE value: {{ .Values.zenml.secretsStore.type | quote }} {{- if eq .Values.zenml.secretsStore.type "aws" }} + - name: ZENML_SECRETS_STORE_AUTH_METHOD + value: {{ .Values.zenml.secretsStore.aws.authMethod | quote }} + {{- if .Values.zenml.secretsStore.aws.region_name }} - name: ZENML_SECRETS_STORE_REGION_NAME value: {{ .Values.zenml.secretsStore.aws.region_name | quote }} + {{- endif }} - name: ZENML_SECRETS_STORE_SECRET_LIST_REFRESH_TIMEOUT value: {{ .Values.zenml.secretsStore.aws.secret_list_refresh_timeout | quote }} {{- else if eq .Values.zenml.secretsStore.type "gcp" }} + - name: ZENML_SECRETS_STORE_AUTH_METHOD + value: {{ .Values.zenml.secretsStore.gcp.authMethod | quote }} + {{- if .Values.zenml.secretsStore.gcp.project_id }} - name: ZENML_SECRETS_STORE_PROJECT_ID value: {{ .Values.zenml.secretsStore.gcp.project_id | quote }} + {{- endif }} {{- if .Values.zenml.secretsStore.gcp.google_application_credentials }} - name: GOOGLE_APPLICATION_CREDENTIALS value: /gcp-credentials/credentials.json {{- end }} {{- else if eq .Values.zenml.secretsStore.type "azure" }} + - name: ZENML_SECRETS_STORE_AUTH_METHOD + value: {{ .Values.zenml.secretsStore.azure.authMethod | quote }} - name: ZENML_SECRETS_STORE_KEY_VAULT_NAME value: {{ .Values.zenml.secretsStore.azure.key_vault_name | quote }} {{- else if eq .Values.zenml.secretsStore.type "hashicorp" }} diff --git a/src/zenml/zen_server/deploy/helm/templates/server-secret.yaml b/src/zenml/zen_server/deploy/helm/templates/server-secret.yaml index d3f514be68a..d1a50989240 100644 --- a/src/zenml/zen_server/deploy/helm/templates/server-secret.yaml +++ b/src/zenml/zen_server/deploy/helm/templates/server-secret.yaml @@ -31,6 +31,9 @@ data: ZENML_SECRETS_STORE_ENCRYPTION_KEY: {{ .Values.zenml.secretsStore.encryptionKey | b64enc | quote }} {{- end }} {{- else if eq .Values.zenml.secretsStore.type "aws" }} + {{- if .Values.zenml.secretsStore.aws.authConfig }} + ZENML_SECRETS_STORE_AUTH_CONFIG: {{ .Values.zenml.secretsStore.aws.authConfig | toJson | b64enc | quote }} + {{- end }} {{- if .Values.zenml.secretsStore.aws.aws_access_key_id }} ZENML_SECRETS_STORE_AWS_ACCESS_KEY_ID: {{ .Values.zenml.secretsStore.aws.aws_access_key_id | b64enc | quote }} {{- end }} @@ -41,6 +44,9 @@ data: ZENML_SECRETS_STORE_AWS_SESSION_TOKEN: {{ .Values.zenml.secretsStore.aws.aws_session_token | b64enc | quote }} {{- end }} {{- else if eq .Values.zenml.secretsStore.type "azure" }} + {{- if .Values.zenml.secretsStore.azure.authConfig }} + ZENML_SECRETS_STORE_AUTH_CONFIG: {{ .Values.zenml.secretsStore.azure.authConfig | toJson | b64enc | quote }} + {{- end }} {{- if .Values.zenml.secretsStore.azure.azure_client_id }} ZENML_SECRETS_STORE_AZURE_CLIENT_ID: {{ .Values.zenml.secretsStore.azure.azure_client_id | b64enc | quote }} {{- end }} @@ -51,6 +57,9 @@ data: ZENML_SECRETS_STORE_AZURE_TENANT_ID: {{ .Values.zenml.secretsStore.azure.azure_tenant_id | b64enc | quote }} {{- end }} {{- else if eq .Values.zenml.secretsStore.type "gcp" }} + {{- if .Values.zenml.secretsStore.gcp.authConfig }} + ZENML_SECRETS_STORE_AUTH_CONFIG: {{ .Values.zenml.secretsStore.gcp.authConfig | toJson | b64enc | quote }} + {{- end }} {{- if .Values.zenml.secretsStore.gcp.google_application_credentials }} GOOGLE_APPLICATION_CREDENTIALS_FILE: {{ .Values.zenml.secretsStore.gcp.google_application_credentials | b64enc | quote }} {{- end }} diff --git a/src/zenml/zen_server/deploy/helm/values.yaml b/src/zenml/zen_server/deploy/helm/values.yaml index 41a1bcec522..6eb2974fa55 100644 --- a/src/zenml/zen_server/deploy/helm/values.yaml +++ b/src/zenml/zen_server/deploy/helm/values.yaml @@ -250,19 +250,71 @@ zenml: # AWS secrets store configuration. Only relevant if the `aws` secrets store # type is configured. + # + # The AWS secrets store uses the AWS Service Connector under the hood to + # authenticate with the AWS Secrets Manager API. This means that you can + # use the same authentication methods and configuration as you would use for + # the AWS Service Connector. Just set the `authMethod` field to the + # authentication method that you want to use and set the required + # configuration attributes under the `authConfig` field. + # + # 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 + # + # You can also use the ZenML CLI to get the list of supported authentication + # methods and their configuration options, e.g.: + # + # ```shell + # zenml service-connector describe-type aws + # zenml service-connector describe-type aws --auth-method secret-key + # ``` aws: + # The AWS Service Connector authentication method to use. The currently + # supported authentication methods are: + # + # - implicit - Use the IAM role attached to the ZenML server pod or + # environment variables to authenticate with the AWS Secrets + # Manager API + # - secret-key - Use an AWS secret key + # - iam-role - Use an IAM role + # - session-token - Use an AWS session token derived from an AWS secret + # key + # - federation-token - Use an AWS federation token derived from an AWS + # secret key + authMethod: secret-key + + # The AWS Service Connector authentication configuration. This should + # include the corresponding authentication configuration attributes for + # the `authMethod` that you have chosen above. + authConfig: + # The AWS region to use. This must be set to the region where the AWS + # Secrets Manager service that you want to use is located. Mandatory + # for all authentication methods. + region: + # The AWS access key and secret key to use to authenticate with the AWS + # Secrets Manager instance. Both are required if the `authMethod` is set + # to `secret-key`, `sts-token`, `iam-role`, or `federation-token`. + 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`. + role_arn: + # The AWS region to use. This must be set to the region where the AWS # Secrets Manager service that you want to use is located. - region_name: us-east-1 + # + # NOTE: deprecated; use `authConfig.region` instead. + region_name: # The AWS credentials to use to authenticate with the AWS Secrets # Manager instance. You can omit these if you are running the ZenML server # in an AWS EKS cluster that has an IAM role attached to it that has # permissions to access the AWS Secrets Manager instance. - # NOTE: setting this is the same as setting the AWS_ACCESS_KEY_ID, - # AWS_SECRET_ACCESS_KEY, and AWS_SESSION_TOKEN environment variables - # in the zenml.secretEnvironment variable. + # + # NOTE: deprecated; use `authConfig.aws_access_key_id`, + # and `authConfig.aws_secret_access_key` instead. aws_access_key_id: aws_secret_access_key: aws_session_token: @@ -280,22 +332,120 @@ zenml: # GCP secrets store configuration. Only relevant if the `gcp` secrets store # type is configured. + # + # The GCP secrets store uses the GCP Service Connector under the hood to + # authenticate with the GCP Secrets Manager API. This means that you can + # use the same authentication methods and configuration as you would use for + # the GCP Service Connector. Just set the `authMethod` field to the + # authentication method that you want to use and set the required + # configuration attributes under the `authConfig` field. + # + # 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 + # + # You can also use the ZenML CLI to get the list of supported authentication + # methods and their configuration options, e.g.: + # + # ```shell + # zenml service-connector describe-type gcp + # zenml service-connector describe-type gcp --auth-method service-account + # ``` gcp: + # The GCP Service Connector authentication method to use. The currently + # supported authentication methods are: + # + # - implicit - Use the GCP service account attached to the ZenML server + # pod or environment variables to authenticate with the GCP + # Secrets Manager API + # - user-account - Use a GCP user account + # - service-account - Use a GCP service account + # - impersonation - Use the GCP service account impersonation feature + authMethod: service-account + + # The GCP Service Connector authentication configuration. This should + # include the corresponding authentication configuration attributes for + # the `authMethod` that you have chosen above. + authConfig: + # The GCP project ID to use. This must be set to the project ID where + # the GCP Secrets Manager service that you want to use is located. + # Mandatory for all authentication methods. + project_id: + + # The GCP user account credentials to use to authenticate with the GCP + # Secrets Manager instance. Required if the `authMethod` is set to + # `user-account`. + user_account_json: + + # The GCP service account credentials to use to authenticate with the + # GCP Secrets Manager instance. Required if the `authMethod` is set to + # `service-account` or `impersonation`. + service_account_json: + + # The GCP service account to impersonate when authenticating with the + # GCP Secrets Manager instance. Required if the `authMethod` is set to + # `impersonation`. + target_principal: + # The GCP project ID to use. This must be set to the project ID where the # GCP Secrets Manager service that you want to use is located. - project_id: my-gcp-project + # + # NOTE: deprecated; use `authConfig.project_id` instead. + project_id: # The JSON content of the GCP credentials file to use to authenticate with # the GCP Secrets Manager instance. You can omit this if you are running # the ZenML server in a GCP GKE cluster that uses workload identity to # authenticate with GCP services without the need for credentials. + # + # NOTE: deprecated; use `authConfig.service_account_json` instead. google_application_credentials: - # AWS Key Vault secrets store configuration. Only relevant if the `azure` + # Azure Key Vault secrets store configuration. Only relevant if the `azure` # secrets store type is configured. + # + # The Azure secrets store uses the Azure Service Connector under the hood to + # authenticate with the Azure Key Vault API. This means that you can + # use the same authentication methods and configuration as you would use for + # the Azure Service Connector. Just set the `authMethod` field to the + # authentication method that you want to use and set the required + # configuration attributes under the `authConfig` field. + # + # 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 + # + # You can also use the ZenML CLI to get the list of supported authentication + # methods and their configuration options, e.g.: + # + # ```shell + # zenml service-connector describe-type azure + # zenml service-connector describe-type azure --auth-method service-principal + # ``` azure: + # The Azure Service Connector authentication method to use. The currently + # supported authentication methods are: + # + # - implicit - Use the Azure managed identity attached to the ZenML server + # pod or environment variables to authenticate with the Azure + # Key Vault API + # - service-principal - Use an Azure service principal + authMethod: service-principal + + # The Azure Service Connector authentication configuration. This should + # include the corresponding authentication configuration attributes for + # the `authMethod` that you have chosen above. + authConfig: + + # The Azure service principal credentials to use to authenticate with + # the Azure Key Vault API. All three are Required if the `authMethod` is + # set to `service-principal`. + client_id: + client_secret: + tenant_id: + # The name of the Azure Key Vault. This must be set to point to the Azure # Key Vault instance that you want to use. key_vault_name: @@ -304,9 +454,9 @@ zenml: # authenticate with the Azure Key Vault API. You can omit these if you are # running the ZenML server hosted in Azure and are using a managed # identity to access the Azure Key Vault service. - # NOTE: setting this is the same as setting the AZURE_CLIENT_ID, - # AZURE_CLIENT_SECRET, and AZURE_TENANT_ID environment variables - # in the zenml.secretEnvironment variable. + # + # NOTE: deprecated; use `authConfig.client_id`, `authConfig.client_secret`, + # and `authConfig.tenant_id` instead. azure_client_id: azure_client_secret: azure_tenant_id: diff --git a/src/zenml/zen_stores/secrets_stores/aws_secrets_store.py b/src/zenml/zen_stores/secrets_stores/aws_secrets_store.py index 95b07ba6235..b639e8790dc 100644 --- a/src/zenml/zen_stores/secrets_stores/aws_secrets_store.py +++ b/src/zenml/zen_stores/secrets_stores/aws_secrets_store.py @@ -105,9 +105,7 @@ def region(self) -> str: if region: return str(region) - raise ValueError( - "AWS `region` must be specified in the auth_config." - ) + raise ValueError("AWS `region` must be specified in the auth_config.") @root_validator(pre=True) def populate_config(cls, values: Dict[str, Any]) -> Dict[str, Any]: @@ -122,7 +120,19 @@ def populate_config(cls, values: Dict[str, Any]) -> Dict[str, Any]: Raises: ValueError: If the connector attribute is not set. """ - if "auth_method" not in values or "auth_config" not in values: + # Search for legacy attributes and populate the connector configuration + # from them, if they exist. + if ( + values.get("aws_access_key_id") + and values.get("aws_secret_access_key") + and values.get("region_name") + ): + logger.warning( + "The `aws_access_key_id`, `aws_secret_access_key` and " + "`region_name` AWS secrets store attributes are deprecated and " + "will be removed in a future version of ZenML. Please use the " + "`auth_method` and `auth_config` attributes instead." + ) values["auth_method"] = AWSAuthenticationMethods.SECRET_KEY values["auth_config"] = dict( aws_access_key_id=values.get("aws_access_key_id"), diff --git a/src/zenml/zen_stores/secrets_stores/azure_secrets_store.py b/src/zenml/zen_stores/secrets_stores/azure_secrets_store.py index eb379be7cb7..cf457e171e8 100644 --- a/src/zenml/zen_stores/secrets_stores/azure_secrets_store.py +++ b/src/zenml/zen_stores/secrets_stores/azure_secrets_store.py @@ -96,7 +96,19 @@ def populate_config(cls, values: Dict[str, Any]) -> Dict[str, Any]: Raises: ValueError: If the connector attribute is not set. """ - if "auth_method" not in values or "auth_config" not in values: + # Search for legacy attributes and populate the connector configuration + # from them, if they exist. + if ( + values.get("azure_client_id") + and values.get("azure_client_secret") + and values.get("azure_tenant_id") + ): + logger.warning( + "The `azure_client_id`, `azure_client_secret` and " + "`azure_tenant_id` attributes are deprecated and will be " + "removed in a future version or ZenML. Please use the " + "`auth_method` and `auth_config` attributes instead." + ) values[ "auth_method" ] = AzureAuthenticationMethods.SERVICE_PRINCIPAL diff --git a/src/zenml/zen_stores/secrets_stores/gcp_secrets_store.py b/src/zenml/zen_stores/secrets_stores/gcp_secrets_store.py index 22559e08020..3bc049c5906 100644 --- a/src/zenml/zen_stores/secrets_stores/gcp_secrets_store.py +++ b/src/zenml/zen_stores/secrets_stores/gcp_secrets_store.py @@ -114,15 +114,25 @@ def populate_config(cls, values: Dict[str, Any]) -> Dict[str, Any]: Raises: ValueError: If the connector attribute is not set. """ - if "auth_method" not in values or "auth_config" not in values: + # Search for legacy attributes and populate the connector configuration + # from them, if they exist. + if values.get("project_id") and os.environ.get( + "GOOGLE_APPLICATION_CREDENTIALS" + ): + logger.warning( + "The `project_id` GCP secrets store attribute and the " + "`GOOGLE_APPLICATION_CREDENTIALS` environment variable are " + "deprecated and will be removed in a future version of ZenML. " + "Please use the `auth_method` and `auth_config` attributes " + "instead." + ) values["auth_method"] = GCPAuthenticationMethods.SERVICE_ACCOUNT values["auth_config"] = dict( project_id=values.get("project_id"), ) - if os.environ.get("GOOGLE_APPLICATION_CREDENTIALS"): - # Load the service account credentials from the file - with open(os.environ["GOOGLE_APPLICATION_CREDENTIALS"]) as f: - values["auth_config"]["service_account_json"] = f.read() + # Load the service account credentials from the file + with open(os.environ["GOOGLE_APPLICATION_CREDENTIALS"]) as f: + values["auth_config"]["service_account_json"] = f.read() return values diff --git a/src/zenml/zen_stores/secrets_stores/service_connector_secrets_store.py b/src/zenml/zen_stores/secrets_stores/service_connector_secrets_store.py index 14134bcc292..c33bc38d6cf 100644 --- a/src/zenml/zen_stores/secrets_stores/service_connector_secrets_store.py +++ b/src/zenml/zen_stores/secrets_stores/service_connector_secrets_store.py @@ -78,7 +78,7 @@ class ServiceConnectorSecretsStore(BaseSecretsStore): _connector: Optional[ServiceConnector] = None _client: Optional[Any] = None - _lock: Lock = Lock() + _lock: Lock = Field(default_factory=Lock) def _initialize(self) -> None: """Initialize the secrets store.""" From b685e0bdb0207587c51cfae8974917bd3b9a9a2a Mon Sep 17 00:00:00 2001 From: Stefan Nica Date: Fri, 15 Dec 2023 22:49:51 +0100 Subject: [PATCH 4/7] Tested all secrets stores and fixed remaining bugs --- .../zenml-self-hosted/deploy-with-helm.md | 14 +++++++- .../deploy/helm/templates/server-db-job.yaml | 4 +-- .../helm/templates/server-deployment.yaml | 4 +-- .../secrets_stores/gcp_secrets_store.py | 6 +--- .../service_connector_secrets_store.py | 32 +++++++++++++++++-- 5 files changed, 47 insertions(+), 13 deletions(-) diff --git a/docs/book/deploying-zenml/zenml-self-hosted/deploy-with-helm.md b/docs/book/deploying-zenml/zenml-self-hosted/deploy-with-helm.md index 2614bd8fcb5..852938a014d 100644 --- a/docs/book/deploying-zenml/zenml-self-hosted/deploy-with-helm.md +++ b/docs/book/deploying-zenml/zenml-self-hosted/deploy-with-helm.md @@ -386,7 +386,19 @@ The GCP Secrets Store uses the ZenML GCP Service Connector under the hood to aut # GCP credentials JSON to use to authenticate with the GCP Secrets # Manager instance. - google_application_credentials: '{...}' + google_application_credentials: | + { + "type": "service_account", + "project_id": "my-project", + "private_key_id": "...", + "private_key": "-----BEGIN PRIVATE KEY-----\n...=\n-----END PRIVATE KEY-----\n", + "client_email": "...", + "client_id": "...", + "auth_uri": "https://accounts.google.com/o/oauth2/auth", + "token_uri": "https://oauth2.googleapis.com/token", + "auth_provider_x509_cert_url": "https://www.googleapis.com/oauth2/v1/certs", + "client_x509_cert_url": "..." + } serviceAccount: diff --git a/src/zenml/zen_server/deploy/helm/templates/server-db-job.yaml b/src/zenml/zen_server/deploy/helm/templates/server-db-job.yaml index 5d47721425a..d3fbff15fc6 100644 --- a/src/zenml/zen_server/deploy/helm/templates/server-db-job.yaml +++ b/src/zenml/zen_server/deploy/helm/templates/server-db-job.yaml @@ -88,7 +88,7 @@ spec: {{- if .Values.zenml.secretsStore.aws.region_name }} - name: ZENML_SECRETS_STORE_REGION_NAME value: {{ .Values.zenml.secretsStore.aws.region_name | quote }} - {{- endif }} + {{- end }} - name: ZENML_SECRETS_STORE_SECRET_LIST_REFRESH_TIMEOUT value: {{ .Values.zenml.secretsStore.aws.secret_list_refresh_timeout | quote }} {{- else if eq .Values.zenml.secretsStore.type "gcp" }} @@ -97,7 +97,7 @@ spec: {{- if .Values.zenml.secretsStore.gcp.project_id }} - name: ZENML_SECRETS_STORE_PROJECT_ID value: {{ .Values.zenml.secretsStore.gcp.project_id | quote }} - {{- endif }} + {{- end }} {{- if .Values.zenml.secretsStore.gcp.google_application_credentials }} - name: GOOGLE_APPLICATION_CREDENTIALS value: /gcp-credentials/credentials.json diff --git a/src/zenml/zen_server/deploy/helm/templates/server-deployment.yaml b/src/zenml/zen_server/deploy/helm/templates/server-deployment.yaml index d16dca8726a..1ff1a3ad563 100644 --- a/src/zenml/zen_server/deploy/helm/templates/server-deployment.yaml +++ b/src/zenml/zen_server/deploy/helm/templates/server-deployment.yaml @@ -166,7 +166,7 @@ spec: {{- if .Values.zenml.secretsStore.aws.region_name }} - name: ZENML_SECRETS_STORE_REGION_NAME value: {{ .Values.zenml.secretsStore.aws.region_name | quote }} - {{- endif }} + {{- end }} - name: ZENML_SECRETS_STORE_SECRET_LIST_REFRESH_TIMEOUT value: {{ .Values.zenml.secretsStore.aws.secret_list_refresh_timeout | quote }} {{- else if eq .Values.zenml.secretsStore.type "gcp" }} @@ -175,7 +175,7 @@ spec: {{- if .Values.zenml.secretsStore.gcp.project_id }} - name: ZENML_SECRETS_STORE_PROJECT_ID value: {{ .Values.zenml.secretsStore.gcp.project_id | quote }} - {{- endif }} + {{- end }} {{- if .Values.zenml.secretsStore.gcp.google_application_credentials }} - name: GOOGLE_APPLICATION_CREDENTIALS value: /gcp-credentials/credentials.json diff --git a/src/zenml/zen_stores/secrets_stores/gcp_secrets_store.py b/src/zenml/zen_stores/secrets_stores/gcp_secrets_store.py index 3bc049c5906..06d5b4a649f 100644 --- a/src/zenml/zen_stores/secrets_stores/gcp_secrets_store.py +++ b/src/zenml/zen_stores/secrets_stores/gcp_secrets_store.py @@ -35,7 +35,6 @@ from google.api_core import exceptions as google_exceptions from google.cloud.secretmanager import SecretManagerServiceClient -from google.oauth2 import service_account as gcp_service_account from pydantic import root_validator from zenml.analytics.enums import AnalyticsEvent @@ -183,10 +182,7 @@ def _initialize_client_from_connector(self, client: Any) -> Any: Returns: The GCP Secrets Manager client. """ - assert isinstance(client, gcp_service_account.Credentials) - return SecretManagerServiceClient( - project=self.config.project_id, credentials=client - ) + return SecretManagerServiceClient(credentials=client) # ------ # Secrets diff --git a/src/zenml/zen_stores/secrets_stores/service_connector_secrets_store.py b/src/zenml/zen_stores/secrets_stores/service_connector_secrets_store.py index c33bc38d6cf..32372ae9a28 100644 --- a/src/zenml/zen_stores/secrets_stores/service_connector_secrets_store.py +++ b/src/zenml/zen_stores/secrets_stores/service_connector_secrets_store.py @@ -14,6 +14,7 @@ """Base secrets store class used for all secrets stores that use a service connector.""" +import json from abc import abstractmethod from threading import Lock from typing import ( @@ -25,7 +26,7 @@ ) from uuid import uuid4 -from pydantic import Field +from pydantic import Field, root_validator from zenml.config.secrets_store_config import SecretsStoreConfiguration from zenml.logger import get_logger @@ -54,6 +55,20 @@ class ServiceConnectorSecretsStoreConfiguration(SecretsStoreConfiguration): auth_method: str auth_config: Dict[str, Any] = Field(default_factory=dict) + @root_validator(pre=True) + def validate_auth_config(cls, values: Dict[str, Any]) -> Dict[str, Any]: + """Convert the authentication configuration if given in JSON format. + + Args: + values: The configuration values. + + Returns: + The validated configuration values. + """ + if isinstance(values.get("auth_config"), str): + values["auth_config"] = json.loads(values["auth_config"]) + return values + class ServiceConnectorSecretsStore(BaseSecretsStore): """Base secrets store class for service connector-based secrets stores. @@ -78,10 +93,11 @@ class ServiceConnectorSecretsStore(BaseSecretsStore): _connector: Optional[ServiceConnector] = None _client: Optional[Any] = None - _lock: Lock = Field(default_factory=Lock) + _lock: Optional[Lock] = None def _initialize(self) -> None: """Initialize the secrets store.""" + self._lock = Lock() # Initialize the client early, just to catch any configuration or # authentication errors early, before the Secrets Store is used. _ = self.client @@ -126,6 +142,16 @@ def _get_client(self) -> Any: self._client = self._initialize_client_from_connector(client) return self._client + @property + def lock(self) -> Lock: + """Get the lock used to treat the client initialization as a critical section. + + Returns: + The lock instance. + """ + assert self._lock is not None + return self._lock + @property def client(self) -> Any: """Get the secrets store API client. @@ -137,7 +163,7 @@ def client(self) -> Any: # of different threads. We want to make sure that we only initialize # the client once, and then reuse it. We have to use a lock to treat # this method as a critical section. - with self._lock: + with self.lock: return self._get_client() @abstractmethod From d3bf7427e73c6fb89862910d9fa5f2d2ed6551b1 Mon Sep 17 00:00:00 2001 From: Stefan Nica Date: Mon, 18 Dec 2023 10:55:14 +0100 Subject: [PATCH 5/7] Fix code repository created and updated fields --- src/zenml/zen_stores/schemas/code_repository_schemas.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/zenml/zen_stores/schemas/code_repository_schemas.py b/src/zenml/zen_stores/schemas/code_repository_schemas.py index 7a9a15e9846..0db52ceb869 100644 --- a/src/zenml/zen_stores/schemas/code_repository_schemas.py +++ b/src/zenml/zen_stores/schemas/code_repository_schemas.py @@ -109,13 +109,13 @@ def to_model(self, hydrate: bool = False) -> "CodeRepositoryResponse": user=self.user.to_model() if self.user else None, source=json.loads(self.source), logo_url=self.logo_url, + created=self.created, + updated=self.updated, ) metadata = None if hydrate: metadata = CodeRepositoryResponseMetadata( workspace=self.workspace.to_model(), - created=self.created, - updated=self.updated, config=json.loads(self.config), description=self.description, ) From 1c2ad8a7fd67e3d744920e83808ecd4adee7d467 Mon Sep 17 00:00:00 2001 From: Stefan Nica Date: Tue, 19 Dec 2023 10:21:41 +0100 Subject: [PATCH 6/7] Apply suggestions from code review Co-authored-by: Jayesh Sharma --- src/zenml/zen_stores/secrets_stores/aws_secrets_store.py | 4 ++-- .../secrets_stores/service_connector_secrets_store.py | 2 +- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/src/zenml/zen_stores/secrets_stores/aws_secrets_store.py b/src/zenml/zen_stores/secrets_stores/aws_secrets_store.py index b639e8790dc..0d2eb14d404 100644 --- a/src/zenml/zen_stores/secrets_stores/aws_secrets_store.py +++ b/src/zenml/zen_stores/secrets_stores/aws_secrets_store.py @@ -212,14 +212,14 @@ class AWSSecretsStore(ServiceConnectorSecretsStore): # -------------------------------- def _initialize_client_from_connector(self, client: Any) -> Any: - """Initialize the GCP Secrets Manager client from the service connector client. + """Initialize the AWS Secrets Manager client from the service connector client. Args: client: The authenticated client object returned by the service connector. Returns: - The GCP Secrets Manager client. + The AWS Secrets Manager client. """ assert isinstance(client, boto3.Session) return client.client( diff --git a/src/zenml/zen_stores/secrets_stores/service_connector_secrets_store.py b/src/zenml/zen_stores/secrets_stores/service_connector_secrets_store.py index 32372ae9a28..250eafa1038 100644 --- a/src/zenml/zen_stores/secrets_stores/service_connector_secrets_store.py +++ b/src/zenml/zen_stores/secrets_stores/service_connector_secrets_store.py @@ -115,7 +115,7 @@ def _get_client(self) -> Any: self._client = None if self._connector is None: - # Initialize a base AWS service connector with the credentials from + # Initialize a base service connector with the credentials from # the configuration. request = ServiceConnectorRequest( name="secrets-store", From 0b16e202ceb098821647328dcfd3af70f1ee0fd4 Mon Sep 17 00:00:00 2001 From: GitHub Actions Date: Tue, 19 Dec 2023 09:34:28 +0000 Subject: [PATCH 7/7] Auto-update of E2E template --- examples/e2e/.copier-answers.yml | 2 +- examples/e2e/steps/deployment/deployment_deploy.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/examples/e2e/.copier-answers.yml b/examples/e2e/.copier-answers.yml index 70e3f86e323..04637068e52 100644 --- a/examples/e2e/.copier-answers.yml +++ b/examples/e2e/.copier-answers.yml @@ -1,5 +1,5 @@ # Changes here will be overwritten by Copier -_commit: 2023.12.06 +_commit: 2023.12.12 _src_path: gh:zenml-io/template-e2e-batch data_quality_checks: true email: '' diff --git a/examples/e2e/steps/deployment/deployment_deploy.py b/examples/e2e/steps/deployment/deployment_deploy.py index b0b905b52c9..109f35b028e 100644 --- a/examples/e2e/steps/deployment/deployment_deploy.py +++ b/examples/e2e/steps/deployment/deployment_deploy.py @@ -37,7 +37,7 @@ def deployment_deploy() -> ( Annotated[ Optional[MLFlowDeploymentService], - ArtifactConfig(name="mlflow_deployment", is_endpoint_artifact=True), + ArtifactConfig(name="mlflow_deployment", is_deployment_artifact=True), ] ): """Predictions step.