refactor(adaptive-core): improved readability recommendations

distributed/deploy/adaptive.py

@@ -8,7 +8,7 @@
 import dask.config
 from dask.utils import parse_timedelta
 
-from distributed.deploy.adaptive_core import AdaptiveCore
+from distributed.deploy.adaptive_core import AdaptiveCore, Recommendation
 from distributed.protocol import pickle
 from distributed.utils import log_errors
 
@@ -152,7 +152,7 @@ async def target(self):
             target_duration=self.target_duration
         )
 
-    async def recommendations(self, target: int) -> dict:
+    async def recommendations(self, target: int) -> Recommendation:
         if len(self.plan) != len(self.requested):
             # Ensure that the number of planned and requested workers
             # are in sync before making recommendations.

distributed/deploy/adaptive_core.py

@@ -4,7 +4,9 @@
 import math
 from collections import defaultdict, deque
 from collections.abc import Iterable
+from dataclasses import asdict, dataclass
 from datetime import timedelta
+from enum import Enum, auto
 from typing import TYPE_CHECKING, cast
 
 import tlz as toolz
@@ -23,6 +25,22 @@
 logger = logging.getLogger(__name__)
 
 
+class RecommendationStatus(Enum):
+    Up = auto()
+    Down = auto()
+    Same = auto()
+
+
+@dataclass
+class Recommendation:
+    status: RecommendationStatus
+    workers: list = []
+    n: int = 0
+
+
+RS = RecommendationStatus
+
+
 class AdaptiveCore:
     """
     The core logic for adaptive deployments, with none of the cluster details
@@ -169,13 +187,13 @@ async def safe_target(self) -> int:
 
         return n
 
-    async def scale_down(self, n: int) -> None:
+    async def scale_down(self, workers: Iterable) -> None:
         raise NotImplementedError()
 
-    async def scale_up(self, workers: Iterable) -> None:
+    async def scale_up(self, n: int) -> None:
         raise NotImplementedError()
 
-    async def recommendations(self, target: int) -> dict:
+    async def recommendations(self, target: int) -> Recommendation:
         """
         Make scale up/down recommendations based on current state and target
         """
@@ -185,11 +203,11 @@ async def recommendations(self, target: int) -> dict:
 
         if target == len(plan):
             self.close_counts.clear()
-            return {"status": "same"}
+            return Recommendation(status=RS.Same)
 
         if target > len(plan):
             self.close_counts.clear()
-            return {"status": "up", "n": target}
+            return Recommendation(status=RS.Up, n=target)
 
         # target < len(plan)
         not_yet_arrived = requested - observed
@@ -212,9 +230,11 @@ async def recommendations(self, target: int) -> dict:
                 del self.close_counts[k]
 
         if firmly_close:
-            return {"status": "down", "workers": list(firmly_close)}
+            return Recommendation(
+                status=RecommendationStatus.Down, workers=list(firmly_close)
+            )
         else:
-            return {"status": "same"}
+            return Recommendation(status=RS.Same)
 
     async def adapt(self) -> None:
         """
@@ -229,18 +249,16 @@ async def adapt(self) -> None:
 
         try:
             target = await self.safe_target()
-            recommendations = await self.recommendations(target)
+            recommendation = await self.recommendations(target)
 
-            if recommendations["status"] != "same":
-                self.log.append((time(), dict(recommendations)))
-
-            status = recommendations.pop("status")
-            if status == "same":
+            if recommendation.status == RS.Same:
                 return
-            if status == "up":
-                await self.scale_up(**recommendations)
-            if status == "down":
-                await self.scale_down(**recommendations)
+            else:
+                self.log.append((time(), asdict(recommendation)))
+                if recommendation.status == RS.Up:
+                    await self.scale_up(recommendation.n)
+                elif recommendation.status == RS.Down:
+                    await self.scale_down(recommendation.workers)
         except OSError:
             if status != "down":
                 logger.error("Adaptive stopping due to error", exc_info=True)