Get flavor GPUs

scripts/azure/api/list-flavors.sh

@@ -32,15 +32,26 @@ json=$(curl --request ${REQUEST}\
      --data-raw "${BODY}" \
      ${URL} 2>/dev/null)
 
-table=$(echo "$json" | jq -r '.flavors[] | [.id, .name, .cpus, .price] | @tsv')
+table=$(echo "$json" | jq -r '.flavors[] | [.id, .name, .cpus, .gpus, .price] | @tsv')
 
 id_width=$(echo "$table" | awk -F'\t' '{print length($1)}' | sort -nr | head -1)
 name_width=$(echo "$table" | awk -F'\t' '{print length($2)}' | sort -nr | head -1)
 cpus_width=$(echo "$table" | awk -F'\t' '{print length($3)}' | sort -nr | head -1)
-price_width=$(echo "$table" | awk -F'\t' '{print length($4)}' | sort -nr | head -1)
-
-printf "%-${id_width}s  %-${name_width}s %-${cpus_width}s %-${price_width}s\n" "ID" "Name" "CPUs" "Price"
-printf "%${id_width}s  %${name_width}s %${cpus_width}s %${price_width}s\n" "$(printf '%*s' "$id_width" | tr ' ' '-')" "$(printf '%*s' "$name_width" | tr ' ' '-')" "$(printf '%*s' "$cpus_width" | tr ' ' '-')" "$(printf '%*s' "$price_width" | tr ' ' '-')"
-echo "$table" | while IFS=$'\t' read -r id name cpus price; do
-    printf "%-${id_width}s  %-${name_width}s %-${cpus_width}s %-${price_width}s\n" "$id" "$name" "$cpus" "$price"
+gpus_width=$(echo "$table" | awk -F'\t' '{print length($4)}' | sort -nr | head -1)
+price_width=$(echo "$table" | awk -F'\t' '{print length($5)}' | sort -nr | head -1)
+
+if [ "$cpus_width" -lt 5 ]; then
+    cpus_width=5
+fi
+if [ "$gpus_width" -lt 5 ]; then
+    gpus_width=5
+fi
+if [ "$price_width" -lt 6 ]; then
+    price_width=6
+fi
+
+printf "%-${id_width}s  %-${name_width}s %-${cpus_width}s %-${gpus_width}s %-${price_width}s\n" "ID" "Name" "CPUs" "GPUs" "Price"
+printf "%${id_width}s  %${name_width}s %${cpus_width}s %${gpus_width}s %${price_width}s\n" "$(printf '%*s' "$id_width" | tr ' ' '-')" "$(printf '%*s' "$name_width" | tr ' ' '-')" "$(printf '%*s' "$cpus_width" | tr ' ' '-')" "$(printf '%*s' "$gpus_width" | tr ' ' '-')" "$(printf '%*s' "$price_width" | tr ' ' '-')"
+echo "$table" | while IFS=$'\t' read -r id name cpus gpus price; do
+    printf "%-${id_width}s  %-${name_width}s %-${cpus_width}s %-${gpus_width}s %-${price_width}s\n" "$id" "$name" "$cpus" "$gpus" "$price"
 done

swmcloudgate/routers/azure/connector.py

@@ -182,7 +182,7 @@ def list_sizes(self, location: str) -> list[VirtualMachineSize]:
                     resource_disk_size_in_mb=it["resource_disk_size_in_mb"],
                     memory_in_mb=it["memory_in_mb"],
                 )
-                vm_size.extra = {"price": it["price"], "description": "test vm size"}
+                vm_size.extra = {"gpus": 0, "price": it["price"], "description": "test vm size"}
                 node_sizes.append(vm_size)
             return node_sizes
 
@@ -192,8 +192,10 @@ def list_sizes(self, location: str) -> list[VirtualMachineSize]:
 
         size_map: dict[str, VirtualMachineSize] = {}
         for size in self._compute_client.virtual_machine_sizes.list(location):
+            size.extra: dict[str, str] = {}
             size_map[size.name] = size
         LOG.debug(f"Retrieved {len(size_map)} flavors from Azure")
+        self._add_gpus(location, size_map)
         result = self._add_prices(location, size_map)
 
         changed, deleted = cache.data_cache("flavors", "azure").update([location], result)
@@ -202,6 +204,20 @@ def list_sizes(self, location: str) -> list[VirtualMachineSize]:
 
         return result
 
+    def _add_gpus(self, location: str, size_map: dict[str, VirtualMachineSize]) -> None:
+        LOG.debug(f"Retriev GPU flavors information from Azure")
+        skus = self._compute_client.resource_skus.list()
+        for sku in skus:
+            if sku.resource_type.lower() != "virtualmachines":
+                continue
+            if sku.name not in size_map.keys():
+                continue
+            location_from_sku = sku.locations[0] if sku.locations else "Unknown"
+            if location != location_from_sku:
+                continue
+            if gpu_count := next((int(c.value) for c in sku.capabilities if c.name.lower() == "gpus"), 0):
+                size_map[sku.name].extra["gpus"] = gpu_count
+
     def _add_prices(self, location: str, size_map: dict[str, VirtualMachineSize]) -> list[VirtualMachineSize]:
         results: list[VirtualMachineSize] = []
 
@@ -238,7 +254,8 @@ def _add_prices(self, location: str, size_map: dict[str, VirtualMachineSize]) ->
             for meter_name in meter.meter_name.split("/"):
                 size_name = f"Standard_{meter_name.replace(' ', '_')}"
                 if vm_size := size_map.get(size_name):
-                    vm_size.extra = {"price": meter.meter_rates["0"], "description": meter.meter_sub_category}
+                    vm_size.extra["price"] = meter.meter_rates["0"]
+                    vm_size.extra["description"] = meter.meter_sub_category
                     if vm_size.name not in already_added:
                         already_added.add(vm_size.name)
                         results.append(vm_size)

swmcloudgate/routers/azure/converters.py

@@ -31,6 +31,7 @@ def convert_to_flavor(data: VirtualMachineSize) -> Flavor:
         id=str(image_id),
         name=data.name,
         cpus=data.number_of_cores,
+        gpus=data.extra.get("gpus", 0),
         mem=data.memory_in_mb,
         storage=data.resource_disk_size_in_mb,
         price=data.extra.get("price", 0.0),

swmcloudgate/routers/models.py

@@ -25,6 +25,7 @@ class Flavor(BaseModel):
     id: str
     name: str
     cpus: int
+    gpus: int
     mem: int
     storage: int
     price: float