Make Ray compatible with pydantic>=2.5.0 (#40451)

This change makes Ray compatible with pydantic>=2.5.0. It makes all Ray unit tests use pydantic==2.5.0. It also adds a test suite to check that Ray Serve works with pydantic<2.

Additional Changes:

This PR also updates the following packages:
mypy: 0.982 -> 1.7.0
deepspeed: 0.8.3 -> 0.12.3
Follow-up:

We need to make a follow-up PR that makes Ray release tests run with pydantic>=2.5.0.

---------

Signed-off-by: Edward Oakes <ed.nmi.oakes@gmail.com>
Signed-off-by: Shreyas Krishnaswamy <shrekris@anyscale.com>
Co-authored-by: Shreyas Krishnaswamy <shrekris@anyscale.com>

.buildkite/serve.rayci.yml

@@ -5,8 +5,8 @@ steps:
     wanda: ci/docker/serve.build.wanda.yaml
     depends_on: oss-ci-base_build
 
-  - name: servepydantic2build
-    wanda: ci/docker/servepydantic2.build.wanda.yaml
+  - name: servepydantic1build
+    wanda: ci/docker/servepydantic1.build.wanda.yaml
     depends_on: oss-ci-base_build
 
   - name: servepython311build
@@ -28,7 +28,7 @@ steps:
     depends_on: servebuild
     job_env: forge
 
-  - label: ":ray-serve: serve: pydantic >= 2.0 tests"
+  - label: ":ray-serve: serve: pydantic < 2.0 tests"
     parallelism: 2
     tags: 
       - serve
@@ -39,8 +39,8 @@ steps:
       - bazel run //ci/ray_ci:test_in_docker -- //python/ray/serve/... //python/ray/tests/...  serve 
         --except-tags post_wheel_build,gpu,worker-container,xcommit
         --workers "$${BUILDKITE_PARALLEL_JOB_COUNT}" --worker-id "$${BUILDKITE_PARALLEL_JOB}" --parallelism-per-worker 3
-        --build-name servepydantic2build --test-env=EXPECTED_PYTHON_VERSION=3.8 --test-env=EXPECTED_PYDANTIC_VERSION=2.4.2
-    depends_on: servepydantic2build
+        --build-name servepydantic1build --test-env=EXPECTED_PYTHON_VERSION=3.8 --test-env=EXPECTED_PYDANTIC_VERSION=1.10.12
+    depends_on: servepydantic1build
     job_env: forge
 
   - label: ":ray-serve: serve: Python 3.11 tests"

ci/docker/serve.build.Dockerfile

@@ -4,7 +4,7 @@ ARG DOCKER_IMAGE_BASE_BUILD=cr.ray.io/rayproject/oss-ci-base_build
 FROM $DOCKER_IMAGE_BASE_BUILD
 
 ARG PYTHON_VERSION
-ARG PYDANTIC_BUILD_COMMIT
+ARG PYDANTIC_VERSION
 
 # Unset dind settings; we are using the host's docker daemon.
 ENV DOCKER_TLS_CERTDIR=
@@ -45,8 +45,8 @@ pip install -c python/requirements_compiled.txt aioboto3
 git clone https://github.com/wg/wrk.git /tmp/wrk && pushd /tmp/wrk && make -j && sudo cp wrk /usr/local/bin && popd
 
 # Install custom Pydantic version if requested.
-if [[ -n "${PYDANTIC_BUILD_COMMIT-}" ]]; then 
-  pip install -U git+https://github.com/pydantic/pydantic@$PYDANTIC_BUILD_COMMIT
+if [[ -n "${PYDANTIC_VERSION-}" ]]; then 
+  pip install -U pydantic==$PYDANTIC_VERSION
 else 
   echo "Not installing Pydantic from source"
 fi

ci/docker/servepydantic1.build.wanda.yaml

@@ -0,0 +1,11 @@
+name: "servepydantic1build"
+froms: ["cr.ray.io/rayproject/oss-ci-base_build"]
+dockerfile: ci/docker/serve.build.Dockerfile
+srcs:
+  - python/requirements.txt
+  - python/requirements_compiled.txt
+  - python/requirements/test-requirements.txt
+build_args:
+  - PYDANTIC_VERSION=1.10.12
+tags:
+  - cr.ray.io/rayproject/servepydantic1build

doc/requirements-doc.txt

@@ -19,10 +19,10 @@ sphinx-sitemap==2.2.0
 sphinxcontrib-redoc==1.6.0
 sphinx-tabs==3.4.0
 sphinx-remove-toctrees==0.0.3
-autodoc_pydantic==1.6.1
+autodoc_pydantic==2.0.1
 sphinx_design==0.4.1
 
-pydantic<2 # Pydantic is required by autodoc_pydantic, but must be <2 for ray
+pydantic!=2.0.*,!=2.1.*,!=2.2.*,!=2.3.*,!=2.4.*,<3
 
 # MyST
 myst-parser==0.15.2

doc/source/_templates/autosummary/class_with_inherited_members.rst

@@ -0,0 +1,7 @@
+{{ fullname | escape | underline}}
+
+.. currentmodule:: {{ module }}
+
+.. autoclass:: {{ objname }}
+    :members:
+    :inherited-members: pydantic.BaseModel

doc/source/ray-overview/pip_freeze_ray-ml-py39-cpu.txt

@@ -96,7 +96,7 @@ databricks-cli==0.18.0
 datasets==2.0.0
 debugpy==1.8.0
 decorator==5.1.1
-deepspeed==0.8.3
+deepspeed==0.12.3
 defusedxml==0.7.1
 Deprecated==1.2.14
 dill==0.3.7

doc/source/serve/api/index.md

@@ -354,6 +354,7 @@ Content-Type: application/json
 .. autosummary::
    :nosignatures:
    :toctree: doc/
+   :template: autosummary/class_with_inherited_members.rst
 
    schema.ServeInstanceDetails
    schema.ApplicationDetails

python/ray/_private/pydantic_compat.py

@@ -57,3 +57,35 @@
         root_validator,
         validator,
     )
