Migrate Pydantic v1 -> v2

bzl/mypy/mypy.ini

@@ -67,6 +67,9 @@ ignore_missing_imports = True
 [mypy-packaging.*]
 ignore_missing_imports = True
 
+[mypy-pydantic_settings.*]
+ignore_missing_imports = True
+
 [mypy-prometheus_client.*]
 ignore_missing_imports = True
 

docs/docs_requirements.txt

@@ -20,9 +20,8 @@
 # pip-compile ./docs_requirements.txt -o locked_requirements.txt --allow-unsafe -v --generate-hashes
 
 # Building documentation
-# Pin pydantic to match src/requirements.txt
-pydantic==1.10.26
-autodoc_pydantic==1.9.1
+pydantic==2.12.5
+autodoc_pydantic==2.2.0
 commonmark==0.9.1
 nvidia-sphinx-theme==0.0.8
 sphinx==7.4.7

src/cli/login.py

@@ -86,7 +86,7 @@ class UrlValidator(pydantic.BaseModel):
         url: pydantic.AnyHttpUrl
     try:
         _ = UrlValidator(url=url)
-    except pydantic.error_wrappers.ValidationError as error:
+    except pydantic.ValidationError as error:
         raise osmo_errors.OSMOUserError(f'Bad url {url}: {error}')
 
     print(f'Logging in to {url}')

src/cli/version.py

