Bump pydantic to 2.8.2 (#561)

* Bump pydantic to 2.8.2

* Run pydantic-bump

* More updates

* Fix openapi generation

* Fix tests

* black formatting

* Fix integration tests

Author: dmchoiboi

.circleci/config.yml

@@ -74,7 +74,7 @@ jobs:
     resource_class: small
     parallelism: 1
     steps:
-      - add_ssh_keys:  # gives write access to CircleCI worker
+      - add_ssh_keys: # gives write access to CircleCI worker
           fingerprints:
             - "76:0c:1b:9e:e3:6a:c3:5c:6f:24:91:ef:7c:54:d2:7a"
       - checkout # checkout source code to working directory
@@ -157,10 +157,10 @@ jobs:
             DOCKER_BUILDKIT=1 docker build -f model-engine/model_engine_server/inference/pytorch_or_tf.user.Dockerfile \
             --build-arg BASE_IMAGE=temp:1.7.1-cuda11.0-cudnn8-runtime-$CIRCLE_SHA1 \
             --build-arg REQUIREMENTS_FILE="$CIRCLE_SHA1-requirements.txt" \
-            -t $CIRCLECI_AWS_ACCOUNT_ID.dkr.ecr.us-west-2.amazonaws.com/hosted-model-inference/async-pytorch:1.7.1-cuda11.0-cudnn8-runtime-$CIRCLE_SHA1-021694 .
+            -t $CIRCLECI_AWS_ACCOUNT_ID.dkr.ecr.us-west-2.amazonaws.com/hosted-model-inference/async-pytorch:1.7.1-cuda11.0-cudnn8-runtime-$CIRCLE_SHA1-b8c25b .
             rm $CIRCLE_SHA1-requirements.txt
 
-            minikube --logtostderr -v 1 image load $CIRCLECI_AWS_ACCOUNT_ID.dkr.ecr.us-west-2.amazonaws.com/hosted-model-inference/async-pytorch:1.7.1-cuda11.0-cudnn8-runtime-$CIRCLE_SHA1-021694
+            minikube --logtostderr -v 1 image load $CIRCLECI_AWS_ACCOUNT_ID.dkr.ecr.us-west-2.amazonaws.com/hosted-model-inference/async-pytorch:1.7.1-cuda11.0-cudnn8-runtime-$CIRCLE_SHA1-b8c25b
       - run:
           name: Install helm chart
           command: |
@@ -207,23 +207,23 @@ commands:
   install_server:
     description: Installs LLM Engine server
     steps:
-     - python/install-packages:
-         pkg-manager: pip
-         app-dir: model-engine
-     - python/install-packages:
-         pkg-manager: pip
-         app-dir: model-engine
-         pip-dependency-file: requirements-test.txt
-     - python/install-packages:
-         pkg-manager: pip
-         app-dir: model-engine
-         pip-dependency-file: requirements_override.txt
-     - run:
-         name: Install Server
-         command: |
-           pushd model-engine
-           pip install -e .
-           popd
+      - python/install-packages:
+          pkg-manager: pip
+          app-dir: model-engine
+      - python/install-packages:
+          pkg-manager: pip
+          app-dir: model-engine
+          pip-dependency-file: requirements-test.txt
+      - python/install-packages:
+          pkg-manager: pip
+          app-dir: model-engine
+          pip-dependency-file: requirements_override.txt
+      - run:
+          name: Install Server
+          command: |
+            pushd model-engine
+            pip install -e .
+            popd
   install_client:
     description: Install LLM Engine client
     steps:

clients/python/llmengine/data_types.py

@@ -6,9 +6,10 @@
 from enum import Enum
 from typing import Any, Dict, List, Literal, Optional, Union
 
-import pydantic
+from pydantic.version import VERSION as PYDANTIC_VERSION
 
-if int(pydantic.__version__.split(".")[0]) > 1:
+PYDANTIC_V2 = PYDANTIC_VERSION.startswith("2.")
+if PYDANTIC_V2:
     from pydantic.v1 import BaseModel, Field, HttpUrl
 else:
     from pydantic import BaseModel, Field, HttpUrl  # type: ignore

integration_tests/rest_api_utils.py

@@ -59,7 +59,12 @@ def my_model(**keyword_args):
             "framework_type": "pytorch",
             "pytorch_image_tag": "1.7.1-cuda11.0-cudnn8-runtime",
         },
