Annotation Task Lifecycle Management

The Problem with Uncontrolled States

Right now, your AnnotationTask model has a status field. A client can send a PATCH request with {"status": "completed"} and the API will obediently update the record, regardless of its current state. In a real system, processes have rules. A package isn't "delivered" before it's "shipped," and an annotation task isn't "completed" before it's been worked on.

A Finite State Machine (FSM) is a model designed to solve exactly this problem. It's based on a simple idea: an object can only exist in one of a finite number of states at any given time, and it can only move between states through predefined, explicit "transitions."

The main benefits of using an FSM are:

• Predictability: The system's behavior becomes explicit and easy to understand.
• Robustness: It's impossible to put an object into an invalid state.
• Maintainability: All the logic for state changes is centralized in one place, not scattered across different views and functions.

Implementing a State Machine with django-fsm

While you can build an FSM from scratch, the Django ecosystem has an excellent library called django-fsm that makes this process declarative and elegant. It integrates directly into your models.

First, let's install it. Stop your development server and run:

pip install django-fsm

And don't forget to add django-fsm to your requirements.txt file.

Now, we'll modify our AnnotationTask model to use it. The process involves two main changes: updating the status field and defining methods for each valid transition.

Step 1: Update the Model

Let's apply these concepts to our datasets/models.py file.

# datasets/models.py

from django.db import models
from django.conf import settings
from django_fsm import FSMField, transition # Import FSMField and transition




# ... your existing Image model ...

class AnnotationTask(models.Model):
    """
    Represents a single annotation task for an image.
    """
    class Status(models.TextChoices):
        PENDING = 'pending', 'Pending'
        IN_PROGRESS = 'in_progress', 'In Progress'
        COMPLETED = 'completed', 'Completed'
        REJECTED = 'rejected', 'Rejected'

    image = models.ForeignKey(Image, on_delete=models.CASCADE, related_name='annotation_tasks')
    annotator = models.ForeignKey(settings.AUTH_USER_MODEL, on_delete=models.SET_NULL, null=True, blank=True, related_name='tasks')
    



    # Replace models.CharField with FSMField
    status = FSMField(
        max_length=20,
        choices=Status.choices,
        default=Status.PENDING,
        protected=True, # This is important!
        db_index=True
    )
    
    annotation_data = models.JSONField(null=True, blank=True)

    created_at = models.DateTimeField(auto_now_add=True)
    updated_at = models.DateTimeField(auto_now=True)




    # --- Add Transition Methods below ---
    @transition(field=status, source=Status.PENDING, target=Status.IN_PROGRESS)
    def start(self):
        """
        Move the task from 'pending' to 'in_progress'.
        This is where you could, for example, assign the task to a specific user.
        """
        pass

    @transition(field=status, source=Status.IN_PROGRESS, target=Status.COMPLETED)
    def complete(self, annotation_data=None):
        """
        Move the task from 'in_progress' to 'completed'.
        """
        if annotation_data:
            self.annotation_data = annotation_data

    @transition(field=status, source=Status.IN_PROGRESS, target=Status.REJECTED)
    def reject(self, reason=""):
        """
        Move the task from 'in_progress' to 'rejected'.
        You could add a field to the model to store the rejection reason.
        """
        pass

    class Meta:
        ordering = ['-created_at']

    def __str__(self):
        return f"Task for Image {self.image.id} ({self.status})"

Notice a few key changes:

1. We replaced models.CharField with FSMField. FSMField behaves like a CharField but has special state-tracking powers.
2. We added protected=True. This is crucial. It tells django-fsm to block any direct changes to the status field (e.g., task.status = 'completed'). The only way to change the state is through a decorated transition method.
3. We added the start, complete, and reject methods, each decorated with @transition. Each decorator clearly defines the valid source state and the resulting target state for that action.

Since we only changed a field type but not its fundamental nature (CharField-like), you don't need to run migrations for this specific change.

Step 2: Test in the Django Shell

Let's see our state machine enforce its rules. Open the Django shell:
python manage.py shell

Now, import your model and TransitionNotAllowed exception, and create a task.

from datasets.models import Image, AnnotationTask
from django_fsm import TransitionNotAllowed




# Make sure you have at least one Image object, with id=1
# If not, create one: Image.objects.create(filename='test.jpg', url='http://example.com/test.jpg')
img = Image.objects.get(id=1) 
task = AnnotationTask.objects.create(image=img)

print(f"Initial state: {task.status}")



# >>> Initial state: pending




# Let's try an invalid transition: completing a pending task
try:
    task.complete({"boxes": []})
except TransitionNotAllowed as e:
    print(f"Error: {e}")



# >>> Error: Can't switch from state 'pending' using method 'complete'




# Now, the correct transition: start the task
task.start()
task.save()
print(f"State after starting: {task.status}")



# >>> State after starting: in_progress




# Now that it's in progress, we can complete it
task.complete({"boxes": [(10, 10, 50, 50)]})
task.save()
print(f"Final state: {task.status}")



# >>> Final state: completed
print(f"Annotation data: {task.annotation_data}")



# >>> Annotation data: {'boxes': [[10, 10, 50, 50]]}

This is the magic of django-fsm. It raises a TransitionNotAllowed exception when you try to perform an invalid action, making your application's logic far more robust.

Exposing Transitions in the API

Our model is now protected, but how do clients use these transitions? Our old PUT/PATCH endpoint is now problematic. If a client tries to PATCH the status field, it will fail because the field is protected.

