From 137d92ac40a05d44b522e6594ce4fbb853268483 Mon Sep 17 00:00:00 2001
From: Pete <epwalsh10@gmail.com>
Date: Fri, 17 Feb 2023 16:14:39 -0800
Subject: [PATCH] add new fields (#211)

* add new fields

* fix

* fix

* add missing Session.result
---
 CHANGELOG.md                    |  5 +++++
 beaker/data_model/experiment.py |  1 +
 beaker/data_model/job.py        | 13 +++++++++++--
 beaker/services/experiment.py   |  6 +-----
 beaker/services/job.py          |  9 ++++++---
 5 files changed, 24 insertions(+), 10 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 1ec2aa6..87e22a9 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -9,6 +9,11 @@ use patch releases for compatibility fixes instead.
 
 ## Unreleased
 
+### Added
+
+- Added new field `Job.result`, `Job.execution.result` is deprecated.
+- Added new field `Task.replica_rank`.
+
 ## [v1.17.4](https://github.com/allenai/beaker-py/releases/tag/v1.17.4) - 2023-02-17
 
 ### Fixed
diff --git a/beaker/data_model/experiment.py b/beaker/data_model/experiment.py
index 7c0e774..464905c 100644
--- a/beaker/data_model/experiment.py
+++ b/beaker/data_model/experiment.py
@@ -39,6 +39,7 @@ class Task(BaseModel):
     schedulable: bool = False
     jobs: Tuple[Job, ...] = Field(default_factory=tuple)
     owner: Optional[Account] = None
+    replica_rank: Optional[int] = None
 
     @property
     def display_name(self) -> str:
diff --git a/beaker/data_model/job.py b/beaker/data_model/job.py
index d366095..7c3b8f0 100644
--- a/beaker/data_model/job.py
+++ b/beaker/data_model/job.py
@@ -5,7 +5,14 @@
 
 from .account import Account
 from .base import BaseModel, IntEnum, StrEnum
-from .experiment_spec import DataMount, EnvVar, ImageSource, Priority, TaskSpec
+from .experiment_spec import (
+    DataMount,
+    EnvVar,
+    ImageSource,
+    Priority,
+    ResultSpec,
+    TaskSpec,
+)
 
 __all__ = [
     "CurrentJobStatus",
@@ -100,7 +107,7 @@ def current(self) -> CurrentJobStatus:
 
 
 class ExecutionResult(BaseModel):
-    beaker: str
+    beaker: Optional[str] = None
 
 
 class JobRequests(BaseModel):
@@ -145,6 +152,7 @@ class Session(BaseModel):
     work_dir: Optional[str] = None
     identity: Optional[str] = None
     constraints: Optional[Dict[str, List[str]]] = None
+    result: Optional[ResultSpec] = None
 
 
 class Job(BaseModel):
@@ -169,6 +177,7 @@ class Job(BaseModel):
     session: Optional[Session] = None
     host_networking: bool = False
     port_mappings: Optional[Dict[str, int]] = None
+    result: Optional[ExecutionResult] = None
 
     @property
     def display_name(self) -> str:
diff --git a/beaker/services/experiment.py b/beaker/services/experiment.py
index 058afc9..e0b2264 100644
--- a/beaker/services/experiment.py
+++ b/beaker/services/experiment.py
@@ -352,11 +352,7 @@ def results(
         if job is None:
             return None
         else:
-            assert job.execution is not None  # for mypy
-            try:
-                return self.beaker.dataset.get(job.execution.result.beaker)
-            except DatasetNotFound:
-                return None
+            return self.beaker.job.results(job)
 
     def wait_for(
         self,
diff --git a/beaker/services/job.py b/beaker/services/job.py
index bad691c..d52f4bd 100644
--- a/beaker/services/job.py
+++ b/beaker/services/job.py
@@ -203,11 +203,14 @@ def results(self, job: Union[str, Job]) -> Optional[Dataset]:
         :raises RequestException: Any other exception that can occur when contacting the
             Beaker server.
         """
-        job = self.get(job.id if isinstance(job, Job) else job)
-        if job.execution is None:
+        job = job if isinstance(job, Job) else self.get(job)
+        if job.result is None or job.result.beaker is None:
             return None
         else:
-            return self.beaker.dataset.get(job.execution.result.beaker)
+            try:
+                return self.beaker.dataset.get(job.result.beaker)
+            except DatasetNotFound:
+                return None
 
     def finalize(self, job: Union[str, Job]) -> Job:
         """