Add CatalogItem.provider_data + reserved GCP A4 (#184)

Add a catalog item field for holding arbitrary
provider-specific data. As an example, use this
field for reserved GCP A4 instances.

Below is a comparison of different field types
that could be used for this field in both gpuhunt
and dstack. This commit suggests using typed
dicts.

JSON string (parsing into a Pydantic model on the
provider/backend side)

Cons:
- Inefficient, inconvenient, and error-prone for
  writing - unnecessary serialization and
  deserialization when updating an attribute of an
  already serialized object.
- Possibility of incorrect usage (writing non-JSON
  data).

`dict` / `TypedDict`

Cons:
- No validation, which means potential errors
  occur at attribute access time rather than at
  model loading time.

Pydantic models with a `type` discriminator

Cons:
- Extra difficulties maintaining backward
  compatibility, as the models are passed from
  gpuhunt to dstack server, from server to client,
  and from client to server, all with validation.
- Duplication of backend type in the
  backend-specific field and in other fields of
  the offer or catalog (e.g.,
  `InstanceOffer.backend` and
  `InstanceOffer.backend_data.type`).
- Discriminators require declaring all possible
  discriminator values, which in the future will
  hinder the transition to a more modular
  architecture with backend plugins.
- Backward compatibility issues when a new
  discriminator value (a new backend) is
  introduced.

Pydantic models + custom deserialization logic
(e.g., custom `InstanceOffer` deserializer that
determines the `InstanceOffer.backend_data` model
based on InstanceOffer.backend)

Cons:
- Extra difficulties maintaining backward
  compatibility as the models are passed from
  gpuhunt to dstack server, from server to client,
  and from client to server, all with validation.
- The need to duplicate deserialization logic in
  all model that hold the backend-specific field -
  at least in RawCatalogItem and `InstanceOffer`.

Author: jvstme

src/gpuhunt/_internal/models.py

@@ -1,4 +1,5 @@
 import enum
+import json
 from collections.abc import Container
 from dataclasses import asdict, dataclass, field, fields
 from typing import (
@@ -9,6 +10,17 @@
 
 from gpuhunt._internal.utils import empty_as_none
 
+JSONType = Union[
+    None,
+    bool,
+    int,
+    float,
+    str,
+    list["JSONType"],
+    "JSONObject",
+]
+JSONObject = dict[str, JSONType]
+
 
 def bool_loader(x: Union[bool, str]) -> bool:
     if isinstance(x, bool):
@@ -73,6 +85,7 @@ class RawCatalogItem:
     gpu_vendor: Optional[str] = None
     flags: list[str] = field(default_factory=list)
     cpu_arch: Optional[str] = None
+    provider_data: JSONObject = field(default_factory=dict)
 
     def __post_init__(self) -> None:
         self._process_gpu_vendor()
@@ -121,12 +134,14 @@ def from_dict(v: dict) -> "RawCatalogItem":
             spot=empty_as_none(v.get("spot"), loader=bool_loader),
             disk_size=empty_as_none(v.get("disk_size"), loader=float),
             flags=v.get("flags", "").split(),
+            provider_data=json.loads(v.get("provider_data", "{}")),
         )
 
     def dict(self) -> dict[str, Union[str, int, float, bool, None]]:
         return {
             **asdict(self),
             "flags": " ".join(self.flags),
+            "provider_data": json.dumps(self.provider_data),
         }
 
 
@@ -153,6 +168,8 @@ class CatalogItem:
             will have to request this flag explicitly to get the catalog item.
             If you are adding a new provider, leave the flags empty.
             Flag names should be in kebab-case.
+        provider_data: dict with provider-specific properties.
+            Prefer defining a TypedDict within provider implementation.
     """
 
     instance_name: str
@@ -169,6 +186,7 @@ class CatalogItem:
     gpu_vendor: Optional[AcceleratorVendor] = None
     flags: list[str] = field(default_factory=list)
     cpu_arch: Optional[CPUArchitecture] = None
+    provider_data: JSONObject = field(default_factory=dict)
 
     def __post_init__(self) -> None:
         self._process_gpu_vendor()

src/gpuhunt/providers/gcp.py

@@ -1,4 +1,5 @@
 import copy
+import enum
 import importlib.resources
 import json
 import logging
@@ -7,14 +8,15 @@
 from collections.abc import Iterable
 from concurrent.futures import ThreadPoolExecutor, as_completed
 from dataclasses import dataclass
-from typing import Optional
+from typing import Optional, cast
 
 import google.cloud.billing_v1 as billing_v1
 import google.cloud.compute_v1 as compute_v1
 from google.cloud import tpu_v2
 from google.cloud.billing_v1 import CloudCatalogClient, ListSkusRequest
 from google.cloud.billing_v1.types.cloud_catalog import Sku
 from google.cloud.location import locations_pb2
+from typing_extensions import NotRequired, TypedDict
 
 from gpuhunt._internal.models import AcceleratorVendor, QueryFilter, RawCatalogItem
 from gpuhunt.providers import AbstractProvider
@@ -243,14 +245,17 @@ def fill_prices(self, instances: list[RawCatalogItem]) -> list[RawCatalogItem]:
 
         offers = []
         for instance in instances:
-            for spot in (False, True):
-                price = prices.get_instance_price(instance, spot)
+            for capacity_type in CapacityType:
+                price = prices.get_instance_price(instance, capacity_type)
                 if price is None:
                     continue
 
                 offer = copy.deepcopy(instance)
                 offer.price = round(price, 6)
-                offer.spot = spot
+                offer.spot = capacity_type is CapacityType.SPOT
+                cast(GCPCatalogItemProviderData, offer.provider_data)["is_dws_calendar_mode"] = (
+                    capacity_type is CapacityType.DWS_CALENDAR_MODE
+                )
                 offers.append(offer)
         return offers
 
@@ -307,6 +312,9 @@ def filter(cls, offers: list[RawCatalogItem]) -> list[RawCatalogItem]:
                 # Filter out on-demand offers that are not actually available on demand.
                 # https://cloud.google.com/compute/docs/accelerator-optimized-machines#consumption_option_availability_by_machine_type
                 i.spot == False
+                and not cast(GCPCatalogItemProviderData, i.provider_data).get(
+                    "is_dws_calendar_mode"
+                )
                 and (
                     i.instance_name.startswith("a4x-")
                     or i.instance_name.startswith("a4-")
@@ -320,17 +328,27 @@ def filter(cls, offers: list[RawCatalogItem]) -> list[RawCatalogItem]:
         ]
 
 
-RegionSpot = tuple[str, bool]
-PricePerRegionSpot = dict[RegionSpot, float]
+class GCPCatalogItemProviderData(TypedDict):
+    is_dws_calendar_mode: NotRequired[bool]
+
+
+class CapacityType(enum.Enum):
+    ON_DEMAND = enum.auto()
+    SPOT = enum.auto()
+    DWS_CALENDAR_MODE = enum.auto()
+
+
+RegionCapacityType = tuple[str, CapacityType]
+PricePerRegionCapacityType = dict[RegionCapacityType, float]
 
 
 class Prices:
     def __init__(self):
-        self.cpu: defaultdict[str, PricePerRegionSpot] = defaultdict(dict)
-        self.gpu: defaultdict[str, PricePerRegionSpot] = defaultdict(dict)
-        self.ram: defaultdict[str, PricePerRegionSpot] = defaultdict(dict)
-        self.local_ssd: PricePerRegionSpot = dict()
-        self.gpu_slice: defaultdict[str, PricePerRegionSpot] = defaultdict(dict)
+        self.cpu: defaultdict[str, PricePerRegionCapacityType] = defaultdict(dict)
+        self.gpu: defaultdict[str, PricePerRegionCapacityType] = defaultdict(dict)
+        self.ram: defaultdict[str, PricePerRegionCapacityType] = defaultdict(dict)
+        self.local_ssd: PricePerRegionCapacityType = dict()
+        self.gpu_slice: defaultdict[str, PricePerRegionCapacityType] = defaultdict(dict)
 
     def add_skus(self, skus: Iterable[Sku]) -> None:
         for sku in skus:
@@ -395,7 +413,9 @@ def add_compute_sku(self, sku: Sku) -> None:
         self._add_price(sku, resource_prices[family], price)
 
     def add_compute_gpu_slice_sku(self, sku: Sku) -> None:
-        if sku.description.startswith("Spot Preemptible A4 Nvidia B200"):
+        if sku.description.startswith(
+            "Spot Preemptible A4 Nvidia B200"
+        ) or sku.description.startswith("DWS Calendar Mode A4 Nvidia B200"):
             gpu = "nvidia-b200"
         else:
             return
@@ -413,34 +433,43 @@ def _calculate_sku_price(sku: Sku) -> float:
         return price.units + price.nanos / 1e9
 
     @staticmethod
-    def _add_price(sku: Sku, family_prices: PricePerRegionSpot, price: float) -> None:
-        spot = sku.category.usage_type == "Preemptible"
+    def _add_price(sku: Sku, family_prices: PricePerRegionCapacityType, price: float) -> None:
+        if sku.category.usage_type == "Preemptible":
+            capacity_type = CapacityType.SPOT
+        elif "DWS Calendar Mode" in sku.description:
+            capacity_type = CapacityType.DWS_CALENDAR_MODE
+        else:
+            capacity_type = CapacityType.ON_DEMAND
         for region in sku.service_regions:
-            family_prices[(region, spot)] = price
+            family_prices[(region, capacity_type)] = price
 
-    def get_instance_price(self, instance: RawCatalogItem, spot: bool) -> Optional[float]:
+    def get_instance_price(
+        self, instance: RawCatalogItem, capacity_type: CapacityType
+    ) -> Optional[float]:
         vm_family = self.get_vm_family(instance.instance_name)
         if vm_family in ["g1", "f1", "m2"]:  # shared-core and reservation-only
             return None
 
-        region_spot = (instance.location[:-2], spot)
+        region_capacity_type = (instance.location[:-2], capacity_type)
 
         # For some instances, the price is proportional to the number of GPUs
-        if instance.gpu_name and region_spot in self.gpu_slice[instance.gpu_name]:
-            return instance.gpu_count * self.gpu_slice[instance.gpu_name][region_spot]
+        if instance.gpu_name and region_capacity_type in self.gpu_slice[instance.gpu_name]:
+            return instance.gpu_count * self.gpu_slice[instance.gpu_name][region_capacity_type]
 
         # For others, the price consists of several components
         price = 0
-        if region_spot not in self.cpu[vm_family]:
+        if region_capacity_type not in self.cpu[vm_family]:
             return None
-        price += instance.cpu * self.cpu[vm_family][region_spot]
-        price += instance.memory * self.ram[vm_family][region_spot]
+        price += instance.cpu * self.cpu[vm_family][region_capacity_type]
+        price += instance.memory * self.ram[vm_family][region_capacity_type]
         if instance.gpu_name:
-            if region_spot not in self.gpu[instance.gpu_name]:
+            if region_capacity_type not in self.gpu[instance.gpu_name]:
                 return None
-            price += instance.gpu_count * self.gpu[instance.gpu_name][region_spot]
+            price += instance.gpu_count * self.gpu[instance.gpu_name][region_capacity_type]
         if instance.instance_name in local_ssd_sizes_gib:
-            price += local_ssd_sizes_gib[instance.instance_name] * self.local_ssd[region_spot]
+            price += (
+                local_ssd_sizes_gib[instance.instance_name] * self.local_ssd[region_capacity_type]
+            )
 
         return price
 
@@ -454,6 +483,8 @@ def get_vm_family(instance_name: str) -> str:
 
 def set_flags(catalog_items: list[RawCatalogItem]) -> None:
     for item in catalog_items:
+        if cast(GCPCatalogItemProviderData, item.provider_data).get("is_dws_calendar_mode"):
+            item.flags.append("gcp-dws-calendar-mode")
         if item.instance_name.startswith("a4-"):
             item.flags.append("gcp-a4")
         elif item.instance_name.startswith("g4-standard-") and item.price == 0:

src/tests/_internal/test_models.py

@@ -151,6 +151,7 @@ def test_raw_catalog_item_to_from_dict() -> None:
         spot=False,
         disk_size=100.0,
         flags=["f1", "f2", "f3"],
+        provider_data={"custom_prop": 42},
     )
     item_dict = item.dict()
     assert item_dict == {
@@ -167,5 +168,6 @@ def test_raw_catalog_item_to_from_dict() -> None:
         "spot": False,
         "disk_size": 100.0,
         "flags": "f1 f2 f3",
+        "provider_data": '{"custom_prop": 42}',
     }
     assert RawCatalogItem.from_dict(item_dict) == item