Split compute and storage tasks (#4064)

See DIAGNijmegen/rse-roadmap#408

Author: jmsmkn

app/config/settings.py

@@ -1262,6 +1262,14 @@ def sentry_before_send(event, hint):
         "task": "grandchallenge.invoices.tasks.send_challenge_invoice_overdue_reminder_emails",
         "schedule": crontab(day_of_month=1, hour=6, minute=0),
     },
+    "update_challenge_storage_size": {
+        "task": "grandchallenge.challenges.tasks.update_challenge_storage_size",
+        "schedule": crontab(hour=6, minute=15),
+    },
+    "update_challenge_compute_costs": {
+        "task": "grandchallenge.challenges.tasks.update_challenge_compute_costs",
+        "schedule": crontab(minute=45),
+    },
     "delete_users_who_dont_login": {
         "task": "grandchallenge.profiles.tasks.delete_users_who_dont_login",
         "schedule": timedelta(hours=1),
@@ -1286,10 +1294,6 @@ def sentry_before_send(event, hint):
         "task": "grandchallenge.core.tasks.cleanup_celery_backend",
         "schedule": timedelta(hours=1),
     },
-    "update_compute_costs_and_storage_size": {
-        "task": "grandchallenge.challenges.tasks.update_compute_costs_and_storage_size",
-        "schedule": timedelta(hours=2),
-    },
     "logout_privileged_users": {
         "task": "grandchallenge.browser_sessions.tasks.logout_privileged_users",
         "schedule": timedelta(hours=1),

app/grandchallenge/challenges/costs.py

@@ -19,10 +19,23 @@
 )
 
 
-def annotate_job_duration_and_compute_costs(*, phase):
+def annotate_compute_costs(*, challenge):
     algorithm_job_utilizations = JobUtilization.objects.filter(
-        phase=phase,
+        challenge=challenge
+    )
+    evaluation_job_utilizations = EvaluationUtilization.objects.filter(
+        challenge=challenge
+    )
+
+    update_compute_cost_euro_millicents(
+        obj=challenge,
+        algorithm_job_utilizations=algorithm_job_utilizations,
+        evaluation_job_utilizations=evaluation_job_utilizations,
     )
+
+
+def annotate_job_duration_and_compute_costs(*, phase):
+    algorithm_job_utilizations = JobUtilization.objects.filter(phase=phase)
     evaluation_job_utilizations = EvaluationUtilization.objects.filter(
         phase=phase, external_evaluation=False
     )
@@ -38,7 +51,25 @@ def annotate_job_duration_and_compute_costs(*, phase):
     )
 
 
-def annotate_compute_costs_and_storage_size(*, challenge):
+def update_compute_cost_euro_millicents(
+    *, obj, algorithm_job_utilizations, evaluation_job_utilizations
+):
+    algorithm_job_costs = algorithm_job_utilizations.aggregate(
+        Sum("compute_cost_euro_millicents")
+    )
+
+    evaluation_costs = evaluation_job_utilizations.aggregate(
+        Sum("compute_cost_euro_millicents")
+    )
+
+    items = [algorithm_job_costs, evaluation_costs]
+
+    obj.compute_cost_euro_millicents = sum(
+        item["compute_cost_euro_millicents__sum"] or 0 for item in items
+    )
+
+
+def annotate_storage_size(*, challenge):
     permission = Permission.objects.get(
         codename="view_job",
         content_type__app_label="algorithms",
@@ -61,19 +92,6 @@ def annotate_compute_costs_and_storage_size(*, challenge):
         evaluation_jobs=evaluation_jobs,
     )
 
