django-q2/django_q/worker.py

import multiprocessing
from queue import Queue
from queue import Empty
from typing import Optional, Tuple, Union
from django_q.queue_task import QueueTask
from django.utils import timezone
import traceback

from multiprocessing import Process, Value, current_process
from django_q.utils import close_old_django_connections
from django.utils.translation import gettext_lazy as _
from django_q.conf import Conf, logger, setproctitle, error_reporter, resource, psutil
from django_q.signals import pre_execute
from django_q.exceptions import TimeoutException, TimeoutHandler
from django_q.process_manager import ProcessManager



class Worker(ProcessManager):

    def spawn_process(self) -> Process:
        """
        :type target: function or class
        """
        self.status = Value("i", Worker.Status.IDLE.value)
        self.manager_pipe, worker_process_pipe = multiprocessing.Pipe(duplex=True)
        p = WorkerProcess(args=(self.status, worker_process_pipe))
        p.daemon = Conf.DAEMONIZE_WORKERS
        p.start()
        return p

    def start_task(self, task) -> None:
        # send task to worker
        try:
            self.manager_pipe.send(task)
        except BrokenPipeError:
            # recycle process if pipe is broken
            self.status.value = ProcessManager.Status.RECYCLE.value


class Pool:
    """This will manager the individual workers"""
    def __init__(self, workers=Conf.WORKERS):
        self.amount_workers = workers
        self.workers = []
        self.task_queue = Queue()
        self.start_workers()

    def start_workers(self):
        for __ in range(self.amount_workers):
            self.workers.append(Worker())

    def get_worker(self, id) -> Optional[Worker]:
        worker = next((worker for worker in self.workers if worker.id == id), None)
        if worker is None:
            logger.error("Couldn't find worker")
            return
        return worker

    @property
    def is_healthy(self):
        """Checks if all workers are still operating"""
        return all(worker.is_alive for worker in self.workers)

    @property
    def is_idle(self):
        """Checks if all workers are idle"""
        return all(worker.is_idle for worker in self.workers)

    @property
    def is_done(self):
        """Checks if all workers are idle and task queue is empty"""
        return self.is_idle and self.task_queue.empty()

    def reincarnate_stopped_workers(self):
        """Reincarnates workers that are not alive anymore"""
        stopped_workers = [worker for worker in self.workers if not worker.is_alive]
        for worker in stopped_workers:
            worker.reincarnate_process()

    def add_task(self, task):
        self.task_queue.put(task)

    def get_done_workers(self):
        """Worker tasks that have been completed, but need to be saved to cache/db - to be processed by monitor worker"""
        return [worker for worker in self.workers if worker.is_done]

    def mark_workers_idle(self, worker_ids):
        """Mark workers idle when they are ready to be used again"""
        for worker_id in worker_ids:
            # We are going to process the result, mark them idle, so they can be used for a different task
            worker = self.get_worker(id=worker_id)
            if worker is not None:
                worker.mark_idle()

    def delegate_tasks(self):
        available_workers = [worker for worker in self.workers if worker.is_idle]
        for worker in available_workers:
            try:
                task = self.task_queue.get_nowait()
            except Empty:
                # if the queue is empty, then just stop
                break
            worker.start_task(task)


class WorkerProcess(Process):

    @staticmethod
    def run_task(task) -> Tuple[QueueTask, bool]:
        # signal execution
        pre_execute.send(sender="django_q", func=task.func, task=task)
        task.started_at = timezone.now()
        try:
            with TimeoutHandler(timeout=task.timeout):
                func = task.callable_func()
                res = func(*task.args, **task.kwargs)
                result = res
        except (TimeoutException, Exception) as e:
            if isinstance(e, TimeoutException):
                task.status = QueueTask.Status.TIMEOUT
            else:
                task.status = QueueTask.Status.FAILED
            result = f"{e} : {traceback.format_exc()}"

            if error_reporter:
                error_reporter.report()
            if task.sync:
                raise
            return task
        else:
            # succeeded
            task.status = QueueTask.Status.SUCCESS
        finally:
            task.result = result
            task.finished_at = timezone.now()
        return task

    def __init__(self, group=None, name=None, args=(), kwargs={}, daemon=None):
        target = self.processing_tasks
        super().__init__(group=group, target=target, name=name, args=args, kwargs=kwargs, daemon=daemon)

    def mark_ready(self):
        self.process_name = current_process().name
        self.process_id = current_process().pid
        self.task_count = 0
        logger.info(
            _("%(proc_name)s ready for work at %(id)s")
            % {"proc_name": self.process_name, "id": self.process_id}
        )


    def mark_start_task(self, task):
        # Log task creation and set process name
        task_desc = (
            _("%(proc_name)s processing %(task_name)s '%(func_name)s'")
            % {
                "proc_name": self.process_name,
                "func_name": task.func_name,
                "task_name": task.name,
            }
        )
        if task.group is not None:
            task_desc += f" [{task.group}]"
        logger.info(task_desc)

        if setproctitle:
            proc_title = f"qcluster {self.process_name} processing {task.name} '{task.func_name}'"
            if task.group is not None:
                proc_title += f" [{task.group}]"
            setproctitle.setproctitle(proc_title)

    def processing_tasks(self, status: Value, pipe):
        self.mark_ready()
        while True:
            task = pipe.recv()
            if task == "STOP":
                logger.info(f"Worker {self.process_name} stopped processing")
                break
            # got a new task, let's mark it starting
            self.mark_start_task(task)

            # make sure the function actually exists, before we try to run it
            try:
                if not task.is_callable:
                    raise ValueError(f"Function {task.func_name} is not defined")
            except Exception as e:
                result = (f"{e} : {traceback.format_exc()}", False)
                if error_reporter:
                    error_reporter.report()
                if task.sync:
                    raise
                # stop here, move on to the next one
                continue

            close_old_django_connections()

            status.value = ProcessManager.Status.BUSY.value
            task = WorkerProcess.run_task(task)
            # Add task towards total
            self.task_count += 1

            # Set to DONE so main process can pick it up
            status.value = ProcessManager.Status.DONE.value
            if setproctitle:
                setproctitle.setproctitle(f"qcluster {self.process_name} completed with task")

            # Recreate a new process if this task has had the max amount of runs or exceeded resources
            if self.task_count == Conf.RECYCLE or self.rss_check():
                status.value = ProcessManager.Status.RECYCLE
                break

            pipe.send(task)

    def rss_check(self):
        if Conf.MAX_RSS:
            if resource:
                return resource.getrusage(resource.RUSAGE_SELF).ru_maxrss >= Conf.MAX_RSS
            elif psutil:
                return psutil.Process().memory_info().rss >= Conf.MAX_RSS * 1024
        return False