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

zenml/config/server_config.py

@@ -13,24 +13,42 @@
 #  permissions and limitations under the License.
 """Functionality to support ZenML GlobalConfiguration."""
 
+import json
 import os
 from secrets import token_hex
-from typing import Any, Dict, List, Optional
+from typing import Any, Dict, List, Optional, Union
 from uuid import UUID
 
-from pydantic import BaseModel, Field, SecretStr, root_validator
+from pydantic import BaseModel, ConfigDict, Field, PositiveInt, model_validator
 
 from zenml.constants import (
     DEFAULT_ZENML_JWT_TOKEN_ALGORITHM,
     DEFAULT_ZENML_JWT_TOKEN_LEEWAY,
     DEFAULT_ZENML_SERVER_DEVICE_AUTH_POLLING,
     DEFAULT_ZENML_SERVER_DEVICE_AUTH_TIMEOUT,
+    DEFAULT_ZENML_SERVER_GENERIC_API_TOKEN_LIFETIME,
+    DEFAULT_ZENML_SERVER_GENERIC_API_TOKEN_MAX_LIFETIME,
+    DEFAULT_ZENML_SERVER_LOGIN_RATE_LIMIT_DAY,
+    DEFAULT_ZENML_SERVER_LOGIN_RATE_LIMIT_MINUTE,
     DEFAULT_ZENML_SERVER_MAX_DEVICE_AUTH_ATTEMPTS,
+    DEFAULT_ZENML_SERVER_MAX_REQUEST_BODY_SIZE_IN_BYTES,
+    DEFAULT_ZENML_SERVER_NAME,
+    DEFAULT_ZENML_SERVER_PIPELINE_RUN_AUTH_WINDOW,
+    DEFAULT_ZENML_SERVER_SECURE_HEADERS_CACHE,
+    DEFAULT_ZENML_SERVER_SECURE_HEADERS_CONTENT,
+    DEFAULT_ZENML_SERVER_SECURE_HEADERS_CSP,
+    DEFAULT_ZENML_SERVER_SECURE_HEADERS_HSTS,
+    DEFAULT_ZENML_SERVER_SECURE_HEADERS_PERMISSIONS,
+    DEFAULT_ZENML_SERVER_SECURE_HEADERS_REFERRER,
+    DEFAULT_ZENML_SERVER_SECURE_HEADERS_XFO,
+    DEFAULT_ZENML_SERVER_SECURE_HEADERS_XXP,
+    DEFAULT_ZENML_SERVER_THREAD_POOL_SIZE,
     ENV_ZENML_SERVER_PREFIX,
 )
 from zenml.enums import AuthScheme
 from zenml.logger import get_logger
 from zenml.models import ServerDeploymentType
+from zenml.utils.pydantic_utils import before_validator_handler
 
 logger = get_logger(__name__)
 
