diff --git a/model_manager/activities/experiment_tracking.py b/model_manager/activities/experiment_tracking.py index 26ff2db..027b69c 100644 --- a/model_manager/activities/experiment_tracking.py +++ b/model_manager/activities/experiment_tracking.py @@ -110,6 +110,17 @@ class ExperimentTracking(Postgres): pass async def _execute_update(self, query: str, params: Mapping[str, Any]) -> dict[str, Any]: + """ + Execute an UPDATE SQL statement asynchronously. + + Args: + query: Parameterized SQL string to execute. + params: Mapping of parameters for the SQL query. + + Returns: + dict: A dictionary containing the affected row count: {'rowcount': int}. + """ + def _run() -> dict[str, Any]: with self.engine.begin() as connection: result = connection.execute(text(query), params) diff --git a/model_manager/activities/training.py b/model_manager/activities/training.py index 6a6272c..89bd537 100644 --- a/model_manager/activities/training.py +++ b/model_manager/activities/training.py @@ -113,32 +113,23 @@ class Training(BaseActivity): """ Train a machine learning model. - This activity orchestrates the complete ML training pipeline: - 1. Validates input parameters - 2. Trains the model using TrainingRepository - 3. Performs post-training calculations + This activity orchestrates the ML training pipeline: + 1. Validate input parameters. + 2. Train the model via TrainingRepository. + 3. Perform post-training calculations. Args: - input_data: Configuration for model training operation - Required keys: - - metadata (dict): Workflow execution metadata - - uploaded_file (BytesIO): Training data file (already downloaded from MinIO) - - train_params (TrainModelParams): Training parameters object + input_data: Training configuration containing: + - metadata (dict): Workflow execution metadata. + - uploaded_file (BytesIO): Training data already downloaded from MinIO. + - train_params (TrainModelParams | dict): Training parameters. Returns: - TrainModelResult: Training result with model, metrics, and data + dict: Keys `run_name` and `run_dir` when training and saving succeed. Raises: - ValueError: If input validation fails - Exception: If training fails (after sending notification) - - Example: - result = await train_model({ - 'metadata': {'workflow_id': 'train-123', 'experiment_run_id': 456}, - 'uploaded_file': BytesIO(csv_data), - 'train_params': TrainModelParams(...) - }) - # Returns: TrainModelResult(...) + ValueError: If input validation fails. + Exception: If training fails (after sending notification). """ metadata = input_data.get('metadata', {}) train_params = input_data['train_params'] @@ -195,6 +186,19 @@ class Training(BaseActivity): @activity.defn(name='cleanup_resources') async def cleanup_resources(self, input_data: dict[str, Any]) -> None: + """ + Cleanup temporary resources created during training. + + Args: + input_data: Cleanup configuration containing: + - metadata (dict): Workflow execution metadata. + - run_dir (str): Temporary directory to remove. + - bucket_name (str): MinIO bucket of the uploaded file. + - file_name (str): MinIO object key to delete. + + Raises: + Exception: If cleanup fails (after sending notification). + """ metadata = input_data.get('metadata', {}) run_dir = input_data.get('run_dir', '') bucket_name = input_data.get('bucket_name', '') diff --git a/model_manager/sientia/model_serving.py b/model_manager/sientia/model_serving.py index c5bcc03..fe3b560 100644 --- a/model_manager/sientia/model_serving.py +++ b/model_manager/sientia/model_serving.py @@ -66,8 +66,8 @@ class ModelServing: Returns: pandas.DataFrame: A DataFrame containing run information. - Raise: - SientiaMlException if unable to search runs + Raises: + SientiaMlException: If unable to search runs. """ try: runs = mlflow.search_runs(experiment_names=experiment_names, order_by=order_by) @@ -82,6 +82,9 @@ class ModelServing: Args: experiment_identifier (str): name or id of the experiment to be setted + + Raises: + Exception: If setting the experiment fails. """ mlflow.set_experiment(experiment_identifier) @@ -99,6 +102,9 @@ class ModelServing: Security Warning: The GitHub token is hardcoded. Consider moving to environment variable or using a secure secret management solution (e.g., K8s secrets). + + Raises: + Exception: If logging the model fails. """ mlflow.sklearn.log_model( sk_model, @@ -117,6 +123,9 @@ class ModelServing: Returns: None + + Raises: + Exception: If logging the parameter fails. """ mlflow.log_param(key, value) @@ -130,6 +139,9 @@ class ModelServing: Returns: None + + Raises: + Exception: If logging the metric fails. """ mlflow.log_metric(key, value) @@ -146,6 +158,9 @@ class ModelServing: Returns: None + + Raises: + Exception: If logging the artifact fails. """ mlflow.log_artifact(local_path=local_path, artifact_path=artifact_path, run_id=run_id) @@ -178,11 +193,8 @@ class ModelServing: Yields: ActiveRun: object that acts as a context manager wrapping the run's state. - Example: - with model_serving.save_experiment(run_name="my_run") as run: - model_serving.log_param("param1", value1) - model_serving.log_metric("metric1", value2) - # Run is automatically closed here, even if an exception occurs + Raises: + Exception: If starting or ending the MLflow run fails. """ run = mlflow.start_run( run_id=run_id, diff --git a/model_manager/utils/connectors_config.py b/model_manager/utils/connectors_config.py index 4f7a9b8..b42d017 100644 --- a/model_manager/utils/connectors_config.py +++ b/model_manager/utils/connectors_config.py @@ -42,8 +42,7 @@ def build_mlflow_config() -> dict[str, Any]: It handles server connection and authentication parameters. Environment Variables: - MLFLOW_HOST: MLFlow server hostname (default: http://localhost) - MLFLOW_PORT: MLFlow server port (default: 5080) + MLFLOW_URL: MLFlow server hostname (default: http://localhost:5080) MLFLOW_USERNAME: MLFlow username (default: aignosi) MLFLOW_PASSWORD: MLFlow password (default: aignosi) diff --git a/model_manager/utils/exceptions.py b/model_manager/utils/exceptions.py index eb1f775..8f9c4a8 100644 --- a/model_manager/utils/exceptions.py +++ b/model_manager/utils/exceptions.py @@ -1,5 +1,31 @@ +""" +Custom exception types for the Model Manager. + +This module defines domain-specific exceptions used across the training +workflow to convey additional context (e.g., flags indicating which steps +completed successfully) without altering control flow semantics. +""" + + class ModelTrainingError(Exception): + """ + Exception raised when the model training workflow fails. + + This exception carries flags indicating whether the model was trained + and/or saved successfully, enabling the workflow to map errors to + appropriate experiment statuses. + """ + def __init__(self, model_trained: bool, model_saved: bool, message: str | None = None): + """ + Initialize ModelTrainingError with training state flags. + + Args: + model_trained: True if the training step completed successfully. + model_saved: True if the model saving step completed successfully. + message: Optional custom error message. If None, a default message + including the state flags is generated. + """ self.model_trained = model_trained self.model_saved = model_saved diff --git a/model_manager/utils/models/experiment_status.py b/model_manager/utils/models/experiment_status.py index 2cfcbca..b7975eb 100644 --- a/model_manager/utils/models/experiment_status.py +++ b/model_manager/utils/models/experiment_status.py @@ -14,6 +14,7 @@ class ExperimentStatus(str, Enum): to maintain compatibility with existing database records and monitoring systems. Attributes: + ORCHESTRATOR_VALIDATION_ERROR: Error in the parameters validation. 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. diff --git a/model_manager/utils/models/train_model_params.py b/model_manager/utils/models/train_model_params.py index 34109a7..790d8ec 100644 --- a/model_manager/utils/models/train_model_params.py +++ b/model_manager/utils/models/train_model_params.py @@ -78,14 +78,6 @@ class TrainModelParams: ValueError: If any required field is missing or None TypeError: If any field has an incorrect type KeyError: If any required key is missing from the dictionary - - Example: - >>> input_data = { - ... 'variable_columns': ['var1', 'var2'], - ... 'lag_train': 5, - ... # ... other fields - ... } - >>> params = TrainModelParams.from_dict(input_data) """ return cls( variable_columns=cls._check_none( @@ -179,10 +171,6 @@ class TrainModelParams: Raises: ValueError: If any business rule is violated - - Example: - >>> params = TrainModelParams.from_dict(data) - >>> params.validate_business_rules() # Raises ValueError if invalid """ # Validate train_size range (10-100%) if not 10 <= self.train_size <= 100: diff --git a/todo-list.txt b/todo-list.txt index b1f6ae9..45c2558 100644 --- a/todo-list.txt +++ b/todo-list.txt @@ -1,3 +1,4 @@ - remover os testes das classes alteradas e refazer de novo depois -- alterar todos os comentários dos métodos que foram alterados - no final alterar o readme +- verificar erro de logs duplicados +- Remover do log do metadata o model name e id