Stop keep alives when worker reconnecting to the scheduler
@@ -382,7 +382,6 @@ class Worker(ServerNode):
self.executed_count = 0
self.long_running = set()
- self.batched_stream = None
self.recent_messages_log = deque(
maxlen=dask.config.get("distributed.comm.recent-messages-log-length")
)
@@ -573,6 +572,7 @@ class Worker(ServerNode):
self.actor_executor = ThreadPoolExecutor(
1, thread_name_prefix="Dask-Actor-Threads"
)
+ self.batched_stream = BatchedSend(interval="2ms", loop=self.loop)
self.name = name
self.scheduler_delay = 0
self.stream_comms = dict()
@@ -650,6 +650,13 @@ class Worker(ServerNode):
pc = PeriodicCallback(self.heartbeat, 1000, io_loop=self.io_loop)
self.periodic_callbacks["heartbeat"] = pc
+ pc = PeriodicCallback(
+ lambda: self.batched_stream.send({"op": "keep-alive"}),
+ 60000,
+ io_loop=self.io_loop,
+ )
+ self.periodic_callbacks["keep-alive"] = pc
+
self._address = contact_address
if self.memory_limit:
@@ -797,6 +804,7 @@ class Worker(ServerNode):
#####################
async def _register_with_scheduler(self):
+ self.periodic_callbacks["keep-alive"].stop()
self.periodic_callbacks["heartbeat"].stop()
start = time()
if self.contact_address is None:
@@ -863,15 +871,8 @@ class Worker(ServerNode):
logger.info(" Registered to: %26s", self.scheduler.address)
logger.info("-" * 49)
- self.batched_stream = BatchedSend(interval="2ms", loop=self.loop)
self.batched_stream.start(comm)
- pc = PeriodicCallback(
- lambda: self.batched_stream.send({"op": "keep-alive"}),
- 60000,
- io_loop=self.io_loop,
- )
- self.periodic_callbacks["keep-alive"] = pc
- pc.start()
+ self.periodic_callbacks["keep-alive"].start()
self.periodic_callbacks["heartbeat"].start()
self.loop.add_callback(self.handle_scheduler, comm)
@@ -1112,7 +1113,11 @@ class Worker(ServerNode):
for k, v in self.services.items():
v.stop()
- if self.batched_stream and not self.batched_stream.comm.closed():
+ if (
+ self.batched_stream
+ and self.batched_stream.comm
+ and not self.batched_stream.comm.closed()
+ ):
self.batched_stream.send({"op": "close-stream"})
if self.batched_stream:
