from __future__ import annotations import functools from collections import namedtuple from collections.abc import Iterable from typing import cast import django_rq from django.conf import settings from django.http import HttpResponseBadRequest, HttpResponseNotFound from django.utils.decorators import method_decorator from django.views.decorators.cache import never_cache from django_rq.queues import DjangoRQ from drf_spectacular.utils import ( OpenApiParameter, OpenApiResponse, OpenApiTypes, extend_schema, extend_schema_view, ) from redis.exceptions import ConnectionError as RedisConnectionError from rest_framework import status, viewsets from rest_framework.decorators import action from rest_framework.response import Response from rq.command import send_stop_job_command from rq.exceptions import InvalidJobOperation from rq.job import Job as RQJob from rq.job import JobStatus as RQJobStatus from rq.worker import Worker as RQWorker from cvat.apps.engine.filters import ( NonModelJsonLogicFilter, NonModelOrderingFilter, NonModelSimpleFilter, ) from cvat.apps.engine.log import ServerLogManager from cvat.apps.engine.models import RequestAction, RequestSubresource from cvat.apps.engine.rq import is_rq_job_owner from cvat.apps.engine.types import ExtendedRequest from cvat.apps.redis_handler.apps import SELECTOR_TO_QUEUE from cvat.apps.redis_handler.permissions import RequestPermission from cvat.apps.redis_handler.rq import CustomRQJob, RequestId from cvat.apps.redis_handler.serializers import RequestSerializer, RequestStatus slogger = ServerLogManager(__name__) CVAT_CAN_STOP_STARTED_JOBS_KEY = "cvat_can_stop_started_jobs" @extend_schema(tags=["requests"]) @extend_schema_view( list=extend_schema( summary="List requests", responses={ "200": RequestSerializer(many=True), }, ), retrieve=extend_schema( summary="Get request details", parameters=[ OpenApiParameter( "id", OpenApiTypes.STR, OpenApiParameter.PATH, ) ], responses={ "200": RequestSerializer, }, ), ) class RequestViewSet(viewsets.GenericViewSet): serializer_class = RequestSerializer iam_supports_organization_params = False iam_permission_class = RequestPermission filter_backends = [ NonModelSimpleFilter, NonModelJsonLogicFilter, NonModelOrderingFilter, ] ordering_fields = ["created_date", "status", "action"] ordering = "-created_date" simple_filters = [ # RQ job fields "status", # derivatives fields (from meta) "project_id", "task_id", "job_id", # derivatives fields (from parsed rq_id) "action", "target", "subresource", "format", # request context "org", "org_id", ] filter_fields = [*simple_filters] lookup_fields = { "created_date": "created_at", "action": "parsed_id.action", "target": "parsed_id.target", "subresource": "parsed_id.subresource", "format": "parsed_id.format", "status": "get_status", "project_id": "meta.project_id", "task_id": "meta.task_id", "job_id": "meta.job_id", "org": "meta.org_slug", "org_id": "meta.org_id", } SchemaField = namedtuple("SchemaField", ["type", "choices"], defaults=(None,)) simple_filters_schema = { "status": SchemaField("string", RequestStatus.choices), "project_id": SchemaField("integer"), "task_id": SchemaField("integer"), "job_id": SchemaField("integer"), "action": SchemaField("string"), "target": SchemaField("string"), "subresource": SchemaField("string"), "format": SchemaField("string"), "org": SchemaField("string"), "org_id": SchemaField("integer"), } class _SchemaQuerySet: # drf-spectacular expects schema views to expose a queryset-like object # with a model attribute. Requests are backed by RQ jobs, not Django models. # See https://drf-spectacular.readthedocs.io/en/latest/faq.html#my-get-queryset-depends-on-some-attributes-not-available-at-schema-generation-time model = None def get_queryset(self): if getattr(self, "swagger_fake_view", False): return self._SchemaQuerySet() return None @property def queues(self) -> Iterable[DjangoRQ]: return (django_rq.get_queue(queue_name) for queue_name in set(SELECTOR_TO_QUEUE.values())) def _get_rq_jobs_from_queue(self, queue: DjangoRQ, user_id: int) -> list[RQJob]: job_ids = set( queue.get_job_ids() + queue.started_job_registry.get_job_ids() + queue.finished_job_registry.get_job_ids() + queue.failed_job_registry.get_job_ids() + queue.deferred_job_registry.get_job_ids() ) jobs = [] for job in queue.job_class.fetch_many(job_ids, queue.connection): if job and is_rq_job_owner(job, user_id): if job.get_status(refresh=False) in {RQJobStatus.CANCELED, RQJobStatus.STOPPED}: continue job = cast(CustomRQJob, job) try: parsed_request_id = RequestId.parse_and_validate_queue( job.id, expected_queue=queue.name ) except Exception: # nosec B112 continue job.parsed_id = parsed_request_id jobs.append(job) return jobs def _get_rq_jobs(self, user_id: int) -> list[RQJob]: """ Get all RQ jobs for a specific user and return them as a list of RQJob objects. Parameters: user_id (int): The ID of the user for whom to retrieve jobs. Returns: list[RQJob]: A list of RQJob objects representing all jobs for the specified user. """ all_jobs = [] for queue in self.queues: jobs = self._get_rq_jobs_from_queue(queue, user_id) all_jobs.extend(jobs) return all_jobs def _get_rq_job_by_id(self, rq_id: str) -> CustomRQJob | None: """ Get a RQJob by its ID from the queues. Args: rq_id (str): The ID of the RQJob to retrieve. Returns: RQJob | None: The retrieved RQJob, or None if not found. """ try: parsed_request_id, queue_name = RequestId.parse(rq_id, try_legacy_format=True) rq_id = parsed_request_id.render() except Exception: return None queue: DjangoRQ = django_rq.get_queue(queue_name) job: CustomRQJob | None = queue.fetch_job(rq_id) if job: job.parsed_id = parsed_request_id return job @staticmethod def _is_export_request(rq_job: CustomRQJob) -> bool: return rq_job.parsed_id.action == RequestAction.EXPORT and rq_job.parsed_id.subresource in { RequestSubresource.ANNOTATIONS, RequestSubresource.DATASET, RequestSubresource.BACKUP, } @classmethod def _is_started_process_cancellable(cls, rq_job: CustomRQJob) -> bool: if not cls._is_export_request(rq_job): return False worker_key = f"{RQWorker.redis_worker_namespace_prefix}{rq_job.worker_name}" # The marker is CVAT-specific and can be absent on already running or older production # workers. Only SimpleWorker writes "0", so default to "can stop" for backward # compatibility. can_stop_started_jobs = ( rq_job.connection.hget(worker_key, CVAT_CAN_STOP_STARTED_JOBS_KEY) or "1" ) if isinstance(can_stop_started_jobs, bytes): can_stop_started_jobs = can_stop_started_jobs.decode() return can_stop_started_jobs != "0" def _handle_redis_exceptions(func): @functools.wraps(func) def wrapper(*args, **kwargs): try: return func(*args, **kwargs) except RedisConnectionError as ex: msg = "Redis service is not available" slogger.glob.exception(f"{msg}: {str(ex)}") return Response(msg, status=status.HTTP_503_SERVICE_UNAVAILABLE) return wrapper @method_decorator(never_cache) @_handle_redis_exceptions def retrieve(self, request: ExtendedRequest, pk: str): job = self._get_rq_job_by_id(pk) if not job: return HttpResponseNotFound("There is no request with specified id") self.check_object_permissions(request, job) if job.get_status(refresh=False) in {RQJobStatus.CANCELED, RQJobStatus.STOPPED}: job.delete() return HttpResponseNotFound("There is no request with specified id") serializer = self.get_serializer(job, context={"request": request}) return Response(data=serializer.data, status=status.HTTP_200_OK) @method_decorator(never_cache) @_handle_redis_exceptions def list(self, request: ExtendedRequest): user_id = request.user.id user_jobs = self._get_rq_jobs(user_id) filtered_jobs = self.filter_queryset(user_jobs) page = self.paginate_queryset(filtered_jobs) if page is not None: serializer = self.get_serializer(page, many=True, context={"request": request}) return self.get_paginated_response(serializer.data) serializer = self.get_serializer(filtered_jobs, many=True, context={"request": request}) return Response(data=serializer.data, status=status.HTTP_200_OK) @extend_schema( summary="Cancel request", parameters=[ OpenApiParameter( "id", OpenApiTypes.STR, OpenApiParameter.PATH, ) ], request=None, responses={ "200": OpenApiResponse(description="The request has been cancelled"), }, ) @method_decorator(never_cache) @action(detail=True, methods=["POST"], url_path="cancel") @_handle_redis_exceptions def cancel(self, request: ExtendedRequest, pk: str): rq_job = self._get_rq_job_by_id(pk) if not rq_job: return HttpResponseNotFound("There is no request with specified id") self.check_object_permissions(request, rq_job) rq_job_status = rq_job.get_status(refresh=False) # Terminal canceled jobs are not useful for users in the requests API. if rq_job_status in {RQJobStatus.CANCELED, RQJobStatus.STOPPED}: return HttpResponseNotFound("There is no request with specified id") # Jobs that have not started yet are safe to cancel directly. RQ will also enqueue # dependents when ONE_RUNNING_JOB_IN_QUEUE_PER_USER is enabled. if rq_job_status in {RQJobStatus.QUEUED, RQJobStatus.DEFERRED}: # FUTURE-TODO: race condition is possible here rq_job.cancel(enqueue_dependents=settings.ONE_RUNNING_JOB_IN_QUEUE_PER_USER) rq_job.delete() return Response(status=status.HTTP_200_OK) # Finished, failed, scheduled, and other non-started states cannot be interrupted. if rq_job_status != RQJobStatus.STARTED: return HttpResponseBadRequest( f"Requests with status {rq_job_status!r} cannot be cancelled" ) # Started cancellation is intentionally limited to non-mutating export jobs running in # production-style workers. if not self._is_started_process_cancellable(rq_job): return HttpResponseBadRequest( "Cancellation of started requests is supported only for export requests " "executed by production workers" ) # RQ will stop the forked work horse and enqueue dependent jobs from its stopped-job path. queue = django_rq.get_queue(rq_job.origin) try: send_stop_job_command(queue.connection, rq_job.id, serializer=queue.serializer) except InvalidJobOperation as ex: slogger.glob.warning("Failed to stop RQ job %s: %s", rq_job.id, ex) return HttpResponseBadRequest("The request cannot be cancelled") return Response(status=status.HTTP_200_OK)