Add retry field to TaskSpec

CHANGELOG.md

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

beaker/data_model/experiment_spec.py

@@ -16,6 +16,7 @@
     "TaskResources",
     "Priority",
     "TaskContext",
+    "TaskRetrySpec",
     "TaskSpec",
     "SpecVersion",
     "ExperimentSpec",
@@ -319,6 +320,18 @@ def __setitem__(self, key: str, val: List[Any]) -> None:
         setattr(self, key, val)
 
 
+class TaskRetrySpec(BaseModel, frozen=False):
+    """
+    Defines the retry behavior of a task.
+    """
+
+    allowed_task_retries: int
+    """
+    A positive integer specifying the maximum number of task retries allowed for the experiment,
+    with a max limit of 10.
+    """
+
+
 class TaskSpec(BaseModel, frozen=False):
     """
     A :class:`TaskSpec` defines a :class:`~beaker.data_model.experiment.Task` within an :class:`ExperimentSpec`.
@@ -345,6 +358,11 @@ class TaskSpec(BaseModel, frozen=False):
     Context describes how and where this task should run.
     """
 
+    retry: Optional[TaskRetrySpec] = None
+    """
+    Defines the retry behavior of a task.
+    """
+
     constraints: Optional[Constraints] = None
     """
     Each task can have many constraints. And each constraint can have many values.
@@ -559,6 +577,14 @@ def with_context(self, **kwargs) -> "TaskSpec":
         """
         return self.model_copy(deep=True, update={"context": TaskContext(**kwargs)})
 
+    def with_retries(self, allowed_task_retries: int) -> "TaskSpec":
+        """
+        Return a new :class:`TaskSpec` with the given number of retries.
+        """
+        return self.model_copy(
+            deep=True, update={"retries": TaskRetrySpec(allowed_task_retries=allowed_task_retries)}
+        )
+
     def with_name(self, name: str) -> "TaskSpec":
         """
         Return a new :class:`TaskSpec` with the given :data:`name`.