Pydantic v2 Overhaul [DNM yet]

[Lnaden]

Description

This PR upgrades the entire QCElemental package to use the Pydantic v2 system, entirely removing any and all calls to the v1 API.

This PR also preserves QCElemental's API as best as it can, throwing deprecation warnings in places where we may want to upgrade; which can be a point of discussion.

All tests pass locally, however, this should NOT be merged yet until a formal discussion can be had. Given the massive change this PR introduces, and considering how many packages depend on this repo, I want everyone to be sure about a full stack upgrade before this is merged. This is a medium priority, but no rush PR (@loriab and @bennybp specifically)

Happy to talk and explain any and everything.

Changelog description

Updates to Pydantic v2, rendering v1 incompatible.

Status

• [x] Code base linted
• [x] Ready to go

[codecov[bot]]

Codecov Report

Merging #321 (75e6aee) into master (c9e308c) will increase coverage by 0.04%. Report is 2 commits behind head on master. The diff coverage is 82.91%.

Additional details and impacted files

[Lnaden]

e32570b might be the most frustrating little bug i fixed ever because we and pydantic do such a good job of trapping bugs all the way up the stack, and the actual issue was Python less than 3.10.

[Lnaden]

I have just spend 2 whole days chasing down the two most subtle interactions ever: One was in the autodoc code (pre Python 3.10 issue with annotations), and one was with Annotated + Numpy.typing.NDArray in pre Python 3.9 and how Annotated does instance checking against Python's native type _GenericAlias objects but NumPy implemented their own. GYAH.

I'll be able to cycle back around to your comments now @loriab.

[loriab]

Thanks for all your work on this so far, @Lnaden. I've updated through psi4 with pydantic v2 API and hit a couple serialization bugs/changes that I hope you can either fix or advise on.

(1) The Datum class doesn't like to be json serialized. (This shows up in psi4 vibrational analysis -- test tu4.) In qcelemental/tests/test_datum.py, this can be recreated by switching out .json() for .dict(). The error is shown below.

def test_to_dict(dataset):
    listans = [i * 4 / 3 for i in range(4)]
    ans = {"label": "an array", "units": "cm^-1", "data": listans, "comment": "freqs", "numeric": True}

    #dicary = dataset["ndarray"].dict()  # fine
    #dicary = dataset["ndarray"].model_dump()  # fine but fails on compare_recursive. error below
    #dicary = dataset["ndarray"].json()  # error below
    dicary = dataset["ndarray"].model_dump_json()  # error below
    assert compare_recursive(ans, dicary, 9)

# .json() or .model_dump_json() error
#E       pydantic_core._pydantic_core.PydanticSerializationError: Error serializing to JSON: PydanticSerializationError: Unable to serialize unknown type: array([0.        , 1.33333333, 2.66666667, 4.        ])

# .model_dump() error
#E       AssertionError: assert False
#E        +  where False = compare_recursive({'comment': 'freqs', 'data': [0.0, 1.3333333333333333, 2.6666666666666665, 4.0], 'label': 'an array', 'numeric': True, ...}, {'comment': 'freqs', 'data': array([0.        , 1.33333333, 2.66666667, 4.        ]), 'doi': None, 'glossary': '', ...}, 9)

(2) optking is using qcel.util.serialization.json_dumps to make an AtomicResult json serializable. https://github.com/psi-rking/optking/blob/master/optking/compute_wrappers.py#L84-L85 If I switch it to ret = json.loads(ret.model_dump_json()), it's fine (that is, the optimization succeeds, but I didn't examine the intermediates). Should this line need to change? It'd be nice for optking (no pydantic dependence) to not need to branch based on qcel.

[Lnaden]

@loriab Thanks for testing the downstream things for me to find some of these bugs.

(1) That was indeed a bug, two separate ones that. I had updated the dict method but not the model_dump method for Datum which is not a derivative of ProtoModel, so the code paths were different depending on which you called. I also learned that you cannot chain together WrapSerializer objects the way I had them, so I combined them into one function which fixes your use cases there.

(2) Not sure on what exactly to do here because I don't understand your expected outputs. Can you tell me what the function was expecting type wise and what it gets out in this v2 version? Explicit outputs would be super helpful.

I'm almost certain I can fix this on my end, I just need more details. I know that one difference in Pydantic between v1 and v2 that might impact this is in JSON encoding. In v1, the .json methods took the model and cast it to a full string-ified dict with output type str. in v2, it casts its model to a JSON-compatible dict where all keys and values are str, but returns a dict overall. I don't know if that's related to what you need here?

Thanks again for testing downstream things, this is a huge help in finding subtle interactions like the one from your (1)

[loriab]

Thanks very much!

(1b) Still on the Datum part, now it gets through all the serializing all the real ndarrays, but it balks on the complex (freq for transition states). If you change dict -> json in test_complex_array, it gives E pydantic_core._pydantic_core.PydanticSerializationError: Error serializing to JSON: PydanticSerializationError: Unable to serialize unknown type: 1j