-    algorithm_job_utilizations = JobUtilization.objects.filter(
-        challenge=challenge,
-    )
-    evaluation_job_utilizations = EvaluationUtilization.objects.filter(
-        challenge=challenge
-    )
-
-    update_compute_cost_euro_millicents(
-        obj=challenge,
-        algorithm_job_utilizations=algorithm_job_utilizations,
-        evaluation_job_utilizations=evaluation_job_utilizations,
-    )
-
 
 def update_size_in_storage_and_registry(
     *, challenge, algorithm_jobs, evaluation_jobs
@@ -170,21 +188,3 @@ def update_size_in_storage_and_registry(
     challenge.size_in_registry = sum(
         item.get("size_in_registry__sum") or 0 for item in items
     )
-
-
-def update_compute_cost_euro_millicents(
-    *, obj, algorithm_job_utilizations, evaluation_job_utilizations
-):
-    algorithm_job_costs = algorithm_job_utilizations.aggregate(
-        Sum("compute_cost_euro_millicents")
-    )
-
-    evaluation_costs = evaluation_job_utilizations.aggregate(
-        Sum("compute_cost_euro_millicents")
-    )
-
-    items = [algorithm_job_costs, evaluation_costs]
-
-    obj.compute_cost_euro_millicents = sum(
-        item["compute_cost_euro_millicents__sum"] or 0 for item in items
-    )

app/grandchallenge/challenges/tasks.py

@@ -10,8 +10,9 @@
 from psycopg.errors import LockNotAvailable
 
 from grandchallenge.challenges.costs import (
-    annotate_compute_costs_and_storage_size,
+    annotate_compute_costs,
     annotate_job_duration_and_compute_costs,
+    annotate_storage_size,
 )
 from grandchallenge.challenges.emails import (
     send_onboarding_task_due_reminder,
@@ -94,20 +95,14 @@ def wrapper(*args, **kwargs):
 
 
 @acks_late_2xlarge_task
-def update_compute_costs_and_storage_size():
+def update_challenge_compute_costs():
     for challenge in Challenge.objects.with_available_compute().iterator():
         with transaction.atomic():
-            annotate_compute_costs_and_storage_size(challenge=challenge)
+            annotate_compute_costs(challenge=challenge)
 
             @retry_with_backoff((LockNotAvailable,))
             def save_challenge():
-                challenge.save(
-                    update_fields=(
-                        "size_in_storage",
-                        "size_in_registry",
-                        "compute_cost_euro_millicents",
-                    )
-                )
+                challenge.save(update_fields=("compute_cost_euro_millicents",))
 
             save_challenge()
 
@@ -128,6 +123,24 @@ def save_phase():
             save_phase()
 
 
+@acks_late_2xlarge_task
+def update_challenge_storage_size():
+    for challenge in Challenge.objects.iterator():
+        with transaction.atomic():
+            annotate_storage_size(challenge=challenge)
+
+            @retry_with_backoff((LockNotAvailable,))
+            def save_challenge():
+                challenge.save(
+                    update_fields=(
+                        "size_in_storage",
+                        "size_in_registry",
+                    )
+                )
+
+            save_challenge()
+
+
 class OnboardingTaskInfo(NamedTuple):
     challenge: str
     num_is_overdue: int

app/tests/challenges_tests/test_tasks.py

@@ -11,8 +11,8 @@
 )
 from grandchallenge.challenges.tasks import (
     send_onboarding_task_reminder_emails,
+    update_challenge_compute_costs,
     update_challenge_results_cache,
-    update_compute_costs_and_storage_size,
 )
 from grandchallenge.invoices.models import PaymentStatusChoices
 from tests.evaluation_tests.factories import EvaluationFactory, PhaseFactory
@@ -164,7 +164,7 @@ def test_challenge_budget_alert_email(settings):
 
     evaluation.utilization.compute_cost_euro_millicents = 500000
     evaluation.utilization.save()
-    update_compute_costs_and_storage_size()
+    update_challenge_compute_costs()
 
     # Budget alert threshold not exceeded
     assert len(mail.outbox) == 0
@@ -175,7 +175,7 @@ def test_challenge_budget_alert_email(settings):
     )
     evaluation.utilization.compute_cost_euro_millicents = 300000
     evaluation.utilization.save()
-    update_compute_costs_and_storage_size()
+    update_challenge_compute_costs()
 
     # Budget alert threshold exceeded
     assert len(mail.outbox) == 3
@@ -205,7 +205,7 @@ def test_challenge_budget_alert_email(settings):
     )
     evaluation.utilization.compute_cost_euro_millicents = 100000
     evaluation.utilization.save()
-    update_compute_costs_and_storage_size()
+    update_challenge_compute_costs()
 
     # Next budget alert threshold not exceeded
     assert len(mail.outbox) == 0
@@ -216,7 +216,7 @@ def test_challenge_budget_alert_email(settings):
     )
     evaluation.utilization.compute_cost_euro_millicents = 1
     evaluation.utilization.save()
-    update_compute_costs_and_storage_size()
+    update_challenge_compute_costs()
 
     # Next budget alert threshold exceeded
     assert len(mail.outbox) != 0
@@ -252,7 +252,7 @@ def test_challenge_budget_alert_two_thresholds_one_email(settings):
     )
     evaluation.utilization.compute_cost_euro_millicents = 950000
     evaluation.utilization.save()
-    update_compute_costs_and_storage_size()
+    update_challenge_compute_costs()
 
     # Two budget alert thresholds exceeded, alert only sent for last one.
     assert len(mail.outbox) == 3
@@ -279,7 +279,7 @@ def test_challenge_budget_alert_no_budget():
     evaluation.utilization.compute_cost_euro_millicents = 1
     evaluation.utilization.save()
     assert len(mail.outbox) == 0
-    update_compute_costs_and_storage_size()
+    update_challenge_compute_costs()
     assert len(mail.outbox) != 0
     assert "Budget Consumed Alert" in mail.outbox[0].subject