Cost logs for API Deployments

backend/usage_v2/helper.py

@@ -1,6 +1,6 @@
 import logging
 from datetime import datetime
-from typing import Any, Optional
+from typing import Any
 
 from django.db.models import QuerySet, Sum
 from rest_framework.exceptions import APIException
@@ -65,46 +65,6 @@ def get_aggregated_token_count(run_id: str) -> dict:
             logger.error(f"An unexpected error occurred for run_id {run_id}: {str(e)}")
             raise APIException("Error while aggregating token counts")
 
-    @staticmethod
-    def get_aggregated_cost(execution_id: str) -> Optional[float]:
-        """Retrieve aggregated cost for the given execution_id.
-
-        Args:
-            execution_id (str): The identifier for the
-            total cost of a particular execution.
-
-        Returns:
-        Optional[float]: The total cost in dollars if available, else None.
-
-        Raises:
-            APIException: For unexpected errors during database operations.
-        """
-        try:
-            # Aggregate the cost for the given execution_id
-            total_cost = Usage.objects.filter(execution_id=execution_id).aggregate(
-                cost_in_dollars=Sum(UsageKeys.COST_IN_DOLLARS)
-            )[UsageKeys.COST_IN_DOLLARS]
-
-            logger.debug(
-                f"Cost aggregated successfully for execution_id: {execution_id}"
-                f", Total cost: {total_cost}"
-            )
-
-            return total_cost
-
-        except Usage.DoesNotExist:
-            # Handle the case where no usage data is found for the given execution_id
-            logger.warning(
-                f"Usage data not found for the specified execution_id: {execution_id}"
-            )
-            return None
-        except Exception as e:
-            # Handle any other exceptions that might occur during the execution
-            logger.error(
-                f"An unexpected error occurred for execution {execution_id}: {str(e)}"
-            )
-            raise APIException("Error while aggregating cost")
-
     @staticmethod
     def aggregate_usage_metrics(queryset: QuerySet) -> dict[str, Any]:
         """

backend/workflow_manager/workflow_v2/execution.py

@@ -290,15 +290,24 @@ def publish_initial_workflow_logs(self, total_files: int) -> None:
         )
 
     def publish_final_workflow_logs(
-        self, total_files: int, successful_files: int, failed_files: int
+        self,
+        total_files: int,
+        successful_files: int,
+        failed_files: int,
+        execution_id: str,
     ) -> None:
         """Publishes the final logs for the workflow.
 
         Returns:
             None
         """
+        self.publish_average_cost_log(
+            execution_id=execution_id, total_files=successful_files
+        )
+
         # To not associate final logs with a file execution
         self.file_execution_id = None
+
         self.publish_update_log(LogState.END_WORKFLOW, "1", LogComponent.STATUS_BAR)
         self.publish_update_log(
             LogState.SUCCESS, "Executed successfully", LogComponent.WORKFLOW
@@ -311,25 +320,24 @@ def publish_final_workflow_logs(
     def publish_average_cost_log(self, execution_id, total_files):
 
         try:
-            total_cost = UsageHelper.get_aggregated_cost(execution_id)
+            execution: WorkflowExecution = WorkflowExecution.objects.get(
+                pk=execution_id
+            )
+
+            total_cost = execution.get_aggregated_usage_cost
             average_cost = round(total_cost / total_files, 5)
+
             self.publish_log(
                 message=(
                     f"The average cost per file for execution '{execution_id}' "
                     f"is '${average_cost:}'. Total cost: '${total_cost:}'"
                 )
             )
         except TypeError as e:
-            self.publish_log(
-                message=(
-                    f"Unable to calculate cost for execution '{execution_id}'. "
-                    f"Cost data may be unavailable or incomplete."
-                )
-            )
-
-            logger.error(
-                f"Error calculating cost for execution '{execution_id}': "
-                f"{str(e)}. Total cost: {total_cost}, Total files: {total_files}"
+            logger.warning(
+                f"Error calculating cost for execution '{execution_id}' "
+                f"of '{total_files}' files : "
+                f"{str(e)}.\nContinuing execution"
             )
 
     def log_total_cost_per_file(self, run_id, file_name):

backend/workflow_manager/workflow_v2/models/execution.py

@@ -6,8 +6,11 @@
 from api_v2.models import APIDeployment
 from django.core.exceptions import ObjectDoesNotExist
 from django.db import models
+from django.db.models import Sum
 from pipeline_v2.models import Pipeline
 from tags.models import Tag
+from usage_v2.constants import UsageKeys
+from usage_v2.models import Usage
 from utils.common_utils import CommonUtils
 from utils.models.base_model import BaseModel
 from workflow_manager.workflow_v2.enums import ExecutionStatus
@@ -157,6 +160,45 @@ def pretty_execution_time(self) -> str:
         )
         return str(timedelta(seconds=time_in_secs)).split(".")[0]
 
+    @property
+    def get_aggregated_usage_cost(self) -> Optional[float]:
+        """Retrieve aggregated cost for the given execution_id.
+
+        Args:
+            execution_id (str): The identifier for the
+            total cost of a particular execution.
+
+        Returns:
+        Optional[float]: The total cost in dollars if available, else None.
+
+        Raises:
+            APIException: For unexpected errors during database operations.
+        """
+        try:
+            # Aggregate the cost for the given execution_id
+            total_cost = Usage.objects.filter(execution_id=self.id).aggregate(
+                cost_in_dollars=Sum(UsageKeys.COST_IN_DOLLARS)
+            )[UsageKeys.COST_IN_DOLLARS]
+
+            logger.debug(
+                f"Cost aggregated successfully for execution_id: {self.id}"
+                f", Total cost: {total_cost}"
+            )
+
+            return total_cost
+
+        except Usage.DoesNotExist:
+            # Handle the case where no usage data is found for the given execution_id
+            logger.warning(
+                f"Usage data not found for the specified execution_id: {self.id}"
+            )
+            return None
+        except Exception as e:
+            logger.warning(
+                f"An unexpected error occurred for execution {self.id}: {str(e)}"
+            )
+            return None
+
     def __str__(self) -> str:
         return (
             f"Workflow execution: {self.id} ("

backend/workflow_manager/workflow_v2/workflow_helper.py

@@ -213,14 +213,12 @@ def process_input_files(
             execution_service.update_execution(ExecutionStatus.COMPLETED)
 
         workflow_execution = execution_service.get_execution_instance()
-        execution_service.publish_average_cost_log(
-            execution_id=workflow_execution.id, total_files=successful_files
-        )
 
         execution_service.publish_final_workflow_logs(
             total_files=total_files,
             successful_files=successful_files,
             failed_files=failed_files,
+            execution_id=workflow_execution.id,
         )
 
         return workflow_execution