-        "requirements": ["cloudpickle==2.1.0", "pyyaml==6.0"],
+        "requirements": [
+            "cloudpickle==2.1.0",
+            "pyyaml==6.0",
+            "pydantic==2.8.2",
+            "fastapi==0.110.0",
+        ],
         "location": "s3://model-engine-integration-tests/model_bundles/echo_bundle",
     },
 }

integration_tests/test_endpoints.py

@@ -232,7 +232,7 @@ def test_sync_streaming_model_endpoint(capsys):
             for response in task_responses:
                 assert (
                     response.strip()
-                    == 'data: {"status": "SUCCESS", "result": {"result": {"y": 1}}, "traceback": null}'
+                    == 'data: {"status":"SUCCESS","result":{"result":{"y":1}},"traceback":null}'
                 )
         finally:
             delete_model_endpoint(create_endpoint_request["name"], user)

model-engine/model_engine_server/common/dtos/batch_jobs.py

@@ -13,20 +13,21 @@
     GpuType,
     StorageSpecificationType,
 )
-from pydantic import BaseModel, root_validator
+from pydantic import BaseModel, ConfigDict, model_validator
 
 
 class CreateBatchJobResourceRequests(BaseModel):
-    cpus: Optional[CpuSpecificationType]
-    memory: Optional[StorageSpecificationType]
-    gpus: Optional[int]
-    gpu_type: Optional[GpuType]
-    storage: Optional[StorageSpecificationType]
-    max_workers: Optional[int]
-    per_worker: Optional[int]
+    cpus: Optional[CpuSpecificationType] = None
+    memory: Optional[StorageSpecificationType] = None
+    gpus: Optional[int] = None
+    gpu_type: Optional[GpuType] = None
+    storage: Optional[StorageSpecificationType] = None
+    max_workers: Optional[int] = None
+    per_worker: Optional[int] = None
 
 
 class CreateBatchJobV1Request(BaseModel):
+    model_config = ConfigDict(protected_namespaces=())
     model_bundle_id: str
     input_path: str
     serialization_format: BatchJobSerializationFormat
@@ -41,10 +42,10 @@ class CreateBatchJobV1Response(BaseModel):
 
 class GetBatchJobV1Response(BaseModel):
     status: BatchJobStatus
-    result: Optional[str]
+    result: Optional[str] = None
     duration: timedelta
-    num_tasks_pending: Optional[int]
-    num_tasks_completed: Optional[int]
+    num_tasks_pending: Optional[int] = None
+    num_tasks_completed: Optional[int] = None
 
 
 class UpdateBatchJobV1Request(BaseModel):
@@ -64,9 +65,7 @@ class CreateDockerImageBatchJobResourceRequests(BaseModel):
     gpus: Optional[int] = None
     gpu_type: Optional[GpuType] = None
     storage: Optional[StorageSpecificationType] = None
