Add setting and serialization of exec and invoke durations (#4385)

Closes DIAGNijmegen/rse-roadmap#441

Author: jmsmkn

app/config/settings.py

@@ -896,6 +896,7 @@ def sentry_before_send(event, hint):
     "PAGE_SIZE": 100,
     "UNAUTHENTICATED_USER": "guardian.utils.get_anonymous_user",
     "DEFAULT_SCHEMA_CLASS": "drf_spectacular.openapi.AutoSchema",
+    "DURATION_FORMAT": "iso-8601",
 }
 
 SPECTACULAR_SETTINGS = {

app/grandchallenge/algorithms/admin.py

@@ -202,6 +202,7 @@ class JobAdmin(admin.ModelAdmin):
         "is_complimentary",
         "credits_consumed",
         "time_limit",
+        "exec_duration",
         "requires_gpu_type",
         "requires_memory_gb",
         "use_warm_pool",
@@ -235,6 +236,8 @@ class JobAdmin(admin.ModelAdmin):
         "runtime_metrics",
         "algorithm_interface",
         "time_limit",
+        "exec_duration",
+        "invoke_duration",
         "job_utilization",
         "public",
         "algorithm_model",

app/grandchallenge/algorithms/serializers.py

@@ -32,6 +32,7 @@
     ComponentInterfaceValueSerializer,
     HyperlinkedComponentInterfaceValueSerializer,
 )
+from grandchallenge.core.drf_fields import ISODurationField
 from grandchallenge.core.guardian import filter_by_permission
 from grandchallenge.core.templatetags.remove_whitespace import oxford_comma
 from grandchallenge.hanging_protocols.serializers import (
@@ -143,6 +144,9 @@ class JobSerializer(serializers.ModelSerializer):
         source="algorithm_image.algorithm.view_content", read_only=True
     )
 
+    exec_duration = ISODurationField(read_only=True)
+    invoke_duration = ISODurationField(read_only=True)
+
     class Meta:
         model = Job
         fields = [
@@ -156,6 +160,8 @@ class Meta:
             "hanging_protocol",
             "optional_hanging_protocols",
             "view_content",
+            "exec_duration",
+            "invoke_duration",
         ]
 
 

app/grandchallenge/components/admin.py

@@ -113,13 +113,18 @@ def requeue_jobs(modeladmin, request, queryset):
     jobs = []
 
     for job in queryset:
+        # If adding fields here remember to add them to the bulk_update fields below
         job.status = ComponentJob.RETRY
         job.attempt += 1
-        job.utilization.duration = None
-        job.utilization.save()
+        job.exec_duration = None
+        job.invoke_duration = None
         job.use_warm_pool = False
         job.error_message = ""
         job.detailed_error_message = {}
+
+        job.utilization.duration = None
+        job.utilization.save()
+
         jobs.append(job)
 
         on_commit(job.execute)
@@ -129,6 +134,9 @@ def requeue_jobs(modeladmin, request, queryset):
         fields=[
             "status",
             "attempt",
+            "exec_duration",
+            "invoke_duration",
+            "use_warm_pool",
             "error_message",
             "detailed_error_message",
         ],

app/grandchallenge/components/backends/amazon_sagemaker_base.py

@@ -357,7 +357,7 @@ def _stop_job_boto(self):
     def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
 
-        self.__duration = None
+        self.__utilization_duration = None
         self.__runtime_metrics = {}
 
         self.__sagemaker_client = None
@@ -413,8 +413,8 @@ def _cloudwatch_client(self):
         return self.__cloudwatch_client
 
     @property
-    def duration(self):
-        return self.__duration
+    def utilization_duration(self):
+        return self.__utilization_duration
 
     @property
     def runtime_metrics(self):
@@ -536,10 +536,10 @@ def _set_duration(self, *, event):
                 self._get_start_time(event=event)
             )
             stopped = ms_timestamp_to_datetime(self._get_end_time(event=event))
-            self.__duration = stopped - started
+            self.__utilization_duration = stopped - started
         except TypeError:
             logger.warning("Invalid start or end time, duration undetermined")
-            self.__duration = None
+            self.__utilization_duration = None
 
     def _get_log_stream_name(self, *, data_log=False):
         response = self._logs_client.describe_log_streams(
@@ -718,7 +718,9 @@ def _handle_failed_job(self, *, event):
             "Out of Memory. Please use a larger instance",
         ):
             try:
-                users_process_exit_code = self._get_task_return_code()
+                users_process_exit_code = (
+                    self._get_inference_result().return_code
+                )
             except UncleanExit:
                 users_process_exit_code = None
 

app/grandchallenge/components/backends/base.py

@@ -21,6 +21,7 @@
 import boto3
 import botocore
 import httpx
+import pydantic
 from asgiref.sync import async_to_sync
 from botocore.auth import SigV4Auth
 from botocore.awsrequest import AWSRequest
@@ -288,6 +289,17 @@ class InferenceTask(BaseModel):
     timeout: timedelta
 
 
+class InferenceResult(BaseModel):
+    model_config = ConfigDict(frozen=True)
+
+    pk: str
+    return_code: int
+    exec_duration: timedelta | None
+    invoke_duration: timedelta | None
+    outputs: list[InferenceIO]
+    sagemaker_shim_version: str
+
+
 class Executor(ABC):
     def __init__(
         self,
@@ -313,11 +325,15 @@ def __init__(
             use_warm_pool and settings.COMPONENTS_USE_WARM_POOL
         )
         self._signing_key = signing_key
+        self._algorithm_model = algorithm_model
+        self._ground_truth = ground_truth
+
+        self._exec_duration = None
+        self._invoke_duration = None
         self._stdout = []
         self._stderr = []
+
         self.__s3_client = None
-        self._algorithm_model = algorithm_model
-        self._ground_truth = ground_truth
 
     def provision(self, *, input_civs, input_prefixes):
         # We cannot run everything async as it requires database access.
@@ -387,7 +403,15 @@ def stderr(self):
 
     @property
     @abstractmethod
-    def duration(self): ...
+    def utilization_duration(self): ...
+
+    @property
+    def exec_duration(self):
+        return self._exec_duration
+
+    @property
+    def invoke_duration(self):
+        return self._invoke_duration
 
     @property
     @abstractmethod
@@ -437,12 +461,13 @@ def _max_memory_mb(self):
 
     @property
     def compute_cost_euro_millicents(self):
-        duration = self.duration
-        if duration is None:
+        utilization_duration = self.utilization_duration
+        if utilization_duration is None:
             return None
         else:
             return duration_to_millicents(
-                duration=duration, usd_cents_per_hour=self.usd_cents_per_hour
+                duration=utilization_duration,
+                usd_cents_per_hour=self.usd_cents_per_hour,
             )
 
     @property
@@ -467,7 +492,7 @@ def _invocation_key(self):
         return safe_join(self._invocation_prefix, "invocation.json")
 
     @property
-    def _result_key(self):
+    def _inference_result_key(self):
         return safe_join(
             self._io_prefix, ".sagemaker_shim", "inference_result.json"
         )
@@ -802,11 +827,11 @@ def _get_upload_input_content_task(*, content, key):
             key=key,
         )
 
-    def _get_task_return_code(self):
+    def _get_inference_result(self):
         try:
             response = self._s3_client.get_object(
                 Bucket=settings.COMPONENTS_OUTPUT_BUCKET_NAME,
-                Key=self._result_key,
+                Key=self._inference_result_key,
             )
         except botocore.exceptions.ClientError as error:
             if error.response["Error"]["Code"] == "404":
@@ -834,26 +859,29 @@ def _get_task_return_code(self):
             )
 
         try:
-            result = json.loads(body.decode("utf-8"))
-        except JSONDecodeError:
+            inference_result = InferenceResult.model_validate_json(
+                json_data=body
+            )
+        except pydantic.ValidationError as error:
+            logger.error(error, exc_info=True)
             raise ComponentException(
                 "The invocation request did not return valid json"
             )
 
-        logger.info(f"{result=}")
+        logger.info(f"{inference_result=}")
 
-        if result["pk"] != self._job_id:
+        if inference_result.pk != self._job_id:
             raise RuntimeError("Wrong result key for this job")
 
-        try:
-            return int(result["return_code"])
-        except (KeyError, ValueError):
-            raise ComponentException(
-                "The invocation response object is not valid"
-            )
+        return inference_result
 
     def _handle_completed_job(self):
-        users_process_exit_code = self._get_task_return_code()
+        inference_result = self._get_inference_result()
+
+        self._exec_duration = inference_result.exec_duration
+        self._invoke_duration = inference_result.invoke_duration
+
+        users_process_exit_code = inference_result.return_code
 
         if users_process_exit_code == 0:
             # Job's a good un

app/grandchallenge/components/models.py

@@ -1736,15 +1736,17 @@ def save(self, *args, **kwargs):
         if adding:
             self.create_utilization()
 
-    def update_status(
+    def update_status(  # noqa:C901
         self,
         *,
         status: STATUS_CHOICES,
         stdout: str = "",
         stderr: str = "",
         error_message="",
         detailed_error_message=None,
-        duration=None,
+        utilization_duration=None,
+        exec_duration=None,
+        invoke_duration=None,
         compute_cost_euro_millicents=None,
         runtime_metrics=None,
     ):
@@ -1765,10 +1767,16 @@ def update_status(
                 for key, value in detailed_error_message.items()
             }
 
-        if duration is not None:
-            self.utilization.duration = duration
+        if utilization_duration is not None:
+            self.utilization.duration = utilization_duration
             self.utilization.save(update_fields=["duration"])
 
+        if exec_duration is not None:
+            self.exec_duration = exec_duration
+
+        if invoke_duration is not None:
+            self.invoke_duration = invoke_duration
+
         if compute_cost_euro_millicents is not None:
             self.utilization.compute_cost_euro_millicents = (
                 compute_cost_euro_millicents

app/grandchallenge/components/tasks.py

@@ -857,7 +857,9 @@ def get_update_status_kwargs(*, executor=None):
         return {
             "stdout": executor.stdout,
             "stderr": executor.stderr,
-            "duration": executor.duration,
+            "utilization_duration": executor.utilization_duration,
+            "exec_duration": executor.exec_duration,
+            "invoke_duration": executor.invoke_duration,
             "compute_cost_euro_millicents": executor.compute_cost_euro_millicents,
             "runtime_metrics": executor.runtime_metrics,
         }

app/grandchallenge/core/drf_fields.py

@@ -0,0 +1,8 @@
+from django.utils.duration import duration_iso_string
+from rest_framework.fields import DurationField
+
+
+class ISODurationField(DurationField):
+    # TODO - This functionality will in DRF 3.17, remove when released
+    def to_representation(self, value):
+        return duration_iso_string(value)

app/grandchallenge/evaluation/admin.py

@@ -180,6 +180,7 @@ class EvaluationAdmin(admin.ModelAdmin):
         "created",
         "submission",
         "time_limit",
+        "exec_duration",
         "requires_gpu_type",
         "requires_memory_gb",
         "use_warm_pool",
@@ -221,6 +222,8 @@ class EvaluationAdmin(admin.ModelAdmin):
         "runtime_metrics",
         "claimed_by",
         "ground_truth",
+        "exec_duration",
+        "invoke_duration",
         "evaluation_utilization",
     )
     actions = (requeue_jobs, cancel_jobs, deprovision_jobs)