wip

Author: cunla

scheduler/worker/commands/kill_worker.py

@@ -1,11 +1,11 @@
 import errno
 import os
 import signal
-from typing import Optional
+from typing import Any
 
-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
 
 
@@ -14,27 +14,26 @@ class KillWorkerCommand(WorkerCommand):
 
     command_name = "kill-worker"
 
-    def __init__(self, *args, **kwargs) -> None:
+    def __init__(self, *args: Any, **kwargs: Any) -> None:
         super().__init__(*args, **kwargs)
-        self.worker_pid: Optional[int] = None
 
     def process_command(self, connection: ConnectionType) -> None:
         from scheduler.worker import Worker
-
-        logger.info("Received kill-worker command.")
+        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)
-        self.worker_pid = worker_model.pid
-        if self.worker_pid is None:
+        if worker_model is None or worker_model.pid is None:
             raise ValueError("Worker PID is not set")
-        logger.info(f"Killing worker main process {self.worker_pid}...")
+        logger.info(f"Killing worker main process {worker_model.pid}...")
         try:
             Worker.from_model(worker_model).request_stop(signal.SIGTERM, None)
-            os.killpg(os.getpgid(self.worker_pid), signal.SIGTERM)
-            logger.info(f"Killed worker main process pid {self.worker_pid}")
+            os.killpg(os.getpgid(worker_model.pid), signal.SIGTERM)
+            logger.info(f"Killed worker main process pid {worker_model.pid}")
         except OSError as e:
             if e.errno == errno.ESRCH:
                 logger.debug(
-                    f"Worker main process for {self.worker_name}:{self.worker_pid} already dead"
+                    f"Worker main process for {self.worker_name}:{worker_model.pid} already dead"
                 )  # "No such process" is fine with us
             else:
                 raise

scheduler/worker/commands/stop_job.py

@@ -19,7 +19,7 @@ def __init__(self, *args, job_name: str, worker_name: Optional[str], **kwargs) -
         if self.job_name is None:
             raise WorkerCommandError("job_name for kill-job command is required")
 
-    def command_payload(self) -> Dict[str, Any]:
+    def command_payload(self,  **kwargs) -> 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

@@ -14,6 +14,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 worker_model.is_suspended:
             logger.warning(f"Worker {self.worker_name} already suspended")
             return

scheduler/worker/scheduler.py

@@ -23,7 +23,7 @@ class SchedulerStatus(str, Enum):
     STOPPED = "stopped"
 
 
-def _reschedule_tasks():
+def _reschedule_tasks() -> None:
     enabled_jobs = list(Task.objects.filter(enabled=True))
     for item in enabled_jobs:
         logger.debug(f"Rescheduling {str(item)}")
@@ -32,11 +32,11 @@ def _reschedule_tasks():
 
 class WorkerScheduler:
     def __init__(
-        self,
-        queues: Sequence[Queue],
-        connection: ConnectionType,
-        worker_name: str,
-        interval: Optional[int] = None,
+            self,
+            queues: Sequence[Queue],
+            connection: ConnectionType,
+            worker_name: str,
+            interval: Optional[int] = None,
     ) -> None:
         interval = interval or SCHEDULER_CONFIG.SCHEDULER_INTERVAL
         self._queues = queues

scheduler/worker/worker.py

@@ -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):
+    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