align with core changes

[avishniakov]

Describe changes

zenml-io/zenml#2375

Pre-requisites

Please ensure you have done the following:

• If my change requires a change to docs, I have updated the documentation accordingly.
• I have added tests to cover my changes.
• I have updated ref-zenml in .github/workflows/ci.yml accordingly (if you don't know - main would be a solid choice)

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)
• Other (add details above)