Add restart endpoint

arkindex/ponos/api.py

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

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"),
 ]