@@ -133,7 +141,7 @@ def test_complex_array():
         "numeric": True,
     }
 
-    dicary = datum1.dict()
+    dicary = datum1.json()
     assert compare_recursive(ans, dicary, 9)

(1c) This next one (still on Datum), I don't need, and I don't know if its a behavior change, I just thought I'd run it by you since you've got a feel for serialization. It doesn't look like one can recreate the model once serialized to json because arrays end up as strings.

(in test_datum.py)

def test_to_dict(dataset):
    listans = [i * 4 / 3 for i in range(4)]
    ans = {"label": "an array", "units": "cm^-1", "data": listans, "comment": "freqs", "numeric": True}

    jstr = dataset["ndarray"].model_dump_json()
    print(jstr)
    dicary = json.loads(jstr)
#    dicary = dataset["ndarray"].model_dump()
    print(dicary)
    model = qcel.Datum(**dicary)
    print(model.model_dump())
    print(ans)

    assert compare_recursive(ans, dicary, 9)

self = <[AttributeError("'Datum' object has no attribute '__pydantic_extra__'") raised in repr()] Datum object at 0x7fa0246782d0>, label = 'an array', units = 'cm^-1'
data = '[0.0, 1.3333333333333333, 2.6666666666666665, 4.0]'

    def __init__(self, label, units, data, *, comment=None, doi=None, glossary=None, numeric=True):
        kwargs = {"label": label, "units": units, "data": data, "numeric": numeric}
        if comment is not None:
            kwargs["comment"] = comment
        if doi is not None:
            kwargs["doi"] = doi
        if glossary is not None:
            kwargs["glossary"] = glossary
    
>       super().__init__(**kwargs)
E       pydantic_core._pydantic_core.ValidationError: 1 validation error for Datum
E       data
E         Value error, Datum data should be float, Decimal, or np.ndarray, not <class 'str'>. [type=value_error, input_value='[0.0, 1.3333333333333333....6666666666666665, 4.0]', input_type=str]
E           For further information visit https://errors.pydantic.dev/2.1/v/value_error

qcelemental/datum.py:90: ValidationError
-------------------------------------------------------------------------------------------- Captured stdout call ---------------------------------------------------------------------------------------------
{"numeric":true,"label":"an array","units":"cm^-1","data":"[0.0, 1.3333333333333333, 2.6666666666666665, 4.0]","comment":"freqs"}
{'numeric': True, 'label': 'an array', 'units': 'cm^-1', 'data': '[0.0, 1.3333333333333333, 2.6666666666666665, 4.0]', 'comment': 'freqs'}

[Lnaden]

@loriab

Should have resolved 1b and 1c now.

(1b) Since serialization doesnt chain anymore, I have to handle all cases in one function instead of letting the type detection handle it. Complex NumPy arrays now are cast to list, and then the complexs are also cast to list. So this should be handled.

(1c) The output of JSON strings is now correctly a list instead of a string-ified list. Technically speaking, your code wouldn't work on master anyways since Dataum doesn't take a list input, but the error you pointed out where the json version of data was a str(list(...)) should now be a list(...) correctly.

[loriab]

1b definitely resolved, thanks!

(1d) What's your view on numpy floats? I can easily solve this in my code by casting to float, but as-is, one can create a Datum with a numpy float but it won't json-ize.

>>> import qcelemental as qcel
>>> rT=0.5
>>> import numpy as np
>>> asdf = qcel.Datum("a np float", "J/K", np.sum(rT / np.expm1(rT) - np.log(1 - np.exp(-rT))))
>>> print(asdf)
----------------------------------------
            Datum a np float            
----------------------------------------
Data:     1.7034991708355878
Units:    [J/K]
doi:      None
Comment:  
Glossary: 
----------------------------------------
>>> type(asdf.data)
<class 'numpy.float64'>
>>> asdf.model_dump()
{'numeric': True, 'label': 'a np float', 'units': 'J/K', 'data': 1.7034991708355878}
>>> asdf.model_dump_json()
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "/psi/toolchainconda/envs/py311pyd2/lib/python3.11/site-packages/pydantic/main.py", line 345, in model_dump_json
    return self.__pydantic_serializer__.to_json(
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
pydantic_core._pydantic_core.PydanticSerializationError: Error serializing to JSON: PydanticSerializationError: Unable to serialize unknown type: 1.7034991708355878
>>> 

[Lnaden]

I think thats something I can fix on my end, it should be casting to float correctly, dunno why its not. Probably has to do with the "automatically inferred output type of the serialization function based on the type hint" and I'm not being generous enough.

[Lnaden]

@loriab

(1d) Fixed this by forcing the NumPy scalars to be cast to Python types. Wheeeeeeeeeeeeeee edge cases

This had nothing to do with the type hint return idea unfortunately.

[loriab]

Confirmed (1d) fixed! thanks, @Lnaden

[loriab]

Nice strategy with the deprecation warning on .dict and .json, btw. The first time through the full psi4 test suite, I triggered 89k of them :-) . Now fixed except for some choice leftovers so we know it won't break for the user.

