Feature/1097 track subtask retries

src/model_execution_worker/celery_request_handler.py

@@ -1,5 +1,5 @@
 import logging
-from celery import Task
+from celery import Task, signature
 from celery.worker.request import Request as CeleryRequest
 from celery.exceptions import WorkerLostError
 
@@ -20,6 +20,7 @@ class CustomRequest(CeleryRequest):
     def on_failure(self, exc_info, send_failed_event=True, return_ok=False):
         """Handler called if the task raised an exception."""
         exc = exc_info.exception
+
         if isinstance(exc, ExceptionWithTraceback):
             exc = exc.exc
 
@@ -38,6 +39,21 @@ def on_failure(self, exc_info, send_failed_event=True, return_ok=False):
             else:
                 self.task.reject_on_worker_lost = False
 
+        # report retry attempt
+        try:
+            task_id = self.task_id
+            task_args = self.kwargs
+            signature('subtask_retry_log').delay(
+                task_args.get('analysis_id'),
+                task_args.get('initiator_id'),
+                task_args.get('slug'),
+                task_id,
+                exc_info.traceback,
+            )
+        except Exception as e:
+            logger.error(f'Falied to store retry attempt logs: {exc_info}')
+            logger.exception(e)
+
         super().on_failure(exc_info, send_failed_event, return_ok)
 
 

src/server/oasisapi/analyses/migrations/0016_auto_20250115_1304.py

@@ -0,0 +1,25 @@
+# Generated by Django 3.2.25 on 2025-01-15 13:04
+
+from django.db import migrations, models
+import django.db.models.deletion
+
+
+class Migration(migrations.Migration):
+
+    dependencies = [
+        ('files', '0014_merge_0006_relatedfile_groups_0013_mappingfile_groups'),
+        ('analyses', '0015_auto_20241125_1407'),
+    ]
+
+    operations = [
+        migrations.AddField(
+            model_name='analysistaskstatus',
+            name='retry_count',
+            field=models.IntegerField(default=0, editable=False),
+        ),
+        migrations.AddField(
+            model_name='analysistaskstatus',
+            name='retry_log',
+            field=models.ForeignKey(blank=True, default=None, editable=False, null=True, on_delete=django.db.models.deletion.SET_NULL, related_name='analysis_run_status_retry_logs', to='files.relatedfile'),
+        ),
+    ]

src/server/oasisapi/analyses/models.py

@@ -102,6 +102,7 @@ class AnalysisTaskStatus(models.Model):
     end_time = models.DateTimeField(null=True, default=None, editable=False)
     name = models.CharField(max_length=255, editable=False)
     slug = models.SlugField(max_length=255, editable=False)
+    retry_count = models.IntegerField(editable=False, default=0)
 
     output_log = models.ForeignKey(
         RelatedFile,
@@ -120,6 +121,15 @@ class AnalysisTaskStatus(models.Model):
         related_name='analysis_run_status_error_logs',
         editable=False,
     )
+    retry_log = models.ForeignKey(
+        RelatedFile,
+        on_delete=models.SET_NULL,
+        blank=True,
+        null=True,
+        default=None,
+        related_name='analysis_run_status_retry_logs',
+        editable=False,
+    )
 
     objects = AnalysisTaskStatusQuerySet.as_manager()
 
@@ -140,6 +150,10 @@ def get_error_log_url(self, request=None, namespace=None):
         override_ns = f'{namespace}:' if namespace else 'v2-analyses:'
         return reverse(f'{override_ns}analysis-task-status-error-log', kwargs={'pk': self.pk}, request=request)
 
