Upgrading to pydantic v2

[bcdurak]

⚠️ Since this PR is turning out to be a mega PR, I will try to explain everything in a very detailed manner. Please go through it before diving into the changes and feel free to contact me if you have any further questions. (Keep in mind that we have a dedicated Notion page and a recorded meeting regarding this upgrade as well.)

Why?

Why do we upgrade pydantic to v2?

When it comes to the advantages of upgrading our pydantic dependency, there are mainly two big ones:

• A significant performance boost
• Being able to keep up with our integrations that already upgraded to pydantic v2

At the same time, we face the following challenges:

• A very significant change to our codebase
• A slightly more challenging debugging process due to the backend changes in pydantic

Moreover, the pydantic team stopped the active development of V1. They still make releases with critical bug fixes and some security maintenance but that will also stop these at the end of June 2024.

Why do we upgrade sqlmodel and sqlalchemy along with pydantic?

• Our current version of sqlmodel (0.0.8) does not support pydantic v2.
• They started supporting pydantic with version (0.0.14). Since then, they have supported both v1 and v2.
• However, at 0.0.12, they made a hard switch to sqlalchemy v2 as well.
• So, any version of sqlmodel that supports pydantic v2 requires you to work with sqlalchemy v2.
• Since we can not drop our sqlmodel dependency, we need to upgrade all of these packages altogether.

Other packages

Due to the pydantic v2 support, there are a few more important dependencies, like fastapi, that were affected by this upgrade. To see the full list, check the changes in the pyproject.toml.

How?

Migration Guides

Before I explain the changes in our codebase, I want to mention that both the pydantic and sqlalchemy upgrades come with significant changes. You can check the respective migration guides and get more info with these links: pydantic v2 migration guide, sqlalchemy v2 migration guide

pydantic was also kind enough to offer a tool called bump-pydantic which helped a lot at the start. It roughly modified 80 files or so, mostly focused on the configuration of models and some validators. But, as you can see from the number of changed files, there were a lot of things that we still had to adopt after the tool did its migration.

The most critical changes w.r.t. the pydantic upgrade

• Configuration for models has been reworked.

  In Pydantic V2, to specify config on a model, you should set a class attribute called model_config to be a dictionary with the key/value pairs you want to be used as the config. The Pydantic V1 behavior to create a class called Config in the namespace of the parent BaseModel subclass is now deprecated.

• Many configuration options have been either deprecated or removed. The most important ones include:

  • allow_mutation is now called frozen and it is set to False by default.

  • underscore_attrs_are_private is removed and the models behave in a way like this value is set to True.

  • The smart_union configuration parameter is now removed. Now, the default behavior is smart_union, which means if we had a Union field that was not parameterized with smart_union before, we have to manually set union_mode='left_to_right' to keep the same behavior. Check here for more details.

  • json_encoders have been removed first, added back afterward, and deprecated later.

    • This was mainly used for pydantic.SecretStrs in our codebase. Now, we have replaced this functionality with a custom type annotation called ZenSecretStr, which serves simply as a SecretStr with a custom pydantic serializer.

  • The regex parameter is removed and a new parameter called pattern is now introduced:

    Pydantic V1 used Python's regex library. Pydantic V2 uses the Rust [regex crate](https://github.com/rust-lang/regex). This crate is not just a "Rust version of regular expressions", it's a completely different approach to regular expressions. In particular, it promises linear time searching of strings in exchange for dropping a couple of features (namely look arounds and backreferences).

• The __fields__ have been replaced by model_fields. Previously in V1, you were to be able to get a .type_ for each field but this is not the case anymore. The replacement is called .annotation. However, it acts in a slightly different way. For instance, an Optional[int] field previously had a .type_ int but now the .annotation is Optional[int].

• Validators have been heavily reworked. There are no @validators or @root_validators anymore. The new validators are called @field_validator and @model_validator. They now feature a lot more flexibility and functionality. IMO, this change is one of the most critical ones and it has a lot of implications when it comes to our codebase. So, if you would like to get a detailed explanation, you can check all the changes here.

• The skip_on_failure parameter in the validator decorator has been removed. The only validator of this type that we had before now throws a warning instead of failing.

• There is an important change when it comes to the serialization of subclasses. Check this issue I have created on their GitHub page for more detail. TLDR, if you are using subclasses in your models, do not forget to use the SerializeAsAny[NameOfBaseModel] as the annotation to keep the same serialization behavior as v1.

• You can not use subclasses so easily anymore. For instance, if you subclass int, you can not directly use it as a type in a pydantic class. This is by design. You need to define a method called get_pydantic_core_schema in this new class in order to be able to use it as an annotation.

• Let’s say you define a new pydantic class A, you annotate on of its fields with another pydantic class B. Now, you subclass A, call it A’, and it requires a subclass version of B, let’s call that B’. Previously, you could use an instance of B to create an instance of A’, but this is not the case anymore. You have to explicitly convert B to B’ before you can pass it to the constructor of A’. Check the base_zen_store and base_secret_store implementations for more details.

• pydantic definition of generic models have been removed as well.

  The pydantic.generics.GenericModel class is no longer necessary, and has been removed. Instead, you can now create generic BaseModel subclasses by just adding Generic as a parent class on a BaseModel subclass directly. This looks like class MyGenericModel(BaseModel, Generic[T]): ....

• Fields do not have a required field anymore, instead they have an is_required() method. Due to this, if you would like to make a field non-required, you have to set the default value or the default factory.

• There is also a very significant change with regards to the optional and nullable fields. Most importantly in our case, if you want to define an optional value, you have to provide at least None as a default value. Otherwise, in contrast to V1, even if you do Optional[int], it will still be a required field.

• parse_obj and parse_raw have been deprecated, instead, the recommendation is to use model_validate. However, this method is functioning in a slightly different way. In contrast to V1, if you feed it an instance of a subclass it fails with a validation error:

  from pydantic import BaseModel
  
  class A(BaseModel):
      a: int = 3
  
  class B(A):
      b: str = "str"
  
  b1 = B.model_validate(A(a=2))  # Fails with a validation error
  
  b2 = B.model_validate_json(A().model_dump_json())  # Works

• The update_forward_refs method has been reworked and renamed. Now it is enough to just do MyModel.model_rebuild().

• There is a new Python package called pydantic-settings. Classes such as the SettingsConfigDict are now a part of this package.

• ValidatedFunction has been deprecated. Check the utils/pydantic_utils.py for further info and see if we can remove this. (tagging @schustmi here)

• ModelMetaclass has been moved to pydantic._internal module. Check the global_config.py and typed_model.py in our codebase.

• They removed their collection of utility methods in their typing module (including functions such as get_args and get_origin). Since our codebase heavily used these functions, I carried the original versions over to work in our codebase.

• There were instances where we used some_model_instance.json(). This behavior is now replaced with the some_model_instance.model_dump_json(). However, if you would like to parameterize this process by using keys like sort_keys, this is unfortunately not possible anymore. As an alternative, I have applied some_model_instance.model_dump() before and then used the json package manually to dump it with the sort_keys parameter.

• pydantic.Fields with the max_length setting now fail if they have UUID in their annotation. In such cases, I have separated the validation function.

• When it comes to fields, field_info.extra has been renamed to field.json_schema_extra. You can find an example how this is being used by check the changes in the zenml.utils.secret_utils.

• This one was a bit interesting and hard to figure out. When you do zenml up, if there is a response model that has an Enum field defined with a pydantic.Field and the field is parameterized with max_length, the local server deployment will fail. Still, I can not figure out the root cause of this issue. However, this is not a critical use case so I removed these instances and now we can do zenml up successfully.

• With pydantic V2, the issue regarding multiple inherited config classes is now resolved. The related ignore tags have been removed.

• The schema_json method is deprecated; we are using model_json_schema and json.dumps instead.

• The copy method is deprecated; we are using model_copy instead. You can check the docstring of BaseModel.copy for details about how to handle include and exclude.

• Our update model decorator has been removed. At first, this change was mainly triggered by various failing mypy linting issues because they changed the way of defining required/optional values in pydantic v2. However, soon it helped us reveal some linting issues that were suppressed by the relationship between our previous ...Request and ...Update models. Each update model is now implemented properly with optional annotations.

• With pydantic v2, the error handling within the validators has been reworked as well:

  As mentioned in the previous sections you can raise either a ValueError or AssertionError(including ones generated by assert ... statements) within a validator to indicate validation failed. You can also raise a PydanticCustomError which is a bit more verbose but gives you extra flexibility. Any other errors (including TypeError) are bubbled up and not wrapped in a ValidationError.

• The following code block used to execute successfully in pydantic v1, but this behavior has changed in pydantic v2 and it now throws a ValidationError:

  from pydantic import BaseModel
  
  class MyModel(BaseModel):
      a: str
  
  MyModel(a=2.3)

• This following code block used to print out True and True but with the new changes in pydantic now it outputs False and False:

  from pydantic import BaseModel
  
  class One(BaseModel):
      a: int = 3
  
  class Two(One):
      pass
  
  print(One() == Two())
  print(One() == {"a": 3})

