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: