Source code for beaker.data_model.cluster

from datetime import datetime
from typing import Optional, Tuple

from .base import BaseModel, StrEnum, field_validator
from .job import Job
from .node import NodeResources, NodeUtilization

__all__ = ["ClusterStatus", "Cluster", "ClusterUtilization", "ClusterSpec", "ClusterPatch"]


[docs]class ClusterStatus(StrEnum): """ Current status of a cluster. """ pending = "pending" active = "active" terminated = "terminated" failed = "failed"
[docs]class Cluster(BaseModel): id: str name: str full_name: str created: datetime autoscale: bool capacity: int preemptible: bool status: ClusterStatus status_message: Optional[str] = None node_spec: Optional[NodeResources] = None """ The requested node configuration. """ node_shape: Optional[NodeResources] = None """ The actual node configuration. """ node_cost: Optional[str] = None validated: Optional[datetime] = None user_restrictions: Optional[Tuple[str, ...]] = None allow_preemptible_restriction_exceptions: Optional[bool] = None compute_source: Optional[str] = None max_job_timeout: Optional[int] = None require_preemptible_tasks: Optional[bool] = None @field_validator("validated") def _validate_datetime(cls, v: Optional[datetime]) -> Optional[datetime]: if v is not None and v.year == 1: return None return v @field_validator("node_spec") def _validate_node_spec(cls, v: Optional[NodeResources]) -> Optional[NodeResources]: if v is not None and not v.to_json(): return None return v @property def is_cloud(self) -> bool: """ Returns ``True`` is the cluster is a cloud cluster, otherwise ``False``. """ return self.node_shape is not None and self.node_spec is not None @property def is_active(self) -> bool: """ Returns ``True`` if the cluster is ready to be used. """ return not self.is_cloud or self.status == ClusterStatus.active
[docs]class ClusterUtilization(BaseModel): cluster: Cluster running_jobs: int queued_jobs: int running_preemptible_jobs: int nodes: Tuple[NodeUtilization, ...] jobs: Tuple[Job, ...] @property def id(self) -> str: return self.cluster.id
[docs]class ClusterSpec(BaseModel): name: str capacity: int preemptible: bool spec: NodeResources
[docs]class ClusterPatch(BaseModel): capacity: Optional[int] = None allow_preemptible_restriction_exceptions: Optional[bool] = None