SIENTIAPDE-1249: Implement data models for Model Manager and add unit tests. This commit introduces data transfer objects (DTOs) and model classes for experiment status, training parameters, and training results, along with corresponding unit tests to ensure their correct behavior.

This commit is contained in:
Bruno Domingues
2025-10-06 16:49:18 -03:00
parent 18ccfed78a
commit d8583cdae7
9 changed files with 940 additions and 0 deletions

View File

@@ -0,0 +1,32 @@
from enum import Enum
class ExperimentStatus(str, Enum):
"""
Status values for experiment run lifecycle.
This enum defines all possible status values that an experiment run can have
throughout its lifecycle, from initialization through training, model saving,
and cleanup. These statuses are used to track progress and identify failures
in the training pipeline.
The status values follow the naming convention from the original Mage pipeline
to maintain compatibility with existing database records and monitoring systems.
Attributes:
MAGE_WAITING_PROC: Initial status indicating experiment is registered and waiting for processing.
TRAINING_SUCCESS: Training completed successfully with model and metrics calculated.
TRAINING_ERROR: Training failed due to data issues, model errors, or other exceptions.
MLFLOW_SENT: Model successfully saved to MLFlow.
MLFLOW_SEND_ERROR: Model saving to MLFlow failed due to connection or serialization errors.
FILE_DELETED: Cleanup completed successfully with all artifacts removed.
FILE_DELETE_ERROR: Cleanup failed due to file system or MinIO errors.
"""
MAGE_WAITING_PROC = 'MAGE_WAITING_PROC'
TRAINING_SUCCESS = 'TRAINING_SUCCESS'
TRAINING_ERROR = 'TRAINING_ERROR'
MLFLOW_SENT = 'MLFLOW_SENT'
MLFLOW_SEND_ERROR = 'MLFLOW_SEND_ERROR'
FILE_DELETED = 'FILE_DELETED'
FILE_DELETE_ERROR = 'FILE_DELETE_ERROR'