From 0edc31267eeff577bf1052bb41408c588876a34e Mon Sep 17 00:00:00 2001 From: Victor Skvortsov Date: Tue, 28 Jul 2026 15:31:37 +0500 Subject: [PATCH] Upgrade to pydantic v2 --- pyproject.toml | 2 +- src/gpuhunt/providers/oci.py | 45 +++++++++-------- src/tests/providers/test_oci.py | 86 +++++++++++++++++++++++++++------ 3 files changed, 96 insertions(+), 37 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 1a21b9c..739e9b9 100644 --- a/pyproject.toml +++ b/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', diff --git a/src/gpuhunt/providers/oci.py b/src/gpuhunt/providers/oci.py index 92ab2fd..1ab5396 100644 --- a/src/gpuhunt/providers/oci.py +++ b/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): diff --git a/src/tests/providers/test_oci.py b/src/tests/providers/test_oci.py index 5516580..edc9803 100644 --- a/src/tests/providers/test_oci.py +++ b/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