12318338392

Committed 13 Dec 2024 03:21PM UTC coverage: 86.388% (-0.001%) from 86.389%

Build # 12318338392

Build Type

Pull #572

github

Committed by

web-flow

Commit Message

Merge 0eb4f5dce into 9cb5d8146

Pull Request Pull Request #572: feat: manage v1 sessions

Run Details

17 of 32 new or added lines in 5 files covered. (53.13%)

4 existing lines in 4 files now uncovered.

14641 of 16948 relevant lines covered (86.39%)

1.52 hits per line

Source File
Press 'n' to go to next uncovered line, 'b' for previous

85.35

/components/renku_data_services/crc/models.py

"""Domain models for the application."""

from collections.abc import Callable
from dataclasses import asdict, dataclass, field
from enum import StrEnum
from typing import Any, Optional, Protocol
from uuid import uuid4

from renku_data_services.errors import ValidationError


class ResourcesProtocol(Protocol):
    """Used to represent resource values present in a resource class or quota."""

    @property
    def cpu(self) -> float:
        """Cpu in fractional cores."""
        ...

    @property
    def gpu(self) -> int:
        """Number of GPUs."""
        ...

    @property
    def memory(self) -> int:
        """Memory in gigabytes."""
        ...

    @property
    def max_storage(self) -> Optional[int]:
        """Maximum allowable storeage in gigabytes."""
        ...


class ResourcesCompareMixin:
    """A mixin that adds comparison operator support on ResourceClasses and Quotas."""

    def __compare(
        self,
        other: ResourcesProtocol,
        compare_func: Callable[[int | float, int | float], bool],
    ) -> bool:
        results = [
            compare_func(self.cpu, other.cpu),  # type: ignore[attr-defined]
            compare_func(self.memory, other.memory),  # type: ignore[attr-defined]
            compare_func(self.gpu, other.gpu),  # type: ignore[attr-defined]
        ]
        self_storage = getattr(self, "max_storage", 99999999999999999999999)
        other_storage = getattr(other, "max_storage", 99999999999999999999999)
        results.append(compare_func(self_storage, other_storage))
        return all(results)

    def __ge__(self, other: ResourcesProtocol) -> bool:
        return self.__compare(other, lambda x, y: x >= y)

    def __gt__(self, other: ResourcesProtocol) -> bool:
        return self.__compare(other, lambda x, y: x > y)

    def __lt__(self, other: ResourcesProtocol) -> bool:
        return self.__compare(other, lambda x, y: x < y)

    def __le__(self, other: ResourcesProtocol) -> bool:
        return self.__compare(other, lambda x, y: x <= y)


@dataclass(frozen=True, eq=True, kw_only=True)
class NodeAffinity:
    """Used to set the node affinity when scheduling sessions."""

    key: str
    required_during_scheduling: bool = False

    @classmethod
    def from_dict(cls, data: dict) -> "NodeAffinity":
        """Create a node affinity from a dictionary."""
        return cls(**data)


@dataclass(frozen=True, eq=True, kw_only=True)
class ResourceClass(ResourcesCompareMixin):
    """Resource class model."""

    name: str
    cpu: float
    memory: int
    max_storage: int
    gpu: int
    id: Optional[int] = None
    default: bool = False
    default_storage: int = 1
    matching: Optional[bool] = None
    node_affinities: list[NodeAffinity] = field(default_factory=list)
    tolerations: list[str] = field(default_factory=list)
    quota: str | None = None

    def __post_init__(self) -> None:
        if len(self.name) > 40:
            raise ValidationError(message="'name' cannot be longer than 40 characters.")
        if self.default_storage > self.max_storage:
            raise ValidationError(message="The default storage cannot be larger than the max allowable storage.")
        # We need to sort node affinities and tolerations to make '__eq__' reliable
        object.__setattr__(
            self, "node_affinities", sorted(self.node_affinities, key=lambda x: (x.key, x.required_during_scheduling))
        )
        object.__setattr__(self, "tolerations", sorted(self.tolerations))

    @classmethod
    def from_dict(cls, data: dict) -> "ResourceClass":
        """Create the model from a plain dictionary."""
        node_affinities: list[NodeAffinity] = []
        tolerations: list[str] = []
        quota: str | None = None
        if data.get("node_affinities"):
            node_affinities = [
                NodeAffinity.from_dict(affinity) if isinstance(affinity, dict) else affinity
                for affinity in data.get("node_affinities", [])
            ]
        if isinstance(data.get("tolerations"), list):
            tolerations = [toleration for toleration in data["tolerations"]]
        if data_quota := data.get("quota"):
            if isinstance(data_quota, str):
                quota = data_quota
            elif isinstance(data_quota, Quota):
                quota = data_quota.id
        return cls(**{**data, "tolerations": tolerations, "node_affinities": node_affinities, "quota": quota})

    def is_quota_valid(self, quota: "Quota") -> bool:
        """Determine if a quota is compatible with the resource class."""
        return quota >= self

    def update(self, **kwargs: dict) -> "ResourceClass":
        """Update a field of the resource class and return a new copy."""
        if not kwargs:
            return self
        return ResourceClass.from_dict({**asdict(self), **kwargs})