+
+
+def register_pydantic_serializers(serialization_context):
+    if not PYDANTIC_INSTALLED:
+        return
+
+    if IS_PYDANTIC_2:
+        # TODO(edoakes): compare against the version that has the fixes.
+        from pydantic.v1.fields import ModelField
+    else:
+        from pydantic.fields import ModelField
+
+    # Pydantic's Cython validators are not serializable.
+    # https://github.com/cloudpipe/cloudpickle/issues/408
+    serialization_context._register_cloudpickle_serializer(
+        ModelField,
+        custom_serializer=lambda o: {
+            "name": o.name,
+            # outer_type_ is the original type for ModelFields,
+            # while type_ can be updated later with the nested type
+            # like int for List[int].
+            "type_": o.outer_type_,
+            "class_validators": o.class_validators,
+            "model_config": o.model_config,
+            "default": o.default,
+            "default_factory": o.default_factory,
+            "required": o.required,
+            "alias": o.alias,
+            "field_info": o.field_info,
+        },
+        custom_deserializer=lambda kwargs: ModelField(**kwargs),
+    )

python/ray/autoscaler/_private/commands.py

@@ -112,7 +112,9 @@ def try_reload_log_state(provider_config: Dict[str, Any], log_state: dict) -> No
         return reload_log_state(log_state)
 
 
-def debug_status(status, error, verbose: bool = False, address: str = None) -> str:
+def debug_status(
+    status, error, verbose: bool = False, address: Optional[str] = None
+) -> str:
     """
     Return a debug string for the autoscaler.
 
@@ -1058,7 +1060,7 @@ def attach_cluster(
 def exec_cluster(
     config_file: str,
     *,
-    cmd: str = None,
+    cmd: Optional[str] = None,
     run_env: str = "auto",
     screen: bool = False,
     tmux: bool = False,

python/ray/autoscaler/sdk/sdk.py

@@ -128,7 +128,7 @@ def rsync(
     source: Optional[str],
     target: Optional[str],
     down: bool,
-    ip_address: str = None,
+    ip_address: Optional[str] = None,
     use_internal_ip: bool = False,
     no_config_cache: bool = False,
     should_bootstrap: bool = True

python/ray/tests/test_client.py

@@ -6,9 +6,12 @@
 import threading
 import time
 from unittest.mock import Mock, patch
+from typing import Type
 
 import numpy as np
 import pytest
+from pydantic import BaseModel as BaseModelV2
+from pydantic.v1 import BaseModel as BaseModelV1
 
 import ray.cloudpickle as cloudpickle
 import ray.util.client.server.server as ray_client_server
@@ -686,10 +689,9 @@ def test_client_gpu_ids(call_ray_start_shared):
             assert ray.get_gpu_ids() == []
 
 
-def test_client_serialize_addon(call_ray_start_shared):
-    import pydantic
-
-    class User(pydantic.BaseModel):
+@pytest.mark.parametrize("BaseModel", [BaseModelV1, BaseModelV2])
+def test_client_serialize_addon(call_ray_start_shared, BaseModel: Type):
+    class User(BaseModel):
         name: str
 
     with ray_start_client_server_for_address(call_ray_start_shared) as ray: