1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
|
from __future__ import annotations
import enum
import math
from typing import Literal, Optional, TypeVar, Union
import pydantic as pd
from kubernetes.client.models import V1Container
from robusta_krr.utils import resource_units
class ResourceType(str, enum.Enum):
"""The type of resource.
Just add new types here and they will be automatically supported.
"""
CPU = "cpu"
Memory = "memory"
RecommendationValue = Union[float, Literal["?"], None]
RecommendationValueRaw = Union[float, str, None]
Self = TypeVar("Self", bound="ResourceAllocations")
class ResourceAllocations(pd.BaseModel):
requests: dict[ResourceType, RecommendationValue]
limits: dict[ResourceType, RecommendationValue]
info: dict[ResourceType, Optional[str]] = {}
@staticmethod
def __parse_resource_value(value: RecommendationValueRaw) -> RecommendationValue:
if value is None:
return None
if isinstance(value, str):
return float(resource_units.parse(value))
if math.isnan(value):
return "?"
return float(value)
@pd.validator("requests", "limits", pre=True)
def validate_requests(
cls, value: dict[ResourceType, RecommendationValueRaw]
) -> dict[ResourceType, RecommendationValue]:
return {
resource_type: cls.__parse_resource_value(resource_value) for resource_type, resource_value in value.items()
}
@classmethod
def from_container(cls: type[Self], container: V1Container) -> Self:
"""Get the resource allocations from a Kubernetes container.
Args:
container: The Kubernetes container.
Returns:
The resource allocations.
"""
return cls(
requests={
ResourceType.CPU: container.resources.requests.get("cpu")
if container.resources and container.resources.requests
else None,
ResourceType.Memory: container.resources.requests.get("memory")
if container.resources and container.resources.requests
else None,
},
limits={
ResourceType.CPU: container.resources.limits.get("cpu")
if container.resources and container.resources.limits
else None,
ResourceType.Memory: container.resources.limits.get("memory")
if container.resources and container.resources.limits
else None,
},
)
|