Coverage for backend/django/core/auxiliary/viewsets/MLViewSet.py: 77%
306 statements
« prev ^ index » next coverage.py v7.10.7, created at 2026-07-22 05:22 +0000
1import json
2import logging
3from typing import Any
5from django.http import HttpResponse
6from django.shortcuts import get_object_or_404
7from django.views.decorators.csrf import csrf_exempt
8from django.db import transaction
9from drf_spectacular.utils import OpenApiParameter, OpenApiTypes, extend_schema
10from pydantic import ValidationError as PydanticValidationError
11from rest_framework import serializers, status
12from rest_framework.decorators import action, api_view, authentication_classes, permission_classes
13from authentication.custom_drf_authentication import DaprApiTokenAuthentication
14from rest_framework.exceptions import NotFound, ValidationError
15from rest_framework.parsers import JSONParser
16from rest_framework.permissions import IsAuthenticated
17from rest_framework.response import Response
19from common.models.idaes.payloads.ml_request_schema import MLTrainingCompletionEvent
20from core.auxiliary.models.MLModel import MLModel
21from core.auxiliary.models.MLWizard import train
22from core.auxiliary.serializers.MLModelSerializer import MLModelSerializer
23from core.auxiliary.services.ml_column_mapping_updates import (
24 clear_ml_model_column_mappings,
25 restore_ml_model_update_snapshot,
26)
27from core.auxiliary.services.object_storage.s3 import presign_download_url
28from core.auxiliary.services.uploads import attach_upload_to_ml_model, inspect_upload_session
29from core.viewset import ModelViewSet
30from flowsheetInternals.unitops.models import SimulationObject
31from idaes_factory.endpoints import process_ml_training_response
34logger = logging.getLogger(__name__)
37def _parse_ml_training_completion_event(data) -> MLTrainingCompletionEvent | None:
38 try:
39 return MLTrainingCompletionEvent.model_validate(data)
40 except PydanticValidationError:
41 logger.warning(
42 "Discarding malformed ML training completion event.",
43 exc_info=True,
44 )
45 return None
48class GetCsvHeaderSerializer(serializers.Serializer):
49 headers = serializers.ListField(child=serializers.CharField())
52class OperationMessageSerializer(serializers.Serializer):
53 message = serializers.CharField()
56class CreateSurrogateModelFromColumnSerializer(serializers.Serializer):
57 model = serializers.IntegerField()
60class DownloadTestResultsSerializer(serializers.Serializer):
61 url = serializers.URLField()
64class UploadModelSerializer(serializers.Serializer):
65 json_data = serializers.JSONField()
66 model = serializers.IntegerField()
68 def validate_json_data(self, value):
69 """Require the labels needed to map an imported surrogate to flowsheet ports."""
70 if not isinstance(value, dict): 70 ↛ 71line 70 didn't jump to line 71 because the condition on line 70 was never true
71 raise serializers.ValidationError("The imported model must be a JSON object.")
73 for label_key in ("input_labels", "output_labels"):
74 labels = value.get(label_key)
75 if not isinstance(labels, list) or not labels or not all( 75 ↛ 78line 75 didn't jump to line 78 because the condition on line 75 was never true
76 isinstance(label, str) and label for label in labels
77 ):
78 raise serializers.ValidationError(
79 f"The imported model must include non-empty {label_key}."
80 )
82 return value
85class CreateMLModelSerializer(serializers.Serializer):
86 simulationObject = serializers.IntegerField()
87 surrogate_model = serializers.JSONField(required=False, default=dict)
90class PatchMLModelSerializer(serializers.ModelSerializer):
91 """Allow editors to change only ML model presentation and algorithm choices."""
93 csv_upload_session = serializers.IntegerField(required=False, write_only=True)
94 flowsheet = serializers.IntegerField(required=False, write_only=True)
96 class Meta:
97 model = MLModel
98 fields = ("displayName", "model_type", "csv_upload_session", "flowsheet")
100 def to_internal_value(self, data):
101 """Reject fields outside the PATCH allow-list instead of silently ignoring them."""
102 unknown_fields = set(data.keys()) - set(self.fields)
103 if unknown_fields:
104 raise serializers.ValidationError(
105 {
106 field: "This field may not be updated."
107 for field in sorted(unknown_fields)
108 }
109 )
110 return super().to_internal_value(data)
112 def validate(self, attrs):
113 """Validate request context while keeping model ownership backend-controlled."""
114 flowsheet_id = attrs.pop("flowsheet", None)
115 if (
116 flowsheet_id is not None
117 and flowsheet_id != self.instance.flowsheet_state.flowsheet_id
118 ):
119 raise serializers.ValidationError(
120 {"flowsheet": "This ML model does not belong to that flowsheet."}
121 )
123 if "csv_upload_session" in attrs and len(attrs) > 1: 123 ↛ 124line 123 didn't jump to line 124 because the condition on line 123 was never true
124 raise serializers.ValidationError(
125 "Attach a CSV separately from other ML model changes."
126 )
127 return attrs
130class MLWizardTransitionSerializer(serializers.Serializer):
131 """Describe an allowed user-driven transition through the ML wizard."""
133 model = serializers.IntegerField()
134 transition = serializers.ChoiceField(
135 choices=(
136 "advance_to_mappings",
137 "start_reset",
138 "start_mapping_update",
139 "cancel",
140 "complete",
141 "select_step",
142 )
143 )
144 active_step = serializers.IntegerField(required=False)
146 def validate(self, attrs):
147 """Require a destination only for an explicit wizard navigation request."""
148 is_selecting_step = attrs["transition"] == "select_step"
149 if is_selecting_step and "active_step" not in attrs: 149 ↛ 150line 149 didn't jump to line 150 because the condition on line 149 was never true
150 raise serializers.ValidationError(
151 {"active_step": "An active step is required when selecting a step."}
152 )
153 if not is_selecting_step and "active_step" in attrs: 153 ↛ 154line 153 didn't jump to line 154 because the condition on line 153 was never true
154 raise serializers.ValidationError(
155 {"active_step": "Only select_step may provide an active step."}
156 )
157 return attrs
160class MLViewSet(ModelViewSet):
161 """Manage ML CSV uploads, header lookup, model import, and surrogate training."""
163 serializer_class = MLModelSerializer
164 parser_classes = [JSONParser]
166 def get_queryset(self):
167 return MLModel.objects.all()
169 @staticmethod
170 def _validate_single_model_rule(simulation_object: SimulationObject) -> None:
171 if simulation_object.objectType != "machineLearningBlock":
172 return
173 if MLModel.objects.filter(simulationObject=simulation_object).exists(): 173 ↛ exitline 173 didn't return from function '_validate_single_model_rule' because the condition on line 173 was always true
174 raise ValidationError(
175 {"simulationObject": "machineLearningBlock can only have one ML model at a time."}
176 )
178 @extend_schema(request=PatchMLModelSerializer, responses=MLModelSerializer)
179 def update(self, request, *args, **kwargs) -> Response:
180 """Apply an allow-listed PATCH or attach a completed CSV upload session."""
181 partial = kwargs.pop("partial", False)
182 instance = self.get_object()
183 serializer = PatchMLModelSerializer(
184 instance,
185 data=request.data,
186 partial=partial,
187 )
188 serializer.is_valid(raise_exception=True)
190 upload_session_id = serializer.validated_data.pop("csv_upload_session", None)
191 if upload_session_id is not None:
192 instance = attach_upload_to_ml_model(
193 ml_model=instance,
194 upload_session_id=upload_session_id,
195 user_id=request.user.id,
196 )
197 else:
198 serializer.save()
200 return Response(self.get_serializer(instance).data, status=200)
202 @extend_schema(request=PatchMLModelSerializer, responses=MLModelSerializer)
203 def partial_update(self, request, *args, **kwargs) -> Response:
204 """Expose the same allow-listed contract for the generated PATCH endpoint."""
205 kwargs["partial"] = True
206 return self.update(request, *args, **kwargs)
208 @extend_schema(request=MLWizardTransitionSerializer, responses=MLModelSerializer)
209 @action(detail=False, methods=["post"], url_path="transition-wizard")
210 @transaction.atomic
211 def transition_wizard(self, request):
212 """Apply a validated ML-wizard transition without exposing state fields to PATCH."""
213 serializer = MLWizardTransitionSerializer(data=request.data)
214 serializer.is_valid(raise_exception=True)
215 transition_data = serializer.validated_data
217 try:
218 ml_model = MLModel.objects.select_for_update().get(
219 id=transition_data["model"]
220 )
221 except MLModel.DoesNotExist:
222 raise NotFound({"model": "ML model not found."})
224 transition = transition_data["transition"]
225 if transition == "advance_to_mappings":
226 if not ml_model.csv_bucket or not ml_model.csv_object_key: 226 ↛ 227line 226 didn't jump to line 227 because the condition on line 226 was never true
227 raise ValidationError({"model": "Attach a CSV before mapping columns."})
228 ml_model.active_step = 1
229 ml_model.completed_steps = [0]
230 update_fields = ["active_step", "return_step", "completed_steps"]
231 elif transition == "start_reset": 231 ↛ 232line 231 didn't jump to line 232 because the condition on line 231 was never true
232 ml_model.active_step = 0
233 ml_model.completed_steps = []
234 ml_model.is_resetting = True
235 ml_model.is_updating = False
236 update_fields = [
237 "active_step",
238 "return_step",
239 "completed_steps",
240 "is_resetting",
241 "is_updating",
242 ]
243 elif transition == "start_mapping_update": 243 ↛ 244line 243 didn't jump to line 244 because the condition on line 243 was never true
244 ml_model.active_step = 1
245 ml_model.completed_steps = [1]
246 ml_model.is_resetting = False
247 ml_model.is_updating = True
248 update_fields = [
249 "active_step",
250 "return_step",
251 "completed_steps",
252 "is_resetting",
253 "is_updating",
254 ]
255 elif transition == "cancel":
256 if restore_ml_model_update_snapshot(ml_model): 256 ↛ 259line 256 didn't jump to line 259 because the condition on line 256 was always true
257 ml_model.refresh_from_db()
258 return Response(self.get_serializer(ml_model).data, status=200)
259 ml_model.active_step = 2
260 ml_model.completed_steps = [0, 1, 2]
261 ml_model.is_resetting = False
262 ml_model.is_updating = False
263 update_fields = [
264 "active_step",
265 "return_step",
266 "completed_steps",
267 "is_resetting",
268 "is_updating",
269 ]
270 elif transition == "complete": 270 ↛ 275line 270 didn't jump to line 275 because the condition on line 270 was always true
271 ml_model.is_resetting = False
272 ml_model.is_updating = False
273 update_fields = ["is_resetting", "is_updating"]
274 else:
275 requested_step = transition_data["active_step"]
276 available_steps = set(ml_model.completed_steps or []) | {
277 ml_model.active_step
278 }
279 if requested_step not in available_steps:
280 raise ValidationError({"active_step": "That wizard step is not available."})
281 ml_model.active_step = requested_step
282 update_fields = ["active_step", "return_step"]
284 ml_model.save(update_fields=update_fields)
285 return Response(self.get_serializer(ml_model).data, status=200)
288 @extend_schema(
289 parameters=[
290 OpenApiParameter(
291 name="simulationObject", required=True, type=OpenApiTypes.INT
292 ),
293 ]
294 )
296 def list(self, request, *args, **kwargs):
297 queryset = MLModel.objects.all().filter(simulationObject=self.request.query_params.get("simulationObject", None))
298 serializer = self.get_serializer(queryset, many=True)
299 return Response(serializer.data, status=200)
301 @extend_schema(
302 request=CreateMLModelSerializer,
303 responses=MLModelSerializer,
304 )
305 def create(self, request, *args, **kwargs):
306 """Create an ML model from a completed object-storage upload session."""
307 serializer = CreateMLModelSerializer(data=request.data)
308 serializer.is_valid(raise_exception=True)
309 validated_data = serializer.validated_data
311 simulation_object_id = validated_data["simulationObject"]
312 surrogate_model = validated_data.get("surrogate_model", {})
313 try:
314 simulation_object = SimulationObject.objects.get(id=simulation_object_id)
315 except SimulationObject.DoesNotExist:
316 raise NotFound({"error": "SimulationObject not found."})
318 self._validate_single_model_rule(simulation_object)
320 instance = MLModel.objects.create(
321 flowsheet_state=simulation_object.flowsheet_state,
322 simulationObject=simulation_object,
323 surrogate_model=surrogate_model,
324 active_step=1 if surrogate_model else 0,
325 )
327 response_serializer = self.get_serializer(instance)
328 return Response(
329 response_serializer.data,
330 status=201,
331 )
334 @extend_schema(
335 parameters=[
336 OpenApiParameter(name="model", required=True,
337 type=OpenApiTypes.INT),
338 ],
339 responses=GetCsvHeaderSerializer,
340 )
341 @action(
342 detail=False,
343 methods=["get"],
344 url_path="get-csv-header",
345 url_name="get-csv-header",
346 )
347 def get_csv_header(self, request):
348 """Return the authoritative CSV headers for the selected ML model."""
349 model = self.request.query_params.get("model")
350 if not model: 350 ↛ 351line 350 didn't jump to line 351 because the condition on line 350 was never true
351 raise ValidationError({"error": "model is required."})
353 try:
354 model = MLModel.objects.get(id=model)
355 except MLModel.DoesNotExist:
356 raise NotFound({"error": "MLModel not found."})
358 headers = model.csv_headers
359 if headers:
360 return Response(GetCsvHeaderSerializer({"headers": headers}).data, status=200)
362 if model.surrogate_model != {}: 362 ↛ 370line 362 didn't jump to line 370 because the condition on line 362 was always true
363 input_labels = model.surrogate_model.get("input_labels")
364 output_labels = model.surrogate_model.get("output_labels")
365 return Response(
366 GetCsvHeaderSerializer({"headers": input_labels + output_labels}).data,
367 status=200,
368 )
370 if not headers and model.csv_upload_session_id:
371 inspection = inspect_upload_session(model.csv_upload_session)
372 headers = inspection.headers
373 model.csv_headers = headers
374 model.csv_delimiter = inspection.delimiter
375 model.save(update_fields=["csv_headers", "csv_delimiter"])
376 if not headers:
377 raise ValidationError({"error": "No CSV headers found for this ML model."})
378 return Response(GetCsvHeaderSerializer({"headers": headers}).data, status=200)
380 @extend_schema(request=UploadModelSerializer, responses=MLModelSerializer)
381 @action(
382 detail=False,
383 methods=["post"],
384 url_path="upload-ml-model",
385 url_name="upload-ml-model",
386 )
387 @transaction.atomic
388 def upload_model(self, request):
389 """Replace an existing ML model with an imported surrogate and mapping state."""
390 serializer = UploadModelSerializer(data=request.data)
391 serializer.is_valid(raise_exception=True)
392 validated_data = serializer.validated_data
393 model_id = validated_data["model"]
395 try:
396 ml_model = MLModel.objects.select_for_update().get(id=model_id)
397 except MLModel.DoesNotExist:
398 raise NotFound({"model": "ML model not found."})
400 clear_ml_model_column_mappings(ml_model)
401 ml_model.surrogate_model = validated_data["json_data"]
402 ml_model.csv_file_name = ""
403 ml_model.csv_bucket = ""
404 ml_model.csv_object_key = ""
405 ml_model.csv_headers = []
406 ml_model.csv_delimiter = ""
407 ml_model.csv_upload_session = None
408 ml_model.charts = []
409 ml_model.metrics = []
410 ml_model.test_results_bucket = ""
411 ml_model.test_results_key = ""
412 ml_model.mapping_update_snapshot = None
413 ml_model.active_step = 1
414 ml_model.completed_steps = [0]
415 ml_model.result_state = MLModel.ResultState.PENDING
416 ml_model.is_resetting = False
417 ml_model.is_updating = False
418 ml_model.save(
419 update_fields=[
420 "surrogate_model",
421 "csv_file_name",
422 "csv_bucket",
423 "csv_object_key",
424 "csv_headers",
425 "csv_delimiter",
426 "csv_upload_session",
427 "charts",
428 "metrics",
429 "test_results_bucket",
430 "test_results_key",
431 "mapping_update_snapshot",
432 "active_step",
433 "return_step",
434 "completed_steps",
435 "result_state",
436 "is_resetting",
437 "is_updating",
438 ]
439 )
441 return Response(self.get_serializer(ml_model).data, status=200)
443 @extend_schema(request=CreateSurrogateModelFromColumnSerializer, responses=None)
444 @action(detail=False, methods=["post"], url_path="create-surrogate-model")
445 def create_surrogate_model(self, request):
446 """Start surrogate-model training once the column mappings are complete."""
447 serializer = CreateSurrogateModelFromColumnSerializer(
448 data=request.data)
449 serializer.is_valid(raise_exception=True)
450 validated_data = serializer.validated_data
451 model = validated_data.get("model")
452 model_instance = get_object_or_404(MLModel.objects, id=model)
454 # if already have one (imported model), skip training
455 if model_instance.surrogate_model != {}: 455 ↛ 456line 455 didn't jump to line 456 because the condition on line 455 was never true
456 model_instance.return_step = model_instance.active_step
457 model_instance.active_step = 2
458 model_instance.result_state = MLModel.ResultState.READY_WITHOUT_DIAGNOSTICS
459 model_instance.save(
460 update_fields=["active_step", "return_step", "result_state"]
461 )
462 return Response(
463 OperationMessageSerializer({"message": "successfully trained"}).data,
464 status=200,
465 )
467 return train(request.user, model_instance)
469 @extend_schema(
470 parameters=[
471 OpenApiParameter(name="model", required=True,
472 type=OpenApiTypes.INT),
473 ]
474 )
475 @action(
476 detail=False,
477 methods=["get"],
478 url_path="export-ml-model",
479 url_name="export-ml-model",
480 )
481 def export_flowsheet(self, request):
482 """Export the serialized surrogate model as a downloadable JSON file."""
483 model_id = request.query_params.get("model")
484 if not model_id: 484 ↛ 485line 484 didn't jump to line 485 because the condition on line 484 was never true
485 raise ValidationError({"error": "model_id parameter is required."})
487 try:
488 ml_model = MLModel.objects.get(id=model_id)
489 except MLModel.DoesNotExist:
490 raise NotFound({"error": "MLModel not found."})
492 data = ml_model.surrogate_model
493 response = HttpResponse(
494 json.dumps(data, indent=4), content_type="application/json"
495 )
496 response["Content-Disposition"] = f'attachment; filename="model.json"'
498 return response
500 @extend_schema(
501 parameters=[
502 OpenApiParameter(name="model", required=True, type=OpenApiTypes.INT),
503 ],
504 responses=DownloadTestResultsSerializer,
505 )
506 @action(
507 detail=False,
508 methods=["get"],
509 url_path="download-test-results",
510 url_name="download-test-results",
511 )
512 def download_test_results(self, request):
513 """Return a presigned URL for downloading the full ML test-results CSV."""
514 raw_model_id = request.query_params.get("model")
515 if not raw_model_id: 515 ↛ 516line 515 didn't jump to line 516 because the condition on line 515 was never true
516 raise ValidationError({"error": "model parameter is required."})
518 try:
519 model_id = int(raw_model_id)
520 except (TypeError, ValueError) as exc:
521 raise ValidationError({"error": "model parameter must be an integer."}) from exc
523 try:
524 ml_model = MLModel.objects.get(id=model_id)
525 except MLModel.DoesNotExist as exc:
526 raise NotFound({"error": "MLModel not found."}) from exc
528 if not ml_model.test_results_bucket or not ml_model.test_results_key:
529 raise ValidationError({"error": "No test results available for this model."})
531 url = presign_download_url(
532 bucket=ml_model.test_results_bucket,
533 key=ml_model.test_results_key,
534 filename=f"test-results-{ml_model.id}.csv",
535 expires_seconds=3600,
536 )
537 return Response(DownloadTestResultsSerializer({"url": url}).data, status=200)
540 def destroy(self,request, *args, **kwargs):
541 # Delete custom properties associated with the ML model
542 # This is necessary to avoid leaving orphaned custom properties that reference the deleted ML model
543 try:
544 ml_model: MLModel = self.get_object()
545 for columnMapping in ml_model.MLColumnMappings.all():
546 if columnMapping.portIndex == -1: # custom property 546 ↛ 545line 546 didn't jump to line 545 because the condition on line 546 was always true
547 columnMapping.propertyInfo.delete() # delete the custom property
549 except MLModel.DoesNotExist:
550 raise NotFound({"error": "MLModel not found."})
551 return super().destroy(request, *args, **kwargs)
553@extend_schema(exclude=True)
554@api_view(["POST"])
555@authentication_classes([DaprApiTokenAuthentication])
556@permission_classes([IsAuthenticated])
557@csrf_exempt
558def process_ml_training_event(request) -> Response:
559 """Handle Dapr-delivered ML completion events and update the stored task/model."""
560 training_response = _parse_ml_training_completion_event(request.data)
561 if training_response is None: 561 ↛ 562line 561 didn't jump to line 562 because the condition on line 561 was never true
562 return Response(status=200)
564 process_ml_training_response(training_response.data)
565 return Response(status=200)