+    def get_retry_log_url(self, request=None, namespace=None):
+        override_ns = f'{namespace}:' if namespace else 'v2-analyses:'
+        return reverse(f'{override_ns}analysis-task-status-retry-log', kwargs={'pk': self.pk}, request=request)
+
 
 class Analysis(TimeStampedModel):
     status_choices = Choices(
@@ -805,6 +819,7 @@ def delete_connected_task_logs(sender, instance, **kwargs):
     for_removal = [
         'output_log',
         'error_log',
+        'retry_log',
     ]
     for ref in for_removal:
         obj_ref = getattr(instance, ref)

src/server/oasisapi/analyses/v2_api/serializers.py

@@ -39,6 +39,7 @@ def create_settings_file(data, user):
 class AnalysisTaskStatusSerializer(serializers.ModelSerializer):
     output_log = serializers.SerializerMethodField()
     error_log = serializers.SerializerMethodField()
+    retry_log = serializers.SerializerMethodField()
 
     class Meta:
         model = AnalysisTaskStatus
@@ -55,6 +56,8 @@ class Meta:
             'end_time',
             'output_log',
             'error_log',
+            'retry_log',
+            'retry_count',
         )
 
     @swagger_serializer_method(serializer_or_field=serializers.URLField)
@@ -67,6 +70,11 @@ def get_error_log(self, instance):
         request = self.context.get('request')
         return instance.get_error_log_url(request=request) if instance.error_log else None
 
+    @swagger_serializer_method(serializer_or_field=serializers.URLField)
+    def get_retry_log(self, instance):
+        request = self.context.get('request')
+        return instance.get_retry_log_url(request=request) if instance.retry_log else None
+
 
 class AnalysisListSerializer(serializers.Serializer):
     """ Read Only Analyses Deserializer for efficiently returning a list of all

src/server/oasisapi/analyses/v2_api/task_controller.py

@@ -578,7 +578,6 @@ def generate_losses(cls, analysis: 'Analysis', initiator: User, events_total: in
         logger.debug(f"'generate_losses' - canvas dispatched, analyses={analysis.pk}, run_uuid={run_data_uuid}")
 
         # Update sub-task ids
-        # from celery.contrib import rdb; rdb.set_trace()
         celery_tasks_list = cls.extract_celery_task_ids(task)
         for sub_t in analysis.sub_task_statuses.all():
             try:

src/server/oasisapi/analyses/v2_api/tasks.py

@@ -764,6 +764,29 @@ def subtask_error_log(analysis_id, initiator_id, slug, task_id, log_file):
     )
 
 
+@celery_app_v2.task(name='subtask_retry_log')
+def subtask_retry_log(analysis_id, initiator_id, slug, task_id, error_trace):
+    # Increment retry counter and append to retry log
+    task = AnalysisTaskStatus.objects.get(analysis_id=analysis_id, slug=slug)
+
+    with TemporaryFile() as tmp_file:
+        if hasattr(task.retry_log, 'read'):
+            tmp_file.write(task.retry_log.read())
+
+        tmp_file.write(f'\n=== {slug}, try={task.retry_count}  ===\n'.encode('utf-8'))
+        tmp_file.write(error_trace.encode('utf-8'))
+        tmp_file.seek(0)
+
+        setattr(task, 'retry_log', RelatedFile.objects.create(
+            file=File(tmp_file, name='{}.txt'.format(uuid.uuid4().hex)),
+            filename='{}-retry.log'.format(task_id),
+            content_type='text/plain',
+            creator_id=initiator_id,
+        ))
+        task.retry_count += 1
+        task.save(update_fields=["retry_count", "retry_log"])
+
+
 @celery_app_v2.task(name='set_task_status_v2')
 def set_task_status(analysis_pk, task_status, dt):
     try:

src/server/oasisapi/analyses/v2_api/viewsets.py

@@ -699,3 +699,15 @@ def error_log(self, request, pk=None, version=None):
         Disassociates the task status' `error_log` contents
         """
         return handle_related_file(self.get_object(), 'error_log', request, ['text/plain'])
+
+    @swagger_auto_schema(methods=['get'], responses={200: FILE_RESPONSE})
+    @action(methods=['get', 'delete'], detail=True)
+    def retry_log(self, request, pk=None, version=None):
+        """
+        get:
+        Gets the task status' `retry_log` contents
+
+        delete:
+        Disassociates the task status' `retry_log` contents
+        """
+        return handle_related_file(self.get_object(), 'retry_log', request, ['text/plain'])