Coverage for backend/django/idaes_factory/endpoints.py: 84%
609 statements
« prev ^ index » next coverage.py v7.10.7, created at 2026-07-22 05:22 +0000
1import json
2import logging
3import os
4from typing import TypedDict
5from common.models.general import TaskPayload
6import requests
7import traceback
9from django.db import IntegrityError, transaction
10from django.utils import timezone
11from opentelemetry import trace
12from rest_framework.exceptions import ValidationError
13from rest_framework.response import Response
14from pydantic import JsonValue
16from CoreRoot import settings
17from authentication.user.models import User
18from ahuora_builder_types.payloads.build_state_request_schema import (
19 BuildStateCompletionPayload,
20 BuildStateRequestContext,
21 BuildStateRequestSchema,
22)
23from ahuora_builder_types.payloads.solve_request_schema import (
24 IdaesSolveRequestPayload,
25 IdaesSolveCompletionPayload,
26 MultiSolvePayload,
27)
28from common.models.idaes.payloads.solve_request_schema import CompletionStatus
29from ahuora_builder_types.payloads.ml_request_schema import (
30 MLTrainRequestPayload,
31 MLTrainingCompletionPayload,
32)
33from core.auxiliary.models.MLModel import MLModel
34from core.auxiliary.models.UploadSession import UploadSessionPurpose
35from core.auxiliary.services.csv_lifecycle import (
36 completed_csv_lifecycle_ttl,
37 expires_at_from_ttl,
38)
39from core.auxiliary.services.object_storage.s3 import schedule_object_expiration
40from common.models.notifications.payloads import (
41 BuildStateCompletedPayload,
42 TaskCompletedPayload,
43 NotificationServiceMessageType,
44 NotificationServiceMessage,
45)
46from core.auxiliary.enums.generalEnums import TaskStatus
47from core.auxiliary.models.Task import Task, TaskSolveMode, TaskType
48from core.auxiliary.models.BuildStateRequestVersion import BuildStateRequestVersion
49from core.auxiliary.models.PropertySet import PropertySet
50from core.auxiliary.models.PropertyValue import PropertyValue
51from core.auxiliary.models.Flowsheet import Flowsheet
52from core.auxiliary.serializers import TaskSerializer
53from common.services.task_cancellation_state import cache_cancelled_tasks
54from core.auxiliary.services.ml_column_mapping_updates import (
55 restore_ml_model_update_snapshot,
56)
57from core.auxiliary.services.solve_completion_email import (
58 queue_solve_completion_email_for_task,
59)
60from core.auxiliary.services.parameter_sweep import validate_parameter_sweep_solve_ready
61from core.auxiliary.services.result_summary.cache import (
62 ResultSummaryCacheInvalidationReason,
63 build_result_summary_cache_for_successful_mss_task,
64 invalidate_result_summary_cache_for_scenario,
65)
66from Economics.scheduling.events import (
67 schedule_scenario_solve_completed,
68 schedule_scenario_solve_started,
69)
70from core.auxiliary.models.FlowsheetState import FlowsheetRevisionKind
71from core.auxiliary.services.flowsheet_states import ActiveTasksError, create_revision
72from core.exceptions import DetailedException
73from .idaes_factory import (
74 IdaesFactory,
75 save_all_initial_values,
76 store_properties_schema,
77)
78from flowsheetInternals.unitops.models.SimulationObject import SimulationObject
79from flowsheetInternals.unitops.services.edit_operations.recorder import (
80 tracked_bulk_update,
81)
82from flowsheetInternals.graphicData.models.groupingModel import Grouping
83from .adapters.stream_properties import serialise_stream
84from .adapters.property_package_adapter import PropertyPackageAdapter
85from .idaes_factory_context import IdaesFactoryContext
86from core.auxiliary.models.Scenario import (
87 Scenario,
88 ScenarioTabTypeEnum,
89 SOLVE_TIMEOUT_DEFAULT_SECONDS,
90)
91from common.services import messaging
92from diagnostics.methods.update_diagnostic_result import update_diagnostics_results
94logger = logging.getLogger(__name__)
95tracer = trace.get_tracer(settings.OPEN_TELEMETRY_TRACER_NAME)
96PARENT_TASK_CANCEL_CHECK_BATCH_SIZE = 25
99def _task_targets_current_state(task: Task) -> bool:
100 """Return whether task output may still mutate the active aggregate."""
102 return (
103 task.flowsheet_state_id is not None
104 and task.flowsheet.current_state_id == task.flowsheet_state_id
105 )
108def _create_auto_revision_after_solve(
109 task: Task,
110 *,
111 solve_index: int | None = None,
112) -> None:
113 """Save a qualifying solve revision without risking solved-value rollback."""
115 if (
116 task.solve_mode != TaskSolveMode.SINGLE_STEADY_STATE
117 or task.parent_id is not None
118 or solve_index is not None
119 or task.children.exists()
120 or not task.flowsheet.auto_snapshot_after_single_solve
121 or not _task_targets_current_state(task)
122 ):
123 return
124 scenario_id = (task.debug or {}).get("scenario_id")
125 if ( 125 ↛ 134line 125 didn't jump to line 134 because the condition on line 125 was never true
126 scenario_id is not None
127 and not Scenario.objects.filter(
128 pk=scenario_id,
129 flowsheet_state_id=task.flowsheet_state_id,
130 state_name=ScenarioTabTypeEnum.SteadyState,
131 enable_dynamics=False,
132 ).exists()
133 ):
134 return
135 try:
136 # This savepoint contains every revision write. Catching outside it
137 # leaves successful solve writes in the surrounding transaction intact.
138 with transaction.atomic():
139 create_revision(
140 flowsheet_id=task.flowsheet_id,
141 kind=FlowsheetRevisionKind.AUTO_SOLVE,
142 created_by=task.creator,
143 )
144 except ActiveTasksError:
145 logger.info(
146 "Auto revision skipped because active tasks remain task_id=%s.",
147 task.pk,
148 )
149 flowsheet_id = task.flowsheet_id
150 task_id = task.pk
151 transaction.on_commit(
152 lambda: messaging.send_flowsheet_notification_message(
153 flowsheet_id,
154 {"task_id": task_id},
155 NotificationServiceMessageType.AUTO_REVISION_SKIPPED_ACTIVE_TASKS,
156 ),
157 robust=True,
158 )
159 return
160 except Exception:
161 logger.exception(
162 "Auto revision failed after successful solve task_id=%s.",
163 task.pk,
164 )
165 return
168def _resolve_solve_timeout_seconds(
169 scenario: Scenario | None,
170 *,
171 parent_task_id: int | None,
172) -> int | None:
173 """Return the solve timeout that should be forwarded to one IDAES solve.
175 MSS parent tasks are orchestration-only and are not forwarded to the IDAES
176 service as solve requests. Their child solves still inherit the scenario
177 timeout, while direct solves fall back to the platform default
178 when no scenario was selected.
179 """
180 if scenario is None:
181 return SOLVE_TIMEOUT_DEFAULT_SECONDS
182 return scenario.solve_timeout_seconds
185class IdaesServiceRequestException(Exception):
186 def __init__(self, message: str) -> None:
187 super().__init__(message)
188 self.message = message
191class SolveFlowsheetError(DetailedException):
192 pass
195class ResponseType(TypedDict, total=False):
196 status: str
197 error: dict | None
198 log: str | None
199 debug: dict | None
202def idaes_service_request(endpoint: str, data: JsonValue) -> JsonValue:
203 """Send a JSON payload to the configured IDAES service endpoint.
205 Args:
206 endpoint: Relative path of the IDAES service endpoint to call.
207 data: Serialised payload that conforms to the endpoint schema.
209 Returns:
210 Parsed JSON response returned by the IDAES service.
212 Raises:
213 IdaesServiceRequestException: If the service responds with a non-200 status.
214 """
215 url = (os.getenv("IDAES_SERVICE_URL") or "http://localhost:8080") + "/" + endpoint
216 result = requests.post(url, json=data)
217 if result.status_code != 200:
218 raise IdaesServiceRequestException(result.json())
220 return result.json()
223@tracer.start_as_current_span("send_flowsheet_solve_request")
224def _solve_flowsheet_request(
225 task_id: int,
226 built_factory: IdaesFactory,
227 parent_task_id: int | None = None,
228 perform_diagnostics: bool = False,
229 high_priority: bool = False,
230):
231 """Queue an IDAES solve request for the provided flowsheet build.
233 Args:
234 task_id: Identifier of the task tracking the solve request.
235 built_factory: Fully built factory containing flowsheet data to solve.
236 perform_diagnostics: Whether to request diagnostic output from IDAES.
237 high_priority: Whether the message should be prioritised over normal solves.
239 Raises:
240 SolveFlowsheetError: If the message cannot be dispatched to the queue.
241 """
243 try:
244 solve_timeout_seconds = _resolve_solve_timeout_seconds(
245 built_factory.scenario,
246 parent_task_id=parent_task_id,
247 )
248 idaes_payload = IdaesSolveRequestPayload(
249 flowsheet=built_factory.flowsheet,
250 solve_index=built_factory.solve_index,
251 scenario_id=(built_factory.scenario.id if built_factory.scenario else None),
252 task_id=task_id,
253 parent_task_id=parent_task_id,
254 perform_diagnostics=perform_diagnostics,
255 solve_timeout_seconds=solve_timeout_seconds,
256 )
258 messaging.send_idaes_solve_message(idaes_payload, high_priority=high_priority)
259 except Exception as e:
260 raise SolveFlowsheetError(e, "idaes_factory_solve_message")
263def _mark_task_dispatched_to_idaes(task: Task) -> None:
264 """Persist that a solve task was successfully queued to the IDAES service."""
265 task.debug = {
266 **(task.debug or {}),
267 "idaes_dispatched": True,
268 }
269 task.save(update_fields=["debug"])
272def start_flowsheet_solve_event(
273 flowsheet_id: int,
274 group_id: int,
275 user: User,
276 scenario: Scenario = None,
277 perform_diagnostics: bool = False,
278) -> Response:
279 """Start a single solve for the given flowsheet and return the tracking task.
281 Args:
282 flowsheet_id: Identifier of the flowsheet that should be solved.
283 user: User initiating the solve request.
284 scenario: Optional scenario providing context for the solve.
285 perform_diagnostics: Whether the solve should run with diagnostic output enabled.
287 Returns:
288 REST response containing the serialised `Task` used to track the solve.
289 """
290 if ( 290 ↛ 296line 290 didn't jump to line 296 because the condition on line 290 was never true
291 scenario
292 and scenario.state_name != ScenarioTabTypeEnum.SteadyState
293 and not scenario.dataRows.exists()
294 ):
295 # if not steady state solves and no data rows exist, cannot proceed
296 return Response(
297 status=400, data=f"No data was provided for {scenario.state_name} scenario."
298 )
300 try:
301 # Remove all previously created DynamicResults for this scenario
302 scenario.solutions.all().delete()
303 except:
304 # there's no solution yet, that's fine
305 pass
307 solve_mode = (
308 TaskSolveMode.DYNAMIC
309 if scenario is not None
310 and (
311 scenario.enable_dynamics
312 or scenario.state_name == ScenarioTabTypeEnum.Dynamic
313 )
314 else TaskSolveMode.SINGLE_STEADY_STATE
315 )
316 expected_state_id = (
317 scenario.flowsheet_state_id
318 if scenario is not None
319 else Grouping.objects.values_list("flowsheet_state_id", flat=True).get(
320 pk=group_id
321 )
322 )
323 solve_task = Task.create(
324 user,
325 flowsheet_id,
326 status=TaskStatus.Pending,
327 solve_mode=solve_mode,
328 expected_flowsheet_state_id=expected_state_id,
329 save=True,
330 )
331 # Persist the user's intent alongside the task so downstream consumers
332 # can tell whether this solve was launched with diagnostics enabled.
333 debug: dict[str, JsonValue] = {
334 **(solve_task.debug or {}),
335 "perform_diagnostics": bool(perform_diagnostics),
336 "solve_timeout_seconds": (
337 scenario.solve_timeout_seconds
338 if scenario is not None
339 else SOLVE_TIMEOUT_DEFAULT_SECONDS
340 ),
341 }
342 if scenario is not None:
343 debug["scenario_id"] = scenario.id
344 solve_task.debug = debug
345 solve_task.save(update_fields=["debug"])
347 try:
348 factory = IdaesFactory(group_id=group_id, scenario=scenario)
349 factory.build()
351 # We send single solve requests as high priority to ensure
352 # they are not blocked by large multi-solve requests.
353 _solve_flowsheet_request(
354 solve_task.id,
355 factory,
356 perform_diagnostics=perform_diagnostics,
357 high_priority=True,
358 )
359 _mark_task_dispatched_to_idaes(solve_task)
360 except DetailedException as e:
361 solve_task.set_failure_with_exception(e, save=True)
363 task_serializer = TaskSerializer(solve_task)
365 return Response(task_serializer.data, status=200)
368def start_multi_steady_state_solve_event(
369 flowsheet_id: int, user: User, scenario: Scenario
370) -> Response:
371 """Kick off a multi steady-state solve and return the parent tracking task.
373 Args:
374 flowsheet_id: Identifier of the flowsheet being solved.
375 user: User who requested the multi-solve.
376 scenario: Scenario containing the steady-state configurations to solve.
378 Returns:
379 REST response containing the parent `Task` that aggregates child solves.
380 """
381 if not scenario.dataRows.exists(): # empty rows 381 ↛ 382line 381 didn't jump to line 382 because the condition on line 381 was never true
382 return Response(
383 status=400, data="No data was provided for multi steady-state scenario."
384 )
385 try:
386 validate_parameter_sweep_solve_ready(scenario)
387 except ValidationError as exc:
388 return Response(status=400, data=exc.detail)
390 # Remove all previously created DynamicResults for this scenario
391 scenario.solutions.all().delete()
392 schedule_scenario_solve_started(scenario_id=scenario.id)
393 invalidate_result_summary_cache_for_scenario(
394 scenario_id=scenario.id,
395 reason=ResultSummaryCacheInvalidationReason.mss_solve_started,
396 )
397 solve_iterations = scenario.dataRows.count()
399 with transaction.atomic():
400 parent_task = Task.create_parent_task(
401 creator=user,
402 flowsheet_id=flowsheet_id,
403 scheduled_tasks=solve_iterations,
404 status=TaskStatus.Running,
405 solve_mode=TaskSolveMode.MULTI_STEADY_STATE,
406 expected_flowsheet_state_id=scenario.flowsheet_state_id,
407 )
408 parent_task.debug = {"scenario_id": scenario.id}
409 parent_task.save(update_fields=["debug"])
410 child_tasks = [
411 Task(
412 creator=user,
413 flowsheet_id=flowsheet_id,
414 flowsheet_state_id=parent_task.flowsheet_state_id,
415 parent=parent_task,
416 status=TaskStatus.Pending,
417 solve_mode=TaskSolveMode.MULTI_STEADY_STATE,
418 )
419 for _ in range(solve_iterations)
420 ]
421 Task.objects.bulk_create(child_tasks)
423 messaging.send_dispatch_multi_solve_message(
424 MultiSolvePayload(task_id=parent_task.id, scenario_id=scenario.id)
425 )
427 return Response(TaskSerializer(parent_task).data, status=200)
430def dispatch_multi_solves(parent_task_id: int, scenario_id: int):
431 """Build and dispatch queued steady-state solves for each child task.
433 Args:
434 parent_task_id: Identifier of the parent multi-solve task.
435 scenario_id: Scenario containing the data rows to iterate through.
436 """
437 parent_task = Task.objects.get(id=parent_task_id)
438 if parent_task.status in ( 438 ↛ 444line 438 didn't jump to line 444 because the condition on line 438 was never true
439 TaskStatus.Completed,
440 TaskStatus.Failed,
441 TaskStatus.Cancelling,
442 TaskStatus.Cancelled,
443 ):
444 return
445 try:
446 scenario = Scenario.objects.get(id=scenario_id)
447 except Scenario.DoesNotExist:
448 if not _task_targets_current_state(parent_task): 448 ↛ 451line 448 didn't jump to line 451 because the condition on line 448 was always true
449 cancel_idaes_solve(parent_task_id)
450 return
451 raise
453 rootGroup = scenario.flowsheet_state.root_grouping
454 factory = IdaesFactory(
455 group_id=rootGroup.id,
456 scenario=scenario,
457 )
459 child_tasks = list(parent_task.children.order_by("start_time"))
460 last_child_index = len(child_tasks) - 1
462 for solve_index, task in enumerate(child_tasks):
463 should_check_parent_status = (
464 solve_index % PARENT_TASK_CANCEL_CHECK_BATCH_SIZE == 0
465 or solve_index == last_child_index
466 )
467 if should_check_parent_status:
468 parent_task.refresh_from_db(fields=["status"])
469 if parent_task.status in (TaskStatus.Cancelling, TaskStatus.Cancelled):
470 logger.info(
471 "Stopping multi-solve dispatch for parent task %s because it is now %s.",
472 parent_task.id,
473 parent_task.status,
474 )
475 break
477 try:
478 factory.clear_flowsheet()
479 factory.use_with_solve_index(solve_index)
480 factory.build()
482 _solve_flowsheet_request(
483 task.id,
484 factory,
485 parent_task_id=parent_task.id,
486 )
487 _mark_task_dispatched_to_idaes(task)
488 except DetailedException as e:
489 task.refresh_from_db(fields=["status", "flowsheet_state"])
490 if task.status == TaskStatus.Cancelling or not _task_targets_current_state( 490 ↛ 500line 490 didn't jump to line 500 because the condition on line 490 was always true
491 task
492 ):
493 task.status = TaskStatus.Cancelled
494 task.completed_time = timezone.now()
495 task.save(update_fields=["status", "completed_time"])
496 Task.increment_cancelled_children_for_parent(parent_task.pk)
497 _send_task_cancelled_notification(task)
498 mark_parent_cancelled(task, scenario_id=scenario.id)
499 break
500 task.set_failure_with_exception(exception=e, save=True)
501 parent_completed = parent_task.update_status_from_child(task)
502 if parent_completed:
503 queue_solve_completion_email_for_task(
504 parent_task, scenario_id=scenario.id
505 )
507 flowsheet_messages = [
508 NotificationServiceMessage(
509 data=TaskSerializer(task).data,
510 message_type=NotificationServiceMessageType.TASK_UPDATED,
511 )
512 for task in [task, parent_task]
513 ]
515 messaging.send_flowsheet_notification_messages(
516 parent_task.flowsheet_id, flowsheet_messages
517 )
520def start_ml_training_event(
521 csv_bucket: str,
522 csv_key: str,
523 csv_delimiter: str | None,
524 input_labels: list[str],
525 output_labels: list[str],
526 user: User,
527 flowsheet_id: int,
528 model_type: str,
529 model_id: int | None = None,
530):
531 """Queue an asynchronous machine-learning training job for the given dataset.
533 Args:
534 csv_bucket: Bucket containing the training CSV.
535 csv_key: Object key containing the training CSV.
536 csv_delimiter: Optional delimiter for the CSV object.
537 input_labels: Names of the input features.
538 output_labels: Names of the predicted outputs.
539 user: User requesting the training run.
540 flowsheet_id: Flowsheet the training run is associated with.
542 Returns:
543 REST response containing the serialised `Task` for the training job.
544 """
545 expected_state_id = (
546 MLModel.objects.values_list("flowsheet_state_id", flat=True).get(pk=model_id)
547 if model_id is not None
548 else None
549 )
550 training_task = Task.create(
551 user,
552 flowsheet_id,
553 task_type=TaskType.ML_TRAINING,
554 status=TaskStatus.Pending,
555 expected_flowsheet_state_id=expected_state_id,
556 save=True,
557 )
558 if model_id is not None: 558 ↛ 562line 558 didn't jump to line 562 because the condition on line 558 was always true
559 MLModel.objects.filter(id=model_id).update(
560 result_state=MLModel.ResultState.TRAINING
561 )
562 training_task.debug = {"model_id": model_id, "model_type": model_type}
563 training_task.save(update_fields=["debug"])
565 try:
566 payload = MLTrainRequestPayload(
567 csv_bucket=csv_bucket,
568 csv_key=csv_key,
569 csv_delimiter=csv_delimiter,
570 input_labels=input_labels,
571 output_labels=output_labels,
572 task_id=training_task.id,
573 model_type=model_type,
574 )
575 logger.warning("payload: %s", payload)
576 messaging.send_ml_training_message(payload)
578 except DetailedException as e:
579 training_task.set_failure_with_exception(e, save=True)
580 if model_id is not None: 580 ↛ 591line 580 didn't jump to line 591 because the condition on line 580 was always true
581 ml_model = MLModel.objects.filter(id=model_id).first()
582 if ml_model: 582 ↛ 591line 582 didn't jump to line 591 because the condition on line 582 was always true
583 restored = restore_ml_model_update_snapshot(ml_model)
584 if not restored:
585 ml_model.return_step = ml_model.active_step
586 ml_model.active_step = 3
587 ml_model.result_state = MLModel.ResultState.FAILED
588 ml_model.save(
589 update_fields=["active_step", "return_step", "result_state"]
590 )
591 _send_task_notifications(training_task)
593 task_serializer = TaskSerializer(training_task)
594 return Response(task_serializer.data, status=200)
597def _send_task_notifications(task: Task, scenario_id: int | None = None):
598 """Broadcast task completion or status updates to interested flowsheet clients.
600 Args:
601 task: Task whose status change should be pushed to subscribers.
602 """
603 flowsheet_messages = []
605 # If this is a child task, update the parent task status
606 if task.parent:
607 parent_completed = task.parent.update_status_from_child(task)
608 if parent_completed and task.parent.task_type == TaskType.IDAES_SOLVE:
609 queue_solve_completion_email_for_task(task.parent, scenario_id=scenario_id)
610 if scenario_id is not None: 610 ↛ 620line 610 didn't jump to line 620 because the condition on line 610 was always true
611 schedule_scenario_solve_completed(scenario_id=scenario_id)
612 parent_task_id = task.parent_id
613 transaction.on_commit(
614 lambda: build_result_summary_cache_for_successful_mss_task(
615 parent_task_id=parent_task_id,
616 scenario_id=scenario_id,
617 )
618 )
620 message_type = (
621 NotificationServiceMessageType.TASK_COMPLETED
622 if task.parent.status == TaskStatus.Completed
623 else NotificationServiceMessageType.TASK_UPDATED
624 )
626 flowsheet_messages.append(
627 NotificationServiceMessage(
628 data=TaskSerializer(task.parent).data, message_type=message_type
629 )
630 )
631 else:
632 if task.task_type == TaskType.IDAES_SOLVE:
633 queue_solve_completion_email_for_task(task, scenario_id=scenario_id)
635 flowsheet_messages.append(
636 NotificationServiceMessage(
637 data=TaskSerializer(task).data,
638 message_type=NotificationServiceMessageType.TASK_COMPLETED,
639 )
640 )
642 messaging.send_flowsheet_notification_messages(
643 task.flowsheet_id, flowsheet_messages
644 )
647def _send_task_cancelled_notification(task: Task):
648 """Broadcast that a task settled in the cancelled state."""
649 messaging.send_flowsheet_notification_messages(
650 task.flowsheet_id,
651 [
652 NotificationServiceMessage(
653 data=TaskSerializer(task).data,
654 message_type=NotificationServiceMessageType.TASK_CANCELLED,
655 )
656 ],
657 )
660def process_idaes_solve_response(solve_response: IdaesSolveCompletionPayload):
661 """Persist the outcome of a completed IDAES solve and notify listeners.
663 Args:
664 solve_response: Payload describing the finished solve result.
665 """
666 # Use a transaction to ensure that either everything succeeds or nothing does
667 should_mark_parent_cancelled = False
668 with transaction.atomic():
669 try:
670 task = (
671 Task.objects.select_related("flowsheet")
672 .select_for_update()
673 .get(id=solve_response.task_id)
674 )
675 except Task.DoesNotExist:
676 logger.info(
677 "Discarding solve completion for missing task_id=%s.",
678 solve_response.task_id,
679 )
680 return
682 # Silently ignore if the task has already been marked as completed.
683 # This allows us to simulate exactly-once delivery semantics (only process
684 # a finished task once).
685 if task.status == TaskStatus.Completed or task.status == TaskStatus.Cancelled: 685 ↛ 686line 685 didn't jump to line 686 because the condition on line 685 was never true
686 return
688 # If cancellation won the race before the completion payload arrived,
689 # keep the task cancelled rather than reviving it back to completed/failed.
690 if task.status == TaskStatus.Cancelling or not _task_targets_current_state(
691 task
692 ):
693 task.status = TaskStatus.Cancelled
694 task.completed_time = timezone.now()
695 task.log = solve_response.log
696 task.debug = {
697 **(task.debug or {}),
698 "timing": solve_response.timing or {},
699 }
700 task.save(update_fields=["status", "completed_time", "log", "debug"])
701 if task.parent_id: 701 ↛ 702line 701 didn't jump to line 702 because the condition on line 701 was never true
702 Task.increment_cancelled_children_for_parent(task.parent_id)
703 should_mark_parent_cancelled = True
704 if not task.parent_id: 704 ↛ 755line 704 didn't jump to line 755
705 queue_solve_completion_email_for_task(
706 task, scenario_id=solve_response.scenario_id
707 )
708 else:
709 task.completed_time = timezone.now()
710 task.log = solve_response.log
711 task.debug = {
712 **(task.debug or {}),
713 "timing": solve_response.timing or {},
714 }
716 if solve_response.status == CompletionStatus.SUCCESS:
717 task.status = TaskStatus.Completed
718 else:
719 task.status = TaskStatus.Failed
720 task.error = {
721 "message": solve_response.error["message"],
722 "cause": "idaes_service_request",
723 "traceback": solve_response.traceback,
724 }
726 task.save(
727 update_fields=["status", "completed_time", "log", "debug", "error"]
728 )
730 # Save the solved flowsheet values
731 if task.status == TaskStatus.Completed:
732 store_properties_schema(
733 solve_response.flowsheet.properties,
734 task.flowsheet_state_id,
735 solve_response.scenario_id,
736 solve_response.solve_index,
737 )
739 # For now, only save initial values for single and dynamic solves, not MSS
740 # In future, we may need some more complex logic to handle MSS initial values
741 if solve_response.solve_index is None:
742 save_all_initial_values(solve_response.flowsheet.initial_values)
743 if task.flowsheet_state_id is not None: 743 ↛ 748line 743 didn't jump to line 748 because the condition on line 743 was always true
744 update_diagnostics_results(
745 task.flowsheet_state,
746 solve_response.unit_diagnostics,
747 )
748 _create_auto_revision_after_solve(
749 task,
750 solve_index=solve_response.solve_index,
751 )
753 _send_task_notifications(task, scenario_id=solve_response.scenario_id)
755 if should_mark_parent_cancelled:
756 _send_task_cancelled_notification(task)
758 if solve_response.scenario_id is None:
759 mark_parent_cancelled(task)
760 else:
761 mark_parent_cancelled(task, scenario_id=solve_response.scenario_id)
764def mark_parent_cancelled(task: Task, scenario_id: int | None = None):
765 """Resolve a task or parent task to ``cancelled`` once cancellation has won.
767 A parent may legitimately end with zero explicitly cancelled children if the
768 last in-flight solves complete before their kill signals land. In that case
769 an already-``cancelling`` parent is still resolved to ``cancelled`` once all
770 children have reached any terminal state.
771 """
772 with transaction.atomic():
773 try:
774 parent = (
775 Task.objects.select_for_update().get(id=task.parent_id)
776 if task.parent_id
777 else None
778 )
779 except Task.DoesNotExist:
780 logger.info(
781 "Discarding parent cancellation update for missing parent task_id=%s.",
782 task.parent_id,
783 )
784 return
786 if parent is None:
787 parent = task
789 has_children = parent.children.exists()
790 has_in_flight_children = parent.children.filter(
791 status__in=[
792 TaskStatus.Running,
793 TaskStatus.Pending,
794 TaskStatus.Cancelling,
795 ]
796 ).exists()
798 has_cancelled_children = parent.children.filter(
799 status=TaskStatus.Cancelled
800 ).exists()
802 # When a cancelling parent has children, resolve it to cancelled once all
803 # children have reached any terminal state. This covers races where child
804 # solves finish before their kill signals land, so none end up explicitly
805 # marked as cancelled even though the parent cancellation should still win.
806 should_cancel_parent = not has_in_flight_children and (
807 has_cancelled_children
808 or (has_children and parent.status == TaskStatus.Cancelling)
809 )
811 if should_cancel_parent:
812 parent.status = TaskStatus.Cancelled
813 parent.completed_time = timezone.now()
814 parent.save(update_fields=["status", "completed_time"])
816 messaging.send_flowsheet_notification_message(
817 parent.flowsheet_id,
818 TaskSerializer(parent).data,
819 NotificationServiceMessageType.TASK_CANCELLED,
820 )
821 queue_solve_completion_email_for_task(parent, scenario_id=scenario_id)
824def process_failed_idaes_solve_response(solve_response: IdaesSolveCompletionPayload):
825 """Handle final failure notifications for solves that could not be processed.
827 Args:
828 solve_response: Completion payload received from the dead-letter queue.
829 """
830 # Use a transaction to ensure that either everything succeeds or nothing does
831 should_mark_parent_cancelled = False
832 with transaction.atomic():
833 try:
834 task = Task.objects.select_for_update().get(id=solve_response.task_id)
835 except Task.DoesNotExist:
836 logger.info(
837 "Discarding failed solve completion for missing task_id=%s.",
838 solve_response.task_id,
839 )
840 return
842 # Silently ignore if the task has already been marked as failed.
843 # This allows us to simulate exactly-once delivery semantics (only process
844 # a failed task once). Our dead letter queue is configured in "at least once" delivery mode.
845 if task.status == TaskStatus.Failed or task.status == TaskStatus.Cancelled: 845 ↛ 846line 845 didn't jump to line 846 because the condition on line 845 was never true
846 return
848 if task.status == TaskStatus.Cancelling or not _task_targets_current_state( 848 ↛ 864line 848 didn't jump to line 864 because the condition on line 848 was always true
849 task
850 ):
851 task.status = TaskStatus.Cancelled
852 task.completed_time = timezone.now()
853 task.log = solve_response.log
854 task.save(update_fields=["status", "completed_time", "log"])
855 if task.parent_id: 855 ↛ 856line 855 didn't jump to line 856 because the condition on line 855 was never true
856 Task.increment_cancelled_children_for_parent(task.parent_id)
857 should_mark_parent_cancelled = True
858 if not task.parent_id: 858 ↛ 878line 858 didn't jump to line 878
859 queue_solve_completion_email_for_task(
860 task,
861 scenario_id=getattr(solve_response, "scenario_id", None),
862 )
863 else:
864 task.completed_time = timezone.now()
865 task.log = solve_response.log
867 task.status = TaskStatus.Failed
868 task.error = {
869 "message": f"Internal server error: several attempts to process finished solve failed. {json.dumps(solve_response.error)}"
870 }
871 task.save()
873 _send_task_notifications(
874 task,
875 scenario_id=getattr(solve_response, "scenario_id", None),
876 )
878 if should_mark_parent_cancelled: 878 ↛ exitline 878 didn't return from function 'process_failed_idaes_solve_response' because the condition on line 878 was always true
879 _send_task_cancelled_notification(task)
880 scenario_id = getattr(solve_response, "scenario_id", None)
881 if scenario_id is None: 881 ↛ 884line 881 didn't jump to line 884 because the condition on line 881 was always true
882 mark_parent_cancelled(task)
883 else:
884 mark_parent_cancelled(task, scenario_id=scenario_id)
887def process_ml_training_response(ml_training_response: MLTrainingCompletionPayload):
888 """Persist the result of a machine-learning training job and send updates.
890 Args:
891 ml_training_response: Completion payload returned by the ML service.
892 """
893 was_cancelled = False
894 with transaction.atomic():
895 try:
896 task = Task.objects.select_for_update().get(id=ml_training_response.task_id)
897 except Task.DoesNotExist:
898 logger.info(
899 "Discarding ML training completion for missing task_id=%s.",
900 ml_training_response.task_id,
901 )
902 return
904 # Silently ignore if the task has already been marked as completed.
905 # This allows us to simulate exactly-once delivery semantics (only process
906 # a finished task once).
907 if task.status in { 907 ↛ 912line 907 didn't jump to line 912 because the condition on line 907 was never true
908 TaskStatus.Completed,
909 TaskStatus.Cancelled,
910 TaskStatus.Failed,
911 }:
912 return
914 if task.status == TaskStatus.Cancelling or not _task_targets_current_state( 914 ↛ 917line 914 didn't jump to line 917 because the condition on line 914 was never true
915 task
916 ):
917 task.status = TaskStatus.Cancelled
918 task.completed_time = timezone.now()
919 task.log = ml_training_response.log
920 task.debug = {
921 **(task.debug or {}),
922 "timing": (
923 ml_training_response.json_response.timing
924 if ml_training_response.json_response is not None
925 else {}
926 ),
927 }
928 task.save(update_fields=["status", "completed_time", "log", "debug"])
929 was_cancelled = True
930 else:
931 task.completed_time = timezone.now()
932 task.log = ml_training_response.log
933 existing_debug = task.debug or {}
934 model_id = existing_debug.get("model_id")
936 if ml_training_response.status == "success":
937 if ml_training_response.json_response is None: 937 ↛ 938line 937 didn't jump to line 938 because the condition on line 937 was never true
938 raise ValueError(
939 "Successful ML training payloads must include json_response."
940 )
941 result = ml_training_response.json_response
942 task.status = TaskStatus.Completed
943 task.debug = {
944 **existing_debug,
945 "timing": result.timing,
946 }
947 else:
948 task.status = TaskStatus.Failed
949 task.error = {
950 "message": ml_training_response.error,
951 "traceback": ml_training_response.traceback,
952 }
953 if model_id is not None: 953 ↛ 971line 953 didn't jump to line 971 because the condition on line 953 was always true
954 ml_model = (
955 MLModel.objects.select_for_update().filter(id=model_id).first()
956 )
957 if ml_model:
958 restored = restore_ml_model_update_snapshot(ml_model)
959 if not restored:
960 ml_model.return_step = ml_model.active_step
961 ml_model.active_step = 3
962 ml_model.result_state = MLModel.ResultState.FAILED
963 ml_model.save(
964 update_fields=[
965 "active_step",
966 "return_step",
967 "result_state",
968 ]
969 )
971 task.save(
972 update_fields=["status", "completed_time", "log", "debug", "error"]
973 )
975 if task.status == TaskStatus.Completed and model_id is not None:
976 ml_model = (
977 MLModel.objects.select_for_update().filter(id=model_id).first()
978 )
979 if ml_model:
980 ml_model.surrogate_model = result.surrogate_model
981 ml_model.charts = [chart.model_dump() for chart in result.charts]
982 ml_model.metrics = [
983 metric.model_dump() for metric in result.metrics
984 ]
985 ml_model.result_state = (
986 MLModel.ResultState.READY
987 if ml_model.charts or ml_model.metrics
988 else MLModel.ResultState.READY_WITHOUT_DIAGNOSTICS
989 )
990 ml_model.test_results_bucket = result.test_results_bucket
991 ml_model.test_results_key = result.test_results_key
992 ml_model.completed_steps = [0, 1, 2]
993 ml_model.mapping_update_snapshot = None
994 ml_model.return_step = ml_model.active_step
995 ml_model.active_step = 2
996 ml_model.save(
997 update_fields=[
998 "surrogate_model",
999 "charts",
1000 "metrics",
1001 "result_state",
1002 "test_results_bucket",
1003 "test_results_key",
1004 "completed_steps",
1005 "mapping_update_snapshot",
1006 "active_step",
1007 "return_step",
1008 ]
1009 )
1010 if result.test_results_bucket and result.test_results_key: 1010 ↛ 1021line 1010 didn't jump to line 1021 because the condition on line 1010 was always true
1011 schedule_object_expiration(
1012 bucket=result.test_results_bucket,
1013 key=result.test_results_key,
1014 expires_at=expires_at_from_ttl(
1015 completed_csv_lifecycle_ttl(
1016 UploadSessionPurpose.ML_TRAINING_CSV
1017 ),
1018 reference_time=task.completed_time,
1019 ),
1020 )
1021 _send_task_notifications(task)
1022 if was_cancelled: 1022 ↛ 1023line 1022 didn't jump to line 1023 because the condition on line 1022 was never true
1023 _send_task_cancelled_notification(task)
1026def cancel_idaes_solve(task_id: int):
1027 """Mark an in-flight solve task as cancelled and notify subscribers.
1029 For parent tasks with children (MSS scenarios), child status updates use
1030 queryset-level UPDATEs (O(1) queries regardless of child count) and the
1031 frontend is told to invalidate its child-task cache with a single
1032 ``TASK_CHILDREN_CANCELLED`` notification instead of one message per child.
1034 Args:
1035 task_id: Identifier of the `Task` being cancelled.
1036 """
1037 to_cache_cancel: set[int] = set()
1038 to_send_cancel: list[int] = []
1039 has_children = False
1040 cancelled_at = timezone.now()
1041 with transaction.atomic():
1042 task = Task.objects.select_for_update().get(id=task_id)
1044 # Ignore cancellation request if a final status (e.g. completed or failed) has already been set
1045 if task.status in (
1046 TaskStatus.Completed,
1047 TaskStatus.Failed,
1048 TaskStatus.Cancelled,
1049 TaskStatus.Cancelling,
1050 ):
1051 return
1053 if task.parent is None and task.children.exists():
1054 has_children = True
1056 # Use queryset-level updates to avoid loading potentially
1057 # hundreds of thousands of child Task objects into memory.
1058 # PostgreSQL guarantees atomicity of each UPDATE statement,
1059 # so the WHERE clause filters at execution time.
1060 non_terminal_children = task.children.exclude(
1061 status__in=[
1062 TaskStatus.Completed,
1063 TaskStatus.Failed,
1064 TaskStatus.Cancelled,
1065 TaskStatus.Cancelling,
1066 ]
1067 )
1068 # Immediately cancel pending children; they haven't been dispatched yet.
1069 cancelled_pending_child_count = non_terminal_children.filter(
1070 status=TaskStatus.Pending
1071 ).update(status=TaskStatus.Cancelled, completed_time=cancelled_at)
1072 # Mark remaining non-terminal children (e.g. Running) as Cancelling.
1073 non_terminal_children.exclude(status=TaskStatus.Pending).update(
1074 status=TaskStatus.Cancelling
1075 )
1077 Task.increment_cancelled_children_for_parent(
1078 task.id, cancelled_pending_child_count
1079 )
1081 to_cache_cancel.add(task.id)
1082 to_send_cancel.append(task.id)
1083 task.status = TaskStatus.Cancelling
1084 else:
1085 task.status = TaskStatus.Cancelling
1086 to_cache_cancel.add(task.id)
1087 to_send_cancel.append(task.id)
1089 task.save()
1091 cache_cancelled_tasks(to_cache_cancel)
1092 for cancel_task_id in to_send_cancel:
1093 # For MSS, broadcasting just the parent task ID is enough: every child
1094 # solve consults shared cancellation state for both its own ID and its
1095 # parent ID before starting work.
1096 messaging.send_task_cancel_message(cancel_task_id)
1098 # Send a single parent-level notification; for parent tasks with children,
1099 # include a children-invalidated event so the frontend drops its cached
1100 # child pages and refetches on demand rather than processing one message
1101 # per child task.
1102 notification_messages = [
1103 NotificationServiceMessage(
1104 data=TaskSerializer(task).data,
1105 message_type=NotificationServiceMessageType.TASK_CANCELLING,
1106 )
1107 ]
1108 if has_children:
1109 notification_messages.append(
1110 NotificationServiceMessage(
1111 data=TaskSerializer(task).data,
1112 message_type=NotificationServiceMessageType.TASK_CHILDREN_CANCELLED,
1113 )
1114 )
1115 messaging.send_flowsheet_notification_messages(
1116 task.flowsheet_id,
1117 notification_messages,
1118 )
1119 mark_parent_cancelled(task)
1122def process_cancel_solve_response(cancel_response: TaskPayload):
1123 """Persist a remote cancellation acknowledgement from the IDAES service.
1125 The IDAES service emits this only after it has successfully interrupted the
1126 active solver, so Django can safely settle the task to a terminal
1127 ``cancelled`` state and then check whether the parent can also be finalised.
1128 """
1129 notification_message: NotificationServiceMessage | None = None
1130 # Use a transaction to ensure that either everything succeeds or nothing does
1131 with transaction.atomic():
1132 try:
1133 task = Task.objects.select_for_update().get(id=cancel_response.task_id)
1134 except Task.DoesNotExist:
1135 logger.info(
1136 "Discarding cancel acknowledgement for missing task_id=%s.",
1137 cancel_response.task_id,
1138 )
1139 return
1141 # Timeout expiry can arrive before Django transitions a task to
1142 # `Cancelling`, so accept pending/running states for timeout paths.
1143 allowed_statuses = {TaskStatus.Cancelling}
1144 if cancel_response.timed_out:
1145 allowed_statuses.update((TaskStatus.Pending, TaskStatus.Running))
1147 if task.status not in allowed_statuses:
1148 return
1150 task.status = TaskStatus.Cancelled
1151 task.completed_time = timezone.now()
1152 # Preserve timeout metadata on the task debug payload so downstream
1153 # notifications have the value that actually governed the solve.
1154 solve_timeout_seconds = (task.debug or {}).get(
1155 "solve_timeout_seconds",
1156 SOLVE_TIMEOUT_DEFAULT_SECONDS,
1157 )
1158 task.debug = {
1159 **(task.debug or {}),
1160 "timed_out": bool(cancel_response.timed_out),
1161 "solve_timeout_seconds": solve_timeout_seconds,
1162 }
1163 if cancel_response.timed_out:
1164 timeout_message = f"Solve timed out after {solve_timeout_seconds} seconds."
1165 task.log = f"{task.log}\n{timeout_message}" if task.log else timeout_message
1166 task.save(update_fields=["status", "completed_time", "debug", "log"])
1167 if task.parent_id:
1168 Task.increment_cancelled_children_for_parent(task.parent_id)
1169 notification_message = NotificationServiceMessage(
1170 data=TaskSerializer(task).data,
1171 message_type=NotificationServiceMessageType.TASK_CANCELLED,
1172 )
1174 if notification_message is not None: 1174 ↛ 1180line 1174 didn't jump to line 1180 because the condition on line 1174 was always true
1175 messaging.send_flowsheet_notification_messages(
1176 task.flowsheet_id, [notification_message]
1177 )
1178 queue_solve_completion_email_for_task(task)
1180 mark_parent_cancelled(task)
1183@transaction.atomic
1184def process_build_state_response(build_state_response: BuildStateCompletionPayload):
1185 """
1186 Apply an asynchronous build-state completion payload from IDAES.
1188 Successful responses update property values for the originating flowsheet.
1189 App-level failures are intentionally acknowledged without applying partial
1190 data and then surfaced to frontend subscribers.
1191 """
1192 context = build_state_response.context
1193 if context is None:
1194 logger.warning("Discarding build-state response without request context.")
1195 return
1197 if not _settle_build_state_task_from_completion(build_state_response):
1198 return
1200 if not _is_current_build_state_response(context): 1200 ↛ 1201line 1200 didn't jump to line 1201 because the condition on line 1200 was never true
1201 return
1203 if build_state_response.status == CompletionStatus.SUCCESS:
1204 task_state_id = (
1205 Task.objects.only("flowsheet_state_id")
1206 .get(pk=context.task_id)
1207 .flowsheet_state_id
1208 )
1209 if task_state_id is None: 1209 ↛ 1210line 1209 didn't jump to line 1210 because the condition on line 1209 was never true
1210 return
1211 store_properties_schema(
1212 build_state_response.properties,
1213 task_state_id,
1214 )
1215 _send_build_state_completion_notification(build_state_response)
1216 return
1218 logger.warning(
1219 "Build-state request failed for flowsheet_id=%s stream_id=%s "
1220 "property_set_id=%s: %s",
1221 context.flowsheet_id,
1222 context.stream_id,
1223 context.property_set_id,
1224 build_state_response.error,
1225 )
1226 _send_build_state_completion_notification(build_state_response)
1229def _settle_build_state_task_from_completion(
1230 build_state_response: BuildStateCompletionPayload,
1231) -> bool:
1232 """Record the final status for a tracked build-state request.
1234 Returns ``False`` when a terminal task already exists for this payload. That
1235 lets duplicate completion deliveries remain idempotent, matching solve and
1236 ML task handling.
1237 """
1238 context = build_state_response.context
1239 if context is None: 1239 ↛ 1240line 1239 didn't jump to line 1240 because the condition on line 1239 was never true
1240 return True
1242 with transaction.atomic():
1243 try:
1244 task = Task.objects.select_for_update().get(id=context.task_id)
1245 except Task.DoesNotExist:
1246 logger.info(
1247 "Discarding build-state completion for missing task_id=%s.",
1248 context.task_id,
1249 )
1250 return False
1252 if task.status in ( 1252 ↛ 1257line 1252 didn't jump to line 1257 because the condition on line 1252 was never true
1253 TaskStatus.Completed,
1254 TaskStatus.Failed,
1255 TaskStatus.Cancelled,
1256 ):
1257 return False
1259 if task.status == TaskStatus.Cancelling or not _task_targets_current_state( 1259 ↛ 1262line 1259 didn't jump to line 1262 because the condition on line 1259 was never true
1260 task
1261 ):
1262 task.status = TaskStatus.Cancelled
1263 elif build_state_response.status == CompletionStatus.SUCCESS:
1264 task.status = TaskStatus.Completed
1265 else:
1266 _revert_build_state_task_values(task)
1267 task.status = TaskStatus.Failed
1268 task.error = {
1269 "message": build_state_response.error,
1270 "cause": "build_state_request",
1271 "traceback": build_state_response.traceback,
1272 }
1274 task.completed_time = timezone.now()
1275 task.log = build_state_response.log
1276 task.debug = {
1277 **(task.debug or {}),
1278 "context": context.model_dump(mode="json"),
1279 }
1280 task.save(update_fields=["status", "completed_time", "log", "debug", "error"])
1282 if task.status == TaskStatus.Cancelled: 1282 ↛ 1283line 1282 didn't jump to line 1283 because the condition on line 1282 was never true
1283 _send_task_cancelled_notification(task)
1284 else:
1285 _send_task_notifications(task)
1286 return task.status != TaskStatus.Cancelled
1289def _is_current_build_state_response(context: BuildStateRequestContext) -> bool:
1290 """Return whether a build-state response still matches the latest request."""
1291 return _is_current_build_state_context(context, "build-state response")
1294def _is_current_build_state_context(
1295 context: BuildStateRequestContext,
1296 event_label: str,
1297) -> bool:
1298 """Return whether a build-state event context still matches the latest request."""
1299 try:
1300 latest_request_version = (
1301 BuildStateRequestVersion.objects.only("version")
1302 .get(
1303 property_set_id=context.property_set_id,
1304 )
1305 .version
1306 )
1307 except BuildStateRequestVersion.DoesNotExist:
1308 logger.warning(
1309 "Discarding %s because no request version tracker exists "
1310 "for flowsheet_id=%s stream_id=%s property_set_id=%s request_version=%s.",
1311 event_label,
1312 context.flowsheet_id,
1313 context.stream_id,
1314 context.property_set_id,
1315 context.request_version,
1316 )
1317 return False
1319 if latest_request_version != context.request_version:
1320 logger.info(
1321 "Discarding stale %s for flowsheet_id=%s stream_id=%s "
1322 "property_set_id=%s request_version=%s latest_request_version=%s.",
1323 event_label,
1324 context.flowsheet_id,
1325 context.stream_id,
1326 context.property_set_id,
1327 context.request_version,
1328 latest_request_version,
1329 )
1330 return False
1332 return True
1335def _send_build_state_completion_notification(
1336 build_state_response: BuildStateCompletionPayload,
1337):
1338 """Broadcast a build-state completion update to flowsheet subscribers."""
1339 context = build_state_response.context
1340 if context is None: 1340 ↛ 1341line 1340 didn't jump to line 1341 because the condition on line 1340 was never true
1341 return
1343 payload = BuildStateCompletedPayload(
1344 context=context,
1345 status=build_state_response.status,
1346 properties=build_state_response.properties,
1347 error=build_state_response.error,
1348 failure_kind=(
1349 "idaes-error"
1350 if build_state_response.status == CompletionStatus.ERROR
1351 else None
1352 ),
1353 )
1354 messaging.send_flowsheet_notification_message(
1355 context.flowsheet_id,
1356 payload.model_dump(mode="json"),
1357 NotificationServiceMessageType.BUILD_STATE_COMPLETED,
1358 )
1361def process_build_state_request_dead_letter(
1362 build_state_request: BuildStateRequestSchema,
1363):
1364 """Record that a build-state request failed broker delivery.
1366 Dead letters are transport-level failures, not IDAES state-block solve
1367 failures. This handler avoids applying property data and only emits a
1368 diagnostic notification when the dead-lettered request is still the latest
1369 request.
1370 """
1371 context = build_state_request.context
1372 if context is None:
1373 logger.warning(
1374 "Build-state request reached the dead-letter topic without request context."
1375 )
1376 return
1378 if not _settle_build_state_task_from_delivery_failure(context):
1379 return
1381 if not _is_current_build_state_context(context, "build-state request dead letter"): 1381 ↛ 1382line 1381 didn't jump to line 1382 because the condition on line 1381 was never true
1382 return
1384 logger.warning(
1385 "Build-state request delivery failed for flowsheet_id=%s stream_id=%s "
1386 "property_set_id=%s request_version=%s.",
1387 context.flowsheet_id,
1388 context.stream_id,
1389 context.property_set_id,
1390 context.request_version,
1391 )
1392 _send_build_state_delivery_failure_notification(context)
1395def _settle_build_state_task_from_delivery_failure(
1396 context: BuildStateRequestContext,
1397) -> bool:
1398 """Mark a tracked build-state request as failed when broker delivery fails."""
1399 with transaction.atomic():
1400 try:
1401 task = Task.objects.select_for_update().get(id=context.task_id)
1402 except Task.DoesNotExist:
1403 logger.info(
1404 "Discarding build-state delivery failure for missing task_id=%s.",
1405 context.task_id,
1406 )
1407 return False
1409 if task.status in ( 1409 ↛ 1414line 1409 didn't jump to line 1414 because the condition on line 1409 was never true
1410 TaskStatus.Completed,
1411 TaskStatus.Failed,
1412 TaskStatus.Cancelled,
1413 ):
1414 return False
1416 task.status = (
1417 TaskStatus.Cancelled
1418 if task.status == TaskStatus.Cancelling
1419 or not _task_targets_current_state(task)
1420 else TaskStatus.Failed
1421 )
1422 if task.status == TaskStatus.Failed: 1422 ↛ 1424line 1422 didn't jump to line 1424 because the condition on line 1422 was always true
1423 _revert_build_state_task_values(task)
1424 task.completed_time = timezone.now()
1425 task.error = {
1426 "message": "Build-state request delivery failed.",
1427 "cause": "message_delivery",
1428 "traceback": None,
1429 }
1430 task.debug = {
1431 **(task.debug or {}),
1432 "context": context.model_dump(mode="json"),
1433 }
1434 task.save(update_fields=["status", "completed_time", "error", "debug"])
1436 if task.status == TaskStatus.Cancelled: 1436 ↛ 1437line 1436 didn't jump to line 1437 because the condition on line 1436 was never true
1437 _send_task_cancelled_notification(task)
1438 else:
1439 _send_task_notifications(task)
1440 return task.status != TaskStatus.Cancelled
1443def _send_build_state_delivery_failure_notification(
1444 context: BuildStateRequestContext,
1445):
1446 """Broadcast a silent transport-level build-state request failure event."""
1447 payload = BuildStateCompletedPayload(
1448 context=context,
1449 status="error",
1450 properties=None,
1451 error=None,
1452 failure_kind="delivery-failure",
1453 )
1454 messaging.send_flowsheet_notification_message(
1455 context.flowsheet_id,
1456 payload.model_dump(mode="json"),
1457 NotificationServiceMessageType.BUILD_STATE_COMPLETED,
1458 )
1461def generate_IDAES_python_request(flowsheet_id: int) -> Response:
1462 """Return the flowsheet JSON for the given flowsheet.
1464 The previous behaviour attempted to generate Python source by forwarding
1465 the flowsheet to the IDAES service. Python generation is no longer
1466 supported; this function now always returns the local flowsheet JSON.
1467 """
1468 scenario = None
1469 flowsheet = Flowsheet.objects.get(id=flowsheet_id)
1470 factory = IdaesFactory(
1471 group_id=flowsheet.current_state.root_grouping_id,
1472 scenario=scenario,
1473 require_variables_fixed=False,
1474 )
1475 response_data = ResponseType(status="success", error=None, log=None, debug=None)
1476 try:
1477 factory.build()
1478 data = factory.flowsheet
1479 # use model_dump_json so that pydantic can handle any non-serializable fields instead
1480 # of django's logic which would error out.
1481 return Response(data.model_dump_json(), status=200)
1482 except Exception as e:
1483 response_data["status"] = "error"
1484 response_data["error"] = {
1485 "message": str(e),
1486 "traceback": traceback.format_exc(),
1487 }
1488 return Response(response_data, status=400)
1489 # NOTE: Python generation used to be implemented by forwarding the
1490 # flowsheet to the IDAES service; that pathway has been removed. We
1491 # have already returned above with the JSON representation.
1494class BuildStateSolveError(Exception):
1495 pass
1498def _next_build_state_request_version(property_set: PropertySet) -> int:
1499 """Increment and return the latest build-state request version for a property set."""
1500 try:
1501 return _next_build_state_request_version_locked(property_set)
1502 except IntegrityError:
1503 # A concurrent first request may have created the tracker between the
1504 # get and create inside get_or_create. Retry once and lock the new row.
1505 return _next_build_state_request_version_locked(property_set)
1508def _next_build_state_request_version_locked(property_set: PropertySet) -> int:
1509 """Increment the version while holding a row lock for this property set."""
1510 with transaction.atomic():
1511 tracker, _ = BuildStateRequestVersion.objects.select_for_update().get_or_create(
1512 property_set_id=property_set.id,
1513 defaults={"flowsheet_state": property_set.flowsheet_state},
1514 )
1515 tracker.version += 1
1516 tracker.save(update_fields=["version", "updated_at"])
1517 return tracker.version
1520def state_request_build(
1521 stream: SimulationObject,
1522 *,
1523 task_id: int,
1524) -> BuildStateRequestSchema:
1525 """Build the IDAES service payload for a stream state request."""
1526 ctx = IdaesFactoryContext(stream.flowsheet_state.root_grouping_id)
1527 property_set = stream.properties
1529 port = stream.connectedPorts.get(direction="inlet")
1530 unitop: SimulationObject = port.unitOp
1531 # find the property package key for this port (we have the value, the key of this port)
1532 property_package_ports = unitop.schema.propertyPackagePorts
1533 for key, port_list in property_package_ports.items():
1534 if port.key in port_list:
1535 property_package_key = key
1536 PropertyPackageAdapter(property_package_key).serialise(ctx, unitop)
1538 return BuildStateRequestSchema(
1539 property_package=ctx.property_packages[0],
1540 properties=serialise_stream(ctx, stream, is_inlet=True),
1541 context=BuildStateRequestContext(
1542 flowsheet_id=stream.flowsheet_state.flowsheet_id,
1543 stream_id=stream.id,
1544 task_id=task_id,
1545 property_set_id=property_set.id,
1546 request_version=_next_build_state_request_version(property_set),
1547 ),
1548 )
1551def _create_build_state_task(
1552 stream: SimulationObject,
1553 user: User | int,
1554 rollback_values: dict[int, float] | None = None,
1555) -> Task:
1556 """Create the tracking task used for a build-state request."""
1557 task = Task.create(
1558 user,
1559 stream.flowsheet_state.flowsheet_id,
1560 task_type=TaskType.BUILD_STATE,
1561 status=TaskStatus.Pending,
1562 expected_flowsheet_state_id=stream.flowsheet_state_id,
1563 save=True,
1564 )
1565 task.debug = {
1566 "stream_id": stream.id,
1567 "property_set_id": stream.properties.id,
1568 "rollback_values": (
1569 {str(prop_id): value for prop_id, value in rollback_values.items()}
1570 if rollback_values
1571 else {}
1572 ),
1573 }
1574 task.save(update_fields=["debug"])
1575 return task
1578def _revert_build_state_task_values(task: Task) -> None:
1579 """Restore raw property values captured before the build-state request."""
1580 rollback_values = (task.debug or {}).get("rollback_values")
1581 if not rollback_values:
1582 return
1584 property_values = list(
1585 PropertyValue.objects.filter(
1586 id__in=[int(prop_id) for prop_id in rollback_values]
1587 )
1588 )
1589 for property_value in property_values:
1590 property_value.value = rollback_values[str(property_value.id)]
1592 tracked_bulk_update(PropertyValue.objects, property_values, ["value"])
1595def _fail_build_state_task_before_dispatch(
1596 task: Task,
1597 exception: Exception,
1598 *,
1599 cause: str,
1600) -> None:
1601 """Persist a build-state setup/dispatch failure before re-raising it."""
1602 _revert_build_state_task_values(task)
1603 task.set_failure_with_exception(
1604 DetailedException(exception, source=cause),
1605 save=True,
1606 )
1607 _send_task_notifications(task)
1610def _mark_build_state_task_dispatched(
1611 task: Task,
1612 context: BuildStateRequestContext | None,
1613) -> None:
1614 """Persist dispatch metadata for a queued build-state task."""
1615 task.debug = {
1616 **(task.debug or {}),
1617 "idaes_dispatched": True,
1618 "context": (context.model_dump(mode="json") if context is not None else None),
1619 }
1620 task.save(update_fields=["debug"])
1623def build_state_request(
1624 stream: SimulationObject,
1625 user: User | int,
1626 rollback_values: dict[int, float] | None = None,
1627):
1628 """Queue a state build request for the provided stream.
1630 Args:
1631 stream: Stream object whose inlet properties should be used for the build.
1632 user: User requesting the property update that triggered the build.
1634 Returns:
1635 REST response containing the tracking task. Built properties are applied
1636 later by ``process_build_state_response`` when IDAES publishes the
1637 completion event.
1639 Raises:
1640 BuildStateSolveError: If the request cannot be queued.
1641 Exception: If preparing the payload fails for any reason.
1642 """
1643 task = _create_build_state_task(
1644 stream,
1645 user,
1646 rollback_values=rollback_values,
1647 )
1648 try:
1649 data = state_request_build(stream, task_id=task.id)
1650 messaging.send_idaes_build_state_request_message(data)
1651 _mark_build_state_task_dispatched(task, data.context)
1652 return Response(TaskSerializer(task).data, status=202)
1653 except Exception as e:
1654 _fail_build_state_task_before_dispatch(
1655 task,
1656 e,
1657 cause="build_state_dispatch",
1658 )
1659 raise BuildStateSolveError(str(e))