Replace ClusterUtilization.id with ClusterUtilization.cluster

allenai · May 5, 2022 · 25676aa · 25676aa
1 parent 7533fa8
commit 25676aa
Show file tree

Hide file tree

Showing 4 changed files with 15 additions and 2 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -7,6 +7,11 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## Unreleased
 
+### Changed
+
+- Replaced `id: str` field of `ClusterUtilization` with `cluster: Cluster`. `id` is still available
+  as a property for backwards compatibility.
+
 ## [v0.14.0](https://github.com/allenai/beaker-py/releases/tag/v0.14.0) - 2022-05-04
 
 ### Added

diff --git a/beaker/data_model/cluster.py b/beaker/data_model/cluster.py
@@ -67,11 +67,15 @@ def is_active(self) -> bool:
 
 
 class ClusterUtilization(BaseModel):
-    id: str
+    cluster: Cluster
     running_jobs: int
     queued_jobs: int
     nodes: Tuple[NodeUtilization, ...]
 
+    @property
+    def id(self) -> str:
+        return self.cluster.id
+
 
 class ClusterSpec(BaseModel):
     name: str

diff --git a/beaker/services/cluster.py b/beaker/services/cluster.py
@@ -215,7 +215,7 @@ def utilization(self, cluster: Union[str, Cluster]) -> ClusterUtilization:
                         node_util["cpus_used"] += job.requests.cpu_count
 
         return ClusterUtilization(
-            id=cluster.id,
+            cluster=cluster,
             running_jobs=running_jobs,
             queued_jobs=queued_jobs,
             nodes=[

diff --git a/tests/cluster_test.py b/tests/cluster_test.py
@@ -23,6 +23,10 @@ def test_cluster_get_on_prem(client: Beaker, beaker_on_prem_cluster_name: str):
     assert cluster.node_shape is None
 
 
+def test_cluster_utilization(client: Beaker, beaker_on_prem_cluster_name: str):
+    client.cluster.utilization(beaker_on_prem_cluster_name)
+
+
 def test_cluster_list(client: Beaker, beaker_org: Organization):
     client.cluster.list(beaker_org)