diff --git a/arkindex/ponos/api.py b/arkindex/ponos/api.py index a7f0864565e5ad980716882e12d64213969c84c5..bb15286ce36a0b241ed26de9d99dfcee5e4d3dfb 100644 --- a/arkindex/ponos/api.py +++ b/arkindex/ponos/api.py @@ -1,12 +1,17 @@ +import uuid from textwrap import dedent +from django.db import transaction from django.shortcuts import get_object_or_404, redirect from drf_spectacular.utils import extend_schema, extend_schema_view +from rest_framework import serializers, status from rest_framework.authentication import SessionAuthentication, TokenAuthentication -from rest_framework.generics import ListCreateAPIView, RetrieveUpdateAPIView, UpdateAPIView +from rest_framework.exceptions import NotFound, ValidationError +from rest_framework.generics import CreateAPIView, ListCreateAPIView, RetrieveUpdateAPIView, UpdateAPIView +from rest_framework.response import Response from rest_framework.views import APIView -from arkindex.ponos.models import Artifact, Task +from arkindex.ponos.models import FINAL_STATES, Artifact, State, Task, task_token_default from arkindex.ponos.permissions import ( IsAgentOrArtifactGuest, IsAgentOrTaskGuest, @@ -15,6 +20,9 @@ from arkindex.ponos.permissions import ( IsTaskAdmin, ) from arkindex.ponos.serializers import ArtifactSerializer, TaskSerializer, TaskTinySerializer +from arkindex.project.mixins import ProcessACLMixin +from arkindex.project.permissions import IsVerified +from arkindex.users.models import Role @extend_schema(tags=["ponos"]) @@ -168,3 +176,75 @@ class TaskUpdate(UpdateAPIView): permission_classes = (IsTaskAdmin, ) queryset = Task.objects.select_related("process__corpus") serializer_class = TaskTinySerializer + + +@extend_schema_view( + post=extend_schema( + operation_id="RestartTask", + tags=["ponos"], + description=dedent( + """ + Restart a task by creating a fresh copy and moves dependent tasks to the new one. + + Scenario restarting `my_worker` task: + ``` + init_elements → my_worker → other worker + ``` + ``` + init_elements → my_worker + ↘ + my_worker_2 → other worker + ``` + + Requires an **admin** access to task's process. + The task must be in a final state to be restarted. + """ + ), + responses={201: TaskSerializer}, + ), +) +class TaskRestart(ProcessACLMixin, CreateAPIView): + permission_classes = (IsVerified,) + serializer_class = serializers.Serializer + + def get_task(self): + task = get_object_or_404( + Task.objects.prefetch_related("parents").select_related("process"), + pk=self.kwargs["pk"], + ) + access_level = self.process_access_level(task.process) + if access_level is None: + raise NotFound + if access_level < Role.Admin.value: + raise ValidationError( + detail="You do not have an admin access to the process of this task." + ) + if task.state not in FINAL_STATES: + raise ValidationError( + detail="Task's state must be in a final state to be restarted." + ) + return task + + def increment(self, name): + basename, *suffix = name.rsplit("_restart", 1) + suffix = int(suffix[0]) + 1 if suffix and suffix[0].isdigit() else 1 + return f"{basename}_restart{suffix}" + + @transaction.atomic + def create(self, request, pk=None, **kwargs): + copy = self.get_task() + parents = list(copy.parents.all()) + + copy.id = uuid.uuid4() + copy.state = State.Pending + copy.token = task_token_default() + copy.slug = self.increment(copy.slug) + copy.save() + + # Create links to retried task parents + copy.parents.add(*parents) + + # Move all tasks depending on the retried task to the copy + Task.children.through.objects.filter(to_task_id=pk).update(to_task_id=copy.id) + + return Response(TaskSerializer(copy).data, status=status.HTTP_201_CREATED) diff --git a/arkindex/project/api_v1.py b/arkindex/project/api_v1.py index 81d9724f0485b94a88eeda85aa198fa83e69e79d..e461fc445404bc9b9bcff1430c05b4a39f6a306e 100644 --- a/arkindex/project/api_v1.py +++ b/arkindex/project/api_v1.py @@ -59,7 +59,7 @@ from arkindex.documents.api.ml import ( ) from arkindex.documents.api.search import CorpusSearch, SearchIndexBuild from arkindex.images.api import IIIFInformationCreate, IIIFURLCreate, ImageCreate, ImageElements, ImageRetrieve -from arkindex.ponos.api import TaskArtifactDownload, TaskArtifacts, TaskDetailsFromAgent, TaskUpdate +from arkindex.ponos.api import TaskArtifactDownload, TaskArtifacts, TaskDetailsFromAgent, TaskRestart, TaskUpdate from arkindex.process.api import ( ApplyProcessTemplate, BucketList, @@ -327,4 +327,5 @@ api = [ TaskArtifactDownload.as_view(), name="task-artifact-download", ), + path("task/<uuid:pk>/restart/", TaskRestart.as_view(), name="task-restart"), ]