Add timeouts to Flower workload setup

Author: Madhav

aiopslab/orchestrator/problems/flower_model_misconfig/model_misconfig.py

@@ -3,8 +3,9 @@
 
 """Model misconfiguration fault in the Flower application."""
 
+import os
 import time
-from typing import Any
+from typing import Any, Callable
 
 from aiopslab.orchestrator.tasks import *
 from aiopslab.service.dock import Docker
@@ -13,6 +14,13 @@
 from aiopslab.session import SessionItem
 from aiopslab.generators.fault.inject_virtual import VirtualizationFaultInjector
 
+WORKLOAD_START_TIMEOUT_SECONDS = int(
+    os.getenv("AIOPSLAB_FLOWER_WORKLOAD_START_TIMEOUT_SECONDS", "300")
+)
+FAULT_PROPAGATION_TIMEOUT_SECONDS = int(
+    os.getenv("AIOPSLAB_FLOWER_FAULT_PROPAGATION_TIMEOUT_SECONDS", "300")
+)
+
 
 class FlowerModelMisconfigBaseTask:
     def __init__(self, faulty_service: str = "user-service"):
@@ -25,30 +33,48 @@ def __init__(self, faulty_service: str = "user-service"):
     def start_workload(self):
         print("== Start Workload ==")
         command = "flwr run train local-deployment"
-        self.docker.exec_command(command, cwd=self.train_dir)
+        self.docker.exec_command(
+            command,
+            cwd=self.train_dir,
+            timeout=WORKLOAD_START_TIMEOUT_SECONDS,
+        )
 
         path = "/app/.flwr/apps"
         check = f""" docker exec -it {self.faulty_service} sh -c "test -d {path} && echo 'exists'" """
 
         print("Waiting for workload to start...")
-        while True:
-            exists = self.docker.exec_command(check)
-            if exists.strip() == "exists":
-                break
-            time.sleep(1)
+        self._wait_until(
+            "Flower workload start",
+            lambda: self.docker.exec_command(check).strip() == "exists",
+            WORKLOAD_START_TIMEOUT_SECONDS,
+        )
         print("Workload started successfully.")
 
         # Inject fault after workload starts, since the required files are created during the workload
         print("Injecting fault...")
         self.inject_fault(inject=True)
 
         print("Waiting for faults to propagate...")
-        while True:
-            logs = self.docker.get_logs(self.faulty_service)
-            if "error" in logs.lower():
-                break
-            time.sleep(1)
+        self._wait_until(
+            "Flower fault propagation",
+            lambda: "error" in self.docker.get_logs(self.faulty_service).lower(),
+            FAULT_PROPAGATION_TIMEOUT_SECONDS,
+        )
         print("Faults propagated.")
+
+    def _wait_until(
+        self,
+        description: str,
+        predicate: Callable[[], bool],
+        timeout_seconds: int,
+        interval_seconds: float = 1,
+    ):
+        deadline = time.monotonic() + timeout_seconds
+        while time.monotonic() < deadline:
+            if predicate():
+                return
+            time.sleep(interval_seconds)
+        raise TimeoutError(f"{description} did not complete within {timeout_seconds} seconds.")
 
     def inject_fault(self, inject: bool = False):
         print("== Fault Injection ==")

aiopslab/service/dock.py

@@ -38,7 +38,7 @@ def cleanup(self):
         command = "docker container prune -f"
         return self.exec_command(command)
 
-    def exec_command(self, command: str, input_data=None, cwd=None):
+    def exec_command(self, command: str, input_data=None, cwd=None, timeout=None):
         """Execute an arbitrary command."""
         if input_data is not None:
             input_data = input_data.encode("utf-8")
@@ -50,8 +50,13 @@ def exec_command(self, command: str, input_data=None, cwd=None):
                 stderr=subprocess.PIPE,
                 shell=True,
                 cwd=cwd,
+                timeout=timeout,
             )
             if out is not None:
                 return out.stdout.decode("utf-8")
+        except subprocess.TimeoutExpired as e:
+            raise RuntimeError(
+                f"Timed out after {timeout} seconds executing command: {command}"
+            ) from e
         except subprocess.CalledProcessError as e:
             return e.stderr.decode("utf-8")

tests/problems/test_flower_model_misconfig.py

@@ -0,0 +1,18 @@
+import unittest
+
+from aiopslab.orchestrator.problems.flower_model_misconfig.model_misconfig import (
+    FlowerModelMisconfigBaseTask,
+)
+
+
+class FlowerModelMisconfigTimeoutTest(unittest.TestCase):
+    def test_wait_until_returns_when_predicate_succeeds(self):
+        task = object.__new__(FlowerModelMisconfigBaseTask)
+
+        task._wait_until("ready", lambda: True, timeout_seconds=1, interval_seconds=0)
+
+    def test_wait_until_times_out_when_predicate_never_succeeds(self):
+        task = object.__new__(FlowerModelMisconfigBaseTask)
+
+        with self.assertRaises(TimeoutError):
+            task._wait_until("ready", lambda: False, timeout_seconds=0, interval_seconds=0)

tests/service/test_docker.py

@@ -0,0 +1,11 @@
+import unittest
+
+from aiopslab.service.dock import Docker
+
+
+class DockerExecTimeoutTest(unittest.TestCase):
+    def test_exec_command_times_out(self):
+        docker = object.__new__(Docker)
+
+        with self.assertRaises(RuntimeError):
+            docker.exec_command("sleep 1", timeout=0.01)