[Backend] Add custom sqs retention time (#4185)

* Add custom sqs retention time

* Fix flake8 issues

* Add missing migration

* Add fixes

* Fix errors

* Update common.py

* Fix migrations

Author: gchhablani

apps/base/utils.py

@@ -201,9 +201,10 @@ def get_or_create_sqs_queue(queue_name, challenge=None):
             != "AWS.SimpleQueueService.NonExistentQueue"
         ):
             logger.exception("Cannot get queue: {}".format(queue_name))
+        sqs_retention_period = SQS_RETENTION_PERIOD if challenge is None else str(challenge.sqs_retention_period)
         queue = sqs.create_queue(
             QueueName=queue_name,
-            Attributes={"MessageRetentionPeriod": SQS_RETENTION_PERIOD},
+            Attributes={"MessageRetentionPeriod": sqs_retention_period},
         )
     return queue
 

apps/challenges/aws_utils.py

@@ -771,6 +771,34 @@ def create_ec2_instance(challenge, ec2_storage=None, worker_instance_type=None,
         }
 
 
+def update_sqs_retention_period(challenge):
+    """
+    Update the SQS retention period for a challenge.
+
+    Args:
+        challenge (Challenge): The challenge for which the SQS retention period is to be updated.
+
+    Returns:
+        dict: A dictionary containing the status and message of the operation.
+    """
+    sqs_retention_period = str(challenge.sqs_retention_period)
+    try:
+        sqs = get_boto3_client("sqs", aws_keys)
+        queue_url = sqs.get_queue_url(QueueName=challenge.queue)['QueueUrl']
+        response = sqs.set_queue_attributes(
+            QueueUrl=queue_url,
+            Attributes={
+                'MessageRetentionPeriod': sqs_retention_period
+            }
+        )
+        return {"message": response.json()}
+    except Exception as e:
+        logger.exception(e)
+        return {
+            "error": e.response,
+        }
+
+
 def start_workers(queryset):
     """
     The function called by the admin action method to start all the selected workers.
@@ -1794,3 +1822,16 @@ def setup_ec2(challenge):
     if challenge_obj.ec2_instance_id:
         return start_ec2_instance(challenge_obj)
     return create_ec2_instance(challenge_obj)
+
+
+@app.task
+def update_sqs_retention_period_task(challenge):
+    """
+    Updates sqs retention period for a challenge when the attribute is changed.
+
+    Arguments:
+        challenge {<class 'django.db.models.query.QuerySet'>} -- instance of the model calling the post hook
+    """
+    for obj in serializers.deserialize("json", challenge):
+        challenge_obj = obj.object
+    return update_sqs_retention_period(challenge_obj)

apps/challenges/migrations/0112_challenge_sqs_retention_period.py

@@ -0,0 +1,18 @@
+# Generated by Django 2.2.20 on 2023-12-10 15:16
+
+from django.db import migrations, models
+
+
+class Migration(migrations.Migration):
+
+    dependencies = [
+        ('challenges', '0111_alter_challenge_ephemeral_storage_default'),
+    ]
+
+    operations = [
+        migrations.AddField(
+            model_name='challenge',
+            name='sqs_retention_period',
+            field=models.PositiveIntegerField(default=259200, verbose_name='SQS Retention Period'),
+        ),
+    ]

apps/challenges/models.py

@@ -35,6 +35,7 @@ def __init__(self, *args, **kwargs):
         super(Challenge, self).__init__(*args, **kwargs)
         self._original_evaluation_script = self.evaluation_script
         self._original_approved_by_admin = self.approved_by_admin
+        self._original_sqs_retention_period = self.sqs_retention_period
 
     title = models.CharField(max_length=100, db_index=True)
     short_description = models.TextField(null=True, blank=True)
@@ -126,6 +127,10 @@ def __init__(self, *args, **kwargs):
         verbose_name="SQS queue name",
         db_index=True,
     )
+    sqs_retention_period = models.PositiveIntegerField(
+        default=259200,
+        verbose_name="SQS Retention Period"
+    )
     is_docker_based = models.BooleanField(
         default=False, verbose_name="Is Docker Based", db_index=True
     )
@@ -277,6 +282,21 @@ def create_eks_cluster_or_ec2_for_challenge(sender, instance, created, **kwargs)
     aws.challenge_approval_callback(sender, instance, field_name, **kwargs)
 
 
+@receiver(signals.post_save, sender="challenges.Challenge")
+def update_sqs_retention_period_for_challenge(sender, instance, created, **kwargs):
+    field_name = "sqs_retention_period"
+    import challenges.aws_utils as aws
+
+    if not created and is_model_field_changed(instance, field_name):
+        serialized_obj = serializers.serialize("json", [instance])
+        aws.update_sqs_retention_period_task.delay(serialized_obj)
+        # Update challenge
+        curr = getattr(instance, "{}".format(field_name))
+        challenge = instance
+        challenge._original_sqs_retention_period = curr
+        challenge.save()
+
+
 class DatasetSplit(TimeStampedModel):
     name = models.CharField(max_length=100)
     codename = models.CharField(max_length=100)

apps/jobs/sender.py

@@ -62,9 +62,10 @@ def get_or_create_sqs_queue(queue_name, challenge=None):
             ex.response["Error"]["Code"]
             == "AWS.SimpleQueueService.NonExistentQueue"
         ):
+            sqs_retention_period = SQS_RETENTION_PERIOD if challenge is None else str(challenge.sqs_retention_period)
             queue = sqs.create_queue(
                 QueueName=queue_name,
-                Attributes={"MessageRetentionPeriod": SQS_RETENTION_PERIOD},
+                Attributes={"MessageRetentionPeriod": sqs_retention_period},
             )
         else:
             logger.exception("Cannot get or create Queue")

scripts/workers/submission_worker.py

@@ -806,9 +806,10 @@ def get_or_create_sqs_queue(queue_name, challenge=None):
             != "AWS.SimpleQueueService.NonExistentQueue"
         ):
             logger.exception("Cannot get queue: {}".format(queue_name))
+        sqs_retention_period = SQS_RETENTION_PERIOD if challenge is None else str(challenge.sqs_retention_period)
         queue = sqs.create_queue(
             QueueName=queue_name,
-            Attributes={"MessageRetentionPeriod": SQS_RETENTION_PERIOD},
+            Attributes={"MessageRetentionPeriod": sqs_retention_period},
         )
     return queue
 

settings/common.py

@@ -406,4 +406,4 @@
 }
 
 # SQS Queue Message Retention Period
-SQS_RETENTION_PERIOD = "1209600"
+SQS_RETENTION_PERIOD = "345600"