Introduce pydantic_v1 compatibility module for pydantic>=2.0.0 support

mii/aml_related/templates.py

@@ -393,6 +393,6 @@
 """torch>=2.0.0
 grpcio
 grpcio-tools
-pydantic<2.0.0
+pydantic
 asyncio
 """

mii/config.py

@@ -6,9 +6,9 @@
 import os
 import string
 from typing import List, Optional, Dict, Any
-from pydantic import validator, root_validator, Field
 import mii
 from .constants import DeploymentType, TaskType, MII_MODEL_PATH_DEFAULT
+from .pydantic_v1 import validator, root_validator, Field
 
 from deepspeed.runtime.config_utils import DeepSpeedConfigModel
 from deepspeed.inference.config import DtypeEnum

mii/pydantic_v1.py

@@ -0,0 +1,16 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+"""Pydantic v1 compatibility module.
+
+Pydantic v2 introduced breaking changes that hinder its adoption:
+https://docs.pydantic.dev/latest/migration/. To provide deepspeed users the option to
+migrate to pydantic v2 on their own timeline, deepspeed uses this compatibility module
+as a pydantic-version-agnostic alias for pydantic's v1 API.
+"""
+
+try:
+    from pydantic.v1 import *  # noqa: F401
+except ImportError:
+    from pydantic import *  # noqa: F401

tests/test_config.py

@@ -4,9 +4,9 @@
 # DeepSpeed Team
 
 import pytest
-import pydantic
 
 import mii
+from mii import pydantic_v1
 
 
 @pytest.mark.parametrize("port_number", [12345])
@@ -23,7 +23,7 @@ def test_base_configs(deployment_name, mii_config, model_config):
 @pytest.mark.parametrize("port_number", ["fail"])
 @pytest.mark.parametrize("tensor_parallel", [3.5])
 def test_base_configs_literalfail(deployment_name, mii_config, model_config):
-    with pytest.raises(pydantic.ValidationError):
+    with pytest.raises(pydantic_v1.ValidationError):
         mii_config["deployment_name"] = deployment_name
         mii_config["model_config"] = model_config
         mii_config = mii.config.MIIConfig(**mii_config)

tests/test_deployment_options.py

@@ -6,8 +6,8 @@
 import pytest
 import json
 import requests
-import pydantic
 import mii
+from mii import pydantic_v1
 
 
 @pytest.mark.deepspeed
@@ -81,7 +81,7 @@ def test_zero_config(deployment, query):
 
 
 @pytest.mark.deepspeed
-@pytest.mark.parametrize("expected_failure", [pydantic.ValidationError])
+@pytest.mark.parametrize("expected_failure", [pydantic_v1.ValidationError])
 @pytest.mark.parametrize(
     "enable_deepspeed, enable_zero, dtype",
     [(True,