feat: save FunctionProfilingStep when event is received

Signed-off-by: Guilhem Barthés <[email protected]>

backend/api/events/sync.py

@@ -22,6 +22,7 @@
 from api.models import DataManager
 from api.models import DataSample
 from api.models import Function
+from api.models import FunctionProfilingStep
 from api.models import Model
 from api.serializers import ChannelOrganizationSerializer
 from api.serializers import ComputePlanSerializer
@@ -430,6 +431,13 @@ def _create_failure_report(data: dict) -> None:
         logger.debug("FailureReport already exists", asset_key=data["asset_key"])
 
 
+def _on_create_profiling_step_event(data: dict) -> None:
+    profiling_step = FunctionProfilingStep(
+        function_key=data["asset_key"], duration=data["profiling_step"]["duration"], step=data["profiling_step"]["step"]
+    )
+    profiling_step.save()
+
+
 EVENT_CALLBACKS = {
     common_pb2.ASSET_COMPUTE_PLAN: {
         event_pb2.EVENT_ASSET_CREATED: _on_create_computeplan_event,
@@ -467,6 +475,9 @@ def _create_failure_report(data: dict) -> None:
     common_pb2.ASSET_FAILURE_REPORT: {
         event_pb2.EVENT_ASSET_CREATED: _on_create_failure_report_event,
     },
+    common_pb2.ASSET_PROFILING_STEP: {
+        event_pb2.EVENT_ASSET_CREATED: _on_create_profiling_step_event,
+    },
 }
 
 

backend/api/migrations/0060_functionprofilingstep.py

@@ -0,0 +1,43 @@
+# Generated by Django 4.2.10 on 2024-04-11 06:55
+
+import django.db.models.deletion
+from django.db import migrations
+from django.db import models
+
+import orchestrator.resources
+
+
+class Migration(migrations.Migration):
+    dependencies = [
+        ("api", "0059_remove_datamanager_type"),
+    ]
+
+    operations = [
+        migrations.CreateModel(
+            name="FunctionProfilingStep",
+            fields=[
+                ("id", models.AutoField(auto_created=True, primary_key=True, serialize=False, verbose_name="ID")),
+                (
+                    "step",
+                    models.CharField(
+                        choices=[
+                            ("BUILD_IMAGE", orchestrator.resources.FunctionProfilingStep["BUILD_IMAGE"]),
+                            ("SAVE_FUNCTION", orchestrator.resources.FunctionProfilingStep["SAVE_FUNCTION"]),
+                        ],
+                        max_length=100,
+                    ),
+                ),
+                ("duration", models.DurationField()),
+                (
+                    "function",
+                    models.ForeignKey(
+                        on_delete=django.db.models.deletion.CASCADE, related_name="profiling_steps", to="api.function"
+                    ),
+                ),
+            ],
+            options={
+                "ordering": ["step"],
+                "unique_together": {("function", "step")},
+            },
+        ),
+    ]

backend/api/models/__init__.py

@@ -13,13 +13,15 @@
 from .model import Model
 from .organization import ChannelOrganization
 from .performance import Performance
+from .task_profiling import FunctionProfilingStep
 from .task_profiling import ProfilingStep
 from .task_profiling import TaskProfiling
 
 __all__ = [
     "Function",
     "FunctionInput",
     "FunctionOutput",
+    "FunctionProfilingStep",
     "ComputePlan",
     "ComputeTask",
     "ComputeTaskOutput",

backend/api/models/task_profiling.py

@@ -1,8 +1,10 @@
 from django.db import models
 
+import orchestrator
+
 
 class ProfilingStep(models.Model):
-    """A profiling step"""
+    """A task profiling step"""
 
     compute_task_profile = models.ForeignKey(
         "TaskProfiling", on_delete=models.CASCADE, related_name="execution_rundown"
@@ -25,3 +27,15 @@ class TaskProfiling(models.Model):
 
     class Meta:
         ordering = ["creation_date"]
+
+
+class FunctionProfilingStep(models.Model):
+    """A function profiling step"""
+
+    function = models.ForeignKey("Function", on_delete=models.CASCADE, related_name="profiling_steps")
+    step = models.CharField(max_length=100, choices=orchestrator.FunctionProfilingStep.to_choices())
+    duration = models.DurationField()
+
+    class Meta:
+        unique_together = (("function", "step"),)
+        ordering = ["step"]

backend/orchestrator/__init__.py

@@ -14,6 +14,7 @@
 from .resources import Function
 from .resources import FunctionInput
 from .resources import FunctionOutput
+from .resources import FunctionProfilingStep
 from .resources import FunctionStatus
 from .resources import InvalidInputAsset
 from .resources import Model
@@ -40,6 +41,7 @@
     "OrcError",
     "FunctionInput",
     "FunctionOutput",
+    "FunctionProfilingStep",
     "FunctionStatus",
     "get_orchestrator_client",
 )