The correct RESTful approach is to expose each transition as a distinct action on the resource. Instead of PATCH /tasks/1/ with {"status": "completed"}, the client will POST /tasks/1/complete/. This is explicit and action-oriented. We can achieve this easily with DRF's @action decorator.

First, let's make our status field read-only in the serializer to prevent confusion.

In datasets/serializers.py:

# datasets/serializers.py

class AnnotationTaskSerializer(serializers.ModelSerializer):
    class Meta:
        model = AnnotationTask
        fields = [
            'id', 'image', 'annotator', 'status', 
            'annotation_data', 'created_at', 'updated_at'
        ]



        # Make status read-only and add annotation_data to the read_only_fields
        # as it will now be set via the complete transition.
        read_only_fields = ['id', 'annotator', 'status', 'created_at', 'updated_at', 'annotation_data']




    # The create method is fine as it is.

Now, let's add the actions to the ViewSet in datasets/views.py.

# datasets/views.py
from rest_framework import viewsets, permissions, status
from rest_framework.decorators import action
from rest_framework.response import Response
from django_fsm import TransitionNotAllowed
from .models import AnnotationTask
from .serializers import AnnotationTaskSerializer

class AnnotationTaskViewSet(viewsets.ModelViewSet):
    """
    API endpoint that allows annotation tasks to be viewed or edited.
    Provides actions to transition task states.
    """
    queryset = AnnotationTask.objects.all().select_related('image', 'annotator')
    serializer_class = AnnotationTaskSerializer
    permission_classes = [permissions.IsAuthenticated]




    # ... get_queryset method remains the same ...
    def get_queryset(self):
        user = self.request.user
        if user.is_staff:
            return AnnotationTask.objects.all()
        return AnnotationTask.objects.filter(annotator=user)
        



    # --- Add custom actions for transitions ---

    @action(detail=True, methods=['post'])
    def start(self, request, pk=None):
        task = self.get_object()
        try:
            task.start()
            task.save()
        except TransitionNotAllowed:
            return Response({'error': 'Transition not allowed.'}, status=status.HTTP_400_BAD_REQUEST)
        return Response(self.get_serializer(task).data)

    @action(detail=True, methods=['post'])
    def complete(self, request, pk=None):
        task = self.get_object()
        annotation_data = request.data.get('annotation_data')
        if not annotation_data:
             return Response({'error': 'annotation_data is required.'}, status=status.HTTP_400_BAD_REQUEST)
        
        try:
            task.complete(annotation_data=annotation_data)
            task.save()
        except TransitionNotAllowed:
            return Response({'error': 'Transition not allowed.'}, status=status.HTTP_400_BAD_REQUEST)
        return Response(self.get_serializer(task).data)

    @action(detail=True, methods=['post'])
    def reject(self, request, pk=None):
        task = self.get_object()
        try:
            task.reject()
            task.save()
        except TransitionNotAllowed:
            return Response({'error': 'Transition not allowed.'}, status=status.HTTP_400_BAD_REQUEST)
        return Response(self.get_serializer(task).data)

With these changes, your API now has new endpoints:

• POST /api/datasets/annotation-tasks/{id}/start/
• POST /api/datasets/annotation-tasks/{id}/complete/
• POST /api/datasets/annotation-tasks/{id}/reject/

Restart your development server. Go to the detail view of an AnnotationTask in the browsable API (e.g., http://127.0.0.1:8000/api/datasets/annotation-tasks/1/). Under the "Extra Actions" dropdown, you will now see buttons for "Start", "Complete", and "Reject", which you can use to test your new state transition endpoints.

Visualizing Your State Machine

One of the most powerful features of django-fsm is its ability to automatically generate a diagram of your state machine. This is invaluable for documentation, onboarding new team members, and communicating with product managers.

You'll need to install graphviz on your system.

• On macOS (with Homebrew): brew install graphviz
• On Debian/Ubuntu: sudo apt-get install graphviz
• On Windows: Download from the official website and add it to your system's PATH.

Then, install the Python library:
pip install pydot (or pygraphviz can also be used). Add pydot to your requirements.txt.

Now, run this management command from your project root:

python manage.py graph_transitions -o annotation_lifecycle.png datasets.AnnotationTask

This command inspects your AnnotationTask model, finds all the @transition decorators, and generates an image file named annotation_lifecycle.png in your project's root directory. Open it. You will see a clear, visual representation of the exact workflow you just coded!

Conclusion

You have now transformed your simple CRUD API into a robust, workflow-driven system. By implementing a state machine, you've ensured that your annotation tasks follow a predictable lifecycle, preventing invalid data and making your application's behavior explicit and maintainable. This is a fundamental pattern used in building complex backend systems, especially for pipelines and multi-step processes like those at Infilect.

Key Takeaways:

• Finite State Machines (FSMs) are essential for managing objects that have a defined lifecycle.
• The django-fsm library provides an elegant, declarative way to implement FSMs using the FSMField and @transition decorator.
• The protected=True flag is vital for preventing direct state modification, forcing all changes to go through controlled transitions.
• Exposing FSM transitions as custom @action endpoints on a DRF ViewSet is the standard, RESTful pattern.
• You can automatically generate visual diagrams of your workflows directly from your code, which is excellent for documentation and communication.

In our next module, we will shift gears from building application logic to preparing our project for the real world. You will learn how to containerize your Django application using Docker, the first and most critical step in creating a portable, scalable, and deployable backend system.