Make ProcessBuilder overridable and add helper for WorkerRuns with costs

arkindex/process/models.py

@@ -11,12 +11,12 @@ from django.db.models import Exists, F, OuterRef, Q
 from django.urls import reverse
 from django.utils import timezone
 from django.utils.functional import cached_property
+from django.utils.module_loading import import_string
 from enumfields import Enum, EnumField
 
 import pgtrigger
 from arkindex.documents.models import Classification, Element
 from arkindex.ponos.models import FINAL_STATES, STATES_ORDERING, State, Task, token_default
-from arkindex.process.builder import ProcessBuilder
 from arkindex.process.managers import (
     ActivityManager,
     CorpusWorkerVersionManager,
@@ -32,6 +32,10 @@ from arkindex.project.validators import MaxValueValidator
 from arkindex.training.models import ModelVersion, ModelVersionState
 from arkindex.users.models import Role
 
+ProcessBuilder = import_string(
+    getattr(settings, "PROCESS_BUILDER_CLASS", None) or "arkindex.process.builder.ProcessBuilder"
+)
+
 
 def process_max_chunks():
     """
@@ -964,6 +968,16 @@ class WorkerRun(models.Model):
         # we add the WorkerRun ID at the end of the slug
         return f"{self.version.worker.slug}_{str(self.id)[:6]}"
 
+    @property
+    def has_costs(self) -> bool:
+        """
+        Whether running a task created from this WorkerRun could incur costs in a budget
+        """
+        # Use either the CPU or GPU cost depending on how we are running this worker
+        hourly_cost = self.version.worker.cost_gpu_hour if self.use_gpu else self.version.worker.cost_cpu_hour
+        # Only apply the cost per element when the process might have worker activities
+        return hourly_cost > 0 or (self.process.activity_state != ActivityState.Disabled and self.version.worker.cost_1k_elements > 0)
+
     def build_task(self, process, env, import_task_name, elements_path, run=0, chunk=None, workflow_runs=None, active_gpu_agents=False):
         """
         Build the Task that will represent this WorkerRun in ponos using :