Mock process memory readings in test_worker.py

distributed/nanny.py

@@ -13,7 +13,7 @@
 from inspect import isawaitable
 from queue import Empty
 from time import sleep as sync_sleep
-from typing import ClassVar
+from typing import TYPE_CHECKING, ClassVar, Literal
 
 import psutil
 from tornado import gen
@@ -45,6 +45,9 @@
 )
 from .worker import Worker, parse_memory_limit, run
 
+if TYPE_CHECKING:
+    from .diagnostics.plugin import NannyPlugin
+
 logger = logging.getLogger(__name__)
 
 
@@ -94,6 +97,7 @@ def __init__(
         services=None,
         name=None,
         memory_limit="auto",
+        memory_terminate_fraction: float | Literal[False] | None = None,
         reconnect=True,
         validate=False,
         quiet=False,
@@ -203,8 +207,10 @@ def __init__(
         self.worker_kwargs = worker_kwargs
 
         self.contact_address = contact_address
-        self.memory_terminate_fraction = dask.config.get(
-            "distributed.worker.memory.terminate"
+        self.memory_terminate_fraction = (
+            memory_terminate_fraction
+            if memory_terminate_fraction is not None
+            else dask.config.get("distributed.worker.memory.terminate")
         )
 
         self.services = services
@@ -231,7 +237,7 @@ def __init__(
             "plugin_remove": self.plugin_remove,
         }
 
-        self.plugins = {}
+        self.plugins: dict[str, NannyPlugin] = {}
 
         super().__init__(
             handlers=handlers, io_loop=self.loop, connection_args=self.connection_args

distributed/tests/test_utils_test.py

@@ -15,7 +15,7 @@
 
 from distributed import Client, Nanny, Scheduler, Worker, config, default_client
 from distributed.compatibility import WINDOWS
-from distributed.core import Server, rpc
+from distributed.core import Server, Status, rpc
 from distributed.metrics import time
 from distributed.utils import mp_context
 from distributed.utils_test import (
@@ -28,6 +28,7 @@
     gen_cluster,
     gen_test,
     inc,
+    mock_rss,
     new_config,
     tls_only_security,
 )
@@ -607,3 +608,27 @@ def test_start_failure_scheduler():
     with pytest.raises(TypeError):
         with cluster(scheduler_kwargs={"foo": "bar"}):
             return
+
+
+@gen_cluster(
+    client=True,
+    worker_kwargs={"heartbeat_interval": "10ms", "memory_monitor_interval": "10ms"},
+)
+async def test_mock_rss(c, s, a, b):
+    # Test that it affects the readings sent to the Scheduler
+    mock_rss(a, 2e6)
+    while s.workers[a.address].memory.process != 2_000_000:
+        await asyncio.sleep(0.01)
+
+    # Test that the instance has been mocked, not the class
+    assert s.workers[b.address].memory.process > 10e6
+
+    # Test that it's compatible with Client.run and can be used with Nannies
+    await c.run(mock_rss, nbytes=3e6, workers=[b.address])
+    while s.workers[b.address].memory.process != 3_000_000:
+        await asyncio.sleep(0.01)
+
+    # Test that it affects Worker.memory_monitor
+    mock_rss(a, 100e9)
+    while s.workers[a.address].status != Status.paused:
+        await asyncio.sleep(0.01)