Skip to content

Worker disconnect notification improvements #558

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
magniloquency wants to merge 9 commits into
base: main
Choose a base branch
from
Draft

Worker disconnect notification improvements #558

Changes from all commits
Commits
Show all changes
9 commits
Select commit Hold shift + click to select a range
649d8a8
Improve YMQ error handling
magniloquency Feb 10, 2026
9f788f3
Rename result2 to callbackResult in PyIOSocket
magniloquency Feb 11, 2026
69e35b8
Rename result to callbackResult in PyIOContext and ymq.h
magniloquency Feb 11, 2026
f22adcc
Merge branch 'main' into ymq-callbacks
magniloquency Feb 11, 2026
22ced2b
format
magniloquency Feb 11, 2026
95bd35d
worker: improve disconnect logic and simplify signal registration
magniloquency Feb 11, 2026
b53a797
Merge branch 'main' into worker-disconnect-notification
sharpener6 Feb 12, 2026
c9dd8e2
Merge branch 'main' into worker-disconnect-notification
sharpener6 Feb 15, 2026
5e16fee
Merge branch 'main' into worker-disconnect-notification
magniloquency Feb 18, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
28 changes: 11 additions & 17 deletions src/scaler/worker/worker.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -10,16 +10,11 @@
import zmq.asyncio

from scaler.config.defaults import PROFILING_INTERVAL_SECONDS
from scaler.config.types.network_backend import NetworkBackend
from scaler.config.types.object_storage_server import ObjectStorageAddressConfig
from scaler.config.types.zmq import ZMQConfig, ZMQType
from scaler.io.async_binder import ZMQAsyncBinder
from scaler.io.mixins import AsyncBinder, AsyncConnector, AsyncObjectStorageConnector
from scaler.io.utility import (
create_async_connector,
create_async_object_storage_connector,
get_scaler_network_backend_from_env,
)
from scaler.io.utility import create_async_connector, create_async_object_storage_connector
from scaler.io.ymq import ymq
from scaler.protocol.python.message import (
ClientDisconnect,
Expand Down Expand Up @@ -269,8 +264,15 @@ async def __get_loops(self):
except Exception as e:
logging.exception(f"{self.identity!r}: failed with unhandled exception:\n{e}")

if get_scaler_network_backend_from_env() == NetworkBackend.tcp_zmq:
try:
await self._connector_external.send(DisconnectRequest.new_msg(self.identity))
except (asyncio.CancelledError, ymq.YMQException) as e:

# This can happen if the scheduler has already closed the connection.
if isinstance(e, asyncio.CancelledError) or e.code == ymq.ErrorCode.ConnectorSocketClosedByRemoteEnd:
pass
else:
logging.exception(f"{self.identity!r}: failed to send DisconnectRequest:\n{e}")

self._connector_external.destroy()
self._processor_manager.destroy("quit")
Expand All @@ -281,16 +283,8 @@ async def __get_loops(self):
logging.info(f"{self.identity!r}: quit")

def __register_signal(self):
backend = get_scaler_network_backend_from_env()
if backend == NetworkBackend.tcp_zmq:
self._loop.add_signal_handler(signal.SIGINT, self.__destroy)
self._loop.add_signal_handler(signal.SIGTERM, self.__destroy)
elif backend == NetworkBackend.ymq:
self._loop.add_signal_handler(signal.SIGINT, lambda: asyncio.ensure_future(self.__graceful_shutdown()))
self._loop.add_signal_handler(signal.SIGTERM, lambda: asyncio.ensure_future(self.__graceful_shutdown()))

async def __graceful_shutdown(self):
await self._connector_external.send(DisconnectRequest.new_msg(self.identity))
self._loop.add_signal_handler(signal.SIGINT, self.__destroy)
self._loop.add_signal_handler(signal.SIGTERM, self.__destroy)

def __destroy(self):
self._task.cancel()
Loading