class GpuKind(StrEnum):
    """GPU kinds for k8s."""

    NVIDIA = "nvidia.com"
    AMD = "amd.com"


@dataclass(frozen=True, eq=True, kw_only=True)
class Quota(ResourcesCompareMixin):
    """Quota model."""

    cpu: float
    memory: int
    gpu: int
    gpu_kind: GpuKind = GpuKind.NVIDIA
    id: Optional[str] = None

    @classmethod
    def from_dict(cls, data: dict) -> "Quota":
        """Create the model from a plain dictionary."""
        gpu_kind = GpuKind.NVIDIA
        if "gpu_kind" in data:
            gpu_kind = data["gpu_kind"] if isinstance(data["gpu_kind"], GpuKind) else GpuKind[data["gpu_kind"]]
        return cls(**{**data, "gpu_kind": gpu_kind})

    def is_resource_class_compatible(self, rc: "ResourceClass") -> bool:
        """Determine if a resource class is compatible with the quota."""
        return rc <= self

    def generate_id(self) -> "Quota":
        """Create a new quota with its ID set to a uuid."""
        if self.id is not None:
            return self
        return self.from_dict({**asdict(self), "id": str(uuid4())})


@dataclass(frozen=True, eq=True, kw_only=True)
class ResourcePool:
    """Resource pool model."""

    name: str
    classes: list["ResourceClass"]
    quota: Optional[Quota] = None
    id: Optional[int] = None
    idle_threshold: Optional[int] = None
    hibernation_threshold: Optional[int] = None
    default: bool = False
    public: bool = False

    def __post_init__(self) -> None:
        """Validate the resource pool after initialization."""
        if len(self.name) > 40:
            raise ValidationError(message="'name' cannot be longer than 40 characters.")
        if self.default and not self.public:
            raise ValidationError(message="The default resource pool has to be public.")
        if self.default and self.quota is not None:
            raise ValidationError(message="A default resource pool cannot have a quota.")
        if (self.idle_threshold and self.idle_threshold < 0) or (
            self.hibernation_threshold and self.hibernation_threshold < 0
        ):
            raise ValidationError(message="Idle threshold and hibernation threshold need to be larger than 0.")

        if self.idle_threshold == 0:
            object.__setattr__(self, "idle_threshold", None)
        if self.hibernation_threshold == 0:
            object.__setattr__(self, "hibernation_threshold", None)

        default_classes = []
        for cls in list(self.classes):
            if self.quota and not self.quota.is_resource_class_compatible(cls):
                raise ValidationError(
                    message=f"The resource class with name {cls.name} is not compatible with the quota."
                )
            if cls.default:
                default_classes.append(cls)
        if len(default_classes) != 1:
            raise ValidationError(message="One default class is required in each resource pool.")

    def set_quota(self, val: Quota) -> "ResourcePool":
        """Set the quota for a resource pool."""
        for cls in list(self.classes):
            if not val.is_resource_class_compatible(cls):
                raise ValidationError(
                    message=f"The resource class with name {cls.name} is not compatiable with the quota."
                )
        return self.from_dict({**asdict(self), "quota": val})

    def update(self, **kwargs: Any) -> "ResourcePool":
        """Determine if an update to a resource pool is valid and if valid create new updated resource pool."""
        if self.default and "default" in kwargs and not kwargs["default"]:
            raise ValidationError(message="A default resource pool cannot be made non-default.")
        return ResourcePool.from_dict({**asdict(self), **kwargs})

    @classmethod
    def from_dict(cls, data: dict) -> "ResourcePool":
        """Create the model from a plain dictionary."""
        quota: Optional[Quota] = None
        if "quota" in data and isinstance(data["quota"], dict):
            quota = Quota.from_dict(data["quota"])
        elif "quota" in data and isinstance(data["quota"], Quota):
            quota = data["quota"]
        if "classes" in data and isinstance(data["classes"], set):
            classes = [ResourceClass.from_dict(c) if isinstance(c, dict) else c for c in list(data["classes"])]
        elif "classes" in data and isinstance(data["classes"], list):
            classes = [ResourceClass.from_dict(c) if isinstance(c, dict) else c for c in data["classes"]]
        return cls(
            name=data["name"],
            id=data.get("id"),
            classes=classes,
            quota=quota,
            default=data.get("default", False),
            public=data.get("public", False),
            idle_threshold=data.get("idle_threshold"),
            hibernation_threshold=data.get("hibernation_threshold"),
        )

