[Backend] Add custom worker image for challenges (#4077)

* Add support for defining custom worker Signed-off-by: AyushR1 <[email protected]> * Fix global dict change and only modify worker image Signed-off-by: AyushR1 <[email protected]> * Revert "Fix global dict change and only modify worker image" This reverts commit 1fa5795. * Update according to review Signed-off-by: AyushR1 <[email protected]> --------- Signed-off-by: AyushR1 <[email protected]> Co-authored-by: Gunjan Chhablani <[email protected]>
Cloud-CV · Aug 9, 2023 · 3cf4204 · 3cf4204
1 parent e3f83d7
commit 3cf4204
Show file tree

Hide file tree

Showing 3 changed files with 35 additions and 6 deletions.
diff --git a/apps/challenges/aws_utils.py b/apps/challenges/aws_utils.py
@@ -175,6 +175,11 @@ def register_task_def_by_challenge_pk(client, queue_name, challenge):
     AWS_SES_REGION_NAME = settings.AWS_SES_REGION_NAME
     AWS_SES_REGION_ENDPOINT = settings.AWS_SES_REGION_ENDPOINT
 
+    if challenge.worker_image_url:
+        updated_settings = {**COMMON_SETTINGS_DICT, "WORKER_IMAGE": challenge.worker_image_url}
+    else:
+        updated_settings = COMMON_SETTINGS_DICT
+
     if execution_role_arn:
         from .utils import get_aws_credentials_for_challenge
 
@@ -208,7 +213,7 @@ def register_task_def_by_challenge_pk(client, queue_name, challenge):
                         log_group_name=log_group_name,
                         EVALAI_DNS=EVALAI_DNS,
                         EFS_ID=efs_id,
-                        **COMMON_SETTINGS_DICT,
+                        **updated_settings,
                         **challenge_aws_keys,
                     )
                 )
@@ -221,7 +226,7 @@ def register_task_def_by_challenge_pk(client, queue_name, challenge):
                         log_group_name=log_group_name,
                         AWS_SES_REGION_NAME=AWS_SES_REGION_NAME,
                         AWS_SES_REGION_ENDPOINT=AWS_SES_REGION_ENDPOINT,
-                        **COMMON_SETTINGS_DICT,
+                        **updated_settings,
                         **aws_keys,
                     )
                 )
@@ -231,7 +236,7 @@ def register_task_def_by_challenge_pk(client, queue_name, challenge):
                     submission_container=submission_container,
                     CPU=worker_cpu_cores,
                     MEMORY=worker_memory,
-                    **COMMON_SETTINGS_DICT,
+                    **updated_settings,
                 )
             else:
                 definition = task_definition_code_upload_worker.format(
@@ -248,7 +253,7 @@ def register_task_def_by_challenge_pk(client, queue_name, challenge):
                     log_group_name=log_group_name,
                     EVALAI_DNS=EVALAI_DNS,
                     EFS_ID=efs_id,
-                    **COMMON_SETTINGS_DICT,
+                    **updated_settings,
                     **challenge_aws_keys,
                 )
         else:
@@ -262,7 +267,7 @@ def register_task_def_by_challenge_pk(client, queue_name, challenge):
                 log_group_name=log_group_name,
                 AWS_SES_REGION_NAME=AWS_SES_REGION_NAME,
                 AWS_SES_REGION_ENDPOINT=AWS_SES_REGION_ENDPOINT,
-                **COMMON_SETTINGS_DICT,
+                **updated_settings,
                 **challenge_aws_keys,
             )
         definition = eval(definition)
@@ -669,6 +674,11 @@ def scale_resources(challenge, worker_cpu_cores, worker_memory):
         logger.exception(e)
         return e.response
 
+    if challenge.worker_image_url:
+        updated_settings = {**COMMON_SETTINGS_DICT, "WORKER_IMAGE": challenge.worker_image_url}
+    else:
+        updated_settings = COMMON_SETTINGS_DICT
+
     queue_name = challenge.queue
     container_name = "worker_{}".format(queue_name)
     log_group_name = get_log_group_name(challenge.pk)
@@ -683,7 +693,7 @@ def scale_resources(challenge, worker_cpu_cores, worker_memory):
         log_group_name=log_group_name,
         AWS_SES_REGION_NAME=settings.AWS_SES_REGION_NAME,
         AWS_SES_REGION_ENDPOINT=settings.AWS_SES_REGION_ENDPOINT,
-        **COMMON_SETTINGS_DICT,
+        **updated_settings,
         **challenge_aws_keys,
     )
     task_def = eval(task_def)

diff --git a/apps/challenges/migrations/0099_challenge_worker_image_url.py b/apps/challenges/migrations/0099_challenge_worker_image_url.py
@@ -0,0 +1,18 @@
+# Generated by Django 2.2.20 on 2023-07-31 22:30
+
+from django.db import migrations, models
+
+
+class Migration(migrations.Migration):
+
+    dependencies = [
+        ('challenges', '0098_challenge_tags_domains'),
+    ]
+
+    operations = [
+        migrations.AddField(
+            model_name='challenge',
+            name='worker_image_url',
+            field=models.URLField(blank=True, default=None, null=True),
+        ),
+    ]
diff --git a/apps/challenges/models.py b/apps/challenges/models.py
@@ -198,6 +198,7 @@ def __init__(self, *args, **kwargs):
     job_memory = models.CharField(
         max_length=256, null=True, blank=True, default="8Gi"
     )
+    worker_image_url = models.URLField(max_length=200, blank=True, null=True, default=None)
 
     class Meta:
         app_label = "challenges"