Coverage for backend/django/core/auxiliary/services/ml_column_mapping_updates.py: 88%
142 statements
« prev ^ index » next coverage.py v7.10.7, created at 2026-07-22 05:22 +0000
1from typing import Any, TypedDict
3from django.db import transaction
4from rest_framework.exceptions import ValidationError
6from common.config_types import PropertyType as PropertyTypeObj
7from core.auxiliary.enums.generalEnums import PropertyType
8from core.auxiliary.models.Flowsheet import Flowsheet
9from core.auxiliary.models.MLColumnMapping import (
10 MLColumnMapping,
11 custom_property_port_index,
12)
13from core.auxiliary.models.MLModel import MLModel
14from core.auxiliary.models.PropertyInfo import PropertyInfo
15from core.auxiliary.models.PropertyValue import PropertyValue
16from flowsheetInternals.unitops.config.config_methods import get_property_fields
19class ColumnMapping(TypedDict):
20 portIndex: int
21 propertyKey: str
22 column: str
25PROPERTY_INFO_SNAPSHOT_FIELDS = [
26 "set_id",
27 "type",
28 "unitType",
29 "unit",
30 "key",
31 "displayName",
32 "index",
33 "managed",
34 "managed_source",
35 "can_edit",
36 "can_edit_formula",
37 "can_delete",
38 "formula_incomplete",
39 "formula_incomplete_reason",
40]
43def _snapshot_mapping(mapping: MLColumnMapping) -> dict[str, Any]:
44 property_info = mapping.propertyInfo
45 property_value = property_info.values.first() if property_info else None
46 property_info_snapshot = None
48 if mapping.portIndex == custom_property_port_index and property_info:
49 property_info_snapshot = {
50 field: getattr(property_info, field)
51 for field in PROPERTY_INFO_SNAPSHOT_FIELDS
52 }
53 property_info_snapshot["value"] = property_value.value if property_value else None
54 property_info_snapshot["displayValue"] = (
55 property_value.displayValue if property_value else None
56 )
57 property_info_snapshot["enabled"] = property_value.enabled if property_value else True
59 return {
60 "order": mapping.order,
61 "column": mapping.column,
62 "portIndex": mapping.portIndex,
63 "propertyKey": mapping.propertyKey,
64 "propertyType": mapping.propertyType,
65 "propertyInfo": property_info_snapshot,
66 "propertyInfoId": None
67 if mapping.portIndex == custom_property_port_index
68 else mapping.propertyInfo_id,
69 }
72def build_mapping_update_snapshot(ml_model: MLModel) -> dict[str, Any]:
73 """Backward-compatible wrapper for the generalized ML model snapshot."""
74 return build_ml_model_update_snapshot(ml_model)
77def build_ml_model_update_snapshot(ml_model: MLModel) -> dict[str, Any]:
78 """Capture the current CSV-backed ML model state before a reset/retrain."""
79 mappings = (
80 ml_model.MLColumnMappings.select_related("propertyInfo")
81 .prefetch_related("propertyInfo__values")
82 .order_by("order", "id")
83 )
84 return {
85 "mappings": [_snapshot_mapping(mapping) for mapping in mappings],
86 "surrogate_model": ml_model.surrogate_model,
87 "charts": ml_model.charts,
88 "metrics": ml_model.metrics,
89 "result_state": ml_model.result_state,
90 "test_results_bucket": ml_model.test_results_bucket,
91 "test_results_key": ml_model.test_results_key,
92 "csv_file_name": ml_model.csv_file_name,
93 "csv_bucket": ml_model.csv_bucket,
94 "csv_object_key": ml_model.csv_object_key,
95 "csv_headers": ml_model.csv_headers,
96 "csv_delimiter": ml_model.csv_delimiter,
97 "csv_upload_session_id": ml_model.csv_upload_session_id,
98 "active_step": ml_model.active_step,
99 "return_step": ml_model.return_step,
100 "completed_steps": ml_model.completed_steps or [],
101 "is_resetting": ml_model.is_resetting,
102 "is_updating": ml_model.is_updating,
103 }
106def _ml_model_has_restorable_state(ml_model: MLModel) -> bool:
107 """Return whether the model has meaningful previous state to roll back to."""
108 return bool(
109 ml_model.MLColumnMappings.exists()
110 or ml_model.csv_bucket
111 or ml_model.csv_object_key
112 or ml_model.csv_upload_session_id
113 or ml_model.csv_headers
114 or ml_model.surrogate_model
115 or ml_model.charts
116 or ml_model.metrics
117 or ml_model.test_results_bucket
118 or ml_model.test_results_key
119 )
122def stage_ml_model_update_snapshot(ml_model: MLModel) -> bool:
123 """Stage a rollback snapshot on the model if one is not already pending."""
124 if ml_model.mapping_update_snapshot or not _ml_model_has_restorable_state(ml_model):
125 return False
127 ml_model.mapping_update_snapshot = build_ml_model_update_snapshot(ml_model)
128 return True
131def _mapping_owned_custom_property_ids(ml_model: MLModel) -> list[int]:
132 return list(
133 ml_model.MLColumnMappings.filter(
134 portIndex=custom_property_port_index,
135 propertyInfo__isnull=False,
136 ).values_list("propertyInfo_id", flat=True)
137 )
140def _delete_mapping_owned_custom_properties(property_info_ids: list[int]) -> None:
141 if property_info_ids:
142 PropertyInfo.objects.filter(id__in=property_info_ids).delete()
145def clear_ml_model_column_mappings(ml_model: MLModel) -> None:
146 """Remove an ML model's mappings and the custom properties they own."""
147 mapping_owned_custom_property_ids = _mapping_owned_custom_property_ids(ml_model)
148 ml_model.MLColumnMappings.all().delete()
149 _delete_mapping_owned_custom_properties(mapping_owned_custom_property_ids)
152def _create_mapping_objects(
153 *,
154 flowsheet: Flowsheet,
155 ml_model: MLModel,
156 inlet_mappings: list[ColumnMapping],
157 outlet_mappings: list[ColumnMapping],
158) -> None:
159 if not ml_model.simulationObject: 159 ↛ 160line 159 didn't jump to line 160 because the condition on line 159 was never true
160 raise ValidationError({"model": "ML model must be attached to a simulation object."})
161 if ml_model.flowsheet_state.flowsheet_id != flowsheet.id: 161 ↛ 162line 161 didn't jump to line 162 because the condition on line 161 was never true
162 raise ValidationError({"model": "ML model does not belong to the requested flowsheet."})
164 flowsheet_state = ml_model.flowsheet_state
165 property_set = ml_model.simulationObject.properties
166 bulk_create_objects = []
167 mappings = [
168 (inlet_mappings, PropertyType.InletProperty),
169 (outlet_mappings, PropertyType.OutletProperty),
170 ]
171 header_order_lookup = {
172 header: index for index, header in enumerate(ml_model.csv_headers or [])
173 }
174 next_fallback_order = len(header_order_lookup)
176 for mappings_list, property_type in mappings:
177 for mapping in mappings_list:
178 property_key = mapping.get("propertyKey")
179 port_index = mapping.get("portIndex")
180 column = mapping.get("column")
181 property_info = None
183 if property_key is None or port_index is None or column is None: 183 ↛ 184line 183 didn't jump to line 184 because the condition on line 183 was never true
184 raise ValidationError(
185 {"mappings": "Each mapping requires propertyKey, portIndex, and column."}
186 )
188 if port_index == custom_property_port_index:
189 prop = PropertyTypeObj(
190 displayName=property_key,
191 value="",
192 unitType="ratio",
193 type="numeric",
194 )
195 fields = get_property_fields(property_key, prop, property_set, 0)
196 value = fields.pop("value")
197 property_info = PropertyInfo.objects.create(
198 **fields,
199 flowsheet_state=flowsheet_state,
200 )
201 PropertyValue.objects.create(
202 value=value,
203 property=property_info,
204 enabled=property_type == PropertyType.InletProperty,
205 flowsheet_state=flowsheet_state,
206 )
208 order = header_order_lookup.get(column)
209 if order is None:
210 order = next_fallback_order
211 next_fallback_order += 1
213 bulk_create_objects.append(
214 MLColumnMapping(
215 order=order,
216 model=ml_model,
217 column=column,
218 portIndex=port_index,
219 propertyKey=property_key,
220 propertyType=property_type,
221 propertyInfo=property_info,
222 flowsheet_state=flowsheet_state,
223 )
224 )
226 MLColumnMapping.objects.bulk_create(bulk_create_objects)
229def create_column_mapping(
230 flowsheet: Flowsheet,
231 model: int,
232 inlet_mappings: list[ColumnMapping],
233 outlet_mappings: list[ColumnMapping],
234) -> None:
235 """Create first-time column mappings for an ML model."""
236 ml_model = MLModel.objects.select_related(
237 "simulationObject",
238 "flowsheet_state__flowsheet",
239 ).get(id=model)
240 _create_mapping_objects(
241 flowsheet=flowsheet,
242 ml_model=ml_model,
243 inlet_mappings=inlet_mappings,
244 outlet_mappings=outlet_mappings,
245 )
246 ml_model.return_step = ml_model.active_step
247 ml_model.active_step = 2
248 if (
249 ml_model.surrogate_model
250 and not ml_model.csv_bucket
251 and not ml_model.csv_object_key
252 ):
253 ml_model.result_state = MLModel.ResultState.READY_WITHOUT_DIAGNOSTICS
254 else:
255 ml_model.result_state = MLModel.ResultState.PENDING
256 ml_model.save(update_fields=["active_step", "return_step", "result_state"])
259@transaction.atomic
260def replace_column_mapping(
261 *,
262 flowsheet: Flowsheet,
263 model: int,
264 inlet_mappings: list[ColumnMapping],
265 outlet_mappings: list[ColumnMapping],
266) -> None:
267 """Replace existing mappings and stage a rollback snapshot for retraining."""
268 ml_model = MLModel.objects.select_for_update().get(id=model)
270 if not ml_model.csv_bucket or not ml_model.csv_object_key:
271 raise ValidationError({"model": "Only CSV-trained ML models can update mappings."})
273 if ml_model.mapping_update_snapshot: 273 ↛ 274line 273 didn't jump to line 274 because the condition on line 273 was never true
274 raise ValidationError({"model": "A mapping update is already pending for this model."})
276 if not ml_model.MLColumnMappings.exists():
277 raise ValidationError({"model": "Column mappings must exist before they can be updated."})
279 ml_model.mapping_update_snapshot = build_ml_model_update_snapshot(ml_model)
280 clear_ml_model_column_mappings(ml_model)
281 _create_mapping_objects(
282 flowsheet=flowsheet,
283 ml_model=ml_model,
284 inlet_mappings=inlet_mappings,
285 outlet_mappings=outlet_mappings,
286 )
288 ml_model.surrogate_model = {}
289 ml_model.charts = []
290 ml_model.metrics = []
291 ml_model.result_state = MLModel.ResultState.PENDING
292 ml_model.test_results_bucket = ""
293 ml_model.test_results_key = ""
294 ml_model.return_step = ml_model.active_step
295 ml_model.active_step = 2
296 ml_model.completed_steps = [0, 1, 2]
297 ml_model.save(
298 update_fields=[
299 "mapping_update_snapshot",
300 "surrogate_model",
301 "charts",
302 "metrics",
303 "result_state",
304 "test_results_bucket",
305 "test_results_key",
306 "active_step",
307 "return_step",
308 "completed_steps",
309 ]
310 )
313def _restore_mapping_from_snapshot(
314 *,
315 ml_model: MLModel,
316 mapping: dict[str, Any],
317) -> None:
318 property_info = None
319 property_info_snapshot = mapping.get("propertyInfo")
321 if mapping["portIndex"] == custom_property_port_index and property_info_snapshot:
322 fields = {
323 field: property_info_snapshot[field]
324 for field in PROPERTY_INFO_SNAPSHOT_FIELDS
325 if field in property_info_snapshot
326 }
327 # Snapshots created before property-set ownership was persisted do not
328 # contain set_id. ML custom properties belong to the model's property
329 # set, so retain that ownership when restoring those older snapshots.
330 fields["set_id"] = (
331 property_info_snapshot.get("set_id")
332 or ml_model.simulationObject.properties_id
333 )
334 property_info = PropertyInfo.objects.create(
335 **fields,
336 flowsheet_state=ml_model.flowsheet_state,
337 )
338 PropertyValue.objects.create(
339 value=property_info_snapshot.get("value"),
340 displayValue=property_info_snapshot.get("displayValue"),
341 enabled=property_info_snapshot.get("enabled", True),
342 property=property_info,
343 flowsheet_state=ml_model.flowsheet_state,
344 )
345 elif mapping.get("propertyInfoId"): 345 ↛ 346line 345 didn't jump to line 346 because the condition on line 345 was never true
346 property_info = PropertyInfo.objects.filter(id=mapping["propertyInfoId"]).first()
348 MLColumnMapping.objects.create(
349 order=mapping["order"],
350 model=ml_model,
351 column=mapping["column"],
352 portIndex=mapping["portIndex"],
353 propertyKey=mapping["propertyKey"],
354 propertyType=mapping["propertyType"],
355 propertyInfo=property_info,
356 flowsheet_state=ml_model.flowsheet_state,
357 )
360@transaction.atomic
361def restore_ml_model_update_snapshot(ml_model: MLModel) -> bool:
362 """Restore a failed reset/retrain if a rollback snapshot is present."""
363 ml_model = MLModel.objects.select_for_update().get(id=ml_model.id)
364 snapshot = ml_model.mapping_update_snapshot
365 if not snapshot:
366 return False
368 clear_ml_model_column_mappings(ml_model)
369 for mapping in snapshot.get("mappings", []):
370 _restore_mapping_from_snapshot(
371 ml_model=ml_model,
372 mapping=mapping,
373 )
375 snapshot_result_state = snapshot.get("result_state")
376 if snapshot_result_state not in MLModel.ResultState.values: 376 ↛ 377line 376 didn't jump to line 377 because the condition on line 376 was never true
377 if snapshot.get("charts") or snapshot.get("metrics"):
378 snapshot_result_state = MLModel.ResultState.READY
379 elif (
380 snapshot.get("surrogate_model")
381 and not snapshot.get("csv_bucket")
382 and not snapshot.get("csv_object_key")
383 ):
384 snapshot_result_state = MLModel.ResultState.READY_WITHOUT_DIAGNOSTICS
385 else:
386 snapshot_result_state = MLModel.ResultState.PENDING
388 MLModel.objects.filter(id=ml_model.id).update(
389 mapping_update_snapshot=None,
390 surrogate_model=snapshot["surrogate_model"],
391 charts=snapshot["charts"],
392 metrics=snapshot["metrics"],
393 result_state=snapshot_result_state,
394 test_results_bucket=snapshot["test_results_bucket"],
395 test_results_key=snapshot["test_results_key"],
396 csv_file_name=snapshot.get("csv_file_name", ml_model.csv_file_name),
397 csv_bucket=snapshot.get("csv_bucket", ml_model.csv_bucket),
398 csv_object_key=snapshot.get("csv_object_key", ml_model.csv_object_key),
399 csv_headers=snapshot.get("csv_headers", ml_model.csv_headers),
400 csv_delimiter=snapshot.get("csv_delimiter", ml_model.csv_delimiter),
401 csv_upload_session_id=snapshot.get(
402 "csv_upload_session_id",
403 ml_model.csv_upload_session_id,
404 ),
405 active_step=snapshot["active_step"],
406 return_step=snapshot["return_step"],
407 completed_steps=snapshot["completed_steps"],
408 is_resetting=snapshot.get("is_resetting", False),
409 is_updating=snapshot.get("is_updating", False),
410 )
411 return True
414def restore_mapping_update_snapshot(ml_model: MLModel) -> bool:
415 """Backward-compatible wrapper for restoring the full ML model snapshot."""
416 return restore_ml_model_update_snapshot(ml_model)
419def clear_mapping_update_snapshot(model_id: int) -> None:
420 MLModel.objects.filter(id=model_id).update(mapping_update_snapshot=None)