add new fields (#211)

* add new fields

* fix

* fix

* add missing Session.result

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

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:

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:

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,

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:
         """