Add retry field to ExperimentSpec (#290)

* Add `retry` field to `TaskSpec`

* Fix

* fixes

* let jobs run elsewhere

* cancel in progress jobs

* fix

.github/workflows/main.yml

@@ -2,7 +2,7 @@ name: Main
 
 concurrency:
   group: ${{ github.workflow }}-${{ github.ref }}
-  # cancel-in-progress: true
+  cancel-in-progress: true
 
 on:
   pull_request:

CHANGELOG.md

@@ -9,6 +9,10 @@ use patch releases for compatibility fixes instead.
 
 ## Unreleased
 
+### Added
+
+- Added `retry` field to `ExperimentSpec`.
+
 ## [v1.31.3](https://github.com/allenai/beaker-py/releases/tag/v1.31.3) - 2024-08-30
 
 ### Added

beaker/data_model/experiment_spec.py

@@ -18,6 +18,7 @@
     "TaskContext",
     "TaskSpec",
     "SpecVersion",
+    "RetrySpec",
     "ExperimentSpec",
     "Constraints",
 ]
@@ -705,6 +706,18 @@ class SpecVersion(StrEnum):
     v2_alpha = "v2-alpha"
 
 
+class RetrySpec(BaseModel, frozen=False):
+    """
+    Defines the retry behavior of an experiment.
+    """
+
+    allowed_task_retries: Optional[int] = None
+    """
+    A positive integer specifying the maximum number of task retries allowed for the experiment,
+    with a max limit of 10.
+    """
+
+
 class ExperimentSpec(BaseModel, frozen=False):
     """
     Experiments are the main unit of execution in Beaker.
@@ -749,6 +762,11 @@ class ExperimentSpec(BaseModel, frozen=False):
     Long-form explanation for an experiment.
     """
 
+    retry: Optional[RetrySpec] = None
+    """
+    Defines the retry behavior of an experiment.
+    """
+
     @field_validator("tasks")
     def _validate_tasks(cls, v: List[TaskSpec]) -> List[TaskSpec]:
         task_names = set()
@@ -882,6 +900,14 @@ def with_description(self, description: str) -> "ExperimentSpec":
         """
         return self.model_copy(deep=True, update={"description": description})
 
+    def with_retries(self, allowed_task_retries: int) -> "ExperimentSpec":
+        """
+        Return a new :class:`ExperimentSpec` with the given number of retries.
+        """
+        return self.model_copy(
+            deep=True, update={"retry": RetrySpec(allowed_task_retries=allowed_task_retries)}
+        )
+
     def validate(self):
         for task in self.tasks:
             if (task.image.beaker is None) == (task.image.docker is None):

integration_tests/experiments_test.py

@@ -16,7 +16,6 @@ def test_experiment_workflow(
     client: Beaker,
     experiment_name: str,
     alternate_experiment_name: str,
-    beaker_cluster_name: str,
     hello_world_experiment_name: str,
 ):
     spec = ExperimentSpec(
@@ -25,7 +24,7 @@ def test_experiment_workflow(
             TaskSpec(
                 name="main",
                 image=ImageSource(docker="hello-world"),
-                context=TaskContext(cluster=beaker_cluster_name),
+                context=TaskContext(preemptible=True),
                 result=ResultSpec(path="/unused"),  # required even if the task produces no output.
             ),
         ],

integration_tests/jobs_test.py

@@ -3,12 +3,11 @@
 from beaker import Beaker, ExperimentSpec, TaskSpec
 
 
-def test_job_stop_and_finalize(client: Beaker, experiment_name: str, beaker_cluster_name: str):
+def test_job_stop_and_finalize(client: Beaker, experiment_name: str):
     start = time.time()
     spec = ExperimentSpec(budget="ai2/allennlp").with_task(
         TaskSpec.new(
             "main",
-            beaker_cluster_name,
             docker_image="hello-world",
         ),
     )