Fix Task.replica_rank

allenai · Jul 13, 2024 · e93f710 · e93f710
1 parent 8e86fcc
commit e93f710
Show file tree

Hide file tree

Showing 2 changed files with 15 additions and 1 deletion.
diff --git 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
@@ -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: