From 2bb86d81521d19149448f0fa2385b1b30f0e45ab Mon Sep 17 00:00:00 2001
From: epwalsh <petew@allenai.org>
Date: Wed, 11 Dec 2024 10:02:22 -0800
Subject: [PATCH] Add two missing fields

---
 CHANGELOG.md                 | 5 +++++
 beaker/data_model/cluster.py | 1 +
 beaker/data_model/job.py     | 1 +
 3 files changed, 7 insertions(+)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index e32edb8..89118d1 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -9,6 +9,11 @@ use patch releases for compatibility fixes instead.
 
 ## Unreleased
 
+### Added
+
+- Added new field `Cluster.max_session_timeout: Optional[int]`.
+- Added new field `JobExecution.retry_ancestor: Optional[str]`.
+
 ### Fixed
 
 - Made checking for version upgrades more robust.
diff --git a/beaker/data_model/cluster.py b/beaker/data_model/cluster.py
index 43f8884..f8f0315 100644
--- a/beaker/data_model/cluster.py
+++ b/beaker/data_model/cluster.py
@@ -43,6 +43,7 @@ class Cluster(BaseModel):
     allow_preemptible_restriction_exceptions: Optional[bool] = None
     compute_source: Optional[str] = None
     max_job_timeout: Optional[int] = None
+    max_session_timeout: Optional[int] = None
     require_preemptible_tasks: Optional[bool] = None
 
     @field_validator("validated")
diff --git a/beaker/data_model/job.py b/beaker/data_model/job.py
index 6f4a3f7..3026807 100644
--- a/beaker/data_model/job.py
+++ b/beaker/data_model/job.py
@@ -133,6 +133,7 @@ class JobExecution(BaseModel):
     workspace: Optional[str] = None
     replica_rank: Optional[int] = None
     replica_group_id: Optional[str] = None
+    retry_ancestor: Optional[str] = None
 
 
 class JobKind(StrEnum):