fix: PLT-843: Webhook job memory optimization (#8136)

Co-authored-by: Marcel Canu <[email protected]>
Co-authored-by: triklozoid <[email protected]>
Co-authored-by: farioas <[email protected]>

Author: 3 people

label_studio/core/settings/base.py

@@ -650,7 +650,7 @@ def collect_versions_dummy(**kwargs):
 COLLECT_VERSIONS = collect_versions_dummy
 
 WEBHOOK_TIMEOUT = float(get_env('WEBHOOK_TIMEOUT', 1.0))
-WEBHOOK_BATCH_SIZE = int(get_env('WEBHOOK_BATCH_SIZE', 100))
+WEBHOOK_BATCH_SIZE = int(get_env('WEBHOOK_BATCH_SIZE', 5000))
 WEBHOOK_SERIALIZERS = {
     'project': 'webhooks.serializers_for_hooks.ProjectWebhookSerializer',
     'task': 'webhooks.serializers_for_hooks.TaskWebhookSerializer',

label_studio/tests/webhooks/test_webhooks.py

@@ -440,3 +440,96 @@ def test_start_training_webhook(setup_project_dialog, ml_start_training_webhook,
     assert request_history[0].url == webhook.url
     assert 'project' in request_history[0].json()
     assert request_history[0].json()['action'] == 'START_TRAINING'
+
+
+@pytest.mark.django_db
+def test_webhook_batching_with_feature_flag(configured_project, organization_webhook):
+    """Test that webhooks are sent in batches when feature flag is enabled."""
+    from unittest.mock import patch
+
+    from django.conf import settings
+    from tasks.models import Task
+    from webhooks.utils import emit_webhooks_for_instance_sync
+
+    webhook = organization_webhook
+    project = configured_project
+
+    # Create multiple tasks to test batching
+    tasks = []
+    for i in range(250):  # Create more than WEBHOOK_BATCH_SIZE
+        task = Task.objects.create(data={'text': f'Test task {i}'}, project=project)
+        tasks.append(task)
+
+    # Test with feature flag enabled
+    with patch('webhooks.utils.flag_set') as mock_flag_set:
+        mock_flag_set.return_value = True
+
+        with requests_mock.Mocker(real_http=True) as m:
+            m.register_uri('POST', webhook.url)
+
+            # Set batch size to 100 for testing
+            original_batch_size = getattr(settings, 'WEBHOOK_BATCH_SIZE', 100)
+            settings.WEBHOOK_BATCH_SIZE = 100
+
+            try:
+                emit_webhooks_for_instance_sync(
+                    webhook.organization, webhook.project, WebhookAction.TASKS_CREATED, instance=tasks
+                )
+
+                # Should have 3 requests (250 tasks / 100 batch size = 3 batches)
+                webhook_requests = list(filter(lambda x: x.url == webhook.url, m.request_history))
+                assert len(webhook_requests) == 3
+
+                # Check first batch has 100 tasks
+                first_batch = webhook_requests[0].json()
+                assert 'tasks' in first_batch
+                assert len(first_batch['tasks']) == 100
+
+                # Check second batch has 100 tasks
+                second_batch = webhook_requests[1].json()
+                assert len(second_batch['tasks']) == 100
+
+                # Check third batch has 50 tasks (remaining)
+                third_batch = webhook_requests[2].json()
+                assert len(third_batch['tasks']) == 50
+
+            finally:
+                settings.WEBHOOK_BATCH_SIZE = original_batch_size
+
+
+@pytest.mark.django_db
+def test_webhook_no_batching_without_feature_flag(configured_project, organization_webhook):
+    """Test that webhooks are sent in single request when feature flag is disabled."""
+    from unittest.mock import patch
+
+    from tasks.models import Task
+    from webhooks.utils import emit_webhooks_for_instance_sync
+
+    webhook = organization_webhook
+    project = configured_project
+
+    # Create multiple tasks
+    tasks = []
+    for i in range(150):
+        task = Task.objects.create(data={'text': f'Test task {i}'}, project=project)
+        tasks.append(task)
+
+    # Test with feature flag disabled
+    with patch('webhooks.utils.flag_set') as mock_flag_set:
+        mock_flag_set.return_value = False
+
+        with requests_mock.Mocker(real_http=True) as m:
+            m.register_uri('POST', webhook.url)
+
+            emit_webhooks_for_instance_sync(
+                webhook.organization, webhook.project, WebhookAction.TASKS_CREATED, instance=tasks
+            )
+
+            # Should have only 1 request (no batching)
+            webhook_requests = list(filter(lambda x: x.url == webhook.url, m.request_history))
+            assert len(webhook_requests) == 1
+
+            # Check all 150 tasks are in single request
+            request_data = webhook_requests[0].json()
+            assert 'tasks' in request_data
+            assert len(request_data['tasks']) == 150

label_studio/webhooks/utils.py

@@ -7,9 +7,12 @@
 from core.utils.common import load_func
 from django.conf import settings
 from django.db.models import Q
+from django.db.models.query import QuerySet
 
 from .models import Webhook, WebhookAction
 
+logger = logging.getLogger(__name__)
+
 
 def get_active_webhooks(organization, project, action):
     """Return all active webhooks for organization or project by action.
@@ -71,37 +74,78 @@ def emit_webhooks_sync(organization, project, action, payload):
         run_webhook_sync(wh, action, payload)
 
 
-def emit_webhooks_for_instance_sync(organization, project, action, instance=None):
-    """Run all active webhooks for the action using instances as payload.
+def _process_webhook_batch(webhooks, project, action, batch, action_meta):
+    """Process a single batch of instances for webhooks.
 
-    Be sure WebhookAction.ACTIONS contains all required fields.
+    Args:
+        webhooks: Active webhooks to send
+        project: Project instance (optional)
+        action: Action name
+        batch: Batch of instances to process
+        action_meta: Action metadata from WebhookAction.ACTIONS
     """
-    webhooks = get_active_webhooks(organization, project, action)
-    if not webhooks.exists():
-        return
     payload = {}
-    # if instances and there is a webhook that sends payload
-    # get serialized payload
-    action_meta = WebhookAction.ACTIONS[action]
 
-    if instance and isinstance(instance, list) and isinstance(instance[0], int):
-        instance = action_meta['model'].objects.filter(id__in=instance)
-
-    if instance and webhooks.filter(send_payload=True).exists():
+    if batch and webhooks.filter(send_payload=True).exists():
         serializer_class = action_meta.get('serializer')
         if serializer_class:
-            payload[action_meta['key']] = serializer_class(instance=instance, many=action_meta['many']).data
+            payload[action_meta['key']] = serializer_class(instance=batch, many=action_meta['many']).data
         if project and payload:
             payload['project'] = load_func(settings.WEBHOOK_SERIALIZERS['project'])(instance=project).data
         if payload and 'nested-fields' in action_meta:
             for key, value in action_meta['nested-fields'].items():
                 payload[key] = value['serializer'](
-                    instance=get_nested_field(instance, value['field']), many=value['many']
+                    instance=get_nested_field(batch, value['field']), many=value['many']
                 ).data
+
     for wh in webhooks:
         run_webhook_sync(wh, action, payload)
 
 
+def emit_webhooks_for_instance_sync(organization, project, action, instance=None):
+    """Run all active webhooks for the action using instances as payload.
+
+    Be sure WebhookAction.ACTIONS contains all required fields.
+    """
+    webhooks = get_active_webhooks(organization, project, action)
+    if not webhooks.exists():
+        return
+
+    action_meta = WebhookAction.ACTIONS[action]
+
+    # Convert list of IDs to queryset
+    if instance and isinstance(instance, list) and isinstance(instance[0], int):
+        instance = action_meta['model'].objects.filter(id__in=instance)
+
+    # Check if batching is needed
+    is_batch_collection = isinstance(instance, (list, QuerySet))
+    use_batching = is_batch_collection and flag_set('fflag_fix_back_plt_843_webhook_memory_improvement_12082025_short')
+
+    if use_batching:
+        # Process in batches
+        batch_size = settings.WEBHOOK_BATCH_SIZE
+
+        if isinstance(instance, QuerySet):
+            # For QuerySets, use iterator with chunk_size
+            total_count = instance.count()
+            logger.debug(f'Processing webhook for {total_count} instances in batches of {batch_size}')
+            for i in range(0, total_count, batch_size):
+                batch = instance[i : i + batch_size]
+                logger.debug(f'Processing batch {i // batch_size + 1} with {batch.count()} instances')
+                _process_webhook_batch(webhooks, project, action, batch, action_meta)
+        else:
+            # For lists, slice directly
+            total_count = len(instance)
+            logger.debug(f'Processing webhook for {total_count} instances in batches of {batch_size}')
+            for i in range(0, len(instance), batch_size):
+                batch = instance[i : i + batch_size]
+                logger.debug(f'Processing batch {i // batch_size + 1} with {len(batch)} instances')
+                _process_webhook_batch(webhooks, project, action, batch, action_meta)
+    else:
+        # Original behavior - process all at once
+        _process_webhook_batch(webhooks, project, action, instance, action_meta)
+
+
 def run_webhook(webhook, action, payload=None):
     """Run one webhook for action.