[Lnaden]

Thanks. I tried to preserve the API on our side as much as possible while keeping the deprecation warning. I can remove the warning on our side and just let people use both. There really shouldn't be any harm in it, so I can cut that if its obnoxious.

[loriab]

Thanks. I tried to preserve the API on our side as much as possible while keeping the deprecation warning. I can remove the warning on our side and just let people use both. There really shouldn't be any harm in it, so I can cut that if its obnoxious.

I think it's good to have the warning. It makes for a nice warning hunt :-)

I agree on allowing both for the foreseeable future. It would have been vexing to hunt all the dict()s down before facing all the real pydantic v2 api issues.

And they're not obnoxious -- pytest collects them very nicely.

tests/pytests/test_addons.py: 3 warnings
tests/pytests/test_detci_opdm.py: 1 warning
tests/pytests/test_standard_suite.py: 26288 warnings
tests/pytests/test_raises.py: 4 warnings
tests/pytests/test_mcmurchie_davidson.py: 1 warning
tests/pytests/test_restart.py: 5 warnings
tests/pytests/test_mints.py: 2 warnings
tests/pytests/test_misc.py: 2 warnings
tests/pytests/test_scf_options.py: 16 warnings
tests/pytests/test_addons_qcschema.py: 9 warnings
tests/pytests/test_qcfractal.py: 224 warnings
tests/pytests/test_molden_writer.py: 9 warnings
tests/pytests/test_nbody.py: 708 warnings
tests/pytests/test_mp2.py: 6 warnings
tests/pytests/test_dft_benchmarks.py: 2208 warnings
tests/pytests/test_mp2d.py: 34 warnings
tests/pytests/test_ccresponse.py: 2 warnings
tests/pytests/test_qcschema.py: 26 warnings
tests/pytests/test_composite.py: 26 warnings
tests/pytests/test_compositejk.py: 112 warnings
tests/pytests/test_nbody_multi_level_qcschema.py: 252 warnings
tests/pytests/test_option.py: 1 warning
tests/pytests/test_optking.py: 177 warnings
tests/pytests/test_psi4.py: 30 warnings
tests/pytests/test_psi4_qcschema.py: 34 warnings
tests/pytests/test_pyside_cubegen.py: 1 warning
tests/pytests/test_dft_blocking.py: 10 warnings
tests/pytests/test_dipoles.py: 60 warnings
tests/pytests/test_erisieve.py: 11 warnings
tests/pytests/test_fchk_writer.py: 12 warnings
tests/pytests/test_fcidump_energy.py: 2 warnings
tests/pytests/test_geometric.py: 31 warnings
tests/pytests/test_gradients.py: 403 warnings
tests/pytests/test_tdscf_excitations.py: 70 warnings
tests/pytests/test_vibanalysis.py: 10374 warnings
tests/pytests/test_tdscf_products.py: 10 warnings
tests/pytests/test_weird_basis.py: 12 warnings
tests/pytests/test_wfn.py: 1 warning
  /psi/gits/QCElemental/qcelemental/models/basemodels.py:133: DeprecationWarning: The `dict` method is deprecated; use `model_dump` instead.
    warnings.warn("The `dict` method is deprecated; use `model_dump` instead.", DeprecationWarning)

[awvwgk]

Is this effort still active? Having QCElemental with Pydantic v2 support would be great, however a smooth transition from v1 to v2 API would be appreciated. Currently QCElemental is relying on the v1 API provided in Pydantic v2, instead of directly switching to the v2 API and removing the support for v1, would it be possible to have both available dependent on which Pydantic API is available? An example for this could be the approach taken in the beanie package.

[loriab]

Is this effort still active? ...

It's still active, and it works (except for one line with optking). And downstream v2-proper PRs are largely ready for QCEngine and Psi4. It's frustrating that pydantic hasn't provided a per-model upgrade scheme because as is, whole software stacks need to be migrated together. Last I heard, OpenFF was having trouble with their v2 migration. Any updates, @mattwthompson ?

Thanks for the beanie model. I think that'd work pretty cleanly for QCEngine and maybe Psi4. QCElemental might need some completely separate files for v1/v2 since it access low-level pydantic stuff more. But I agree something needs to be done -- I'm translating new v2 models back to v1 just so I can actually use them in the stack.

[mattwthompson]

Our "proper" v2 migration is tied up in a huge mess of things and stalled - the crux of the thing being getting our custom pint registry to work in Pydantic models without breaking all of our downstream users' APIs

https://github.com/openforcefield/openff-models/blob/main/openff/models/types.py#L20-L37

[loriab]

Thanks for the update, @mattwthompson . After I get through these next papers I'll try a dual v1/v2 version of qcel. But I bet one wants a triple version for complete compatibility -- v1, v2 as v1, v2.