@@ -49,9 +67,21 @@ def generate_jwt_secret_key() -> str:
 class ServerConfiguration(BaseModel):
     """ZenML Server configuration attributes.
 
+    All these attributes can be set through the environment with the `ZENML_SERVER_`-Prefix.
+    The value of the `ZENML_SERVER_DEPLOYMENT_TYPE` environment variable will be extracted to deployment_type.
+
     Attributes:
         deployment_type: The type of ZenML server deployment that is running.
-        root_url_path: The root URL path of the ZenML server.
+        server_url: The URL where the ZenML server API is reachable. Must be
+            configured for features that involve triggering workloads from the
+            ZenML dashboard (e.g., running pipelines). If not specified, the
+            clients will use the same URL used to connect them to the ZenML
+            server.
+        dashboard_url: The URL where the ZenML dashboard is reachable.
+            If not specified, the `server_url` value is used. This should be
+            configured if the dashboard is served from a different URL than the
+            ZenML server.
+        root_url_path: The root URL path for the ZenML API and dashboard.
         auth_scheme: The authentication scheme used by the ZenML server.
         jwt_token_algorithm: The algorithm used to sign and verify JWT tokens.
         jwt_token_issuer: The issuer of the JWT tokens. If not specified, the
@@ -78,23 +108,21 @@ class ServerConfiguration(BaseModel):
             2.0 device authorization request expires.
         device_auth_polling_interval: The polling interval in seconds used to
             poll the OAuth 2.0 device authorization endpoint.
-        dashboard_url: The URL where the ZenML dashboard is hosted. Used to
-            construct the OAuth 2.0 device authorization endpoint. If not set,
-            a partial URL is returned to the client which is used to construct
-            the full URL based on the server's root URL path.
-        device_expiration: The time in minutes that an OAuth 2.0 device is
+        device_expiration_minutes: The time in minutes that an OAuth 2.0 device is
             allowed to be used to authenticate with the ZenML server. If not
             set or if `jwt_token_expire_minutes` is not set, the devices are
             allowed to be used indefinitely. This controls the expiration time
             of the JWT tokens issued to clients after they have authenticated
             with the ZenML server using an OAuth 2.0 device.
-        trusted_device_expiration: The time in minutes that a trusted OAuth 2.0
+        trusted_device_expiration_minutes: The time in minutes that a trusted OAuth 2.0
             device is allowed to be used to authenticate with the ZenML server.
             If not set or if `jwt_token_expire_minutes` is not set, the devices
             are allowed to be used indefinitely. This controls the expiration
             time of the JWT tokens issued to clients after they have
             authenticated with the ZenML server using an OAuth 2.0 device
             that has been marked as trusted.
+        generic_api_token_lifetime: The lifetime in seconds that generic
+            short-lived API tokens issued for automation purposes are valid.
         external_login_url: The login URL of an external authenticator service
             to use with the `EXTERNAL` authentication scheme.
         external_user_info_url: The user info URL of an external authenticator
@@ -106,14 +134,119 @@ class ServerConfiguration(BaseModel):
         external_server_id: The ID of the ZenML server to use with the
             `EXTERNAL` authentication scheme. If not specified, the regular
             ZenML server ID is used.
+        metadata: Additional metadata to be associated with the ZenML server.
         rbac_implementation_source: Source pointing to a class implementing
             the RBAC interface defined by
             `zenml.zen_server.rbac_interface.RBACInterface`. If not specified,
             RBAC will not be enabled for this server.
+        feature_gate_implementation_source: Source pointing to a class
+            implementing the feature gate interface defined by
+            `zenml.zen_server.feature_gate.feature_gate_interface.FeatureGateInterface`.
+            If not specified, feature usage will not be gated/tracked for this
+            server.
+        workload_manager_implementation_source: Source pointing to a class
+            implementing the workload management interface.
+        pipeline_run_auth_window: The default time window in minutes for which
+            a pipeline run action is allowed to authenticate with the ZenML
+            server.
+        login_rate_limit_minute: The number of login attempts allowed per minute.
+        login_rate_limit_day: The number of login attempts allowed per day.
+        secure_headers_server: Custom value to be set in the `Server` HTTP
+            header to identify the server. If not specified, or if set to one of
+            the reserved values `enabled`, `yes`, `true`, `on`, the `Server`
+            header will be set to the default value (ZenML server ID). If set to
+            one of the reserved values `disabled`, `no`, `none`, `false`, `off`
+            or to an empty string, the `Server` header will not be included in
+            responses.
+        secure_headers_hsts: The server header value to be set in the HTTP
+            header `Strict-Transport-Security`. If not specified, or if set to
+            one of the reserved values `enabled`, `yes`, `true`, `on`, the
+            `Strict-Transport-Security` header will be set to the default value
+            (`max-age=63072000; includeSubdomains`). If set to one of
+            the reserved values `disabled`, `no`, `none`, `false`, `off` or to
+            an empty string, the `Strict-Transport-Security` header will not be
+            included in responses.
+        secure_headers_xfo: The server header value to be set in the HTTP
+            header `X-Frame-Options`. If not specified, or if set to one of the
+            reserved values `enabled`, `yes`, `true`, `on`, the `X-Frame-Options`
+            header will be set to the default value (`SAMEORIGIN`). If set to
+            one of the reserved values `disabled`, `no`, `none`, `false`, `off`
+            or to an empty string, the `X-Frame-Options` header will not be
+            included in responses.
+        secure_headers_xxp: The server header value to be set in the HTTP
+            header `X-XSS-Protection`. If not specified, or if set to one of the
+            reserved values `enabled`, `yes`, `true`, `on`, the `X-XSS-Protection`
+            header will be set to the default value (`0`). If set to one of the
+            reserved values `disabled`, `no`, `none`, `false`, `off` or
+            to an empty string, the `X-XSS-Protection` header will not be
+            included in responses. NOTE: this header is deprecated and should
+            always be set to `0`. The `Content-Security-Policy` header should be
+            used instead.
+        secure_headers_content: The server header value to be set in the HTTP
+            header `X-Content-Type-Options`. If not specified, or if set to one
+            of the reserved values `enabled`, `yes`, `true`, `on`, the
+            `X-Content-Type-Options` header will be set to the default value
+            (`nosniff`). If set to one of the reserved values `disabled`, `no`,
+            `none`, `false`, `off` or to an empty string, the
+            `X-Content-Type-Options` header will not be included in responses.
+        secure_headers_csp: The server header value to be set in the HTTP
+            header `Content-Security-Policy`. If not specified, or if set to one
+            of the reserved values `enabled`, `yes`, `true`, `on`, the
+            `Content-Security-Policy` header will be set to a default value
+            that is compatible with the ZenML dashboard. If set to one of the
+            reserved values `disabled`, `no`, `none`, `false`, `off` or to an
+            empty string, the `Content-Security-Policy` header will not be
+            included in responses.
+        secure_headers_referrer: The server header value to be set in the HTTP
+            header `Referrer-Policy`. If not specified, or if set to one of the
+            reserved values `enabled`, `yes`, `true`, `on`, the `Referrer-Policy`
+            header will be set to the default value
+            (`no-referrer-when-downgrade`). If set to one of the reserved values
+            `disabled`, `no`, `none`, `false`, `off` or to an empty string, the
+            `Referrer-Policy` header will not be included in responses.
+        secure_headers_cache: The server header value to be set in the HTTP
+            header `Cache-Control`. If not specified, or if set to one of the
+            reserved values `enabled`, `yes`, `true`, `on`, the `Cache-Control`
+            header will be set to the default value
+            (`no-store, no-cache, must-revalidate`). If set to one of the
+            reserved values `disabled`, `no`, `none`, `false`, `off` or to an
+            empty string, the `Cache-Control` header will not be included in
+            responses.
+        secure_headers_permissions: The server header value to be set in the
+            HTTP header `Permissions-Policy`. If not specified, or if set to one
+            of the reserved values `enabled`, `yes`, `true`, `on`, the
+            `Permissions-Policy` header will be set to the default value
+            (`accelerometer=(), camera=(), geolocation=(), gyroscope=(),
+            magnetometer=(), microphone=(), payment=(), usb=()`). If set to
+            one of the reserved values `disabled`, `no`, `none`, `false`, `off`
+            or to an empty string, the `Permissions-Policy` header will not be
+            included in responses.
+        server_name: The name of the ZenML server. Used only during initial
+            deployment. Can be changed later as a part of the server settings.
+        display_announcements: Whether to display announcements about ZenML in
+            the dashboard. Used only during initial deployment. Can be changed
+            later as a part of the server settings.
+        display_updates: Whether to display notifications about ZenML updates in
+            the dashboard. Used only during initial deployment. Can be changed
+            later as a part of the server settings.
+        auto_activate: Whether to automatically activate the server and create a
+            default admin user account with an empty password during the initial
+            deployment.
+        max_request_body_size_in_bytes: The maximum size of the request body in
+            bytes. If not specified, the default value of 256 Kb will be used.
+        memcache_max_capacity: The maximum number of entries that the memory
+            cache can hold. If not specified, the default value of 1000 will be
+            used.
+        memcache_default_expiry: The default expiry time in seconds for cache
+            entries. If not specified, the default value of 30 seconds will be
+            used.
     """
 
     deployment_type: ServerDeploymentType = ServerDeploymentType.OTHER
