wehub-resource-sync
113 lines
3.3 KiB
Python
113 lines
3.3 KiB
Python
# Source:
|
|
# https://github.com/kubernetes-client/python/blob/master/kubernetes/utils/quantity.py
|
|
from decimal import Decimal, InvalidOperation
|
|
from typing import Optional, Union
|
|
|
|
from ray._private.accelerators.tpu import (
|
|
get_num_chips_from_topology,
|
|
get_tpu_cores_per_chip,
|
|
)
|
|
|
|
# Mapping used to get generation for TPU-{accelerator}-head resource
|
|
# https://cloud.google.com/kubernetes-engine/docs/how-to/tpus#run
|
|
gke_tpu_accelerator_to_generation = {
|
|
"tpu-v4-podslice": "v4",
|
|
"tpu-v5-lite-device": "v5litepod",
|
|
"tpu-v5-lite-podslice": "v5litepod",
|
|
"tpu-v5p-slice": "v5p",
|
|
"tpu-v6e-slice": "v6e",
|
|
"tpu7x": "v7x",
|
|
}
|
|
|
|
|
|
def parse_quantity(quantity: Union[str, int, float, Decimal]) -> Decimal:
|
|
"""Parse kubernetes canonical form quantity like 200Mi to a decimal number.
|
|
Supported SI suffixes:
|
|
base1024: Ki | Mi | Gi | Ti | Pi | Ei
|
|
base1000: n | u | m | "" | k | M | G | T | P | E
|
|
|
|
See
|
|
https://github.com/kubernetes/apimachinery/blob/master/pkg/api/resource/quantity.go
|
|
|
|
Args:
|
|
quantity: kubernetes canonical form quantity.
|
|
|
|
Returns:
|
|
The parsed quantity as a decimal number.
|
|
|
|
Raises:
|
|
ValueError: On invalid or unknown input
|
|
"""
|
|
if isinstance(quantity, (int, float, Decimal)):
|
|
return Decimal(quantity)
|
|
|
|
exponents = {
|
|
"n": -3,
|
|
"u": -2,
|
|
"m": -1,
|
|
"K": 1,
|
|
"k": 1,
|
|
"M": 2,
|
|
"G": 3,
|
|
"T": 4,
|
|
"P": 5,
|
|
"E": 6,
|
|
}
|
|
|
|
quantity = str(quantity)
|
|
number = quantity
|
|
suffix = None
|
|
if len(quantity) >= 2 and quantity[-1] == "i":
|
|
if quantity[-2] in exponents:
|
|
number = quantity[:-2]
|
|
suffix = quantity[-2:]
|
|
elif len(quantity) >= 1 and quantity[-1] in exponents:
|
|
number = quantity[:-1]
|
|
suffix = quantity[-1:]
|
|
|
|
try:
|
|
number = Decimal(number)
|
|
except InvalidOperation:
|
|
raise ValueError("Invalid number format: {}".format(number))
|
|
|
|
if suffix is None:
|
|
return number
|
|
|
|
if suffix.endswith("i"):
|
|
base = 1024
|
|
elif len(suffix) == 1:
|
|
base = 1000
|
|
else:
|
|
raise ValueError("{} has unknown suffix".format(quantity))
|
|
|
|
# handle SI inconsistency
|
|
if suffix == "ki":
|
|
raise ValueError("{} has unknown suffix".format(quantity))
|
|
|
|
if suffix[0] not in exponents:
|
|
raise ValueError("{} has unknown suffix".format(quantity))
|
|
|
|
exponent = Decimal(exponents[suffix[0]])
|
|
return number * (base**exponent)
|
|
|
|
|
|
def tpu_node_selectors_to_type(topology: str, accelerator: str) -> Optional[str]:
|
|
"""Convert Kubernetes gke-tpu nodeSelectors to TPU accelerator_type
|
|
for a kuberay TPU worker group.
|
|
Args:
|
|
topology: value of the cloud.google.com/gke-tpu-topology Kubernetes
|
|
nodeSelector, describes the physical topology of the TPU podslice.
|
|
accelerator: value of the cloud.google.com/gke-tpu-accelerator nodeSelector,
|
|
the name of the TPU accelerator, e.g. tpu-v4-podslice
|
|
Returns:
|
|
A string, accelerator_type, e.g. "v4-8".
|
|
"""
|
|
if topology and accelerator:
|
|
generation = gke_tpu_accelerator_to_generation.get(accelerator)
|
|
if not generation:
|
|
return None
|
|
num_chips = get_num_chips_from_topology(topology)
|
|
num_cores = num_chips * get_tpu_cores_per_chip(generation)
|
|
return f"{generation}-{num_cores}"
|
|
return None
|