• Fields that are provided as extra fields to any model can be accessed by .model_extra now.

• In contrast to pydantic v1, defining any pydantic class without properly annotating its fields will raise a pydantic.errors.PydanticUserError now.

Critical changes w.r.t. the sqlmodel upgrade

The most critical factor in this upgrade stems from the PR right here. With this change, they have changed the way they handle Enum values.

For instance, if you are familiar with our component schema (which we defined through sqlmodel), we have a field called type which was a StackComponentType:

# the schema of the stack component
class StackComponentSchema(...):
    ...
    type: StackComponentType
    ...

# and the stack component type looked like this:
class StackComponentType(StrEnum):
    ...
    ARTIFACT_STORE = "artifact_store"
    ...

With this setup, when we registered, for instance, a new artifact store, we created an entry in the components table of our DB where the column type had the string artifact_store stored in it as a value. However, with the new changes, sqlmodel now gives higher priority to Enum fields and saves the value ARTIFACT_STORE instead. While this is alright if you are starting from scratch, if you have any entry in a table with an Enum field zenml will fail after the upgrade. Instead of taking the migration route, we decided to adjust our schemas to use str fields instead and updated the corresponding to_model, update, and from_request methods.

Critical changes w.r.t. the sqlalchemy upgrade

The new sqlalchemy v2 has a lot of functional and syntactic changes as well. Luckily, most of the pure sqlalchemy code in our codebase can only be found around our sql_zen_store implementation and migration scripts. I have tried my best to fix all the deprecation issues but I ask you to pay extra attention to these changes, especially around the migration scripts.

Other critical changes

• As I mentioned before, pydantic has changed how you can define optional, required, and nullable fields. Moreover, they removed the required field from the FieldInfo. With the new update, there is a function called .is_required() for each field which checks if the field has a default value or default factory. Due to these changes, I had to rework our update_model decorator. However, in my experiments, the new possible solutions created a lot of problems with mypy and I ended up removing this decorator altogether. I implemented the exact equivalent version of these update models. This revealed a bunch of issues that were hidden before (because previously, fields in the update models were considered to be required instead of optional by mypy.). That's why you may see a few updates in the codebase, especially when it comes to the ServiceConnector models.

Integration Corner

I will try to update the following subsections as the fixes come along. Here you will find a list of all the integrations affected by the aforementioned changes to our codebase and dependencies.

AWS @wjayesh

The upgrade to kfp V2 (in integrations like kubeflow, tekton, or gcp) bumps our protobuf dependency from 3.X to 4.X. This is why we need to relax the sagemaker dependency.

• This needs to be tested very thoroughly as it is one of our major integrations.
• Update the docs.

Airflow ✅ @schustmi

I believe this was the most critical update. Airflow still has a dependency on sqlalchemy V1 and this conflicts with this entire PR as we have to migrate to sqlalchemy V1. However, we managed to figure out a way where we can still run pipelines on Airflow by keeping the Airflow and ZenML installation separated.

• Test the execution locally.
• Test the execution on a remote setup.
• Update the docs.

Evidently @safoinme

Relaxing the main dependency here resolved the installation issue. They started supporting pydantic V2 starting from the version 0.4.16. As their latest version is 0.4.22, the dependency is limited between the two. When you install zenml and the evidently integration afterward, it installs 0.4.22. However, if you use the install-zenml-dev script, it ends up installing 0.4.16. This is why it might make sense to test both versions.

• Test a pipeline using the Evidently integration on 0.4.16
• Test a pipeline using the Evidently integration on 0.4.22
• Update the docs.

Feast ✅ @strickvl

To fix the installation issues, we had to remove the redis extra from the feast integration. As the latest version is 0.37.1, the dependency is capped at 0.37.1.

• Test the functionalities of the Feast integration.
• Update the docs.

GCP @safoinme @strickvl

This is also one of the major changes. As they switched to their own V2, the Python SDK of kfp removed their pydantic v1 dependency, which ultimately solved our installation issues. However, this means that we have to adapt our integration accordingly to work with kfp>=2.0. You can find the migration guide for KFP SDK V1 to V2 here. Also, Felix previously worked on this issue and you can find his changes right here in this PR.

• This needs to be tested very thoroughly as it is one of our major integrations.

Great Expectations ✅ @stefannica

Similar to the previous integrations, relaxing the main dependencies here resolved the installation issue. As they started supporting installations with pydantic v2 from 0.17.15, the minimum requirement was changed. There was a note in the requirements of this integration stating that typing_extensions>4.6.0 does not work with GE, and the resolved version is 4.10.0. We need to figure out if this is still an issue. Moreover, they are closing on their 1.0 release. Since this might include major breaking changes, I put the upper limit to <1.0 for now.

