pydantic 2 updates

container/conda.yml

@@ -10,7 +10,7 @@ dependencies:
   - ae5-admin:mlflow-adsp>=0.5,<1.0
   - defaults:tqdm
   - defaults:psutil
-  - defaults:pydantic<2.0
+  - defaults:pydantic>2.0
   - defaults:click
   - defaults:requests
   - defaults:pip

mlflow_adsp/common/scheduler.py

@@ -166,8 +166,7 @@ def execute_step(step: Step) -> ADSPSubmittedRun:
             An instance of `SubmittedRun` for the requested workflow step run.
         """
 
-        step_dict: Dict = step.dict(by_alias=False)
-
+        step_dict: Dict = step.model_dump(by_alias=False)
         message: str = f"Launching new background job for: {step_dict}"
         logger.debug(message)
 

mlflow_adsp/contracts/dto/base_model.py

@@ -30,11 +30,9 @@ def lower_camel_case(string: str) -> str:
 class BaseModel(PydanticBaseModel):
     """BaseModel [Pydantic] Over-Ride"""
 
-    # https://pydantic-docs.helpmanual.io/usage/model_config/#options
+    # Pydantic Config Over-Ride
     class Config:
-        """Pydantic Config Over-Ride"""
-
         alias_generator = lower_camel_case
-        allow_population_by_field_name = True
+        populate_by_name = True
         arbitrary_types_allowed = True
         use_enum_values = True

mlflow_adsp/services/endpoint_manager.py

@@ -53,9 +53,6 @@ def __init__(self, client: MlflowClient, params: EndpointManagerParameters):
         self.client = client
 
         self.params = params
-        params_dict: dict = self.params.dict(by_alias=True)
-        logger.info(params_dict)
-
         self.metadata = TargetMetadata(model_uri=self.params.model_uri)
         if self.metadata.reloadable:
             self.version = self._get_latest_version()
@@ -103,7 +100,7 @@ def _poll_network_service(self, process: subprocess.Popen) -> bool:
 
         try:
             version_endpoint_url: str = f"http://{self.params.host}:{self.params.port}/version"
-            response: Response = requests.get(url=version_endpoint_url)
+            response: Response = requests.get(url=version_endpoint_url, timeout=self.params.timeout)
             if response.status_code != 200:
                 logger.debug("Service not yet healthy, waiting ..")
                 EndpointManager._proc_comm(process=process, timeout=self.params.timeout)

setup.py

@@ -10,7 +10,15 @@
     packages=setuptools.find_packages(),
     author="Joshua C. Burt",
     python_requires=">=3.8",
-    install_requires=["mlflow>=2.3.0", "ae5-tools>=0.6.1", "psutil", "pydantic<2.0", "tqdm", "click", "requests"],
+    install_requires=[
+        "mlflow>=2.3.0",
+        "ae5-tools>=0.7,<1.0",
+        "psutil",
+        "pydantic>=2.0,<3",
+        "tqdm",
+        "click",
+        "requests",
+    ],
     entry_points={
         # Define a MLFlow Project Backend plugin called 'adsp'
         "mlflow.project_backend": "adsp=mlflow_adsp:adsp_backend_builder"