[BUG] MLFlow 2.6.0 PydanticUserError when importing mlflow w/ latest pydantic version

[kparaju]

Issues Policy acknowledgement

• I have read and agree to submit bug reports in accordance with the issues policy

Willingness to contribute

No, I cannot contribute a bug fix at this time.

MLflow version

• Client: 2.6.0
• Tracking server: databricks

System information

• OS Platform and Distribution (e.g., Linux Ubuntu 16.04): Ubuntu 20.04.5 LTS (Databricks 11.3 LTS)
• Python version: 3.9.5
• yarn version, if running the dev UI:

Describe the problem

Exception is thrown when importing mlflow.

/tmp/tmpnmraxlyo.py in <module>
      1 import logging
      2 
----> 3 import mlflow

(....)

/local_disk0/.ephemeral_nfs/cluster_libraries/python/lib/python3.9/site-packages/mlflow/gateway/config.py in OpenAIConfig()
     77         return _resolve_api_key_from_input(value)
     78 
---> 79     @root_validator(pre=False)
     80     def validate_field_compatibility(cls, config: Dict[str, Any]):
     81         api_type = config.get("openai_api_type")

/local_disk0/.ephemeral_nfs/cluster_libraries/python/lib/python3.9/site-packages/pydantic/deprecated/class_validators.py in root_validator(pre, skip_on_failure, allow_reuse, *__args)
    226     mode: Literal['before', 'after'] = 'before' if pre is True else 'after'
    227     if pre is False and skip_on_failure is not True:
--> 228         raise PydanticUserError(
    229             'If you use `@root_validator` with pre=False (the default) you MUST specify `skip_on_failure=True`.'
    230             ' Note that `@root_validator` is deprecated and should be replaced with `@model_validator`.',

PydanticUserError: If you use `@root_validator` with pre=False (the default) you MUST specify `skip_on_failure=True`. Note that `@root_validator` is deprecated and should be replaced with `@model_validator`.

For further information visit https://errors.pydantic.dev/2.1.1/u/root-validator-pre-skip

Locking to mlflow==2.5.0 gets rid of the issue. Lowering to pydantic<2.0.0 also seems to solve the issue but I haven't checked if there is any other issues.

Tracking information

REPLACE_ME

Code to reproduce issue

% pip install --upgrade mlflow==2.6.0 pydantic==2.1.1
import mlflow

Stack trace

---------------------------------------------------------------------------
PydanticUserError                         Traceback (most recent call last)
<command-4153148899206496> in <cell line: 1>()
----> 1 import mlflow

/databricks/python_shell/dbruntime/PythonPackageImportsInstrumentation/__init__.py in import_patch(name, globals, locals, fromlist, level)
    169             # Import the desired module. If you’re seeing this while debugging a failed import,
    170             # look at preceding stack frames for relevant error information.
--> 171             original_result = python_builtin_import(name, globals, locals, fromlist, level)
    172 
    173             is_root_import = thread_local._nest_level == 1

/usr/lib/python3.9/importlib/_bootstrap.py in _find_and_load(name, import_)

/usr/lib/python3.9/importlib/_bootstrap.py in _find_and_load_unlocked(name, import_)

/usr/lib/python3.9/importlib/_bootstrap.py in _load_unlocked(spec)

/usr/lib/python3.9/importlib/_bootstrap.py in _load_backward_compatible(spec)

/databricks/python_shell/dbruntime/PostImportHook.py in load_module(self, fullname)
    216     def load_module(self, fullname):
    217         try:
--> 218             module = self.loader.load_module(fullname)
    219             notify_module_loaded(module)
    220         except (ImportError, AttributeError):

/databricks/python/lib/python3.9/site-packages/mlflow/__init__.py in <module>
    245 # Importing this module fails if they are not installed.
    246 with contextlib.suppress(ImportError):
--> 247     from mlflow import gateway  # noqa: F401
    248 
    249     __all__.append("gateway")

/databricks/python_shell/dbruntime/PythonPackageImportsInstrumentation/__init__.py in import_patch(name, globals, locals, fromlist, level)
    169             # Import the desired module. If you’re seeing this while debugging a failed import,
    170             # look at preceding stack frames for relevant error information.
--> 171             original_result = python_builtin_import(name, globals, locals, fromlist, level)
    172 
    173             is_root_import = thread_local._nest_level == 1

/databricks/python/lib/python3.9/site-packages/mlflow/gateway/__init__.py in <module>
----> 1 from mlflow.gateway.fluent import get_route, search_routes, query, create_route, delete_route
      2 from mlflow.gateway.utils import set_gateway_uri, get_gateway_uri
      3 from mlflow.gateway.client import MlflowGatewayClient
      4 
      5 __all__ = [

/databricks/python_shell/dbruntime/PythonPackageImportsInstrumentation/__init__.py in import_patch(name, globals, locals, fromlist, level)
    169             # Import the desired module. If you’re seeing this while debugging a failed import,
    170             # look at preceding stack frames for relevant error information.
--> 171             original_result = python_builtin_import(name, globals, locals, fromlist, level)
    172 
    173             is_root_import = thread_local._nest_level == 1

/databricks/python/lib/python3.9/site-packages/mlflow/gateway/fluent.py in <module>
      2 from typing import List, Dict, Any
      3 
----> 4 from mlflow.gateway.client import MlflowGatewayClient
      5 from mlflow.gateway.config import Route
      6 from mlflow.gateway.constants import MLFLOW_GATEWAY_SEARCH_ROUTES_PAGE_SIZE

/databricks/python_shell/dbruntime/PythonPackageImportsInstrumentation/__init__.py in import_patch(name, globals, locals, fromlist, level)
    169             # Import the desired module. If you’re seeing this while debugging a failed import,
    170             # look at preceding stack frames for relevant error information.
--> 171             original_result = python_builtin_import(name, globals, locals, fromlist, level)
    172 
    173             is_root_import = thread_local._nest_level == 1

/databricks/python/lib/python3.9/site-packages/mlflow/gateway/client.py in <module>
      6 
      7 from mlflow import MlflowException
----> 8 from mlflow.gateway.config import Route
      9 from mlflow.gateway.constants import (
     10     MLFLOW_GATEWAY_CRUD_ROUTE_BASE,

/databricks/python_shell/dbruntime/PythonPackageImportsInstrumentation/__init__.py in import_patch(name, globals, locals, fromlist, level)
    169             # Import the desired module. If you’re seeing this while debugging a failed import,
    170             # look at preceding stack frames for relevant error information.
--> 171             original_result = python_builtin_import(name, globals, locals, fromlist, level)
    172 
    173             is_root_import = thread_local._nest_level == 1

/databricks/python/lib/python3.9/site-packages/mlflow/gateway/config.py in <module>
     64 
     65 
---> 66 class OpenAIConfig(ConfigModel):
     67     openai_api_key: str
     68     openai_api_type: OpenAIAPIType = OpenAIAPIType.OPENAI

/databricks/python/lib/python3.9/site-packages/mlflow/gateway/config.py in OpenAIConfig()
     77         return _resolve_api_key_from_input(value)
     78 
---> 79     @root_validator(pre=False)
     80     def validate_field_compatibility(cls, config: Dict[str, Any]):
     81         api_type = config.get("openai_api_type")

/local_disk0/.ephemeral_nfs/envs/pythonEnv-a2fe7bb1-c2e4-4142-acab-f370db57c875/lib/python3.9/site-packages/pydantic/deprecated/class_validators.py in root_validator(pre, skip_on_failure, allow_reuse, *__args)
    226     mode: Literal['before', 'after'] = 'before' if pre is True else 'after'
    227     if pre is False and skip_on_failure is not True:
--> 228         raise PydanticUserError(
    229             'If you use `@root_validator` with pre=False (the default) you MUST specify `skip_on_failure=True`.'
    230             ' Note that `@root_validator` is deprecated and should be replaced with `@model_validator`.',

PydanticUserError: If you use `@root_validator` with pre=False (the default) you MUST specify `skip_on_failure=True`. Note that `@root_validator` is deprecated and should be replaced with `@model_validator`.

For further information visit https://errors.pydantic.dev/2.1.1/u/root-validator-pre-skip

Other info / logs

TBD

What component(s) does this bug affect?

• area/artifacts: Artifact stores and artifact logging
• area/build: Build and test infrastructure for MLflow
• area/docs: MLflow documentation pages
• area/examples: Example code
• area/gateway: AI Gateway service, Gateway client APIs, third-party Gateway integrations
• area/model-registry: Model Registry service, APIs, and the fluent client calls for Model Registry
• area/models: MLmodel format, model serialization/deserialization, flavors
• area/recipes: Recipes, Recipe APIs, Recipe configs, Recipe Templates
• area/projects: MLproject format, project running backends
• area/scoring: MLflow Model server, model deployment tools, Spark UDFs
• area/server-infra: MLflow Tracking server backend
• area/tracking: Tracking Service, tracking client APIs, autologging

What interface(s) does this bug affect?

• area/uiux: Front-end, user experience, plotting, JavaScript, JavaScript dev server
• area/docker: Docker use across MLflow's components, such as MLflow Projects and MLflow Models
• area/sqlalchemy: Use of SQLAlchemy in the Tracking Service or Model Registry
• area/windows: Windows support

What language(s) does this bug affect?

• language/r: R APIs and clients
• language/java: Java APIs and clients
• language/new: Proposals for new client languages

What integration(s) does this bug affect?

• integrations/azure: Azure and Azure ML integrations
• integrations/sagemaker: SageMaker integrations
• integrations/databricks: Databricks integrations

[BenWilson2]

Thank you for bringing this up to our attention @kparaju ! I've filed a fix for it and we'll be sure to get a patch out for this. In the meantime, can you pin pydantic to <2.x or use MLflow 2.5 until we get the patch released?

[kparaju]

Thank you for bringing this up to our attention @kparaju ! I've filed a fix for it and we'll be sure to get a patch out for this. In the meantime, can you pin pydantic to <2.x or use MLflow 2.5 until we get the patch released?

Yes, we can and will pin to mlflow 2.5 until patch is released.

[mabreuortega]

I'm seeing this issue with mlflow 2.7.1

[BenWilson2]

@mabreuortega Can you verify the results of:

pip freeze | grep pydantic

and

pip freeze | grep mlflow

?

[mabreuortega]

pip freeze | grep pydantic
pydantic==2.3.0 pydantic_core==2.6.3

pip freeze | grep mlflow

mlflow==2.7.1

[mabreuortega]

disregard, this above, I was facing another issue. ;-)

[harupy]

@mabreuortega I built this dockerfile, but couldn't reproduce the error:

FROM python:3.8

RUN pip install pydantic==2.3.0 pydantic_core==2.6.3 mlflow==2.7.1
RUN python -c 'import mlflow'

Can you run python -c 'import mlflow'?

[harupy]

disregard, this above, I was facing another issue. ;-)

Got it. Feel free to open a new issue.