• Test a pipeline using the Great Expectations integration.
• Update the docs.

Kubeflow @safoinme @strickvl

Similar to the GCP integration, relaxing the kfp python SDK dependency resolved the installation issue, however, the code still needs to be migrated. You can find the migration guide for KFP SDK V1 to V2 here. Also, Felix previously worked on this issue and you can find his changes right here in this PR.

• Test after the migration.
• Update the docs.

mlflow @avishniakov

This was an interesting change. As they stand right now, the dependencies of the mlflow integration are compatible with zenml using pydantic v2. However, if you install zenml first and then do zenml integration install mlflow -y, it downgrades pydantic to v1. (I think this is an important problem that we have to solve separately in a generalized manner!) This is why I had to manually add the same duplicated pydantic requirement in the integration definition as well.

• Test the local case.
• Test the remote case.
• Update the docs.

Label Studio ✅ @strickvl

They still have a hard dependency on pydantic = "<=1.11.0,>=1.7.3" for the label_studio package. @strickvl has opened up an issue on their GitHub page. We decided to remove that and just rely on the label_studio_sdk package as that allows for Pydantic >2.x.

• Test local Label Studio
• Update the docs.
• Test the annotator Python SDK / interface

Skypilot @safoinme

While uv was able to compile a list of requirements using pydantic>=2.7.0 with both skypilot[aws]<=0.5.0and skypilot[gcp]<=0.5.0 respectively, skypilot[azure]<=0.5.0 is still creating issues.

• Test all possible variations of this.
• Update the docs.

Tensorflow @avishniakov

The new version of pydantic creates a drift between the tensorflow and typing_extensions packages. Relaxing the dependencies here resolves the issue, however, there is a known issue between torch and tensorflow and we need to test whether this is still problematic.

Additionally, the upgrade to kfp V2 (in integrations like kubeflow, tekton, or gcp) bumps our protobuf dependency from 3.X to 4.X. This is another reason why the tensorflow upgrade is necessary.

• Test the tensorflow integration.
• Update the docs.

Tekton @safoinme @strickvl

The tekton integration should go through a major change as well, since it is affected by the kfp changes. You can find the migration guide for KFP SDK V1 to V2 here. Also, Felix previously worked on this issue and you can find his changes right here in this PR.

• Needs to be tested.
• Update the docs.

Docs changes

Keep in mind, much like the changes in the airflow integration, some future updates will probably require changes in our documentation.

Special Thanks

Special thanks to the pydantic team (especially @sydney-runkle) for helping us out when we got stuck. It has been a blast to work on this upgrade. Looking forward to V3 😄

Leftover TODOs

• When testing the mlflow integration, we realized the data type in the artifact versions (that have a DistributionPackageSource) was deserialized incorrectly when combined with the tenant setup, leading to a failure when you do artifact_version.load(). We need to investigate and solve this issue.
• Detect and fix any Union fields that were not parameterized with smart_union before, we have to manually set union_mode='left_to_right' to keep the same behavior.
• There is a problem with the serialization of SecretStrs. I have created an issue in the pydantic repository for this problem. Simply put, our service connectors have a configuration field. This field is a dictionary that might or might not have SecretStrs inside. When fastapi tries to serialize the response model of such a service connector, it fails with a pydantic SerializationError.
• Test how overriding serialize_as_any would work within fields of a BaseModel.
• The code is still using pydantic_encoder which is deprecated. Find an alternative solution to it.
• any_pydantic_model.dict() method is now deprecated. Even though, I fixed and removed most of these calls, it is really hard to scan the codebase for similar instances. So, anytime you run into any deprecation warnings, we have to remove these calls.
• There are a few deprecation warnings regarding sqlalchemy as well. We need to find replacements for those as well.
• Refactor the __init__ call from the BaseService.

Pre-requisites

Please ensure you have done the following:

• I have read the CONTRIBUTING.md document.
• If my change requires a change to docs, I have updated the documentation accordingly.
• I have added tests to cover my changes.
• I have based my new branch on develop and the open PR is targeting develop. If your branch wasn't based on develop read Contribution guide on rebasing branch to develop.
• If my changes require changes to the dashboard, these changes are communicated/requested.

Types of changes

• Bug fix (non-breaking change which fixes an issue)
• New feature (non-breaking change which adds functionality)
• Breaking change (fix or feature that would cause existing functionality to change) (UNDERSTATEMENT 😄 )
• Other (add details above)