From 8ca8db2bce90e197ca57064a1be3a2cbe466a563 Mon Sep 17 00:00:00 2001
From: Edward Oakes <ed.nmi.oakes@gmail.com>
Date: Thu, 25 Jan 2024 13:07:55 -0600
Subject: [PATCH] Fix `IS_PYDANTIC_2` logic for `pydantic<1.9.0` (#42704)

The `__version__` attribute didn't exist prior to `1.9.0` and our check does not work properly.

Tested manually with `pip install -U pydantic==1.8.0`.

---------

Signed-off-by: Edward Oakes <ed.nmi.oakes@gmail.com>
---
 python/ray/_private/pydantic_compat.py   |  2 +-
 python/ray/tests/test_minimal_install.py | 57 ++++++++++++++++++++----
 2 files changed, 50 insertions(+), 9 deletions(-)

diff --git a/python/ray/_private/pydantic_compat.py b/python/ray/_private/pydantic_compat.py
index 8a94e5b63f91..af204f55748d 100644
--- a/python/ray/_private/pydantic_compat.py
+++ b/python/ray/_private/pydantic_compat.py
@@ -28,7 +28,7 @@
 # In pydantic <1.9.0, __version__ attribute is missing, issue ref:
 # https://github.com/pydantic/pydantic/issues/2572, so we need to check
 # the existence prior to comparison.
-elif hasattr(pydantic, "__version__") and packaging.version.parse(
+elif not hasattr(pydantic, "__version__") or packaging.version.parse(
     pydantic.__version__
 ) < packaging.version.parse("2.0"):
     IS_PYDANTIC_2 = False
diff --git a/python/ray/tests/test_minimal_install.py b/python/ray/tests/test_minimal_install.py
index 6006677adae7..c83765d34efd 100644
--- a/python/ray/tests/test_minimal_install.py
+++ b/python/ray/tests/test_minimal_install.py
@@ -3,9 +3,14 @@
 Tests that are specific to minimal installations.
 """
 
-import pytest
+import unittest.mock as mock
+import itertools
+import packaging
 import os
 import sys
+from typing import Dict
+
+import pytest
 
 
 @pytest.mark.skipif(
@@ -29,14 +34,45 @@ def test_correct_python_version():
     )
 
 
+class MockBaseModel:
+    def __init__(self, *args, **kwargs):
+        pass
+
+    def __init_subclass__(self, *args, **kwargs):
+        pass
+
+
+def _make_mock_pydantic_modules(pydantic_version: str) -> Dict:
+    """Make a mock for the `pydantic` module.
+
+    This module requires special handling to:
+        - Make `BaseModel` a class object so type hints work.
+        - Set the `__version__` attribute appropriately.
+        - Also mock `pydantic.v1` for `pydantic >= 2.0`.
+        - Also mock `pydantic.dataclasses`.
+
+    Returns a dict of mocked modules.
+    """
+    mock_modules = {
+        "pydantic": mock.MagicMock(),
+        "pydantic.dataclasses": mock.MagicMock(),
+    }
+    mock_modules["pydantic"].BaseModel = MockBaseModel
+    if packaging.version.parse(pydantic_version) >= packaging.version.parse("1.9.0"):
+        mock_modules["pydantic"].__version__ = pydantic_version
+
+    if packaging.version.parse(pydantic_version) >= packaging.version.parse("2.0.0"):
+        mock_modules["pydantic.v1"] = mock_modules["pydantic"]
+
+    return mock_modules
+
+
+@pytest.mark.parametrize("pydantic_version", ["1.8.0", "1.9.0", "2.0.0"])
 @pytest.mark.skipif(
     os.environ.get("RAY_MINIMAL", "0") != "1",
     reason="Skip unless running in a minimal install.",
 )
-def test_module_import_with_various_non_minimal_deps():
-    import unittest.mock as mock
-    import itertools
-
+def test_module_import_with_various_non_minimal_deps(pydantic_version: str):
     optional_modules = [
         "opencensus",
         "prometheus_client",
@@ -48,9 +84,14 @@ def test_module_import_with_various_non_minimal_deps():
     for i in range(len(optional_modules)):
         for install_modules in itertools.combinations(optional_modules, i):
             print(install_modules)
-            with mock.patch.dict(
-                "sys.modules", {mod: mock.Mock() for mod in install_modules}
-            ):
+            mock_modules = {}
+            for mod in install_modules:
+                if mod == "pydantic":
+                    mock_modules.update(**_make_mock_pydantic_modules(pydantic_version))
+                else:
+                    mock_modules[mod] = mock.MagicMock()
+
+            with mock.patch.dict("sys.modules", mock_modules):
                 from ray.dashboard.utils import get_all_modules
                 from ray.dashboard.utils import DashboardHeadModule