refactor: recovery attributes in config (#329)

Since we have added recovery attempts with a previous commit, that was hard coded with a value in the code as a constant. We want to make it configurable. For that reason, we moved it into the job configuration file, we added a new data structure with a default value and we updated all examples to use this new structure.

Author: raresgaia123

examples/llama3/auto/linear-no-recover.yaml

@@ -10,7 +10,9 @@ fwd_policy: "static"
 job_id: "job13"
 # maximum number of requests in flight at any given point in time
 max_inflight: 4
-recover: False
+recovery:
+  enable: false # default: true
+  attempts: 3 # default: 5
 
 flow_graph:
   s-0:

examples/llama3/static/linear-no-recover.yaml

@@ -10,7 +10,9 @@ fwd_policy: "static"
 job_id: "job12"
 # maximum number of requests in flight at any given point in time
 max_inflight: 4
-recover: False
+recovery:
+  enable: false # default: true
+  attempts: 3 # default: 5
 
 # Note: IP addresses should be agents'
 flow_graph:

examples/resnet152/auto/diamond-no-recover.yaml

@@ -11,7 +11,9 @@ fwd_policy: "rr"
 job_id: "job8"
 # maximum number of requests in flight at any given point in time
 max_inflight: 8
-recover: False
+recovery:
+  enable: false # default: true
+  attempts: 3 # default: 5
 
 flow_graph:
   s-0:

examples/resnet152/auto/linear-no-recover.yaml

@@ -11,7 +11,9 @@ fwd_policy: "rr"
 job_id: "job9"
 # maximum number of requests in flight at any given point in time
 max_inflight: 8
-recover: False
+recovery:
+  enable: false # default: true
+  attempts: 3 # default: 5
 
 flow_graph:
   s-0:

examples/resnet152/static/diamond-no-recover.yaml

@@ -14,7 +14,9 @@ fwd_policy: "rr"
 job_id: "job10"
 # maximum number of requests in flight at any given point in time
 max_inflight: 800
-recover: False
+recovery:
+  enable: false # default: true
+  attempts: 3 # default: 5
 
 # Note: IP addresses should be agents'
 flow_graph:

examples/resnet152/static/linear-no-recover.yaml

@@ -11,7 +11,9 @@ fwd_policy: "rr"
 job_id: "job11"
 # maximum number of requests in flight at any given point in time
 max_inflight: 8
-recover: False
+recovery:
+  enable: false # default: true
+  attempts: 3 # default: 5
 
 # Note: IP addresses should be agents'
 flow_graph:

infscale/configs/job.py

@@ -85,6 +85,14 @@ class WorkerData:
     recover: bool = False
 
 
+@dataclass
+class Recovery:
+    """Specification about recovery."""
+
+    enable: bool = True
+    attempts: int = 5
+
+
 @dataclass
 class ServeConfig:
     """Class for keeping config values of serve specification."""
@@ -228,18 +236,23 @@ class JobConfig:
     flow_graph: dict[str, list[WorldInfo]]
     dataset: Dataset
     job_id: str
+    recovery: Recovery | dict | None = None
     nfaults: int = 0
     micro_batch_size: int = 8
     fwd_policy: str = "random"
     max_inflight: int = 1
-    recover: bool = True
     force_terminate: bool = False
 
     # this will be set by controller  based on its configuration
     reqgen_config: GenConfig | None = None
 
     def __post_init__(self) -> None:
         """Handle post init class variables."""
+        if self.recovery is None:
+            self.recovery = Recovery()
+        elif not isinstance(self.recovery, Recovery):
+            self.recovery = Recovery(**self.recovery)
+
         for k in list(self.flow_graph.keys()):
             for i, item in enumerate(self.flow_graph[k]):
                 world_info = item if isinstance(item, WorldInfo) else WorldInfo(**item)

infscale/controller/job_context.py

@@ -50,7 +50,6 @@
     from infscale.controller.controller import Controller
 
 MAX_RES_RECOVER_RETRIES = 8  # max resources retries with exponential backoff.
-MAX_DEPLOY_RECOVER_RETRIES = 5  # maximum deploy recovery retries
 
 
 logger = None
@@ -653,7 +652,7 @@ def cond_stopped(self):
 
     async def cond_recovery(self):
         """Handle the transition to failed."""
-        if self._recovery_count == MAX_DEPLOY_RECOVER_RETRIES:
+        if self._recovery_count == self.context._cur_cfg.recovery.attempts:
             failed_wrk_ids = self._get_failed_wrk_ids()
             await self._remove_pipeline_n_update(failed_wrk_ids, self.context._cur_cfg)
 
@@ -755,7 +754,7 @@ async def do_wrk_cond(self, wid: str, status: WorkerStatus) -> None:
                 await self.send_check_loop_command()
 
                 # if failure happens while starting, current config is None
-                if self._cur_cfg is not None and self._cur_cfg.recover:
+                if self._cur_cfg is not None and self._cur_cfg.recovery.enable:
                     await self.cond_recovery()
 
                     return