1	"""Domain models for the application."""
2
3	from collections.abc import Callable	2✔
4	from dataclasses import asdict, dataclass, field	2✔
5	from enum import StrEnum	2✔
6	from typing import Any, Optional, Protocol	2✔
7	from uuid import uuid4	2✔
8
9	from renku_data_services.errors import ValidationError	2✔
10
11
12	class ResourcesProtocol(Protocol):	2✔
13	"""Used to represent resource values present in a resource class or quota."""
14
15	@property	2✔
16	def cpu(self) -> float:	2✔
17	"""Cpu in fractional cores."""
18	...	×
19
20	@property	2✔
21	def gpu(self) -> int:	2✔
22	"""Number of GPUs."""
23	...	×
24
25	@property	2✔
26	def memory(self) -> int:	2✔
27	"""Memory in gigabytes."""
28	...	×
29
30	@property	2✔
31	def max_storage(self) -> Optional[int]:	2✔
32	"""Maximum allowable storeage in gigabytes."""
33	...	×
34
35
36	class ResourcesCompareMixin:	2✔
37	"""A mixin that adds comparison operator support on ResourceClasses and Quotas."""
38
39	def __compare(	2✔
40	self,
41	other: ResourcesProtocol,
42	compare_func: Callable[[int \| float, int \| float], bool],
43	) -> bool:
44	results = [	2✔
45	compare_func(self.cpu, other.cpu), # type: ignore[attr-defined]
46	compare_func(self.memory, other.memory), # type: ignore[attr-defined]
47	compare_func(self.gpu, other.gpu), # type: ignore[attr-defined]
48	]
49	self_storage = getattr(self, "max_storage", 99999999999999999999999)	2✔
50	other_storage = getattr(other, "max_storage", 99999999999999999999999)	2✔
51	results.append(compare_func(self_storage, other_storage))	2✔
52	return all(results)	2✔
53
54	def __ge__(self, other: ResourcesProtocol) -> bool:	2✔
55	return self.__compare(other, lambda x, y: x >= y)	×
56
57	def __gt__(self, other: ResourcesProtocol) -> bool:	2✔
58	return self.__compare(other, lambda x, y: x > y)	×
59
60	def __lt__(self, other: ResourcesProtocol) -> bool:	2✔
61	return self.__compare(other, lambda x, y: x < y)	×
62
63	def __le__(self, other: ResourcesProtocol) -> bool:	2✔
64	return self.__compare(other, lambda x, y: x <= y)	2✔
65
66
67	@dataclass(frozen=True, eq=True, kw_only=True)	2✔
68	class NodeAffinity:	2✔
69	"""Used to set the node affinity when scheduling sessions."""
70
71	key: str	2✔
72	required_during_scheduling: bool = False	2✔
73
74	@classmethod	2✔
75	def from_dict(cls, data: dict) -> "NodeAffinity":	2✔
76	"""Create a node affinity from a dictionary."""
77	return cls(**data)	2✔
78
79
80	@dataclass(frozen=True, eq=True, kw_only=True)	2✔
81	class ResourceClass(ResourcesCompareMixin):	2✔
82	"""Resource class model."""
83
84	name: str	2✔
85	cpu: float	2✔
86	memory: int	2✔
87	max_storage: int	2✔
88	gpu: int	2✔
89	id: Optional[int] = None	2✔
90	default: bool = False	2✔
91	default_storage: int = 1	2✔
92	matching: Optional[bool] = None	2✔
93	node_affinities: list[NodeAffinity] = field(default_factory=list)	2✔
94	tolerations: list[str] = field(default_factory=list)	2✔
95	quota: str \| None = None	2✔
96
97	def __post_init__(self) -> None:	2✔
98	if len(self.name) > 40:	2✔
99	raise ValidationError(message="'name' cannot be longer than 40 characters.")	1✔
100	if self.default_storage > self.max_storage:	2✔
101	raise ValidationError(message="The default storage cannot be larger than the max allowable storage.")	×
102	# We need to sort node affinities and tolerations to make '__eq__' reliable
103	object.__setattr__(	2✔
104	self, "node_affinities", sorted(self.node_affinities, key=lambda x: (x.key, x.required_during_scheduling))
105	)
106	object.__setattr__(self, "tolerations", sorted(self.tolerations))	2✔
107
108	@classmethod	2✔
109	def from_dict(cls, data: dict) -> "ResourceClass":	2✔
110	"""Create the model from a plain dictionary."""
111	node_affinities: list[NodeAffinity] = []	2✔
112	tolerations: list[str] = []	2✔
113	quota: str \| None = None	2✔
114	if data.get("node_affinities"):	2✔
115	node_affinities = [	2✔
116	NodeAffinity.from_dict(affinity) if isinstance(affinity, dict) else affinity
117	for affinity in data.get("node_affinities", [])
118	]
119	if isinstance(data.get("tolerations"), list):	2✔
120	tolerations = [toleration for toleration in data["tolerations"]]	2✔
121	if data_quota := data.get("quota"):	2✔
122	if isinstance(data_quota, str):	1✔
123	quota = data_quota	1✔
124	elif isinstance(data_quota, Quota):	×
125	quota = data_quota.id	×
126	return cls({data, "tolerations": tolerations, "node_affinities": node_affinities, "quota": quota})	2✔
127
128	def is_quota_valid(self, quota: "Quota") -> bool:	2✔
129	"""Determine if a quota is compatible with the resource class."""
130	return quota >= self	×
131
132	def update(self, **kwargs: dict) -> "ResourceClass":	2✔
133	"""Update a field of the resource class and return a new copy."""
134	if not kwargs:	1✔
135	return self	×
136	return ResourceClass.from_dict({asdict(self), kwargs})	1✔
137
138
139	class GpuKind(StrEnum):	2✔
140	"""GPU kinds for k8s."""
141
142	NVIDIA = "nvidia.com"	2✔
143	AMD = "amd.com"	2✔
144
145
146	@dataclass(frozen=True, eq=True, kw_only=True)	2✔
147	class Quota(ResourcesCompareMixin):	2✔
148	"""Quota model."""
149
150	cpu: float	2✔
151	memory: int	2✔
152	gpu: int	2✔
153	gpu_kind: GpuKind = GpuKind.NVIDIA	2✔
154	id: Optional[str] = None	2✔
155
156	@classmethod	2✔
157	def from_dict(cls, data: dict) -> "Quota":	2✔
158	"""Create the model from a plain dictionary."""
159	gpu_kind = GpuKind.NVIDIA	2✔
160	if "gpu_kind" in data:	2✔
161	gpu_kind = data["gpu_kind"] if isinstance(data["gpu_kind"], GpuKind) else GpuKind[data["gpu_kind"]]	2✔
162	return cls({data, "gpu_kind": gpu_kind})	2✔
163
164	def is_resource_class_compatible(self, rc: "ResourceClass") -> bool:	2✔
165	"""Determine if a resource class is compatible with the quota."""
166	return rc <= self	2✔
167
168	def generate_id(self) -> "Quota":	2✔
169	"""Create a new quota with its ID set to a uuid."""
170	if self.id is not None:	2✔
171	return self	×
172	return self.from_dict({**asdict(self), "id": str(uuid4())})	2✔
173
174
175	@dataclass(frozen=True, eq=True, kw_only=True)	2✔
176	class ResourcePool:	2✔
177	"""Resource pool model."""
178
179	name: str	2✔
180	classes: list["ResourceClass"]	2✔
181	quota: Optional[Quota] = None	2✔
182	id: Optional[int] = None	2✔
183	idle_threshold: Optional[int] = None	2✔
184	hibernation_threshold: Optional[int] = None	2✔
185	default: bool = False	2✔
186	public: bool = False	2✔
187
188	def __post_init__(self) -> None:	2✔
189	"""Validate the resource pool after initialization."""
190	if len(self.name) > 40:	2✔
UNCOV 191	raise ValidationError(message="'name' cannot be longer than 40 characters.")	×
192	if self.default and not self.public:	2✔
193	raise ValidationError(message="The default resource pool has to be public.")	×
194	if self.default and self.quota is not None:	2✔
195	raise ValidationError(message="A default resource pool cannot have a quota.")	×
196	if (self.idle_threshold and self.idle_threshold < 0) or (	2✔
197	self.hibernation_threshold and self.hibernation_threshold < 0
198	):
199	raise ValidationError(message="Idle threshold and hibernation threshold need to be larger than 0.")	×
200
201	if self.idle_threshold == 0:	2✔
202	object.__setattr__(self, "idle_threshold", None)	×
203	if self.hibernation_threshold == 0:	2✔
204	object.__setattr__(self, "hibernation_threshold", None)	×
205
206	default_classes = []	2✔
207	for cls in list(self.classes):	2✔
208	if self.quota and not self.quota.is_resource_class_compatible(cls):	2✔
209	raise ValidationError(	×
210	message=f"The resource class with name {cls.name} is not compatible with the quota."
211	)
212	if cls.default:	2✔
213	default_classes.append(cls)	2✔
214	if len(default_classes) != 1:	2✔
215	raise ValidationError(message="One default class is required in each resource pool.")	1✔
216
217	def set_quota(self, val: Quota) -> "ResourcePool":	2✔
218	"""Set the quota for a resource pool."""
219	for cls in list(self.classes):	2✔
220	if not val.is_resource_class_compatible(cls):	2✔
221	raise ValidationError(	×
222	message=f"The resource class with name {cls.name} is not compatiable with the quota."
223	)
224	return self.from_dict({**asdict(self), "quota": val})	2✔
225
226	def update(self, **kwargs: Any) -> "ResourcePool":	2✔
227	"""Determine if an update to a resource pool is valid and if valid create new updated resource pool."""
228	if self.default and "default" in kwargs and not kwargs["default"]:	1✔
229	raise ValidationError(message="A default resource pool cannot be made non-default.")	×
230	return ResourcePool.from_dict({asdict(self), kwargs})	1✔
231
232	@classmethod	2✔
233	def from_dict(cls, data: dict) -> "ResourcePool":	2✔
234	"""Create the model from a plain dictionary."""
235	quota: Optional[Quota] = None	2✔
236	if "quota" in data and isinstance(data["quota"], dict):	2✔
237	quota = Quota.from_dict(data["quota"])	2✔
238	elif "quota" in data and isinstance(data["quota"], Quota):	2✔
239	quota = data["quota"]	2✔
240	if "classes" in data and isinstance(data["classes"], set):	2✔
241	classes = [ResourceClass.from_dict(c) if isinstance(c, dict) else c for c in list(data["classes"])]	×
242	elif "classes" in data and isinstance(data["classes"], list):	2✔
243	classes = [ResourceClass.from_dict(c) if isinstance(c, dict) else c for c in data["classes"]]	2✔
244	return cls(	2✔
245	name=data["name"],
246	id=data.get("id"),
247	classes=classes,
248	quota=quota,
249	default=data.get("default", False),
250	public=data.get("public", False),
251	idle_threshold=data.get("idle_threshold"),
252	hibernation_threshold=data.get("hibernation_threshold"),
253	)

SwissDataScienceCenter / renku-data-services / 12318338392

Source File Press 'n' to go to next uncovered line, 'b' for previous

Source File
Press 'n' to go to next uncovered line, 'b' for previous