/cvat/apps/redis_handler/background.py

# Copyright (C) CVAT.ai Corporation
#
# SPDX-License-Identifier: MIT

from abc import ABCMeta, abstractmethod
from collections.abc import Callable
from typing import Any, ClassVar

import django_rq
from django.conf import settings
from django.db.models import Model
from django_rq.queues import DjangoRQ, DjangoScheduler
from rest_framework import status
from rest_framework.response import Response
from rq import Callback
from rq.job import Job as RQJob
from rq.job import JobStatus as RQJobStatus

from cvat.apps.engine.log import ServerLogManager
from cvat.apps.engine.models import RequestTarget
from cvat.apps.engine.rq import BaseRQMeta, define_dependent_job
from cvat.apps.engine.types import ExtendedRequest
from cvat.apps.engine.utils import get_rq_lock_by_user, get_rq_lock_for_job
from cvat.apps.redis_handler.serializers import RqIdSerializer

slogger = ServerLogManager(__name__)

REQUEST_TIMEOUT = 60
# it's better to return LockNotAvailableError instead of response with 504 status
LOCK_TTL = REQUEST_TIMEOUT - 5
LOCK_ACQUIRE_TIMEOUT = LOCK_TTL - 5


class AbstractRequestManager(metaclass=ABCMeta):
    SUPPORTED_TARGETS: ClassVar[set[RequestTarget] | None] = None
    QUEUE_NAME: ClassVar[str]
    REQUEST_ID_KEY = "rq_id"

    callback: Callable
    callback_args: tuple | None
    callback_kwargs: dict[str, Any] | None

    job_on_success_callback: Callback | None
    job_on_failure_callback: Callback | None

    def __init__(
        self,
        *,
        request: ExtendedRequest,
        db_instance: Model | None = None,
    ) -> None:
        self.request = request
        self.user_id = request.user.id
        self.db_instance = db_instance

        if db_instance:
            assert self.SUPPORTED_TARGETS, "Should be defined"
            self.target = RequestTarget(db_instance.__class__.__name__.lower())
            assert self.target in self.SUPPORTED_TARGETS, f"Unsupported target: {self.target}"

    @classmethod
    def get_queue(cls) -> DjangoRQ:
        return django_rq.get_queue(cls.QUEUE_NAME)

    @property
    def job_result_ttl(self) -> int | None:
        """
        Time to live for successful job result in seconds,
        if not set, the default result TTL will be used
        """
        return None

    @property
    def job_failed_ttl(self) -> int | None:
        """
        Time to live for failures in seconds,
        if not set, the default failure TTL will be used
        """
        return None

    @abstractmethod
    def build_request_id(self): ...

    def validate_request_id(self, request_id: str, /) -> None: ...

    def get_job_by_id(self, id_: str, /) -> RQJob | None:
        try:
            self.validate_request_id(id_)
        except Exception:
            return None

        queue = self.get_queue()
        return queue.fetch_job(id_)

    def init_request_args(self):
        """
        Hook to initialize operation args based on the request
        """

    @abstractmethod
    def init_callback_with_params(self) -> None:
        """
        Method should initialize callback function with its args/kwargs:

        self.callback = ...
        (optional) self.callback_args = ...
        (optional) self.callback_kwargs = ...
        """

    def _set_default_callback_params(self):
        self.callback_args = None
        self.callback_kwargs = None
        self.job_on_success_callback = None
        self.job_on_failure_callback = None

    def init_job_callbacks(self) -> None:
        """Hook to initialize RQ lifecycle callbacks for the job"""

    def validate_request(self) -> Response | None:
        """Hook to run some validations before processing a request"""

        # prevent architecture bugs
        assert (
            "POST" == self.request.method
        ), "Only POST requests can be used to initiate a background process"

    def handle_existing_job(self, job: RQJob | None, queue: DjangoRQ) -> Response | None:
        if not job:
            return None

        job_status = job.get_status(refresh=False)

        if job_status in {
            # FUTURE-TODO: cancelling and re-enqueuing a started job should probably be allowed
            RQJobStatus.STARTED,
            RQJobStatus.QUEUED,
            RQJobStatus.DEFERRED,
        }:
            return Response(
                RqIdSerializer({"rq_id": job.id}).data,
                status=status.HTTP_409_CONFLICT,
            )

        # RQ jobs can be scheduled only by CVAT internal logic, in that case job has no dependencies
        if job_status == RQJobStatus.SCHEDULED:
            scheduler: DjangoScheduler = django_rq.get_scheduler(queue.name, queue=queue)
            # remove the job id from the set with scheduled keys
            scheduler.cancel(job)
            job.cancel(enqueue_dependents=settings.ONE_RUNNING_JOB_IN_QUEUE_PER_USER)

        job.delete()
        return None

    def build_meta(self, *, request_id: str) -> dict[str, Any]:
        return BaseRQMeta.build(request=self.request, db_obj=self.db_instance)

    def setup_new_job(self, queue: DjangoRQ, request_id: str, /, **kwargs):
        with get_rq_lock_by_user(queue, self.user_id):
            queue.enqueue_call(
                func=self.callback,
                args=self.callback_args,
                kwargs=self.callback_kwargs,
                job_id=request_id,
                meta=self.build_meta(request_id=request_id),
                depends_on=define_dependent_job(queue, self.user_id, rq_id=request_id),
                result_ttl=self.job_result_ttl,
                failure_ttl=self.job_failed_ttl,
                on_success=self.job_on_success_callback,
                on_failure=self.job_on_failure_callback,
                **kwargs,
            )

    def finalize_request(self) -> None:
        """Hook to run some actions (e.g. collect events) after processing a request"""

    def get_response(self, request_id: str) -> Response:
        serializer = RqIdSerializer({"rq_id": request_id})
        return Response(serializer.data, status=status.HTTP_202_ACCEPTED)

    def enqueue_job(self) -> Response:
        self.init_request_args()
        self.validate_request()
        self._set_default_callback_params()
        self.init_callback_with_params()
        self.init_job_callbacks()

        queue: DjangoRQ = django_rq.get_queue(self.QUEUE_NAME)
        request_id = self.build_request_id()

        # ensure that there is no race condition when processing parallel requests
        with get_rq_lock_for_job(queue, request_id):
            job = queue.fetch_job(request_id)

            if response := self.handle_existing_job(job, queue):
                return response

            self.setup_new_job(queue, request_id)

        self.finalize_request()
        return self.get_response(request_id)