Coverage for backend/django/core/auxiliary/models/Task.py: 86%
108 statements
« prev ^ index » next coverage.py v7.10.7, created at 2026-07-22 05:22 +0000
1"""Task model offering tracking and coordination for asynchronous operations."""
3from typing import Optional
4from django.db import models
5from django.db.models import F
6from django.utils import timezone
7from authentication.user.models import User
8from core.auxiliary.enums.generalEnums import TaskStatus
9from core.auxiliary.models import Flowsheet
10from core.exceptions import DetailedException
11from core.managers import IdentityAccessControlManager, StaleFlowsheetState
14class TaskMeta(models.Model):
15 """Aggregate counts that parent tasks use to monitor child progress."""
17 scheduled_tasks = models.PositiveIntegerField(default=0)
18 failed_tasks = models.PositiveIntegerField(default=0)
19 successful_tasks = models.PositiveIntegerField(default=0)
20 cancelled_tasks = models.PositiveIntegerField(default=0)
22class TaskType(models.TextChoices):
23 """Enum of supported task categories."""
25 ML_TRAINING = 'ML Training'
26 IDAES_SOLVE = 'Solve'
27 BUILD_STATE = 'Build State'
28 CSV_IMPORT_SCENARIO = 'CSV Import Scenario'
29 CSV_IMPORT_PINCH_UTILITIES = 'CSV Import Pinch Utilities'
32class TaskSolveMode(models.TextChoices):
33 """Explicit solve classification used by revision and cancellation logic."""
35 SINGLE_STEADY_STATE = "single_steady_state", "Single steady-state"
36 MULTI_STEADY_STATE = "multi_steady_state", "Multi steady-state"
37 DYNAMIC = "dynamic", "Dynamic"
39class Task(models.Model):
40 """Represents a unit of work for platform computation (e.g. solving, ML), the status of which is observable
41 and cancellable by users."""
43 task_type = models.CharField(max_length=50, choices=TaskType.choices , default=TaskType.IDAES_SOLVE)
44 solve_mode = models.CharField(
45 max_length=24,
46 choices=TaskSolveMode.choices,
47 null=True,
48 blank=True,
49 )
50 creator = models.ForeignKey(User, on_delete=models.CASCADE)
51 status = models.CharField(choices=TaskStatus.choices , default=TaskStatus.Pending)
52 flowsheet = models.ForeignKey(Flowsheet, on_delete=models.CASCADE)
53 flowsheet_state = models.ForeignKey(
54 "FlowsheetState",
55 on_delete=models.SET_NULL,
56 related_name="tasks",
57 null=True,
58 blank=True,
59 )
60 start_time = models.DateTimeField(auto_now_add=True)
61 completed_time = models.DateTimeField(null=True)
62 error = models.JSONField(null=True)
63 debug = models.JSONField(null=True)
64 log = models.TextField(null=True)
65 parent = models.ForeignKey('Task', on_delete=models.CASCADE, null=True, related_name='children')
66 metadata = models.OneToOneField(TaskMeta, on_delete=models.CASCADE, null=True, related_name='task')
68 objects = IdentityAccessControlManager()
70 class Meta:
71 ordering = ['-start_time']
73 @classmethod
74 def create(
75 cls,
76 creator: User | int,
77 flowsheet: Flowsheet | int,
78 task_type: TaskType = TaskType.IDAES_SOLVE,
79 parent: Optional['Task'] = None,
80 status: TaskStatus = TaskStatus.Pending,
81 solve_mode: TaskSolveMode | None = None,
82 expected_flowsheet_state_id: int | None = None,
83 save=False
84 ):
85 """Instantiate a task, optionally persisting immediately.
87 Args:
88 creator: User creating the task or their primary key.
89 flowsheet: Flowsheet associated with the task or its primary key.
90 task_type: Category used to tailor downstream behaviour.
91 parent: Parent task when building a hierarchy of work.
92 status: Initial status to set on the task.
93 solve_mode: Explicit solve classification for IDAES work.
94 expected_flowsheet_state_id: Request-captured state that must still
95 be current when the task is bound.
96 save: When true the task is inserted into the database before returning.
98 Returns:
99 The newly constructed `Task` instance.
100 """
102 flowsheet_id = flowsheet.pk if isinstance(flowsheet, Flowsheet) else flowsheet
103 current_flowsheet = Flowsheet.objects.only("pk", "current_state_id").get(
104 pk=flowsheet_id
105 )
106 if expected_flowsheet_state_id is None:
107 from core.validation import get_current_flowsheet
109 expected_flowsheet_state_id = (
110 (get_current_flowsheet() or {}).get("flowsheet_state")
111 )
112 if (
113 expected_flowsheet_state_id is not None
114 and current_flowsheet.current_state_id != expected_flowsheet_state_id
115 ):
116 raise StaleFlowsheetState()
117 task = Task(
118 creator=creator,
119 flowsheet=current_flowsheet,
120 flowsheet_state_id=current_flowsheet.current_state_id,
121 parent=parent,
122 status=status,
123 task_type=task_type,
124 solve_mode=solve_mode,
125 )
126 if save: 126 ↛ 128line 126 didn't jump to line 128 because the condition on line 126 was always true
127 task.save()
128 return task
130 @classmethod
131 def create_parent_task(
132 cls,
133 creator: User | int,
134 flowsheet_id: int,
135 task_type: TaskType = TaskType.IDAES_SOLVE,
136 scheduled_tasks: int = 0,
137 status: TaskStatus = TaskStatus.Pending,
138 solve_mode: TaskSolveMode | None = None,
139 expected_flowsheet_state_id: int | None = None,
140 ):
141 """Create a parent task seeded with metadata describing child workload."""
142 current_flowsheet = Flowsheet.objects.only("pk", "current_state_id").get(
143 pk=flowsheet_id
144 )
145 if expected_flowsheet_state_id is None:
146 from core.validation import get_current_flowsheet
148 expected_flowsheet_state_id = (
149 (get_current_flowsheet() or {}).get("flowsheet_state")
150 )
151 if ( 151 ↛ 155line 151 didn't jump to line 155 because the condition on line 151 was never true
152 expected_flowsheet_state_id is not None
153 and current_flowsheet.current_state_id != expected_flowsheet_state_id
154 ):
155 raise StaleFlowsheetState()
156 parent_task = Task(
157 creator=creator,
158 flowsheet=current_flowsheet,
159 flowsheet_state_id=current_flowsheet.current_state_id,
160 status=status,
161 task_type=task_type,
162 solve_mode=solve_mode,
163 )
164 parent_task.metadata = TaskMeta.objects.create(scheduled_tasks=scheduled_tasks)
165 parent_task.save()
167 return parent_task
169 def clean(self):
170 """Require optional task provenance to match the stable flowsheet."""
172 super().clean()
173 if (
174 self.flowsheet_state_id is not None
175 and self.flowsheet_state.flowsheet_id != self.flowsheet_id
176 ):
177 from django.core.exceptions import ValidationError
179 raise ValidationError(
180 {"flowsheet_state": "Task state must belong to the task flowsheet."}
181 )
183 def update_status_from_child(self, child_task: 'Task') -> bool:
184 """Update parent completion state in response to a child transition.
186 If a user has already requested cancellation of the parent, preserve the
187 ``cancelling`` status even when the final child settles normally. The
188 caller will then resolve the parent to ``cancelled`` once all in-flight
189 children have finished.
191 Returns ``True`` only when this call transitions the parent itself into
192 ``completed``. Callers use that signal to trigger parent-level follow-up
193 work such as summary notifications.
194 """
196 metadata = self.metadata
197 if child_task.status == TaskStatus.Failed:
198 # Use F expressions to avoid race conditions when multiple children update concurrently.
199 metadata.failed_tasks = F('failed_tasks') + 1
200 metadata.save(update_fields=['failed_tasks'])
201 elif child_task.status == TaskStatus.Completed: 201 ↛ 204line 201 didn't jump to line 204 because the condition on line 201 was always true
202 metadata.successful_tasks = F('successful_tasks') + 1
203 metadata.save(update_fields=['successful_tasks'])
204 elif child_task.status == TaskStatus.Cancelled:
205 metadata.cancelled_tasks = F('cancelled_tasks') + 1
206 metadata.save(update_fields=['cancelled_tasks'])
208 # Refresh to obtain concrete integer values for comparison.
209 metadata.refresh_from_db(fields=['failed_tasks', 'successful_tasks', 'cancelled_tasks'])
211 if (
212 metadata.scheduled_tasks
213 == (metadata.successful_tasks + metadata.failed_tasks + metadata.cancelled_tasks)
214 ):
215 completed_time = timezone.now()
216 transitioned = Task.objects.filter(
217 pk=self.pk,
218 status__in=[TaskStatus.Pending, TaskStatus.Running],
219 flowsheet_state_id=F("flowsheet__current_state_id"),
220 ).update(
221 status=TaskStatus.Completed,
222 completed_time=completed_time,
223 )
224 if transitioned:
225 self.status = TaskStatus.Completed
226 self.completed_time = completed_time
227 return True
228 self.refresh_from_db(fields=["status", "completed_time"])
230 return False
232 @classmethod
233 def increment_cancelled_children_for_parent(cls, parent_task_id: int, count: int = 1):
234 """Increment a parent task's cancelled-child counter in one database update.
236 When count is greater than 1, multiple child cancellations are batched into
237 a single update. Returns early without touching the database when count is
238 zero or negative.
239 """
240 if count <= 0: 240 ↛ 241line 240 didn't jump to line 241 because the condition on line 240 was never true
241 return
243 TaskMeta.objects.filter(task__id=parent_task_id).update(
244 cancelled_tasks=F('cancelled_tasks') + count
245 )
247 @staticmethod
248 def failure_details(exception: Exception | DetailedException) -> dict:
249 """Serialize an exception into the durable task error contract."""
250 detailed_exception = (exception
251 if isinstance(exception, DetailedException)
252 else DetailedException(exception, source="")
253 )
254 return {
255 "message": detailed_exception.message,
256 "cause": detailed_exception.source,
257 "traceback": detailed_exception.traceback
258 }
260 def set_failure_with_exception(self, exception: Exception | DetailedException, save: bool = False):
261 """Record failure details on the task, preserving the stack trace."""
262 self.status = TaskStatus.Failed
263 self.error = self.failure_details(exception)
265 if self.completed_time is None: 265 ↛ 269line 265 didn't jump to line 269 because the condition on line 265 was always true
266 # ensure completed timestamp reflects when failure was logged
267 self.completed_time = timezone.now()
269 if save: 269 ↛ exitline 269 didn't return from function 'set_failure_with_exception' because the condition on line 269 was always true
270 self.save()