Upgrade to pydantic v2

pyproject.toml

@@ -44,7 +44,7 @@ nebius = [
 ]
 oci = [
     "oci",
-    "pydantic>=1.10.10,<2.0.0",
+    "pydantic>=2.0.0",
 ]
 verda = [
     'verda>=1.22.0',

src/gpuhunt/providers/oci.py

@@ -3,11 +3,11 @@
 import re
 from collections.abc import Iterable
 from dataclasses import asdict, dataclass
-from typing import Annotated
+from typing import Annotated, TypeVar
 
 import oci
 from oci.identity.models import Region
-from pydantic import BaseModel, Field
+from pydantic import BaseModel, BeforeValidator, ConfigDict, Field
 from requests import Session
 from typing_extensions import TypedDict
 
@@ -20,6 +20,14 @@
 COST_ESTIMATOR_URL_TEMPLATE = "https://www.oracle.com/a/ocom/docs/cloudestimator2/data/{resource}"
 COST_ESTIMATOR_REQUEST_TIMEOUT = 10
 
+ResponseModelT = TypeVar("ResponseModelT", bound=BaseModel)
+
+# The Cost Estimator reports some integral quantities as fractional floats, e.g. the disk
+# `qty` of BM.DenseIO.E5.128 is 81.6. Pydantic v1 truncated such values silently, while v2
+# rejects them, failing the whole document. Truncate to keep the v1 behaviour: an unparsable
+# shape we do not even use would otherwise break catalog collection entirely.
+LaxInt = Annotated[int, BeforeValidator(lambda v: int(v) if isinstance(v, float) else v)]
+
 
 class OCICredentials(TypedDict):
     user: str | None
@@ -113,30 +121,28 @@ class CostEstimatorTypeField(BaseModel):
 
 
 class CostEstimatorShapeProduct(BaseModel):
+    model_config = ConfigDict(alias_generator=to_camel_case)
+
     type: CostEstimatorTypeField
     part_number: str
-    qty: int | None
-
-    class Config:
-        alias_generator = to_camel_case
+    qty: LaxInt | None = None
 
 
 class CostEstimatorShape(BaseModel):
+    model_config = ConfigDict(alias_generator=to_camel_case)
+
     name: str
     hidden: bool
     status: str
     allow_preemptible: bool
-    bundle_memory_qty: int | None
-    gpu_qty: int | None
-    gpu_memory_qty: int | None
+    bundle_memory_qty: LaxInt | None = None
+    gpu_qty: LaxInt | None = None
+    gpu_memory_qty: LaxInt | None = None
     processor_type: CostEstimatorTypeField
     shape_type: CostEstimatorTypeField
     sub_type: CostEstimatorTypeField
     products: list[CostEstimatorShapeProduct]
 
-    class Config:
-        alias_generator = to_camel_case
-
     def is_arm_cpu(self):
         is_ampere_gpu = self.sub_type.value == "gpu" and (
             "GPU4" in self.name or "GPU.A10" in self.name
@@ -160,22 +166,21 @@ class CostEstimatorPrice(BaseModel):
 
 
 class CostEstimatorPriceLocalization(BaseModel):
+    model_config = ConfigDict(alias_generator=to_camel_case)
+
     currency_code: str
     prices: list[CostEstimatorPrice]
 
-    class Config:
-        alias_generator = to_camel_case
-
 
 class CostEstimatorProduct(BaseModel):
+    model_config = ConfigDict(alias_generator=to_camel_case)
+
     part_number: str
     billing_model: str
+    # The explicit alias takes priority over the generated `priceType`
     price_type: Annotated[str, Field(alias="pricetype")]
     currency_code_localizations: list[CostEstimatorPriceLocalization]
 
-    class Config:
-        alias_generator = to_camel_case
-
     def find_price_l10n(self, currency_code: str) -> CostEstimatorPriceLocalization | None:
         return next(
             filter(
@@ -203,11 +208,11 @@ def get_shapes(self) -> CostEstimatorShapeList:
     def get_products(self) -> CostEstimatorProductList:
         return self._get("products.json", CostEstimatorProductList)
 
-    def _get(self, resource: str, ResponseModel: type[BaseModel]):
+    def _get(self, resource: str, ResponseModel: type[ResponseModelT]) -> ResponseModelT:
         url = COST_ESTIMATOR_URL_TEMPLATE.format(resource=resource)
         resp = self.session.get(url, timeout=COST_ESTIMATOR_REQUEST_TIMEOUT)
         resp.raise_for_status()
-        return ResponseModel.parse_raw(resp.content)
+        return ResponseModel.model_validate_json(resp.content)
 
 
 class CostEstimatorDataError(Exception):

src/tests/providers/test_oci.py

@@ -1,20 +1,74 @@
+import json
+
 import pytest
 
-from gpuhunt.providers.oci import get_gpu_name
+from gpuhunt.providers.oci import CostEstimatorShapeList, get_gpu_name
+
+# Trimmed excerpts of the real Cost Estimator payload
+# (https://www.oracle.com/a/ocom/docs/cloudestimator2/data/shapes.json).
+# The `disk` product of BM.DenseIO.E5.128 reports a fractional qty of 81.6.
+SHAPES = {
+    "items": [
+        {
+            "name": "BM.DenseIO.E5.128",
+            "hidden": False,
+            "status": "ACTIVE",
+            "allowPreemptible": False,
+            "bundleMemoryQty": 1536,
+            "gpuQty": None,
+            "gpuMemoryQty": None,
+            "processorType": {"value": "amd"},
+            "shapeType": {"value": "bm"},
+            "subType": {"value": "dense"},
+            "products": [
+                {"partNumber": "B98202", "qty": 128, "type": {"value": "ocpu"}},
+                {"partNumber": "B98203", "qty": 1536, "type": {"value": "memory"}},
+                {"partNumber": "B98204", "qty": 81.6, "type": {"value": "disk"}},
+            ],
+        },
+        {
+            "name": "BM.Standard.E4.128",
+            "hidden": False,
+            "status": "ACTIVE",
+            "allowPreemptible": False,
+            "bundleMemoryQty": 2048,
+            "gpuQty": None,
+            "gpuMemoryQty": None,
+            "processorType": {"value": "amd"},
+            "shapeType": {"value": "bm"},
+            "subType": {"value": "standard"},
+            "products": [
+                {"partNumber": "B93113", "qty": 128, "type": {"value": "ocpu"}},
+                {"partNumber": "B93114", "qty": 2048, "type": {"value": "memory"}},
+            ],
+        },
+    ]
+}
+
+
+class TestCostEstimatorShapeList:
+    def test_fractional_qty_is_truncated(self):
+        """Some integral quantities are reported as floats. Pydantic v1 truncated them, v2
+        errors out, which would fail the whole document and not just the shape we do not
+        even use."""
+        shapes = CostEstimatorShapeList.model_validate_json(json.dumps(SHAPES))
+        assert [product.qty for product in shapes.items[0].products] == [128, 1536, 81]
+        assert shapes.items[1].name == "BM.Standard.E4.128"
 
 
-@pytest.mark.parametrize(
-    ("shape_name", "gpu_name"),
-    [
-        ("VM.GPU.A10.2", "A10"),
-        ("BM.GPU.A100-v2.8", "A100"),
-        ("BM.GPU4.8", "A100"),
-        ("VM.GPU3.4", "V100"),
-        ("VM.GPU2.1", "P100"),
-        ("BM.GPU.H100.8", "H100"),
-        ("VM.Standard2.8", None),
-        ("VM.Notgpu.A10", None),
-    ],
-)
-def test_get_gpu_name(shape_name, gpu_name):
-    assert get_gpu_name(shape_name) == gpu_name
+class TestGetGpuName:
+    @pytest.mark.parametrize(
+        ("shape_name", "gpu_name"),
+        [
+            ("VM.GPU.A10.2", "A10"),
+            ("BM.GPU.A100-v2.8", "A100"),
+            ("BM.GPU4.8", "A100"),
+            ("VM.GPU3.4", "V100"),
+            ("VM.GPU2.1", "P100"),
+            ("BM.GPU.H100.8", "H100"),
+            ("VM.Standard2.8", None),
+            ("VM.Notgpu.A10", None),
+        ],
+    )
+    def test_get_gpu_name(self, shape_name, gpu_name):
+        assert get_gpu_name(shape_name) == gpu_name