+    server_url: Optional[str] = None
+    dashboard_url: Optional[str] = None
     root_url_path: str = ""
+    metadata: Dict[str, Any] = {}
     auth_scheme: AuthScheme = AuthScheme.OAUTH2_PASSWORD_BEARER
     jwt_token_algorithm: str = DEFAULT_ZENML_JWT_TOKEN_ALGORITHM
     jwt_token_issuer: Optional[str] = None
@@ -124,32 +257,100 @@ class ServerConfiguration(BaseModel):
     auth_cookie_name: Optional[str] = None
     auth_cookie_domain: Optional[str] = None
     cors_allow_origins: Optional[List[str]] = None
-    max_failed_device_auth_attempts: (
-        int
-    ) = DEFAULT_ZENML_SERVER_MAX_DEVICE_AUTH_ATTEMPTS
+    max_failed_device_auth_attempts: int = (
+        DEFAULT_ZENML_SERVER_MAX_DEVICE_AUTH_ATTEMPTS
+    )
     device_auth_timeout: int = DEFAULT_ZENML_SERVER_DEVICE_AUTH_TIMEOUT
-    device_auth_polling_interval: (
-        int
-    ) = DEFAULT_ZENML_SERVER_DEVICE_AUTH_POLLING
-    dashboard_url: Optional[str] = None
+    device_auth_polling_interval: int = (
+        DEFAULT_ZENML_SERVER_DEVICE_AUTH_POLLING
+    )
     device_expiration_minutes: Optional[int] = None
     trusted_device_expiration_minutes: Optional[int] = None
 
