def test_get(resource_name: str, quantity: ResourceQuantity):
resource = Resource(
milli_cpu=ResourceQuantity(100),
memory=ResourceQuantity(100),
scalar_resources={"nvidia.com/gpu": ResourceQuantity(100)},
)
result = resource.get(resource_name)
assert result == quantity
def __mul__(self, ratio: int) -> "Resource":
resource: Resource = self.copy(deep=True)
resource.milli_cpu += ResourceQuantity(1)
resource.milli_cpu *= ratio
resource.memory *= ratio
for resource_name, quant in resource.scalar_resources.items():
resource.scalar_resources.setdefault(resource_name,
ResourceQuantity(0))
resource.scalar_resources[resource_name] *= ratio
return resource
def diff(self: "Resource",
other: "Resource") -> Tuple["Resource", "Resource"]:
increase_value = Resource()
decrease_value = Resource()
if self.milli_cpu > other.milli_cpu:
handle_value = increase_value
else:
handle_value = decrease_value
handle_value.milli_cpu += abs(self.milli_cpu - other.milli_cpu)
if self.memory > other.memory:
handle_value = increase_value
else:
handle_value = decrease_value
handle_value.memory += abs(self.memory - other.memory)
for resource_name, quant in self.scalar_resources.items():
other_quant = other.scalar_resources.get(resource_name, 0)
if quant > other_quant:
handle_value = increase_value
else:
handle_value = decrease_value
handle_value.scalar_resources.setdefault(resource_name,
ResourceQuantity(0))
handle_value.scalar_resources[resource_name] += abs(quant -
other_quant)
return increase_value, decrease_value
def new_quota(name: str, weight: int) -> ResourceQuota:
quota: ResourceQuota = ResourceQuota.parse_obj(
{"metadata": {
"name": name
}})
if weight >= 0:
quota.spec.hard[NamespaceWeightKey] = ResourceQuantity(weight)
return quota
def __add__(self, other: "Resource") -> "Resource":
resource: Resource = self.copy(deep=True)
resource.milli_cpu += other.milli_cpu
resource.memory += other.memory
for resource_name, quant in other.scalar_resources.items():
resource.scalar_resources.setdefault(resource_name,
ResourceQuantity(0))
resource.scalar_resources[resource_name] += quant
return resource
def __sub__(self, other: "Resource") -> "Resource":
assert (
self >= other
), f"resource is not sufficient to do operation: {self} sub {other}"
resource: Resource = self.copy(deep=True)
resource.milli_cpu -= other.milli_cpu
resource.memory -= other.memory
for resource_name, quant in other.scalar_resources.items():
resource.scalar_resources.setdefault(resource_name,
ResourceQuantity(0))
resource.scalar_resources[resource_name] -= quant
return resource
def fit_delta(self, other: "Resource") -> "Resource":
if other.milli_cpu > 0:
self.milli_cpu -= other.milli_cpu + MinMilliCpu
if other.memory > 0:
self.memory -= other.memory + MinMemory
for resource_name, other_quant in other.scalar_resources.items():
if other_quant > 0:
self.scalar_resources.setdefault(resource_name,
ResourceQuantity(0))
self.scalar_resources[resource_name] -= (
other_quant + MinMilliScalarResources)
return self
def new(cls: Type["Resource"],
resource_list: Union[ResourceList, Dict[str, str]]) -> "Resource":
resource = cls()
for resource_name, value in resource_list.items():
if not value:
continue
value = ResourceQuantity(value)
if resource_name == "cpu":
resource.milli_cpu += value * 1000
elif resource_name == "memory":
resource.memory += value
elif resource_name == "pods":
resource.max_task_num = int(value)
elif helper.is_scalar_resource_name(resource_name):
resource.scalar_resources[resource_name] = value * 1000
return resource
def test_set_scalar_resource():
empty_resource = Resource()
empty_resource.set_scalar_resource("nvidia.com/gpu", 100)
assert empty_resource == Resource(
scalar_resources={"nvidia.com/gpu": ResourceQuantity(100)})
})
assert resource.resource_names == ["cpu", "memory", "nvidia.com/gpu"]
@parametrize(
"resource_list,resource",
[
(
{
"cpu": "2000m",
"memory": "1G",
"pods": 20,
"nvidia.com/gpu": "1G"
},
Resource(
milli_cpu=ResourceQuantity(2000),
memory=ResourceQuantity(1e9),
max_task_num=20,
scalar_resources={"nvidia.com/gpu": ResourceQuantity(1e12)},
),
),
(
{
"cpu": "2000m",
"memory": "1G",
"pods": 20,
"nvidia.com/gpu": ""
},
Resource(
milli_cpu=ResourceQuantity(2000),
memory=ResourceQuantity(1e9),
from decimal import Decimal
import pytest
from airport.kube.api import ResourceQuantity
parametrize = pytest.mark.parametrize
@parametrize(
"resource_quantity,expected",
[
(ResourceQuantity("1"), "1"),
(ResourceQuantity("1.1"), "1.1"),
(ResourceQuantity("10m"), "0.01"),
(ResourceQuantity("10M"), str(10 * 1000 * 1000)),
(ResourceQuantity("10Mi"), str(10 * 1024 * 1024)),
],
)
def test_resource_quantity(resource_quantity, expected):
assert resource_quantity == Decimal(expected)
@parametrize(
"name,op,expected",
[
("add", ResourceQuantity(100) + 2, ResourceQuantity("102")),
("sub", ResourceQuantity(100) - 2, ResourceQuantity("98")),
("mul", ResourceQuantity(100) * 2, ResourceQuantity("200")),
(
def set_scalar_resource(self, resource_name: str,
quantity: Union[float, ResourceQuantity]):
self.scalar_resources[resource_name] = ResourceQuantity(quantity)
class Resource(BaseModel):
milli_cpu: ResourceQuantity = ResourceQuantity(0)
memory: ResourceQuantity = ResourceQuantity(0)
scalar_resources: Dict[str, ResourceQuantity] = {}
max_task_num: Optional[int] = None
@property
def resource_names(self) -> List[str]:
return ["cpu", "memory", *self.scalar_resources.keys()]
@classmethod
def new(cls: Type["Resource"],
resource_list: Union[ResourceList, Dict[str, str]]) -> "Resource":
resource = cls()
for resource_name, value in resource_list.items():
if not value:
continue
value = ResourceQuantity(value)
if resource_name == "cpu":
resource.milli_cpu += value * 1000
elif resource_name == "memory":
resource.memory += value
elif resource_name == "pods":
resource.max_task_num = int(value)
elif helper.is_scalar_resource_name(resource_name):
resource.scalar_resources[resource_name] = value * 1000
return resource
def get(self, resource_name: str) -> ResourceQuantity:
if resource_name == "cpu":
return self.milli_cpu
elif resource_name == "memory":
return self.memory
else:
try:
return self.scalar_resources[resource_name]
except KeyError:
raise ValueError(f"Unknown resource {resource_name}")
def set_scalar_resource(self, resource_name: str,
quantity: Union[float, ResourceQuantity]):
self.scalar_resources[resource_name] = ResourceQuantity(quantity)
def is_empty(self) -> bool:
"""
Returns bool after checking any of resource is less than min possible value
"""
if self.milli_cpu >= MinMilliCpu or self.memory >= MinMemory:
return False
for quant in self.scalar_resources.values():
if quant >= MinMilliScalarResources:
return False
return True
def is_zero(self, resource_name: str) -> bool:
"""
Checks whether that resource is less than min possible value
"""
if resource_name == "cpu":
return self.milli_cpu < MinMilliCpu
elif resource_name == "memory":
return self.memory < MinMemory
else:
try:
quantity = self.scalar_resources[resource_name]
return quantity < MinMilliScalarResources
except KeyError:
raise ValueError(f"Unknown resource {resource_name}")
def set_max_resource(self, other: "Resource") -> "Resource":
self.milli_cpu = max(self.milli_cpu, other.milli_cpu)
self.memory = max(self.memory, other.memory)
for resource_name, other_quant in other.scalar_resources.items():
self.scalar_resources[resource_name] = max(
self.scalar_resources[resource_name], other_quant)
return self
def fit_delta(self, other: "Resource") -> "Resource":
if other.milli_cpu > 0:
self.milli_cpu -= other.milli_cpu + MinMilliCpu
if other.memory > 0:
self.memory -= other.memory + MinMemory
for resource_name, other_quant in other.scalar_resources.items():
if other_quant > 0:
self.scalar_resources.setdefault(resource_name,
ResourceQuantity(0))
self.scalar_resources[resource_name] -= (
other_quant + MinMilliScalarResources)
return self
def diff(self: "Resource",
other: "Resource") -> Tuple["Resource", "Resource"]:
increase_value = Resource()
decrease_value = Resource()
if self.milli_cpu > other.milli_cpu:
handle_value = increase_value
else:
handle_value = decrease_value
handle_value.milli_cpu += abs(self.milli_cpu - other.milli_cpu)
if self.memory > other.memory:
handle_value = increase_value
else:
handle_value = decrease_value
handle_value.memory += abs(self.memory - other.memory)
for resource_name, quant in self.scalar_resources.items():
other_quant = other.scalar_resources.get(resource_name, 0)
if quant > other_quant:
handle_value = increase_value
else:
handle_value = decrease_value
handle_value.scalar_resources.setdefault(resource_name,
ResourceQuantity(0))
handle_value.scalar_resources[resource_name] += abs(quant -
other_quant)
return increase_value, decrease_value
def less_equal_strict(self, other: "Resource") -> bool:
if self.milli_cpu > other.milli_cpu or self.memory > other.memory:
return False
for resource_name, quant in self.scalar_resources.items():
if (other_quant := other.scalar_resources.get(resource_name)
) is None or quant > other_quant:
return False
return True