wip

Author: cunla

scheduler/worker/commands/kill_worker.py

@@ -19,8 +19,6 @@ def __init__(self, *args: Any, **kwargs: Any) -> None:
 
     def process_command(self, connection: ConnectionType) -> None:
         from scheduler.worker import Worker
-        if self.worker_name is None:
-            raise ValueError("Worker name must be provided")
         logger.info(f"Received kill-worker command for {self.worker_name}")
         worker_model = WorkerModel.get(self.worker_name, connection)
         if worker_model is None or worker_model.pid is None:

scheduler/worker/commands/stop_job.py

@@ -1,10 +1,10 @@
 import os
 import signal
-from typing import Dict, Any, Optional
+from typing import Dict, Any
 
-from scheduler.types import ConnectionType
 from scheduler.redis_models import WorkerModel, JobModel
 from scheduler.settings import logger
+from scheduler.types import ConnectionType
 from scheduler.worker.commands.worker_commands import WorkerCommand, WorkerCommandError
 
 
@@ -13,13 +13,13 @@ class StopJobCommand(WorkerCommand):
 
     command_name = "stop-job"
 
-    def __init__(self, *args, job_name: str, worker_name: Optional[str], **kwargs) -> None:
+    def __init__(self, *args: Any, job_name: str, worker_name: str, **kwargs: Any) -> None:
         super().__init__(*args, worker_name=worker_name, **kwargs)
         self.job_name = job_name
         if self.job_name is None:
             raise WorkerCommandError("job_name for kill-job command is required")
 
-    def command_payload(self,  **kwargs) -> Dict[str, Any]:
+    def command_payload(self, **kwargs: Any) -> Dict[str, Any]:
         return super().command_payload(job_name=self.job_name)
 
     def process_command(self, connection: ConnectionType) -> None:

scheduler/worker/commands/suspend_worker.py

@@ -1,6 +1,6 @@
-from scheduler.types import ConnectionType
 from scheduler.redis_models import WorkerModel
 from scheduler.settings import logger
+from scheduler.types import ConnectionType
 from scheduler.worker.commands.worker_commands import WorkerCommand
 
 
@@ -32,6 +32,7 @@ def process_command(self, connection: ConnectionType) -> None:
         worker_model = WorkerModel.get(self.worker_name, connection)
         if worker_model is None:
             logger.warning(f"Worker {self.worker_name} not found")
+            return
         if not worker_model.is_suspended:
             logger.warning(f"Worker {self.worker_name} not suspended and therefore can't be resumed")
             return

scheduler/worker/commands/worker_commands.py

@@ -1,12 +1,13 @@
 import json
 from abc import ABC
 from datetime import datetime, timezone
-from typing import Type, Dict, Any, Optional
+from typing import Type, Dict, Any
 
 from scheduler.settings import logger
 from scheduler.types import ConnectionType, Self
 
 _PUBSUB_CHANNEL_TEMPLATE: str = ":workers:pubsub:{}"
+_WORKER_COMMANDS_REGISTRY: Dict[str, Type["WorkerCommand"]] = dict()
 
 
 class WorkerCommandError(Exception):
@@ -16,13 +17,12 @@ class WorkerCommandError(Exception):
 class WorkerCommand(ABC):
     """Abstract class for commands to be sent to a worker and processed by worker"""
 
-    _registry: Dict[str, Type[Self]] = dict()
     command_name: str = ""
 
-    def __init__(self, *args, worker_name: Optional[str], **kwargs) -> None:
-        self.worker_name: Optional[str] = worker_name
+    def __init__(self, *args: Any, worker_name: str, **kwargs: Any) -> None:
+        self.worker_name: str = worker_name
 