+    generic_api_token_lifetime: PositiveInt = (
+        DEFAULT_ZENML_SERVER_GENERIC_API_TOKEN_LIFETIME
+    )
+    generic_api_token_max_lifetime: PositiveInt = (
+        DEFAULT_ZENML_SERVER_GENERIC_API_TOKEN_MAX_LIFETIME
+    )
+
     external_login_url: Optional[str] = None
     external_user_info_url: Optional[str] = None
     external_cookie_name: Optional[str] = None
     external_server_id: Optional[UUID] = None
 
     rbac_implementation_source: Optional[str] = None
+    feature_gate_implementation_source: Optional[str] = None
+    workload_manager_implementation_source: Optional[str] = None
+    pipeline_run_auth_window: int = (
+        DEFAULT_ZENML_SERVER_PIPELINE_RUN_AUTH_WINDOW
+    )
+
+    rate_limit_enabled: bool = False
+    login_rate_limit_minute: int = DEFAULT_ZENML_SERVER_LOGIN_RATE_LIMIT_MINUTE
+    login_rate_limit_day: int = DEFAULT_ZENML_SERVER_LOGIN_RATE_LIMIT_DAY
+
+    secure_headers_server: Union[bool, str] = Field(
+        default=True,
+        union_mode="left_to_right",
+    )
+    secure_headers_hsts: Union[bool, str] = Field(
+        default=DEFAULT_ZENML_SERVER_SECURE_HEADERS_HSTS,
+        union_mode="left_to_right",
+    )
+    secure_headers_xfo: Union[bool, str] = Field(
+        default=DEFAULT_ZENML_SERVER_SECURE_HEADERS_XFO,
+        union_mode="left_to_right",
+    )
+    secure_headers_xxp: Union[bool, str] = Field(
+        default=DEFAULT_ZENML_SERVER_SECURE_HEADERS_XXP,
+        union_mode="left_to_right",
+    )
+    secure_headers_content: Union[bool, str] = Field(
+        default=DEFAULT_ZENML_SERVER_SECURE_HEADERS_CONTENT,
+        union_mode="left_to_right",
+    )
+    secure_headers_csp: Union[bool, str] = Field(
+        default=DEFAULT_ZENML_SERVER_SECURE_HEADERS_CSP,
+        union_mode="left_to_right",
+    )
+    secure_headers_referrer: Union[bool, str] = Field(
+        default=DEFAULT_ZENML_SERVER_SECURE_HEADERS_REFERRER,
+        union_mode="left_to_right",
+    )
+    secure_headers_cache: Union[bool, str] = Field(
+        default=DEFAULT_ZENML_SERVER_SECURE_HEADERS_CACHE,
+        union_mode="left_to_right",
+    )
+    secure_headers_permissions: Union[bool, str] = Field(
+        default=DEFAULT_ZENML_SERVER_SECURE_HEADERS_PERMISSIONS,
+        union_mode="left_to_right",
+    )
+
+    server_name: str = DEFAULT_ZENML_SERVER_NAME
+    display_announcements: bool = True
+    display_updates: bool = True
+    auto_activate: bool = False
+
+    thread_pool_size: int = DEFAULT_ZENML_SERVER_THREAD_POOL_SIZE
+
+    max_request_body_size_in_bytes: int = (
+        DEFAULT_ZENML_SERVER_MAX_REQUEST_BODY_SIZE_IN_BYTES
+    )
+
+    memcache_max_capacity: int = 1000
+    memcache_default_expiry: int = 30
 
     _deployment_id: Optional[UUID] = None
 
-    @root_validator(pre=True)
-    def _validate_config(cls, values: Dict[str, Any]) -> Dict[str, Any]:
+    @model_validator(mode="before")
+    @classmethod
+    @before_validator_handler
+    def _validate_config(cls, data: Dict[str, Any]) -> Dict[str, Any]:
         """Validate the server configuration.
 
         Args:
-            values: The server configuration values.
+            data: The server configuration values.
 
         Returns:
             The validated server configuration values.
@@ -157,10 +358,10 @@ class ServerConfiguration(BaseModel):
         Raises:
             ValueError: If the server configuration is invalid.
         """
-        if values.get("auth_scheme") == AuthScheme.EXTERNAL:
+        if data.get("auth_scheme") == AuthScheme.EXTERNAL:
             # If the authentication scheme is set to `EXTERNAL`, the
             # external authenticator URLs must be specified.
-            if not values.get("external_login_url") or not values.get(
+            if not data.get("external_login_url") or not data.get(
                 "external_user_info_url"
             ):
                 raise ValueError(
@@ -171,19 +372,68 @@ class ServerConfiguration(BaseModel):
 
             # If the authentication scheme is set to `EXTERNAL`, the
             # external cookie name must be specified.
-            if not values.get("external_cookie_name"):
+            if not data.get("external_cookie_name"):
                 raise ValueError(
                     "The external cookie name must be specified when "
                     "using the EXTERNAL authentication scheme."
                 )
 
-        if cors_allow_origins := values.get("cors_allow_origins"):
+        if cors_allow_origins := data.get("cors_allow_origins"):
             origins = cors_allow_origins.split(",")
-            values["cors_allow_origins"] = origins
+            data["cors_allow_origins"] = origins
         else:
-            values["cors_allow_origins"] = ["*"]
+            data["cors_allow_origins"] = ["*"]
 
-        return values
+        # if metadata is a string, convert it to a dictionary
+        if isinstance(data.get("metadata"), str):
+            try:
+                data["metadata"] = json.loads(data["metadata"])
+            except json.JSONDecodeError as e:
+                raise ValueError(
+                    f"The server metadata is not a valid JSON string: {e}"
+                )
+
+        # if one of the secure headers options is set to a boolean value, set
+        # the corresponding value
+        for k, v in data.copy().items():
+            if k.startswith("secure_headers_") and isinstance(v, str):
+                if v.lower() in ["disabled", "no", "none", "false", "off", ""]:
+                    data[k] = False
+                if v.lower() in ["enabled", "yes", "true", "on"]:
+                    # Revert to the default value if the header is enabled
+                    del data[k]
+
+        # Handle merging of user-defined secure_headers_csp value with the default value
+        if "secure_headers_csp" in data:
+            user_defined_csp = data["secure_headers_csp"]
+            if isinstance(user_defined_csp, str):
+                # Parse the user-defined CSP string into a dictionary
+                user_defined_csp_dict = {}
+                for directive in user_defined_csp.split(";"):
+                    directive = directive.strip()
+                    if directive:
+                        key, value = directive.split(" ", 1)
+                        user_defined_csp_dict[key] = value.strip("'\"")
+
+                # Merge the user-defined CSP dictionary with the default CSP dictionary
+                default_csp_dict = {}
+                for directive in DEFAULT_ZENML_SERVER_SECURE_HEADERS_CSP.split(
+                    ";"
+                ):
+                    directive = directive.strip()
+                    if directive:
+                        key, value = directive.split(" ", 1)
+                        default_csp_dict[key] = value.strip("'\"")
+
+                merged_csp_dict = {**default_csp_dict, **user_defined_csp_dict}
+
+                # Convert the merged CSP dictionary back to a string
+                merged_csp_str = "; ".join(
+                    f"{key} {value}" for key, value in merged_csp_dict.items()
+                )
+                data["secure_headers_csp"] = merged_csp_str
+
+        return data
 
     @property
     def deployment_id(self) -> UUID:
@@ -212,6 +462,24 @@ class ServerConfiguration(BaseModel):
         """
         return self.rbac_implementation_source is not None
 
+    @property
+    def feature_gate_enabled(self) -> bool:
+        """Whether feature gating is enabled on the server or not.
+
+        Returns:
+            Whether feature gating is enabled on the server or not.
+        """
+        return self.feature_gate_implementation_source is not None
+
+    @property
+    def workload_manager_enabled(self) -> bool:
+        """Whether workload management is enabled on the server or not.
+
+        Returns:
+            Whether workload management is enabled on the server or not.
+        """
+        return self.workload_manager_implementation_source is not None
+
     def get_jwt_token_issuer(self) -> str:
         """Get the JWT token issuer.
 
@@ -291,18 +559,8 @@ class ServerConfiguration(BaseModel):
 
         return ServerConfiguration(**env_server_config)
 
-    class Config:
-        """Pydantic configuration class."""
-
+    model_config = ConfigDict(
         # Allow extra attributes from configs of previous ZenML versions to
         # permit downgrading
-        extra = "allow"
-        # all attributes with leading underscore are private and therefore
-        # are mutable and not included in serialization
-        underscore_attrs_are_private = True
-
-        # This is needed to allow correct handling of SecretStr values during
-        # serialization.
-        json_encoders = {
-            SecretStr: lambda v: v.get_secret_value() if v else None
-        }
+        extra="allow",
+    )

zenml/config/settings_resolver.py

@@ -12,6 +12,7 @@
 #  or implied. See the License for the specific language governing
 #  permissions and limitations under the License.
 """Class for resolving settings."""
+
 from typing import TYPE_CHECKING, Type, TypeVar
 
 from pydantic import ValidationError
@@ -121,7 +122,7 @@ class SettingsResolver:
         Returns:
             The converted settings.
         """
-        settings_dict = self._settings.dict()
+        settings_dict = dict(self._settings)
         try:
             return target_class(**settings_dict)
         except ValidationError:

zenml/config/source.py

@@ -12,11 +12,19 @@
 #  or implied. See the License for the specific language governing
 #  permissions and limitations under the License.
 """Source classes."""
+
 from enum import Enum
-from typing import TYPE_CHECKING, Any, Optional, Type, Union
+from typing import TYPE_CHECKING, Any, Dict, Optional
 from uuid import UUID
 
-from pydantic import BaseModel, Extra, validator
+from pydantic import (
+    BaseModel,
+    BeforeValidator,
+    ConfigDict,
+    SerializeAsAny,
+    field_validator,
+)
+from typing_extensions import Annotated
 
 from zenml.logger import get_logger
 
@@ -34,6 +42,7 @@ class SourceType(Enum):
     INTERNAL = "internal"
     DISTRIBUTION_PACKAGE = "distribution_package"
     CODE_REPOSITORY = "code_repository"
+    NOTEBOOK = "notebook"
     UNKNOWN = "unknown"
 
 
@@ -129,10 +138,29 @@ class Source(BaseModel):
         """
         return self.attribute is None
 
-    class Config:
-        """Pydantic config class."""
+    model_config = ConfigDict(extra="allow")
+
+    def model_dump(self, **kwargs: Any) -> Dict[str, Any]:
+        """Dump the source as a dictionary.
+
+        Args:
+            **kwargs: Additional keyword arguments.
+
+        Returns:
+            The source as a dictionary.
+        """
+        return super().model_dump(serialize_as_any=True, **kwargs)
+
+    def model_dump_json(self, **kwargs: Any) -> str:
+        """Dump the source as a JSON string.
+
+        Args:
+            **kwargs: Additional keyword arguments.
 
-        extra = Extra.allow
+        Returns:
+            The source as a JSON string.
+        """
+        return super().model_dump_json(serialize_as_any=True, **kwargs)
 
 
 class DistributionPackageSource(Source):
@@ -147,7 +175,8 @@ class DistributionPackageSource(Source):
     version: Optional[str] = None
     type: SourceType = SourceType.DISTRIBUTION_PACKAGE
 
-    @validator("type")
+    @field_validator("type")
+    @classmethod
     def _validate_type(cls, value: SourceType) -> SourceType:
         """Validate the source type.
 
@@ -181,7 +210,8 @@ class CodeRepositorySource(Source):
     subdirectory: str
     type: SourceType = SourceType.CODE_REPOSITORY
 
-    @validator("type")
+    @field_validator("type")
+    @classmethod
     def _validate_type(cls, value: SourceType) -> SourceType:
         """Validate the source type.
 
@@ -200,39 +230,75 @@ class CodeRepositorySource(Source):
         return value
 
 
-def convert_source_validator(*attributes: str) -> "AnyClassMethod":
-    """Function to convert pydantic fields containing legacy class paths.
+class NotebookSource(Source):
+    """Source representing an object defined in a notebook.
 
-    In older versions, sources (sometimes also called class paths) like
-    `zenml.materializers.BuiltInMaterializer` were stored as strings in our
-    configuration classes. These strings got replaced by a separate class, and
-    this function returns a validator to convert those old strings to the new
-    classes.
+    Attributes:
+        replacement_module: Name of the module from which this source should
+            be loaded in case the code is not running in a notebook.
+        artifact_store_id: ID of the artifact store in which the replacement
+            module code is stored.
+    """
 
-    Args:
-        *attributes: List of attributes to convert.
+    replacement_module: Optional[str] = None
+    artifact_store_id: Optional[UUID] = None
+    type: SourceType = SourceType.NOTEBOOK
 
-    Returns:
-        Pydantic validator class method to be used on BaseModel subclasses
-        to convert source fields.
-    """
+    @field_validator("type")
+    @classmethod
+    def _validate_type(cls, value: SourceType) -> SourceType:
+        """Validate the source type.
+
+        Args:
+            value: The source type.
+
+        Raises:
+            ValueError: If the source type is not `NOTEBOOK`.
+
+        Returns:
+            The source type.
+        """
+        if value != SourceType.NOTEBOOK:
+            raise ValueError("Invalid source type.")
 
-    @validator(*attributes, pre=True, allow_reuse=True)
-    def _convert_source(
-        cls: Type[BaseModel], value: Union[Source, str, None]
-    ) -> Optional[Source]:
-        """Converts an old source string to a source object.
+        return value
+
+    @field_validator("module")
+    @classmethod
+    def _validate_module(cls, value: str) -> str:
+        """Validate the module.
 
         Args:
-            cls: The class on which the attributes are defined.
-            value: Source string or object.
+            value: The module.
+
+        Raises:
+            ValueError: If the module is not `__main__`.
 
         Returns:
-            The converted source.
+            The module.
         """
-        if isinstance(value, str):
-            value = Source.from_import_path(value)
+        if value != "__main__":
+            raise ValueError("Invalid module for notebook source.")
 
         return value
 
-    return _convert_source
+
+def convert_source(source: Any) -> Any:
+    """Converts an old source string to a source object.
+
+    Args:
+        source: Source string or object.
+
+    Returns:
+        The converted source.
+    """
+    if isinstance(source, str):
+        source = Source.from_import_path(source)
+
+    return source
+
+
+SourceWithValidator = Annotated[
+    SerializeAsAny[Source],
+    BeforeValidator(convert_source),
+]