to_dict

dask · Mar 5, 2022 · 825112a · 825112a
1 parent 41d6e88
commit 825112a
Show file tree

Hide file tree

Showing 3 changed files with 19 additions and 6 deletions.
diff --git a/distributed/tests/test_worker.py b/distributed/tests/test_worker.py
@@ -3272,7 +3272,7 @@ async def test_Worker__to_dict(c, s, a):
     x = c.submit(inc, 1, key="x")
     await wait(x)
     d = a._to_dict()
-    assert d.keys() == {
+    assert set(d) == {
         "type",
         "id",
         "scheduler",
@@ -3288,7 +3288,13 @@ async def test_Worker__to_dict(c, s, a):
         "in_flight_workers",
         "log",
         "tasks",
-        "memory_manager",
+        "data",
+        "max_spill",
+        "memory_limit",
+        "memory_monitor_interval",
+        "memory_pause_fraction",
+        "memory_spill_fraction",
+        "memory_target_fraction",
         "logs",
         "config",
         "incoming_transfer_log",
@@ -3297,6 +3303,7 @@ async def test_Worker__to_dict(c, s, a):
         "pending_data_per_worker",
     }
     assert d["tasks"]["x"]["key"] == "x"
+    assert d["data"] == ["x"]
 
 
 @gen_cluster(client=True, nthreads=[("", 1)])

diff --git a/distributed/worker.py b/distributed/worker.py
@@ -1213,16 +1213,13 @@ def _to_dict(self, *, exclude: Container[str] = ()) -> dict:
             "in_flight_workers": self.in_flight_workers,
             "log": self.log,
             "tasks": self.tasks,
-            "memory_limit": self.memory_manager.memory_limit,
-            "memory_target_fraction": self.memory_target_fraction,
-            "memory_spill_fraction": self.memory_spill_fraction,
-            "memory_pause_fraction": self.memory_pause_fraction,
             "logs": self.get_logs(),
             "config": dask.config.config,
             "incoming_transfer_log": self.incoming_transfer_log,
             "outgoing_transfer_log": self.outgoing_transfer_log,
         }
         info.update(extra)
+        info.update(self.memory_manager._to_dict(exclude=exclude))
         info = {k: v for k, v in info.items() if k not in exclude}
         return recursive_to_dict(info, exclude=exclude)
 

diff --git a/distributed/worker_memory.py b/distributed/worker_memory.py
@@ -264,6 +264,15 @@ async def _maybe_spill(self, worker: Worker, memory: int) -> None:
                 format_bytes(total_spilled),
             )
 
+    def _to_dict(self, *, exclude: Container[str] = ()) -> dict:
+        info = {
+            k: v
+            for k, v in self.__dict__.items()
+            if not k.startswith("_") and k != "data" and k not in exclude
+        }
+        info["data"] = list(self.data)
+        return info
+
 
 class NannyMemoryManager:
     memory_limit: int | None