-    def command_payload(self, **kwargs) -> Dict[str, Any]:
+    def command_payload(self, **kwargs: Any) -> Dict[str, Any]:
         commands_channel = WorkerCommandsChannelListener._commands_channel(self.worker_name)
         payload = {
             "command": self.command_name,
@@ -41,17 +41,19 @@ def process_command(self, connection: ConnectionType) -> None:
         raise NotImplementedError
 
     @classmethod
-    def __init_subclass__(cls, *args, **kwargs):
+    def __init_subclass__(cls, *args: Any, **kwargs: Any) -> None:
         if cls is WorkerCommand:
             return
         if not cls.command_name:
-            raise NotImplementedError(f"{cls.__name__} must have a name attribute")
-        WorkerCommand._registry[cls.command_name] = cls
+            raise NotImplementedError(f"{cls.__name__} must have a command_name attribute")
+        _WORKER_COMMANDS_REGISTRY[cls.command_name] = cls
 
     @classmethod
     def from_payload(cls, payload: Dict[str, Any]) -> Type[Self]:
         command_name = payload.get("command")
-        command_class = WorkerCommand._registry.get(command_name)
+        if command_name is None:
+            raise WorkerCommandError("Payload must contain 'command' key")
+        command_class = _WORKER_COMMANDS_REGISTRY.get(command_name)
         if command_class is None:
             raise WorkerCommandError(f"Invalid command: {command_name}")
         return command_class(**payload)

scheduler/worker/scheduler.py

@@ -41,15 +41,15 @@ def __init__(
         interval = interval or SCHEDULER_CONFIG.SCHEDULER_INTERVAL
         self._queues = queues
         self._scheduled_job_registries: List[ScheduledJobRegistry] = []
-        self.lock_acquisition_time = None
+        self.lock_acquisition_time: Optional[datetime] = None
         self._pool_class = connection.connection_pool.connection_class
         self._pool_kwargs = connection.connection_pool.connection_kwargs.copy()
         self._locks: Dict[str, SchedulerLock] = dict()
         self.connection = connection
         self.interval = interval
         self._stop_requested = False
         self.status = SchedulerStatus.STOPPED
-        self._thread = None
+        self._thread: Optional[Thread] = None
         self._pid: Optional[int] = None
         self.worker_name = worker_name
 
@@ -163,7 +163,7 @@ def enqueue_scheduled_jobs(self) -> None:
         self.status = SchedulerStatus.STARTED
 
 
-def run_scheduler(scheduler: WorkerScheduler):
+def run_scheduler(scheduler: WorkerScheduler) -> None:
     try:
         scheduler.work()
     except:  # noqa

scheduler/worker/worker.py

@@ -14,7 +14,7 @@
 from enum import Enum
 from random import shuffle
 from types import FrameType
-from typing import List, Optional, Tuple, Any, Iterable, Collection
+from typing import List, Optional, Tuple, Any, Iterable, Collection, Union
 
 import scheduler
 from scheduler.helpers.queues import get_queue
@@ -63,7 +63,7 @@ class QueueConnectionDiscrepancyError(Exception):
 )
 
 
-def signal_name(signum) -> str:
+def signal_name(signum: int) -> str:
     try:
         return signal.Signals(signum).name
     except KeyError:
@@ -99,19 +99,19 @@ def from_model(cls, model: WorkerModel) -> Self:
         return res
 
     def __init__(
-        self,
-        queues,
-        name: str,
-        connection: ConnectionType,
-        maintenance_interval: int = SCHEDULER_CONFIG.DEFAULT_MAINTENANCE_TASK_INTERVAL,
-        job_monitoring_interval=SCHEDULER_CONFIG.DEFAULT_JOB_MONITORING_INTERVAL,
-        dequeue_strategy: DequeueStrategy = DequeueStrategy.DEFAULT,
-        disable_default_exception_handler: bool = False,
-        fork_job_execution: bool = True,
-        with_scheduler: bool = True,
-        burst: bool = False,
-        model: Optional[WorkerModel] = None,
-    ):  # noqa
+            self,
+            queues: Iterable[Union[str, Queue]],
+            name: str,
+            connection: ConnectionType,
+            maintenance_interval: int = SCHEDULER_CONFIG.DEFAULT_MAINTENANCE_TASK_INTERVAL,
+            job_monitoring_interval: int = SCHEDULER_CONFIG.DEFAULT_JOB_MONITORING_INTERVAL,
+            dequeue_strategy: DequeueStrategy = DequeueStrategy.DEFAULT,
+            disable_default_exception_handler: bool = False,
+            fork_job_execution: bool = True,
+            with_scheduler: bool = True,
+            burst: bool = False,
+            model: Optional[WorkerModel] = None,
+    ) -> None:
         self.fork_job_execution = fork_job_execution
         self.job_monitoring_interval: int = job_monitoring_interval
         self.maintenance_interval = maintenance_interval
