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

zenml/zen_stores/schemas/workspace_schemas.py

@@ -12,8 +12,9 @@
 #  or implied. See the License for the specific language governing
 #  permissions and limitations under the License.
 """SQL Model Implementations for Workspaces."""
+
 from datetime import datetime
-from typing import TYPE_CHECKING, List
+from typing import TYPE_CHECKING, Any, List
 
 from sqlmodel import Relationship
 
@@ -28,12 +29,12 @@ from zenml.zen_stores.schemas.base_schemas import NamedSchema
 
 if TYPE_CHECKING:
     from zenml.zen_stores.schemas import (
+        ActionSchema,
         ArtifactVersionSchema,
         CodeRepositorySchema,
+        EventSourceSchema,
         FlavorSchema,
         ModelSchema,
-        ModelVersionArtifactSchema,
-        ModelVersionPipelineRunSchema,
         ModelVersionSchema,
         PipelineBuildSchema,
         PipelineDeploymentSchema,
@@ -43,9 +44,11 @@ if TYPE_CHECKING:
         ScheduleSchema,
         SecretSchema,
         ServiceConnectorSchema,
+        ServiceSchema,
         StackComponentSchema,
         StackSchema,
         StepRunSchema,
+        TriggerSchema,
     )
 
 
@@ -100,35 +103,40 @@ class WorkspaceSchema(NamedSchema, table=True):
         back_populates="workspace",
         sa_relationship_kwargs={"cascade": "delete"},
     )
