From e93f7100269b4a5671195cf916785379101377de Mon Sep 17 00:00:00 2001
From: epwalsh <petew@allenai.org>
Date: Sat, 13 Jul 2024 15:26:53 -0700
Subject: [PATCH] Fix `Task.replica_rank`

---
 CHANGELOG.md                    |  4 ++++
 beaker/data_model/experiment.py | 12 +++++++++++-
 2 files changed, 15 insertions(+), 1 deletion(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 69edf55..34c7cfc 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -9,6 +9,10 @@ use patch releases for compatibility fixes instead.
 
 ## Unreleased
 
+### Fixed
+
+- Fixed `Task.replica_rank` to be set to the right value now instead of always `None`.
+
 ## [v1.31.0](https://github.com/allenai/beaker-py/releases/tag/v1.31.0) - 2024-07-08
 
 ### Added
diff --git a/beaker/data_model/experiment.py b/beaker/data_model/experiment.py
index f339b5f..9540a71 100644
--- a/beaker/data_model/experiment.py
+++ b/beaker/data_model/experiment.py
@@ -41,7 +41,17 @@ class Task(BaseModel):
     schedulable: bool = False
     jobs: Tuple[Job, ...] = Field(default_factory=tuple)
     owner: Optional[Account] = None
-    replica_rank: Optional[int] = None
+    #  replica_rank: Optional[int] = None
+
+    @property
+    def replica_rank(self) -> Optional[int]:
+        if (
+            (job := self.latest_job) is not None
+            and job.execution is not None
+            and (replica_rank := job.execution.replica_rank) is not None
+        ):
+            return replica_rank
+        return None
 
     @property
     def display_name(self) -> str: