"""
Entities related to resource management.
"""
from __future__ import annotations
from typing import Annotated
from pydantic import Field, NonNegativeInt, computed_field, model_validator
from ska_sdp_config.entity.common import ProcessingBlockId
from .base import EntityKeyBaseModel, MultiEntityBaseModel
RESOURCE_TYPE_UNIT_MAPPING = {
"performance-buffer-storage": "MB",
"capacity-buffer-storage": "MB",
"processing-node": "",
}
SUPPORTED_KIND = RESOURCE_TYPE_UNIT_MAPPING.keys()
KIND_PATTERN = "|".join(SUPPORTED_KIND)
PHASES_NAME_PATTERN = r"^\d+-[a-zA-Z]+$"
[docs]
class RequestAllocationKey(EntityKeyBaseModel):
"""Primary key for Request or Allocation."""
pb_id: ProcessingBlockId
"""
The ID of the Processing Block that the request/allocation belongs to.
"""
kind: Annotated[str, Field(pattern=KIND_PATTERN)]
"""The kind of this particular request/allocation/resource."""
name: Annotated[str, Field(pattern=r"^[A-Za-z0-9-]{1,96}$")]
"""
Unique name of a request/allocation/resource of
the given kind. In case of allocations, they need to match
the name of the corresponding request.
"""
[docs]
class ResourceBaseModel(MultiEntityBaseModel):
"""
Shared model between resource management entities.
"""
quantity: NonNegativeInt = 0
"""Quantity of request/allocation/resource."""
tags: list[str] = []
"""
Tags specifying the sub-resource types.
E.g. "cpu_node", "gpu_node".
"""
instances: list[str] = []
"""
List of identifiers of a certain resource type,
e.g. node names, subnets, etc.
"""
@computed_field
@property
def unit(self) -> str:
"""Unit of the quantity."""
return RESOURCE_TYPE_UNIT_MAPPING[self.key.kind]
[docs]
def is_same_kind_as(self, other_entity: ResourceBaseModel) -> bool:
"""Does this entity share a kind with another one?"""
return self.key.kind == other_entity.key.kind
[docs]
def is_quantity_subset_of(self, other_entity: ResourceBaseModel) -> bool:
"""Is this entity's quantity <= than the quanity of another?"""
return self.quantity <= other_entity.quantity
[docs]
def is_subset_of(self, other_entity: ResourceBaseModel) -> bool:
"""
Is this entity a subset of another based on
kind, quantity, and tags?
"""
return (
self.is_same_kind_as(other_entity)
and self.is_quantity_subset_of(other_entity)
and self.are_tags_subset_of(other_entity)
)
[docs]
@model_validator(mode="after")
def validate_data(self) -> ResourceBaseModel:
"""Run post-init validation"""
if self.quantity < len(self.instances):
raise ValueError(
"Quantity must be at least as large as the "
"total number of instances."
)
return self
[docs]
class Resource(ResourceBaseModel):
"""
Indicates platform availability of a resource.
The resource manager will update free space quantity
and status on each cycle.
Resources types can include, but not limited to:
buffer, cpu, memory, nodes etc.
"""
[docs]
class Key(EntityKeyBaseModel):
"""A resource primary key."""
kind: Annotated[str, Field(pattern=KIND_PATTERN)]
"""The kind of this particular resource."""
name: Annotated[str, Field(pattern=r"^[A-Za-z0-9-]{1,96}$")]
"""The name of this resource (e.g. a PVC of choice)."""
key: Key
information: dict | None = None
"""Optional specifications such as contingency reserve"""
[docs]
class Request(ResourceBaseModel):
"""
Indicates a request for a resource.
Created by processing scripts to request
resources of a certain kind.
"""
[docs]
class Key(RequestAllocationKey):
"""
A request primary key. No overrides.
Needed to be redefined here for Key pattern validation
"""
key: Key
phases: Annotated[
list[Annotated[str, Field(pattern=PHASES_NAME_PATTERN)]],
Field(default_factory=list),
]
[docs]
class Allocation(ResourceBaseModel):
"""
Allocation for a resource request,
created by the processing controller.
"""
[docs]
class Key(RequestAllocationKey):
"""
A request primary key. No overrides.
Needed to be redefined here for Key pattern validation
"""
key: Key
resource_link: Resource.Key | None = None
"""Key to Resource this Allocation is taken from."""
request_link: Request.Key | None = None
"""Key to the Request this Allocation fulfills."""
[docs]
@model_validator(mode="after")
def validate_data(self) -> Allocation:
super().validate_data()
if self.resource_link and self.request_link:
if (
self.key.kind != self.resource_link.kind
or (self.key.kind != self.request_link.kind)
or (self.resource_link.kind != self.request_link.kind)
):
raise ValueError(
f"Cannot create Allocation of kind {self.key.kind} "
f"from Resource of kind {self.resource_link.kind} to "
f"Request of kind {self.request_link.kind}"
)
if self.resource_link:
if self.key.kind != self.resource_link.kind:
raise ValueError(
f"Cannot create Allocation of kind {self.key.kind} "
f"from Resource of kind {self.resource_link.kind}"
)
if self.request_link and self.resource_link is None:
raise ValueError(
"Cannot create Allocation for a Request without a Resource."
)
return self