-    deployments: List["PipelineDeploymentSchema"] = Relationship(
+    actions: List["ActionSchema"] = Relationship(
         back_populates="workspace",
         sa_relationship_kwargs={"cascade": "delete"},
     )
-    code_repositories: List["CodeRepositorySchema"] = Relationship(
+    triggers: List["TriggerSchema"] = Relationship(
         back_populates="workspace",
         sa_relationship_kwargs={"cascade": "delete"},
     )
-    service_connectors: List["ServiceConnectorSchema"] = Relationship(
+    event_sources: List["EventSourceSchema"] = Relationship(
         back_populates="workspace",
         sa_relationship_kwargs={"cascade": "delete"},
     )
-    models: List["ModelSchema"] = Relationship(
+
+    deployments: List["PipelineDeploymentSchema"] = Relationship(
         back_populates="workspace",
         sa_relationship_kwargs={"cascade": "delete"},
     )
-    model_versions: List["ModelVersionSchema"] = Relationship(
+    code_repositories: List["CodeRepositorySchema"] = Relationship(
         back_populates="workspace",
         sa_relationship_kwargs={"cascade": "delete"},
     )
-    model_versions_artifacts_links: List[
-        "ModelVersionArtifactSchema"
-    ] = Relationship(
+    services: List["ServiceSchema"] = Relationship(
         back_populates="workspace",
         sa_relationship_kwargs={"cascade": "delete"},
     )
-    model_versions_pipeline_runs_links: List[
-        "ModelVersionPipelineRunSchema"
-    ] = Relationship(
+    service_connectors: List["ServiceConnectorSchema"] = Relationship(
+        back_populates="workspace",
+        sa_relationship_kwargs={"cascade": "delete"},
+    )
+    models: List["ModelSchema"] = Relationship(
+        back_populates="workspace",
+        sa_relationship_kwargs={"cascade": "delete"},
+    )
+    model_versions: List["ModelVersionSchema"] = Relationship(
         back_populates="workspace",
         sa_relationship_kwargs={"cascade": "delete"},
     )
@@ -155,24 +163,33 @@ class WorkspaceSchema(NamedSchema, table=True):
         Returns:
             The updated `WorkspaceSchema`.
         """
-        for field, value in workspace_update.dict(exclude_unset=True).items():
+        for field, value in workspace_update.model_dump(
+            exclude_unset=True
+        ).items():
             setattr(self, field, value)
 
         self.updated = datetime.utcnow()
         return self
 
-    def to_model(self, hydrate: bool = False) -> WorkspaceResponse:
+    def to_model(
+        self,
+        include_metadata: bool = False,
+        include_resources: bool = False,
+        **kwargs: Any,
+    ) -> WorkspaceResponse:
         """Convert a `WorkspaceSchema` to a `WorkspaceResponse`.
 
         Args:
-            hydrate: bool to decide whether to return a hydrated version of the
-                model.
+            include_metadata: Whether the metadata will be filled.
+            include_resources: Whether the resources will be filled.
+            **kwargs: Keyword arguments to allow schema specific logic
+
 
         Returns:
             The converted `WorkspaceResponseModel`.
         """
         metadata = None
-        if hydrate:
+        if include_metadata:
             metadata = WorkspaceResponseMetadata(
                 description=self.description,
             )

zenml/zen_stores/secrets_stores/aws_secrets_store.py

@@ -25,7 +25,7 @@ from uuid import UUID
 
 import boto3
 from botocore.exceptions import ClientError
-from pydantic import root_validator
+from pydantic import ConfigDict, model_validator
 
 from zenml.enums import (
     SecretsStoreType,
@@ -38,6 +38,7 @@ from zenml.integrations.aws.service_connectors.aws_service_connector import (
     AWSAuthenticationMethods,
 )
 from zenml.logger import get_logger
+from zenml.utils.pydantic_utils import before_validator_handler
 from zenml.zen_stores.secrets_stores.service_connector_secrets_store import (
     ServiceConnectorSecretsStore,
     ServiceConnectorSecretsStoreConfiguration,
@@ -74,43 +75,51 @@ class AWSSecretsStoreConfiguration(ServiceConnectorSecretsStoreConfiguration):
 
         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]:
+    @model_validator(mode="before")
+    @classmethod
+    @before_validator_handler
+    def populate_config(cls, data: Dict[str, Any]) -> Dict[str, Any]:
         """Populate the connector configuration from legacy attributes.
 
         Args:
-            values: Dict representing user-specified runtime settings.
+            data: Dict representing user-specified runtime settings.
 
         Returns:
             Validated settings.
         """
         # 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"),
-                aws_secret_access_key=values.get("aws_secret_access_key"),
-                region=values.get("region_name"),
-            )
-
-        return values
-
-    class Config:
-        """Pydantic configuration class."""
-
-        # Allow extra attributes set in the class.
-        extra = "allow"
+        if data.get("region_name"):
+            if not data.get("aws_access_key_id") or not data.get(
+                "aws_secret_access_key"
+            ):
+                logger.warning(
+                    "The `region_name` AWS secrets store attribute is deprecated "
+                    "and will be removed in a future version of ZenML. Please use "
+                    "the `auth_method` and `auth_config` attributes instead. "
+                    "Using an implicit authentication method for AWS Secrets."
+                )
+                data["auth_method"] = AWSAuthenticationMethods.IMPLICIT
+                data["auth_config"] = dict(
+                    region=data.get("region_name"),
+                )
+            else:
+                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."
+                )
+                data["auth_method"] = AWSAuthenticationMethods.SECRET_KEY
+                data["auth_config"] = dict(
+                    aws_access_key_id=data.get("aws_access_key_id"),
+                    aws_secret_access_key=data.get("aws_secret_access_key"),
+                    region=data.get("region_name"),
+                )
+
+        return data
+
+    model_config = ConfigDict(extra="allow")
 
 
 class AWSSecretsStore(ServiceConnectorSecretsStore):
@@ -138,9 +147,9 @@ class AWSSecretsStore(ServiceConnectorSecretsStore):
 
     config: AWSSecretsStoreConfiguration
     TYPE: ClassVar[SecretsStoreType] = SecretsStoreType.AWS
-    CONFIG_TYPE: ClassVar[
-        Type[ServiceConnectorSecretsStoreConfiguration]
-    ] = AWSSecretsStoreConfiguration
+    CONFIG_TYPE: ClassVar[Type[ServiceConnectorSecretsStoreConfiguration]] = (
+        AWSSecretsStoreConfiguration
+    )
     SERVICE_CONNECTOR_TYPE: ClassVar[str] = AWS_CONNECTOR_TYPE
     SERVICE_CONNECTOR_RESOURCE_TYPE: ClassVar[str] = AWS_RESOURCE_TYPE
 

zenml/zen_stores/secrets_stores/azure_secrets_store.py

@@ -27,7 +27,7 @@ from uuid import UUID
 from azure.core.credentials import TokenCredential
 from azure.core.exceptions import HttpResponseError, ResourceNotFoundError
 from azure.keyvault.secrets import SecretClient
-from pydantic import root_validator
+from pydantic import ConfigDict, model_validator
 
 from zenml.enums import (
     SecretsStoreType,
@@ -40,6 +40,7 @@ from zenml.integrations.azure.service_connectors.azure_service_connector import
     AzureAuthenticationMethods,
 )
 from zenml.logger import get_logger
+from zenml.utils.pydantic_utils import before_validator_handler
 from zenml.zen_stores.secrets_stores.service_connector_secrets_store import (
     ServiceConnectorSecretsStore,
     ServiceConnectorSecretsStoreConfiguration,
@@ -65,12 +66,14 @@ class AzureSecretsStoreConfiguration(
     type: SecretsStoreType = SecretsStoreType.AZURE
     key_vault_name: str
 
-    @root_validator(pre=True)
-    def populate_config(cls, values: Dict[str, Any]) -> Dict[str, Any]:
+    @model_validator(mode="before")
+    @classmethod
+    @before_validator_handler
+    def populate_config(cls, data: Dict[str, Any]) -> Dict[str, Any]:
         """Populate the connector configuration from legacy attributes.
 
         Args:
-            values: Dict representing user-specified runtime settings.
+            data: Dict representing user-specified runtime settings.
 
         Returns:
             Validated settings.
@@ -78,9 +81,9 @@ class AzureSecretsStoreConfiguration(
         # 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")
+            data.get("azure_client_id")
+            and data.get("azure_client_secret")
+            and data.get("azure_tenant_id")
         ):
             logger.warning(
                 "The `azure_client_id`, `azure_client_secret` and "
@@ -88,22 +91,16 @@ class AzureSecretsStoreConfiguration(
                 "removed in a future version or ZenML. Please use the "
                 "`auth_method` and `auth_config` attributes instead."
             )
-            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"),
+            data["auth_method"] = AzureAuthenticationMethods.SERVICE_PRINCIPAL
+            data["auth_config"] = dict(
+                client_id=data.get("azure_client_id"),
+                client_secret=data.get("azure_client_secret"),
+                tenant_id=data.get("azure_tenant_id"),
             )
 
-        return values
-
-    class Config:
-        """Pydantic configuration class."""
+        return data
 
-        # Forbid extra attributes set in the class.
-        extra = "allow"
+    model_config = ConfigDict(extra="allow")
 
 
 class AzureSecretsStore(ServiceConnectorSecretsStore):
@@ -129,9 +126,9 @@ class AzureSecretsStore(ServiceConnectorSecretsStore):
 
     config: AzureSecretsStoreConfiguration
     TYPE: ClassVar[SecretsStoreType] = SecretsStoreType.AZURE
-    CONFIG_TYPE: ClassVar[
-        Type[ServiceConnectorSecretsStoreConfiguration]
-    ] = AzureSecretsStoreConfiguration
+    CONFIG_TYPE: ClassVar[Type[ServiceConnectorSecretsStoreConfiguration]] = (
+        AzureSecretsStoreConfiguration
+    )
     SERVICE_CONNECTOR_TYPE: ClassVar[str] = AZURE_CONNECTOR_TYPE
     SERVICE_CONNECTOR_RESOURCE_TYPE: ClassVar[str] = AZURE_RESOURCE_TYPE
 

zenml/zen_stores/secrets_stores/base_secrets_store.py

@@ -12,6 +12,7 @@
 #  or implied. See the License for the specific language governing
 #  permissions and limitations under the License.
 """Base Secrets Store implementation."""
+
 from abc import ABC
 from typing import (
     TYPE_CHECKING,
@@ -23,12 +24,13 @@ from typing import (
 )
 from uuid import UUID
 
-from pydantic import BaseModel
+from pydantic import BaseModel, ConfigDict, model_validator
 
 from zenml.config.secrets_store_config import SecretsStoreConfiguration
 from zenml.enums import SecretsStoreType
 from zenml.logger import get_logger
 from zenml.utils import source_utils
+from zenml.utils.pydantic_utils import before_validator_handler
 from zenml.zen_stores.secrets_stores.secrets_store_interface import (
     SecretsStoreInterface,
 )
@@ -57,6 +59,79 @@ class BaseSecretsStore(BaseModel, SecretsStoreInterface, ABC):
     TYPE: ClassVar[SecretsStoreType]
     CONFIG_TYPE: ClassVar[Type[SecretsStoreConfiguration]]
 
+    @model_validator(mode="before")
+    @classmethod
+    @before_validator_handler
+    def convert_config(cls, data: Dict[str, Any]) -> Dict[str, Any]:
+        """Method to infer the correct type of the config and convert.
+
+        Args:
+            data: The provided configuration object, can potentially be a
+                generic object
+
+        Raises:
+            ValueError: If the provided config object's type does not match
+                any of the current implementations.
+
+        Returns:
+            The converted configuration object.
+        """
+        if data["config"].type == SecretsStoreType.SQL:
+            from zenml.zen_stores.secrets_stores.sql_secrets_store import (
+                SqlSecretsStoreConfiguration,
+            )
+
+            data["config"] = SqlSecretsStoreConfiguration(
+                **data["config"].model_dump()
+            )
+
+        elif data["config"].type == SecretsStoreType.GCP:
+            from zenml.zen_stores.secrets_stores.gcp_secrets_store import (
+                GCPSecretsStoreConfiguration,
+            )
+
+            data["config"] = GCPSecretsStoreConfiguration(
+                **data["config"].model_dump()
+            )
+
+        elif data["config"].type == SecretsStoreType.AWS:
+            from zenml.zen_stores.secrets_stores.aws_secrets_store import (
+                AWSSecretsStoreConfiguration,
+            )
+
+            data["config"] = AWSSecretsStoreConfiguration(
+                **data["config"].model_dump()
+            )
+
+        elif data["config"].type == SecretsStoreType.AZURE:
+            from zenml.zen_stores.secrets_stores.azure_secrets_store import (
+                AzureSecretsStoreConfiguration,
+            )
+
+            data["config"] = AzureSecretsStoreConfiguration(
+                **data["config"].model_dump()
+            )
+
+        elif data["config"].type == SecretsStoreType.HASHICORP:
+            from zenml.zen_stores.secrets_stores.hashicorp_secrets_store import (
+                HashiCorpVaultSecretsStoreConfiguration,
+            )
+
+            data["config"] = HashiCorpVaultSecretsStoreConfiguration(
+                **data["config"].model_dump()
+            )
+        elif (
+            data["config"].type == SecretsStoreType.CUSTOM
+            or data["config"].type == SecretsStoreType.NONE
+        ):
+            pass
+        else:
+            raise ValueError(
+                f"Unknown type '{data['config'].type}' for the configuration."
+            )
+
+        return data
+
     # ---------------------------------
     # Initialization and configuration
     # ---------------------------------
@@ -296,14 +371,7 @@ class BaseSecretsStore(BaseModel, SecretsStoreInterface, ABC):
                 f"expected {secret_id}, got {stored_secret_id}"
             )
 
-    class Config:
-        """Pydantic configuration class."""
-
-        # Validate attributes when assigning them. We need to set this in order
-        # to have a mix of mutable and immutable attributes
-        validate_assignment = True
-        # Ignore extra attributes from configs of previous ZenML versions
-        extra = "ignore"
-        # all attributes with leading underscore are private and therefore
-        # are mutable and not included in serialization
-        underscore_attrs_are_private = True
+    model_config = ConfigDict(
+        validate_assignment=True,
+        extra="ignore",
+    )

zenml/zen_stores/secrets_stores/gcp_secrets_store.py

@@ -13,7 +13,6 @@
 #  permissions and limitations under the License.
 """Implementation of the GCP Secrets Store."""
 
-
 import json
 import os
 from typing import (
@@ -28,7 +27,7 @@ from uuid import UUID
 
 from google.api_core import exceptions as google_exceptions
 from google.cloud.secretmanager import SecretManagerServiceClient
-from pydantic import root_validator
+from pydantic import ConfigDict, model_validator
 
 from zenml.enums import (
     SecretsStoreType,
@@ -41,6 +40,7 @@ from zenml.integrations.gcp.service_connectors.gcp_service_connector import (
     GCPAuthenticationMethods,
 )
 from zenml.logger import get_logger
+from zenml.utils.pydantic_utils import before_validator_handler
 from zenml.zen_stores.secrets_stores.service_connector_secrets_store import (
     ServiceConnectorSecretsStore,
     ServiceConnectorSecretsStoreConfiguration,
@@ -79,43 +79,52 @@ class GCPSecretsStoreConfiguration(ServiceConnectorSecretsStoreConfiguration):
 
         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]:
+    @model_validator(mode="before")
+    @classmethod
+    @before_validator_handler
+    def populate_config(cls, data: Dict[str, Any]) -> Dict[str, Any]:
         """Populate the connector configuration from legacy attributes.
 
         Args:
-            values: Dict representing user-specified runtime settings.
+            data: Dict representing user-specified runtime settings.
 
         Returns:
             Validated settings.
         """
         # 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"),
-            )
-            # 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 = "allow"
+        if data.get("project_id"):
+            if not os.environ.get("GOOGLE_APPLICATION_CREDENTIALS"):
+                logger.warning(
+                    "The `project_id` GCP secrets store attribute is "
+                    "deprecated and will be removed in a future version of ZenML. "
+                    "Please use the `auth_method` and `auth_config` attributes "
+                    "instead. Using an implicit GCP authentication to access "
+                    "the GCP Secrets Manager API."
+                )
+                data["auth_method"] = GCPAuthenticationMethods.IMPLICIT
+                data["auth_config"] = dict(
+                    project_id=data.get("project_id"),
+                )
+            else:
+                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."
+                )
+                data["auth_method"] = GCPAuthenticationMethods.SERVICE_ACCOUNT
+                data["auth_config"] = dict(
+                    project_id=data.get("project_id"),
+                )
+                # Load the service account credentials from the file
+                with open(os.environ["GOOGLE_APPLICATION_CREDENTIALS"]) as f:
+                    data["auth_config"]["service_account_json"] = f.read()
+
+        return data
+
+    model_config = ConfigDict(extra="allow")
 
 
 class GCPSecretsStore(ServiceConnectorSecretsStore):
@@ -123,9 +132,9 @@ class GCPSecretsStore(ServiceConnectorSecretsStore):
 
     config: GCPSecretsStoreConfiguration
     TYPE: ClassVar[SecretsStoreType] = SecretsStoreType.GCP
-    CONFIG_TYPE: ClassVar[
-        Type[ServiceConnectorSecretsStoreConfiguration]
-    ] = GCPSecretsStoreConfiguration
+    CONFIG_TYPE: ClassVar[Type[ServiceConnectorSecretsStoreConfiguration]] = (
+        GCPSecretsStoreConfiguration
+    )
     SERVICE_CONNECTOR_TYPE: ClassVar[str] = GCP_CONNECTOR_TYPE
     SERVICE_CONNECTOR_RESOURCE_TYPE: ClassVar[str] = GCP_RESOURCE_TYPE
 

zenml/zen_stores/secrets_stores/hashicorp_secrets_store.py

@@ -26,13 +26,14 @@ from hvac.exceptions import (  # type: ignore[import-untyped]
     InvalidPath,
     VaultError,
 )
-from pydantic import SecretStr
+from pydantic import ConfigDict
 
 from zenml.config.secrets_store_config import SecretsStoreConfiguration
 from zenml.enums import (
     SecretsStoreType,
 )
 from zenml.logger import get_logger
+from zenml.utils.secret_utils import PlainSerializedSecretStr
 from zenml.zen_stores.secrets_stores.base_secrets_store import (
     BaseSecretsStore,
 )
@@ -63,16 +64,11 @@ class HashiCorpVaultSecretsStoreConfiguration(SecretsStoreConfiguration):
     type: SecretsStoreType = SecretsStoreType.HASHICORP
 
     vault_addr: str
-    vault_token: Optional[SecretStr] = None
+    vault_token: Optional[PlainSerializedSecretStr] = None
     vault_namespace: Optional[str] = None
     mount_point: Optional[str] = None
     max_versions: int = 1
-
-    class Config:
-        """Pydantic configuration class."""
-
-        # Forbid extra attributes set in the class.
-        extra = "forbid"
+    model_config = ConfigDict(extra="forbid")
 
 
 class HashiCorpVaultSecretsStore(BaseSecretsStore):
@@ -101,9 +97,9 @@ class HashiCorpVaultSecretsStore(BaseSecretsStore):
 
     config: HashiCorpVaultSecretsStoreConfiguration
     TYPE: ClassVar[SecretsStoreType] = SecretsStoreType.HASHICORP
-    CONFIG_TYPE: ClassVar[
-        Type[SecretsStoreConfiguration]
-    ] = HashiCorpVaultSecretsStoreConfiguration
+    CONFIG_TYPE: ClassVar[Type[SecretsStoreConfiguration]] = (
+        HashiCorpVaultSecretsStoreConfiguration
+    )
 
     _client: Optional[hvac.Client] = None
 

zenml/zen_stores/secrets_stores/secrets_store_interface.py

@@ -12,6 +12,7 @@
 #  or implied. See the License for the specific language governing
 #  permissions and limitations under the License.
 """ZenML secrets store interface."""
+
 from abc import ABC, abstractmethod
 from typing import Dict
 from uuid import UUID