Fix mypy errors, restore web_config, and update WorkerDisconnectNotification handling

- Restore missing web_config in NativeWorkerAdapterConfig
- Fix type incompatibility in Cluster for per_worker_capabilities
- Ensure async worker group operations are properly awaited in scaling policies and adapters
- Update protocol documentation and implementation for WorkerDisconnectNotification (WDN)
- Reformat modified files with black and isort

Author: magniloquency

src/scaler/config/section/native_worker_adapter.py

@@ -11,7 +11,6 @@
 
 @dataclasses.dataclass
 class NativeWorkerAdapterConfig(ConfigClass):
-    web_config: WebConfig
     worker_adapter_config: WorkerAdapterConfig
     worker_config: WorkerConfig = dataclasses.field(default_factory=WorkerConfig)
     logging_config: LoggingConfig = dataclasses.field(default_factory=LoggingConfig)

src/scaler/protocol/worker.md

@@ -227,11 +227,11 @@ consider the worker dead.
 
 When a `DisconnectRequest` message is sent, the worker should disconnect from the scheduler.
 
-### WorkerDisconnectNotification `WD` (optional)
+### WorkerDisconnectNotification `WDN` (optional)
 
 | message_type | worker  |
 |:------------:|:-------:|
-|    b"WD"     | X bytes |
+|    b"WDN"     | X bytes |
 
 * worker: Worker ID
 

src/scaler/scheduler/controllers/mixins.py

@@ -13,6 +13,7 @@
     TaskCancel,
     TaskCancelConfirm,
     TaskResult,
+    WorkerDisconnectNotification,
     WorkerHeartbeat,
 )
 from scaler.utility.identifiers import ClientID, ObjectID, TaskID, WorkerID
@@ -175,6 +176,10 @@ async def on_client_shutdown(self, client_id: ClientID):
     async def on_disconnect(self, worker_id: WorkerID, request: DisconnectRequest):
         raise NotImplementedError()
 
+    @abc.abstractmethod
+    async def on_disconnect_notification(self, worker_id: WorkerID, notification: WorkerDisconnectNotification):
+        raise NotImplementedError()
+
     @abc.abstractmethod
     def has_available_worker(self) -> bool:
         raise NotImplementedError()

src/scaler/scheduler/controllers/worker_controller.py

@@ -11,6 +11,7 @@
     StateWorker,
     Task,
     TaskCancel,
+    WorkerDisconnectNotification,
     WorkerHeartbeat,
     WorkerHeartbeatEcho,
 )
@@ -80,6 +81,9 @@ async def on_disconnect(self, worker_id: WorkerID, request: DisconnectRequest):
         await self.__disconnect_worker(request.worker)
         await self._binder.send(worker_id, DisconnectResponse.new_msg(request.worker))
 
+    async def on_disconnect_notification(self, worker_id: WorkerID, notification: WorkerDisconnectNotification):
+        await self.__disconnect_worker(notification.worker)
+
     async def routine(self):
         await self.__clean_workers()
 

src/scaler/scheduler/scheduler.py

@@ -24,6 +24,7 @@
     TaskCancelConfirm,
     TaskLog,
     TaskResult,
+    WorkerDisconnectNotification,
     WorkerHeartbeat,
 )
 from scaler.protocol.python.mixins import Message
@@ -205,6 +206,10 @@ async def on_receive_message(self, source: bytes, message: Message):
             await self._worker_controller.on_disconnect(WorkerID(source), message)
             return
 
+        if isinstance(message, WorkerDisconnectNotification):
+            await self._worker_controller.on_disconnect_notification(WorkerID(source), message)
+            return
+
         # =====================================================================================
         # object manager
         if isinstance(message, ObjectInstruction):

src/scaler/worker/worker.py

@@ -10,27 +10,22 @@
 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,
-    DisconnectRequest,
     DisconnectResponse,
     ObjectInstruction,
     ProcessorInitialized,
     Task,
     TaskCancel,
     TaskLog,
     TaskResult,
+    WorkerDisconnectNotification,
     WorkerHeartbeatEcho,
 )
 from scaler.protocol.python.mixins import Message
@@ -260,8 +255,7 @@ 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:
-            await self._connector_external.send(DisconnectRequest.new_msg(self.identity))
+        await self._connector_external.send(WorkerDisconnectNotification.new_msg(self.identity))
 
         self._connector_external.destroy()
         self._processor_manager.destroy("quit")
@@ -275,14 +269,7 @@ def __run_forever(self):
         self._loop.run_until_complete(self._task)
 
     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)
-        elif backend == NetworkBackend.ymq:
-            self._loop.add_signal_handler(signal.SIGINT, 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)
 
     def __destroy(self):
         self._task.cancel()