WIP consolidation task

Author: chrisvanrun

app/grandchallenge/challenges/costs.py

@@ -6,6 +6,8 @@
     AlgorithmModel,
     Job,
 )
+from dateutil.utils import today
+
 from grandchallenge.cases.models import ImageFile
 from grandchallenge.components.models import ComponentInterfaceValue
 from grandchallenge.evaluation.models import (
@@ -20,6 +22,51 @@
 )
 
 
+def get_compute_costs_for_invoice(*, invoice, limit_to_covered_period):
+
+    algorithm_job_utilizations = JobUtilization.objects.filter(
+        challenge=invoice.challenge
+    )
+    job_warm_pool_utilizations = JobWarmPoolUtilization.objects.filter(
+        challenge=invoice.challenge
+    )
+    evaluation_job_utilizations = EvaluationUtilization.objects.filter(
+        challenge=invoice.challenge
+    )
+
+    if limit_to_covered_period:
+        algorithm_job_utilizations = algorithm_job_utilizations.filter(
+            created__lte=invoice.expires_on
+        )
+        job_warm_pool_utilizations = job_warm_pool_utilizations.filter(
+            created__lte=invoice.expires_on
+        )
+        evaluation_job_utilizations = evaluation_job_utilizations.filter(
+            created__lte=invoice.expires_on
+        )
+
+    algorithm_job_costs = algorithm_job_utilizations.aggregate(
+        Sum("compute_cost_euro_millicents")
+    )
+    job_warm_pool_costs = job_warm_pool_utilizations.aggregate(
+        Sum("compute_cost_euro_millicents")
+    )
+    evaluation_costs = evaluation_job_utilizations.aggregate(
+        Sum("compute_cost_euro_millicents")
+    )
+
+    items = [algorithm_job_costs, job_warm_pool_costs, evaluation_costs]
+
+    total = sum(
+        item["compute_cost_euro_millicents__sum"] or 0 for item in items
+    )
+
+    if limit_to_covered_period:
+        total = min(total, invoice.compute_costs_euros * 1000 * 100)
+
+    return total
+
+
 def annotate_compute_costs(*, challenge):
     algorithm_job_utilizations = JobUtilization.objects.filter(
         challenge=challenge

app/grandchallenge/challenges/tasks.py

@@ -6,14 +6,15 @@
 from django.conf import settings
 from django.contrib.auth import get_user_model
 from django.db import transaction
-from django.db.models import Count, Max, Min, Q
+from django.db.models import Count, Max, Min, Q, Sum
 from django.utils.timezone import datetime, now
 from psycopg.errors import LockNotAvailable
 
 from grandchallenge.challenges.costs import (
     annotate_compute_costs,
     annotate_job_duration_and_compute_costs,
     annotate_storage_size,
+    get_compute_costs_for_invoice,
 )
 from grandchallenge.challenges.emails import (
     send_challenge_requests_draft_reminder,
@@ -31,6 +32,7 @@
     acks_late_micro_short_task,
 )
 from grandchallenge.evaluation.models import Evaluation, Phase
+from grandchallenge.invoices.models import Invoice
 
 
 @acks_late_2xlarge_task
@@ -102,17 +104,50 @@ def wrapper(*args, **kwargs):
 
 @acks_late_2xlarge_task
 def update_challenge_compute_costs():
-    for challenge in Challenge.objects.with_available_compute().iterator(
-        chunk_size=1000
-    ):
+    challenge_ids = Invoice.objects.values_list(
+        "challenge_id", flat=True
+    ).distinct()
+
+    for challenge_id in challenge_ids:
         with transaction.atomic():
-            annotate_compute_costs(challenge=challenge)
+            invoices = list(
+                Invoice.objects.filter(challenge_id=challenge_id).order_by(
+                    "expires_on", "created"
+                )  # Todo: add queryset-level filter to only return those invoices with authorized budget
+            )
+
+            covered_compute_costs = 0
+            last_index = len(invoices) - 1
+            for idx, invoice in enumerate(invoices):
+                is_last = idx == last_index
+                compute_costs = get_compute_costs_for_invoice(
+                    invoice=invoice,
+                    limit_to_covered_period=not is_last,
+                )
+                invoice.compute_cost_euro_millicents = (
+                    compute_costs - covered_compute_costs
+                )
+                covered_compute_costs = compute_costs
 
             @retry_with_backoff((LockNotAvailable,))
-            def save_challenge():
+            def save_invoices():
+                Invoice.objects.bulk_update(
+                    invoices,
+                    ["compute_cost_euro_millicents"],
+                )
+
+                # TODO: remove this once compute_cost_euro_millicents is removed or deprecated
+                challenge = Challenge.objects.with_available_compute().get(
+                    pk=challenge_id
+                )
+                challenge.compute_cost_euro_millicents = (
+                    Invoice.objects.filter(challenge=challenge_id).aggregate(
+                        total=Sum("compute_cost_euro_millicents")
+                    )["total"]
+                )
                 challenge.save(update_fields=("compute_cost_euro_millicents",))
 
-            save_challenge()
+            save_invoices()
 
     for phase in Phase.objects.iterator(chunk_size=1000):
         with transaction.atomic():

app/tests/challenges_tests/test_tasks.py

@@ -1,9 +1,11 @@
 from zoneinfo import ZoneInfo
 
+from django.utils import timezone
 import pytest
 from django.core import mail
-from django.utils.timezone import datetime, timedelta
+from django.utils.timezone import datetime, now, timedelta
 
+from grandchallenge.algorithms.models import Job
 from grandchallenge.challenges.models import (
     Challenge,
     ChallengeRequest,
@@ -15,7 +17,13 @@
     update_challenge_compute_costs,
     update_challenge_results_cache,
 )
-from grandchallenge.invoices.models import PaymentStatusChoices
+from grandchallenge.invoices.models import (
+    PaymentStatusChoices,
+    PaymentTypeChoices,
+)
+from grandchallenge.utilization.models import JobWarmPoolUtilization
+from grandchallenge.utilization.tasks import create_job_warm_pool_utilizations
+from tests.algorithms_tests.factories import AlgorithmJobFactory
 from tests.evaluation_tests.factories import EvaluationFactory, PhaseFactory
 from tests.factories import (
     ChallengeFactory,
@@ -444,3 +452,136 @@ def test_challenge_request_draft_reminder_emails(
         ]
         assert len(user_emails) == 1
         assert requests[i].title in user_emails[0].body
+
+
+@pytest.mark.django_db
+def test_update_challenge_compute_costs_no_utilization():
+    challenge = ChallengeFactory()
+    invoice = InvoiceFactory(challenge=challenge)
+
+    assert invoice.compute_cost_euro_millicents == 0
+    update_challenge_compute_costs()
+    invoice.refresh_from_db()
+    assert invoice.compute_cost_euro_millicents == 0
+
+
+@pytest.mark.django_db
+def test_update_challenge_compute_costs(settings):
+
+    settings.COMPONENTS_DEFAULT_BACKEND = (
+        "tests.utilization_tests.test_tasks.UtilizationExecutor"
+    )
+
+    challenge = ChallengeFactory()
+    invoice = InvoiceFactory(challenge=challenge)
+    phase = PhaseFactory(challenge=challenge)
+
+    evaluation = EvaluationFactory(
+        submission__phase=phase,
+        time_limit=60,
+    )
+    evaluation_utilization = evaluation.evaluation_utilization
+    evaluation_utilization.invoice = invoice
+    evaluation_utilization.compute_cost_euro_millicents = 1
+    evaluation_utilization.save()
+
+    job = AlgorithmJobFactory(
+        status=Job.SUCCESS,
+        use_warm_pool=True,
+        time_limit=60,
+    )
+
+    job_utilization = job.job_utilization
+    job_utilization.phase = phase
+    job_utilization.challenge = challenge
+    job_utilization.invoice = invoice
+    job_utilization.compute_cost_euro_millicents = 2
+    job_utilization.save()
+
+    job2 = AlgorithmJobFactory(
+        status=Job.SUCCESS,
+        use_warm_pool=True,
+        time_limit=60,
+    )
+
+    job2_utilization = job2.job_utilization
+    job2_utilization.phase = phase
+    job2_utilization.challenge = challenge
+    job2_utilization.invoice = invoice
+    job2_utilization.compute_cost_euro_millicents = 4
+    job2_utilization.save()
+
+    assert not JobWarmPoolUtilization.objects.exists()
+    create_job_warm_pool_utilizations()
+    for job_warm_pool_utilization in JobWarmPoolUtilization.objects.all():
+        job_warm_pool_utilization.compute_cost_euro_millicents = 8
+        job_warm_pool_utilization.save()
+
+    assert invoice.compute_cost_euro_millicents == 0
+
+    update_challenge_compute_costs()
+
+    invoice.refresh_from_db()
+    assert invoice.compute_cost_euro_millicents == 1 + 2 + 4 + 8 + 8
+
+
+@pytest.mark.django_db
+def test_update_challenge_compute_costs_expired_utilization():
+
+    challenge = ChallengeFactory()
+
+    invoice_last_year = InvoiceFactory(
+        challenge=challenge,
+        compute_costs_euros=10,
+        payment_status=PaymentStatusChoices.PAID,
+        payment_type=PaymentTypeChoices.PREPAID,
+        expires_on=now().date() - timedelta(days=365),
+    )
+
+    invoice_this_year = InvoiceFactory(
+        challenge=challenge,
+        compute_costs_euros=20,
+        payment_status=PaymentStatusChoices.PAID,
+        payment_type=PaymentTypeChoices.PREPAID,
+        expires_on=now().date(),
+    )
+
+    invoice_later_this_year = InvoiceFactory(
+        challenge=challenge,
+        compute_costs_euros=40,
+        payment_status=PaymentStatusChoices.PAID,
+        payment_type=PaymentTypeChoices.PREPAID,
+        expires_on=now().date() + timedelta(days=365),
+    )
+
+    job_last_year = AlgorithmJobFactory(
+        status=Job.SUCCESS,
+        time_limit=60,
+    )
+    job_utilization = job_last_year.job_utilization
+    job_utilization.challenge = challenge
+    job_utilization.created = now() - timedelta(days=365)
+    job_utilization.compute_cost_euro_millicents = (
+        100 * 1000 * 100  # Note: Huge
+    )
+    job_utilization.save()
+
+    assert invoice_last_year.compute_cost_euro_millicents == 0
+    assert invoice_this_year.compute_cost_euro_millicents == 0
+    assert invoice_later_this_year.compute_cost_euro_millicents == 0
+
+    update_challenge_compute_costs()
+
+    invoice_last_year.refresh_from_db()
+    assert (  # Fully utilized
+        invoice_last_year.compute_cost_euro_millicents == 10 * 1000 * 100
+    )
+    assert (  # Fully utilized
+        invoice_this_year.compute_cost_euro_millicents == 20 * 1000 * 100
+    )
+    assert (  # Utilization from earlier bubble up to this one
+        invoice_later_this_year.compute_cost_euro_millicents == 70 * 1000 * 100
+    )
+
+
+# TODO: add test that checks if we filter on authorized budget when calculating compute costs for invoices