-
-    class Config:
-        orm_mode = True
+    model_config = ConfigDict(from_attributes=True)
 
     @classmethod
     def merge_requests(
@@ -93,7 +92,7 @@ def common_requests(
 class CreateDockerImageBatchJobV1Request(BaseModel):
     docker_image_batch_job_bundle_name: Optional[str] = None
     docker_image_batch_job_bundle_id: Optional[str] = None
-    job_config: Optional[Dict[str, Any]]
+    job_config: Optional[Dict[str, Any]] = None
     # TODO also expose a separate argument to pass an s3file to the job, as opposed to job_config
     labels: Dict[str, str]  # TODO this probably should go in the bundle
 
@@ -103,7 +102,7 @@ class CreateDockerImageBatchJobV1Request(BaseModel):
 
     override_job_max_runtime_s: Optional[int] = None
 
-    @root_validator
+    @model_validator(mode="before")
     def exactly_one_name_or_id(cls, values):
         bundle_name = values.get("docker_image_batch_job_bundle_name")
         bundle_id = values.get("docker_image_batch_job_bundle_id")
@@ -166,16 +165,14 @@ class DockerImageBatchJobBundleV1Response(BaseModel):
     image_tag: str
     command: List[str]
     env: Dict[str, str]
-    mount_location: Optional[str]
-    cpus: Optional[str]
-    memory: Optional[str]
-    storage: Optional[str]
-    gpus: Optional[int]
-    gpu_type: Optional[str]
-    public: Optional[bool]
-
-    class Config:
-        orm_mode = True
+    mount_location: Optional[str] = None
+    cpus: Optional[str] = None
+    memory: Optional[str] = None
+    storage: Optional[str] = None
+    gpus: Optional[int] = None
+    gpu_type: Optional[str] = None
+    public: Optional[bool] = None
+    model_config = ConfigDict(from_attributes=True)
 
 
 class ListDockerImageBatchJobBundleV1Response(BaseModel):

model-engine/model_engine_server/common/dtos/core.py

@@ -0,0 +1,17 @@
+from pydantic import BaseModel, BeforeValidator, ConfigDict, HttpUrl, TypeAdapter
+from typing_extensions import Annotated
+
+# See: https://github.com/pydantic/pydantic/issues/7186
+# pydantic v2 doesn't treat HttpUrl the same way as in v1 which causes various issue
+# This is an attempt to make it behave as similar as possible
+HttpUrlTypeAdapter = TypeAdapter(HttpUrl)
+HttpUrlStr = Annotated[
+    str,
+    BeforeValidator(lambda value: HttpUrlTypeAdapter.validate_python(value) and value),
+]
+
+
+class LLMEngineModel(BaseModel):
+    """Common pydantic configurations for model engine"""
+
+    model_config = ConfigDict(protected_namespaces=())

model-engine/model_engine_server/common/dtos/docker_repository.py

@@ -10,8 +10,8 @@ class BuildImageRequest(BaseModel):
     base_path: str
     dockerfile: str
     base_image: str
-    requirements_folder: Optional[str]
-    substitution_args: Optional[Dict[str, str]]
+    requirements_folder: Optional[str] = None
+    substitution_args: Optional[Dict[str, str]] = None
 
 
 class BuildImageResponse(BaseModel):

model-engine/model_engine_server/common/dtos/endpoint_builder.py

@@ -20,19 +20,19 @@ class BuildEndpointRequest(BaseModel):
     cpus: CpuSpecificationType
     gpus: int
     memory: StorageSpecificationType
-    gpu_type: Optional[GpuType]
-    storage: Optional[StorageSpecificationType]
+    gpu_type: Optional[GpuType] = None
+    storage: Optional[StorageSpecificationType] = None
     optimize_costs: bool
     aws_role: str
     results_s3_bucket: str
-    child_fn_info: Optional[Dict[str, Any]]  # TODO: remove this if we don't need it.
-    post_inference_hooks: Optional[List[str]]
+    child_fn_info: Optional[Dict[str, Any]] = None  # TODO: remove this if we don't need it.
+    post_inference_hooks: Optional[List[str]] = None
     labels: Dict[str, str]
-    billing_tags: Optional[Dict[str, Any]]
+    billing_tags: Optional[Dict[str, Any]] = None
     prewarm: bool = True
-    high_priority: Optional[bool]
-    default_callback_url: Optional[str]
-    default_callback_auth: Optional[CallbackAuth]
+    high_priority: Optional[bool] = None
+    default_callback_url: Optional[str] = None
+    default_callback_auth: Optional[CallbackAuth] = None
 
 
 class BuildEndpointStatus(str, Enum):