@@ -232,7 +232,7 @@ def work(self, max_jobs: Optional[int] = None, max_idle_time: Optional[int] = No
 
                 timeout = None if self.burst else (SCHEDULER_CONFIG.DEFAULT_WORKER_TTL - 15)
                 job, queue = self.dequeue_job_and_maintain_ttl(timeout, max_idle_time)
-                if job is None:
+                if job is None or queue is None:
                     if self.burst:
                         logger.info(f"[Worker {self.name}/{self._pid}]: done, quitting")
                         break
@@ -267,7 +267,7 @@ def work(self, max_jobs: Optional[int] = None, max_idle_time: Optional[int] = No
             self.teardown()
         return False
 
-    def handle_job_failure(self, job: JobModel, queue: Queue, exc_string="") -> None:
+    def handle_job_failure(self, job: JobModel, queue: Queue, exc_string: str = "") -> None:
         """
         Handles the failure or an executing job by:
             1. Setting the job status to failed
@@ -312,7 +312,7 @@ def handle_job_failure(self, job: JobModel, queue: Queue, exc_string="") -> None
                 # Ensure that custom exception handlers are called even if the Broker is down
                 pass
 
-    def bootstrap(self)-> None:
+    def bootstrap(self) -> None:
         """Bootstraps the worker.
         Runs the basic tasks that should run when the worker actually starts working.
         Used so that new workers can focus on the work loop implementation rather
@@ -327,7 +327,8 @@ def bootstrap(self)-> None:
             self._model.has_scheduler = True
             self._model.save(connection=self.connection)
         if self.with_scheduler and self.burst:
-            self.scheduler.request_stop_and_wait()
+            if self.scheduler is not None:
+                self.scheduler.request_stop_and_wait()
             self._model.has_scheduler = False
             self._model.save(connection=self.connection)
         qnames = [queue.name for queue in self.queues]
@@ -375,8 +376,8 @@ def run_maintenance_tasks(self) -> None:
             self._model.save(connection=self.connection)
 
     def dequeue_job_and_maintain_ttl(
-        self, timeout: Optional[int], max_idle_time: Optional[int] = None
-    ) -> Tuple[JobModel, Queue]:
+            self, timeout: Optional[int], max_idle_time: Optional[int] = None
+    ) -> Tuple[Optional[JobModel], Optional[Queue]]:
         """Dequeues a job while maintaining the TTL.
         :param timeout: The timeout for the dequeue operation.
         :param max_idle_time: The maximum idle time for the worker.
@@ -550,7 +551,7 @@ def reorder_queues(self, reference_queue: Queue) -> None:
             return
         if self._dequeue_strategy == DequeueStrategy.ROUND_ROBIN:
             pos = self._ordered_queues.index(reference_queue)
-            self._ordered_queues = self._ordered_queues[pos + 1 :] + self._ordered_queues[: pos + 1]
+            self._ordered_queues = self._ordered_queues[pos + 1:] + self._ordered_queues[: pos + 1]
             return
         if self._dequeue_strategy == DequeueStrategy.RANDOM:
             shuffle(self._ordered_queues)
@@ -634,7 +635,7 @@ def monitor_job_execution_process(self, job: JobModel, queue: Queue) -> None:
         while True:
             try:
                 with SCHEDULER_CONFIG.DEATH_PENALTY_CLASS(
-                    self.job_monitoring_interval, JobExecutionMonitorTimeoutException
+                        self.job_monitoring_interval, JobExecutionMonitorTimeoutException
                 ):
                     retpid, ret_val = self.wait_for_job_execution_process()
                 break
@@ -877,7 +878,7 @@ class RoundRobinWorker(Worker):
 
     def reorder_queues(self, reference_queue: Queue) -> None:
         pos = self._ordered_queues.index(reference_queue)
-        self._ordered_queues = self._ordered_queues[pos + 1 :] + self._ordered_queues[: pos + 1]
+        self._ordered_queues = self._ordered_queues[pos + 1:] + self._ordered_queues[: pos + 1]
 
 
 class RandomWorker(Worker):