Merge pull request #13 from richardliaw/htop

[htop] autoscaler

python/ray/autoscaler/_private/htop.py

@@ -1,3 +1,4 @@
+import json
 import os  # noqa: F401
 from queue import Empty, Queue
 from typing import Dict, Iterable, Optional, Tuple
@@ -17,6 +18,11 @@
 from rich.table import Table
 from rich.text import Text
 
+import ray.ray_constants as ray_constants
+from ray.autoscaler._private.load_metrics import LoadMetricsSummary
+from ray.autoscaler._private.autoscaler import AutoscalerSummary
+from ray.autoscaler._private.util import DEBUG_AUTOSCALING_STATUS
+
 try:
     # Windows
     import msvcrt
@@ -234,21 +240,34 @@ def __rich__(self) -> Layout:
 
 
 class DataManager:
-    def __init__(self, url: str):
+    def __init__(self, url: str, mock_autoscaler=True):
         self.url = url
 
-        self.mock = None
-        self.mock = os.path.expanduser("~/coding/sandbox/out2.json")
+        self.mock = os.path.expanduser(os.environ.get("RAY_HTOP_MOCK"))
         self.mock_cache = None
 
         self.nodes = []
+
+        # Autoscaler info
+        self.ray_address = None
+        self.redis_client = None
+        self.mock_autoscaler = mock_autoscaler
+        self.autoscaler_summary = None
+        self.lm_summary = None
         self.update()
 
+    def _create_redis_client(
+            self, address,
+            redis_password=ray_constants.REDIS_DEFAULT_PASSWORD):
+        import ray._private.services as services
+        if not address:
+            address = services.get_ray_address_to_use_or_die()
+        redis_client = services.create_redis_client(address, redis_password)
+        self.redis_client = redis_client
+
     def update(self):
         if self.mock:
             if not self.mock_cache:
-                import json
-
                 with open(self.mock, "rt") as f:
                     self.mock_cache = json.load(f)
             resp_json = self.mock_cache
@@ -261,6 +280,30 @@ def update(self):
 
         self.nodes = [Node(node_dict) for node_dict in resp_data["clients"]]
 
+        self._load_autoscaler_state()
+
+    def _load_autoscaler_state(self):
+        as_dict = None
+        if self.mock_autoscaler:
+            if isinstance(self.mock_autoscaler, str):
+                with open(self.mock_autoscaler_data) as f:
+                    as_dict = json.loads(f.read())
+        else:
+            if not self.redis_client:
+                self._create_redis_client(self.address)
+            status = self.redis_client.hget(DEBUG_AUTOSCALING_STATUS, "value")
+            if status:
+                status = status.decode("utf-8")
+                as_dict = json.loads(status)
+
+        if as_dict:
+            self.lm_summary = LoadMetricsSummary(
+                **as_dict["load_metrics_report"])
+            if "autoscaler_report" in as_dict:
+                self.autoscaler_summary = AutoscalerSummary(
+                    **as_dict["autoscaler_report"])
+            # TODO: process the autoscaler data.
+
 
 class Node:
     percent_only = (