9446003648

Committed 10 Jun 2024 09:35AM UTC coverage: 90.298% (+0.06%) from 90.239%

Build # 9446003648

Build Type

Pull #248

github

Committed by

web-flow

Commit Message

Merge 9ff3e8a6c into 1e340ea36

Pull Request Pull Request #248: feat: add support for bitbucket

Run Details

40 of 46 new or added lines in 3 files covered. (86.96%)

4 existing lines in 4 files now uncovered.

8488 of 9400 relevant lines covered (90.3%)

1.6 hits per line

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

84.91

/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)

    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] = []
        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"]]
        return cls(**{**data, "tolerations": tolerations, "node_affinities": node_affinities})

    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.")

        # We need to sort classes to make '__eq__' reliable
        object.__setattr__(
            self, "classes", sorted(self.classes, key=lambda x: (x.default, x.cpu, x.memory, x.default_storage, x.name))
        )

    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✔
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."""	2✔
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."""	2✔
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."""	2✔
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)	1✔
78
79
80	@dataclass(frozen=True, eq=True, kw_only=True)	2✔
81	class ResourceClass(ResourcesCompareMixin):	2✔
82	"""Resource class model."""	2✔
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
96	def __post_init__(self) -> None:	2✔
97	if len(self.name) > 40:	2✔
UNCOV 98	raise ValidationError(message="'name' cannot be longer than 40 characters.")	×
99	if self.default_storage > self.max_storage:	2✔
100	raise ValidationError(message="The default storage cannot be larger than the max allowable storage.")	×
101	# We need to sort node affinities and tolerations to make '__eq__' reliable
102	object.__setattr__(	2✔
103	self, "node_affinities", sorted(self.node_affinities, key=lambda x: (x.key, x.required_during_scheduling))
104	)
105	object.__setattr__(self, "tolerations", sorted(self.tolerations))	2✔
106
107	@classmethod	2✔
108	def from_dict(cls, data: dict) -> "ResourceClass":	2✔
109	"""Create the model from a plain dictionary."""
110	node_affinities: list[NodeAffinity] = []	2✔
111	tolerations: list[str] = []	2✔
112	if data.get("node_affinities"):	2✔
113	node_affinities = [	1✔
114	NodeAffinity.from_dict(affinity) if isinstance(affinity, dict) else affinity
115	for affinity in data.get("node_affinities", [])
116	]
117	if isinstance(data.get("tolerations"), list):	2✔
118	tolerations = [toleration for toleration in data["tolerations"]]	2✔
119	return cls({data, "tolerations": tolerations, "node_affinities": node_affinities})	2✔
120
121	def is_quota_valid(self, quota: "Quota") -> bool:	2✔
122	"""Determine if a quota is compatible with the resource class."""
123	return quota >= self	×
124
125	def update(self, **kwargs: dict) -> "ResourceClass":	2✔
126	"""Update a field of the resource class and return a new copy."""
127	if not kwargs:	×
128	return self	×
129	return ResourceClass.from_dict({asdict(self), kwargs})	×
130
131
132	class GpuKind(StrEnum):	2✔
133	"""GPU kinds for k8s."""	2✔
134
135	NVIDIA = "nvidia.com"	2✔
136	AMD = "amd.com"	2✔
137
138
139	@dataclass(frozen=True, eq=True, kw_only=True)	2✔
140	class Quota(ResourcesCompareMixin):	2✔
141	"""Quota model."""	2✔
142
143	cpu: float	2✔
144	memory: int	2✔
145	gpu: int	2✔
146	gpu_kind: GpuKind = GpuKind.NVIDIA	2✔
147	id: Optional[str] = None	2✔
148
149	@classmethod	2✔
150	def from_dict(cls, data: dict) -> "Quota":	2✔
151	"""Create the model from a plain dictionary."""
152	gpu_kind = GpuKind.NVIDIA	2✔
153	if "gpu_kind" in data:	2✔
154	gpu_kind = data["gpu_kind"] if isinstance(data["gpu_kind"], GpuKind) else GpuKind[data["gpu_kind"]]	2✔
155	return cls({data, "gpu_kind": gpu_kind})	2✔
156
157	def is_resource_class_compatible(self, rc: "ResourceClass") -> bool:	2✔
158	"""Determine if a resource class is compatible with the quota."""
159	return rc <= self	2✔
160
161	def generate_id(self) -> "Quota":	2✔
162	"""Create a new quota with its ID set to a uuid."""
163	if self.id is not None:	2✔
164	return self	×
165	return self.from_dict({**asdict(self), "id": str(uuid4())})	2✔
166
167
168	@dataclass(frozen=True, eq=True, kw_only=True)	2✔
169	class ResourcePool:	2✔
170	"""Resource pool model."""	2✔
171
172	name: str	2✔
173	classes: list["ResourceClass"]	2✔
174	quota: Optional[Quota] = None	2✔
175	id: Optional[int] = None	2✔
176	idle_threshold: Optional[int] = None	2✔
177	hibernation_threshold: Optional[int] = None	2✔
178	default: bool = False	2✔
179	public: bool = False	2✔
180
181	def __post_init__(self) -> None:	2✔
182	"""Validate the resource pool after initialization."""
183	if len(self.name) > 40:	2✔
184	raise ValidationError(message="'name' cannot be longer than 40 characters.")	×
185	if self.default and not self.public:	2✔
186	raise ValidationError(message="The default resource pool has to be public.")	×
187	if self.default and self.quota is not None:	2✔
188	raise ValidationError(message="A default resource pool cannot have a quota.")	×
189	if (self.idle_threshold and self.idle_threshold < 0) or (	2✔
190	self.hibernation_threshold and self.hibernation_threshold < 0
191	):
192	raise ValidationError(message="Idle threshold and hibernation threshold need to be larger than 0.")	×
193
194	if self.idle_threshold == 0:	2✔
195	object.__setattr__(self, "idle_threshold", None)	×
196	if self.hibernation_threshold == 0:	2✔
197	object.__setattr__(self, "hibernation_threshold", None)	×
198
199	default_classes = []	2✔
200	for cls in list(self.classes):	2✔
201	if self.quota and not self.quota.is_resource_class_compatible(cls):	2✔
202	raise ValidationError(	×
203	message=f"The resource class with name {cls.name} is not compatible with the quota."
204	)
205	if cls.default:	2✔
206	default_classes.append(cls)	2✔
207	if len(default_classes) != 1:	2✔
208	raise ValidationError(message="One default class is required in each resource pool.")	1✔
209
210	# We need to sort classes to make '__eq__' reliable
211	object.__setattr__(	2✔
212	self, "classes", sorted(self.classes, key=lambda x: (x.default, x.cpu, x.memory, x.default_storage, x.name))
213	)
214
215	def set_quota(self, val: Quota) -> "ResourcePool":	2✔
216	"""Set the quota for a resource pool."""
217	for cls in list(self.classes):	2✔
218	if not val.is_resource_class_compatible(cls):	2✔
219	raise ValidationError(	×
220	message=f"The resource class with name {cls.name} is not compatiable with the quota."
221	)
222	return self.from_dict({**asdict(self), "quota": val})	2✔
223
224	def update(self, **kwargs: Any) -> "ResourcePool":	2✔
225	"""Determine if an update to a resource pool is valid and if valid create new updated resource pool."""
226	if self.default and "default" in kwargs and not kwargs["default"]:	1✔
227	raise ValidationError(message="A default resource pool cannot be made non-default.")	×
228	return ResourcePool.from_dict({asdict(self), kwargs})	1✔
229
230	@classmethod	2✔
231	def from_dict(cls, data: dict) -> "ResourcePool":	2✔
232	"""Create the model from a plain dictionary."""
233	quota: Optional[Quota] = None	2✔
234	if "quota" in data and isinstance(data["quota"], dict):	2✔
235	quota = Quota.from_dict(data["quota"])	2✔
236	elif "quota" in data and isinstance(data["quota"], Quota):	2✔
237	quota = data["quota"]	2✔
238	if "classes" in data and isinstance(data["classes"], set):	2✔
239	classes = [ResourceClass.from_dict(c) if isinstance(c, dict) else c for c in list(data["classes"])]	×
240	elif "classes" in data and isinstance(data["classes"], list):	2✔
241	classes = [ResourceClass.from_dict(c) if isinstance(c, dict) else c for c in data["classes"]]	2✔
242	return cls(	2✔
243	name=data["name"],
244	id=data.get("id"),
245	classes=classes,
246	quota=quota,
247	default=data.get("default", False),
248	public=data.get("public", False),
249	idle_threshold=data.get("idle_threshold"),
250	hibernation_threshold=data.get("hibernation_threshold"),
251	)

SwissDataScienceCenter / renku-data-services / 9446003648

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