@@ -50,7 +50,7 @@ def _client_version(service_client: client.ServiceClient, args: argparse.Namespa
         pass
     client_version = version.VERSION
     if args.format_type == 'json':
-        output = {'client': client_version.dict()}
+        output = {'client': client_version.model_dump()}
         if result:
             output['service'] = result
         print(json.dumps(output, indent=common.JSON_INDENT_SIZE))

src/cli/workflow.py

@@ -54,7 +54,7 @@
 RESIZE_PREFIX = b'\x00RESIZE:'
 
 
-class TemplateData(pydantic.BaseModel, extra=pydantic.Extra.forbid):
+class TemplateData(pydantic.BaseModel, extra='forbid'):
     """Pydantic model representing parsed template data from workflow files."""
     file: str
     set_variables: List[str]

src/lib/data/dataset/common.py

@@ -196,7 +196,7 @@ class RemoteToRemoteMapping(NamedTuple):
 
 @pydantic.dataclasses.dataclass(
     config=pydantic.ConfigDict(
-        extra=pydantic.Extra.forbid,
+        extra='forbid',
         frozen=True,
     ),
 )
@@ -211,7 +211,7 @@ class UploadStartResult:
 
 @pydantic.dataclasses.dataclass(
     config=pydantic.ConfigDict(
-        extra=pydantic.Extra.forbid,
+        extra='forbid',
         frozen=True,
     ),
 )
@@ -225,7 +225,7 @@ class UploadResult:
 
 @pydantic.dataclasses.dataclass(
     config=pydantic.ConfigDict(
-        extra=pydantic.Extra.forbid,
+        extra='forbid',
         frozen=True,
     ),
 )
@@ -235,14 +235,14 @@ class UpdateStartResult:
     """
     upload_response: UploadResponse
     current_manifest_path: str
-    local_update_paths: List[LocalToRemoteMapping] | None
-    backend_update_paths: List[RemoteToRemoteMapping] | None
-    remove_regex: str | None
+    local_update_paths: List[LocalToRemoteMapping] | None = None
+    backend_update_paths: List[RemoteToRemoteMapping] | None = None
+    remove_regex: str | None = None
 
 
 @pydantic.dataclasses.dataclass(
     config=pydantic.ConfigDict(
-        extra=pydantic.Extra.forbid,
+        extra='forbid',
         frozen=True,
     ),
 )

src/lib/data/dataset/manager.py

@@ -43,11 +43,12 @@ class Manager(pydantic.BaseModel):
     Manager for a dataset.
     """
 
-    class Config:
-        arbitrary_types_allowed = True
-        extra = pydantic.Extra.forbid
-        frozen = True
-        keep_untouched = (functools.cached_property,)
+    model_config = pydantic.ConfigDict(
+        arbitrary_types_allowed=True,
+        extra='forbid',
+        frozen=True,
+        ignored_types=(functools.cached_property,),
+    )
 
     #########################
     #    Required fields    #

src/lib/data/storage/backends/backends.py

@@ -229,9 +229,8 @@ class SwiftBackend(Boto3Backend):
     Swift Backend
     """
 
-    scheme: str = pydantic.Field(
+    scheme: Literal['swift'] = pydantic.Field(
         default=common.StorageBackendType.SWIFT.value,
-        const=True,
         description='The scheme of the Swift backend.',
     )
 
@@ -242,7 +241,6 @@ class SwiftBackend(Boto3Backend):
 
     supports_batch_delete: Literal[True] = pydantic.Field(
         default=True,
-        const=True,
         description='Whether the backend supports batch delete.',
     )
 
@@ -390,21 +388,18 @@ class S3Backend(Boto3Backend):
     AWS S3 Backend
     """
 
-    scheme: str = pydantic.Field(
+    scheme: Literal['s3'] = pydantic.Field(
         default=common.StorageBackendType.S3.value,
-        const=True,
         description='The scheme of the S3 backend.',
     )
 
     supports_batch_delete: Literal[True] = pydantic.Field(
         default=True,
-        const=True,
         description='Whether the backend supports batch delete.',
     )
 
     supports_environment_auth: Literal[True] = pydantic.Field(
         default=True,
-        const=True,
         description='Whether the backend supports environment authentication.',
     )
 
@@ -585,17 +580,15 @@ class GSBackend(Boto3Backend):
     Google Cloud Platform GS Backend
     """
 
-    scheme: str = pydantic.Field(
+    scheme: Literal['gs'] = pydantic.Field(
         default=common.StorageBackendType.GS.value,
-        const=True,
         description='The scheme of the GS backend.',
     )
 
     # Google Cloud Storage does not support batch delete via S3 API:
     # https://issuetracker.google.com/issues/162653700
     supports_batch_delete: Literal[False] = pydantic.Field(
         default=False,
-        const=True,
         description='Whether the backend supports batch delete.',
     )
 
@@ -708,15 +701,13 @@ class TOSBackend(Boto3Backend):
     https://docs.byteplus.com/en/docs/tos/docs-compatibility-with-amazon-s3#appendix-tos-compatible-s3-apis
     """
 
-    scheme: str = pydantic.Field(
+    scheme: Literal['tos'] = pydantic.Field(
         default=common.StorageBackendType.TOS.value,
-        const=True,
         description='The scheme of the TOS backend.',
     )
 
     supports_batch_delete: Literal[True] = pydantic.Field(
         default=True,
-        const=True,
         description='Whether the backend supports batch delete.',
     )
 
@@ -819,9 +810,8 @@ class AzureBlobStorageBackend(common.StorageBackend):
     Azure Blob Storage Backend
     """
 
-    scheme: str = pydantic.Field(
+    scheme: Literal['azure'] = pydantic.Field(
         default=common.StorageBackendType.AZURE.value,
-        const=True,
         description='The scheme of the Azure Blob Storage backend.',
     )
 
@@ -832,7 +822,6 @@ class AzureBlobStorageBackend(common.StorageBackend):
 
     supports_environment_auth: Literal[True] = pydantic.Field(
         default=True,
-        const=True,
         description='Whether the backend supports environment authentication.',
     )
 

src/lib/data/storage/backends/common.py

@@ -99,13 +99,15 @@ class StoragePath:
 class StorageBackend(
     abc.ABC,
     pydantic.BaseModel,
-    extra=pydantic.Extra.forbid,
-    arbitrary_types_allowed=True,
-    keep_untouched=(functools.cached_property,),  # Don't serialize cached properties
 ):
     """
     Represents information about a storage backend.
     """
+    model_config = pydantic.ConfigDict(
+        extra='forbid',
+        arbitrary_types_allowed=True,
+        ignored_types=(functools.cached_property,),  # Don't serialize cached properties
+    )
 
     scheme: str
     uri: str

src/lib/data/storage/client.py

@@ -68,10 +68,11 @@ class Client(pydantic.BaseModel):
     A storage client that can be used to perform data operations against a remote storage.
     """
 
-    class Config:
-        extra = pydantic.Extra.forbid
-        frozen = True
-        keep_untouched = (functools.cached_property,)
+    model_config = pydantic.ConfigDict(
+        extra='forbid',
+        frozen=True,
+        ignored_types=(functools.cached_property,),
+    )
 
     #########################
     #    Factory methods    #
@@ -192,7 +193,7 @@ def create(
 
     storage_uri: str = pydantic.Field(
         ...,
-        regex=constants.STORAGE_BACKEND_REGEX,
+        pattern=constants.STORAGE_BACKEND_REGEX,
         description='The URI of the remote storage this instance of storage Client will '
                     'operate against. Must point to a valid container.',
     )
@@ -231,29 +232,25 @@ def create(
         description='Headers to apply to all requests of this client.',
     )
 
-    @pydantic.root_validator(skip_on_failure=True)
-    @classmethod
-    def validate_data_credential_endpoint(cls, values):
+    @pydantic.model_validator(mode='after')
+    def validate_data_credential_endpoint(self):
         """
         Validates that the data credential endpoint matches the storage backend profile.
         """
-        data_credential_input = values.get('data_credential_input')
-        if data_credential_input is not None:
-            storage_uri = values.get('storage_uri')
-
+        if self.data_credential_input is not None:
             # Construct backends to validate profiles match
             data_cred_backend = backends.construct_storage_backend(
-                uri=data_credential_input.endpoint,
+                uri=self.data_credential_input.endpoint,
             )
             storage_backend = backends.construct_storage_backend(
-                uri=storage_uri,
+                uri=self.storage_uri,
             )
 
             if data_cred_backend.profile != storage_backend.profile:
                 raise osmo_errors.OSMOCredentialError(
                     'Credential endpoint must match the storage backend profile')
 
-        return values
+        return self
 
     @functools.cached_property
     def data_credential(self) -> credentials.DataCredential:
@@ -1195,9 +1192,7 @@ class SingleObjectClient(pydantic.BaseModel):
     interacting with a single object.
     """
 
-    class Config:
-        extra = pydantic.Extra.forbid
-        frozen = True
+    model_config = pydantic.ConfigDict(extra='forbid', frozen=True)
 
     @overload
     @classmethod

src/lib/data/storage/constants/constants.py

@@ -38,7 +38,7 @@
 TOS_REGEX = fr'^tos://{URI_COMPONENT}(/{URI_COMPONENT})+/*$'
 AZURE_REGEX = fr'^azure://{URI_COMPONENT}(/{URI_COMPONENT})+/*$'
 STORAGE_BACKEND_REGEX = fr'({SWIFT_REGEX}|{S3_REGEX}|{GS_REGEX}|{TOS_REGEX}|{AZURE_REGEX})'
-StorageBackendPattern = Annotated[str, pydantic.Field(regex=STORAGE_BACKEND_REGEX)]
+StorageBackendPattern = Annotated[str, pydantic.Field(pattern=STORAGE_BACKEND_REGEX)]
 
 
 # Regex rules for storage profiles
@@ -50,7 +50,7 @@
 STORAGE_PROFILE_REGEX = fr'({SWIFT_PROFILE_REGEX}|{S3_PROFILE_REGEX}|' \
     fr'{GS_PROFILE_REGEX}|{TOS_PROFILE_REGEX}|' \
     fr'{AZURE_PROFILE_REGEX})'
-StorageProfilePattern = Annotated[str, pydantic.Field(regex=STORAGE_PROFILE_REGEX)]
+StorageProfilePattern = Annotated[str, pydantic.Field(pattern=STORAGE_PROFILE_REGEX)]
 
 STORAGE_CREDENTIAL_REGEX = fr'({STORAGE_PROFILE_REGEX}|{STORAGE_BACKEND_REGEX})'
-StorageCredentialPattern = Annotated[str, pydantic.Field(regex=STORAGE_CREDENTIAL_REGEX)]
+StorageCredentialPattern = Annotated[str, pydantic.Field(pattern=STORAGE_CREDENTIAL_REGEX)]

src/lib/data/storage/core/BUILD

@@ -28,6 +28,7 @@ osmo_py_library(
         "//src/lib/utils:logging",
         "//src/lib/utils:osmo_errors",
         requirement("pydantic"),
+        requirement("pydantic-settings"),
         requirement("tqdm"),
     ],
 )

src/lib/data/storage/core/executor.py

@@ -47,6 +47,7 @@
 )
 
 import pydantic
+import pydantic_settings
 
 from . import progress, provider
 from ....utils import common, logging as logging_utils, osmo_errors
@@ -67,32 +68,30 @@
 #   Executor Schemas (External)   #
 ###################################
 
-class ExecutorParameters(pydantic.BaseSettings):
+class ExecutorParameters(pydantic_settings.BaseSettings):
     """
     A class for storing parameters regarding multi-process/thread operations.
 
     Allows for environment variable overrides of the parameters.
     """
 
-    class Config:
-        """
-        Pydantic configuration for the ExecutorParameters class.
-        """
-        env_prefix = 'OSMO_EXECUTOR_'
-
-        @classmethod
-        def customise_sources(
-            cls,
-            init_settings,
-            env_settings,
-            file_secret_settings,
-        ):
-            # Treat explicit None as "unset" so env vars can apply
-            def init_without_none(settings):
-                data = init_settings(settings)
-                return {k: v for k, v in data.items() if v is not None}
-
-            return (init_without_none, env_settings, file_secret_settings)
+    model_config = pydantic_settings.SettingsConfigDict(env_prefix='OSMO_EXECUTOR_')
+
+    @classmethod
+    def settings_customise_sources(
+        cls,
+        settings_cls,  # pylint: disable=unused-argument
+        init_settings,
+        env_settings,
+        dotenv_settings,  # pylint: disable=unused-argument
+        file_secret_settings,
+    ):
+        # Treat explicit None as "unset" so env vars can apply.
+        # Override init_settings to filter out None values.
+        init_settings.init_kwargs = {
+            k: v for k, v in init_settings.init_kwargs.items() if v is not None
+        }
+        return (init_settings, env_settings, file_secret_settings)
 
     num_processes: int | None = pydantic.Field(
         default=None,
@@ -130,10 +129,7 @@ def init_without_none(settings):
         description='The size of the log queue for the executor. Only used for multi-process jobs.',
     )
 
-    @pydantic.validator(
-        'num_threads_inflight_multiplier',
-        'chunk_queue_size_multiplier',
-    )
+    @pydantic.field_validator('num_threads_inflight_multiplier', 'chunk_queue_size_multiplier')
     @classmethod
     def _validate_multiplier_max(cls, v: int) -> int:
         if v > MAX_MULTIPLIER: